From 360643e21315eb02ca790eb22effb87a6d48167b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 24 Feb 2026 16:40:04 -0800 Subject: [PATCH] [Feat] UI - Allow using AI to understand Usage patterns (#22042) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add Ask AI chat component to Usage page - Create UsageAIChatModal component with streaming chat interface - Integrate with existing model hub for model selection - Pass usage data context (spend, models, providers, keys) to AI - Add Ask AI button next to Export Data button in global view - Add tests for the new component and integration Co-authored-by: Ishaan Jaff * Convert Ask AI from modal to right-side sliding panel - Replace UsageAIChatModal with UsageAIChatPanel - Panel slides in from right side, usage page stays visible - Full-height panel with header, model selector, chat area, and input - Smooth CSS transition for open/close animation - Update tests for new panel component (34 tests passing) Co-authored-by: Ishaan Jaff * Remove build output directory from tracking Co-authored-by: Ishaan Jaff * 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 * Make model selection optional, default to gpt-4o-mini on backend Co-authored-by: Ishaan Jaff * 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 * Fix: inject today's date into system prompt so AI resolves relative dates correctly Co-authored-by: Ishaan Jaff * Show tool calls as distinct steps + render markdown in responses - Backend emits tool_call events with tool_name, label, args, and status - Frontend shows each tool call as a step with ✓/spinner/✗ indicator - Tool call steps show icon, label, date range, and filters - AI responses rendered with ReactMarkdown (bold, lists, tables, code) - Cursor-like UX: Thinking → tool calls → Analyzing → streamed answer Co-authored-by: Ishaan Jaff * Refactor backend for code quality: proper types, constants, all functions ≤50 LOC - TypedDict for SSE events (SSEStatusEvent, SSEToolCallEvent, etc.) and ToolHandler - Constants for table names, entity fields, temperature, page sizes, top-N limits - Shared _query_activity() eliminates duplicated fetch logic - _accumulate_breakdown() + _ranked_lines() replace inline aggregation loops - Extracted _process_tool_call() and _stream_final_response() from main stream fn - Black + Ruff clean, all 15 functions verified ≤50 LOC - Replaced Tremor Button with Antd Button in panel (Tremor deprecated per AGENTS.md) Co-authored-by: Ishaan Jaff * Address greptile review: security fixes and input validation - Restrict team/tag tools to admin-only users (non-admins only get get_usage_data) - Constrain ChatMessage.role to Literal['user', 'assistant'] to prevent system prompt injection - Add test for base tools restriction (non-admin gets 1 tool, admin gets 3) - Issues 3 (unused imports) and 4 (inline datetime) were already fixed in prior commit Co-authored-by: Ishaan Jaff * Address greptile round 2: sanitize errors, defense-in-depth allowlist, revert tsconfig - Sanitize error messages: generic 'An internal error occurred' sent to client, full exception logged server-side via verbose_proxy_logger - Defense-in-depth: _process_tool_call validates fn_name against role-based allowlist before dispatch (even though LLM only receives allowed tools) - Revert tsconfig.json jsx back to 'preserve' (Next.js recommended default) Co-authored-by: Ishaan Jaff * Role-scoped system prompt + additional test coverage - System prompt is now role-aware: admin sees all 3 tool descriptions, non-admin only sees get_usage_data (consistent with tool filtering) - Added tests: non-admin prompt excludes team/tag tools, date injection - 15 backend tests, 34 frontend tests passing Co-authored-by: Ishaan Jaff * Fix LLM arg validation + cap conversation size at 20 messages - _resolve_fetch_kwargs uses .get() with ValueError for missing dates (handles malformed LLM tool arguments gracefully) - MAX_CHAT_MESSAGES = 20 constant; backend truncates to last 20 - Frontend also sends only last 20 messages per request - Prevents excessive token usage and context-length errors Co-authored-by: Ishaan Jaff --------- Co-authored-by: Cursor Agent Co-authored-by: Ishaan Jaff --- .../usage_endpoints/__init__.py | 9 + .../usage_endpoints/ai_usage_chat.py | 578 ++++++++++++++++++ .../usage_endpoints/endpoints.py | 65 ++ litellm/proxy/proxy_server.py | 2 + .../usage_endpoints/__init__.py | 0 .../usage_endpoints/test_ai_usage_chat.py | 402 ++++++++++++ ui/litellm-dashboard/package-lock.json | 151 +---- .../components/UsageAIChatPanel.test.tsx | 85 +++ .../UsagePage/components/UsageAIChatPanel.tsx | 402 ++++++++++++ .../components/UsagePageView.test.tsx | 26 + .../UsagePage/components/UsagePageView.tsx | 51 +- .../src/components/networking.tsx | 76 +++ 12 files changed, 1704 insertions(+), 143 deletions(-) create mode 100644 litellm/proxy/management_endpoints/usage_endpoints/__init__.py create mode 100644 litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py create mode 100644 litellm/proxy/management_endpoints/usage_endpoints/endpoints.py create mode 100644 tests/test_litellm/proxy/management_endpoints/usage_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py create mode 100644 ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx create mode 100644 ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx diff --git a/litellm/proxy/management_endpoints/usage_endpoints/__init__.py b/litellm/proxy/management_endpoints/usage_endpoints/__init__.py new file mode 100644 index 0000000000..6e68dcd2a2 --- /dev/null +++ b/litellm/proxy/management_endpoints/usage_endpoints/__init__.py @@ -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, +) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py new file mode 100644 index 0000000000..f156be7d2c --- /dev/null +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -0,0 +1,578 @@ +""" +AI Usage Chat - uses LLM tool calling to answer questions about +usage/spend data by querying the aggregated daily activity endpoints. +""" + +import json +from datetime import date +from typing import Any, AsyncIterator, Callable, Dict, List, Literal, Optional + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + SpendAnalyticsPaginatedResponse, +) + +from typing_extensions import TypedDict + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +USAGE_AI_TEMPERATURE = 0.2 + +TABLE_DAILY_USER_SPEND = "litellm_dailyuserspend" +TABLE_DAILY_TEAM_SPEND = "litellm_dailyteamspend" +TABLE_DAILY_TAG_SPEND = "litellm_dailytagspend" + +ENTITY_FIELD_USER = "user_id" +ENTITY_FIELD_TEAM = "team_id" +ENTITY_FIELD_TAG = "tag" + +PAGINATED_PAGE_SIZE = 200 +MAX_CHAT_MESSAGES = 20 +TOP_N_MODELS = 15 +TOP_N_PROVIDERS = 10 +TOP_N_KEYS = 10 + +# --------------------------------------------------------------------------- +# Types +# --------------------------------------------------------------------------- + + +class SSEStatusEvent(TypedDict): + type: Literal["status"] + message: str + + +class SSEToolCallEvent(TypedDict, total=False): + type: Literal["tool_call"] + tool_name: str + tool_label: str + arguments: Dict[str, str] + status: Literal["running", "complete", "error"] + error: str + + +class SSEChunkEvent(TypedDict): + type: Literal["chunk"] + content: str + + +class SSEDoneEvent(TypedDict): + type: Literal["done"] + + +class SSEErrorEvent(TypedDict): + type: Literal["error"] + message: str + + +SSEEvent = ( + SSEStatusEvent | SSEToolCallEvent | SSEChunkEvent | SSEDoneEvent | SSEErrorEvent +) + + +class ToolHandler(TypedDict): + fetch: Callable[..., Any] + summarise: Callable[[Dict[str, Any]], str] + label: str + + +# --------------------------------------------------------------------------- +# Tool definitions (OpenAI function-calling schema) +# --------------------------------------------------------------------------- + +_DATE_PARAMS = { + "start_date": {"type": "string", "description": "Start date in YYYY-MM-DD format"}, + "end_date": {"type": "string", "description": "End date in YYYY-MM-DD format"}, +} + +_TOOL_USAGE = { + "type": "function", + "function": { + "name": "get_usage_data", + "description": ( + "Fetch aggregated global usage/spend data. Returns daily spend, " + "token counts, request counts, and breakdowns by model, provider, " + "and API key. Use for overall spend, top models, top providers." + ), + "parameters": { + "type": "object", + "properties": { + **_DATE_PARAMS, + "user_id": { + "type": "string", + "description": "Optional user ID filter. Omit for global view.", + }, + }, + "required": ["start_date", "end_date"], + }, + }, +} + +_TOOL_TEAM = { + "type": "function", + "function": { + "name": "get_team_usage_data", + "description": ( + "Fetch usage/spend data broken down by team. Use for questions " + "like 'which team spends the most' or 'show me team X usage'." + ), + "parameters": { + "type": "object", + "properties": { + **_DATE_PARAMS, + "team_ids": { + "type": "string", + "description": "Optional comma-separated team IDs. Omit for all teams.", + }, + }, + "required": ["start_date", "end_date"], + }, + }, +} + +_TOOL_TAG = { + "type": "function", + "function": { + "name": "get_tag_usage_data", + "description": ( + "Fetch usage/spend data broken down by tag. Tags are labels " + "attached to requests (features, environments, credentials)." + ), + "parameters": { + "type": "object", + "properties": { + **_DATE_PARAMS, + "tags": { + "type": "string", + "description": "Optional comma-separated tag names. Omit for all tags.", + }, + }, + "required": ["start_date", "end_date"], + }, + }, +} + +TOOLS_BASE = [_TOOL_USAGE] +TOOLS_ADMIN = [_TOOL_USAGE, _TOOL_TEAM, _TOOL_TAG] + + +def get_tools_for_role(is_admin: bool) -> List[Dict[str, Any]]: + """Return the tool list appropriate for the user's role.""" + return TOOLS_ADMIN if is_admin else TOOLS_BASE + + +_SYSTEM_PROMPT_BASE = ( + "You are an AI assistant embedded in the LiteLLM Usage dashboard. " + "You help users understand their LLM API spend and usage data.\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 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 will be provided below. Use it to interpret relative dates " + "like 'this week', 'this month', 'last 7 days', etc." +) + +_TOOL_DESCRIPTIONS_ADMIN = ( + "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" +) + +_TOOL_DESCRIPTIONS_BASE = ( + "You have access to this tool:\n" + "- `get_usage_data`: Your usage data (spend, models, providers, API keys)\n\n" +) + + +def _build_system_prompt(is_admin: bool) -> str: + """Build role-appropriate system prompt with today's date.""" + tool_desc = _TOOL_DESCRIPTIONS_ADMIN if is_admin else _TOOL_DESCRIPTIONS_BASE + return ( + f"{_SYSTEM_PROMPT_BASE}\n\n{tool_desc}" + f"Today's date: {date.today().isoformat()}" + ) + + +# keep a public reference for test assertions +SYSTEM_PROMPT = _SYSTEM_PROMPT_BASE + +# --------------------------------------------------------------------------- +# Data fetchers +# --------------------------------------------------------------------------- + + +def _parse_csv_ids(raw: Optional[str]) -> Optional[List[str]]: + if not raw: + return None + return [t.strip() for t in raw.split(",") if t.strip()] + + +async def _query_activity( + table_name: str, + entity_id_field: str, + entity_id: Optional[Any], + start_date: str, + end_date: str, + *, + use_aggregated: bool = False, +) -> SpendAnalyticsPaginatedResponse: + """Shared helper that calls the daily activity query layer.""" + from litellm.proxy.management_endpoints.common_daily_activity import ( + get_daily_activity, + get_daily_activity_aggregated, + ) + from litellm.proxy.proxy_server import prisma_client + + if use_aggregated: + return await get_daily_activity_aggregated( + prisma_client=prisma_client, + table_name=table_name, + entity_id_field=entity_id_field, + entity_id=entity_id, + entity_metadata_field=None, + start_date=start_date, + end_date=end_date, + model=None, + api_key=None, + ) + return await get_daily_activity( + prisma_client=prisma_client, + table_name=table_name, + entity_id_field=entity_id_field, + entity_id=entity_id, + entity_metadata_field=None, + start_date=start_date, + end_date=end_date, + model=None, + api_key=None, + page=1, + page_size=PAGINATED_PAGE_SIZE, + ) + + +async def _fetch_usage_data( + start_date: str, end_date: str, user_id: Optional[str] = None +) -> Dict[str, Any]: + resp = await _query_activity( + TABLE_DAILY_USER_SPEND, + ENTITY_FIELD_USER, + user_id, + start_date, + end_date, + use_aggregated=True, + ) + return resp.model_dump(mode="json") + + +async def _fetch_team_usage_data( + start_date: str, end_date: str, team_ids: Optional[str] = None +) -> Dict[str, Any]: + resp = await _query_activity( + TABLE_DAILY_TEAM_SPEND, + ENTITY_FIELD_TEAM, + _parse_csv_ids(team_ids), + start_date, + end_date, + ) + return resp.model_dump(mode="json") + + +async def _fetch_tag_usage_data( + start_date: str, end_date: str, tags: Optional[str] = None +) -> Dict[str, Any]: + resp = await _query_activity( + TABLE_DAILY_TAG_SPEND, + ENTITY_FIELD_TAG, + _parse_csv_ids(tags), + start_date, + end_date, + ) + return resp.model_dump(mode="json") + + +# --------------------------------------------------------------------------- +# Summarisers — convert raw JSON to concise text the LLM can reason over +# --------------------------------------------------------------------------- + + +def _accumulate_breakdown( + results: List[Dict[str, Any]], dimension: str, fields: List[str] +) -> Dict[str, Dict[str, float]]: + """Aggregate a single breakdown dimension across days.""" + totals: Dict[str, Dict[str, float]] = {} + for day in results: + for key, entry in day.get("breakdown", {}).get(dimension, {}).items(): + if key not in totals: + totals[key] = {f: 0.0 for f in fields} + m = entry.get("metrics", {}) + for f in fields: + totals[key][f] += m.get(f, 0) + return totals + + +def _ranked_lines( + totals: Dict[str, Dict[str, float]], + fmt: Callable[[str, Dict[str, float]], str], + limit: int, +) -> List[str]: + """Sort by spend descending, format each entry, and truncate.""" + return [ + fmt(name, vals) + for name, vals in sorted(totals.items(), key=lambda x: -x[1].get("spend", 0))[ + :limit + ] + ] + + +def _summarise_usage_data(data: Dict[str, Any]) -> str: + meta = data.get("metadata", {}) + results = data.get("results", []) + + header = ( + f"Total Spend: ${meta.get('total_spend', 0):.4f}\n" + f"Total Requests: {meta.get('total_api_requests', 0)}\n" + f"Successful: {meta.get('total_successful_requests', 0)} | " + f"Failed: {meta.get('total_failed_requests', 0)}\n" + f"Total Tokens: {meta.get('total_tokens', 0)}" + ) + + models = _accumulate_breakdown( + results, "models", ["spend", "api_requests", "total_tokens"] + ) + providers = _accumulate_breakdown(results, "providers", ["spend", "api_requests"]) + + model_lines = _ranked_lines( + models, + lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs, {int(d['total_tokens'])} tokens)", + TOP_N_MODELS, + ) + provider_lines = _ranked_lines( + providers, + lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs)", + TOP_N_PROVIDERS, + ) + + sections = [header, ""] + sections += ["Top Models by Spend:"] + (model_lines or [" (no data)"]) + [""] + sections += ["Top Providers by Spend:"] + (provider_lines or [" (no data)"]) + return "\n".join(sections) + + +def _summarise_entity_data(data: Dict[str, Any], entity_label: str) -> str: + """Summarise team/tag entity usage data.""" + results = data.get("results", []) + if not results: + return f"No {entity_label} usage data found for the given date range." + + totals: Dict[str, Dict[str, Any]] = {} + for day in results: + for eid, entry in day.get("breakdown", {}).get("entities", {}).items(): + if eid not in totals: + alias = entry.get("metadata", {}).get("alias", eid) + totals[eid] = {"alias": alias, "spend": 0.0, "requests": 0, "tokens": 0} + m = entry.get("metrics", {}) + totals[eid]["spend"] += m.get("spend", 0) + totals[eid]["requests"] += m.get("api_requests", 0) + totals[eid]["tokens"] += m.get("total_tokens", 0) + + lines = [f"{entity_label} Usage ({len(totals)} {entity_label.lower()}s):", ""] + for eid, d in sorted(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} | " + f"{int(d['requests'])} reqs | {int(d['tokens'])} tokens" + ) + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Tool dispatch registry +# --------------------------------------------------------------------------- + +TOOL_HANDLERS: Dict[str, ToolHandler] = { + "get_usage_data": ToolHandler( + fetch=_fetch_usage_data, + summarise=_summarise_usage_data, + label="global usage data", + ), + "get_team_usage_data": ToolHandler( + fetch=_fetch_team_usage_data, + summarise=lambda data: _summarise_entity_data(data, "Team"), + label="team usage data", + ), + "get_tag_usage_data": ToolHandler( + fetch=_fetch_tag_usage_data, + summarise=lambda data: _summarise_entity_data(data, "Tag"), + label="tag usage data", + ), +} + + +# --------------------------------------------------------------------------- +# SSE streaming +# --------------------------------------------------------------------------- + + +def _sse(event: SSEEvent) -> str: + return f"data: {json.dumps(event)}\n\n" + + +def _resolve_fetch_kwargs( + fn_name: str, + fn_args: Dict[str, str], + user_id: Optional[str], + is_admin: bool, +) -> Dict[str, Any]: + """Build keyword arguments for a tool's fetch function.""" + start_date = fn_args.get("start_date", "") + end_date = fn_args.get("end_date", "") + if not start_date or not end_date: + raise ValueError("Missing required start_date or end_date from tool arguments") + kwargs: Dict[str, Any] = {"start_date": start_date, "end_date": end_date} + if fn_name == "get_usage_data": + if not is_admin: + kwargs["user_id"] = user_id + elif fn_args.get("user_id"): + kwargs["user_id"] = fn_args["user_id"] + elif fn_name == "get_team_usage_data" and fn_args.get("team_ids"): + kwargs["team_ids"] = fn_args["team_ids"] + elif fn_name == "get_tag_usage_data" and fn_args.get("tags"): + kwargs["tags"] = fn_args["tags"] + return kwargs + + +async def _execute_tool_call( + handler: ToolHandler, + fn_name: str, + fn_args: Dict[str, str], + user_id: Optional[str], + is_admin: bool, +) -> str: + """Run a single tool and return the summarised result text.""" + kwargs = _resolve_fetch_kwargs(fn_name, fn_args, user_id, is_admin) + raw_data = await handler["fetch"](**kwargs) + return handler["summarise"](raw_data) + + +async def _process_tool_call( + tc: Any, + chat_messages: List[Dict[str, Any]], + user_id: Optional[str], + is_admin: bool, +) -> AsyncIterator[str]: + """Execute a single tool call, yielding SSE events for status.""" + fn_name = tc.function.name + fn_args = json.loads(tc.function.arguments) + + allowed_names = {t["function"]["name"] for t in get_tools_for_role(is_admin)} + handler = TOOL_HANDLERS.get(fn_name) + + if fn_name not in allowed_names or not handler: + chat_messages.append( + { + "role": "tool", + "tool_call_id": tc.id, + "content": f"Tool not available: {fn_name}", + } + ) + return + + tool_event_base = { + "type": "tool_call", + "tool_name": fn_name, + "tool_label": handler["label"], + "arguments": fn_args, + } + yield _sse({**tool_event_base, "status": "running"}) + + try: + tool_result = await _execute_tool_call( + handler, fn_name, fn_args, user_id, is_admin + ) + yield _sse({**tool_event_base, "status": "complete"}) + except Exception as e: + verbose_proxy_logger.error("Tool %s failed: %s", fn_name, e) + tool_result = f"Error fetching {handler['label']}. Please try again." + yield _sse({**tool_event_base, "status": "error"}) + + chat_messages.append( + {"role": "tool", "tool_call_id": tc.id, "content": tool_result} + ) + + +async def _stream_final_response( + model: str, chat_messages: List[Dict[str, Any]] +) -> AsyncIterator[str]: + """Stream the final LLM response after tool results are appended.""" + yield _sse({"type": "status", "message": "Analyzing results..."}) + + response = await litellm.acompletion( + model=model, + messages=chat_messages, + stream=True, + temperature=USAGE_AI_TEMPERATURE, + ) + async for chunk in response: + delta = chunk.choices[0].delta.content + if delta: + yield _sse({"type": "chunk", "content": delta}) + + +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 SSE events: status → tool_call → chunk → done.""" + resolved_model = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL + truncated = ( + messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages + ) + chat_messages: List[Dict[str, Any]] = [ + {"role": "system", "content": _build_system_prompt(is_admin)}, + *truncated, + ] + + try: + yield _sse({"type": "status", "message": "Thinking..."}) + tools = get_tools_for_role(is_admin) + response = await litellm.acompletion( + model=resolved_model, + messages=chat_messages, + tools=tools, + temperature=USAGE_AI_TEMPERATURE, + ) + choice = response.choices[0] # type: ignore + + if not choice.message.tool_calls: + if choice.message.content: + yield _sse({"type": "chunk", "content": choice.message.content}) + yield _sse({"type": "done"}) + return + + chat_messages.append(choice.message.model_dump()) + for tc in choice.message.tool_calls: + async for event in _process_tool_call(tc, chat_messages, user_id, is_admin): + yield event + async for event in _stream_final_response(resolved_model, chat_messages): + yield event + yield _sse({"type": "done"}) + + except Exception as e: + verbose_proxy_logger.error("AI usage chat failed: %s", e) + yield _sse( + { + "type": "error", + "message": "An internal error occurred. Please try again.", + } + ) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py new file mode 100644 index 0000000000..0dbe518afb --- /dev/null +++ b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py @@ -0,0 +1,65 @@ +""" +USAGE AI CHAT ENDPOINTS + +/usage/ai/chat - Stream AI chat responses about usage data +""" + +from typing import List, Literal, Optional + +from fastapi import APIRouter, Depends, Request +from fastapi.responses import StreamingResponse +from pydantic import BaseModel, Field + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router = APIRouter() + + +class ChatMessage(BaseModel): + role: Literal["user", "assistant"] + 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 tools that query aggregated daily activity data. + """ + from litellm.proxy.management_endpoints.common_utils import ( + _user_has_admin_view, + ) + from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( + stream_usage_ai_chat, + ) + + 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"}, + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 36b6bcc770..607306f380 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -392,6 +392,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, ) @@ -12872,6 +12873,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) diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/__init__.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/__init__.py new file mode 100644 index 0000000000..e69de29bb2 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 new file mode 100644 index 0000000000..f9303bd13a --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -0,0 +1,402 @@ +""" +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 ( + TOOL_HANDLERS, + TOOLS_ADMIN, + TOOLS_BASE, + _build_system_prompt, + _summarise_entity_data, + _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, + }, + "metadata": {}, + "api_key_breakdown": {}, + }, + }, + "providers": { + "openai": { + "metrics": {"spend": 50.25, "api_requests": 500}, + "metadata": {}, + "api_key_breakdown": {}, + }, + }, + "api_keys": { + "sk-test123": { + "metrics": {"spend": 50.25}, + "metadata": {"key_alias": "Production Key"}, + }, + }, + "model_groups": {}, + "mcp_servers": {}, + "entities": {}, + }, + }, + ], + "metadata": { + "total_spend": 50.25, + "total_api_requests": 500, + "total_successful_requests": 480, + "total_failed_requests": 20, + "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 TestToolSchemas: + def test_admin_tools_include_all(self): + assert len(TOOLS_ADMIN) == 3 + names = {t["function"]["name"] for t in TOOLS_ADMIN} + assert "get_usage_data" in names + assert "get_team_usage_data" in names + assert "get_tag_usage_data" in names + + def test_base_tools_restricted_to_usage_only(self): + assert len(TOOLS_BASE) == 1 + assert TOOLS_BASE[0]["function"]["name"] == "get_usage_data" + + def test_admin_prompt_mentions_all_tools(self): + prompt = _build_system_prompt(is_admin=True) + assert "get_usage_data" in prompt + assert "get_team_usage_data" in prompt + assert "get_tag_usage_data" in prompt + + def test_non_admin_prompt_only_mentions_usage_tool(self): + prompt = _build_system_prompt(is_admin=False) + assert "get_usage_data" in prompt + assert "get_team_usage_data" not in prompt + assert "get_tag_usage_data" not in prompt + + def test_system_prompt_includes_todays_date(self): + from datetime import date + + prompt = _build_system_prompt(is_admin=True) + assert date.today().isoformat() in prompt + + +class TestSummariseUsageData: + def test_summarise_includes_totals(self): + summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) + assert "$50.25" in summary + assert "500" in summary + + def test_summarise_includes_models(self): + summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) + assert "gpt-4" in summary + + def test_summarise_includes_providers(self): + summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) + assert "openai" 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_emits_status_events(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 = "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(json.loads(event.replace("data: ", "").strip())) + + status_events = [e for e in events if e["type"] == "status"] + tool_call_events = [e for e in events if e["type"] == "tool_call"] + chunk_events = [e for e in events if e["type"] == "chunk"] + done_events = [e for e in events if e["type"] == "done"] + + assert len(status_events) >= 1 + assert "Thinking" in status_events[0]["message"] + assert len(tool_call_events) >= 1 + assert tool_call_events[0]["tool_name"] == "get_usage_data" + assert tool_call_events[0]["status"] in ("running", "complete") + assert len(chunk_events) >= 1 + assert len(done_events) == 1 + + @pytest.mark.asyncio + 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", + } + ) + + 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": "Which team spends the most?"}], + model="gpt-4o-mini", + is_admin=True, + ): + events.append(json.loads(event.replace("data: ", "").strip())) + + 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): + 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(json.loads(event.replace("data: ", "").strip())) + + error_events = [e for e in events if e["type"] == "error"] + assert len(error_events) == 1 + assert "internal error" in error_events[0]["message"].lower() + + @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 = "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.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(), + ] + ) + + events = [] + async for event in stream_usage_ai_chat( + messages=[{"role": "user", "content": "Show 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", + ) diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 489a39a7ee..3787f451ad 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -1757,29 +1757,6 @@ "url": "https://opencollective.com/libvips" } }, - "node_modules/@isaacs/balanced-match": { - "version": "4.0.1", - "resolved": "https://registry.npmjs.org/@isaacs/balanced-match/-/balanced-match-4.0.1.tgz", - "integrity": "sha512-yzMTt9lEb8Gv7zRioUilSglI0c0smZ9k5D65677DLWLtWJaXIS3CqcGyUFByYKlnUj6TkjLVs54fBl6+TiGQDQ==", - "dev": true, - "license": "MIT", - "engines": { - "node": "20 || >=22" - } - }, - "node_modules/@isaacs/brace-expansion": { - "version": "5.0.1", - "resolved": "https://registry.npmjs.org/@isaacs/brace-expansion/-/brace-expansion-5.0.1.tgz", - "integrity": "sha512-WMz71T1JS624nWj2n2fnYAuPovhv7EUhk69R6i9dsVyzxt5eM3bjwvgk9L+APE1TRscGysAVMANkB0jh0LQZrQ==", - "dev": true, - "license": "MIT", - "dependencies": { - "@isaacs/balanced-match": "^4.0.1" - }, - "engines": { - "node": "20 || >=22" - } - }, "node_modules/@istanbuljs/schema": { "version": "0.1.3", "resolved": "https://registry.npmjs.org/@istanbuljs/schema/-/schema-0.1.3.tgz", @@ -3696,32 +3673,6 @@ "typescript": ">=4.8.4 <6.0.0" } }, - "node_modules/@typescript-eslint/typescript-estree/node_modules/brace-expansion": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.0.2.tgz", - "integrity": "sha512-Jt0vHyM+jmUBqojB7E1NIYadt0vI0Qxjxd2TErW94wDz+E2LAm5vKMXXwg6ZZBTHPuUlDgQHKXvjGBdfcF1ZDQ==", - "dev": true, - "license": "MIT", - "dependencies": { - "balanced-match": "^1.0.0" - } - }, - "node_modules/@typescript-eslint/typescript-estree/node_modules/minimatch": { - "version": "9.0.5", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-9.0.5.tgz", - "integrity": "sha512-G6T0ZX48xgozx7587koeX9Ys2NYy6Gmv//P89sEte9V9whIapMNF4idKxnW2QtCcLiTWlb/wfCabAtAFWhhBow==", - "dev": true, - "license": "ISC", - "dependencies": { - "brace-expansion": "^2.0.1" - }, - "engines": { - "node": ">=16 || 14 >=14.17" - }, - "funding": { - "url": "https://github.com/sponsors/isaacs" - } - }, "node_modules/@typescript-eslint/utils": { "version": "8.54.0", "resolved": "https://registry.npmjs.org/@typescript-eslint/utils/-/utils-8.54.0.tgz", @@ -4749,11 +4700,14 @@ } }, "node_modules/balanced-match": { - "version": "1.0.2", - "resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz", - "integrity": "sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw==", + "version": "4.0.4", + "resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-4.0.4.tgz", + "integrity": "sha512-BLrgEcRTwX2o6gGxGOCNyMvGSp35YofuYzw9h1IMTRmKqttAZZVU67bdb9Pr2vUHA8+j3i2tJfjO6C6+4myGTA==", "dev": true, - "license": "MIT" + "license": "MIT", + "engines": { + "node": "18 || 20 || >=22" + } }, "node_modules/baseline-browser-mapping": { "version": "2.9.19", @@ -4787,14 +4741,16 @@ } }, "node_modules/brace-expansion": { - "version": "1.1.12", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.12.tgz", - "integrity": "sha512-9T9UjW3r0UW5c1Q7GTwllptXwhvYmEzFhzMfZ9H7FQWt+uZePjZPjBP/W1ZEyZ1twGWom5/56TF4lPcqjnDHcg==", + "version": "5.0.3", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.3.tgz", + "integrity": "sha512-fy6KJm2RawA5RcHkLa1z/ScpBeA762UF9KmZQxwIbDtRJrgLzM10depAiEQ+CXYcoiqW1/m96OAAoke2nE9EeA==", "dev": true, "license": "MIT", "dependencies": { - "balanced-match": "^1.0.0", - "concat-map": "0.0.1" + "balanced-match": "^4.0.2" + }, + "engines": { + "node": "18 || 20 || >=22" } }, "node_modules/braces": { @@ -5149,13 +5105,6 @@ "integrity": "sha512-VRhuHOLoKYOy4UbilLbUzbYg93XLjv2PncJC50EuTWPA3gaja1UjBsUP/D/9/juV3vQFr6XBEzn9KCAHdUvOHw==", "license": "MIT" }, - "node_modules/concat-map": { - "version": "0.0.1", - "resolved": "https://registry.npmjs.org/concat-map/-/concat-map-0.0.1.tgz", - "integrity": "sha512-/Srv4dswyQNBfohGpz9o6Yb3Gz3SrUDqBH5rTuhGR7ahtlbYKnVxw2bCFMRljaA7EXHaXZ8wsHdodFvbkhKmqg==", - "dev": true, - "license": "MIT" - }, "node_modules/copy-to-clipboard": { "version": "3.3.3", "resolved": "https://registry.npmjs.org/copy-to-clipboard/-/copy-to-clipboard-3.3.3.tgz", @@ -6924,22 +6873,6 @@ "node": ">=10.13.0" } }, - "node_modules/glob/node_modules/minimatch": { - "version": "10.1.1", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.1.1.tgz", - "integrity": "sha512-enIvLvRAFZYXJzkCYG5RKmPfrFArdLv+R+lbQ53BmIMLIry74bjKzX6iHAm8WYamJkhSSEabrWN5D97XnKObjQ==", - "dev": true, - "license": "BlueOak-1.0.0", - "dependencies": { - "@isaacs/brace-expansion": "^5.0.0" - }, - "engines": { - "node": "20 || >=22" - }, - "funding": { - "url": "https://github.com/sponsors/isaacs" - } - }, "node_modules/globals": { "version": "14.0.0", "resolved": "https://registry.npmjs.org/globals/-/globals-14.0.0.tgz", @@ -9035,16 +8968,19 @@ } }, "node_modules/minimatch": { - "version": "3.1.2", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-3.1.2.tgz", - "integrity": "sha512-J7p63hRiAjw1NDEww1W7i37+ByIrOWO5XQQAzZ3VOcL0PNybwpfmV/N05zFAzwQ9USyEcX6t3UO+K5aqBQOIHw==", + "version": "10.2.2", + "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.2.2.tgz", + "integrity": "sha512-+G4CpNBxa5MprY+04MbgOw1v7So6n5JY166pFi9KfYwT78fxScCeSNQSNzp6dpPSW2rONOps6Ocam1wFhCgoVw==", "dev": true, - "license": "ISC", + "license": "BlueOak-1.0.0", "dependencies": { - "brace-expansion": "^1.1.7" + "brace-expansion": "^5.0.2" }, "engines": { - "node": "*" + "node": "18 || 20 || >=22" + }, + "funding": { + "url": "https://github.com/sponsors/isaacs" } }, "node_modules/minimist": { @@ -12004,32 +11940,6 @@ "node": ">=18" } }, - "node_modules/test-exclude/node_modules/brace-expansion": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.0.2.tgz", - "integrity": "sha512-Jt0vHyM+jmUBqojB7E1NIYadt0vI0Qxjxd2TErW94wDz+E2LAm5vKMXXwg6ZZBTHPuUlDgQHKXvjGBdfcF1ZDQ==", - "dev": true, - "license": "MIT", - "dependencies": { - "balanced-match": "^1.0.0" - } - }, - "node_modules/test-exclude/node_modules/minimatch": { - "version": "9.0.5", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-9.0.5.tgz", - "integrity": "sha512-G6T0ZX48xgozx7587koeX9Ys2NYy6Gmv//P89sEte9V9whIapMNF4idKxnW2QtCcLiTWlb/wfCabAtAFWhhBow==", - "dev": true, - "license": "ISC", - "dependencies": { - "brace-expansion": "^2.0.1" - }, - "engines": { - "node": ">=16 || 14 >=14.17" - }, - "funding": { - "url": "https://github.com/sponsors/isaacs" - } - }, "node_modules/thenify": { "version": "3.3.1", "resolved": "https://registry.npmjs.org/thenify/-/thenify-3.3.1.tgz", @@ -13085,21 +12995,6 @@ "type": "github", "url": "https://github.com/sponsors/wooorm" } - }, - "node_modules/@next/swc-win32-ia32-msvc": { - "version": "14.2.33", - "resolved": "https://registry.npmjs.org/@next/swc-win32-ia32-msvc/-/swc-win32-ia32-msvc-14.2.33.tgz", - "integrity": "sha512-pc9LpGNKhJ0dXQhZ5QMmYxtARwwmWLpeocFmVG5Z0DzWq5Uf0izcI8tLc+qOpqxO1PWqZ5A7J1blrUIKrIFc7Q==", - "cpu": [ - "ia32" - ], - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": ">= 10" - } } } } diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx new file mode 100644 index 0000000000..85ce11e605 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx @@ -0,0 +1,85 @@ +import { screen } from "@testing-library/react"; +import { beforeAll, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../tests/test-utils"; +import UsageAIChatPanel from "./UsageAIChatPanel"; + +beforeAll(() => { + if (typeof window !== "undefined" && !window.ResizeObserver) { + window.ResizeObserver = class ResizeObserver { + observe() {} + unobserve() {} + disconnect() {} + } as any; + } +}); + +vi.mock("../../networking", () => ({ + modelHubCall: vi.fn().mockResolvedValue({ + data: [ + { model_group: "gpt-4" }, + { model_group: "claude-3-opus" }, + ], + }), + usageAiChatStream: vi.fn(), +})); + +const defaultProps = { + open: true, + onClose: vi.fn(), + accessToken: "test-token", +}; + +describe("UsageAIChatPanel", () => { + it("should render the panel when open", () => { + renderWithProviders(); + + expect(screen.getByText("Ask AI")).toBeInTheDocument(); + expect( + screen.getByText("Ask about your spend, models, keys, and trends") + ).toBeInTheDocument(); + }); + + it("should render model selector", () => { + renderWithProviders(); + + expect(screen.getByText("Select a model (optional, defaults to gpt-4o-mini)")).toBeInTheDocument(); + }); + + it("should render empty state message when no conversation", () => { + renderWithProviders(); + + expect(screen.getByText("Ask a question about your usage")).toBeInTheDocument(); + }); + + it("should render the send button", () => { + renderWithProviders(); + + expect(screen.getByText("Send")).toBeInTheDocument(); + }); + + it("should render input placeholder", () => { + renderWithProviders(); + + expect(screen.getByPlaceholderText("Ask about your usage...")).toBeInTheDocument(); + }); + + it("should render clear chat button", () => { + renderWithProviders(); + + expect(screen.getByText("Clear chat")).toBeInTheDocument(); + }); + + it("should have the panel element even when closed (just off-screen)", () => { + renderWithProviders(); + + expect(screen.getByTestId("usage-ai-chat-panel")).toBeInTheDocument(); + expect(screen.getByTestId("usage-ai-chat-panel")).toHaveClass("translate-x-full"); + }); + + it("should not have translate-x-full class when open", () => { + renderWithProviders(); + + expect(screen.getByTestId("usage-ai-chat-panel")).not.toHaveClass("translate-x-full"); + expect(screen.getByTestId("usage-ai-chat-panel")).toHaveClass("translate-x-0"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx new file mode 100644 index 0000000000..85f46bfa34 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx @@ -0,0 +1,402 @@ +import React, { useEffect, useRef, useState } from "react"; +import { Button, Select, Input, Spin } from "antd"; +import ReactMarkdown from "react-markdown"; +import { modelHubCall, usageAiChatStream, UsageAiToolCallEvent } from "../../networking"; + +const { TextArea } = Input; + +interface ToolCallStep { + tool_name: string; + tool_label: string; + arguments: Record; + status: "running" | "complete" | "error"; + error?: string; +} + +interface ChatMessage { + role: "user" | "assistant"; + content: string; + toolCalls?: ToolCallStep[]; +} + +interface UsageAIChatPanelProps { + open: boolean; + onClose: () => void; + accessToken: string | null; +} + +const TOOL_ICONS: Record = { + get_usage_data: "📊", + get_team_usage_data: "👥", + get_tag_usage_data: "🏷️", +}; + +const ToolCallDisplay: React.FC<{ step: ToolCallStep }> = ({ step }) => { + const icon = TOOL_ICONS[step.tool_name] || "🔧"; + const args = step.arguments; + const dateRange = args.start_date && args.end_date + ? `${args.start_date} → ${args.end_date}` + : ""; + const filter = args.team_ids || args.tags || args.user_id || ""; + + return ( +
+ + {step.status === "running" ? ( + + ) : step.status === "error" ? ( + ✗ + ) : ( + ✓ + )} + +
+
+ {icon} {step.tool_label} +
+ {dateRange && ( +
{dateRange}
+ )} + {filter && ( +
Filter: {filter}
+ )} + {step.status === "error" && step.error && ( +
{step.error}
+ )} +
+
+ ); +}; + +const MarkdownContent: React.FC<{ content: string }> = ({ content }) => ( +

{children}

, + strong: ({ children }) => {children}, + ul: ({ children }) =>
    {children}
, + ol: ({ children }) =>
    {children}
, + li: ({ children }) =>
  • {children}
    • , + h1: ({ children }) =>

    {children}

    , + h2: ({ children }) =>

    {children}

    , + h3: ({ children }) =>

    {children}

    , + code: ({ children, className }) => { + const isBlock = className?.includes("language-"); + return isBlock ? ( + 
    +            {children}
+
    + ) : ( + {children} + ); + }, + table: ({ children }) => ( +
    + {children}
    +
    + ), + th: ({ children }) => {children}, + td: ({ children }) => {children}, + }} + > + {content} +     +); + +const UsageAIChatPanel: React.FC = ({ + open, + onClose, + accessToken, +}) => { + const [messages, setMessages] = useState([]); + const [inputText, setInputText] = useState(""); + const [isLoading, setIsLoading] = useState(false); + const [selectedModel, setSelectedModel] = useState(undefined); + const [availableModels, setAvailableModels] = useState([]); + const [isLoadingModels, setIsLoadingModels] = useState(false); + const [streamingContent, setStreamingContent] = useState(""); + const [statusMessage, setStatusMessage] = useState(null); + const [activeToolCalls, setActiveToolCalls] = useState([]); + const messagesEndRef = useRef(null); + const abortControllerRef = useRef(null); + + useEffect(() => { + if (open && availableModels.length === 0) { + loadModels(); + } + }, [open]); + + useEffect(() => { + if (typeof messagesEndRef.current?.scrollIntoView === "function") { + messagesEndRef.current.scrollIntoView({ behavior: "smooth" }); + } + }, [messages, streamingContent, activeToolCalls, statusMessage]); + + const loadModels = async () => { + if (!accessToken) return; + setIsLoadingModels(true); + try { + const fetchedModels = await modelHubCall(accessToken); + if (fetchedModels?.data?.length > 0) { + const models = fetchedModels.data + .map((item: any) => item.model_group as string) + .sort(); + setAvailableModels(models); + } + } catch (error) { + console.error("Failed to load models:", error); + } finally { + setIsLoadingModels(false); + } + }; + + const handleSend = async () => { + if (!accessToken || !inputText.trim() || isLoading) return; + + const userMessage: ChatMessage = { role: "user", content: inputText.trim() }; + const updatedMessages = [...messages, userMessage]; + setMessages(updatedMessages); + setInputText(""); + setIsLoading(true); + setStreamingContent(""); + setStatusMessage(null); + setActiveToolCalls([]); + + const abortController = new AbortController(); + abortControllerRef.current = abortController; + + let accumulated = ""; + const toolCalls: ToolCallStep[] = []; + + try { + await usageAiChatStream( + accessToken, + updatedMessages.slice(-20).map((m) => ({ role: m.role, content: m.content })), + selectedModel || "", + (content: string) => { + setStatusMessage(null); + accumulated += content; + setStreamingContent(accumulated); + }, + () => { + setStatusMessage(null); + setActiveToolCalls([]); + setMessages((prev) => [ + ...prev, + { role: "assistant", content: accumulated, toolCalls: toolCalls.length > 0 ? [...toolCalls] : undefined }, + ]); + setStreamingContent(""); + }, + (errorMsg: string) => { + setStatusMessage(null); + setActiveToolCalls([]); + setMessages((prev) => [ + ...prev, + { role: "assistant", content: `Error: ${errorMsg}` }, + ]); + setStreamingContent(""); + }, + (status: string) => { + setStatusMessage(status); + }, + (event: UsageAiToolCallEvent) => { + const idx = toolCalls.findIndex((tc) => tc.tool_name === event.tool_name); + if (idx >= 0) { + toolCalls[idx] = { ...event }; + } else { + toolCalls.push({ ...event }); + } + setActiveToolCalls([...toolCalls]); + }, + abortController.signal, + ); + } catch (error: any) { + if (error?.name === "AbortError" || abortController.signal.aborted) { + return; + } + const errorMsg = error?.message || "Failed to get response. Please try again."; + setMessages((prev) => [ + ...prev, + { role: "assistant", content: `Error: ${errorMsg}` }, + ]); + setStreamingContent(""); + } finally { + setIsLoading(false); + abortControllerRef.current = null; + } + }; + + const handleKeyDown = (e: React.KeyboardEvent) => { + if (e.key === "Enter" && !e.shiftKey) { + e.preventDefault(); + handleSend(); + } + }; + + const handleClose = () => { + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + } + onClose(); + }; + + const handleClear = () => { + setMessages([]); + setStreamingContent(""); + setActiveToolCalls([]); + setStatusMessage(null); + }; + + return ( +
    + {/* Header */} +
    +
    +
    + + + +

    Ask AI

    +
    + +
    +

    + Ask about your spend, models, keys, and trends +

    +
    + + {/* Model selector */} +
    +