feat: add token usage facade (#377)

- 增加 token 用量统计接口

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$
+;
+
+

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]

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

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()

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 {}
+  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)