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 <noreply@anthropic.com>
This commit is contained in:
nickviljoen
parent
62138e9142
commit
cba9e57db9
5 changed files with 41 additions and 11 deletions
15
CLAUDE.md
15
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
|
||||
|
|
|
|||
4
crud.py
4
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
|
||||
|
|
|
|||
20
main.py
20
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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -722,7 +722,7 @@ function displayAgents(agentsToShow) {
|
|||
${agent.quality_audit_status ? '<span class="badge bg-success" title="Quality Audited"><i class="fas fa-certificate"></i></span>' : ''}
|
||||
${agent.quality_audit_status && agent.risk_factor ? getRiskFactorBadge(agent.risk_factor) : ''}
|
||||
${agent.total_messages ? `<span class="badge bg-info text-white" title="Total Messages"><i class="fas fa-comments me-1"></i>${agent.total_messages.toLocaleString()}</span>` : ''}
|
||||
${agent.total_tokens ? `<span class="badge bg-warning text-dark" title="Total Tokens"><i class="fas fa-coins me-1"></i>${agent.total_tokens.toLocaleString()}</span>` : ''}
|
||||
${agent.total_tokens ? `<span class="badge bg-warning text-dark" title="Total Tokens${agent.prompt_tokens != null ? '\nPrompt: ' + agent.prompt_tokens.toLocaleString() : ''}${agent.completion_tokens != null ? '\nCompletion: ' + agent.completion_tokens.toLocaleString() : ''}"><i class="fas fa-coins me-1"></i>${agent.total_tokens.toLocaleString()}</span>` : ''}
|
||||
${agent.unique_users ? `<span class="badge bg-primary" title="Unique Users"><i class="fas fa-users me-1"></i>${agent.unique_users}</span>` : ''}
|
||||
${agent.url ? `<a href="${agent.url}" target="_blank" class="btn btn-sm btn-success" onclick="event.stopPropagation();" title="Start a conversation with this agent"><i class="fas fa-external-link-alt me-1"></i>Open Agent</a>` : ''}
|
||||
</div>
|
||||
|
|
@ -1389,6 +1389,11 @@ async function loadUsageChart(agentName, period = 'daily', startDate = null, end
|
|||
<div class="border rounded p-2">
|
||||
<h6 class="mb-1 text-warning">${stats.total_tokens != null ? stats.total_tokens.toLocaleString() : 'N/A'}</h6>
|
||||
<small class="text-muted">Total Tokens</small>
|
||||
${stats.prompt_tokens != null || stats.completion_tokens != null ? `
|
||||
<div class="mt-1" style="font-size: 0.7rem;">
|
||||
<span class="text-muted">In: ${stats.prompt_tokens != null ? stats.prompt_tokens.toLocaleString() : '—'}</span> |
|
||||
<span class="text-muted">Out: ${stats.completion_tokens != null ? stats.completion_tokens.toLocaleString() : '—'}</span>
|
||||
</div>` : ''}
|
||||
</div>
|
||||
</div>
|
||||
<div class="col">
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue