From cba9e57db9f836739783154703c30826ea49fa05 Mon Sep 17 00:00:00 2001 From: nickviljoen Date: Tue, 24 Feb 2026 09:35:26 +0200 Subject: [PATCH] Add prompt_tokens and completion_tokens to collector API and UI Supports the token breakdown (prompt/completion) now sent by agent-sync from LibreChat's transactions collection. Updates models, collector endpoint, usage stats, CSV export, and agent card/modal display. Co-Authored-By: Claude Opus 4.6 --- CLAUDE.md | 15 ++++++++------- crud.py | 4 ++++ main.py | 20 +++++++++++++++++--- models.py | 6 ++++++ templates/agent_management.html | 7 ++++++- 5 files changed, 41 insertions(+), 11 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index cfa96b0..a746c55 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -60,9 +60,9 @@ Create `.env` file with: - `AiAgent`: Core agent model with comprehensive fields (includes `discipline`, `rating`) - `UsageTimelineEntry`: Daily usage data including `message_count` and `token_count` - `UserCreate/UserResponse`: User management models -- `AiAgentCreate/AiAgentResponse`: API request/response models (includes `total_tokens`, `discipline`, `rating`, `rating_count`) -- `AgentCollectorCreate`: Collector API input model (includes `total_tokens`, `discipline`) -- `AgentUsageStatsResponse`: Usage statistics response (includes `total_tokens`) +- `AiAgentCreate/AiAgentResponse`: API request/response models (includes `total_tokens`, `prompt_tokens`, `completion_tokens`, `discipline`, `rating`, `rating_count`) +- `AgentCollectorCreate`: Collector API input model (includes `total_tokens`, `prompt_tokens`, `completion_tokens`, `discipline`) +- `AgentUsageStatsResponse`: Usage statistics response (includes `total_tokens`, `prompt_tokens`, `completion_tokens`) **crud.py**: Database operations using Motor (async MongoDB driver): - User CRUD: authentication, creation, management @@ -142,13 +142,14 @@ Located in `templates/` directory: - Discipline passed through collector API; rating is human-only (not in collector) ### Token Usage Tracking -- `total_tokens` field on agents tracks cumulative LLM token consumption +- `total_tokens`, `prompt_tokens`, `completion_tokens` fields on agents track cumulative LLM token consumption +- `prompt_tokens` (input) and `completion_tokens` (output) provide cost breakdown detail - `token_count` per day in usage timeline entries alongside `message_count` -- Token badge displayed on agent cards (gold/coins icon) -- Usage modal shows Total Tokens stat alongside messages/conversations/users +- Token badge displayed on agent cards (gold/coins icon) with prompt/completion breakdown in tooltip +- Usage modal shows Total Tokens stat with In/Out breakdown alongside messages/conversations/users - Dual-axis chart (messages left axis, tokens right axis) when token data exists - Sort agents by Total Tokens -- CSV export includes `total_tokens` column +- CSV export includes `total_tokens`, `prompt_tokens`, `completion_tokens` columns ### High Usage Email Notifications - Entirely optional — silently disabled when Mailgun env vars are not set diff --git a/crud.py b/crud.py index 2ac7f1b..3f3bb8f 100644 --- a/crud.py +++ b/crud.py @@ -157,6 +157,10 @@ async def create_agent(agent_data: dict, user_id: str, skip_time_check: bool = F agent_doc["rating"] = None if "total_tokens" not in agent_doc: agent_doc["total_tokens"] = None + if "prompt_tokens" not in agent_doc: + agent_doc["prompt_tokens"] = None + if "completion_tokens" not in agent_doc: + agent_doc["completion_tokens"] = None result = await agents_collection.insert_one(agent_doc) agent_doc["_id"] = result.inserted_id return agent_doc diff --git a/main.py b/main.py index 948049d..f7df4c9 100644 --- a/main.py +++ b/main.py @@ -236,7 +236,9 @@ def create_agent_response(agent: dict) -> models.AiAgentResponse: total_messages=agent.get("total_messages"), first_used=agent.get("first_used"), last_used=agent.get("last_used"), - total_tokens=agent.get("total_tokens") + total_tokens=agent.get("total_tokens"), + prompt_tokens=agent.get("prompt_tokens"), + completion_tokens=agent.get("completion_tokens") ) async def verify_agent_collector_api_key(x_api_key: str = Header(alias="X-API-Key")): @@ -296,6 +298,8 @@ def map_agent_collector_to_internal(collector_data: models.AgentCollectorCreate) "first_used": collector_data.first_used, "last_used": collector_data.last_used, "total_tokens": collector_data.total_tokens, + "prompt_tokens": collector_data.prompt_tokens, + "completion_tokens": collector_data.completion_tokens, } @@ -1150,6 +1154,8 @@ async def export_agents_csv(current_user: dict = Depends(require_admin)): "last_edited_by", "created_by", "total_tokens", + "prompt_tokens", + "completion_tokens", "discipline", "rating", "rating_count" @@ -1188,6 +1194,8 @@ async def export_agents_csv(current_user: dict = Depends(require_admin)): "last_edited_by": agent.get("last_edited_by", ""), "created_by": agent.get("created_by", ""), "total_tokens": str(agent.get("total_tokens", "")) if agent.get("total_tokens") is not None else "", + "prompt_tokens": str(agent.get("prompt_tokens", "")) if agent.get("prompt_tokens") is not None else "", + "completion_tokens": str(agent.get("completion_tokens", "")) if agent.get("completion_tokens") is not None else "", "discipline": agent.get("discipline", ""), "rating": str(agent.get("rating", "")) if agent.get("rating") is not None else "", "rating_count": str(agent.get("rating_count", "")) if agent.get("rating_count") is not None else "" @@ -1403,6 +1411,8 @@ async def create_agent_collector( "first_used": internal_data.get("first_used"), "last_used": internal_data.get("last_used"), "total_tokens": internal_data.get("total_tokens"), + "prompt_tokens": internal_data.get("prompt_tokens"), + "completion_tokens": internal_data.get("completion_tokens"), "updated_at": datetime.utcnow() } @@ -1518,7 +1528,9 @@ async def get_agent_usage( usage_by_period=usage_by_period, conversation_count=agent.get("conversation_count"), unique_users=agent.get("unique_users"), - total_tokens=agent.get("total_tokens") + total_tokens=agent.get("total_tokens"), + prompt_tokens=agent.get("prompt_tokens"), + completion_tokens=agent.get("completion_tokens") ) else: # FALLBACK TO CALCULATED DATA (old system) @@ -1544,7 +1556,9 @@ async def get_agent_usage( usage_by_period=usage_by_period, conversation_count=None, # Old system doesn't track this unique_users=None, # Old system doesn't track this - total_tokens=None # Old system doesn't track this + total_tokens=None, # Old system doesn't track this + prompt_tokens=None, + completion_tokens=None ) except HTTPException: raise diff --git a/models.py b/models.py index 212f5d2..b61ea03 100644 --- a/models.py +++ b/models.py @@ -140,6 +140,8 @@ class AiAgentResponse(BaseModel): first_used: Optional[str] = None last_used: Optional[str] = None total_tokens: Optional[int] = None + prompt_tokens: Optional[int] = None + completion_tokens: Optional[int] = None # Agent Collector API Models (for compatibility with agent_collector app) class AgentCollectorCreate(BaseModel): @@ -169,6 +171,8 @@ class AgentCollectorCreate(BaseModel): first_used: Optional[str] = None # ISO 8601 datetime string last_used: Optional[str] = None # ISO 8601 datetime string total_tokens: Optional[int] = Field(default=None, ge=0) + prompt_tokens: Optional[int] = Field(default=None, ge=0) + completion_tokens: Optional[int] = Field(default=None, ge=0) class AgentCollectorResponse(BaseModel): status: str = "success" @@ -201,3 +205,5 @@ class AgentUsageStatsResponse(BaseModel): conversation_count: Optional[int] = None unique_users: Optional[int] = None total_tokens: Optional[int] = None + prompt_tokens: Optional[int] = None + completion_tokens: Optional[int] = None diff --git a/templates/agent_management.html b/templates/agent_management.html index 3b2fc6e..5e151cc 100644 --- a/templates/agent_management.html +++ b/templates/agent_management.html @@ -722,7 +722,7 @@ function displayAgents(agentsToShow) { ${agent.quality_audit_status ? '' : ''} ${agent.quality_audit_status && agent.risk_factor ? getRiskFactorBadge(agent.risk_factor) : ''} ${agent.total_messages ? `${agent.total_messages.toLocaleString()}` : ''} - ${agent.total_tokens ? `${agent.total_tokens.toLocaleString()}` : ''} + ${agent.total_tokens ? `${agent.total_tokens.toLocaleString()}` : ''} ${agent.unique_users ? `${agent.unique_users}` : ''} ${agent.url ? `Open Agent` : ''} @@ -1389,6 +1389,11 @@ async function loadUsageChart(agentName, period = 'daily', startDate = null, end
${stats.total_tokens != null ? stats.total_tokens.toLocaleString() : 'N/A'}
Total Tokens + ${stats.prompt_tokens != null || stats.completion_tokens != null ? ` +
+ In: ${stats.prompt_tokens != null ? stats.prompt_tokens.toLocaleString() : '—'} | + Out: ${stats.completion_tokens != null ? stats.completion_tokens.toLocaleString() : '—'} +
` : ''}