mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Add backend AI usage chat endpoint with tool calling
Backend: - New /usage/ai/chat SSE streaming endpoint - AI agent has get_usage_data tool that queries /user/daily/activity/aggregated - Follows same architecture as policy AI suggest (litellm.acompletion + tools) - Non-admin users are restricted to their own data - 12 backend unit tests Frontend: - Panel now calls /usage/ai/chat backend endpoint via SSE - Removed direct OpenAI client calls from frontend - Added usageAiChatStream networking function following enrichPolicyTemplateStream pattern Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>
This commit is contained in:
Cursor Agent
parent
1a6ded461a
commit
fb8eaf2f7d
10 changed files with 751 additions and 242 deletions
|
|
@ -0,0 +1,9 @@
|
|||
"""
|
||||
Usage endpoints package.
|
||||
|
||||
Re-exports the router from endpoints module.
|
||||
"""
|
||||
|
||||
from litellm.proxy.management_endpoints.usage_endpoints.endpoints import ( # noqa: F401
|
||||
router,
|
||||
)
|
||||
|
|
@ -0,0 +1,254 @@
|
|||
"""
|
||||
AI Usage Chat - uses LLM tool calling to answer questions about
|
||||
usage/spend data from the /user/daily/activity endpoints.
|
||||
"""
|
||||
|
||||
import json
|
||||
from datetime import date, timedelta
|
||||
from typing import Any, AsyncIterator, Dict, List, Optional
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL
|
||||
|
||||
GET_USAGE_DATA_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_usage_data",
|
||||
"description": (
|
||||
"Fetch aggregated 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."
|
||||
),
|
||||
"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",
|
||||
},
|
||||
"user_id": {
|
||||
"type": "string",
|
||||
"description": "Optional user ID to filter by a specific user. Omit for global view.",
|
||||
},
|
||||
},
|
||||
"required": ["start_date", "end_date"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
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"
|
||||
"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."
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
response = await get_daily_activity_aggregated(
|
||||
prisma_client=prisma_client,
|
||||
table_name="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
entity_id=user_id,
|
||||
entity_metadata_field=None,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
model=None,
|
||||
api_key=None,
|
||||
)
|
||||
return response.model_dump(mode="json")
|
||||
|
||||
|
||||
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", [])
|
||||
|
||||
lines = [
|
||||
f"Date Range: {results[0]['date'] if results else 'N/A'} to {results[-1]['date'] if results else 'N/A'}",
|
||||
f"Total Spend: ${meta.get('total_spend', 0):.4f}",
|
||||
f"Total Requests: {meta.get('total_api_requests', 0)}",
|
||||
f"Successful Requests: {meta.get('total_successful_requests', 0)}",
|
||||
f"Failed Requests: {meta.get('total_failed_requests', 0)}",
|
||||
f"Total Tokens: {meta.get('total_tokens', 0)}",
|
||||
"",
|
||||
]
|
||||
|
||||
model_spend: Dict[str, Dict[str, float]] = {}
|
||||
provider_spend: Dict[str, Dict[str, float]] = {}
|
||||
key_spend: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
for day in results:
|
||||
breakdown = day.get("breakdown", {})
|
||||
for model, metrics in breakdown.get("models", {}).items():
|
||||
if model not in model_spend:
|
||||
model_spend[model] = {"spend": 0, "requests": 0, "tokens": 0}
|
||||
m = metrics.get("metrics", {})
|
||||
model_spend[model]["spend"] += m.get("spend", 0)
|
||||
model_spend[model]["requests"] += m.get("api_requests", 0)
|
||||
model_spend[model]["tokens"] += m.get("total_tokens", 0)
|
||||
|
||||
for provider, metrics in breakdown.get("providers", {}).items():
|
||||
if provider not in provider_spend:
|
||||
provider_spend[provider] = {"spend": 0, "requests": 0}
|
||||
m = metrics.get("metrics", {})
|
||||
provider_spend[provider]["spend"] += m.get("spend", 0)
|
||||
provider_spend[provider]["requests"] += m.get("api_requests", 0)
|
||||
|
||||
for key, metrics in breakdown.get("api_keys", {}).items():
|
||||
if key not in key_spend:
|
||||
alias = metrics.get("metadata", {}).get("key_alias")
|
||||
key_spend[key] = {"spend": 0, "alias": alias}
|
||||
key_spend[key]["spend"] += metrics.get("metrics", {}).get("spend", 0)
|
||||
|
||||
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)")
|
||||
else:
|
||||
lines.append("Models: (no data)")
|
||||
|
||||
lines.append("")
|
||||
|
||||
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)")
|
||||
else:
|
||||
lines.append("Providers: (no data)")
|
||||
|
||||
lines.append("")
|
||||
|
||||
if key_spend:
|
||||
lines.append("Top API Keys by Spend:")
|
||||
for key, d in sorted(key_spend.items(), key=lambda x: -x[1]["spend"])[:10]:
|
||||
label = d["alias"] or key
|
||||
lines.append(f" - {label}: ${d['spend']:.4f}")
|
||||
else:
|
||||
lines.append("API Keys: (no data)")
|
||||
|
||||
lines.append("")
|
||||
|
||||
if results:
|
||||
lines.append("Daily Spend:")
|
||||
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)")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
async def stream_usage_ai_chat(
|
||||
messages: List[Dict[str, str]],
|
||||
model: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
Stream an AI chat response about usage data.
|
||||
|
||||
Yields SSE-formatted events:
|
||||
data: {"type": "chunk", "content": "..."}
|
||||
data: {"type": "done"}
|
||||
data: {"type": "error", "message": "..."}
|
||||
"""
|
||||
model = model or DEFAULT_COMPETITOR_DISCOVERY_MODEL
|
||||
|
||||
chat_messages: List[Dict[str, Any]] = [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
*messages,
|
||||
]
|
||||
|
||||
try:
|
||||
response = await litellm.acompletion(
|
||||
model=model,
|
||||
messages=chat_messages,
|
||||
tools=[GET_USAGE_DATA_TOOL],
|
||||
temperature=0.2,
|
||||
)
|
||||
|
||||
choice = response.choices[0] # type: ignore
|
||||
tool_calls = choice.message.tool_calls
|
||||
|
||||
if tool_calls:
|
||||
chat_messages.append(choice.message.model_dump())
|
||||
|
||||
for tool_call in tool_calls:
|
||||
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(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"content": tool_result,
|
||||
}
|
||||
)
|
||||
|
||||
final_response = await litellm.acompletion(
|
||||
model=model,
|
||||
messages=chat_messages,
|
||||
stream=True,
|
||||
temperature=0.2,
|
||||
)
|
||||
|
||||
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"
|
||||
else:
|
||||
content = choice.message.content or ""
|
||||
if content:
|
||||
event = json.dumps({"type": "chunk", "content": content})
|
||||
yield f"data: {event}\n\n"
|
||||
|
||||
yield f"data: {json.dumps({'type': 'done'})}\n\n"
|
||||
|
||||
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"
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
"""
|
||||
USAGE AI CHAT ENDPOINTS
|
||||
|
||||
/usage/ai/chat - Stream AI chat responses about usage data
|
||||
"""
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class ChatMessage(BaseModel):
|
||||
role: str
|
||||
content: str
|
||||
|
||||
|
||||
class UsageAIChatRequest(BaseModel):
|
||||
messages: List[ChatMessage] = Field(
|
||||
..., description="Chat messages (user/assistant history)"
|
||||
)
|
||||
model: Optional[str] = Field(
|
||||
default=None, description="Model to use for AI chat"
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/usage/ai/chat",
|
||||
tags=["Budget & Spend Tracking"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def usage_ai_chat(
|
||||
data: UsageAIChatRequest,
|
||||
request: Request,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""
|
||||
AI chat about usage data.
|
||||
|
||||
Streams SSE events with the AI response. The AI agent has access
|
||||
to the `get_usage_data` tool which queries the aggregated daily
|
||||
activity endpoint internally.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import (
|
||||
stream_usage_ai_chat,
|
||||
)
|
||||
|
||||
from litellm.proxy.management_endpoints.common_utils import (
|
||||
_user_has_admin_view,
|
||||
)
|
||||
|
||||
is_admin = _user_has_admin_view(user_api_key_dict)
|
||||
user_id = user_api_key_dict.user_id
|
||||
|
||||
messages = [{"role": m.role, "content": m.content} for m in data.messages]
|
||||
|
||||
return StreamingResponse(
|
||||
stream_usage_ai_chat(
|
||||
messages=messages,
|
||||
model=data.model,
|
||||
user_id=user_id,
|
||||
is_admin=is_admin,
|
||||
),
|
||||
media_type="text/event-stream",
|
||||
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
||||
)
|
||||
|
|
@ -391,6 +391,7 @@ from litellm.proxy.management_endpoints.organization_endpoints import (
|
|||
router as organization_router,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.policy_endpoints import router as policy_router
|
||||
from litellm.proxy.management_endpoints.usage_endpoints import router as usage_ai_router
|
||||
from litellm.proxy.management_endpoints.project_endpoints import (
|
||||
router as project_router,
|
||||
)
|
||||
|
|
@ -12844,6 +12845,7 @@ app.include_router(caching_router)
|
|||
app.include_router(analytics_router)
|
||||
app.include_router(guardrails_router)
|
||||
app.include_router(policy_router)
|
||||
app.include_router(usage_ai_router)
|
||||
app.include_router(policy_crud_router)
|
||||
app.include_router(policy_resolve_router)
|
||||
app.include_router(search_tool_management_router)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,325 @@
|
|||
"""
|
||||
Tests for AI Usage Chat module.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import (
|
||||
GET_USAGE_DATA_TOOL,
|
||||
SYSTEM_PROMPT,
|
||||
_summarise_usage_data,
|
||||
stream_usage_ai_chat,
|
||||
)
|
||||
|
||||
|
||||
SAMPLE_AGGREGATED_RESPONSE = {
|
||||
"results": [
|
||||
{
|
||||
"date": "2025-01-15",
|
||||
"metrics": {
|
||||
"spend": 50.25,
|
||||
"prompt_tokens": 20000,
|
||||
"completion_tokens": 10000,
|
||||
"total_tokens": 30000,
|
||||
"api_requests": 500,
|
||||
"successful_requests": 480,
|
||||
"failed_requests": 20,
|
||||
"cache_read_input_tokens": 0,
|
||||
"cache_creation_input_tokens": 0,
|
||||
},
|
||||
"breakdown": {
|
||||
"models": {
|
||||
"gpt-4": {
|
||||
"metrics": {
|
||||
"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": {},
|
||||
},
|
||||
},
|
||||
"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,
|
||||
},
|
||||
"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,
|
||||
},
|
||||
"metadata": {"key_alias": "Production Key"},
|
||||
},
|
||||
},
|
||||
"model_groups": {},
|
||||
"mcp_servers": {},
|
||||
"entities": {},
|
||||
},
|
||||
},
|
||||
],
|
||||
"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,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
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):
|
||||
assert "get_usage_data" in SYSTEM_PROMPT
|
||||
assert "usage" in SYSTEM_PROMPT.lower()
|
||||
|
||||
|
||||
class TestSummariseUsageData:
|
||||
def test_summarise_includes_totals(self):
|
||||
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)
|
||||
assert "openai" in summary
|
||||
|
||||
def test_summarise_includes_api_keys(self):
|
||||
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 TestStreamUsageAiChat:
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_with_tool_call(self):
|
||||
mock_tool_call = MagicMock()
|
||||
mock_tool_call.id = "call_123"
|
||||
mock_tool_call.function.name = "get_usage_data"
|
||||
mock_tool_call.function.arguments = json.dumps({
|
||||
"start_date": "2025-01-01",
|
||||
"end_date": "2025-01-31",
|
||||
})
|
||||
|
||||
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_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"
|
||||
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_usage_data", new_callable=AsyncMock) as mock_fetch:
|
||||
|
||||
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": "What is my total spend?"}],
|
||||
model="gpt-4o-mini",
|
||||
user_id="user-123",
|
||||
is_admin=True,
|
||||
):
|
||||
events.append(event)
|
||||
|
||||
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"]
|
||||
|
||||
done_event = json.loads(events[-1].replace("data: ", "").strip())
|
||||
assert done_event["type"] == "done"
|
||||
|
||||
@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."
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm") as mock_litellm:
|
||||
mock_litellm.acompletion = AsyncMock(return_value=mock_response)
|
||||
|
||||
events = []
|
||||
async for event in stream_usage_ai_chat(
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
):
|
||||
events.append(event)
|
||||
|
||||
assert len(events) >= 2
|
||||
chunk_event = json.loads(events[0].replace("data: ", "").strip())
|
||||
assert chunk_event["type"] == "chunk"
|
||||
assert "context" in chunk_event["content"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_handles_error(self):
|
||||
with patch("litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm") as mock_litellm:
|
||||
mock_litellm.acompletion = AsyncMock(side_effect=Exception("LLM error"))
|
||||
|
||||
events = []
|
||||
async for event in stream_usage_ai_chat(
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
):
|
||||
events.append(event)
|
||||
|
||||
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"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_admin_enforces_user_id(self):
|
||||
mock_tool_call = MagicMock()
|
||||
mock_tool_call.id = "call_456"
|
||||
mock_tool_call.function.name = "get_usage_data"
|
||||
mock_tool_call.function.arguments = json.dumps({
|
||||
"start_date": "2025-01-01",
|
||||
"end_date": "2025-01-31",
|
||||
"user_id": "other-user",
|
||||
})
|
||||
|
||||
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_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."
|
||||
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_usage_data", new_callable=AsyncMock) as mock_fetch:
|
||||
|
||||
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"}],
|
||||
model="gpt-4o-mini",
|
||||
user_id="my-user-id",
|
||||
is_admin=False,
|
||||
):
|
||||
events.append(event)
|
||||
|
||||
mock_fetch.assert_called_once_with(
|
||||
start_date="2025-01-01",
|
||||
end_date="2025-01-31",
|
||||
user_id="my-user-id",
|
||||
)
|
||||
|
|
@ -20,116 +20,13 @@ vi.mock("../../networking", () => ({
|
|||
{ model_group: "claude-3-opus" },
|
||||
],
|
||||
}),
|
||||
getProxyBaseUrl: vi.fn().mockReturnValue("http://localhost:4000"),
|
||||
usageAiChatStream: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("openai", () => {
|
||||
return {
|
||||
default: {
|
||||
OpenAI: vi.fn().mockImplementation(() => ({
|
||||
chat: {
|
||||
completions: {
|
||||
create: vi.fn(),
|
||||
},
|
||||
},
|
||||
})),
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
const mockUserSpendData = {
|
||||
results: [
|
||||
{
|
||||
date: "2025-01-01",
|
||||
metrics: {
|
||||
spend: 100.5,
|
||||
api_requests: 1000,
|
||||
successful_requests: 950,
|
||||
failed_requests: 50,
|
||||
total_tokens: 50000,
|
||||
prompt_tokens: 30000,
|
||||
completion_tokens: 20000,
|
||||
cache_read_input_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
},
|
||||
breakdown: {
|
||||
models: {
|
||||
"gpt-4": {
|
||||
metrics: {
|
||||
spend: 80.0,
|
||||
api_requests: 800,
|
||||
successful_requests: 780,
|
||||
failed_requests: 20,
|
||||
total_tokens: 40000,
|
||||
prompt_tokens: 24000,
|
||||
completion_tokens: 16000,
|
||||
cache_read_input_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
},
|
||||
metadata: {},
|
||||
api_key_breakdown: {},
|
||||
},
|
||||
},
|
||||
model_groups: {},
|
||||
mcp_servers: {},
|
||||
providers: {
|
||||
openai: {
|
||||
metrics: {
|
||||
spend: 100.5,
|
||||
api_requests: 1000,
|
||||
successful_requests: 950,
|
||||
failed_requests: 50,
|
||||
total_tokens: 50000,
|
||||
prompt_tokens: 30000,
|
||||
completion_tokens: 20000,
|
||||
cache_read_input_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
},
|
||||
metadata: {},
|
||||
api_key_breakdown: {},
|
||||
},
|
||||
},
|
||||
api_keys: {
|
||||
"sk-test": {
|
||||
metrics: {
|
||||
spend: 100.5,
|
||||
api_requests: 1000,
|
||||
successful_requests: 950,
|
||||
failed_requests: 50,
|
||||
total_tokens: 50000,
|
||||
prompt_tokens: 30000,
|
||||
completion_tokens: 20000,
|
||||
cache_read_input_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
},
|
||||
metadata: {
|
||||
key_alias: "Test Key",
|
||||
team_id: null,
|
||||
},
|
||||
},
|
||||
},
|
||||
entities: {},
|
||||
},
|
||||
},
|
||||
],
|
||||
metadata: {
|
||||
total_spend: 100.5,
|
||||
total_api_requests: 1000,
|
||||
total_successful_requests: 950,
|
||||
total_failed_requests: 50,
|
||||
total_tokens: 50000,
|
||||
},
|
||||
};
|
||||
|
||||
const defaultProps = {
|
||||
open: true,
|
||||
onClose: vi.fn(),
|
||||
accessToken: "test-token",
|
||||
userSpendData: mockUserSpendData,
|
||||
dateRange: {
|
||||
from: new Date("2025-01-01"),
|
||||
to: new Date("2025-01-07"),
|
||||
},
|
||||
};
|
||||
|
||||
describe("UsageAIChatPanel", () => {
|
||||
|
|
|
|||
|
|
@ -1,9 +1,7 @@
|
|||
import React, { useEffect, useMemo, useRef, useState } from "react";
|
||||
import React, { useEffect, useRef, useState } from "react";
|
||||
import { Select, Input, Spin } from "antd";
|
||||
import { Button } from "@tremor/react";
|
||||
import { getProxyBaseUrl, modelHubCall } from "../../networking";
|
||||
import { DailyData } from "../types";
|
||||
import openai from "openai";
|
||||
import { modelHubCall, usageAiChatStream } from "../../networking";
|
||||
|
||||
const { TextArea } = Input;
|
||||
|
||||
|
|
@ -16,103 +14,12 @@ interface UsageAIChatPanelProps {
|
|||
open: boolean;
|
||||
onClose: () => void;
|
||||
accessToken: string | null;
|
||||
userSpendData: {
|
||||
results: DailyData[];
|
||||
metadata: any;
|
||||
};
|
||||
dateRange: {
|
||||
from?: Date;
|
||||
to?: Date;
|
||||
};
|
||||
}
|
||||
|
||||
function buildUsageSummary(
|
||||
userSpendData: UsageAIChatPanelProps["userSpendData"],
|
||||
dateRange: UsageAIChatPanelProps["dateRange"]
|
||||
): string {
|
||||
const meta = userSpendData.metadata || {};
|
||||
const results = userSpendData.results || [];
|
||||
|
||||
const fromStr = dateRange.from?.toLocaleDateString() ?? "N/A";
|
||||
const toStr = dateRange.to?.toLocaleDateString() ?? "N/A";
|
||||
|
||||
const modelSpend: Record<string, { spend: number; requests: number; tokens: number }> = {};
|
||||
const providerSpend: Record<string, { spend: number; requests: number }> = {};
|
||||
const keySpend: Record<string, { spend: number; alias: string | null }> = {};
|
||||
|
||||
for (const day of results) {
|
||||
for (const [model, metrics] of Object.entries(day.breakdown.models || {})) {
|
||||
if (!modelSpend[model]) modelSpend[model] = { spend: 0, requests: 0, tokens: 0 };
|
||||
modelSpend[model].spend += metrics.metrics.spend;
|
||||
modelSpend[model].requests += metrics.metrics.api_requests;
|
||||
modelSpend[model].tokens += metrics.metrics.total_tokens;
|
||||
}
|
||||
for (const [provider, metrics] of Object.entries(day.breakdown.providers || {})) {
|
||||
if (!providerSpend[provider]) providerSpend[provider] = { spend: 0, requests: 0 };
|
||||
providerSpend[provider].spend += metrics.metrics.spend;
|
||||
providerSpend[provider].requests += metrics.metrics.api_requests;
|
||||
}
|
||||
for (const [key, metrics] of Object.entries(day.breakdown.api_keys || {})) {
|
||||
if (!keySpend[key]) keySpend[key] = { spend: 0, alias: metrics.metadata.key_alias };
|
||||
keySpend[key].spend += metrics.metrics.spend;
|
||||
}
|
||||
}
|
||||
|
||||
const topModels = Object.entries(modelSpend)
|
||||
.sort((a, b) => b[1].spend - a[1].spend)
|
||||
.slice(0, 10)
|
||||
.map(([name, d]) => ` - ${name}: $${d.spend.toFixed(4)} (${d.requests} requests, ${d.tokens} tokens)`)
|
||||
.join("\n");
|
||||
|
||||
const topProviders = Object.entries(providerSpend)
|
||||
.sort((a, b) => b[1].spend - a[1].spend)
|
||||
.slice(0, 10)
|
||||
.map(([name, d]) => ` - ${name}: $${d.spend.toFixed(4)} (${d.requests} requests)`)
|
||||
.join("\n");
|
||||
|
||||
const topKeys = Object.entries(keySpend)
|
||||
.sort((a, b) => b[1].spend - a[1].spend)
|
||||
.slice(0, 10)
|
||||
.map(([key, d]) => ` - ${d.alias || key}: $${d.spend.toFixed(4)}`)
|
||||
.join("\n");
|
||||
|
||||
const dailySummary = results
|
||||
.sort((a, b) => new Date(a.date).getTime() - new Date(b.date).getTime())
|
||||
.map((d) => ` - ${d.date}: $${d.metrics.spend.toFixed(4)} (${d.metrics.api_requests} requests)`)
|
||||
.join("\n");
|
||||
|
||||
return `Date Range: ${fromStr} to ${toStr}
|
||||
Total Spend: $${(meta.total_spend || 0).toFixed(4)}
|
||||
Total Requests: ${meta.total_api_requests || 0}
|
||||
Successful Requests: ${meta.total_successful_requests || 0}
|
||||
Failed Requests: ${meta.total_failed_requests || 0}
|
||||
Total Tokens: ${meta.total_tokens || 0}
|
||||
|
||||
Top Models by Spend:
|
||||
${topModels || " (no data)"}
|
||||
|
||||
Top Providers by Spend:
|
||||
${topProviders || " (no data)"}
|
||||
|
||||
Top API Keys by Spend:
|
||||
${topKeys || " (no data)"}
|
||||
|
||||
Daily Spend:
|
||||
${dailySummary || " (no data)"}`;
|
||||
}
|
||||
|
||||
const SYSTEM_PROMPT = `You are an AI assistant that helps users understand their LLM API usage data. You are embedded in the LiteLLM Usage dashboard.
|
||||
|
||||
You have access to the user's current usage data which is provided below. Use it to answer questions about their spending, model usage, API key activity, provider costs, request volumes, and trends.
|
||||
|
||||
Be concise and helpful. Use specific numbers from the data. When discussing costs, format them as dollar amounts. If the user asks about something not available in the data, let them know.`;
|
||||
|
||||
const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({
|
||||
open,
|
||||
onClose,
|
||||
accessToken,
|
||||
userSpendData,
|
||||
dateRange,
|
||||
}) => {
|
||||
const [messages, setMessages] = useState<ChatMessage[]>([]);
|
||||
const [inputText, setInputText] = useState("");
|
||||
|
|
@ -154,11 +61,6 @@ const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({
|
|||
}
|
||||
};
|
||||
|
||||
const usageSummary = useMemo(
|
||||
() => buildUsageSummary(userSpendData, dateRange),
|
||||
[userSpendData, dateRange]
|
||||
);
|
||||
|
||||
const handleSend = async () => {
|
||||
if (!accessToken || !inputText.trim() || !selectedModel || isLoading) return;
|
||||
|
||||
|
|
@ -172,44 +74,30 @@ const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({
|
|||
const abortController = new AbortController();
|
||||
abortControllerRef.current = abortController;
|
||||
|
||||
let accumulated = "";
|
||||
|
||||
try {
|
||||
const proxyBaseUrl = getProxyBaseUrl();
|
||||
const client = new openai.OpenAI({
|
||||
apiKey: accessToken,
|
||||
baseURL: proxyBaseUrl,
|
||||
dangerouslyAllowBrowser: true,
|
||||
});
|
||||
|
||||
const chatHistory = [
|
||||
{
|
||||
role: "system" as const,
|
||||
content: `${SYSTEM_PROMPT}\n\nCurrent Usage Data:\n${usageSummary}`,
|
||||
await usageAiChatStream(
|
||||
accessToken,
|
||||
updatedMessages.map((m) => ({ role: m.role, content: m.content })),
|
||||
selectedModel,
|
||||
(content: string) => {
|
||||
accumulated += content;
|
||||
setStreamingContent(accumulated);
|
||||
},
|
||||
...updatedMessages.map((m) => ({
|
||||
role: m.role as "user" | "assistant",
|
||||
content: m.content,
|
||||
})),
|
||||
];
|
||||
|
||||
const response = await client.chat.completions.create(
|
||||
{
|
||||
model: selectedModel,
|
||||
stream: true,
|
||||
messages: chatHistory,
|
||||
() => {
|
||||
setMessages((prev) => [...prev, { role: "assistant", content: accumulated }]);
|
||||
setStreamingContent("");
|
||||
},
|
||||
{ signal: abortController.signal }
|
||||
(errorMsg: string) => {
|
||||
setMessages((prev) => [
|
||||
...prev,
|
||||
{ role: "assistant", content: `Error: ${errorMsg}` },
|
||||
]);
|
||||
setStreamingContent("");
|
||||
},
|
||||
abortController.signal,
|
||||
);
|
||||
|
||||
let fullContent = "";
|
||||
for await (const chunk of response) {
|
||||
if (chunk.choices[0]?.delta?.content) {
|
||||
fullContent += chunk.choices[0].delta.content;
|
||||
setStreamingContent(fullContent);
|
||||
}
|
||||
}
|
||||
|
||||
setMessages((prev) => [...prev, { role: "assistant", content: fullContent }]);
|
||||
setStreamingContent("");
|
||||
} catch (error: any) {
|
||||
if (error?.name === "AbortError" || abortController.signal.aborted) {
|
||||
return;
|
||||
|
|
@ -293,7 +181,7 @@ const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({
|
|||
/>
|
||||
</div>
|
||||
|
||||
{/* Chat messages — fills remaining space */}
|
||||
{/* Chat messages */}
|
||||
<div className="flex-1 overflow-y-auto p-4 space-y-3 bg-gray-50">
|
||||
{messages.length === 0 && !streamingContent && (
|
||||
<div className="flex flex-col items-center justify-center h-full text-gray-400">
|
||||
|
|
@ -341,7 +229,7 @@ const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({
|
|||
<div ref={messagesEndRef} />
|
||||
</div>
|
||||
|
||||
{/* Input area — pinned to bottom */}
|
||||
{/* Input area */}
|
||||
<div className="px-4 py-3 border-t border-gray-200 bg-white flex-shrink-0">
|
||||
<div className="flex gap-2">
|
||||
<TextArea
|
||||
|
|
|
|||
|
|
@ -945,8 +945,6 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
|
|||
open={isAiChatOpen}
|
||||
onClose={() => setIsAiChatOpen(false)}
|
||||
accessToken={accessToken}
|
||||
userSpendData={userSpendData}
|
||||
dateRange={dateValue}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -5907,6 +5907,68 @@ export const enrichPolicyTemplateStream = async (
|
|||
}
|
||||
};
|
||||
|
||||
export const usageAiChatStream = async (
|
||||
accessToken: string,
|
||||
messages: { role: string; content: string }[],
|
||||
model: string,
|
||||
onChunk: (content: string) => void,
|
||||
onDone: () => void,
|
||||
onError?: (error: string) => void,
|
||||
signal?: AbortSignal,
|
||||
) => {
|
||||
const url = proxyBaseUrl
|
||||
? `${proxyBaseUrl}/usage/ai/chat`
|
||||
: `/usage/ai/chat`;
|
||||
|
||||
const response = await fetch(url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify({ messages, model }),
|
||||
signal,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const errorData = await response.json();
|
||||
const errorMessage = deriveErrorMessage(errorData);
|
||||
handleError(errorMessage);
|
||||
throw new Error(errorMessage);
|
||||
}
|
||||
|
||||
const reader = response.body?.getReader();
|
||||
if (!reader) throw new Error("No response body");
|
||||
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = "";
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
const lines = buffer.split("\n");
|
||||
buffer = lines.pop() || "";
|
||||
|
||||
for (const line of lines) {
|
||||
if (!line.startsWith("data: ")) continue;
|
||||
try {
|
||||
const event = JSON.parse(line.slice(6));
|
||||
if (event.type === "chunk") {
|
||||
onChunk(event.content);
|
||||
} else if (event.type === "done") {
|
||||
onDone();
|
||||
} else if (event.type === "error") {
|
||||
onError?.(event.message);
|
||||
}
|
||||
} catch {
|
||||
// skip malformed events
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
export const createPolicyCall = async (accessToken: string, policyData: any) => {
|
||||
try {
|
||||
const url = proxyBaseUrl ? `${proxyBaseUrl}/policies` : `/policies`;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue