From b7399de05bdcaf79ca4931c1f505d669e908aebc Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 24 Feb 2026 21:17:59 +0000 Subject: [PATCH] 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 --- .../usage_endpoints/ai_usage_chat.py | 333 +++++++++++++++--- .../usage_endpoints/test_ai_usage_chat.py | 262 +++++++------- .../UsagePage/components/UsageAIChatPanel.tsx | 13 +- .../src/components/networking.tsx | 3 + 4 files changed, 432 insertions(+), 179 deletions(-) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index fe36ced62bb..7331e4f86d4 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -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)}) diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py index e28798ea816..967fed1b489 100644 --- a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py +++ b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -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, diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx index 7ddce26d8b7..6209313c097 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx @@ -28,6 +28,7 @@ const UsageAIChatPanel: React.FC = ({ const [availableModels, setAvailableModels] = useState([]); const [isLoadingModels, setIsLoadingModels] = useState(false); const [streamingContent, setStreamingContent] = useState(""); + const [statusMessage, setStatusMessage] = useState(null); const messagesEndRef = useRef(null); const abortControllerRef = useRef(null); @@ -70,6 +71,7 @@ const UsageAIChatPanel: React.FC = ({ setInputText(""); setIsLoading(true); setStreamingContent(""); + setStatusMessage(null); const abortController = new AbortController(); abortControllerRef.current = abortController; @@ -82,20 +84,26 @@ const UsageAIChatPanel: React.FC = ({ 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 = ({ {isLoading && !streamingContent && (
-
+
+ {statusMessage && ( + {statusMessage} + )}
)} diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 4359ac91021..57e4030b818 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -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") {