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 from app.services.user_quota_service import UserQuotaService 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, "by_day": summary.by_day, } @router.get("/quota") async def get_quota( current_user: User = Depends(get_current_user), ): tracker = get_usage_tracker() if tracker._session: quota_service = UserQuotaService(session=tracker._session) return await quota_service.check_quota_with_plan( user_id=_user_id(current_user), user_plan=current_user.plan, ) return 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}