mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-23 05:34:41 +00:00
fix(proxy): convert Union to modern syntax for managed vector stores lint
This commit is contained in:
Arjun Pakhan
parent
3ce2c87268
commit
fa3ba72841
1 changed files with 75 additions and 99 deletions
|
|
@ -3,7 +3,7 @@
|
|||
## It allows creating vector stores across multiple models and managing them with unified IDs
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Union, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Union, cast
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
|
@ -29,7 +29,7 @@ if TYPE_CHECKING:
|
|||
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
|
||||
from litellm.proxy.utils import PrismaClient as _PrismaClient
|
||||
|
||||
Span = Union[_Span, Any]
|
||||
Span = _Span | Any
|
||||
InternalUsageCache = _InternalUsageCache
|
||||
PrismaClient = _PrismaClient
|
||||
else:
|
||||
|
|
@ -38,12 +38,10 @@ else:
|
|||
PrismaClient = Any
|
||||
|
||||
|
||||
class _PROXY_LiteLLMManagedVectorStores(
|
||||
CustomLogger, BaseManagedResource[VectorStoreCreateResponse]
|
||||
):
|
||||
class _PROXY_LiteLLMManagedVectorStores(CustomLogger, BaseManagedResource[VectorStoreCreateResponse]):
|
||||
"""
|
||||
Managed vector stores with target_model_names support.
|
||||
|
||||
|
||||
This class provides functionality to:
|
||||
- Create vector stores across multiple models
|
||||
- Retrieve vector stores by unified ID
|
||||
|
|
@ -51,9 +49,7 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
- List vector stores created by a user
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient
|
||||
):
|
||||
def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient):
|
||||
CustomLogger.__init__(self)
|
||||
BaseManagedResource.__init__(self, internal_usage_cache, prisma_client)
|
||||
|
||||
|
|
@ -75,18 +71,18 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
def get_unified_resource_id_format(
|
||||
self,
|
||||
resource_object: VectorStoreCreateResponse,
|
||||
target_model_names_list: List[str],
|
||||
target_model_names_list: list[str],
|
||||
) -> str:
|
||||
"""
|
||||
Generate the format string for the unified vector store ID.
|
||||
|
||||
|
||||
Format:
|
||||
litellm_proxy:vector_store;unified_id,<uuid>;target_model_names,<models>;resource_id,<vs_id>;model_id,<model_id>
|
||||
"""
|
||||
# VectorStoreCreateResponse is a TypedDict, so resource_object is a dictionary
|
||||
# Extract provider resource ID from the response
|
||||
provider_resource_id = resource_object.get("id", "")
|
||||
|
||||
|
||||
# Model ID is stored in hidden params if the response object supports it
|
||||
# For TypedDict responses, we need to check if _hidden_params was added
|
||||
hidden_params: Mapping[str, object] = {}
|
||||
|
|
@ -111,13 +107,13 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
) -> VectorStoreCreateResponse:
|
||||
"""
|
||||
Create a vector store for a specific model.
|
||||
|
||||
|
||||
Args:
|
||||
llm_router: LiteLLM router instance
|
||||
model: Model name to create vector store for
|
||||
request_data: Request data for vector store creation
|
||||
litellm_parent_otel_span: OpenTelemetry span for tracing
|
||||
|
||||
|
||||
Returns:
|
||||
VectorStoreCreateResponse from the provider
|
||||
"""
|
||||
|
|
@ -133,30 +129,28 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
self,
|
||||
create_request: VectorStoreCreateOptionalRequestParams,
|
||||
llm_router: Router,
|
||||
target_model_names_list: List[str],
|
||||
target_model_names_list: list[str],
|
||||
litellm_parent_otel_span: Span,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> VectorStoreCreateResponse:
|
||||
"""
|
||||
Create a vector store across multiple models.
|
||||
|
||||
|
||||
Args:
|
||||
create_request: Vector store creation request parameters
|
||||
llm_router: LiteLLM router instance
|
||||
target_model_names_list: List of target model names
|
||||
litellm_parent_otel_span: OpenTelemetry span for tracing
|
||||
user_api_key_dict: User API key authentication details
|
||||
|
||||
|
||||
Returns:
|
||||
VectorStoreCreateResponse with unified ID
|
||||
"""
|
||||
verbose_logger.info(
|
||||
f"Creating managed vector store for models: {target_model_names_list}"
|
||||
)
|
||||
verbose_logger.info(f"Creating managed vector store for models: {target_model_names_list}")
|
||||
|
||||
# Create vector store for each model
|
||||
# Convert TypedDict to Dict[str, Any] for base class compatibility
|
||||
request_data_dict: Dict[str, object] = dict(create_request)
|
||||
request_data_dict: dict[str, object] = dict(create_request)
|
||||
responses = await self.create_resource_for_each_model(
|
||||
llm_router=llm_router,
|
||||
request_data=request_data_dict,
|
||||
|
|
@ -171,7 +165,7 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
)
|
||||
|
||||
# Extract model mappings from responses
|
||||
model_mappings: Dict[str, str] = {}
|
||||
model_mappings: dict[str, str] = {}
|
||||
for response in responses:
|
||||
hidden_params = getattr(response, "_hidden_params", {}) or {}
|
||||
model_id = hidden_params.get("model_id")
|
||||
|
|
@ -179,9 +173,7 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
# VectorStoreCreateResponse is a TypedDict, use dict access
|
||||
model_mappings[model_id] = response["id"]
|
||||
|
||||
verbose_logger.debug(
|
||||
f"Created vector stores with model mappings: {model_mappings}"
|
||||
)
|
||||
verbose_logger.debug(f"Created vector stores with model mappings: {model_mappings}")
|
||||
|
||||
# Store in database
|
||||
await self.store_unified_resource_id(
|
||||
|
|
@ -196,29 +188,27 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
# VectorStoreCreateResponse is a TypedDict, so we need to create a new dict with the unified ID
|
||||
response = responses[0].copy()
|
||||
response["id"] = unified_id
|
||||
|
||||
verbose_logger.info(
|
||||
f"Successfully created managed vector store with unified ID: {unified_id}"
|
||||
)
|
||||
|
||||
verbose_logger.info(f"Successfully created managed vector store with unified ID: {unified_id}")
|
||||
|
||||
return response
|
||||
|
||||
async def alist_vector_stores(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
limit: Optional[int] = None,
|
||||
after: Optional[str] = None,
|
||||
order: Optional[str] = None,
|
||||
) -> Dict[str, object]:
|
||||
limit: int | None = None,
|
||||
after: str | None = None,
|
||||
order: str | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
List vector stores created by a user.
|
||||
|
||||
|
||||
Args:
|
||||
user_api_key_dict: User API key authentication details
|
||||
limit: Maximum number of vector stores to return
|
||||
after: Cursor for pagination
|
||||
order: Sort order ('asc' or 'desc')
|
||||
|
||||
|
||||
Returns:
|
||||
Dictionary with list of vector stores and pagination info
|
||||
"""
|
||||
|
|
@ -233,65 +223,55 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
# ACCESS CONTROL
|
||||
# ============================================================================
|
||||
|
||||
async def check_vector_store_access(
|
||||
self, vector_store_id: str, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> bool:
|
||||
async def check_vector_store_access(self, vector_store_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
"""
|
||||
Check if user has access to a vector store.
|
||||
|
||||
|
||||
Args:
|
||||
vector_store_id: The unified vector store ID
|
||||
user_api_key_dict: User API key authentication details
|
||||
|
||||
|
||||
Returns:
|
||||
True if user has access, False otherwise
|
||||
"""
|
||||
is_unified_id = is_base64_encoded_unified_id(vector_store_id)
|
||||
|
||||
|
||||
if is_unified_id:
|
||||
# Check access for managed vector store
|
||||
return await self.can_user_access_unified_resource_id(
|
||||
vector_store_id,
|
||||
user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
# Not a managed vector store, allow access
|
||||
return True
|
||||
|
||||
async def check_managed_vector_store_access(
|
||||
self, data: Dict, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> bool:
|
||||
async def check_managed_vector_store_access(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
"""
|
||||
Check if user has access to a managed vector store in request data.
|
||||
|
||||
|
||||
Args:
|
||||
data: Request data containing vector_store_id
|
||||
user_api_key_dict: User API key authentication details
|
||||
|
||||
|
||||
Returns:
|
||||
True if this is a managed vector store and user has access
|
||||
|
||||
|
||||
Raises:
|
||||
HTTPException: If user doesn't have access
|
||||
"""
|
||||
vector_store_id = cast(Optional[str], data.get("vector_store_id"))
|
||||
is_unified_id = (
|
||||
is_base64_encoded_unified_id(vector_store_id)
|
||||
if vector_store_id
|
||||
else False
|
||||
)
|
||||
|
||||
vector_store_id = cast(str | None, data.get("vector_store_id"))
|
||||
is_unified_id = is_base64_encoded_unified_id(vector_store_id) if vector_store_id else False
|
||||
|
||||
if is_unified_id and vector_store_id:
|
||||
if await self.can_user_access_unified_resource_id(
|
||||
vector_store_id, user_api_key_dict
|
||||
):
|
||||
if await self.can_user_access_unified_resource_id(vector_store_id, user_api_key_dict):
|
||||
return True
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"User {user_api_key_dict.user_id} does not have access to vector store {vector_store_id}",
|
||||
)
|
||||
|
||||
|
||||
return False
|
||||
|
||||
# ============================================================================
|
||||
|
|
@ -302,23 +282,23 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: "DualCache",
|
||||
data: Dict,
|
||||
data: dict,
|
||||
call_type: str,
|
||||
) -> Union[Exception, str, Dict, None]:
|
||||
) -> Exception | str | dict | None:
|
||||
"""
|
||||
Pre-call hook to handle vector store operations.
|
||||
|
||||
|
||||
This hook intercepts vector store requests and:
|
||||
- Validates access for managed vector stores
|
||||
- Transforms unified IDs to provider-specific IDs
|
||||
- Adds model routing information
|
||||
|
||||
|
||||
Args:
|
||||
user_api_key_dict: User API key authentication details
|
||||
cache: Cache instance
|
||||
data: Request data
|
||||
call_type: Type of call being made
|
||||
|
||||
|
||||
Returns:
|
||||
Modified request data or None
|
||||
"""
|
||||
|
|
@ -330,40 +310,36 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
# Handle vector store search operations
|
||||
if call_type == "avector_store_search":
|
||||
vector_store_id = data.get("vector_store_id")
|
||||
|
||||
|
||||
if vector_store_id:
|
||||
# Check if it's a managed vector store ID
|
||||
decoded_id = is_base64_encoded_unified_id(vector_store_id)
|
||||
|
||||
|
||||
if decoded_id:
|
||||
verbose_logger.debug(
|
||||
f"Processing managed vector store search: {vector_store_id}"
|
||||
)
|
||||
|
||||
verbose_logger.debug(f"Processing managed vector store search: {vector_store_id}")
|
||||
|
||||
# Check access
|
||||
has_access = await self.can_user_access_unified_resource_id(
|
||||
vector_store_id, user_api_key_dict
|
||||
)
|
||||
|
||||
has_access = await self.can_user_access_unified_resource_id(vector_store_id, user_api_key_dict)
|
||||
|
||||
if not has_access:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"User {user_api_key_dict.user_id} does not have access to vector store {vector_store_id}",
|
||||
)
|
||||
|
||||
|
||||
# Parse the unified ID to extract components
|
||||
parsed_id = parse_unified_id(vector_store_id)
|
||||
|
||||
|
||||
if parsed_id:
|
||||
# Extract the model ID and provider resource 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}"
|
||||
)
|
||||
|
||||
|
||||
# Determine which model to use for routing
|
||||
# Priority: model_id (deployment ID) > first target_model_name
|
||||
routing_model = None
|
||||
|
|
@ -371,28 +347,26 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
routing_model = model_id
|
||||
elif target_model_names and len(target_model_names) > 0:
|
||||
routing_model = target_model_names[0]
|
||||
|
||||
|
||||
# Set the model for routing
|
||||
if routing_model:
|
||||
data["model"] = routing_model
|
||||
verbose_logger.info(
|
||||
f"Routing vector store search to model: {routing_model}"
|
||||
)
|
||||
|
||||
verbose_logger.info(f"Routing vector store search to model: {routing_model}")
|
||||
|
||||
# Replace the unified ID with the provider-specific ID
|
||||
if provider_resource_id:
|
||||
data["vector_store_id"] = provider_resource_id
|
||||
verbose_logger.debug(
|
||||
f"Replaced unified ID with provider resource ID: {provider_resource_id}"
|
||||
)
|
||||
|
||||
|
||||
# Handle vector store retrieve/delete operations
|
||||
elif call_type in ("avector_store_retrieve", "avector_store_delete"):
|
||||
await self.check_managed_vector_store_access(data, user_api_key_dict)
|
||||
|
||||
|
||||
# If it's a managed vector store, we'll handle it in the endpoint
|
||||
# No need to transform here as the endpoint will route to the hook
|
||||
|
||||
|
||||
return data
|
||||
|
||||
# ============================================================================
|
||||
|
|
@ -401,21 +375,21 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
|
||||
async def async_post_call_success_hook(
|
||||
self,
|
||||
data: Dict,
|
||||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: LLMResponseTypes,
|
||||
) -> LLMResponseTypes:
|
||||
"""
|
||||
Post-call hook to transform responses.
|
||||
|
||||
|
||||
This hook can be used to transform responses if needed.
|
||||
For now, it just passes through the response.
|
||||
|
||||
|
||||
Args:
|
||||
data: Request data
|
||||
user_api_key_dict: User API key authentication details
|
||||
response: Response from the provider
|
||||
|
||||
|
||||
Returns:
|
||||
Potentially modified response
|
||||
"""
|
||||
|
|
@ -429,28 +403,28 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
async def async_filter_deployments( # type: ignore[override]
|
||||
self,
|
||||
model: str,
|
||||
healthy_deployments: List,
|
||||
messages: Optional[List] = None,
|
||||
request_kwargs: Optional[Dict] = None,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
) -> List[Dict]:
|
||||
healthy_deployments: list,
|
||||
messages: list | None = None,
|
||||
request_kwargs: dict | None = None,
|
||||
parent_otel_span: Span | None = None,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
Filter deployments based on vector store availability.
|
||||
|
||||
|
||||
This is used by the router to select only deployments that have
|
||||
the vector store available.
|
||||
|
||||
|
||||
Note: This method signature is a compromise between CustomLogger and BaseManagedResource
|
||||
parent classes which have incompatible signatures. The type: ignore[override] is necessary
|
||||
due to this multiple inheritance conflict.
|
||||
|
||||
|
||||
Args:
|
||||
model: Model name
|
||||
healthy_deployments: List of healthy deployments
|
||||
messages: Messages (unused for vector stores, required by CustomLogger interface)
|
||||
request_kwargs: Request kwargs containing vector_store_id and mappings
|
||||
parent_otel_span: OpenTelemetry span for tracing
|
||||
|
||||
|
||||
Returns:
|
||||
Filtered list of deployments
|
||||
"""
|
||||
|
|
@ -462,4 +436,6 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
parent_otel_span=parent_otel_span,
|
||||
resource_id_key="vector_store_id",
|
||||
)
|
||||
|
||||
|
||||
PROXY_LiteLLMManagedVectorStores = _PROXY_LiteLLMManagedVectorStores
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue