From 070a3163d28d43b753ec6997aa70bd05c4bec5cd Mon Sep 17 00:00:00 2001 From: "raoha.rh" Date: Tue, 24 Dec 2024 11:01:24 +0800 Subject: [PATCH] feat: add bot token usage analyzer --- .../20241224030938_remote_schema.sql | 24 +++++++++++++++++++ server/core/dao/userTokenUsageDAO.py | 10 +++++++- server/core/models/user_token_usage.py | 8 +++++++ server/core/service/user_token_usage.py | 3 +++ server/event_handler/pull_request.py | 1 - server/user/router.py | 12 ++++++++-- 6 files changed, 54 insertions(+), 4 deletions(-) create mode 100644 migrations/supabase/migrations/20241224030938_remote_schema.sql diff --git a/migrations/supabase/migrations/20241224030938_remote_schema.sql b/migrations/supabase/migrations/20241224030938_remote_schema.sql new file mode 100644 index 00000000..8de95576 --- /dev/null +++ b/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$ +; + + diff --git a/server/core/dao/userTokenUsageDAO.py b/server/core/dao/userTokenUsageDAO.py index 7d200040..ce9024e3 100644 --- a/server/core/dao/userTokenUsageDAO.py +++ b/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] \ No newline at end of file diff --git a/server/core/models/user_token_usage.py b/server/core/models/user_token_usage.py index 3a6cd527..e7479f5f 100644 --- a/server/core/models/user_token_usage.py +++ b/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 diff --git a/server/core/service/user_token_usage.py b/server/core/service/user_token_usage.py index 7413ffdd..4cde43da 100644 --- a/server/core/service/user_token_usage.py +++ b/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() diff --git a/server/event_handler/pull_request.py b/server/event_handler/pull_request.py index 4eeeabe9..d98fc508 100644 --- a/server/event_handler/pull_request.py +++ b/server/event_handler/pull_request.py @@ -12,7 +12,6 @@ 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 ( diff --git a/server/user/router.py b/server/user/router.py index c30a0dca..df34a891 100644 --- a/server/user/router.py +++ b/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) \ No newline at end of file + 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) \ No newline at end of file