Add using managed vector store creds for vector store files endpoint

This commit is contained in:
Sameer Kankute 2026-02-13 10:33:49 +05:30
parent 26c0624fd7
commit be69ba270e
3 changed files with 96 additions and 13 deletions

View File

@ -1,9 +1,6 @@
from typing import Dict, Literal, Type, Union
from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
from litellm_enterprise.proxy.hooks.managed_vector_store_files import (
_PROXY_LiteLLMManagedVectorStoreFiles,
)
from litellm_enterprise.proxy.hooks.managed_vector_stores import (
_PROXY_LiteLLMManagedVectorStores,
)
@ -13,7 +10,6 @@ from litellm.integrations.custom_logger import CustomLogger
ENTERPRISE_PROXY_HOOKS: Dict[str, Type[CustomLogger]] = {
"managed_files": _PROXY_LiteLLMManagedFiles,
"managed_vector_stores": _PROXY_LiteLLMManagedVectorStores,
"managed_vector_store_files": _PROXY_LiteLLMManagedVectorStoreFiles,
}
@ -22,7 +18,6 @@ def get_enterprise_proxy_hook(
Literal[
"managed_files",
"managed_vector_stores",
"managed_vector_store_files",
"max_parallel_requests",
],
str,

View File

@ -3945,6 +3945,18 @@ class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase):
file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, ResponsesAPIResponse]
class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase):
"""Table for managing vector stores with target_model_names support."""
unified_resource_id: str
resource_object: Optional[Any] = None # VectorStoreCreateResponse
model_mappings: Dict[str, str]
flat_model_resource_ids: List[str]
created_by: Optional[str]
updated_by: Optional[str]
storage_backend: Optional[str] = None
storage_url: Optional[str] = None
class EnterpriseLicenseData(TypedDict, total=False):
expiration_date: str
user_id: str

View File

@ -1,9 +1,9 @@
from typing import Dict, Optional
from typing import TYPE_CHECKING, Dict, Optional
import litellm
from fastapi import APIRouter, Depends, Request, Response
from fastapi.responses import ORJSONResponse
import litellm
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
@ -17,13 +17,88 @@ from litellm.proxy.vector_store_endpoints.utils import (
)
from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.router import Router
router = APIRouter()
def _update_request_data_with_litellm_managed_vector_store_registry(
data: Dict,
vector_store_id: str,
llm_router: Optional["Router"] = None,
) -> Dict:
"""
Update request data with model routing information from managed vector store.
This function handles two types of vector stores:
1. Legacy vector stores from registry (non-managed)
2. Managed vector stores with unified IDs (requires decoding)
For managed vector stores, this function:
- Decodes the unified vector store ID
- Extracts the model_id and provider resource ID
- Sets data["model"] so the router can use the correct deployment credentials
- Replaces the unified ID with the provider-specific ID
Args:
data: Request data to update
vector_store_id: Vector store ID (can be unified or legacy)
llm_router: LiteLLM router for credential lookup (required for managed vector stores)
Returns:
Updated request data with model routing information
"""
from litellm import verbose_logger
from litellm.llms.base_llm.managed_resources.utils import (
is_base64_encoded_unified_id,
parse_unified_id,
)
# Check if this is a managed vector store ID (base64 encoded unified ID)
decoded_id = is_base64_encoded_unified_id(vector_store_id)
if decoded_id:
# This is a managed vector store - decode and extract routing information
verbose_logger.debug(
f"Processing managed vector store ID: {vector_store_id}"
)
parsed_id = parse_unified_id(vector_store_id)
if parsed_id:
model_id = parsed_id.get("model_id")
provider_resource_id = parsed_id.get("provider_resource_id")
target_model_names = parsed_id.get("target_model_names", [])
verbose_logger.debug(
f"Decoded vector store - model_id: {model_id}, provider_resource_id: {provider_resource_id}, target_model_names: {target_model_names}"
)
# Set the model for routing - this tells the router which deployment to use
# The router will automatically get the credentials from the deployment
routing_model = None
if model_id:
routing_model = model_id
elif target_model_names and len(target_model_names) > 0:
routing_model = target_model_names[0]
if routing_model:
data["model"] = routing_model
verbose_logger.info(
f"Routing vector store files operation to model: {routing_model}"
)
# Replace unified vector store ID with provider resource ID
if provider_resource_id:
data["vector_store_id"] = provider_resource_id
verbose_logger.debug(
f"Replaced unified vector store ID with provider resource ID: {provider_resource_id}"
)
return data
# Legacy path: Check vector store registry for non-managed vector stores
if litellm.vector_store_registry is not None:
vector_store_to_run = (
litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry(
@ -42,6 +117,7 @@ def _update_request_data_with_litellm_managed_vector_store_registry(
if "litellm_params" in vector_store_to_run:
litellm_params = vector_store_to_run.get("litellm_params", {}) or {}
data.update(litellm_params)
return data
@ -129,7 +205,7 @@ async def vector_store_file_create(
data["vector_store_id"] = vector_store_id
data = _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id
data=data, vector_store_id=vector_store_id, llm_router=llm_router
)
provider_enum = await _resolve_provider(data=data, request=request)
@ -209,7 +285,7 @@ async def vector_store_file_list(
data.update(query_params)
data = _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id
data=data, vector_store_id=vector_store_id, llm_router=llm_router
)
provider_enum = await _resolve_provider(data=data, request=request)
@ -291,7 +367,7 @@ async def vector_store_file_retrieve(
}
data = _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id
data=data, vector_store_id=vector_store_id, llm_router=llm_router
)
provider_enum = await _resolve_provider(data=data, request=request)
@ -373,7 +449,7 @@ async def vector_store_file_content(
}
data = _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id
data=data, vector_store_id=vector_store_id, llm_router=llm_router
)
provider_enum = await _resolve_provider(data=data, request=request)
@ -455,7 +531,7 @@ async def vector_store_file_update(
data["file_id"] = file_id
data = _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id
data=data, vector_store_id=vector_store_id, llm_router=llm_router
)
provider_enum = await _resolve_provider(data=data, request=request)
@ -537,7 +613,7 @@ async def vector_store_file_delete(
}
data = _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id
data=data, vector_store_id=vector_store_id, llm_router=llm_router
)
provider_enum = await _resolve_provider(data=data, request=request)