diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index 32ad556277c..d07a31618fd 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -425,7 +425,6 @@ async def stream_usage_ai_chat( handler = TOOL_HANDLERS.get(fn_name) if not handler: - yield _sse({"type": "status", "message": f"Unknown tool: {fn_name}"}) chat_messages.append({ "role": "tool", "tool_call_id": tool_call.id, @@ -434,8 +433,11 @@ async def stream_usage_ai_chat( continue yield _sse({ - "type": "status", - "message": f"Fetching {handler['label']} ({fn_args.get('start_date', '')} to {fn_args.get('end_date', '')})..." + "type": "tool_call", + "tool_name": fn_name, + "tool_label": handler["label"], + "arguments": fn_args, + "status": "running", }) try: @@ -458,8 +460,24 @@ async def stream_usage_ai_chat( raw_data = await handler["fetch"](**fetch_kwargs) tool_result = handler["summarise"](raw_data) + + yield _sse({ + "type": "tool_call", + "tool_name": fn_name, + "tool_label": handler["label"], + "arguments": fn_args, + "status": "complete", + }) except Exception as e: tool_result = f"Error fetching {handler['label']}: {str(e)}" + yield _sse({ + "type": "tool_call", + "tool_name": fn_name, + "tool_label": handler["label"], + "arguments": fn_args, + "status": "error", + "error": str(e), + }) chat_messages.append({ "role": "tool", diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py index 967fed1b489..23ae5f0dd50 100644 --- a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py +++ b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -200,12 +200,15 @@ class TestStreamUsageAiChat: 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) >= 2 + assert len(status_events) >= 1 assert "Thinking" in status_events[0]["message"] - assert "Fetching" in status_events[1]["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 diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx index 206c14673e4..85ce11e605a 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx @@ -42,7 +42,7 @@ describe("UsageAIChatPanel", () => { it("should render model selector", () => { renderWithProviders(); - expect(screen.getByText("Select a model")).toBeInTheDocument(); + expect(screen.getByText("Select a model (optional, defaults to gpt-4o-mini)")).toBeInTheDocument(); }); it("should render empty state message when no conversation", () => { diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx index 6209313c097..7249474b8de 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx @@ -1,13 +1,23 @@ import React, { useEffect, useRef, useState } from "react"; import { Select, Input, Spin } from "antd"; import { Button } from "@tremor/react"; -import { modelHubCall, usageAiChatStream } from "../../networking"; +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 { @@ -16,6 +26,83 @@ interface UsageAIChatPanelProps { 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, @@ -29,6 +116,7 @@ const UsageAIChatPanel: React.FC = ({ 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); @@ -42,7 +130,7 @@ const UsageAIChatPanel: React.FC = ({ if (typeof messagesEndRef.current?.scrollIntoView === "function") { messagesEndRef.current.scrollIntoView({ behavior: "smooth" }); } - }, [messages, streamingContent]); + }, [messages, streamingContent, activeToolCalls, statusMessage]); const loadModels = async () => { if (!accessToken) return; @@ -72,11 +160,13 @@ const UsageAIChatPanel: React.FC = ({ setIsLoading(true); setStreamingContent(""); setStatusMessage(null); + setActiveToolCalls([]); const abortController = new AbortController(); abortControllerRef.current = abortController; let accumulated = ""; + const toolCalls: ToolCallStep[] = []; try { await usageAiChatStream( @@ -90,11 +180,16 @@ const UsageAIChatPanel: React.FC = ({ }, () => { setStatusMessage(null); - setMessages((prev) => [...prev, { role: "assistant", content: accumulated }]); + 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}` }, @@ -104,6 +199,15 @@ const UsageAIChatPanel: React.FC = ({ (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) { @@ -139,6 +243,8 @@ const UsageAIChatPanel: React.FC = ({ const handleClear = () => { setMessages([]); setStreamingContent(""); + setActiveToolCalls([]); + setStatusMessage(null); }; return ( @@ -175,11 +281,12 @@ const UsageAIChatPanel: React.FC = ({ {/* Model selector */}