diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatModal.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatModal.test.tsx new file mode 100644 index 00000000000..fb0c0ef267d --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatModal.test.tsx @@ -0,0 +1,193 @@ +import { act, fireEvent, screen, waitFor } from "@testing-library/react"; +import { beforeAll, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../tests/test-utils"; +import UsageAIChatModal from "./UsageAIChatModal"; + +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" }, + ], + }), + getProxyBaseUrl: vi.fn().mockReturnValue("http://localhost:4000"), +})); + +vi.mock("openai", () => { + return { + default: { + OpenAI: vi.fn().mockImplementation(() => ({ + chat: { + completions: { + create: vi.fn(), + }, + }, + })), + }, + }; +}); + +const mockUserSpendData = { + results: [ + { + date: "2025-01-01", + metrics: { + spend: 100.5, + api_requests: 1000, + successful_requests: 950, + failed_requests: 50, + total_tokens: 50000, + prompt_tokens: 30000, + completion_tokens: 20000, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + breakdown: { + models: { + "gpt-4": { + metrics: { + spend: 80.0, + api_requests: 800, + successful_requests: 780, + failed_requests: 20, + total_tokens: 40000, + prompt_tokens: 24000, + completion_tokens: 16000, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + metadata: {}, + api_key_breakdown: {}, + }, + }, + model_groups: {}, + mcp_servers: {}, + providers: { + openai: { + metrics: { + spend: 100.5, + api_requests: 1000, + successful_requests: 950, + failed_requests: 50, + total_tokens: 50000, + prompt_tokens: 30000, + completion_tokens: 20000, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + metadata: {}, + api_key_breakdown: {}, + }, + }, + api_keys: { + "sk-test": { + metrics: { + spend: 100.5, + api_requests: 1000, + successful_requests: 950, + failed_requests: 50, + total_tokens: 50000, + prompt_tokens: 30000, + completion_tokens: 20000, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + metadata: { + key_alias: "Test Key", + team_id: null, + }, + }, + }, + entities: {}, + }, + }, + ], + metadata: { + total_spend: 100.5, + total_api_requests: 1000, + total_successful_requests: 950, + total_failed_requests: 50, + total_tokens: 50000, + }, +}; + +const defaultProps = { + visible: true, + onCancel: vi.fn(), + accessToken: "test-token", + userSpendData: mockUserSpendData, + dateRange: { + from: new Date("2025-01-01"), + to: new Date("2025-01-07"), + }, +}; + +describe("UsageAIChatModal", () => { + it("should render the modal when visible", () => { + renderWithProviders(); + + expect(screen.getByText("Ask AI about Usage")).toBeInTheDocument(); + expect( + screen.getByText("Ask questions about your spend, models, API keys, and usage trends") + ).toBeInTheDocument(); + }); + + it("should render model selector", () => { + renderWithProviders(); + + expect(screen.getByText("Model")).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 conversation button", () => { + renderWithProviders(); + + expect(screen.getByText("Clear conversation")).toBeInTheDocument(); + }); + + it("should not render when not visible", () => { + renderWithProviders(); + + expect(screen.queryByText("Ask AI about Usage")).not.toBeInTheDocument(); + }); + + it("should call onCancel when modal is closed", () => { + const onCancel = vi.fn(); + renderWithProviders(); + + const closeButtons = document.querySelectorAll(".ant-modal-close"); + if (closeButtons.length > 0) { + act(() => { + fireEvent.click(closeButtons[0]); + }); + expect(onCancel).toHaveBeenCalled(); + } + }); +}); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatModal.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatModal.tsx new file mode 100644 index 00000000000..45d8ecb552c --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatModal.tsx @@ -0,0 +1,381 @@ +import React, { useEffect, useMemo, useRef, useState } from "react"; +import { Modal, Select, Input, Spin } from "antd"; +import { Button } from "@tremor/react"; +import { getProxyBaseUrl, modelHubCall } from "../../networking"; +import { DailyData } from "../types"; +import openai from "openai"; + +const { TextArea } = Input; + +interface ChatMessage { + role: "user" | "assistant"; + content: string; +} + +interface UsageAIChatModalProps { + visible: boolean; + onCancel: () => void; + accessToken: string | null; + userSpendData: { + results: DailyData[]; + metadata: any; + }; + dateRange: { + from?: Date; + to?: Date; + }; +} + +function buildUsageSummary( + userSpendData: UsageAIChatModalProps["userSpendData"], + dateRange: UsageAIChatModalProps["dateRange"] +): string { + const meta = userSpendData.metadata || {}; + const results = userSpendData.results || []; + + const fromStr = dateRange.from?.toLocaleDateString() ?? "N/A"; + const toStr = dateRange.to?.toLocaleDateString() ?? "N/A"; + + const modelSpend: Record = {}; + const providerSpend: Record = {}; + const keySpend: Record = {}; + + for (const day of results) { + for (const [model, metrics] of Object.entries(day.breakdown.models || {})) { + if (!modelSpend[model]) modelSpend[model] = { spend: 0, requests: 0, tokens: 0 }; + modelSpend[model].spend += metrics.metrics.spend; + modelSpend[model].requests += metrics.metrics.api_requests; + modelSpend[model].tokens += metrics.metrics.total_tokens; + } + for (const [provider, metrics] of Object.entries(day.breakdown.providers || {})) { + if (!providerSpend[provider]) providerSpend[provider] = { spend: 0, requests: 0 }; + providerSpend[provider].spend += metrics.metrics.spend; + providerSpend[provider].requests += metrics.metrics.api_requests; + } + for (const [key, metrics] of Object.entries(day.breakdown.api_keys || {})) { + if (!keySpend[key]) keySpend[key] = { spend: 0, alias: metrics.metadata.key_alias }; + keySpend[key].spend += metrics.metrics.spend; + } + } + + const topModels = Object.entries(modelSpend) + .sort((a, b) => b[1].spend - a[1].spend) + .slice(0, 10) + .map(([name, d]) => ` - ${name}: $${d.spend.toFixed(4)} (${d.requests} requests, ${d.tokens} tokens)`) + .join("\n"); + + const topProviders = Object.entries(providerSpend) + .sort((a, b) => b[1].spend - a[1].spend) + .slice(0, 10) + .map(([name, d]) => ` - ${name}: $${d.spend.toFixed(4)} (${d.requests} requests)`) + .join("\n"); + + const topKeys = Object.entries(keySpend) + .sort((a, b) => b[1].spend - a[1].spend) + .slice(0, 10) + .map(([key, d]) => ` - ${d.alias || key}: $${d.spend.toFixed(4)}`) + .join("\n"); + + const dailySummary = results + .sort((a, b) => new Date(a.date).getTime() - new Date(b.date).getTime()) + .map((d) => ` - ${d.date}: $${d.metrics.spend.toFixed(4)} (${d.metrics.api_requests} requests)`) + .join("\n"); + + return `Date Range: ${fromStr} to ${toStr} +Total Spend: $${(meta.total_spend || 0).toFixed(4)} +Total Requests: ${meta.total_api_requests || 0} +Successful Requests: ${meta.total_successful_requests || 0} +Failed Requests: ${meta.total_failed_requests || 0} +Total Tokens: ${meta.total_tokens || 0} + +Top Models by Spend: +${topModels || " (no data)"} + +Top Providers by Spend: +${topProviders || " (no data)"} + +Top API Keys by Spend: +${topKeys || " (no data)"} + +Daily Spend: +${dailySummary || " (no data)"}`; +} + +const SYSTEM_PROMPT = `You are an AI assistant that helps users understand their LLM API usage data. You are embedded in the LiteLLM Usage dashboard. + +You have access to the user's current usage data which is provided below. Use it to answer questions about their spending, model usage, API key activity, provider costs, request volumes, and trends. + +Be concise and helpful. Use specific numbers from the data. When discussing costs, format them as dollar amounts. If the user asks about something not available in the data, let them know.`; + +const UsageAIChatModal: React.FC = ({ + visible, + onCancel, + accessToken, + userSpendData, + dateRange, +}) => { + 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 messagesEndRef = useRef(null); + const abortControllerRef = useRef(null); + + useEffect(() => { + if (visible && availableModels.length === 0) { + loadModels(); + } + }, [visible]); + + useEffect(() => { + if (typeof messagesEndRef.current?.scrollIntoView === "function") { + messagesEndRef.current.scrollIntoView({ behavior: "smooth" }); + } + }, [messages, streamingContent]); + + 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 usageSummary = useMemo( + () => buildUsageSummary(userSpendData, dateRange), + [userSpendData, dateRange] + ); + + const handleSend = async () => { + if (!accessToken || !inputText.trim() || !selectedModel || isLoading) return; + + const userMessage: ChatMessage = { role: "user", content: inputText.trim() }; + const updatedMessages = [...messages, userMessage]; + setMessages(updatedMessages); + setInputText(""); + setIsLoading(true); + setStreamingContent(""); + + const abortController = new AbortController(); + abortControllerRef.current = abortController; + + try { + const proxyBaseUrl = getProxyBaseUrl(); + const client = new openai.OpenAI({ + apiKey: accessToken, + baseURL: proxyBaseUrl, + dangerouslyAllowBrowser: true, + }); + + const chatHistory = [ + { + role: "system" as const, + content: `${SYSTEM_PROMPT}\n\nCurrent Usage Data:\n${usageSummary}`, + }, + ...updatedMessages.map((m) => ({ + role: m.role as "user" | "assistant", + content: m.content, + })), + ]; + + const response = await client.chat.completions.create( + { + model: selectedModel, + stream: true, + messages: chatHistory, + }, + { signal: abortController.signal } + ); + + let fullContent = ""; + for await (const chunk of response) { + if (chunk.choices[0]?.delta?.content) { + fullContent += chunk.choices[0].delta.content; + setStreamingContent(fullContent); + } + } + + setMessages((prev) => [...prev, { role: "assistant", content: fullContent }]); + setStreamingContent(""); + } catch (error: any) { + if (error?.name === "AbortError" || abortController.signal.aborted) { + return; + } + 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 handleCancel = () => { + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + } + onCancel(); + }; + + const handleClear = () => { + setMessages([]); + setStreamingContent(""); + }; + + return ( + + {/* Header */} +
+
+ + + +

Ask AI about Usage

+
+

+ Ask questions about your spend, models, API keys, and usage trends +

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