74 lines
1.8 KiB
Python
74 lines
1.8 KiB
Python
import logging
|
|
|
|
from fastapi import APIRouter, Depends, Query
|
|
|
|
from app.api.deps import get_current_user
|
|
from app.models.user import User
|
|
from app.services.usage_tracker import UsageTracker
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
_usage_tracker: UsageTracker = UsageTracker()
|
|
|
|
|
|
def set_usage_tracker(tracker: UsageTracker) -> None:
|
|
global _usage_tracker
|
|
_usage_tracker = tracker
|
|
|
|
|
|
def get_usage_tracker() -> UsageTracker:
|
|
return _usage_tracker
|
|
|
|
|
|
def _user_id(current_user: User) -> str:
|
|
return str(current_user.id)
|
|
|
|
|
|
@router.get("/summary")
|
|
async def get_usage_summary(
|
|
period: str = Query("month", pattern=r"^(day|week|month)$"),
|
|
current_user: User = Depends(get_current_user),
|
|
):
|
|
summary = get_usage_tracker().get_summary(
|
|
user_id=_user_id(current_user),
|
|
period=period,
|
|
)
|
|
return {
|
|
"period": summary.period,
|
|
"start_date": summary.start_date,
|
|
"end_date": summary.end_date,
|
|
"total_queries": summary.total_queries,
|
|
"total_input_tokens": summary.total_input_tokens,
|
|
"total_output_tokens": summary.total_output_tokens,
|
|
"total_cost": summary.total_cost,
|
|
"by_engine": summary.by_engine,
|
|
}
|
|
|
|
|
|
@router.get("/quota")
|
|
async def get_quota(
|
|
current_user: User = Depends(get_current_user),
|
|
):
|
|
return get_usage_tracker().check_quota(user_id=_user_id(current_user))
|
|
|
|
|
|
@router.get("/by-engine")
|
|
async def get_usage_by_engine(
|
|
current_user: User = Depends(get_current_user),
|
|
):
|
|
summary = get_usage_tracker().get_summary(
|
|
user_id=_user_id(current_user),
|
|
period="month",
|
|
)
|
|
engines = [
|
|
{
|
|
"type": engine_type,
|
|
"queries": data["queries"],
|
|
"cost": data["cost"],
|
|
}
|
|
for engine_type, data in summary.by_engine.items()
|
|
]
|
|
return {"engines": engines}
|