diff --git a/litellm/proxy/_experimental/out/assets/aria-agent.png b/litellm/proxy/_experimental/out/assets/aria-agent.png new file mode 100644 index 00000000000..870773933e3 Binary files /dev/null and b/litellm/proxy/_experimental/out/assets/aria-agent.png differ diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 0e2eed73364..28c701e0fc2 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -299,11 +299,6 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( module_path="litellm.proxy.spend_tracking.vantage_endpoints", path_prefixes=("/vantage",), ), - LazyFeature( - name="usage_ai", - module_path="litellm.proxy.management_endpoints.usage_endpoints", - path_prefixes=("/usage/ai",), - ), LazyFeature( name="prompts", module_path="litellm.proxy.prompts.prompt_endpoints", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index e5e0027a808..0dd978ee855 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -57898,163 +57898,6 @@ } } }, - "usage_ai": { - "components": { - "schemas": { - "ChatMessage": { - "properties": { - "content": { - "title": "Content", - "type": "string" - }, - "role": { - "enum": [ - "user", - "assistant" - ], - "title": "Role", - "type": "string" - } - }, - "required": [ - "role", - "content" - ], - "title": "ChatMessage", - "type": "object" - }, - "HTTPValidationError": { - "properties": { - "detail": { - "items": { - "$ref": "#/components/schemas/ValidationError" - }, - "title": "Detail", - "type": "array" - } - }, - "title": "HTTPValidationError", - "type": "object" - }, - "UsageAIChatRequest": { - "properties": { - "messages": { - "description": "Chat messages (user/assistant history)", - "items": { - "$ref": "#/components/schemas/ChatMessage" - }, - "title": "Messages", - "type": "array" - }, - "model": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "description": "Model to use for AI chat", - "title": "Model" - } - }, - "required": [ - "messages" - ], - "title": "UsageAIChatRequest", - "type": "object" - }, - "ValidationError": { - "properties": { - "ctx": { - "title": "Context", - "type": "object" - }, - "input": { - "title": "Input" - }, - "loc": { - "items": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "integer" - } - ] - }, - "title": "Location", - "type": "array" - }, - "msg": { - "title": "Message", - "type": "string" - }, - "type": { - "title": "Error Type", - "type": "string" - } - }, - "required": [ - "loc", - "msg", - "type" - ], - "title": "ValidationError", - "type": "object" - } - } - }, - "paths": { - "/usage/ai/chat": { - "post": { - "description": "AI chat about usage data. Streams SSE events with the AI response.\nThe AI agent has access to tools that query aggregated daily activity data.", - "operationId": "usage_ai_chat_usage_ai_chat_post", - "requestBody": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/UsageAIChatRequest" - } - } - }, - "required": true - }, - "responses": { - "200": { - "content": { - "application/json": { - "schema": {} - } - }, - "description": "Successful Response" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - } - }, - "security": [ - { - "APIKeyHeader": [] - } - ], - "summary": "Usage Ai Chat", - "tags": [ - "usage_ai" - ] - } - } - } - }, "vantage": { "components": { "schemas": { diff --git a/litellm/proxy/management_endpoints/usage_endpoints/__init__.py b/litellm/proxy/management_endpoints/usage_endpoints/__init__.py deleted file mode 100644 index 6e68dcd2a2e..00000000000 --- a/litellm/proxy/management_endpoints/usage_endpoints/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -""" -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 deleted file mode 100644 index 7a45d0438de..00000000000 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ /dev/null @@ -1,627 +0,0 @@ -""" -AI Usage Chat - uses LLM tool calling to answer questions about -usage/spend data by querying the aggregated daily activity endpoints. -""" - -import json -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence -from datetime import date -from typing import Final, Literal, NamedTuple, Protocol, cast, overload - -from fastapi import HTTPException -from typing_extensions import ReadOnly, TypedDict - -import litellm -from litellm._logging import verbose_proxy_logger -from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL -from litellm.proxy._types import CommonProxyErrors -from litellm.types.proxy.management_endpoints.common_daily_activity import ( - SpendAnalyticsPaginatedResponse, -) -from litellm.types.utils import ChatCompletionMessageToolCall - -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- - -USAGE_AI_TEMPERATURE: Final = 0.2 - -TABLE_DAILY_USER_SPEND: Final = "litellm_dailyuserspend" -TABLE_DAILY_TEAM_SPEND: Final = "litellm_dailyteamspend" -TABLE_DAILY_TAG_SPEND: Final = "litellm_dailytagspend" - -ENTITY_FIELD_USER: Final = "user_id" -ENTITY_FIELD_TEAM: Final = "team_id" -ENTITY_FIELD_TAG: Final = "tag" - -PAGINATED_PAGE_SIZE: Final = 200 -MAX_CHAT_MESSAGES: Final = 20 -TOP_N_MODELS: Final = 15 -TOP_N_PROVIDERS: Final = 10 -TOP_N_KEYS: Final = 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 _EntityEntry(TypedDict, total=False): - metrics: ReadOnly[Mapping[str, float]] - metadata: ReadOnly[Mapping[str, str]] - - -class _DayDump(TypedDict, total=False): - breakdown: ReadOnly[Mapping[str, Mapping[str, _EntityEntry]]] - - -class _EntityTotal(NamedTuple): - """Running per-entity totals accumulated while summarising a usage dump.""" - - alias: str - spend: float - requests: float - tokens: float - - -class _UsageDump(Protocol): - @overload - def get(self, key: Literal["metadata"], default: Mapping[str, float], /) -> Mapping[str, float]: ... - @overload - def get(self, key: Literal["results"], default: Sequence[_DayDump], /) -> Sequence[_DayDump]: ... - - -class _ToolFunctionDef(TypedDict): - name: ReadOnly[str] - description: ReadOnly[str] - parameters: ReadOnly[Mapping[str, object]] - - -class _ToolDef(TypedDict): - type: ReadOnly[str] - function: ReadOnly[_ToolFunctionDef] - - -class ToolHandler(TypedDict): - fetch: Callable[..., Awaitable[_UsageDump]] - summarise: Callable[[_UsageDump], str] - label: str - - -# --------------------------------------------------------------------------- -# Tool definitions (OpenAI function-calling schema) -# --------------------------------------------------------------------------- - -_DATE_PARAMS: Final = { - "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: Final[_ToolDef] = { - "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: Final[_ToolDef] = { - "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: Final[_ToolDef] = { - "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: Final = [_TOOL_USAGE] -TOOLS_ADMIN: Final = [_TOOL_USAGE, _TOOL_TEAM, _TOOL_TAG] - - -def get_tools_for_role(is_admin: bool) -> list[_ToolDef]: - """Return the tool list appropriate for the user's role.""" - return TOOLS_ADMIN if is_admin else TOOLS_BASE - - -_SYSTEM_PROMPT_BASE: Final = ( - "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: Final = ( - "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: Final = ( - "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: Final = _TOOL_DESCRIPTIONS_ADMIN if is_admin else _TOOL_DESCRIPTIONS_BASE - return f"{_SYSTEM_PROMPT_BASE}\n\n{tool_desc}Today's date: {date.today().isoformat()}" - - -# keep a public reference for test assertions -SYSTEM_PROMPT: Final = _SYSTEM_PROMPT_BASE - -# --------------------------------------------------------------------------- -# Data fetchers -# --------------------------------------------------------------------------- - - -def _parse_csv_ids(raw: str | None) -> list[str] | None: - 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: str | list[str] | None, - 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 ( - daily_activity_repository, - daily_activity_scope, - get_daily_activity, - get_daily_activity_aggregated, - ) - from litellm.proxy.proxy_server import prisma_client - - if use_aggregated: - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - repository: Final = daily_activity_repository(prisma_client) - scope: Final = daily_activity_scope( - table_name, - entity_id_field, - entity_id, - None, - None, - start_date, - end_date, - None, - None, - ) - return await get_daily_activity_aggregated( - repository, - scope, - ) - 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: str | None = None) -> _UsageDump: - resp: Final = 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: str | None = None) -> _UsageDump: - resp: Final = 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: str | None = None) -> _UsageDump: - resp: Final = 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: Sequence[_DayDump], dimension: str, fields: Sequence[str] -) -> dict[str, dict[str, float]]: - """Aggregate a single breakdown dimension across days.""" - totals: Final[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: _UsageDump) -> str: - meta: Final = data.get("metadata", {}) - results: Final = data.get("results", []) - - header: Final = ( - 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: Final = _accumulate_breakdown(results, "models", ["spend", "api_requests", "total_tokens"]) - providers: Final = _accumulate_breakdown(results, "providers", ["spend", "api_requests"]) - - model_lines: Final = _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: Final = _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: _UsageDump, entity_label: str) -> str: - """Summarise team/tag entity usage data.""" - results: Final = data.get("results", []) - if not results: - return f"No {entity_label} usage data found for the given date range." - - totals: Final[dict[str, _EntityTotal]] = {} - for day in results: - for eid, entry in day.get("breakdown", {}).get("entities", {}).items(): - previous = totals.get(eid) - m = entry.get("metrics", {}) - totals[eid] = _EntityTotal( - alias=previous.alias if previous is not None else entry.get("metadata", {}).get("alias", eid), - spend=(previous.spend if previous is not None else 0.0) + m.get("spend", 0), - requests=(previous.requests if previous is not None else 0) + m.get("api_requests", 0), - tokens=(previous.tokens if previous is not None else 0) + m.get("total_tokens", 0), - ) - - lines: Final = [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} | {int(d.requests)} reqs | {int(d.tokens)} tokens") - return "\n".join(lines) - - -# --------------------------------------------------------------------------- -# Tool dispatch registry -# --------------------------------------------------------------------------- - -TOOL_HANDLERS: Final[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: Mapping[str, str], - user_id: str | None, - is_admin: bool, -) -> dict[str, str]: - """Build keyword arguments for a tool's fetch function.""" - start_date: Final = fn_args.get("start_date", "") - end_date: Final = 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: Final[dict[str, str]] = {"start_date": start_date, "end_date": end_date} - if fn_name == "get_usage_data": - if not is_admin: - if user_id is None: - # Defense-in-depth: the endpoint guard in usage_endpoints/endpoints.py - # should have already rejected this. If we ever reach here it means - # a future caller invoked the helper without scoping — fail loudly - # rather than issuing an unfiltered global query. - raise ValueError( - "Non-admin caller has user_id=None; refusing to issue an " - "unscoped query. Endpoint-level guard missing." - ) - 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: Mapping[str, str], - user_id: str | None, - is_admin: bool, -) -> str: - """Run a single tool and return the summarised result text, or the reason the query - layer rejected its dates so the model can correct them.""" - kwargs: Final = _resolve_fetch_kwargs(fn_name, fn_args, user_id, is_admin) - try: - raw_data: Final = await handler["fetch"](**kwargs) - except HTTPException as e: - reason: Final = _invalid_date_reason(e) - if reason is None: - raise - return f"Invalid date range: {reason}" - return handler["summarise"](raw_data) - - -def _invalid_date_reason(error: HTTPException) -> str | None: - """The reason in the daily activity query layer's 400 for invalid or reversed dates.""" - detail: Final = cast(object, error.detail) # cast-ok: starlette types detail as str; the query layer passes a dict - match detail: - case {"error": str() as reason} if error.status_code == 400: - return reason - case _: - return None - - -async def _process_tool_call( - tc: ChatCompletionMessageToolCall, - chat_messages: list[Mapping[str, object]], - user_id: str | None, - is_admin: bool, -) -> AsyncIterator[str]: - """Execute a single tool call, yielding SSE events for status.""" - fn_name: Final = tc.function.name - fn_args: Final[Mapping[str, str]] = json.loads(tc.function.arguments) - - allowed_names: Final = {t["function"]["name"] for t in get_tools_for_role(is_admin)} - handler: Final = TOOL_HANDLERS.get(fn_name) if fn_name is not None else None - - if fn_name is None or 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: Final = { - "type": "tool_call", - "tool_name": fn_name, - "tool_label": handler["label"], - "arguments": fn_args, - } - yield _sse(cast(SSEToolCallEvent, {**tool_event_base, "status": "running"})) - - try: - tool_result = await _execute_tool_call(handler, fn_name, fn_args, user_id, is_admin) - yield _sse(cast(SSEToolCallEvent, {**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(cast(SSEToolCallEvent, {**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[Mapping[str, object]]) -> AsyncIterator[str]: - """Stream the final LLM response after tool results are appended.""" - yield _sse({"type": "status", "message": "Analyzing results..."}) - - response: Final = 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: str | None = None, - user_id: str | None = None, - is_admin: bool = False, -) -> AsyncGenerator[str, None]: - """Stream SSE events: status → tool_call → chunk → done.""" - resolved_model: Final = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL - truncated: Final = messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages - chat_messages: Final[list[Mapping[str, object]]] = [ - {"role": "system", "content": _build_system_prompt(is_admin)}, - *truncated, - ] - - try: - yield _sse({"type": "status", "message": "Thinking..."}) - tools: Final = get_tools_for_role(is_admin) - response: Final = await litellm.acompletion( - model=resolved_model, - messages=chat_messages, - tools=tools, - temperature=USAGE_AI_TEMPERATURE, - ) - choice: Final = response.choices[0] - - 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 deleted file mode 100644 index cff5bef05a0..00000000000 --- a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py +++ /dev/null @@ -1,77 +0,0 @@ -""" -USAGE AI CHAT ENDPOINTS - -/usage/ai/chat - Stream AI chat responses about usage data -""" - -from typing import Final, Literal - -from fastapi import APIRouter, Depends, Request -from fastapi.responses import StreamingResponse -from pydantic import Field - -import litellm -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.sse_keepalive import ( - SSE_COMMENT_PING, - wrap_sse_stream_with_keepalive_pings, -) -from litellm.types.llms.base import LiteLLMBaseModel - -router: Final = APIRouter() - - -class ChatMessage(LiteLLMBaseModel): - role: Literal["user", "assistant"] - content: str - - -class UsageAIChatRequest(LiteLLMBaseModel): - messages: list[ChatMessage] = Field(..., description="Chat messages (user/assistant history)") - model: str | None = 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 ( - require_caller_user_id_for_non_admin, - user_api_key_has_admin_view, - ) - from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( - stream_usage_ai_chat, - ) - - is_admin: Final = user_api_key_has_admin_view(user_api_key_dict) - if is_admin: - user_id = user_api_key_dict.user_id - else: - user_id = require_caller_user_id_for_non_admin(user_api_key_dict) - messages: Final = [{"role": m.role, "content": m.content} for m in data.messages] - - return StreamingResponse( - wrap_sse_stream_with_keepalive_pings( - stream_usage_ai_chat( - messages=messages, - model=data.model, - user_id=user_id, - is_admin=is_admin, - ), - ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, - ping_chunk=SSE_COMMENT_PING, - ), - media_type="text/event-stream", - headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, - ) diff --git a/tests/integration/compatibility/test_status_code_followups_usage_chat.py b/tests/integration/compatibility/test_status_code_followups_usage_chat.py deleted file mode 100644 index de9ec16cf30..00000000000 --- a/tests/integration/compatibility/test_status_code_followups_usage_chat.py +++ /dev/null @@ -1,125 +0,0 @@ -"""The /usage/ai/chat tools hand the model the reason the daily activity query rejected its dates, the same check -the daily activity routes answer 400 with, instead of a generic fetch error.""" - -from __future__ import annotations - -import json -import uuid -from collections.abc import Iterator -from dataclasses import dataclass -from typing import Final - -import pytest -from integration._support.client import JSON_OBJECT, Gateway, object_value -from integration._support.upstream import delete_scenario, register_scenario -from integration.compatibility._status_code_audit import Upstream, json_response, owned_gateway -from pydantic import JsonValue - -pytestmark: Final = pytest.mark.timeout(180) - -_TOOLS: Final = ("get_usage_data", "get_team_usage_data", "get_tag_usage_data") - - -@dataclass(frozen=True, slots=True) -class _UsageChat: - gateway: Gateway - identity: str - - -@pytest.fixture(scope="module") -def usage_chat(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_UsageChat]: - """A proxy whose usage chat model (`openai/...`, called straight through litellm) is the scripted upstream.""" - identity: Final = f"audit-usage-chat-{uuid.uuid4().hex}" - handle: Final = register_scenario(identity, json_response({})) - try: - with owned_gateway( - tmp_path_factory.mktemp("audit-usage-chat"), - {"model_list": [], "router_settings": {"num_retries": 0}}, - environment={"OPENAI_API_BASE": handle.api_base(), "OPENAI_API_KEY": identity}, - ) as owned: - yield _UsageChat(owned, identity) - finally: - delete_scenario(handle) - - -def _tool_calls(start_date: str, end_date: str) -> dict[str, JsonValue]: - arguments: Final = json.dumps({"start_date": start_date, "end_date": end_date}) - return { - "id": "chatcmpl-usage", - "object": "chat.completion", - "created": 1, - "model": "gpt-4o-mini", - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": None, - "tool_calls": [ - {"id": f"call-{tool}", "type": "function", "function": {"name": tool, "arguments": arguments}} - for tool in _TOOLS - ], - }, - "finish_reason": "tool_calls", - } - ], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, - } - - -def _tool_results(usage_chat: _UsageChat, start_date: str, end_date: str) -> tuple[dict[str, str], tuple[str, ...]]: - """Runs one chat turn whose model calls every tool with these dates; returns the tool messages the follow-up - model call carried, keyed by tool, and the tool statuses the SSE stream reported.""" - register_scenario(usage_chat.identity, json_response(_tool_calls(start_date, end_date))) - Upstream(usage_chat.gateway.upstream_url).reset() - with usage_chat.gateway.client.stream( - "POST", - "/usage/ai/chat", - json={"model": "openai/gpt-4o-mini", "messages": [{"role": "user", "content": "How much did we spend?"}]}, - headers={"Authorization": f"Bearer {usage_chat.gateway.key}"}, - timeout=60, - ) as response: - assert response.status_code == 200 - events: Final = tuple( - object_value(JSON_OBJECT.validate_json(line.removeprefix("data: "))) - for line in response.iter_lines() - if line.startswith("data: ") - ) - calls: Final = Upstream(usage_chat.gateway.upstream_url).calls(usage_chat.identity) - assert len(calls) == 2, calls - messages: Final = object_value(calls[1]["body"]).get("messages") - assert isinstance(messages, list), calls[1] - tool_messages: Final = { - str(message["tool_call_id"]).removeprefix("call-"): str(message["content"]) - for message in messages - if isinstance(message, dict) and message.get("role") == "tool" - } - statuses: Final = tuple(str(event.get("status")) for event in events if event.get("type") == "tool_call") - return tool_messages, statuses - - -_REVERSED: Final = "Invalid date range: end_date must be on or after start_date" -_NOT_CANONICAL: Final = "Invalid date range: start_date and end_date must be valid YYYY-MM-DD dates" -_NO_TEAM_DATA: Final = "No Team usage data found for the given date range." -_NO_TAG_DATA: Final = "No Tag usage data found for the given date range." - - -@pytest.mark.parametrize( - ("start_date", "end_date", "team_message", "tag_message"), - ( - pytest.param("2026-10-09", "2026-10-01", _REVERSED, _REVERSED, id="reversed"), - pytest.param("2026-13-01", "2026-10-09", _NOT_CANONICAL, _NOT_CANONICAL, id="invalid-month"), - pytest.param("2001/01/01", "2001-01-09", _NOT_CANONICAL, _NOT_CANONICAL, id="slashes"), - pytest.param("2001-01-01", "2001-01-09", _NO_TEAM_DATA, _NO_TAG_DATA, id="valid"), - ), -) -def test_usage_chat_tools_answer_with_the_query_layers_date_check( - usage_chat: _UsageChat, start_date: str, end_date: str, team_message: str, tag_message: str -) -> None: - """The team and tag tools give the model the daily activity routes' 400 reason; the global tool, whose query - takes these dates without that check, still answers with its usage summary.""" - tool_messages, statuses = _tool_results(usage_chat, start_date, end_date) - assert tool_messages["get_usage_data"].startswith("Total Spend: $"), tool_messages - assert tool_messages["get_team_usage_data"] == team_message, tool_messages - assert tool_messages["get_tag_usage_data"] == tag_message, tool_messages - assert statuses == ("running", "complete") * len(_TOOLS), statuses diff --git a/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py b/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py index 780ee43fb3f..0c18f2a4b6f 100644 --- a/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py +++ b/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py @@ -99,16 +99,9 @@ class _Activity(BaseModel): results: tuple[_Day, ...] -class _UpstreamMessage(BaseModel): - model_config = ConfigDict(frozen=True) - role: str | None = None - - class _UpstreamRequest(BaseModel): model_config = ConfigDict(frozen=True) stream: bool | None = None - tools: tuple[JsonValue, ...] | None = None - messages: tuple[_UpstreamMessage, ...] = () @dataclass(frozen=True, slots=True) @@ -171,17 +164,6 @@ def _stream_frames() -> tuple[bytes, ...]: return (*(f"data: {json.dumps({**envelope, **frame})}\n\n".encode() for frame in frames), b"data: [DONE]\n\n") -def _usage_tool_call() -> Mapping[str, JsonValue]: - arguments: Final = json.dumps({"start_date": _day(-1), "end_date": _day(1)}) - return { - "role": "assistant", - "content": None, - "tool_calls": [ - {"id": "call_usage", "type": "function", "function": {"name": "get_usage_data", "arguments": arguments}} - ], - } - - def _response_object() -> bytes: return json.dumps( { @@ -216,9 +198,6 @@ def _respond(request: Request) -> Reply: body: Final = _UpstreamRequest.model_validate_json(request.body or b"{}") if body.stream: return Reply(content_type="text/event-stream", chunks=_stream_frames()) - roles: Final = frozenset(message.role for message in body.messages) - if body.tools and "tool" not in roles: - return Reply(body=_completion(_usage_tool_call(), "tool_calls")) return Reply(body=_completion({"role": "assistant", "content": REPLY_TEXT}, "stop")) @@ -459,12 +438,6 @@ def _restore(database_url: str, digest: str) -> None: ) -def _events(response: httpx.Response) -> tuple[Mapping[str, JsonValue], ...]: - return tuple( - json.loads(line.removeprefix("data: ")) for line in response.text.splitlines() if line.startswith("data: ") - ) - - @dataclass(frozen=True, slots=True) class Tenant: label: str @@ -714,39 +687,19 @@ def _rig(gateway: Gateway, directory: Path) -> Generator[tuple[Gateway, str]]: @pytest.mark.timeout(360) -def test_usage_ai_chat_timeout_then_second_timeout_still_recovers_key_alias_after_database_frees( +def test_usage_page_timeout_then_second_timeout_still_recovers_key_alias_after_database_frees( gateway: Gateway, tmp_path: Path ) -> None: with _rig(gateway, tmp_path) as (proxy, database_url): - spender: Final = _spender(proxy, database_url, "ai-chat") + spender: Final = _spender(proxy, database_url, "repeated-timeout") with _pinned(proxy) as pinned: with _locked_spend_logs(database_url): - with _recording(database_url) as during_chat: - chat: Final = pinned.request( - "POST", - "/usage/ai/chat", - body={ - "messages": [{"role": "user", "content": "What did we spend?"}], - "model": "openai/gpt-4o-mini", - }, - ) - assert chat.status_code == 200, chat.text - assert chat.elapsed >= FAILED_LOOKUP_FLOOR, chat.elapsed - tool_call: Final = { - "type": "tool_call", - "tool_name": "get_usage_data", - "tool_label": "global usage data", - "arguments": {"start_date": _day(-1), "end_date": _day(1)}, - } - assert _events(chat) == ( - {"type": "status", "message": "Thinking..."}, - {**tool_call, "status": "running"}, - {**tool_call, "status": "complete"}, - {"type": "status", "message": "Analyzing results..."}, - {"type": "chunk", "content": REPLY_TEXT}, - {"type": "done"}, - ), chat.text - assert len(_lookups(during_chat)) == 1 + with _recording(database_url) as during_timeout: + timed_out: Final = _read(pinned, spender.digest) + assert timed_out.status_code == 200, timed_out.text + assert timed_out.elapsed >= FAILED_LOOKUP_FLOOR, timed_out.elapsed + assert _aliases(timed_out, spender.digest) == (None,), timed_out.text + assert len(_lookups(during_timeout)) == 1 cached_miss: Final = _read(pinned, spender.digest) assert cached_miss.elapsed < FAILED_LOOKUP_FLOOR, cached_miss.elapsed assert _aliases(cached_miss, spender.digest) == (None,), cached_miss.text @@ -1087,44 +1040,3 @@ def test_usage_page_rejects_bad_key_filters_and_unrelated_routes_ignore_a_locked assert gateway_activity == 200, gateway_activity chat: Final = proxy.chat(MODEL, text=f"locked {uuid.uuid4().hex}") assert object_value(chat["usage"]) == dict(USAGE), chat - - -@pytest.mark.timeout(300) -def test_usage_ai_chat_survives_a_dropped_database_connection_during_alias_recovery( - gateway: Gateway, tmp_path: Path -) -> None: - with ( - scratch_database() as database_url, - database_relay(database_url, b"AS " + LOOKUP_MARKER.encode()) as (relay, relayed_url), - wire_server(_respond) as wire, - _proxy(gateway, tmp_path, relayed_url, wire.url) as proxy, - ): - spender: Final = _spender(proxy, database_url, "dropped-ai-chat") - with _pinned(proxy) as pinned: - relay.arm() - chat: Final = pinned.request( - "POST", - "/usage/ai/chat", - body={ - "messages": [{"role": "user", "content": "What did we spend?"}], - "model": "openai/gpt-4o-mini", - }, - ) - assert chat.status_code == 200, chat.text - assert relay.tripped.is_set(), chat.text - tool_call: Final = { - "type": "tool_call", - "tool_name": "get_usage_data", - "tool_label": "global usage data", - "arguments": {"start_date": _day(-1), "end_date": _day(1)}, - } - assert _events(chat) == ( - {"type": "status", "message": "Thinking..."}, - {**tool_call, "status": "running"}, - {**tool_call, "status": "complete"}, - {"type": "status", "message": "Analyzing results..."}, - {"type": "chunk", "content": REPLY_TEXT}, - {"type": "done"}, - ), chat.text - recovered: Final = _named_once_reconnected(pinned, spender) - assert _metadata(recovered, spender.digest) == _full_metadata(spender), recovered.text diff --git a/tests/integration/spend/test_daily_activity_key_owner_traffic.py b/tests/integration/spend/test_daily_activity_key_owner_traffic.py index b8f113aca49..5284ad631de 100644 --- a/tests/integration/spend/test_daily_activity_key_owner_traffic.py +++ b/tests/integration/spend/test_daily_activity_key_owner_traffic.py @@ -6,7 +6,6 @@ from collections.abc import Iterable from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime, timedelta from hashlib import sha256 -from pathlib import Path from queue import SimpleQueue from typing import Final @@ -31,8 +30,7 @@ from integration._support.daily_activity import ( user_row, user_with_an_email, ) -from integration._support.database import read_rows, scratch_database -from integration._support.process import owned_proxy +from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue @@ -51,22 +49,6 @@ UNIFIED_ENDPOINTS: Final = ("/v1/chat/completions", "/v1/messages", "/v1/respons REQUESTS_OF_A_BURST: Final = 21 READS_DURING_A_BURST: Final = 30 TOKEN_LIMIT_DISCOVERY: Final = ("GET", "/v1/models") -TOOL_CALL: Final = "call_integration_usage" -ANSWER: Final = "One request cost $0.25" -SUMMARY_OF_ONE_SEEDED_ROW: Final = "\n".join( - ( - "Total Spend: $0.2500", - "Total Requests: 1", - "Successful: 1 | Failed: 0", - "Total Tokens: 15", - "", - "Top Models by Spend:", - " - gpt-4o-mini: $0.2500 (1 reqs, 15 tokens)", - "", - "Top Providers by Spend:", - " - openai: $0.2500 (1 reqs)", - ) -) def _chat_completion() -> dict[str, JsonValue]: @@ -105,55 +87,6 @@ def _provider(request: Request) -> Reply: return Reply(body=json.dumps(body).encode()) -def _usage_tool_call() -> dict[str, JsonValue]: - call: Final[dict[str, JsonValue]] = { - "id": TOOL_CALL, - "type": "function", - "function": { - "name": "get_usage_data", - "arguments": json.dumps({"start_date": DAY, "end_date": DAY}), - }, - } - return { - "id": f"chatcmpl-{uuid.uuid4().hex}", - "object": "chat.completion", - "created": 1, - "model": "gpt-4o-mini", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": None, "tool_calls": [call]}, - "finish_reason": "tool_calls", - } - ], - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, - } - - -def _streamed_chunk(delta: dict[str, JsonValue], finish_reason: str | None) -> bytes: - chunk: Final = { - "id": "chatcmpl-integration-usage", - "object": "chat.completion.chunk", - "created": 1, - "model": "gpt-4o-mini", - "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], - } - return f"data: {json.dumps(chunk)}\n\n".encode() - - -def _usage_analyst(request: Request) -> Reply: - if json.loads(request.body).get("stream"): - return Reply( - chunks=( - _streamed_chunk({"role": "assistant", "content": ANSWER}, None), - _streamed_chunk({}, "stop"), - b"data: [DONE]\n\n", - ), - content_type="text/event-stream", - ) - return Reply(body=json.dumps(_usage_tool_call()).encode()) - - def _sent_for_callers(requests: Iterable[Request]) -> tuple[Request, ...]: return tuple(request for request in requests if (request.method, request.target) != TOKEN_LIMIT_DISCOVERY) @@ -322,65 +255,6 @@ def test_cli_session_spend_is_reported_with_the_user_and_team_of_the_session( ) -@pytest.mark.timeout(300) -def test_usage_ai_chat_hands_the_model_the_usage_summary_without_any_key_owner( - gateway: Gateway, tmp_path: Path -) -> None: - question: Final = f"what did we spend {uuid.uuid4().hex}" - owner: Final = f"integration-{uuid.uuid4().hex}" - ownerless_key: Final = f"integration-ownerless-{uuid.uuid4().hex}" - with ( - scratch_database() as scratch_url, - wire_server(_usage_analyst) as wire, - owned_proxy( - gateway, - tmp_path, - { - "DATABASE_URL": scratch_url, - "OPENAI_API_BASE": f"{wire.url}/v1", - "OPENAI_BASE_URL": f"{wire.url}/v1", - "OPENAI_API_KEY": "integration-provider-key", - }, - remove_environment=("DATABASE_URL_READ_REPLICA",), - ) as candidate, - ): - candidate.post("/user/new", {"user_id": owner, "user_email": f"{owner}@example.com", "auto_create_key": False}) - with daily_rows((user_row(owner, ownerless_key, DAY),), database_url=scratch_url): - answer: Final = candidate.request( - "POST", - "/usage/ai/chat", - {"messages": [{"role": "user", "content": question}], "model": "openai/gpt-4o-mini"}, - ) - assert answer.status_code == 200, answer.text - tool_call: Final = { - "type": "tool_call", - "tool_name": "get_usage_data", - "tool_label": "global usage data", - "arguments": {"start_date": DAY, "end_date": DAY}, - } - events: Final = [ - json.loads(line.removeprefix("data: ")) for line in answer.text.splitlines() if line.startswith("data: ") - ] - assert events == [ - {"type": "status", "message": "Thinking..."}, - {**tool_call, "status": "running"}, - {**tool_call, "status": "complete"}, - {"type": "status", "message": "Analyzing results..."}, - {"type": "chunk", "content": ANSWER}, - {"type": "done"}, - ], answer.text - asked, analysed = wire.drain() - assert [asked.target, analysed.target] == ["/v1/chat/completions", "/v1/chat/completions"] - assert json.loads(asked.body)["messages"][-1] == {"role": "user", "content": question} - assert json.loads(analysed.body)["messages"][-1] == { - "role": "tool", - "tool_call_id": TOOL_CALL, - "content": SUMMARY_OF_ONE_SEEDED_ROW, - } - assert owner not in analysed.body.decode() - assert ownerless_key not in analysed.body.decode() - - @pytest.mark.timeout(300) def test_owner_is_reported_on_every_route_while_a_burst_of_requests_waits_on_the_provider(gateway: Gateway) -> None: released: Final = threading.Event() diff --git a/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py deleted file mode 100644 index 4daddb89350..00000000000 --- a/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py +++ /dev/null @@ -1,657 +0,0 @@ -""" -Tests for AI Usage Chat module. -""" - -import json -from collections.abc import AsyncIterator -from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest -from fastapi import HTTPException - -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 - @pytest.mark.parametrize( - "reason", - [ - "end_date must be on or after start_date", - "start_date and end_date must be valid YYYY-MM-DD dates", - ], - ids=["reversed", "invalid-format"], - ) - @pytest.mark.parametrize( - "tool_name", ["get_usage_data", "get_team_usage_data", "get_tag_usage_data"] - ) - async def test_invalid_tool_dates_are_reported_to_the_model( - self, tool_name: str, reason: str - ) -> None: - fetch: Final = AsyncMock( - side_effect=HTTPException(status_code=400, detail={"error": reason}) - ) - - tool_message, statuses = await _run_one_tool( - tool_name, "2025-01-31", "2025-01-01", fetch - ) - - fetch.assert_awaited_once() - assert tool_message == f"Invalid date range: {reason}" - assert statuses == ["running", "complete"] - - @pytest.mark.asyncio - async def test_tool_dates_the_query_layer_accepts_are_queried_as_sent( - self, - ) -> None: - fetch: Final = AsyncMock(return_value=SAMPLE_AGGREGATED_RESPONSE) - - tool_message, statuses = await _run_one_tool( - "get_usage_data", "2025/01/01", "2025-01-31T00:00:00", fetch - ) - - fetch.assert_awaited_once_with( - start_date="2025/01/01", end_date="2025-01-31T00:00:00" - ) - assert tool_message.startswith("Total Spend: $") - assert statuses == ["running", "complete"] - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "error", - [ - HTTPException(status_code=500, detail={"error": "db down"}), - HTTPException(status_code=400, detail="page and page_size must be >= 1"), - ], - ids=["server-error", "other-400"], - ) - async def test_other_tool_errors_stay_a_generic_fetch_error( - self, error: HTTPException - ) -> None: - tool_message, statuses = await _run_one_tool( - "get_team_usage_data", - "2025-01-01", - "2025-01-31", - AsyncMock(side_effect=error), - ) - - assert tool_message == "Error fetching team usage data. Please try again." - assert statuses == ["running", "error"] - - @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", - ) - - -class TestUsageAiChatServiceAccountGuard: - """ - Security regression: a non-admin caller with user_id=None (service-account - key) must be rejected at the endpoint boundary, before any tool dispatch. - """ - - @pytest.mark.asyncio - async def test_non_admin_with_user_id_none_is_rejected(self): - from fastapi import HTTPException - - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - from litellm.proxy.management_endpoints.usage_endpoints.endpoints import ( - ChatMessage, - UsageAIChatRequest, - usage_ai_chat, - ) - - service_account_key = UserAPIKeyAuth( - user_id=None, - user_role=LitellmUserRoles.INTERNAL_USER, - ) - request = MagicMock() - body = UsageAIChatRequest( - messages=[ChatMessage(role="user", content="hi")], - model="gpt-4o-mini", - ) - - with pytest.raises(HTTPException) as exc_info: - await usage_ai_chat( - data=body, - request=request, - user_api_key_dict=service_account_key, - ) - - assert exc_info.value.status_code == 403 - assert "Service-account keys" in str(exc_info.value.detail) - - def test_resolve_fetch_kwargs_tripwire_fires_on_none_user_id(self): - """ - Defense-in-depth: if a future endpoint forgets the entry guard and - a non-admin caller with user_id=None reaches _resolve_fetch_kwargs, - the tripwire must fire rather than issuing an unscoped query. - """ - from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( - _resolve_fetch_kwargs, - ) - - with pytest.raises(ValueError, match='Non-admin caller has user_id=None; refusing to issue an') as exc_info: - _resolve_fetch_kwargs( - fn_name="get_usage_data", - fn_args={"start_date": "2025-01-01", "end_date": "2025-01-31"}, - user_id=None, - is_admin=False, - ) - assert "Endpoint-level guard missing" in str(exc_info.value) - - -class TestUsageAiChatKeepalive: - async def _collect_endpoint_body(self, monkeypatch, interval, delay=0.3) -> tuple[list[bytes], dict]: - import asyncio - - import litellm - from fastapi.responses import StreamingResponse - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - from litellm.proxy.management_endpoints.usage_endpoints.endpoints import ( - ChatMessage, - UsageAIChatRequest, - usage_ai_chat, - ) - - monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval) - - async def slow_acompletion(**kwargs): - await asyncio.sleep(delay) - response = MagicMock() - response.choices = [MagicMock()] - response.choices[0].message.tool_calls = None - response.choices[0].message.content = "Total spend is $50.25" - return response - - with patch( # test-quality-ok: the stream calls the module-level litellm.acompletion directly; no injection seam - "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm.acompletion", - new=AsyncMock(side_effect=slow_acompletion), - ): - response = await usage_ai_chat( - data=UsageAIChatRequest(messages=[ChatMessage(role="user", content="hi")], model="gpt-4o-mini"), - request=MagicMock(), - user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), - ) - assert isinstance(response, StreamingResponse) - chunks = [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator] - return chunks, dict(response.headers) - - @pytest.mark.asyncio - async def test_endpoint_pings_while_the_planning_completion_is_still_running(self, monkeypatch): - chunks, headers = await self._collect_endpoint_body(monkeypatch, interval=0.05) - - assert headers["content-type"].startswith("text/event-stream") - assert headers["cache-control"] == "no-cache" - assert headers["x-accel-buffering"] == "no" - assert chunks[0].startswith(b'data: {"type": "status"') - assert chunks[1] == b": ping\n\n" - assert chunks.count(b": ping\n\n") >= 3 - assert b'"content": "Total spend is $50.25"' in b"".join(chunks) - assert chunks[-1] == b'data: {"type": "done"}\n\n' - - @pytest.mark.asyncio - async def test_endpoint_stream_is_untouched_while_keepalives_are_unconfigured(self, monkeypatch): - chunks, _ = await self._collect_endpoint_body(monkeypatch, interval=None, delay=0.15) - - assert b": ping\n\n" not in chunks - assert chunks[0].startswith(b'data: {"type": "status"') - assert chunks[-1] == b'data: {"type": "done"}\n\n' - - -async def _run_one_tool( - tool_name: str, start_date: str, end_date: str, fetch: AsyncMock -) -> tuple[str, list[object]]: - """Runs one admin chat turn whose model calls `tool_name` with these dates; returns the - tool message the follow-up model call carried and the tool statuses streamed.""" - arguments: Final = json.dumps({"start_date": start_date, "end_date": end_date}) - mock_tool_call: Final = MagicMock() - mock_tool_call.id = "call_dates" - mock_tool_call.function.name = tool_name - mock_tool_call.function.arguments = arguments - - mock_first_response: Final = 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_dates", - "type": "function", - "function": {"name": tool_name, "arguments": arguments}, - } - ], - } - - async def mock_stream() -> AsyncIterator[MagicMock]: - chunk: Final = MagicMock() - chunk.choices = [MagicMock()] - chunk.choices[0].delta.content = "Done." - yield chunk - - with ( - patch( - "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" - ) as mock_litellm, - patch.dict(TOOL_HANDLERS[tool_name], {"fetch": fetch}), - ): - mock_litellm.acompletion = AsyncMock( - side_effect=[mock_first_response, mock_stream()] - ) - events: Final[tuple[dict[str, object], ...]] = tuple( - [ - json.loads(event.replace("data: ", "").strip()) - async for event in stream_usage_ai_chat( - messages=[{"role": "user", "content": "How much did we spend?"}], - model="gpt-4o-mini", - is_admin=True, - ) - ] - ) - final_messages: Final = mock_litellm.acompletion.call_args_list[1].kwargs[ - "messages" - ] - - tool_messages: Final = [m for m in final_messages if m["role"] == "tool"] - assert len(tool_messages) == 1, tool_messages - assert tool_messages[0]["tool_call_id"] == "call_dates" - return str(tool_messages[0]["content"]), [ - e["status"] for e in events if e["type"] == "tool_call" - ] diff --git a/tests/unit/proxy/test__lazy_features.py b/tests/unit/proxy/test__lazy_features.py index adc4de8c8fa..0749fc12ea9 100644 --- a/tests/unit/proxy/test__lazy_features.py +++ b/tests/unit/proxy/test__lazy_features.py @@ -93,6 +93,20 @@ def test_retired_roi_calculator_cannot_be_loaded_or_called( assert response.status_code == 404, f"{method} {path}: {response.text}" +@pytest.mark.parametrize("disable_lazy_routes", ("false", "true")) +@pytest.mark.parametrize("path", ("/lazy/warm/usage_ai", "/usage/ai/chat")) +def test_retired_usage_chat_cannot_be_loaded_or_called( + monkeypatch: pytest.MonkeyPatch, disable_lazy_routes: str, path: str +) -> None: + monkeypatch.setenv(FLAG, disable_lazy_routes) + app: Final = FastAPI() + attach_lazy_features(app) + + with TestClient(app) as client: + response: Final = client.post(path) + assert response.status_code == 404, f"{path}: {response.text}" + + @pytest.mark.parametrize("value", ("1", "true", "TRUE", "yes", "on")) def test_flag_registers_every_feature_at_startup(monkeypatch: pytest.MonkeyPatch, value: str) -> None: monkeypatch.setenv(FLAG, value) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 1451256631f..38d4c435e9a 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1098,14 +1098,6 @@ "count": 1 } }, - "src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "react-hooks/immutability": { - "count": 1 - } - }, "src/app/(dashboard)/usage/_components/components/UsagePageView.tsx": { "react-hooks/purity": { "count": 1 @@ -1767,13 +1759,13 @@ "count": 1 }, "max-params": { - "count": 15 + "count": 14 }, "no-nested-ternary": { "count": 5 }, "no-restricted-syntax": { - "count": 140 + "count": 139 }, "prefer-const": { "count": 28 diff --git a/ui/litellm-dashboard/public/assets/aria-agent.png b/ui/litellm-dashboard/public/assets/aria-agent.png new file mode 100644 index 00000000000..870773933e3 Binary files /dev/null and b/ui/litellm-dashboard/public/assets/aria-agent.png differ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx index b957393be64..edde511d1cd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx @@ -141,7 +141,7 @@ describe("HomePage", () => { expect(screen.queryByRole("link", { name: /create new virtual key/i })).not.toBeInTheDocument(); }); - it("offers LiteAdmin on Home inside the dashboard frame", async () => { + it("should offer Aria Agent on Home inside the dashboard frame", async () => { render( @@ -149,7 +149,7 @@ describe("HomePage", () => { , ); - const entry = screen.getByRole("region", { name: "Ask LiteAdmin" }); + const entry = screen.getByRole("region", { name: "Ask Aria Agent" }); expect(await within(entry).findByPlaceholderText("What would you like to do?")).toBeInTheDocument(); expect(within(entry).getByRole("button", { name: "Send message" })).toBeDisabled(); expect(await screen.findByText("$501.74")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx deleted file mode 100644 index d971f5d0506..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx +++ /dev/null @@ -1,83 +0,0 @@ -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("@/components/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(); - - // One library paints the prompt as its own text node and the other leaves it on the input's - // placeholder attribute, so either one means the user is being told what to pick. - const prompt = "Select a model (optional, defaults to gpt-4o-mini)"; - expect(screen.queryAllByText(prompt).length + screen.queryAllByPlaceholderText(prompt).length).toBeGreaterThan(0); - }); - - 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/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx deleted file mode 100644 index dc4291d25e1..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx +++ /dev/null @@ -1,385 +0,0 @@ -import React, { useEffect, useRef, useState } from "react"; -import ReactMarkdown from "react-markdown"; -import { MessageSquare, Sparkles, X } from "lucide-react"; -import { Button } from "@/components/ui/button"; -import { - Combobox, - ComboboxContent, - ComboboxEmpty, - ComboboxInput, - ComboboxItem, - ComboboxList, -} from "@/components/ui/combobox"; -import { Textarea } from "@/components/ui/textarea"; -import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; -import { modelHubCall, usageAiChatStream, UsageAiToolCallEvent } from "@/components/networking"; - -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 about your spend, models, keys, and trends

    -
    - -
    - - {/* Model selector */} -
    - setSelectedModel(value ?? undefined)} - > - - - {isLoadingModels ? "Loading models…" : "No models found"} - - {(model: string) => ( - - {model} - - )} - - - -
    - - {/* Chat messages */} -
    - {messages.length === 0 && !streamingContent && !isLoading && ( -
    -
    - )} - - {messages.map((msg, idx) => ( -
    - {msg.role === "user" ? ( -
    -
    - {msg.content} -
    -
    - ) : ( -
    - {/* Tool calls for this message */} - {msg.toolCalls && msg.toolCalls.length > 0 && ( -
    - {msg.toolCalls.map((tc, tcIdx) => ( - - ))} -
    - )} - {/* Response */} -
    - -
    -
    - )} -
    - ))} - - {/* Active tool calls (in-progress) */} - {isLoading && activeToolCalls.length > 0 && ( -
    - {activeToolCalls.map((tc, idx) => ( - - ))} -
    - )} - - {/* Status / spinner */} - {isLoading && !streamingContent && ( -
    - - {statusMessage || "Thinking..."} -
    - )} - - {/* Streaming response */} - {streamingContent && ( -
    - -
    - )} - -
    -
    - - {/* Input area */} -
    -
    -