diff --git a/migrations/supabase/migrations/20240912023152_remote_schema.sql b/migrations/supabase/migrations/20240912023152_remote_schema.sql new file mode 100644 index 00000000..d51bfe06 --- /dev/null +++ b/migrations/supabase/migrations/20240912023152_remote_schema.sql @@ -0,0 +1,23 @@ +set check_function_bodies = off; + +CREATE OR REPLACE FUNCTION public.get_user_stats(filter_user_id text, start_date date, end_date date) + RETURNS TABLE(usage_date date, input_tokens bigint, output_tokens bigint, total_tokens bigint) + LANGUAGE plpgsql +AS $function$ +BEGIN + RETURN QUERY + SELECT + u.date AS usage_date, -- 使用别名来避免歧义 + SUM(u.input_token)::BIGINT AS input_tokens, -- 将结果转换为 BIGINT + SUM(u.output_token)::BIGINT AS output_tokens, -- 将结果转换为 BIGINT + SUM(u.total_token)::BIGINT AS total_tokens -- 将结果转换为 BIGINT + FROM user_token_usage u + WHERE u.user_id = filter_user_id -- 使用别名来引用参数 + AND u.date >= start_date + AND u.date <= end_date + GROUP BY u.date; +END; +$function$ +; + + diff --git a/server/core/dao/userTokenUsageDAO.py b/server/core/dao/userTokenUsageDAO.py index c8a6fd11..7d200040 100644 --- a/server/core/dao/userTokenUsageDAO.py +++ b/server/core/dao/userTokenUsageDAO.py @@ -1,6 +1,7 @@ +from datetime import datetime from supabase.client import Client -from core.models.user_token_usage import UserTokenUsage +from core.models.user_token_usage import UserTokenUsage, UserTokenUsageStats from petercat_utils.db.client.supabase import get_client from core.dao.BaseDAO import BaseDAO @@ -18,3 +19,13 @@ def create(self, token_usage: UserTokenUsage): .insert(token_usage.model_dump(exclude=["id"])) .execute() ) + + + def stats(self, user_id: str, start_date: datetime, end_date: datetime): + resp = self.client.rpc("get_user_stats", { + "filter_user_id": user_id, + "start_date": start_date.strftime("%Y-%m-%d"), + "end_date": end_date.strftime("%Y-%m-%d"), + }).execute() + + return [UserTokenUsageStats(**stats) for stats in resp.data] diff --git a/server/core/models/user_token_usage.py b/server/core/models/user_token_usage.py index c5a4312d..3a6cd527 100644 --- a/server/core/models/user_token_usage.py +++ b/server/core/models/user_token_usage.py @@ -12,3 +12,10 @@ class UserTokenUsage(BaseModel): input_token: Optional[int] = 0 output_token: Optional[int] = 0 total_token: Optional[int] = 0 + + +class UserTokenUsageStats(BaseModel): + usage_date: datetime + input_tokens: Optional[int] = 0 + output_tokens: Optional[int] = 0 + total_tokens: Optional[int] = 0 \ No newline at end of file diff --git a/server/core/service/user_token_usage.py b/server/core/service/user_token_usage.py index 2ad528ca..884a3692 100644 --- a/server/core/service/user_token_usage.py +++ b/server/core/service/user_token_usage.py @@ -1,4 +1,5 @@ +from datetime import datetime from typing import AsyncGenerator, Dict, Optional from pydantic import BaseModel @@ -15,6 +16,16 @@ class TokenUsageVO(BaseModel): output_token: Optional[int] = 0 total_token: Optional[int] = 0 +class UserTokenUsageService(): + user_token_usage_dao: UserTokenUsageDAO + def __init__(self) -> None: + self.user_token_usage_dao = UserTokenUsageDAO() + + def usage_stats(self, user_id: str, start_date: datetime.date, end_date: datetime.date): + return self.user_token_usage_dao.stats(user_id=user_id, start_date=start_date, end_date=end_date) + +def get_user_token_usage_service(): + return UserTokenUsageService() def create_token_recorder(user: User, bot: Bot): user_token_usage_dao = UserTokenUsageDAO() diff --git a/server/user/router.py b/server/user/router.py index 1155bba2..c30a0dca 100644 --- a/server/user/router.py +++ b/server/user/router.py @@ -1,11 +1,13 @@ +from datetime import datetime from typing import Annotated from fastapi import APIRouter, Depends from agent.llm import get_registered_llm_client, import_clients from auth.get_user_info import get_user_id from core.service.user_llm_token import UserLLMTokenService, CreateUserLLMTokenVO, get_llm_token_service +from core.service.user_token_usage import UserTokenUsageService, get_user_token_usage_service router = APIRouter( prefix="/api/user", @@ -54,4 +56,13 @@ def delete_token( ): print(f"delete_llm_token, token={token_id}, user_id={user_id}") llm_service.delete_llm_token(id=token_id, user_id=user_id) - return {} \ No newline at end of file + return {} + +@router.get("/llm_token_usages") +def token_usage( + start_date: datetime, + end_date: datetime, + user_id: Annotated[str | None, Depends(get_user_id)] = None, + user_token_usage_service: Annotated[UserTokenUsageService | None, Depends(get_user_token_usage_service)] = None, +): + return user_token_usage_service.usage_stats(user_id=user_id, start_date=start_date, end_date=end_date) \ No newline at end of file