mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Add team/tag tools, status indicators, and improved AI agent
- AI agent now has 3 tools: get_usage_data, get_team_usage_data, get_tag_usage_data - Stream status events (Thinking... Fetching... Analyzing...) to UI - Frontend shows spinner + status text during tool execution - Better system prompt guiding tool selection - Entity summariser for team/tag data with ranked breakdowns - 13 backend tests, 34 frontend tests passing Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>
This commit is contained in:
Cursor Agent
parent
4f63135e1d
commit
b7399de05b
4 changed files with 432 additions and 179 deletions
|
|
@ -1,25 +1,28 @@
|
|||
"""
|
||||
AI Usage Chat - uses LLM tool calling to answer questions about
|
||||
usage/spend data from the /user/daily/activity endpoints.
|
||||
usage/spend data by querying the aggregated daily activity endpoints.
|
||||
"""
|
||||
|
||||
import json
|
||||
from datetime import date, timedelta
|
||||
from typing import Any, AsyncIterator, Dict, List, Optional
|
||||
from typing import Any, AsyncIterator, Dict, List, Optional, Union
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool definitions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
GET_USAGE_DATA_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_usage_data",
|
||||
"description": (
|
||||
"Fetch aggregated usage/spend data for the LiteLLM proxy. "
|
||||
"Fetch aggregated global usage/spend data for the LiteLLM proxy. "
|
||||
"Returns daily spend, token usage, request counts, and breakdowns "
|
||||
"by model, provider, and API key for the given date range. "
|
||||
"Always call this tool first to get data before answering."
|
||||
"Use this for questions about overall spend, top models, top providers, etc."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
|
|
@ -42,28 +45,96 @@ GET_USAGE_DATA_TOOL = {
|
|||
},
|
||||
}
|
||||
|
||||
GET_TEAM_USAGE_DATA_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_team_usage_data",
|
||||
"description": (
|
||||
"Fetch usage/spend data broken down by team. "
|
||||
"Returns each team's spend, requests, tokens, model breakdown, and provider breakdown. "
|
||||
"Use this for questions like 'which team is spending the most' or 'show me team X usage'."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"start_date": {
|
||||
"type": "string",
|
||||
"description": "Start date in YYYY-MM-DD format",
|
||||
},
|
||||
"end_date": {
|
||||
"type": "string",
|
||||
"description": "End date in YYYY-MM-DD format",
|
||||
},
|
||||
"team_ids": {
|
||||
"type": "string",
|
||||
"description": "Optional comma-separated team IDs to filter by. Omit for all teams.",
|
||||
},
|
||||
},
|
||||
"required": ["start_date", "end_date"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
GET_TAG_USAGE_DATA_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_tag_usage_data",
|
||||
"description": (
|
||||
"Fetch usage/spend data broken down by tag. "
|
||||
"Tags are labels attached to requests (e.g. feature names, environments, credentials). "
|
||||
"Use this for questions about tag-level spend or 'top tags for team X'."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"start_date": {
|
||||
"type": "string",
|
||||
"description": "Start date in YYYY-MM-DD format",
|
||||
},
|
||||
"end_date": {
|
||||
"type": "string",
|
||||
"description": "End date in YYYY-MM-DD format",
|
||||
},
|
||||
"tags": {
|
||||
"type": "string",
|
||||
"description": "Optional comma-separated tag names to filter. Omit for all tags.",
|
||||
},
|
||||
},
|
||||
"required": ["start_date", "end_date"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ALL_TOOLS = [GET_USAGE_DATA_TOOL, GET_TEAM_USAGE_DATA_TOOL, GET_TAG_USAGE_DATA_TOOL]
|
||||
|
||||
SYSTEM_PROMPT = (
|
||||
"You are an AI assistant embedded in the LiteLLM Usage dashboard. "
|
||||
"You help users understand their LLM API spend and usage data.\n\n"
|
||||
"You have access to a tool called `get_usage_data` which fetches "
|
||||
"aggregated usage data from the LiteLLM database. You MUST call "
|
||||
"this tool first to retrieve the data, then answer the user's question "
|
||||
"based on the returned data.\n\n"
|
||||
"You have access to these tools:\n"
|
||||
"- `get_usage_data`: Global/user-level usage (spend, models, providers, API keys)\n"
|
||||
"- `get_team_usage_data`: Team-level usage breakdown\n"
|
||||
"- `get_tag_usage_data`: Tag-level usage breakdown\n\n"
|
||||
"ALWAYS call the appropriate tool(s) first to fetch data before answering. "
|
||||
"You may call multiple tools if the question spans different dimensions.\n\n"
|
||||
"Guidelines:\n"
|
||||
"- Be concise and specific. Use exact numbers from the data.\n"
|
||||
"- Format costs as dollar amounts (e.g. $12.34).\n"
|
||||
"- When comparing models or providers, show a ranked list.\n"
|
||||
"- If data is empty, say so clearly.\n"
|
||||
"- Do not hallucinate data — only use what the tool returns."
|
||||
"- When comparing entities, show a ranked list.\n"
|
||||
"- If data is empty or no results found, say so clearly.\n"
|
||||
"- Do not hallucinate data — only use what the tools return.\n"
|
||||
"- Today's date is available in the data context. Use recent dates if the user says 'this week' or 'this month'."
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data fetchers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def _fetch_usage_data(
|
||||
start_date: str,
|
||||
end_date: str,
|
||||
user_id: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Call the aggregated daily activity query and return serialisable dict."""
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
get_daily_activity_aggregated,
|
||||
)
|
||||
|
|
@ -83,8 +154,71 @@ async def _fetch_usage_data(
|
|||
return response.model_dump(mode="json")
|
||||
|
||||
|
||||
async def _fetch_team_usage_data(
|
||||
start_date: str,
|
||||
end_date: str,
|
||||
team_ids: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
get_daily_activity,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
team_ids_list: Optional[List[str]] = None
|
||||
if team_ids:
|
||||
team_ids_list = [t.strip() for t in team_ids.split(",") if t.strip()]
|
||||
|
||||
response = await get_daily_activity(
|
||||
prisma_client=prisma_client,
|
||||
table_name="litellm_dailyteamspend",
|
||||
entity_id_field="team_id",
|
||||
entity_id=team_ids_list,
|
||||
entity_metadata_field=None,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
model=None,
|
||||
api_key=None,
|
||||
page=1,
|
||||
page_size=200,
|
||||
)
|
||||
return response.model_dump(mode="json")
|
||||
|
||||
|
||||
async def _fetch_tag_usage_data(
|
||||
start_date: str,
|
||||
end_date: str,
|
||||
tags: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
get_daily_activity,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
tag_list: Optional[List[str]] = None
|
||||
if tags:
|
||||
tag_list = [t.strip() for t in tags.split(",") if t.strip()]
|
||||
|
||||
response = await get_daily_activity(
|
||||
prisma_client=prisma_client,
|
||||
table_name="litellm_dailytagspend",
|
||||
entity_id_field="tag",
|
||||
entity_id=tag_list,
|
||||
entity_metadata_field=None,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
model=None,
|
||||
api_key=None,
|
||||
page=1,
|
||||
page_size=200,
|
||||
)
|
||||
return response.model_dump(mode="json")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Summarisers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _summarise_usage_data(data: Dict[str, Any]) -> str:
|
||||
"""Convert the raw aggregated response into a concise text summary for the LLM."""
|
||||
meta = data.get("metadata", {})
|
||||
results = data.get("results", [])
|
||||
|
||||
|
|
@ -128,7 +262,7 @@ def _summarise_usage_data(data: Dict[str, Any]) -> str:
|
|||
if model_spend:
|
||||
lines.append("Top Models by Spend:")
|
||||
for name, d in sorted(model_spend.items(), key=lambda x: -x[1]["spend"])[:15]:
|
||||
lines.append(f" - {name}: ${d['spend']:.4f} ({int(d['requests'])} requests, {int(d['tokens'])} tokens)")
|
||||
lines.append(f" - {name}: ${d['spend']:.4f} ({int(d['requests'])} reqs, {int(d['tokens'])} tokens)")
|
||||
else:
|
||||
lines.append("Models: (no data)")
|
||||
|
||||
|
|
@ -137,7 +271,7 @@ def _summarise_usage_data(data: Dict[str, Any]) -> str:
|
|||
if provider_spend:
|
||||
lines.append("Top Providers by Spend:")
|
||||
for name, d in sorted(provider_spend.items(), key=lambda x: -x[1]["spend"])[:10]:
|
||||
lines.append(f" - {name}: ${d['spend']:.4f} ({int(d['requests'])} requests)")
|
||||
lines.append(f" - {name}: ${d['spend']:.4f} ({int(d['requests'])} reqs)")
|
||||
else:
|
||||
lines.append("Providers: (no data)")
|
||||
|
||||
|
|
@ -158,11 +292,88 @@ def _summarise_usage_data(data: Dict[str, Any]) -> str:
|
|||
sorted_days = sorted(results, key=lambda x: x["date"])
|
||||
for day in sorted_days:
|
||||
m = day.get("metrics", {})
|
||||
lines.append(f" - {day['date']}: ${m.get('spend', 0):.4f} ({m.get('api_requests', 0)} requests)")
|
||||
lines.append(f" - {day['date']}: ${m.get('spend', 0):.4f} ({m.get('api_requests', 0)} reqs)")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _summarise_entity_data(data: Dict[str, Any], entity_label: str) -> str:
|
||||
"""Summarise team/tag/org/customer entity usage data."""
|
||||
results = data.get("results", [])
|
||||
if not results:
|
||||
return f"No {entity_label} usage data found for the given date range."
|
||||
|
||||
entity_totals: Dict[str, Dict[str, Any]] = {}
|
||||
for day in results:
|
||||
breakdown = day.get("breakdown", {})
|
||||
for entity_id, entity_data in breakdown.get("entities", {}).items():
|
||||
if entity_id not in entity_totals:
|
||||
alias = entity_data.get("metadata", {}).get("alias", entity_id)
|
||||
entity_totals[entity_id] = {
|
||||
"alias": alias,
|
||||
"spend": 0,
|
||||
"requests": 0,
|
||||
"tokens": 0,
|
||||
"models": {},
|
||||
}
|
||||
m = entity_data.get("metrics", {})
|
||||
entity_totals[entity_id]["spend"] += m.get("spend", 0)
|
||||
entity_totals[entity_id]["requests"] += m.get("api_requests", 0)
|
||||
entity_totals[entity_id]["tokens"] += m.get("total_tokens", 0)
|
||||
|
||||
for model_name, model_data in entity_data.get("api_key_breakdown", {}).items():
|
||||
models_dict = entity_totals[entity_id]["models"]
|
||||
if model_name not in models_dict:
|
||||
models_dict[model_name] = 0
|
||||
models_dict[model_name] += model_data.get("metrics", {}).get("spend", 0)
|
||||
|
||||
lines = [f"{entity_label} Usage Summary ({len(entity_totals)} {entity_label.lower()}s):", ""]
|
||||
|
||||
for eid, d in sorted(entity_totals.items(), key=lambda x: -x[1]["spend"]):
|
||||
label = d["alias"] if d["alias"] != eid else eid
|
||||
lines.append(f"- {label} (ID: {eid}): ${d['spend']:.4f} | {int(d['requests'])} reqs | {int(d['tokens'])} tokens")
|
||||
if d["models"]:
|
||||
for model, spend in sorted(d["models"].items(), key=lambda x: -x[1])[:5]:
|
||||
lines.append(f" Model: {model}: ${spend:.4f}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool dispatcher
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
TOOL_HANDLERS = {
|
||||
"get_usage_data": {
|
||||
"fetch": _fetch_usage_data,
|
||||
"summarise": _summarise_usage_data,
|
||||
"label": "global usage data",
|
||||
},
|
||||
"get_team_usage_data": {
|
||||
"fetch": _fetch_team_usage_data,
|
||||
"summarise": lambda data: _summarise_entity_data(data, "Team"),
|
||||
"label": "team usage data",
|
||||
},
|
||||
"get_tag_usage_data": {
|
||||
"fetch": _fetch_tag_usage_data,
|
||||
"summarise": lambda data: _summarise_entity_data(data, "Tag"),
|
||||
"label": "tag usage data",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SSE helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _sse(event: dict) -> str:
|
||||
return f"data: {json.dumps(event)}\n\n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main streaming function
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def stream_usage_ai_chat(
|
||||
messages: List[Dict[str, str]],
|
||||
model: Optional[str] = None,
|
||||
|
|
@ -172,10 +383,11 @@ async def stream_usage_ai_chat(
|
|||
"""
|
||||
Stream an AI chat response about usage data.
|
||||
|
||||
Yields SSE-formatted events:
|
||||
data: {"type": "chunk", "content": "..."}
|
||||
data: {"type": "done"}
|
||||
data: {"type": "error", "message": "..."}
|
||||
Yields SSE events:
|
||||
{"type": "status", "message": "..."} - thinking/tool status
|
||||
{"type": "chunk", "content": "..."} - streamed response text
|
||||
{"type": "done"} - stream finished
|
||||
{"type": "error", "message": "..."} - error
|
||||
"""
|
||||
model = model.strip() if model else ""
|
||||
model = model or DEFAULT_COMPETITOR_DISCOVERY_MODEL
|
||||
|
|
@ -186,10 +398,12 @@ async def stream_usage_ai_chat(
|
|||
]
|
||||
|
||||
try:
|
||||
yield _sse({"type": "status", "message": "Thinking..."})
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model=model,
|
||||
messages=chat_messages,
|
||||
tools=[GET_USAGE_DATA_TOOL],
|
||||
tools=ALL_TOOLS,
|
||||
temperature=0.2,
|
||||
)
|
||||
|
||||
|
|
@ -203,32 +417,51 @@ async def stream_usage_ai_chat(
|
|||
fn_name = tool_call.function.name
|
||||
fn_args = json.loads(tool_call.function.arguments)
|
||||
|
||||
if fn_name == "get_usage_data":
|
||||
effective_user = None
|
||||
if not is_admin:
|
||||
effective_user = user_id
|
||||
elif fn_args.get("user_id"):
|
||||
effective_user = fn_args["user_id"]
|
||||
|
||||
try:
|
||||
raw_data = await _fetch_usage_data(
|
||||
start_date=fn_args["start_date"],
|
||||
end_date=fn_args["end_date"],
|
||||
user_id=effective_user,
|
||||
)
|
||||
tool_result = _summarise_usage_data(raw_data)
|
||||
except Exception as e:
|
||||
tool_result = f"Error fetching usage data: {str(e)}"
|
||||
else:
|
||||
tool_result = f"Unknown tool: {fn_name}"
|
||||
|
||||
chat_messages.append(
|
||||
{
|
||||
handler = TOOL_HANDLERS.get(fn_name)
|
||||
if not handler:
|
||||
yield _sse({"type": "status", "message": f"Unknown tool: {fn_name}"})
|
||||
chat_messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"content": tool_result,
|
||||
"content": f"Unknown tool: {fn_name}",
|
||||
})
|
||||
continue
|
||||
|
||||
yield _sse({
|
||||
"type": "status",
|
||||
"message": f"Fetching {handler['label']} ({fn_args.get('start_date', '')} to {fn_args.get('end_date', '')})..."
|
||||
})
|
||||
|
||||
try:
|
||||
fetch_kwargs: Dict[str, Any] = {
|
||||
"start_date": fn_args["start_date"],
|
||||
"end_date": fn_args["end_date"],
|
||||
}
|
||||
)
|
||||
|
||||
if fn_name == "get_usage_data":
|
||||
if not is_admin:
|
||||
fetch_kwargs["user_id"] = user_id
|
||||
elif fn_args.get("user_id"):
|
||||
fetch_kwargs["user_id"] = fn_args["user_id"]
|
||||
elif fn_name == "get_team_usage_data":
|
||||
if fn_args.get("team_ids"):
|
||||
fetch_kwargs["team_ids"] = fn_args["team_ids"]
|
||||
elif fn_name == "get_tag_usage_data":
|
||||
if fn_args.get("tags"):
|
||||
fetch_kwargs["tags"] = fn_args["tags"]
|
||||
|
||||
raw_data = await handler["fetch"](**fetch_kwargs)
|
||||
tool_result = handler["summarise"](raw_data)
|
||||
except Exception as e:
|
||||
tool_result = f"Error fetching {handler['label']}: {str(e)}"
|
||||
|
||||
chat_messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"content": tool_result,
|
||||
})
|
||||
|
||||
yield _sse({"type": "status", "message": "Analyzing results..."})
|
||||
|
||||
final_response = await litellm.acompletion(
|
||||
model=model,
|
||||
|
|
@ -240,16 +473,14 @@ async def stream_usage_ai_chat(
|
|||
async for chunk in final_response:
|
||||
delta_content = chunk.choices[0].delta.content
|
||||
if delta_content:
|
||||
event = json.dumps({"type": "chunk", "content": delta_content})
|
||||
yield f"data: {event}\n\n"
|
||||
yield _sse({"type": "chunk", "content": delta_content})
|
||||
else:
|
||||
content = choice.message.content or ""
|
||||
if content:
|
||||
event = json.dumps({"type": "chunk", "content": content})
|
||||
yield f"data: {event}\n\n"
|
||||
yield _sse({"type": "chunk", "content": content})
|
||||
|
||||
yield f"data: {json.dumps({'type': 'done'})}\n\n"
|
||||
yield _sse({"type": "done"})
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("AI usage chat failed: %s", e)
|
||||
yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
|
||||
yield _sse({"type": "error", "message": str(e)})
|
||||
|
|
|
|||
|
|
@ -8,8 +8,9 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
import pytest
|
||||
|
||||
from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import (
|
||||
GET_USAGE_DATA_TOOL,
|
||||
ALL_TOOLS,
|
||||
SYSTEM_PROMPT,
|
||||
_summarise_entity_data,
|
||||
_summarise_usage_data,
|
||||
stream_usage_ai_chat,
|
||||
)
|
||||
|
|
@ -37,27 +38,6 @@ SAMPLE_AGGREGATED_RESPONSE = {
|
|||
"spend": 40.0,
|
||||
"api_requests": 300,
|
||||
"total_tokens": 25000,
|
||||
"prompt_tokens": 15000,
|
||||
"completion_tokens": 10000,
|
||||
"successful_requests": 290,
|
||||
"failed_requests": 10,
|
||||
"cache_read_input_tokens": 0,
|
||||
"cache_creation_input_tokens": 0,
|
||||
},
|
||||
"metadata": {},
|
||||
"api_key_breakdown": {},
|
||||
},
|
||||
"gpt-3.5-turbo": {
|
||||
"metrics": {
|
||||
"spend": 10.25,
|
||||
"api_requests": 200,
|
||||
"total_tokens": 5000,
|
||||
"prompt_tokens": 3000,
|
||||
"completion_tokens": 2000,
|
||||
"successful_requests": 190,
|
||||
"failed_requests": 10,
|
||||
"cache_read_input_tokens": 0,
|
||||
"cache_creation_input_tokens": 0,
|
||||
},
|
||||
"metadata": {},
|
||||
"api_key_breakdown": {},
|
||||
|
|
@ -65,34 +45,14 @@ SAMPLE_AGGREGATED_RESPONSE = {
|
|||
},
|
||||
"providers": {
|
||||
"openai": {
|
||||
"metrics": {
|
||||
"spend": 50.25,
|
||||
"api_requests": 500,
|
||||
"total_tokens": 30000,
|
||||
"prompt_tokens": 20000,
|
||||
"completion_tokens": 10000,
|
||||
"successful_requests": 480,
|
||||
"failed_requests": 20,
|
||||
"cache_read_input_tokens": 0,
|
||||
"cache_creation_input_tokens": 0,
|
||||
},
|
||||
"metrics": {"spend": 50.25, "api_requests": 500},
|
||||
"metadata": {},
|
||||
"api_key_breakdown": {},
|
||||
},
|
||||
},
|
||||
"api_keys": {
|
||||
"sk-test123": {
|
||||
"metrics": {
|
||||
"spend": 50.25,
|
||||
"api_requests": 500,
|
||||
"total_tokens": 30000,
|
||||
"prompt_tokens": 20000,
|
||||
"completion_tokens": 10000,
|
||||
"successful_requests": 480,
|
||||
"failed_requests": 20,
|
||||
"cache_read_input_tokens": 0,
|
||||
"cache_creation_input_tokens": 0,
|
||||
},
|
||||
"metrics": {"spend": 50.25},
|
||||
"metadata": {"key_alias": "Production Key"},
|
||||
},
|
||||
},
|
||||
|
|
@ -104,35 +64,55 @@ SAMPLE_AGGREGATED_RESPONSE = {
|
|||
],
|
||||
"metadata": {
|
||||
"total_spend": 50.25,
|
||||
"total_prompt_tokens": 20000,
|
||||
"total_completion_tokens": 10000,
|
||||
"total_tokens": 30000,
|
||||
"total_api_requests": 500,
|
||||
"total_successful_requests": 480,
|
||||
"total_failed_requests": 20,
|
||||
"total_cache_read_input_tokens": 0,
|
||||
"total_cache_creation_input_tokens": 0,
|
||||
"page": 1,
|
||||
"total_pages": 1,
|
||||
"has_more": False,
|
||||
"total_tokens": 30000,
|
||||
},
|
||||
}
|
||||
|
||||
SAMPLE_TEAM_RESPONSE = {
|
||||
"results": [
|
||||
{
|
||||
"date": "2025-01-15",
|
||||
"metrics": {"spend": 100.0, "api_requests": 1000, "total_tokens": 50000},
|
||||
"breakdown": {
|
||||
"entities": {
|
||||
"team-1": {
|
||||
"metrics": {"spend": 60.0, "api_requests": 600, "total_tokens": 30000},
|
||||
"metadata": {"alias": "Engineering"},
|
||||
"api_key_breakdown": {},
|
||||
},
|
||||
"team-2": {
|
||||
"metrics": {"spend": 40.0, "api_requests": 400, "total_tokens": 20000},
|
||||
"metadata": {"alias": "Marketing"},
|
||||
"api_key_breakdown": {},
|
||||
},
|
||||
},
|
||||
"models": {},
|
||||
"providers": {},
|
||||
"api_keys": {},
|
||||
"model_groups": {},
|
||||
"mcp_servers": {},
|
||||
},
|
||||
},
|
||||
],
|
||||
"metadata": {"total_spend": 100.0, "total_api_requests": 1000},
|
||||
}
|
||||
|
||||
class TestToolSchema:
|
||||
def test_tool_schema_is_valid(self):
|
||||
assert GET_USAGE_DATA_TOOL["type"] == "function"
|
||||
func = GET_USAGE_DATA_TOOL["function"]
|
||||
assert func["name"] == "get_usage_data"
|
||||
params = func["parameters"]
|
||||
assert "start_date" in params["properties"]
|
||||
assert "end_date" in params["properties"]
|
||||
assert "user_id" in params["properties"]
|
||||
assert params["required"] == ["start_date", "end_date"]
|
||||
|
||||
def test_system_prompt_mentions_tool(self):
|
||||
class TestToolSchemas:
|
||||
def test_all_tools_defined(self):
|
||||
assert len(ALL_TOOLS) == 3
|
||||
names = {t["function"]["name"] for t in ALL_TOOLS}
|
||||
assert "get_usage_data" in names
|
||||
assert "get_team_usage_data" in names
|
||||
assert "get_tag_usage_data" in names
|
||||
|
||||
def test_system_prompt_mentions_all_tools(self):
|
||||
assert "get_usage_data" in SYSTEM_PROMPT
|
||||
assert "usage" in SYSTEM_PROMPT.lower()
|
||||
assert "get_team_usage_data" in SYSTEM_PROMPT
|
||||
assert "get_tag_usage_data" in SYSTEM_PROMPT
|
||||
|
||||
|
||||
class TestSummariseUsageData:
|
||||
|
|
@ -140,14 +120,10 @@ class TestSummariseUsageData:
|
|||
summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE)
|
||||
assert "$50.25" in summary
|
||||
assert "500" in summary
|
||||
assert "480" in summary
|
||||
assert "20" in summary
|
||||
|
||||
def test_summarise_includes_models(self):
|
||||
summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE)
|
||||
assert "gpt-4" in summary
|
||||
assert "gpt-3.5-turbo" in summary
|
||||
assert "$40.0" in summary
|
||||
|
||||
def test_summarise_includes_providers(self):
|
||||
summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE)
|
||||
|
|
@ -157,19 +133,29 @@ class TestSummariseUsageData:
|
|||
summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE)
|
||||
assert "Production Key" in summary
|
||||
|
||||
def test_summarise_includes_daily(self):
|
||||
summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE)
|
||||
assert "2025-01-15" in summary
|
||||
|
||||
def test_summarise_handles_empty_data(self):
|
||||
empty = {"results": [], "metadata": {}}
|
||||
summary = _summarise_usage_data(empty)
|
||||
assert "no data" in summary.lower()
|
||||
|
||||
|
||||
class TestSummariseEntityData:
|
||||
def test_team_summary_includes_teams(self):
|
||||
summary = _summarise_entity_data(SAMPLE_TEAM_RESPONSE, "Team")
|
||||
assert "Engineering" in summary
|
||||
assert "Marketing" in summary
|
||||
assert "$60.0" in summary
|
||||
assert "$40.0" in summary
|
||||
|
||||
def test_team_summary_empty(self):
|
||||
empty = {"results": [], "metadata": {}}
|
||||
summary = _summarise_entity_data(empty, "Team")
|
||||
assert "No Team usage data" in summary
|
||||
|
||||
|
||||
class TestStreamUsageAiChat:
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_with_tool_call(self):
|
||||
async def test_stream_emits_status_events(self):
|
||||
mock_tool_call = MagicMock()
|
||||
mock_tool_call.id = "call_123"
|
||||
mock_tool_call.function.name = "get_usage_data"
|
||||
|
|
@ -182,24 +168,17 @@ class TestStreamUsageAiChat:
|
|||
mock_first_response.choices = [MagicMock()]
|
||||
mock_first_response.choices[0].message.tool_calls = [mock_tool_call]
|
||||
mock_first_response.choices[0].message.model_dump.return_value = {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_123",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_usage_data",
|
||||
"arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}',
|
||||
},
|
||||
}
|
||||
],
|
||||
"role": "assistant", "content": None,
|
||||
"tool_calls": [{"id": "call_123", "type": "function", "function": {
|
||||
"name": "get_usage_data",
|
||||
"arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}',
|
||||
}}],
|
||||
}
|
||||
|
||||
async def mock_stream():
|
||||
chunk = MagicMock()
|
||||
chunk.choices = [MagicMock()]
|
||||
chunk.choices[0].delta.content = "Your total spend is $50.25"
|
||||
chunk.choices[0].delta.content = "Total spend is $50.25"
|
||||
yield chunk
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm") as mock_litellm, \
|
||||
|
|
@ -218,36 +197,65 @@ class TestStreamUsageAiChat:
|
|||
user_id="user-123",
|
||||
is_admin=True,
|
||||
):
|
||||
events.append(event)
|
||||
events.append(json.loads(event.replace("data: ", "").strip()))
|
||||
|
||||
assert len(events) >= 2
|
||||
chunk_event = json.loads(events[0].replace("data: ", "").strip())
|
||||
assert chunk_event["type"] == "chunk"
|
||||
assert "$50.25" in chunk_event["content"]
|
||||
status_events = [e for e in events if e["type"] == "status"]
|
||||
chunk_events = [e for e in events if e["type"] == "chunk"]
|
||||
done_events = [e for e in events if e["type"] == "done"]
|
||||
|
||||
done_event = json.loads(events[-1].replace("data: ", "").strip())
|
||||
assert done_event["type"] == "done"
|
||||
assert len(status_events) >= 2
|
||||
assert "Thinking" in status_events[0]["message"]
|
||||
assert "Fetching" in status_events[1]["message"]
|
||||
assert len(chunk_events) >= 1
|
||||
assert len(done_events) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_without_tool_call(self):
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.tool_calls = None
|
||||
mock_response.choices[0].message.content = "I need more context."
|
||||
async def test_stream_handles_team_tool(self):
|
||||
mock_tool_call = MagicMock()
|
||||
mock_tool_call.id = "call_team"
|
||||
mock_tool_call.function.name = "get_team_usage_data"
|
||||
mock_tool_call.function.arguments = json.dumps({
|
||||
"start_date": "2025-01-01",
|
||||
"end_date": "2025-01-31",
|
||||
})
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm") as mock_litellm:
|
||||
mock_litellm.acompletion = AsyncMock(return_value=mock_response)
|
||||
mock_first_response = MagicMock()
|
||||
mock_first_response.choices = [MagicMock()]
|
||||
mock_first_response.choices[0].message.tool_calls = [mock_tool_call]
|
||||
mock_first_response.choices[0].message.model_dump.return_value = {
|
||||
"role": "assistant", "content": None,
|
||||
"tool_calls": [{"id": "call_team", "type": "function", "function": {
|
||||
"name": "get_team_usage_data",
|
||||
"arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}',
|
||||
}}],
|
||||
}
|
||||
|
||||
async def mock_stream():
|
||||
chunk = MagicMock()
|
||||
chunk.choices = [MagicMock()]
|
||||
chunk.choices[0].delta.content = "Engineering is the top team."
|
||||
yield chunk
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm") as mock_litellm, \
|
||||
patch("litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._fetch_team_usage_data", new_callable=AsyncMock) as mock_fetch:
|
||||
|
||||
mock_litellm.acompletion = AsyncMock(side_effect=[
|
||||
mock_first_response,
|
||||
mock_stream(),
|
||||
])
|
||||
mock_fetch.return_value = SAMPLE_TEAM_RESPONSE
|
||||
|
||||
events = []
|
||||
async for event in stream_usage_ai_chat(
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
messages=[{"role": "user", "content": "Which team spends the most?"}],
|
||||
model="gpt-4o-mini",
|
||||
is_admin=True,
|
||||
):
|
||||
events.append(event)
|
||||
events.append(json.loads(event.replace("data: ", "").strip()))
|
||||
|
||||
assert len(events) >= 2
|
||||
chunk_event = json.loads(events[0].replace("data: ", "").strip())
|
||||
assert chunk_event["type"] == "chunk"
|
||||
assert "context" in chunk_event["content"]
|
||||
chunk_events = [e for e in events if e["type"] == "chunk"]
|
||||
assert len(chunk_events) >= 1
|
||||
assert "Engineering" in chunk_events[0]["content"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_handles_error(self):
|
||||
|
|
@ -258,12 +266,11 @@ class TestStreamUsageAiChat:
|
|||
async for event in stream_usage_ai_chat(
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
):
|
||||
events.append(event)
|
||||
events.append(json.loads(event.replace("data: ", "").strip()))
|
||||
|
||||
assert len(events) == 1
|
||||
error_event = json.loads(events[0].replace("data: ", "").strip())
|
||||
assert error_event["type"] == "error"
|
||||
assert "LLM error" in error_event["message"]
|
||||
error_events = [e for e in events if e["type"] == "error"]
|
||||
assert len(error_events) == 1
|
||||
assert "LLM error" in error_events[0]["message"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_admin_enforces_user_id(self):
|
||||
|
|
@ -280,38 +287,39 @@ class TestStreamUsageAiChat:
|
|||
mock_first_response.choices = [MagicMock()]
|
||||
mock_first_response.choices[0].message.tool_calls = [mock_tool_call]
|
||||
mock_first_response.choices[0].message.model_dump.return_value = {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_456",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_usage_data",
|
||||
"arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31","user_id":"other-user"}',
|
||||
},
|
||||
}
|
||||
],
|
||||
"role": "assistant", "content": None,
|
||||
"tool_calls": [{"id": "call_456", "type": "function", "function": {
|
||||
"name": "get_usage_data",
|
||||
"arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31","user_id":"other-user"}',
|
||||
}}],
|
||||
}
|
||||
|
||||
async def mock_stream():
|
||||
chunk = MagicMock()
|
||||
chunk.choices = [MagicMock()]
|
||||
chunk.choices[0].delta.content = "Here is your data."
|
||||
chunk.choices[0].delta.content = "Data."
|
||||
yield chunk
|
||||
|
||||
mock_fetch = AsyncMock(return_value=SAMPLE_AGGREGATED_RESPONSE)
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm") as mock_litellm, \
|
||||
patch("litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._fetch_usage_data", new_callable=AsyncMock) as mock_fetch:
|
||||
patch.dict(
|
||||
"litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.TOOL_HANDLERS",
|
||||
{"get_usage_data": {
|
||||
"fetch": mock_fetch,
|
||||
"summarise": _summarise_usage_data,
|
||||
"label": "global usage data",
|
||||
}},
|
||||
):
|
||||
|
||||
mock_litellm.acompletion = AsyncMock(side_effect=[
|
||||
mock_first_response,
|
||||
mock_stream(),
|
||||
])
|
||||
mock_fetch.return_value = SAMPLE_AGGREGATED_RESPONSE
|
||||
|
||||
events = []
|
||||
async for event in stream_usage_ai_chat(
|
||||
messages=[{"role": "user", "content": "Show me other user data"}],
|
||||
messages=[{"role": "user", "content": "Show data"}],
|
||||
model="gpt-4o-mini",
|
||||
user_id="my-user-id",
|
||||
is_admin=False,
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({
|
|||
const [availableModels, setAvailableModels] = useState<string[]>([]);
|
||||
const [isLoadingModels, setIsLoadingModels] = useState(false);
|
||||
const [streamingContent, setStreamingContent] = useState("");
|
||||
const [statusMessage, setStatusMessage] = useState<string | null>(null);
|
||||
const messagesEndRef = useRef<HTMLDivElement>(null);
|
||||
const abortControllerRef = useRef<AbortController | null>(null);
|
||||
|
||||
|
|
@ -70,6 +71,7 @@ const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({
|
|||
setInputText("");
|
||||
setIsLoading(true);
|
||||
setStreamingContent("");
|
||||
setStatusMessage(null);
|
||||
|
||||
const abortController = new AbortController();
|
||||
abortControllerRef.current = abortController;
|
||||
|
|
@ -82,20 +84,26 @@ const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({
|
|||
updatedMessages.map((m) => ({ role: m.role, content: m.content })),
|
||||
selectedModel || "",
|
||||
(content: string) => {
|
||||
setStatusMessage(null);
|
||||
accumulated += content;
|
||||
setStreamingContent(accumulated);
|
||||
},
|
||||
() => {
|
||||
setStatusMessage(null);
|
||||
setMessages((prev) => [...prev, { role: "assistant", content: accumulated }]);
|
||||
setStreamingContent("");
|
||||
},
|
||||
(errorMsg: string) => {
|
||||
setStatusMessage(null);
|
||||
setMessages((prev) => [
|
||||
...prev,
|
||||
{ role: "assistant", content: `Error: ${errorMsg}` },
|
||||
]);
|
||||
setStreamingContent("");
|
||||
},
|
||||
(status: string) => {
|
||||
setStatusMessage(status);
|
||||
},
|
||||
abortController.signal,
|
||||
);
|
||||
} catch (error: any) {
|
||||
|
|
@ -220,8 +228,11 @@ const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({
|
|||
|
||||
{isLoading && !streamingContent && (
|
||||
<div className="flex justify-start">
|
||||
<div className="rounded-xl px-3.5 py-2 bg-white border border-gray-200">
|
||||
<div className="rounded-xl px-3.5 py-2 bg-white border border-gray-200 flex items-center gap-2">
|
||||
<Spin size="small" />
|
||||
{statusMessage && (
|
||||
<span className="text-xs text-gray-500 italic">{statusMessage}</span>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
|
|
|||
|
|
@ -5914,6 +5914,7 @@ export const usageAiChatStream = async (
|
|||
onChunk: (content: string) => void,
|
||||
onDone: () => void,
|
||||
onError?: (error: string) => void,
|
||||
onStatus?: (message: string) => void,
|
||||
signal?: AbortSignal,
|
||||
) => {
|
||||
const url = proxyBaseUrl
|
||||
|
|
@ -5957,6 +5958,8 @@ export const usageAiChatStream = async (
|
|||
const event = JSON.parse(line.slice(6));
|
||||
if (event.type === "chunk") {
|
||||
onChunk(event.content);
|
||||
} else if (event.type === "status") {
|
||||
onStatus?.(event.message);
|
||||
} else if (event.type === "done") {
|
||||
onDone();
|
||||
} else if (event.type === "error") {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue