From 58878aa0bb5a1afedbec6845ba4b7adf57ac0061 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 26 Apr 2025 16:02:58 -0700 Subject: [PATCH] [UI] Fixes for sessions on UI - ensure errors have a session and use 1 session for test key (#10342) * ui fix use same trace id on test key page * _get_session_id_for_spend_log * fix linting --- .../spend_tracking/spend_tracking_utils.py | 33 ++++++++++++++++--- ui/litellm-dashboard/package-lock.json | 21 +++++++++++- ui/litellm-dashboard/package.json | 4 ++- .../src/components/chat_ui.tsx | 15 +++++++-- .../chat_ui/llm_calls/chat_completion.tsx | 13 ++++++-- .../chat_ui/llm_calls/responses_api.tsx | 14 ++++++-- 6 files changed, 87 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index b6046f1fa31..cde2c43559e 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -290,10 +290,9 @@ def get_logging_payload( # noqa: PLR0915 proxy_server_request=_get_proxy_server_request_for_spend_logs_payload( metadata=metadata, litellm_params=litellm_params ), - session_id=( - standard_logging_payload.get("trace_id") - if standard_logging_payload is not None - else None + session_id=_get_session_id_for_spend_log( + kwargs=kwargs, + standard_logging_payload=standard_logging_payload, ), ) @@ -310,6 +309,32 @@ def get_logging_payload( # noqa: PLR0915 raise e +def _get_session_id_for_spend_log( + kwargs: dict, + standard_logging_payload: Optional[StandardLoggingPayload], +) -> str: + """ + Get the session id for the spend log. + + This ensures each spend log is associated with a unique session id. + + """ + import uuid + + if ( + standard_logging_payload is not None + and standard_logging_payload.get("trace_id") is not None + ): + return str(standard_logging_payload.get("trace_id")) + + # Users can dynamically set the trace_id for each request by passing `litellm_trace_id` in kwargs + if kwargs.get("litellm_trace_id") is not None: + return str(kwargs.get("litellm_trace_id")) + + # Ensure we always have a session id, if none is provided + return str(uuid.uuid4()) + + def _ensure_datetime_utc(timestamp: datetime) -> datetime: """Helper to ensure datetime is in UTC""" timestamp = timestamp.astimezone(timezone.utc) diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 921f4d2603d..92cad020414 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -17,6 +17,7 @@ "@tanstack/react-table": "^8.20.6", "@tremor/react": "^3.13.3", "@types/papaparse": "^5.3.15", + "@types/uuid": "^10.0.0", "antd": "^5.13.2", "cva": "^1.0.0-beta.3", "fs": "^0.0.1-security", @@ -31,7 +32,8 @@ "react-dom": "^18", "react-markdown": "^9.0.1", "react-syntax-highlighter": "^15.6.1", - "tailwind-merge": "^3.2.0" + "tailwind-merge": "^3.2.0", + "uuid": "^11.1.0" }, "devDependencies": { "@tailwindcss/forms": "^0.5.7", @@ -1100,6 +1102,11 @@ "resolved": "https://registry.npmjs.org/@types/unist/-/unist-3.0.2.tgz", "integrity": "sha512-dqId9J8K/vGi5Zr7oo212BGii5m3q5Hxlkwy3WpYuKPklmBEvsbMYYyLxAQpSffdLl/gdW0XUpKWFvYmyoWCoQ==" }, + "node_modules/@types/uuid": { + "version": "10.0.0", + "resolved": "https://registry.npmjs.org/@types/uuid/-/uuid-10.0.0.tgz", + "integrity": "sha512-7gqG38EyHgyP1S+7+xomFtL+ZNHcKv6DwNaCZmJmo1vgMugyF3TCnXVg4t1uk89mLNwnLtnY3TpOpCOyp1/xHQ==" + }, "node_modules/@typescript-eslint/parser": { "version": "6.19.1", "resolved": "https://registry.npmjs.org/@typescript-eslint/parser/-/parser-6.19.1.tgz", @@ -7687,6 +7694,18 @@ "resolved": "https://registry.npmjs.org/util-deprecate/-/util-deprecate-1.0.2.tgz", "integrity": "sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==" }, + "node_modules/uuid": { + "version": "11.1.0", + "resolved": "https://registry.npmjs.org/uuid/-/uuid-11.1.0.tgz", + "integrity": "sha512-0/A9rDy9P7cJ+8w1c9WD9V//9Wj15Ce2MPz8Ri6032usz+NfePxx5AcN3bN+r6ZL6jEo066/yNYB3tn4pQEx+A==", + "funding": [ + "https://github.com/sponsors/broofa", + "https://github.com/sponsors/ctavan" + ], + "bin": { + "uuid": "dist/esm/bin/uuid" + } + }, "node_modules/vfile": { "version": "6.0.1", "resolved": "https://registry.npmjs.org/vfile/-/vfile-6.0.1.tgz", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 60801371646..b242a39bccb 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -18,6 +18,7 @@ "@tanstack/react-table": "^8.20.6", "@tremor/react": "^3.13.3", "@types/papaparse": "^5.3.15", + "@types/uuid": "^10.0.0", "antd": "^5.13.2", "cva": "^1.0.0-beta.3", "fs": "^0.0.1-security", @@ -32,7 +33,8 @@ "react-dom": "^18", "react-markdown": "^9.0.1", "react-syntax-highlighter": "^15.6.1", - "tailwind-merge": "^3.2.0" + "tailwind-merge": "^3.2.0", + "uuid": "^11.1.0" }, "devDependencies": { "@tailwindcss/forms": "^0.5.7", diff --git a/ui/litellm-dashboard/src/components/chat_ui.tsx b/ui/litellm-dashboard/src/components/chat_ui.tsx index 6f9801c6326..bb4b088324e 100644 --- a/ui/litellm-dashboard/src/components/chat_ui.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui.tsx @@ -22,6 +22,7 @@ import { Button, Divider, } from "@tremor/react"; +import { v4 as uuidv4 } from 'uuid'; import { message, Select, Spin, Typography, Tooltip, Input } from "antd"; import { makeOpenAIChatCompletionRequest } from "./chat_ui/llm_calls/chat_completion"; @@ -82,6 +83,7 @@ const ChatUI: React.FC = ({ const [isLoading, setIsLoading] = useState(false); const abortControllerRef = useRef(null); const [selectedTags, setSelectedTags] = useState([]); + const [messageTraceId, setMessageTraceId] = useState(null); const chatEndRef = useRef(null); @@ -300,6 +302,12 @@ const ChatUI: React.FC = ({ // Create message object without model field for API call const newUserMessage = { role: "user", content: inputMessage }; + // Generate new trace ID for a new conversation or use existing one + const traceId = messageTraceId || uuidv4(); + if (!messageTraceId) { + setMessageTraceId(traceId); + } + // Update UI with full message object setChatHistory([...chatHistory, newUserMessage]); setIsLoading(true); @@ -319,7 +327,8 @@ const ChatUI: React.FC = ({ signal, updateReasoningContent, updateTimingData, - updateUsageData + updateUsageData, + traceId ); } else if (endpointType === EndpointType.IMAGE) { // For image generation @@ -344,7 +353,8 @@ const ChatUI: React.FC = ({ signal, updateReasoningContent, updateTimingData, - updateUsageData + updateUsageData, + traceId ); } } @@ -365,6 +375,7 @@ const ChatUI: React.FC = ({ const clearChatHistory = () => { setChatHistory([]); + setMessageTraceId(null); message.success("Chat history cleared."); }; diff --git a/ui/litellm-dashboard/src/components/chat_ui/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/chat_ui/llm_calls/chat_completion.tsx index 010025f0baa..8a157bb8f41 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/llm_calls/chat_completion.tsx @@ -12,7 +12,8 @@ export async function makeOpenAIChatCompletionRequest( signal?: AbortSignal, onReasoningContent?: (content: string) => void, onTimingData?: (timeToFirstToken: number) => void, - onUsageData?: (usage: TokenUsage) => void + onUsageData?: (usage: TokenUsage) => void, + traceId?: string ) { // base url should be the current base_url const isLocal = process.env.NODE_ENV === "development"; @@ -23,11 +24,18 @@ export async function makeOpenAIChatCompletionRequest( const proxyBaseUrl = isLocal ? "http://localhost:4000" : window.location.origin; + + // Prepare headers with tags and trace ID + const headers: Record = {}; + if (tags && tags.length > 0) { + headers['x-litellm-tags'] = tags.join(','); + } + const client = new openai.OpenAI({ apiKey: accessToken, baseURL: proxyBaseUrl, dangerouslyAllowBrowser: true, - defaultHeaders: tags && tags.length > 0 ? { 'x-litellm-tags': tags.join(',') } : undefined, + defaultHeaders: headers, }); try { @@ -45,6 +53,7 @@ export async function makeOpenAIChatCompletionRequest( stream_options: { include_usage: true, }, + litellm_trace_id: traceId, messages: chatHistory as ChatCompletionMessageParam[], }, { signal }); diff --git a/ui/litellm-dashboard/src/components/chat_ui/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/chat_ui/llm_calls/responses_api.tsx index 744935159b9..29e66cdbee4 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/llm_calls/responses_api.tsx @@ -12,7 +12,8 @@ export async function makeOpenAIResponsesRequest( signal?: AbortSignal, onReasoningContent?: (content: string) => void, onTimingData?: (timeToFirstToken: number) => void, - onUsageData?: (usage: TokenUsage) => void + onUsageData?: (usage: TokenUsage) => void, + traceId?: string ) { if (!accessToken) { throw new Error("API key is required"); @@ -28,11 +29,17 @@ export async function makeOpenAIResponsesRequest( ? "http://localhost:4000" : window.location.origin; + // Prepare headers with tags and trace ID + const headers: Record = {}; + if (tags && tags.length > 0) { + headers['x-litellm-tags'] = tags.join(','); + } + const client = new openai.OpenAI({ apiKey: accessToken, baseURL: proxyBaseUrl, dangerouslyAllowBrowser: true, - defaultHeaders: tags && tags.length > 0 ? { 'x-litellm-tags': tags.join(',') } : undefined, + defaultHeaders: headers, }); try { @@ -52,6 +59,7 @@ export async function makeOpenAIResponsesRequest( model: selectedModel, input: formattedInput, stream: true, + litellm_trace_id: traceId, }, { signal }); for await (const event of response) { @@ -60,7 +68,7 @@ export async function makeOpenAIResponsesRequest( // Use a type-safe approach to handle events if (typeof event === 'object' && event !== null) { // Handle output text delta - // 1) drop any “role” streams + // 1) drop any "role" streams if (event.type === "response.role.delta") { continue; }