From a843e876a80764f2cef84f4ffac1bf71bef8b4c3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 8 Aug 2025 18:28:10 -0700 Subject: [PATCH] [Feat] Working e2e flow for Responses API session management with media (#13456) * add MultimodalContent on chat UI * add multi modal img on chat ui * utils for responses API imgs * add code snippet with imgs * chat UI add imgs * add imge upload * chat ui allow adding images * fix chat send button * fix button styles * fix clear chat * fixes session management * fixes for session management * QA fix _should_check_cold_storage_for_full_payload * test_should_check_cold_storage_for_full_payload --- litellm/proxy/proxy_config.yaml | 3 +- .../session_handler.py | 2 +- .../test_session_handler.py | 71 +++++++++++ .../src/components/chat_ui.tsx | 55 +++++++- .../components/chat_ui/SessionManagement.tsx | 119 ++++++++++++++++++ .../chat_ui/llm_calls/responses_api.tsx | 13 +- 6 files changed, 256 insertions(+), 7 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/chat_ui/SessionManagement.tsx diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index ae709c83f2..e1d784bb02 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -11,4 +11,5 @@ litellm_settings: general_settings: cold_storage_custom_logger: s3_v2 - store_prompts_in_cold_storage: true \ No newline at end of file + store_prompts_in_cold_storage: true + store_prompts_in_spend_logs: true diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 28e76ca893..886c46e737 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -245,7 +245,7 @@ class ResponsesSessionHandler: return True if len(proxy_server_request_dict) == 0: return True - if LITELLM_TRUNCATED_PAYLOAD_FIELD in proxy_server_request_dict: + if LITELLM_TRUNCATED_PAYLOAD_FIELD in str(proxy_server_request_dict): return True return False diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py index 27edfef3eb..1cc159bf96 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py @@ -293,3 +293,74 @@ async def test_e2e_cold_storage_fallback_to_truncated_payload(): # Verify result structure assert result.get("litellm_session_id") == "session-999" assert len(result.get("messages", [])) >= 1 # At least the assistant response + + +@pytest.mark.asyncio +async def test_should_check_cold_storage_for_full_payload(): + """ + Test _should_check_cold_storage_for_full_payload returns True for proxy server requests with truncated content + """ + + # Test case 1: Proxy server request with truncated PDF content (should return True) + proxy_request_with_truncated_pdf = { + "input": [ + { + "role": "user", + "type": "message", + "content": [ + { + "text": "what was datadogs largest source of operating cash ? quote the section you saw ", + "type": "input_text" + }, + { + "type": "input_image", + "image_url": "data:application/pdf;base64,JVBERi0xLjcKJYGBgYEKCjcgMCBvYmoKPDwKL0ZpbHRlciAvRmxhdGVEZWNvZGUKL0xlbmd0aCA1NjcxCj4+CnN0cmVhbQp4nO1dW4/cthV+31+h5wKVeb8AhoG9Bn0I0DYL9NlInQBFHKSpA+Tnl5qRNNRIn8ij4WpnbdqAsRaX90Oe23cOWyH94U/Dwt+/ttF/neKt59675sfPN/+9Ubp1MvwRjfAtN92fRkgn2+5jI5Xyre9++fdPN//6S/NrqCFax4XqvnVtn/631FLogjfd339+1xx/+P3nm3ffyebn/92ww2Bc46zRrGv/p5vWMOmb+N9Qb/4xtOEazn2oH3rjfV3fDTj+t6s7+zjU5XFdFxo9fPs8/CiaX26cYmc/svDjhlF+Pv7QNdT30/9wbI8dFjK0cfzhUO8wPjaOr/Eq/v/d8827vzfv37/7/v5vD6HKhw93D/c3755UI3jYuOb5p7Dsh53nYQtZqyUXugn71Dx/vnnPmHQfmuf/3HDdKhY2z8jwq8//broSjkrE/aHEtZIxZhQ/VbHHKqoVQhsv7KmKgyUWdcPkoUSHaYgwmmhk5liFt85w45WcDUC0RgntpTp1o44lMhCpl95eNOZ+AI/f3988Pp9tAV/dAu5VKz0Ls+SB0vstgNNZWTUDtw1vqC65oahKctUW5gl3etg1EnXSCWplzHewG7gDoq9jWu28DYuzPB3NucmZzm3UmvFcLI/aMpcxT0w1AvUi2aFEsJZLprxMF0xIQzmXQQArRwA2Fo/YCm7P13LxdIrodIa7QIojL5zekqblloeuwkiGI3o7jk+zMEJ9vqW+NWFDtTDnU7KtCpet1WZGHqyVxjLNxPlcdcudsU648xnNOxmIY95Wf3JduAidF3q+bgtVUC/s6VAgZ/TUE9q8oD+DCwM2qA/UFD/eWqY1RrFQ61QgUYFDBRYV3GOKkSPFaEQxQqqWWRcGzZ3u... (litellm_truncated 1197576 chars)" + } + ] + } + ], + "model": "anthropic/claude-3-7-sonnet-20250219", + "stream": True, + "litellm_trace_id": "16b86861-c120-4ecb-865b-4d2238bfd8f0" + } + + # Test case 2: Regular proxy request without truncation (should return False) + proxy_request_regular = { + "input": [ + { + "role": "user", + "type": "message", + "content": "Hello, this is a regular message" + } + ], + "model": "anthropic/claude-3-7-sonnet-20250219", + "stream": True + } + + # Test case 3: Empty request (should return True) + proxy_request_empty = {} + + # Test case 4: None request (should return True) + proxy_request_none = None + + with patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value="s3"): + # Test case 1: Should return True for truncated content + result1 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf) + assert result1 == True, "Should return True for proxy request with truncated PDF content" + + # Test case 2: Should return False for regular content + result2 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_regular) + assert result2 == False, "Should return False for regular proxy request without truncation" + + # Test case 3: Should return True for empty request + result3 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_empty) + assert result3 == True, "Should return True for empty proxy request" + + # Test case 4: Should return True for None request + result4 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_none) + assert result4 == True, "Should return True for None proxy request" + + # Test case 5: Should return False when cold storage is not configured + with patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value=None): + result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf) + assert result5 == False, "Should return False when cold storage is not configured, even with truncated content" diff --git a/ui/litellm-dashboard/src/components/chat_ui.tsx b/ui/litellm-dashboard/src/components/chat_ui.tsx index 170eb8251b..c7aae0d896 100644 --- a/ui/litellm-dashboard/src/components/chat_ui.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui.tsx @@ -51,6 +51,7 @@ import { convertImageToBase64, createMultimodalMessage, createDisplayMessage } f import ChatImageUpload from "./chat_ui/ChatImageUpload"; import ChatImageRenderer from "./chat_ui/ChatImageRenderer"; import { createChatMultimodalMessage, createChatDisplayMessage } from "./chat_ui/ChatImageUtils"; +import SessionManagement from "./chat_ui/SessionManagement"; import { SendOutlined, ApiOutlined, @@ -163,6 +164,11 @@ const ChatUI: React.FC = ({ } }); const [messageTraceId, setMessageTraceId] = useState(() => sessionStorage.getItem('messageTraceId') || null); + const [responsesSessionId, setResponsesSessionId] = useState(() => sessionStorage.getItem('responsesSessionId') || null); + const [useApiSessionManagement, setUseApiSessionManagement] = useState(() => { + const saved = sessionStorage.getItem('useApiSessionManagement'); + return saved ? JSON.parse(saved) : true; // Default to API session management + }); const [uploadedImage, setUploadedImage] = useState(null); const [imagePreviewUrl, setImagePreviewUrl] = useState(null); const [responsesUploadedImage, setResponsesUploadedImage] = useState(null); @@ -245,7 +251,13 @@ const ChatUI: React.FC = ({ } else { sessionStorage.removeItem('messageTraceId'); } - }, [apiKeySource, apiKey, selectedModel, endpointType, selectedTags, selectedVectorStores, selectedGuardrails, messageTraceId, selectedMCPTools]); + if (responsesSessionId) { + sessionStorage.setItem('responsesSessionId', responsesSessionId); + } else { + sessionStorage.removeItem('responsesSessionId'); + } + sessionStorage.setItem('useApiSessionManagement', JSON.stringify(useApiSessionManagement)); + }, [apiKeySource, apiKey, selectedModel, endpointType, selectedTags, selectedVectorStores, selectedGuardrails, messageTraceId, responsesSessionId, useApiSessionManagement, selectedMCPTools]); useEffect(() => { let userApiKey = apiKeySource === 'session' ? accessToken : apiKey; @@ -415,6 +427,21 @@ const ChatUI: React.FC = ({ }); }; + const handleResponseId = (responseId: string) => { + console.log("Received response ID for session management:", responseId); + if (useApiSessionManagement) { + setResponsesSessionId(responseId); + } + }; + + const handleToggleSessionManagement = (useApi: boolean) => { + setUseApiSessionManagement(useApi); + if (!useApi) { + // Clear API session when switching to UI mode + setResponsesSessionId(null); + } + }; + const updateImageUI = (imageUrl: string, model: string) => { setChatHistory((prevHistory) => [ ...prevHistory, @@ -602,7 +629,15 @@ const ChatUI: React.FC = ({ } } else if (endpointType === EndpointType.RESPONSES) { // Create chat history for API call - strip out model field and isImage field - const apiChatHistory = [...chatHistory.filter(msg => !msg.isImage).map(({ role, content }) => ({ role, content })), newUserMessage]; + let apiChatHistory; + + if (useApiSessionManagement && responsesSessionId) { + // When using API session management with existing session, only send the new message + apiChatHistory = [newUserMessage]; + } else { + // When using UI session management or starting new API session, send full history + apiChatHistory = [...chatHistory.filter(msg => !msg.isImage).map(({ role, content }) => ({ role, content })), newUserMessage]; + } await makeOpenAIResponsesRequest( apiChatHistory, @@ -617,7 +652,9 @@ const ChatUI: React.FC = ({ traceId, selectedVectorStores.length > 0 ? selectedVectorStores : undefined, selectedGuardrails.length > 0 ? selectedGuardrails : undefined, - selectedMCPTools // Pass the selected tool directly + selectedMCPTools, // Pass the selected tool directly + useApiSessionManagement ? responsesSessionId : null, // Only pass session ID if API mode is enabled + handleResponseId // Pass callback to capture new response ID ); } else if (endpointType === EndpointType.ANTHROPIC_MESSAGES) { const apiChatHistory = [...chatHistory.filter(msg => !msg.isImage).map(({ role, content }) => ({ role, content })), newUserMessage]; @@ -669,11 +706,13 @@ const ChatUI: React.FC = ({ const clearChatHistory = () => { setChatHistory([]); setMessageTraceId(null); + setResponsesSessionId(null); // Clear responses session ID handleRemoveImage(); // Clear any uploaded images for image edits handleRemoveResponsesImage(); // Clear any uploaded images for responses handleRemoveChatImage(); // Clear any uploaded images for chat completions sessionStorage.removeItem('chatHistory'); sessionStorage.removeItem('messageTraceId'); + sessionStorage.removeItem('responsesSessionId'); message.success("Chat history cleared."); }; @@ -790,7 +829,15 @@ const ChatUI: React.FC = ({ } }} className="mb-4" - /> + /> + + {/* Session Management Component */} +
diff --git a/ui/litellm-dashboard/src/components/chat_ui/SessionManagement.tsx b/ui/litellm-dashboard/src/components/chat_ui/SessionManagement.tsx new file mode 100644 index 0000000000..541cee3320 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat_ui/SessionManagement.tsx @@ -0,0 +1,119 @@ +import React from "react"; +import { Switch, Tooltip, message } from "antd"; +import { InfoCircleOutlined, CopyOutlined } from "@ant-design/icons"; +import { EndpointType } from "./mode_endpoint_mapping"; + +interface SessionManagementProps { + endpointType: string; + responsesSessionId: string | null; + useApiSessionManagement: boolean; + onToggleSessionManagement: (useApi: boolean) => void; +} + +const SessionManagement: React.FC = ({ + endpointType, + responsesSessionId, + useApiSessionManagement, + onToggleSessionManagement, +}) => { + if (endpointType !== EndpointType.RESPONSES) { + return null; + } + + const handleCopySessionId = () => { + if (responsesSessionId) { + navigator.clipboard.writeText(responsesSessionId); + message.success("Response ID copied to clipboard!"); + } + }; + + const getSessionDisplay = () => { + if (!responsesSessionId) { + return useApiSessionManagement ? 'API Session: Ready' : 'UI Session: Ready'; + } + + const sessionPrefix = useApiSessionManagement ? 'Response ID' : 'UI Session'; + const truncatedId = responsesSessionId.slice(0, 10); + return `${sessionPrefix}: ${truncatedId}...`; + }; + + const getSessionDescription = () => { + if (!responsesSessionId) { + return useApiSessionManagement + ? 'LiteLLM will manage session using previous_response_id' + : 'UI will manage session using chat history'; + } + + return useApiSessionManagement + ? 'LiteLLM API session active - context maintained server-side' + : 'UI session active - context maintained client-side'; + }; + + return ( +
+ {/* Session Management Toggle */} +
+
+ Session Management + + + +
+ +
+ + {/* Session Status Indicator */} +
+
+
+ + {getSessionDisplay()} +
+ {responsesSessionId && ( + +
Copy response ID to continue session:
+
+{`curl -X POST "your-proxy-url/v1/responses" \\ + -H "Authorization: Bearer your-api-key" \\ + -H "Content-Type: application/json" \\ + -d '{ + "model": "your-model", + "input": [{"role": "user", "content": "your message", "type": "message"}], + "previous_response_id": "${responsesSessionId}", + "stream": true + }'`} +
+
+ } + overlayStyle={{ maxWidth: '500px' }} + > + + + )} +
+
+ {getSessionDescription()} +
+
+
+ ); +}; + +export default SessionManagement; \ No newline at end of file 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 26e1bd6e28..71033e5d9d 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 @@ -18,7 +18,9 @@ export async function makeOpenAIResponsesRequest( traceId?: string, vector_store_ids?: string[], guardrails?: string[], - selectedMCPTool?: string + selectedMCPTool?: string, + previousResponseId?: string | null, + onResponseId?: (responseId: string) => void ) { if (!accessToken) { throw new Error("API key is required"); @@ -84,6 +86,7 @@ export async function makeOpenAIResponsesRequest( input: formattedInput, stream: true, litellm_trace_id: traceId, + ...(previousResponseId ? { previous_response_id: previousResponseId } : {}), ...(vector_store_ids ? { vector_store_ids } : {}), ...(guardrails ? { guardrails } : {}), ...(tools ? { tools, tool_choice: "required" } : {}), @@ -144,6 +147,14 @@ export async function makeOpenAIResponsesRequest( const response_obj = event.response; const usage = response_obj.usage; console.log("Usage data:", usage); + console.log("Response completed event:", response_obj); + + // Extract response_id for session management + if (response_obj.id && onResponseId) { + console.log("Response ID for session management:", response_obj.id); + onResponseId(response_obj.id); + } + if (usage && onUsageData) { console.log("Usage data:", usage);