Feat/add bot token usage

migrations/supabase/migrations/20241224030938_remote_schema.sql

@@ -0,0 +1,24 @@
+set check_function_bodies = off;
+
+CREATE OR REPLACE FUNCTION public.analyze_user_token_usage(start_date date, end_date date)
+ RETURNS TABLE(bot_id text, usage_date date, input_tokens bigint, output_tokens bigint, total_tokens bigint)
+ LANGUAGE plpgsql
+AS $function$
+BEGIN
+    RETURN QUERY
+    SELECT
+        u.bot_id AS bot_id,
+        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.date >= start_date AND
+        u.date <= end_date
+    GROUP BY u.date, u.bot_id;
+END;
+$function$
+;
+
+

server/core/dao/userTokenUsageDAO.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 from supabase.client import Client
 
-from core.models.user_token_usage import UserTokenUsage, UserTokenUsageStats
+from core.models.user_token_usage import BotTokenUsageStats, UserTokenUsage, UserTokenUsageStats
 from petercat_utils.db.client.supabase import get_client
 from core.dao.BaseDAO import BaseDAO
 
@@ -29,3 +29,11 @@ def stats(self, user_id: str, start_date: datetime, end_date: datetime):
         }).execute()
 
         return [UserTokenUsageStats(**stats) for stats in resp.data]
+
+    def analyze(self, start_date: datetime, end_date: datetime):
+        resp = self.client.rpc("analyze_user_token_usage", {
+            "start_date": start_date.strftime("%Y-%m-%d"),
+            "end_date": end_date.strftime("%Y-%m-%d"),
+        }).execute()
+
+        return [BotTokenUsageStats(**stats) for stats in resp.data]

server/core/models/user_token_usage.py

@@ -15,6 +15,14 @@ class UserTokenUsage(BaseModel):
 
 
 class UserTokenUsageStats(BaseModel):
+    usage_date: datetime
+    input_tokens: Optional[int] = 0
+    output_tokens: Optional[int] = 0
+    total_tokens: Optional[int] = 0
+
+
+class BotTokenUsageStats(BaseModel):
+    bot_id: str
     usage_date: datetime
     input_tokens: Optional[int] = 0
     output_tokens: Optional[int] = 0

server/core/service/user_token_usage.py

@@ -24,6 +24,9 @@ def __init__(self) -> None:
   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 analyze_token_usage(self, start_date: datetime.date, end_date: datetime.date):
+    return self.user_token_usage_dao.analyze(start_date=start_date, end_date=end_date)
+
 def get_user_token_usage_service():
   return UserTokenUsageService()
 

server/event_handler/pull_request.py

@@ -12,8 +12,7 @@
 from agent.bot.get_bot import get_bot_by_id
 from core.models.bot import BotModel
 
-from utils.fuzzy_match import contains_keyword_fuzzy
 from utils.path_to_hunk import convert_patch_to_hunk
 from utils.random_str import random_str
 from agent.prompts.pull_request import (
     get_role_prompt,

server/user/router.py

@@ -1,6 +1,6 @@
 
 
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Annotated
 from fastapi import APIRouter, Depends
 
@@ -65,4 +65,12 @@ def token_usage(
   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)
+  return user_token_usage_service.usage_stats(user_id=user_id, start_date=start_date, end_date=end_date)
+
+@router.get("/llm_token_usages/analyzer")
+def token_usage_analyze(
+  start_date: datetime = datetime.now() - timedelta(days=7),
+  end_date: datetime = datetime.now(),
+  user_token_usage_service: Annotated[UserTokenUsageService | None, Depends(get_user_token_usage_service)] = None,
+):
+  return user_token_usage_service.analyze_token_usage(start_date=start_date, end_date=end_date)