[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
2026-08-02 10:21:52 +00:00 · 2025-08-08 18:28:10 -07:00
parent 1d514cc68b
commit a843e876a8
6 changed files with 256 additions and 7 deletions
@@ -11,4 +11,5 @@ litellm_settings:

 general_settings:
  cold_storage_custom_logger: s3_v2
-  store_prompts_in_cold_storage: true
+  store_prompts_in_cold_storage: true
+  store_prompts_in_spend_logs: true
@@ -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

@@ -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"
@@ -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<ChatUIProps> = ({
    }
  });
  const [messageTraceId, setMessageTraceId] = useState<string | null>(() => sessionStorage.getItem('messageTraceId') || null);
+  const [responsesSessionId, setResponsesSessionId] = useState<string | null>(() => sessionStorage.getItem('responsesSessionId') || null);
+  const [useApiSessionManagement, setUseApiSessionManagement] = useState<boolean>(() => {
+    const saved = sessionStorage.getItem('useApiSessionManagement');
+    return saved ? JSON.parse(saved) : true; // Default to API session management
+  });
  const [uploadedImage, setUploadedImage] = useState<File | null>(null);
  const [imagePreviewUrl, setImagePreviewUrl] = useState<string | null>(null);
  const [responsesUploadedImage, setResponsesUploadedImage] = useState<File | null>(null);
@@ -245,7 +251,13 @@ const ChatUI: React.FC<ChatUIProps> = ({
    } 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<ChatUIProps> = ({
    });
  };

+  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<ChatUIProps> = ({
          }
        } 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<ChatUIProps> = ({
            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<ChatUIProps> = ({
  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<ChatUIProps> = ({
                    }
                  }}
                  className="mb-4"
-                />  
+                />
+                
+                {/* Session Management Component */}
+                <SessionManagement
+                  endpointType={endpointType}
+                  responsesSessionId={responsesSessionId}
+                  useApiSessionManagement={useApiSessionManagement}
+                  onToggleSessionManagement={handleToggleSessionManagement}
+                />
              </div>

              <div>
@@ -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<SessionManagementProps> = ({
+  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 (
+    <div className="mb-4">
+      {/* Session Management Toggle */}
+      <div className="flex items-center justify-between mb-2">
+        <div className="flex items-center gap-2">
+          <span className="text-sm font-medium text-gray-700">Session Management</span>
+          <Tooltip title="Choose between LiteLLM API session management (using previous_response_id) or UI-based session management (using chat history)">
+            <InfoCircleOutlined className="text-gray-400" style={{ fontSize: '12px' }} />
+          </Tooltip>
+        </div>
+        <Switch
+          checked={useApiSessionManagement}
+          onChange={onToggleSessionManagement}
+          checkedChildren="API"
+          unCheckedChildren="UI"
+          size="small"
+        />
+      </div>
+
+      {/* Session Status Indicator */}
+      <div className={`text-xs p-2 rounded-md ${
+        responsesSessionId 
+          ? 'bg-green-50 text-green-700 border border-green-200' 
+          : 'bg-blue-50 text-blue-700 border border-blue-200'
+      }`}>
+        <div className="flex items-center justify-between">
+          <div className="flex items-center gap-1">
+            <InfoCircleOutlined style={{ fontSize: '12px' }} />
+            {getSessionDisplay()}
+          </div>
+          {responsesSessionId && (
+            <Tooltip 
+              title={
+                <div className="text-xs">
+                  <div className="mb-1">Copy response ID to continue session:</div>
+                  <div className="bg-gray-800 text-gray-100 p-2 rounded font-mono text-xs whitespace-pre-wrap">
+{`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
+  }'`}
+                  </div>
+                </div>
+              }
+              overlayStyle={{ maxWidth: '500px' }}
+            >
+              <button
+                onClick={handleCopySessionId}
+                className="ml-2 p-1 hover:bg-green-100 rounded transition-colors"
+              >
+                <CopyOutlined style={{ fontSize: '12px' }} />
+              </button>
+            </Tooltip>
+          )}
+        </div>
+        <div className="text-xs opacity-75 mt-1">
+          {getSessionDescription()}
+        </div>
+      </div>
+    </div>
+  );
+};
+
+export default SessionManagement;
@@ -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);