Add support for managed file ids

litellm/proxy/vector_store_files_endpoints/endpoints.py

@@ -12,6 +12,10 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
     get_custom_llm_provider_from_request_headers,
     get_custom_llm_provider_from_request_query,
 )
+from litellm.proxy.openai_files_endpoints.common_utils import (
+    handle_model_based_routing,
+    prepare_data_with_credentials,
+)
 from litellm.proxy.vector_store_endpoints.utils import (
     is_allowed_to_call_vector_store_files_endpoint,
 )
@@ -23,6 +27,161 @@ if TYPE_CHECKING:
 router = APIRouter()
 
 
+def _update_request_data_with_managed_file_id(
+    data: Dict,
+    file_id: str,
+    request: Request,
+    llm_router: Optional["Router"] = None,
+) -> tuple[Dict, Optional[str]]:
+    """
+    Update request data with model routing information from managed file ID.
+    
+    This function handles two types of file IDs:
+    1. Simple encoded file IDs (format: litellm:{file_id};model,{model})
+    2. Unified managed file IDs (format: litellm_proxy:{mime};unified_id,{uuid};...;llm_output_file_id,{file_id};...)
+    
+    For unified managed file IDs, it:
+    - Decodes the unified ID to extract the actual provider file ID (llm_output_file_id)
+    - Extracts the model routing information (target_model_names)
+    - Updates data with credentials for the correct deployment
+    
+    Args:
+        data: Request data to update
+        file_id: File ID (can be managed/encoded or regular)
+        request: FastAPI request object
+        llm_router: LiteLLM router for credential lookup (required for managed files)
+        
+    Returns:
+        Tuple of (updated request data, original_managed_file_id)
+        - original_managed_file_id is the original file_id if it was managed/encoded, None otherwise
+    """
+    import re
+
+    from litellm import verbose_logger
+    from litellm.llms.base_llm.managed_resources.utils import (
+        is_base64_encoded_unified_id,
+        parse_unified_id,
+    )
+
+    # First, check if this is a unified managed file ID (base64 encoded)
+    decoded_id = is_base64_encoded_unified_id(file_id)
+    
+    if decoded_id:
+        # This is a unified managed file ID
+        verbose_logger.debug(
+            f"Processing unified managed file ID: {file_id}"
+        )
+        
+        # Parse the unified ID to extract components
+        parsed_id = parse_unified_id(file_id)
+        
+        if parsed_id:
+            target_model_names = parsed_id.get("target_model_names", [])
+            
+            # Extract the actual provider file ID from llm_output_file_id field
+            # Format: litellm_proxy:...;llm_output_file_id,{actual_file_id};...
+            llm_output_file_id = None
+            try:
+                match = re.search(r"llm_output_file_id,([^;]+)", decoded_id)
+                if match:
+                    llm_output_file_id = match.group(1).strip()
+            except Exception:
+                pass
+            
+            verbose_logger.debug(
+                f"Decoded unified file ID - target_model_names: {target_model_names}, llm_output_file_id: {llm_output_file_id}"
+            )
+            
+            # Set the model for routing
+            if target_model_names and len(target_model_names) > 0:
+                routing_model = target_model_names[0]
+                data["model"] = routing_model
+                
+                # Get credentials for the model
+                if llm_router:
+                    credentials = llm_router.get_deployment_credentials_with_provider(
+                        model_id=routing_model
+                    )
+                    if credentials:
+                        prepare_data_with_credentials(
+                            data=data,
+                            credentials=credentials,
+                            file_id=llm_output_file_id,  # Use the actual provider file ID
+                        )
+                        verbose_logger.info(
+                            f"Routing vector store file operation to model: {routing_model}, file_id: {file_id} -> {llm_output_file_id}"
+                        )
+                        return data, file_id  # Return original managed file ID
+            
+            # If we extracted the provider file ID but no routing, still use it
+            if llm_output_file_id:
+                data["file_id"] = llm_output_file_id
+                verbose_logger.debug(
+                    f"Replaced unified file ID with provider file ID: {llm_output_file_id}"
+                )
+                return data, file_id  # Return original managed file ID
+        
+        return data, file_id if decoded_id else None
+    
+    # Fall back to simple encoded file ID handling (format: litellm:{file_id};model,{model})
+    should_route, model_used, original_file_id, credentials = handle_model_based_routing(
+        file_id=file_id,
+        request=request,
+        llm_router=llm_router,
+        data=data,
+        check_file_id_encoding=True,
+    )
+    
+    if should_route:
+        # Use model-based routing with credentials from config
+        prepare_data_with_credentials(
+            data=data,
+            credentials=credentials,  # type: ignore
+            file_id=original_file_id,  # Use decoded file ID if from encoded ID
+        )
+        
+        verbose_logger.debug(
+            f"Routing vector store file operation using model: {model_used}"
+            + (f", file_id: {file_id} -> {original_file_id}" if original_file_id else "")
+        )
+        return data, file_id  # Return original file ID for response replacement
+    
+    return data, None
+
+
+def _replace_file_id_in_response(response, original_file_id: str):
+    """
+    Replace the provider file ID in the response with the original managed file ID.
+    
+    This ensures that when a user sends a managed file ID, they get back the same
+    managed file ID in the response, not the decoded provider file ID.
+    
+    Args:
+        response: The response object from the provider
+        original_file_id: The original managed file ID to restore
+        
+    Returns:
+        Modified response with original file ID
+    """
+    if response is None:
+        return response
+    
+    # Handle different response types
+    if isinstance(response, dict):
+        # For dict responses (e.g., VectorStoreFileDeleteResponse)
+        if "id" in response:
+            response["id"] = original_file_id
+        if "file_id" in response:
+            response["file_id"] = original_file_id
+    elif hasattr(response, "id"):
+        # For object responses (e.g., VectorStoreFileObject)
+        response.id = original_file_id
+    elif hasattr(response, "file_id"):
+        response.file_id = original_file_id
+    
+    return response
+
+
 def _update_request_data_with_litellm_managed_vector_store_registry(
     data: Dict,
     vector_store_id: str,
@@ -204,6 +363,14 @@ async def vector_store_file_create(
     if "vector_store_id" not in data:
         data["vector_store_id"] = vector_store_id
 
+    # Handle managed file IDs if present in request body
+    original_managed_file_id = None
+    if "file_id" in data:
+        data, original_managed_file_id = _update_request_data_with_managed_file_id(
+            data=data, file_id=data["file_id"], request=request, llm_router=llm_router
+        )
+
+    # Then handle managed vector store IDs
     data = _update_request_data_with_litellm_managed_vector_store_registry(
         data=data, vector_store_id=vector_store_id, llm_router=llm_router
     )
@@ -221,7 +388,7 @@ async def vector_store_file_create(
 
     processor = ProxyBaseLLMRequestProcessing(data=data)
     try:
-        return await processor.base_process_llm_request(
+        response = await processor.base_process_llm_request(
             request=request,
             fastapi_response=fastapi_response,
             user_api_key_dict=user_api_key_dict,
@@ -239,6 +406,12 @@ async def vector_store_file_create(
             user_api_base=user_api_base,
             version=version,
         )
+        
+        # Replace provider file ID with original managed file ID in response
+        if original_managed_file_id:
+            response = _replace_file_id_in_response(response, original_managed_file_id)
+        
+        return response
     except Exception as e:  # noqa: BLE001
         raise await processor._handle_llm_api_exception(
             e=e,
@@ -366,6 +539,12 @@ async def vector_store_file_retrieve(
         "file_id": file_id,
     }
 
+    # Handle managed file IDs first
+    data, original_managed_file_id = _update_request_data_with_managed_file_id(
+        data=data, file_id=file_id, request=request, llm_router=llm_router
+    )
+
+    # Then handle managed vector store IDs
     data = _update_request_data_with_litellm_managed_vector_store_registry(
         data=data, vector_store_id=vector_store_id, llm_router=llm_router
     )
@@ -383,7 +562,7 @@ async def vector_store_file_retrieve(
 
     processor = ProxyBaseLLMRequestProcessing(data=data)
     try:
-        return await processor.base_process_llm_request(
+        response = await processor.base_process_llm_request(
             request=request,
             fastapi_response=fastapi_response,
             user_api_key_dict=user_api_key_dict,
@@ -401,6 +580,12 @@ async def vector_store_file_retrieve(
             user_api_base=user_api_base,
             version=version,
         )
+        
+        # Replace provider file ID with original managed file ID in response
+        if original_managed_file_id:
+            response = _replace_file_id_in_response(response, original_managed_file_id)
+        
+        return response
     except Exception as e:  # noqa: BLE001
         raise await processor._handle_llm_api_exception(
             e=e,
@@ -448,6 +633,12 @@ async def vector_store_file_content(
         "file_id": file_id,
     }
 
+    # Handle managed file IDs first
+    data, original_managed_file_id = _update_request_data_with_managed_file_id(
+        data=data, file_id=file_id, request=request, llm_router=llm_router
+    )
+
+    # Then handle managed vector store IDs
     data = _update_request_data_with_litellm_managed_vector_store_registry(
         data=data, vector_store_id=vector_store_id, llm_router=llm_router
     )
@@ -465,7 +656,7 @@ async def vector_store_file_content(
 
     processor = ProxyBaseLLMRequestProcessing(data=data)
     try:
-        return await processor.base_process_llm_request(
+        response = await processor.base_process_llm_request(
             request=request,
             fastapi_response=fastapi_response,
             user_api_key_dict=user_api_key_dict,
@@ -483,6 +674,12 @@ async def vector_store_file_content(
             user_api_base=user_api_base,
             version=version,
         )
+        
+        # Replace provider file ID with original managed file ID in response
+        if original_managed_file_id:
+            response = _replace_file_id_in_response(response, original_managed_file_id)
+        
+        return response
     except Exception as e:  # noqa: BLE001
         raise await processor._handle_llm_api_exception(
             e=e,
@@ -530,6 +727,12 @@ async def vector_store_file_update(
     data["vector_store_id"] = vector_store_id
     data["file_id"] = file_id
 
+    # Handle managed file IDs first
+    data, original_managed_file_id = _update_request_data_with_managed_file_id(
+        data=data, file_id=file_id, request=request, llm_router=llm_router
+    )
+
+    # Then handle managed vector store IDs
     data = _update_request_data_with_litellm_managed_vector_store_registry(
         data=data, vector_store_id=vector_store_id, llm_router=llm_router
     )
@@ -547,7 +750,7 @@ async def vector_store_file_update(
 
     processor = ProxyBaseLLMRequestProcessing(data=data)
     try:
-        return await processor.base_process_llm_request(
+        response = await processor.base_process_llm_request(
             request=request,
             fastapi_response=fastapi_response,
             user_api_key_dict=user_api_key_dict,
@@ -565,6 +768,12 @@ async def vector_store_file_update(
             user_api_base=user_api_base,
             version=version,
         )
+        
+        # Replace provider file ID with original managed file ID in response
+        if original_managed_file_id:
+            response = _replace_file_id_in_response(response, original_managed_file_id)
+        
+        return response
     except Exception as e:  # noqa: BLE001
         raise await processor._handle_llm_api_exception(
             e=e,
@@ -612,6 +821,12 @@ async def vector_store_file_delete(
         "file_id": file_id,
     }
 
+    # Handle managed file IDs first
+    data, original_managed_file_id = _update_request_data_with_managed_file_id(
+        data=data, file_id=file_id, request=request, llm_router=llm_router
+    )
+
+    # Then handle managed vector store IDs
     data = _update_request_data_with_litellm_managed_vector_store_registry(
         data=data, vector_store_id=vector_store_id, llm_router=llm_router
     )
@@ -629,7 +844,7 @@ async def vector_store_file_delete(
 
     processor = ProxyBaseLLMRequestProcessing(data=data)
     try:
-        return await processor.base_process_llm_request(
+        response = await processor.base_process_llm_request(
             request=request,
             fastapi_response=fastapi_response,
             user_api_key_dict=user_api_key_dict,
@@ -647,6 +862,12 @@ async def vector_store_file_delete(
             user_api_base=user_api_base,
             version=version,
         )
+        
+        # Replace provider file ID with original managed file ID in response
+        if original_managed_file_id:
+            response = _replace_file_id_in_response(response, original_managed_file_id)
+        
+        return response
     except Exception as e:  # noqa: BLE001
         raise await processor._handle_llm_api_exception(
             e=e,