diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index aac4469be20..626a3b0bb83 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -78,6 +78,7 @@ from litellm.proxy.litellm_pre_call_utils import ( from litellm.proxy.openai_files_endpoints.common_utils import ( BATCH_CREATE_HIDDEN_PARAM, FILE_LIST_CONTINUATION_CHUNK_SIZE, + TeamProviderDeployment, _is_base64_encoded_unified_file_id, apply_unified_file_ids, decode_model_from_file_id, @@ -1426,6 +1427,43 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): user_api_key_dict=owner, ) + async def register_provider_file( + self, + *, + provider_file_id: str, + content_type: str, + deployment: TeamProviderDeployment, + owner: UserAPIKeyAuth, + llm_router: Router | None, + ) -> str | None: + unified_file_id: Final = self.get_unified_output_file_id( + provider_file_id, deployment.deployment_id, deployment.public_model_name, file_type=content_type + ) + provider_file_object, _ = await self._fetch_provider_file_object( + unified_file_id=unified_file_id, + provider_file_id=provider_file_id, + model_id=deployment.deployment_id, + model_name=None, + llm_router=llm_router, + ) + if provider_file_object is None: + return None + try: + await self.store_unified_file_id( + file_id=unified_file_id, + file_object=provider_file_object, + litellm_parent_otel_span=None, + model_mappings={deployment.deployment_id: provider_file_id}, + user_api_key_dict=owner, + ) + except Exception as error: # noqa: BLE001 # the upload succeeded, so the caller keeps the provider id + verbose_logger.warning( + f"Failed to register provider_file_id={sanitize_for_log(provider_file_id)} as a managed file: " + f"{type(error).__name__} {sanitize_for_log(error)}" + ) + return None + return unified_file_id + async def list_user_batches( self, user_api_key_dict: UserAPIKeyAuth, @@ -2253,10 +2291,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id) return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=") - def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: Optional[str]) -> str: + def get_unified_output_file_id( + self, output_file_id: str, model_id: str, model_name: Optional[str], file_type: str = "application/json" + ) -> str: deterministic_uuid: Final = uuid5(uuid5(NAMESPACE_URL, model_id), output_file_id) unified_output_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( - "application/json", + file_type, str(deterministic_uuid), model_name or "", output_file_id, diff --git a/litellm/litellm_core_utils/credential_accessor.py b/litellm/litellm_core_utils/credential_accessor.py index b2310894c3a..bab250fe583 100644 --- a/litellm/litellm_core_utils/credential_accessor.py +++ b/litellm/litellm_core_utils/credential_accessor.py @@ -1,6 +1,7 @@ """Utils for accessing credentials.""" from collections.abc import Mapping +from copy import deepcopy from typing import Final import litellm @@ -27,7 +28,9 @@ class CredentialAccessor: """Safe accessor for credentials.""" credential: Final = CredentialAccessor.find_credential(credential_name) - return {} if credential is None else dict(credential_api_key_secret_reference(credential.credential_values)) + if credential is None: + return {} + return deepcopy(dict(credential_api_key_secret_reference(credential.credential_values))) @staticmethod def upsert_credentials(credentials: list[CredentialItem]): diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 1ae0d475d8c..9a044f1dee3 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -21,10 +21,13 @@ from typing import ( from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict +from litellm._logging import verbose_proxy_logger from litellm.constants import MAX_FILE_LIST_LIMIT from litellm.litellm_core_utils.batch_utils import batch_cost_is_final from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.llms.base_llm.managed_resources.utils import is_base64_encoded_unified_id from litellm.proxy._types import ProxyException +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.repositories.managed_file_repository import ManagedFileRepository from litellm.repositories.table_repositories import ( ManagedObjectRepository, @@ -104,6 +107,13 @@ class ManagedResourceAccessChecker(Protocol): ) -> bool: ... +@dataclass(frozen=True, slots=True) +class TeamProviderDeployment: + deployment_id: str + public_model_name: str + credentials: Mapping[str, object] + + @runtime_checkable class ManagedFileIdResolver(Protocol): async def get_unified_file_ids_for_provider_file_ids( @@ -138,6 +148,86 @@ class ManagedBatchOutputFileWriter(Protocol): ... +@runtime_checkable +class ManagedFileModelMappingReader(Protocol): + async def get_model_file_id_mapping( + self, + file_ids: Sequence[str], + litellm_parent_otel_span: "Span | None", + ) -> Mapping[str, Mapping[str, str]]: ... + + +@runtime_checkable +class ManagedProviderFileRegistrar(Protocol): + async def register_provider_file( + self, + *, + provider_file_id: str, + content_type: str, + deployment: TeamProviderDeployment, + owner: "UserAPIKeyAuth", + llm_router: "Router | None", + ) -> str | None: ... + + +async def managed_file_copy_on_deployment( + file_id: str, + deployment: TeamProviderDeployment | None, + managed_files_obj: object | None, + litellm_parent_otel_span: "Span | None", +) -> str | None: + if deployment is None or not isinstance(managed_files_obj, ManagedFileModelMappingReader): + return None + mappings: Final = await managed_files_obj.get_model_file_id_mapping((file_id,), litellm_parent_otel_span) + model_mappings: Final = mappings.get(file_id) + return None if model_mappings is None else model_mappings.get(deployment.deployment_id) + + +def provider_file_id_from_managed_id(managed_file_id: str) -> str: + decoded_id: Final = is_base64_encoded_unified_id(managed_file_id) + if not decoded_id: + return managed_file_id + + match: Final = re.search(r"(?:^|;)llm_output_file_id,([^;]+)", decoded_id) + return match.group(1).strip() if match else managed_file_id + + +async def _unified_file_ids_or_empty_on_db_failure( + managed_files_obj: ManagedFileIdResolver, + provider_file_ids: Sequence[str], + user_api_key_dict: "UserAPIKeyAuth", +) -> Mapping[str, str]: + try: + return await managed_files_obj.get_unified_file_ids_for_provider_file_ids( + provider_file_ids=provider_file_ids, + user_api_key_dict=user_api_key_dict, + ) + except Exception as e: + if not PrismaDBExceptionHandler.is_database_infrastructure_error(e): + raise + verbose_proxy_logger.warning("Managed file id lookup failed, keeping provider file ids: %s", e) + return MappingProxyType({}) + + +async def managed_file_ids_for_provider_file_ids( + provider_file_ids: Sequence[str], + managed_files_obj: object | None, + user_api_key_dict: "UserAPIKeyAuth", +) -> Mapping[str, str]: + if not provider_file_ids or not isinstance(managed_files_obj, ManagedFileIdResolver): + return MappingProxyType({}) + id_map: Final = await _unified_file_ids_or_empty_on_db_failure( + managed_files_obj, tuple(dict.fromkeys(provider_file_ids)), user_api_key_dict + ) + return MappingProxyType( + { + provider_file_id: managed_file_id + for provider_file_id, managed_file_id in id_map.items() + if provider_file_id_from_managed_id(managed_file_id) == provider_file_id + } + ) + + def is_base64_encoded_unified_file_id(b64_uid: object) -> str | Literal[False]: # Ensure b64_uid is a string and not a mock object if not isinstance(b64_uid, str): @@ -464,6 +554,30 @@ async def authorize_model_for_key( ) +async def team_deployment_the_caller_may_use( + deployment: TeamProviderDeployment | None, + llm_router: Optional["Router"], + user_api_key_dict: "UserAPIKeyAuth", +) -> TeamProviderDeployment | None: + if deployment is None: + return None + return ( + deployment if await caller_may_use_model(deployment.public_model_name, llm_router, user_api_key_dict) else None + ) + + +async def caller_may_use_model( + public_model_name: str, + llm_router: Optional["Router"], + user_api_key_dict: "UserAPIKeyAuth", +) -> bool: + try: + await authorize_model_for_key(public_model_name, llm_router, user_api_key_dict) + except ProxyException: + return False + return True + + async def get_authorized_credentials_for_model( llm_router: Optional["Router"], model_id: str, @@ -522,8 +636,20 @@ def _team_deployment_rank(model_info: Mapping[str, object], deployment_id: str) return (created_at if created_at.tzinfo else created_at.replace(tzinfo=timezone.utc)), deployment_id -def _by_team_deployment_rank(order: tuple[tuple[datetime, str], str]) -> tuple[datetime, str]: - return order[0] +@dataclass(frozen=True, slots=True) +class _TeamDeploymentOrder: + rank: tuple[datetime, str] + deployment_id: str + public_model_name: str + + +def _by_team_deployment_rank(order: _TeamDeploymentOrder) -> tuple[datetime, str]: + return order.rank + + +@dataclass(frozen=True, slots=True) +class _AccessibleProviderCredentials: + credentials: Mapping[str, object] def get_team_provider_credentials( @@ -558,6 +684,33 @@ def get_team_provider_credentials( Returns None when the router is unavailable or no authorized deployment matches, so the caller can fall back to default credential resolution. """ + resolved: Final = _resolve_provider_credentials( + llm_router, + user_api_key_dict, + custom_llm_provider, + include_accessible_deployments=include_accessible_deployments, + ) + return None if resolved is None else dict(resolved.credentials) + + +def get_team_provider_deployment( + llm_router: Optional["Router"], + user_api_key_dict: "UserAPIKeyAuth", + custom_llm_provider: str, +) -> TeamProviderDeployment | None: + resolved: Final = _resolve_provider_credentials( + llm_router, user_api_key_dict, custom_llm_provider, include_accessible_deployments=False + ) + return resolved if isinstance(resolved, TeamProviderDeployment) else None + + +def _resolve_provider_credentials( + llm_router: Optional["Router"], + user_api_key_dict: "UserAPIKeyAuth", + custom_llm_provider: str, + *, + include_accessible_deployments: bool, +) -> TeamProviderDeployment | _AccessibleProviderCredentials | None: if llm_router is None: return None @@ -603,16 +756,17 @@ def get_team_provider_credentials( return {key: value for key, value in credentials.items() if key != "model"} return None - def _usable_team_deployment_order(deployment: _RouterDeployment) -> tuple[tuple[datetime, str], str] | None: + def _usable_team_deployment_order(deployment: _RouterDeployment) -> _TeamDeploymentOrder | None: model_info: Final[Mapping[str, object]] = deployment["model_info"] if model_info.get("team_id") != team_id: return None - if not _key_may_use(model_info.get("team_public_model_name") or deployment.get("model_name")): + public_model_name: Final = model_info.get("team_public_model_name") or deployment.get("model_name") + if not isinstance(public_model_name, str) or not _key_may_use(public_model_name): return None deployment_id: Final = model_info.get("id") if not isinstance(deployment_id, str): return None - return _team_deployment_rank(model_info, deployment_id), deployment_id + return _TeamDeploymentOrder(_team_deployment_rank(model_info, deployment_id), deployment_id, public_model_name) # 1. Prefer the team's own BYOK deployment, matched by model_info.team_id. Config file # deployments win, kept in file order by the stable sort, then DB ones oldest first, then @@ -627,13 +781,16 @@ def get_team_provider_credentials( ), key=_by_team_deployment_rank, ) - team_deployment_ids: Final = tuple(deployment_id for _, deployment_id in team_deployment_orders) - team_credentials: Final = next( - (credentials for credentials in map(_provider_credentials, team_deployment_ids) if credentials is not None), + team_deployment: Final = next( + ( + TeamProviderDeployment(order.deployment_id, order.public_model_name, MappingProxyType(credentials)) + for order in team_deployment_orders + if (credentials := _provider_credentials(order.deployment_id)) is not None + ), None, ) - if team_credentials is not None: - return team_credentials + if team_deployment is not None: + return team_deployment if not include_accessible_deployments: return None @@ -669,7 +826,7 @@ def get_team_provider_credentials( is_model_allowed_by_pattern(model=model_name, allowed_model_pattern=granted) for granted in granted_models ) ) - return next( + accessible_credentials: Final = next( ( credentials for credentials in map( @@ -679,6 +836,7 @@ def get_team_provider_credentials( ), None, ) + return None if accessible_credentials is None else _AccessibleProviderCredentials(accessible_credentials) HTTP_HEADER_VALUE: Final = re.compile(r"[\x21-\x7e]+(?:[ \t]+[\x21-\x7e]+)*") @@ -707,6 +865,10 @@ def _request_api_key(data: Mapping[str, object]) -> str | None: ) +def carries_explicit_credentials(data: Mapping[str, object]) -> bool: + return any(_non_blank_string(data.get(key)) is not None for key in EXPLICIT_CREDENTIAL_KEYS) + + def explicit_request_credentials(data: Mapping[str, object]) -> Mapping[str, str]: api_key: Final = _request_api_key(data) if api_key is None: @@ -1244,6 +1406,17 @@ async def validate_managed_id_requirement( ), ) + await authorize_managed_resource_id(resource_id, resource_kind, user_api_key_dict, managed_files_obj) + + +async def authorize_managed_resource_id( + resource_id: str, + resource_kind: Literal["file", "batch", "fine-tuning job"], + user_api_key_dict: "UserAPIKeyAuth", + managed_files_obj: object | None, +) -> None: + from fastapi import HTTPException + if not isinstance(managed_files_obj, ManagedResourceAccessChecker): raise HTTPException( status_code=500, diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 2d81dc32181..f4710fa30e3 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -11,6 +11,7 @@ import json from collections.abc import Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final +from urllib.parse import urlsplit import orjson from fastapi import APIRouter, Depends, HTTPException, Request, Response @@ -41,8 +42,17 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg ) from litellm.proxy.openai_files_endpoints.common_utils import ( EXPLICIT_CREDENTIAL_KEYS, + ManagedProviderFileRegistrar, + ManagedResourceAccessChecker, + TeamProviderDeployment, + authorize_managed_resource_id, + detect_content_type_from_filename, explicit_request_credentials, - get_team_provider_credentials, + get_team_provider_deployment, + is_base64_encoded_unified_file_id, + managed_file_copy_on_deployment, + provider_file_id_from_managed_id, + team_deployment_the_caller_may_use, ) from litellm.proxy.rag_endpoints.upload_security import ( MAX_UPLOAD_SIZE_BYTES, @@ -58,10 +68,13 @@ from litellm.proxy.vector_store_endpoints.endpoints import ( ) from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, + with_managed_search_file_ids, ) from litellm.rag.main import get_ingestion_class +from litellm.rag.rag_query import RAGQuery from litellm.repositories.table_repositories import ManagedVectorStoresRepository -from litellm.types.utils import ModelResponse +from litellm.types.rag import RAGIngestResponse +from litellm.types.utils import Choices, ModelResponse if TYPE_CHECKING: from litellm import Router @@ -260,31 +273,115 @@ def _with_request_credentials( return MappingProxyType({**team_credentials, **explicit_request_credentials(request_data)}) -def _team_provider_credentials( +def _team_provider_deployment( effective_config: Mapping[str, object], llm_router: "Router | None", user_api_key_dict: UserAPIKeyAuth, -) -> Mapping[str, object]: +) -> TeamProviderDeployment | None: if not EXPLICIT_CREDENTIAL_KEYS.isdisjoint(effective_config): - return _EMPTY_TEAM_PROVIDER_CREDENTIALS + return None if _store_provider(effective_config) != _TEAM_CREDENTIAL_PROVIDER: - return _EMPTY_TEAM_PROVIDER_CREDENTIALS + return None - validated_credentials: Final = _TEAM_PROVIDER_CREDENTIALS_ADAPTER.validate_python( - get_team_provider_credentials( - llm_router=llm_router, - user_api_key_dict=user_api_key_dict, - custom_llm_provider=_TEAM_CREDENTIAL_PROVIDER, - include_accessible_deployments=False, - ) - or {} + return get_team_provider_deployment( + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=_TEAM_CREDENTIAL_PROVIDER, ) + + +def _team_provider_credentials(deployment: TeamProviderDeployment | None) -> Mapping[str, object]: + if deployment is None: + return _EMPTY_TEAM_PROVIDER_CREDENTIALS return MappingProxyType( - {key: value for key, value in validated_credentials.items() if key not in {"custom_llm_provider", "model"}} + {key: value for key, value in deployment.credentials.items() if key not in {"custom_llm_provider", "model"}} ) +async def _provider_file_id_for_ingest( + file_id: str | None, + deployment: TeamProviderDeployment | None, + managed_files_obj: object | None, + user_api_key_dict: UserAPIKeyAuth, +) -> str | None: + if file_id is None or not is_base64_encoded_unified_file_id(file_id): + return file_id + if not isinstance(managed_files_obj, ManagedResourceAccessChecker): + raise HTTPException( + status_code=400, + detail={ + "error": "This proxy does not accept LiteLLM managed file ids. Send the provider's file id instead." + }, + ) + await authorize_managed_resource_id(file_id, "file", user_api_key_dict, managed_files_obj) + deployment_copy: Final = await managed_file_copy_on_deployment( + file_id, deployment, managed_files_obj, user_api_key_dict.parent_otel_span + ) + return deployment_copy or provider_file_id_from_managed_id(file_id) + + +def _first_choice_search_results(response: ModelResponse) -> Mapping[str, object] | None: + first_choice: Final = response.choices[0] if response.choices else None + if not isinstance(first_choice, Choices) or first_choice.message.provider_specific_fields is None: + return None + return _as_string_keyed_mapping(first_choice.message.provider_specific_fields.get("search_results")) + + +async def _with_managed_search_result_ids( + response: ModelResponse, + managed_files_obj: object | None, + user_api_key_dict: UserAPIKeyAuth, +) -> ModelResponse: + search_results: Final = _first_choice_search_results(response) + if search_results is None: + return response + managed_results: Final = await with_managed_search_file_ids(search_results, managed_files_obj, user_api_key_dict) + if managed_results is search_results: + return response + return RAGQuery.add_search_results_to_response(response, managed_results) + + +def _uploaded_content_type(file_data: tuple[str, bytes, str] | None, file_url: str | None) -> str: + if file_data is not None: + return file_data[2] + return detect_content_type_from_filename(urlsplit(file_url or "").path.rsplit("/", 1)[-1]) + + +def _with_caller_file_id(response: RAGIngestResponse, caller_file_id: str | None) -> RAGIngestResponse: + caller_response: Final[RAGIngestResponse] = {**response, "file_id": caller_file_id} + return caller_response + + +async def _with_registered_upload( + response: RAGIngestResponse, + *, + upload_deployment: TeamProviderDeployment | None, + content_type: str, + managed_files_obj: object | None, + user_api_key_dict: UserAPIKeyAuth, + llm_router: "Router | None", +) -> RAGIngestResponse: + provider_file_id: Final = response.get("file_id") + if ( + upload_deployment is None + or provider_file_id is None + or not isinstance(managed_files_obj, ManagedProviderFileRegistrar) + ): + return response + managed_file_id: Final = await managed_files_obj.register_provider_file( + provider_file_id=provider_file_id, + content_type=content_type, + deployment=upload_deployment, + owner=user_api_key_dict, + llm_router=llm_router, + ) + if managed_file_id is None: + return response + managed_response: Final[RAGIngestResponse] = {**response, "file_id": managed_file_id} + return managed_response + + def _build_file_metadata_entry( response: object, file_data: tuple[str, bytes, str] | None = None, @@ -663,6 +760,7 @@ async def rag_ingest( llm_router, prisma_client, proxy_config, + proxy_logging_obj, version, ) @@ -704,8 +802,9 @@ async def rag_ingest( forwarded_credentials: Final = _forwarded_store_credentials( base, _with_body_drop_params(request_data, await _rag_ingest_body(request)) ) + team_deployment: Final = _team_provider_deployment(base, llm_router, user_api_key_dict) merged_vector_store_config: Final = { - **_team_provider_credentials(base, llm_router, user_api_key_dict), + **_team_provider_credentials(team_deployment), **forwarded_credentials, **base, } @@ -727,20 +826,41 @@ async def rag_ingest( merged_vector_store_config.get("custom_llm_provider", "openai"), ) + managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") + credential_deployment: Final = None if forwarded_credentials else team_deployment + upload_deployment: Final = await team_deployment_the_caller_may_use( + credential_deployment, llm_router, user_api_key_dict + ) + ingest_file_id: Final = await _provider_file_id_for_ingest( + file_id, credential_deployment, managed_files_obj, user_api_key_dict + ) + # Call ingest - response: Final = await litellm.aingest( + ingest_response: Final = await litellm.aingest( ingest_options=merged_ingest_options, file_data=file_data, file_url=file_url, - file_id=file_id, + file_id=ingest_file_id, router=llm_router, **request_data, ) - if response.get("status") == "failed": + if ingest_response.get("status") == "failed": raise HTTPException( - status_code=response.get("error_status_code") or 500, - detail={"id": response.get("id"), "status": "failed", "error": response.get("error")}, + status_code=ingest_response.get("error_status_code") or 500, + detail={"id": ingest_response.get("id"), "status": "failed", "error": ingest_response.get("error")}, ) + response: Final = ( + await _with_registered_upload( + ingest_response, + upload_deployment=upload_deployment, + content_type=_uploaded_content_type(file_data, file_url), + managed_files_obj=managed_files_obj, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + ) + if file_data is not None or bool(file_url) + else _with_caller_file_id(ingest_response, file_id) + ) # Save vector store to database if it was newly created and prisma_client is available verbose_proxy_logger.debug( @@ -845,6 +965,7 @@ async def rag_query( general_settings, llm_router, proxy_config, + proxy_logging_obj, select_data_generator, version, ) @@ -935,7 +1056,9 @@ async def rag_query( proxy_config=proxy_config, ) team_credentials: Final = _with_request_credentials( - _team_provider_credentials(effective_vector_store_config, llm_router, user_api_key_dict), + _team_provider_credentials( + _team_provider_deployment(effective_vector_store_config, llm_router, user_api_key_dict) + ), _with_body_drop_params(request_data, data), ) vector_store_params: Final = ( @@ -996,7 +1119,9 @@ async def rag_query( response: Final = await query() fastapi_response.headers.update(custom_headers_for(response)) - return response + return await _with_managed_search_result_ids( + response, proxy_logging_obj.get_proxy_hook("managed_files"), user_api_key_dict + ) except HTTPException: raise diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index ffc0e31408d..b1ebd042f3d 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -14,13 +14,20 @@ from fastapi import HTTPException, status from pydantic import BeforeValidator, TypeAdapter, ValidationError import litellm -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.routing_common_utils import is_proxy_admin_request from litellm.litellm_core_utils.total_deadline import enforce_total_timeout_scope from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.openai_files_endpoints.common_utils import ( + TeamProviderDeployment, + carries_explicit_credentials, + get_team_provider_deployment, + team_deployment_the_caller_may_use, +) from litellm.router import router_arguments from litellm.router_api import router_init_args +from litellm.router_attempt import metadata_field from litellm.router_config import build_router_config # Client-supplied params that make the router or the call path fabricate a @@ -715,6 +722,73 @@ RouteType = Literal[ ] +_MODEL_LESS_VECTOR_STORE_FILE_ROUTES: Final[frozenset[str]] = frozenset( + { + "avector_store_file_create", + "avector_store_file_list", + "avector_store_file_retrieve", + "avector_store_file_content", + "avector_store_file_update", + "avector_store_file_delete", + } +) + + +@dataclass(frozen=True, slots=True) +class _ModelLessTeamFileCall: + caller: UserAPIKeyAuth + custom_llm_provider: str + + +def _model_less_team_vector_store_file_call( + route_type: str, + request: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth | None, +) -> _ModelLessTeamFileCall | None: + custom_llm_provider: Final = request.get("custom_llm_provider") + if ( + route_type not in _MODEL_LESS_VECTOR_STORE_FILE_ROUTES + or user_api_key_dict is None + or user_api_key_dict.team_id is None + or request.get("model") + or carries_explicit_credentials(request) + or not isinstance(custom_llm_provider, str) + ): + return None + return _ModelLessTeamFileCall(user_api_key_dict, custom_llm_provider) + + +async def _authorized_team_file_deployment( + llm_router: LitellmRouter | None, call: _ModelLessTeamFileCall +) -> TeamProviderDeployment | None: + return await team_deployment_the_caller_may_use( + get_team_provider_deployment(llm_router, call.caller, call.custom_llm_provider), llm_router, call.caller + ) + + +def _with_team_deployment_attribution( + kwargs: Mapping[str, object], + team_deployment: TeamProviderDeployment, + llm_router: LitellmRouter | None, + route_type: str, +) -> Mapping[str, object]: + deployment: Final = None if llm_router is None else llm_router.registry().deployment(team_deployment.deployment_id) + if deployment is None: + return kwargs + field: Final = metadata_field(route_type) + return { + **kwargs, + field: { + **(as_str_mapping(kwargs.get(field)) or {}), + "model_group": team_deployment.public_model_name, + "deployment": deployment.litellm_params.model, + "model_info": deployment.model_info.model_dump(), + "api_base": deployment.litellm_params.api_base, + "deployment_model_name": deployment.model_name, + }, + } + + def _merge_router_settings_override( data: dict[str, object], # mutable-ok: route_request mutates the shared request body in place ) -> frozenset[str]: @@ -920,6 +994,23 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr "alist_agent_versions", ]: return getattr(llm_router, f"{route_type}")(**data) + model_less_team_file_call: Final = _model_less_team_vector_store_file_call( + route_type, JSON_OBJECT_ADAPTER.validate_python(data), user_api_key_dict + ) + team_file_deployment: Final = ( + None + if model_less_team_file_call is None + else await _authorized_team_file_deployment(llm_router, model_less_team_file_call) + ) + if team_file_deployment is not None: + return getattr(litellm, f"{route_type}")( + **_with_team_deployment_attribution( + {**direct_litellm_kwargs(), **team_file_deployment.credentials}, + team_file_deployment, + llm_router, + route_type, + ) + ) if route_type in [ "avideo_list", "avideo_status", diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index c0d7ab8d6e1..a6f9ba2ef14 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -15,7 +15,12 @@ from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) -from litellm.litellm_core_utils.core_helpers import as_object_list, decoded_json_object, require_attributes +from litellm.litellm_core_utils.core_helpers import ( + as_object_list, + as_str_mapping, + decoded_json_object, + require_attributes, +) from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -26,6 +31,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, assert_user_can_access_vector_store, get_litellm_managed_vector_store, + with_managed_search_file_ids, ) from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository from litellm.types.utils import LlmProviders @@ -198,7 +204,7 @@ async def vector_store_search( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + search_response: Final = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -216,6 +222,12 @@ async def vector_store_search( user_api_base=user_api_base, version=version, ) + search_page: Final = as_str_mapping(search_response) + if search_page is None: + return search_response + return await with_managed_search_file_ids( + search_page, proxy_logging_obj.get_proxy_hook("managed_files"), user_api_key_dict + ) except Exception as e: raise await processor.handle_llm_api_exception( e=e, diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 4dec4c7b38b..9004d1951b8 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -9,6 +9,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LitellmUserRoles, @@ -18,6 +19,7 @@ from litellm.proxy.mcp_server.ui_session_utils import ( is_ui_session_credential, resolve_ui_session_team_ids, ) +from litellm.proxy.openai_files_endpoints.common_utils import managed_file_ids_for_provider_file_ids from litellm.types.utils import LlmProviders from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager @@ -581,3 +583,35 @@ def is_allowed_to_call_vector_store_files_endpoint( ) return has_permission + + +def _search_hit_file_id(hit: object) -> str | None: + fields: Final = as_str_mapping(hit) + file_id: Final = None if fields is None else fields.get("file_id") + return file_id if isinstance(file_id, str) else None + + +def _with_managed_search_file_id(hit: object, id_map: Mapping[str, str]) -> object: + fields: Final = as_str_mapping(hit) + file_id: Final = _search_hit_file_id(hit) + if fields is None or file_id is None or file_id not in id_map: + return hit + return {**fields, "file_id": id_map[file_id]} + + +async def with_managed_search_file_ids( + search_page: Mapping[str, object], + managed_files_obj: object | None, + user_api_key_dict: UserAPIKeyAuth, +) -> Mapping[str, object]: + hits: Final = as_object_list(search_page.get("data")) + if hits is None: + return search_page + id_map: Final = await managed_file_ids_for_provider_file_ids( + tuple(file_id for hit in hits if (file_id := _search_hit_file_id(hit)) is not None), + managed_files_obj, + user_api_key_dict, + ) + if not id_map: + return search_page + return {**search_page, "data": [_with_managed_search_file_id(hit, id_map) for hit in hits]} diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index dbe08dbd98e..a561362a307 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -1,13 +1,20 @@ import re from collections.abc import Mapping +from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional, cast +import openai from fastapi import APIRouter, Depends, Request, Response from fastapi.responses import ORJSONResponse +from pydantic import TypeAdapter import litellm -from litellm.llms.base_llm.managed_resources.utils import is_base64_encoded_unified_id +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.llms.base_llm.managed_resources.utils import ( + extract_target_model_names_from_unified_id, + is_base64_encoded_unified_id, +) 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,23 +24,33 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_query, ) from litellm.proxy.openai_files_endpoints.common_utils import ( - ManagedFileIdResolver, authorize_model_for_key, + caller_may_use_model, + carries_explicit_credentials, + explicit_request_credentials, get_credentials_for_model, + get_team_provider_deployment, handle_model_based_routing, + managed_file_copy_on_deployment, + managed_file_ids_for_provider_file_ids, prepare_data_with_credentials, + provider_file_id_from_managed_id, + team_deployment_the_caller_may_use, ) from litellm.proxy.rag_endpoints.upload_security import safe_download_headers from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, is_allowed_to_call_vector_store_files_endpoint, ) +from litellm.router_rust.invoke.endpoints import PINNED_DEPLOYMENT_KEY +from litellm.types.router import Deployment from litellm.types.utils import LlmProviders from litellm.types.vector_store_files import ( VectorStoreFileListResponse, VectorStoreFileObject, ) from litellm.types.vector_stores import LiteLLM_ManagedVectorStore +from litellm.vector_store_files.main import aretrieve as vector_store_file_aretrieve if TYPE_CHECKING: from litellm.router import Router @@ -41,24 +58,12 @@ if TYPE_CHECKING: router: Final = APIRouter() -def _provider_file_id_from_managed_id(managed_file_id: str | None) -> str | None: - if managed_file_id is None: - return None - - decoded_id: Final = is_base64_encoded_unified_id(managed_file_id) - if not decoded_id: - return managed_file_id - - match: Final = re.search(r"(?:^|;)llm_output_file_id,([^;]+)", decoded_id) - return match.group(1).strip() if match else managed_file_id - - def _with_provider_file_id_cursors( query_params: Mapping[str, str], ) -> Mapping[str, str | None]: return MappingProxyType( { - key: (_provider_file_id_from_managed_id(value) if key in {"after", "before"} else value) + key: (provider_file_id_from_managed_id(value) if key in {"after", "before"} else value) for key, value in query_params.items() } ) @@ -108,24 +113,114 @@ async def _with_managed_file_list_ids( user_api_key_dict: UserAPIKeyAuth, ) -> VectorStoreFileListResponse: data: Final = response.get("data") - if not data or not isinstance(managed_files_obj, ManagedFileIdResolver): + if not data: return response - provider_file_ids: Final = tuple( - dict.fromkeys(provider_file_id for file in data if isinstance(provider_file_id := file.get("id"), str)) + id_map: Final = await managed_file_ids_for_provider_file_ids( + tuple(provider_file_id for file in data if isinstance(provider_file_id := file.get("id"), str)), + managed_files_obj, + user_api_key_dict, ) - id_map: Final = await managed_files_obj.get_unified_file_ids_for_provider_file_ids( - provider_file_ids=provider_file_ids, - user_api_key_dict=user_api_key_dict, + return _with_managed_file_ids(response, id_map) + + +@dataclass(frozen=True, slots=True) +class _ManagedFileCopy: + deployment_id: str + provider_file_id: str + + +_PROBE_ARGUMENTS: Final = frozenset( + {"vector_store_id", "file_id", "extra_headers", "timeout", "custom_llm_provider", "no-log"} +) +_EXTRA_HEADERS: Final = TypeAdapter(dict[str, str] | None) + + +async def _store_lacks_file( + provider: LlmProviders, vector_store_id: str, file_id: str, credentials: Mapping[str, object] +) -> bool: + try: + await vector_store_file_aretrieve( + vector_store_id=vector_store_id, + file_id=file_id, + extra_headers=_EXTRA_HEADERS.validate_python(credentials.get("extra_headers")), + timeout=None, + custom_llm_provider=provider.value, + **{key: value for key, value in credentials.items() if key not in _PROBE_ARGUMENTS}, + **{"no-log": True}, + ) + except openai.NotFoundError: + return True + except openai.APIError: + return False + return False + + +def _raw_vector_store_id(data: Mapping[str, object]) -> str | None: + vector_store_id: Final = data.get("vector_store_id") + if not isinstance(vector_store_id, str) or is_base64_encoded_unified_id(vector_store_id): + return None + return vector_store_id + + +async def _caller_team_copy_of_managed_file( + provider: LlmProviders | None, + file_id: str, + upload_copy: str | None, + upload_credentials: Mapping[str, object] | None, + vector_store_id: str, + user_api_key_dict: UserAPIKeyAuth, + managed_files_obj: object | None, + llm_router: Optional["Router"], +) -> _ManagedFileCopy | None: + if provider is None: + return None + team_deployment: Final = await team_deployment_the_caller_may_use( + get_team_provider_deployment(llm_router, user_api_key_dict, provider.value), llm_router, user_api_key_dict ) - round_trippable_id_map: Final = MappingProxyType( - { - provider_file_id: managed_file_id - for provider_file_id, managed_file_id in id_map.items() - if _provider_file_id_from_managed_id(managed_file_id) == provider_file_id - } + provider_file_id: Final = await managed_file_copy_on_deployment( + file_id, team_deployment, managed_files_obj, user_api_key_dict.parent_otel_span ) - return _with_managed_file_ids(response, round_trippable_id_map) + if team_deployment is None or provider_file_id is None: + return None + team_copy: Final = _ManagedFileCopy(team_deployment.deployment_id, provider_file_id) + if provider_file_id == upload_copy: + return team_copy + store_lacks_upload_copy: Final = ( + upload_copy is None + or upload_credentials is None + or await _store_lacks_file(provider, vector_store_id, upload_copy, upload_credentials) + ) + return team_copy if store_lacks_upload_copy else None + + +def _store_carries_credentials(store: LiteLLM_ManagedVectorStore | None) -> bool: + return store is not None and ( + carries_explicit_credentials(store.get("litellm_params") or MappingProxyType({})) + or carries_explicit_credentials( + MappingProxyType({"litellm_credential_name": store.get("litellm_credential_name")}) + ) + ) + + +_CALLER_DESTINATION_KEYS: Final = frozenset({"api_base", "base_url"}) + + +def _keys_a_pinned_call_drops(data: Mapping[str, object]) -> frozenset[str]: + if explicit_request_credentials(data): + return frozenset({"custom_llm_provider"}) + return frozenset({"custom_llm_provider", *_CALLER_DESTINATION_KEYS}) + + +def _public_model_name_reaching(llm_router: "Router", deployment: Deployment, team_id: str | None) -> str | None: + public_model_name: Final = deployment.model_info.team_public_model_name or deployment.model_name + reachable_ids: Final = frozenset( + (as_str_mapping(row.get("model_info")) or {}).get("id") + for row in llm_router.registry().deployments( + public_model_name, source="for_request", request={"metadata": {"user_api_key_team_id": team_id}} + ) + ) + return public_model_name if deployment.model_info.id in reachable_ids else None async def _update_request_data_with_managed_file_id( @@ -135,6 +230,7 @@ async def _update_request_data_with_managed_file_id( user_api_key_dict: UserAPIKeyAuth, managed_files_obj: object | None, llm_router: Optional["Router"] = None, + managed_vector_store: LiteLLM_ManagedVectorStore | None = None, ) -> tuple[dict, str | None]: """ Update request data with model routing information from managed file ID. @@ -184,7 +280,7 @@ async def _update_request_data_with_managed_file_id( parsed_id: Final = parse_unified_id(file_id) if parsed_id: - target_model_names: Final = parsed_id.get("target_model_names", []) + target_model_names: Final = extract_target_model_names_from_unified_id(decoded_id) # Extract the actual provider file ID from llm_output_file_id field # Format: litellm_proxy:...;llm_output_file_id,{actual_file_id};... @@ -202,9 +298,73 @@ async def _update_request_data_with_managed_file_id( llm_output_file_id, ) - # Set the model for routing - if target_model_names and len(target_model_names) > 0: - routing_model: Final = target_model_names[0] + stored_deployment_match: Final = re.search(r";llm_output_file_model_id,([^;]+)", decoded_id) + stored_deployment_id: Final = stored_deployment_match.group(1).strip() if stored_deployment_match else None + stored_deployment: Final = ( + llm_router.registry().credentials(stored_deployment_id, form="deployment") + if stored_deployment_id and llm_router + else None + ) + raw_vector_store_id: Final = _raw_vector_store_id(data) + upload_deployment_id: Final = ( + stored_deployment_id + if raw_vector_store_id is not None + and stored_deployment is not None + and stored_deployment.model_info.team_id in (None, user_api_key_dict.team_id) + and await caller_may_use_model( + stored_deployment.model_info.team_public_model_name or stored_deployment.model_name, + llm_router, + user_api_key_dict, + ) + else None + ) + upload_routing_model: Final = upload_deployment_id or ( + target_model_names[0] if target_model_names else None + ) + team_copy: Final = ( + None + if raw_vector_store_id is None + or carries_explicit_credentials(data) + or _store_carries_credentials(managed_vector_store) + else await _caller_team_copy_of_managed_file( + await _resolve_provider(data=data, request=request), + file_id, + llm_output_file_id, + ( + llm_router.registry().credentials(upload_routing_model, form="with_provider") + if llm_router and upload_routing_model + else None + ), + raw_vector_store_id, + user_api_key_dict, + managed_files_obj, + llm_router, + ) + ) + routing_deployment_id: Final = team_copy.deployment_id if team_copy is not None else upload_deployment_id + routing_deployment: Final = ( + llm_router.registry().credentials(routing_deployment_id, form="deployment") + if llm_router and routing_deployment_id + else None + ) + provider_file_id: Final = team_copy.provider_file_id if team_copy is not None else llm_output_file_id + if llm_router and routing_deployment is not None: + verbose_logger.info( + "Routing vector store file operation to deployment: %s, file_id: %s -> %s", + routing_deployment_id, + file_id, + provider_file_id, + ) + dropped_keys: Final = _keys_a_pinned_call_drops(data) + return { + **{key: value for key, value in data.items() if key not in dropped_keys}, + "model": _public_model_name_reaching(llm_router, routing_deployment, user_api_key_dict.team_id) + or routing_deployment_id, + "file_id": provider_file_id, + PINNED_DEPLOYMENT_KEY: routing_deployment_id, + }, file_id + routing_model: Final = team_copy.deployment_id if team_copy is not None else upload_routing_model + if routing_model: data["model"] = routing_model # Get credentials for the model @@ -214,13 +374,13 @@ async def _update_request_data_with_managed_file_id( prepare_data_with_credentials( data=data, credentials=credentials, - file_id=llm_output_file_id, # Use the actual provider file ID + file_id=provider_file_id, ) verbose_logger.info( "Routing vector store file operation to model: %s, file_id: %s -> %s", routing_model, file_id, - llm_output_file_id, + provider_file_id, ) return data, file_id # Return original managed file ID @@ -450,7 +610,7 @@ def _update_request_data_with_litellm_managed_vector_store_registry( if parsed_id: model_id: Final = parsed_id.get("model_id") provider_resource_id: Final = parsed_id.get("provider_resource_id") - target_model_names: Final = parsed_id.get("target_model_names", []) + target_model_names: Final = extract_target_model_names_from_unified_id(decoded_id) verbose_logger.debug( "Decoded vector store - model_id: %s, provider_resource_id: %s, target_model_names: %s", @@ -593,6 +753,7 @@ async def vector_store_file_create( user_api_key_dict=user_api_key_dict, managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), llm_router=llm_router, + managed_vector_store=managed_vector_store, ) # Then handle managed vector store IDs @@ -807,6 +968,7 @@ async def vector_store_file_retrieve( user_api_key_dict=user_api_key_dict, managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), llm_router=llm_router, + managed_vector_store=managed_vector_store, ) # Then handle managed vector store IDs @@ -914,6 +1076,7 @@ async def vector_store_file_content( user_api_key_dict=user_api_key_dict, managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), llm_router=llm_router, + managed_vector_store=managed_vector_store, ) # Then handle managed vector store IDs @@ -1024,6 +1187,7 @@ async def vector_store_file_update( user_api_key_dict=user_api_key_dict, managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), llm_router=llm_router, + managed_vector_store=managed_vector_store, ) # Then handle managed vector store IDs @@ -1131,6 +1295,7 @@ async def vector_store_file_delete( user_api_key_dict=user_api_key_dict, managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), llm_router=llm_router, + managed_vector_store=managed_vector_store, ) # Then handle managed vector store IDs diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index 33994133ac0..66686beb982 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -1,4 +1,4 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Final, TypeVar from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -114,7 +114,7 @@ class RAGQuery: @staticmethod def add_search_results_to_response( response: ModelResponse, - search_results: VectorStoreSearchResponse, + search_results: Mapping[str, object], rerank_results: object = None, ) -> ModelResponse: """ diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 69f8f70f793..8ff0798cafa 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -104,6 +104,7 @@ from litellm.router_rust.invoke.endpoints import ( generic_plan, log_refused_call, log_refused_call_sync, + pinned_vector_store_file_deployment, pinned_video_deployment, plan_for, reports_own_success, @@ -2701,6 +2702,7 @@ class RustCore( fields: dict[str, object] = {"model": model} # mutable-ok: built up below, later keys win # rebind-ok: same fields.update(self._routing_fields(prepared, plan.pick_input(prepared)[0])) fields.update(pinned_video_deployment(plan, request)) + fields.update(pinned_vector_store_file_deployment(plan, request)) fields.update(unrouted_fields(plan)) fields.update(failover_fields(plan)) fields.update(scoped_fallback_fields(plan, request)) diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index 0bd522ba06d..d728583cb31 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -193,10 +193,13 @@ def _user_message(content: object) -> list[dict[str, object]]: # mutable-ok: a return [{"role": "user", "content": content}] +PINNED_DEPLOYMENT_KEY: Final = "_pinned_deployment_id" + # Request keys the router reads to route the call; the deployment call never receives them. ROUTING_KEYS: Final = frozenset( { "model", + PINNED_DEPLOYMENT_KEY, "num_retries", "disable_fallbacks", "specific_deployment", @@ -1386,6 +1389,13 @@ def pinned_video_deployment(plan: Plan, request: Mapping[str, object]) -> Mappin return {"pinned_deployment_id": deployment_id} if deployment_id else {} +def pinned_vector_store_file_deployment(plan: Plan, request: Mapping[str, object]) -> Mapping[str, object]: + deployment_id: Final = request.get(PINNED_DEPLOYMENT_KEY) + if not plan.endpoint.startswith("avector_store_file_") or not isinstance(deployment_id, str) or not deployment_id: + return {} + return {"pinned_deployment_id": deployment_id} + + async def _await(value: object) -> object: """``await value``, raising the TypeError Python raises for a value that is not awaitable.""" return await value # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType] # deliberately awaited as the Python router does @@ -2493,6 +2503,7 @@ if TYPE_CHECKING: __all__ = [ + "PINNED_DEPLOYMENT_KEY", "PLANS", "ROUTING_KEYS", "AttemptCall", @@ -2509,6 +2520,7 @@ __all__ = [ "generic_plan", "log_refused_call", "log_refused_call_sync", + "pinned_vector_store_file_deployment", "pinned_video_deployment", "plan_for", "prepare_request", diff --git a/litellm/utils.py b/litellm/utils.py index 5a1765692ae..7f7c6ffbc4d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -809,7 +809,7 @@ def _named_credential_values(kwargs: Mapping[str, object]) -> Mapping[str, objec return MappingProxyType({}) return MappingProxyType( { - key: value + key: copy.deepcopy(value) for key, value in credential_api_key_secret_reference(credential.credential_values).items() if key not in kwargs } diff --git a/tests/integration/authorization/test_rag_team_provider_credentials.py b/tests/integration/authorization/test_rag_team_provider_credentials.py index 61bab74ecaf..0c1f47cd2a3 100644 --- a/tests/integration/authorization/test_rag_team_provider_credentials.py +++ b/tests/integration/authorization/test_rag_team_provider_credentials.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import base64 import json import re import threading @@ -811,6 +812,27 @@ def _ingest_signatures(prefix: str, api_key: str, store_id: str) -> tuple[Signat ) +def _ingested_file_id(response: httpx.Response) -> str: + return string_value(JSON_OBJECT.validate_json(response.content)["file_id"]) + + +def _provider_file_id(response: httpx.Response) -> str: + managed_file_id: Final = _ingested_file_id(response) + decoded: Final = base64.urlsafe_b64decode(managed_file_id + "=" * (-len(managed_file_id) % 4)).decode() + provider_file_id: Final = re.search(r";llm_output_file_id,([^;]+)", decoded) + assert provider_file_id is not None, decoded + return provider_file_id[1] + + +def _registered_ingest_signatures( + prefix: str, api_key: str, store_id: str, response: httpx.Response +) -> tuple[Signature, ...]: + return ( + *_ingest_signatures(prefix, api_key, store_id), + (prefix, "GET", f"/v1/files/{_provider_file_id(response)}", f"Bearer {api_key}"), + ) + + def _search_signature(prefix: str, api_key: str, store_id: str) -> Signature: return (prefix, "POST", f"/v1/vector_stores/{store_id}/search", f"Bearer {api_key}") @@ -842,8 +864,9 @@ def _closed_url() -> str: def test_viewer_ingest_and_team_and_jwt_query_use_team_provider_credentials(rag_rig: RagRig) -> None: tenants: Final = rag_rig.tenants - store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, tenants.viewer_key)) - assert rag_rig.signatures() == _ingest_signatures("team", TEAM_API_KEY, store_id) + ingested: Final = _ingest(rag_rig.gateway, tenants.viewer_key) + store_id: Final = _completed_store(rag_rig, ingested) + assert rag_rig.signatures() == _registered_ingest_signatures("team", TEAM_API_KEY, store_id, ingested) team_query: Final = _query( rag_rig.gateway, tenants.team_a_key, tenants.team_a_model, {**OPENAI, "vector_store_id": store_id} @@ -898,10 +921,9 @@ def test_unreachable_credential_api_base_fails_ingest_without_saving_a_store(rag def test_rag_alias_ingest_and_query_use_the_team_wildcard_deployment(rag_rig: RagRig) -> None: - store_id: Final = _completed_store( - rag_rig, _ingest(rag_rig.gateway, rag_rig.tenants.team_w_key, path="/rag/ingest") - ) - assert rag_rig.signatures() == _ingest_signatures("wildcard", WILDCARD_API_KEY, store_id) + ingested: Final = _ingest(rag_rig.gateway, rag_rig.tenants.team_w_key, path="/rag/ingest") + store_id: Final = _completed_store(rag_rig, ingested) + assert rag_rig.signatures() == _registered_ingest_signatures("wildcard", WILDCARD_API_KEY, store_id, ingested) rows: Final = _store_rows(store_id) assert len(rows) == 1, rows assert WILDCARD_API_KEY not in json.dumps(rows), rows @@ -924,9 +946,10 @@ def test_rag_alias_ingest_and_query_use_the_team_wildcard_deployment(rag_rig: Ra def test_jwt_team_member_ingest_uses_team_provider_credentials(rag_rig: RagRig) -> None: token: Final = rag_rig.jwt(rag_rig.tenants.jwt_member, rag_rig.tenants.team_a) - store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, token)) + ingested: Final = _ingest(rag_rig.gateway, token) + store_id: Final = _completed_store(rag_rig, ingested) - assert rag_rig.signatures() == _ingest_signatures("team", TEAM_API_KEY, store_id) + assert rag_rig.signatures() == _registered_ingest_signatures("team", TEAM_API_KEY, store_id, ingested) rows: Final = _store_rows(store_id) assert [row["team_id"] for row in rows] == [rag_rig.tenants.team_a], rows @@ -963,10 +986,9 @@ def test_caller_api_key_or_api_base_on_ingest_is_rejected_before_any_upstream_ca def test_blank_caller_credential_name_falls_back_to_the_team_deployment(rag_rig: RagRig, blank: str | None) -> None: tenants: Final = rag_rig.tenants - store_id: Final = _completed_store( - rag_rig, _ingest(rag_rig.gateway, tenants.team_a_key, {**OPENAI, "litellm_credential_name": blank}) - ) - assert rag_rig.signatures() == _ingest_signatures("team", TEAM_API_KEY, store_id) + ingested: Final = _ingest(rag_rig.gateway, tenants.team_a_key, {**OPENAI, "litellm_credential_name": blank}) + store_id: Final = _completed_store(rag_rig, ingested) + assert rag_rig.signatures() == _registered_ingest_signatures("team", TEAM_API_KEY, store_id, ingested) response: Final = _query( rag_rig.gateway, tenants.team_a_key, tenants.team_a_model, {**OPENAI, "vector_store_id": store_id} @@ -1015,8 +1037,10 @@ def test_key_restricted_to_a_model_under_the_team_wildcard_uses_the_team_deploym ) -> None: key: Final = getattr(rag_rig.tenants, caller) - store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, key)) + ingested: Final = _ingest(rag_rig.gateway, key) + store_id: Final = _completed_store(rag_rig, ingested) assert rag_rig.signatures() == _ingest_signatures("wildcard", WILDCARD_API_KEY, store_id) + assert _ingested_file_id(ingested).startswith("file-wildcard-"), ingested.text rows: Final = _store_rows(store_id) assert [row["team_id"] for row in rows] == [rag_rig.tenants.team_w], rows assert WILDCARD_API_KEY not in json.dumps(rows), rows @@ -1033,8 +1057,10 @@ def test_key_restricted_to_a_model_under_the_team_wildcard_uses_the_team_deploym def test_key_restricted_to_a_global_model_under_the_team_wildcard_splits_rag_and_chat(rag_rig: RagRig) -> None: tenants: Final = rag_rig.tenants - store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, tenants.wildcard_global_key)) + ingested: Final = _ingest(rag_rig.gateway, tenants.wildcard_global_key) + store_id: Final = _completed_store(rag_rig, ingested) assert rag_rig.signatures() == _ingest_signatures("wildcard", WILDCARD_API_KEY, store_id) + assert _ingested_file_id(ingested).startswith("file-wildcard-"), ingested.text response: Final = _query( rag_rig.gateway, tenants.wildcard_global_key, tenants.global_model, {**OPENAI, "vector_store_id": store_id} @@ -1252,8 +1278,9 @@ def test_azure_query_never_receives_the_team_openai_key(rag_rig: RagRig) -> None def test_another_team_cannot_query_a_team_store(rag_rig: RagRig) -> None: tenants: Final = rag_rig.tenants - store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, tenants.team_a_key)) - assert rag_rig.signatures() == _ingest_signatures("team", TEAM_API_KEY, store_id) + ingested: Final = _ingest(rag_rig.gateway, tenants.team_a_key) + store_id: Final = _completed_store(rag_rig, ingested) + assert rag_rig.signatures() == _registered_ingest_signatures("team", TEAM_API_KEY, store_id, ingested) response: Final = _query( rag_rig.gateway, tenants.team_w_key, WILDCARD_MODEL, {**OPENAI, "vector_store_id": store_id} @@ -1380,8 +1407,12 @@ def _burst_names() -> tuple[str, ...]: return tuple(f"rag-burst-{uuid.uuid4().hex}" for _ in range(10)) -def _all_ingest_signatures(store_ids: Iterable[str]) -> list[Signature]: - return sorted(chain.from_iterable(_ingest_signatures("team", TEAM_API_KEY, store_id) for store_id in store_ids)) +def _all_ingest_signatures(ingests: Iterable[tuple[str, httpx.Response]]) -> list[Signature]: + return sorted( + chain.from_iterable( + _registered_ingest_signatures("team", TEAM_API_KEY, store_id, response) for store_id, response in ingests + ) + ) async def test_ingest_burst_during_a_provider_outage_fails_each_request_then_recovers(rag_rig: RagRig) -> None: @@ -1401,7 +1432,7 @@ async def test_ingest_burst_during_a_provider_outage_fails_each_request_then_rec store_ids: Final = tuple(_completed_store(rag_rig, response) for response in recovered) assert len(frozenset(store_ids)) == 10, store_ids - assert sorted(rag_rig.signatures()) == _all_ingest_signatures(store_ids) + assert sorted(rag_rig.signatures()) == _all_ingest_signatures(zip(store_ids, recovered, strict=True)) assert [[row["team_id"] for row in _store_rows(store_id)] for store_id in store_ids] == [ [rag_rig.tenants.team_a] ] * 10 @@ -1502,19 +1533,19 @@ async def test_worker_sigkill_mid_query_burst_leaves_the_sibling_on_team_credent def _ingest_and_drain( gateway: Gateway, scenario: Scenario, key: str, upstream: Wire -) -> tuple[str, tuple[Signature, ...]]: +) -> tuple[str, httpx.Response, tuple[Signature, ...]]: upstream.drain() response: Final = _ingest(gateway, key) assert response.status_code == 200, response.text store_id: Final = string_value(object_value(response.json())["vector_store_id"]) scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": store_id}) calls: Final = tuple(_call(request) for request in upstream.drain()) - return store_id, tuple(call.signature for call in calls if call.route != "/v1/models") + return store_id, response, tuple(call.signature for call in calls if call.route != "/v1/models") def _ingests_after_editing( gateway: Gateway, scenario: Scenario, key: str, upstream: Wire, model_id: str -) -> tuple[tuple[str, tuple[Signature, ...]], ...]: +) -> tuple[tuple[str, httpx.Response, tuple[Signature, ...]], ...]: edited: Final = gateway.request("PATCH", f"/model/{model_id}/update", {"litellm_params": {"timeout": 30}}) assert edited.status_code == 200, edited.text return tuple(_ingest_and_drain(gateway, scenario, key, upstream) for _ in range(4)) @@ -1577,8 +1608,9 @@ def test_a_team_with_two_openai_deployments_keeps_ingesting_through_the_older_on *_ingests_after_editing(gateway, owned_scenario, key, upstream, other_id), ) - assert [calls for _, calls in ingests] == [ - _ingest_signatures(picked_prefix, picked_key, store_id) for store_id, _ in ingests + assert [calls for _, _, calls in ingests] == [ + _registered_ingest_signatures(picked_prefix, picked_key, store_id, response) + for store_id, response, _ in ingests ] @@ -1604,9 +1636,13 @@ def test_ingest_into_a_managed_store_uses_its_own_key_else_the_team_key( assert response.status_code == 200, response.text assert object_value(response.json())["vector_store_id"] == store_id, response.text bearer: Final = f"Bearer {api_key}" + team_key_registration: Final = ( + ((prefix, "GET", f"/v1/files/{_provider_file_id(response)}", bearer),) if stored == "null_key" else () + ) assert rag_rig.signatures() == ( (prefix, "POST", "/v1/files", bearer), (prefix, "POST", f"/v1/vector_stores/{store_id}/files", bearer), + *team_key_registration, ) diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 747c09b9108..8c8e970b80d 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -24,12 +24,18 @@ from litellm_enterprise.proxy.hooks.managed_files import ( _provider_file_retrieve_credentials, ) +from litellm import Router from litellm.caching import DualCache +from litellm.litellm_core_utils.prompt_templates.file_id_utils import get_format_from_file_id from litellm.proxy._types import CallTypes, LiteLLM_ManagedFileTable, UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import ( BATCH_CREATE_HIDDEN_PARAM, + TeamProviderDeployment, + convert_b64_uid_to_unified_uid, encode_file_id_with_model, + get_models_from_unified_file_id, is_base64_encoded_unified_file_id, + provider_file_id_from_managed_id, ) from litellm.types.llms.openai import OpenAIFileObject from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder @@ -6319,3 +6325,132 @@ async def test_only_a_streamed_background_create_is_marked_for_the_managed_bill( assert isinstance(result, dict) assert (MANAGED_BACKGROUND_RESPONSE_METADATA_KEY in result["litellm_metadata"]) is marked assert result["metadata"] == {"run": "nightly"} + + +_TEAM_DEPLOYMENT_API_BASE: Final = "https://team-openai.example/v1" +_TEAM_DEPLOYMENT: Final = TeamProviderDeployment( + deployment_id="team-a-openai", public_model_name="openai/*", credentials={} +) + + +def _team_openai_wildcard_router() -> Router: + return Router( + model_list=[ + { + "model_name": "model_name_team-a_openai", + "litellm_params": { + "model": "openai/*", + "api_key": "team-openai-key", + "api_base": _TEAM_DEPLOYMENT_API_BASE, + }, + "model_info": {"id": "team-a-openai", "team_id": "team-a", "team_public_model_name": "openai/*"}, + } + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_respx_sees_provider_calls") +@respx.mock +async def test_register_provider_file_stores_a_managed_file_its_owner_can_call() -> None: + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma() + provider_retrieve: Final = respx.get(f"{_TEAM_DEPLOYMENT_API_BASE}/files/file-ingested").respond( + 200, + json={ + "id": "file-ingested", + "object": "file", + "bytes": 12, + "created_at": 456, + "filename": "sample.txt", + "purpose": "assistants", + "status": "processed", + }, + ) + owner: Final = UserAPIKeyAuth(user_id="viewer", team_id="team-a") + + managed_file_id: Final = await proxy_managed_files.register_provider_file( + provider_file_id="file-ingested", + content_type="text/plain", + deployment=_TEAM_DEPLOYMENT, + owner=owner, + llm_router=_team_openai_wildcard_router(), + ) + + assert managed_file_id is not None + assert is_base64_encoded_unified_file_id(managed_file_id) + assert get_format_from_file_id(managed_file_id) == "text/plain" + assert get_models_from_unified_file_id(convert_b64_uid_to_unified_uid(managed_file_id)) == ["openai/*"] + assert provider_file_id_from_managed_id(managed_file_id) == "file-ingested" + assert provider_retrieve.calls.last.request.headers["authorization"] == "Bearer team-openai-key" + stored: Final = managed_file_table.rows[managed_file_id] + assert stored.model_mappings == {"team-a-openai": "file-ingested"} + assert stored.file_object is not None + assert stored.file_object.id == managed_file_id + assert stored.file_object.filename == "sample.txt" + assert await proxy_managed_files.can_user_call_unified_file_id(managed_file_id, owner) is True + assert ( + await proxy_managed_files.can_user_call_unified_file_id( + managed_file_id, UserAPIKeyAuth(user_id="outsider", team_id="team-b") + ) + is False + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_respx_sees_provider_calls") +@respx.mock +async def test_register_provider_file_returns_none_when_the_provider_cannot_find_the_file() -> None: + proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma() + respx.get(f"{_TEAM_DEPLOYMENT_API_BASE}/files/file-ingested").respond( + 404, json={"error": {"message": "No such File object: file-ingested", "type": "invalid_request_error"}} + ) + + managed_file_id: Final = await proxy_managed_files.register_provider_file( + provider_file_id="file-ingested", + content_type="text/plain", + deployment=_TEAM_DEPLOYMENT, + owner=UserAPIKeyAuth(user_id="viewer", team_id="team-a"), + llm_router=_team_openai_wildcard_router(), + ) + + assert managed_file_id is None + assert managed_file_table.rows == {} + + +class _UnwritableManagedFileTable(_InMemoryManagedFileTable): + async def upsert( + self, + where: Mapping[str, object], + data: Mapping[str, Mapping[str, object]], + ) -> LiteLLM_ManagedFileTable: + raise RuntimeError("database unavailable") + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_respx_sees_provider_calls") +@respx.mock +async def test_register_provider_file_returns_none_when_the_managed_row_cannot_be_written() -> None: + proxy_managed_files: Final = PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=_InMemoryManagedFilesPrismaClient(_UnwritableManagedFileTable()) + ) + respx.get(f"{_TEAM_DEPLOYMENT_API_BASE}/files/file-ingested").respond( + 200, + json={ + "id": "file-ingested", + "object": "file", + "bytes": 12, + "created_at": 456, + "filename": "sample.txt", + "purpose": "assistants", + }, + ) + + managed_file_id: Final = await proxy_managed_files.register_provider_file( + provider_file_id="file-ingested", + content_type="text/plain", + deployment=_TEAM_DEPLOYMENT, + owner=UserAPIKeyAuth(user_id="viewer", team_id="team-a"), + llm_router=_team_openai_wildcard_router(), + ) + + assert managed_file_id is None diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py index b3323de3de5..bbe7f4383b6 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py @@ -18,6 +18,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( apply_unified_file_ids, get_credentials_for_model, get_team_provider_credentials, + get_team_provider_deployment, is_litellm_executed_batch, map_raw_file_ids_to_unified, ) @@ -770,6 +771,61 @@ def test_a_team_keeps_its_older_openai_deployment_when_a_newer_one_has_a_lower_i assert credentials["api_key"] == "older-key" +@pytest.mark.parametrize("newer_first", [True, False]) +def test_the_team_provider_deployment_is_the_one_whose_credentials_the_team_resolves(newer_first: bool): + older: Final = _team_openai_deployment("deployment-z", "older-key", datetime(2026, 1, 1, tzinfo=timezone.utc)) + newer: Final = _team_openai_deployment("deployment-a", "newer-key", datetime(2026, 6, 1, tzinfo=timezone.utc)) + router: Final = Router(model_list=[newer, older] if newer_first else [older, newer]) + team_member: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="team-a") + + deployment: Final = get_team_provider_deployment( + llm_router=router, user_api_key_dict=team_member, custom_llm_provider="openai" + ) + + assert deployment is not None + assert (deployment.deployment_id, deployment.public_model_name) == ("deployment-z", "team-gpt-deployment-z") + assert dict(deployment.credentials) == get_team_provider_credentials( + llm_router=router, + user_api_key_dict=team_member, + custom_llm_provider="openai", + include_accessible_deployments=False, + ) + assert deployment.credentials["api_key"] == "older-key" + + +def test_the_team_provider_deployment_carries_the_team_public_model_name(): + deployment: Final = get_team_provider_deployment( + llm_router=_team_openai_router("openai/*", "openai/*"), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="team-a"), + custom_llm_provider="openai", + ) + + assert deployment is not None + assert (deployment.deployment_id, deployment.public_model_name) == ("team-a-openai", "openai/*") + + +def test_a_key_with_only_an_accessible_non_team_deployment_has_no_team_provider_deployment(): + router: Final = Router( + model_list=[ + { + "model_name": "openai/gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "openai-deployment-key"}, + } + ] + ) + caller: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", models=["openai/*"]) + + accessible_credentials: Final = get_team_provider_credentials( + llm_router=router, user_api_key_dict=caller, custom_llm_provider="openai", include_accessible_deployments=True + ) + + assert accessible_credentials is not None + assert accessible_credentials["api_key"] == "openai-deployment-key" + assert ( + get_team_provider_deployment(llm_router=router, user_api_key_dict=caller, custom_llm_provider="openai") is None + ) + + def test_a_team_keeps_its_first_config_file_openai_deployment_over_older_db_ones(): db_deployment: Final = _team_openai_deployment("deployment-a", "db-key", datetime(2020, 1, 1, tzinfo=timezone.utc)) config_deployments: Final = tuple( diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index bf4b5d5d77a..6eff950dddf 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -2,9 +2,12 @@ Tests for RAG proxy endpoints. """ +import base64 import io import json from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -17,7 +20,12 @@ import litellm from litellm import Router from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.openai_files_endpoints.common_utils import ( + TeamProviderDeployment, + detect_content_type_from_filename, +) from litellm.proxy.proxy_server import app +from litellm.types.utils import SpecialEnums from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.vector_stores.vector_store_registry import VectorStoreRegistry @@ -637,6 +645,474 @@ def test_rag_ingest_blank_caller_credential_name_uses_team_provider_credentials( assert forwarded["api_base"] == _TEAM_OPENAI_BASE +@dataclass(frozen=True, slots=True) +class _ProviderFileRegistration: + provider_file_id: str + content_type: str + deployment: TeamProviderDeployment + owner: UserAPIKeyAuth + + +class _ManagedFilesHookStub: + def __init__( + self, + *, + managed_file_id: str | None, + can_access: bool = True, + model_mappings: Mapping[str, Mapping[str, str]] = MappingProxyType({}), + unified_file_ids: Mapping[str, str] = MappingProxyType({}), + ) -> None: + self.managed_file_id: Final = managed_file_id + self.can_access: Final = can_access + self.model_mappings: Final = model_mappings + self.unified_file_ids: Final = unified_file_ids + self.registrations: tuple[_ProviderFileRegistration, ...] = () + + async def register_provider_file( + self, + *, + provider_file_id: str, + content_type: str, + deployment: TeamProviderDeployment, + owner: UserAPIKeyAuth, + llm_router: Router | None, + ) -> str | None: + self.registrations = ( + *self.registrations, + _ProviderFileRegistration(provider_file_id, content_type, deployment, owner), + ) + return self.managed_file_id + + async def can_user_call_unified_file_id(self, unified_file_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + return self.can_access + + async def can_user_call_unified_object_id(self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + return self.can_access + + async def get_model_file_id_mapping( + self, file_ids: Sequence[str], litellm_parent_otel_span: object + ) -> Mapping[str, Mapping[str, str]]: + return {file_id: self.model_mappings[file_id] for file_id in file_ids if file_id in self.model_mappings} + + async def get_unified_file_ids_for_provider_file_ids( + self, provider_file_ids: Sequence[str], user_api_key_dict: UserAPIKeyAuth + ) -> Mapping[str, str]: + return {file_id: self.unified_file_ids[file_id] for file_id in provider_file_ids if file_id in self.unified_file_ids} + + +def _managed_file_id(provider_file_id: str) -> str: + unified_file_id: Final = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "unified-1", "openai/team-provider", provider_file_id, "team-deployment" + ) + return base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") + + +def _installed_managed_files_hook(hook: _ManagedFilesHookStub | None): + from litellm.proxy import proxy_server + + proxy_hook_mapping: Final = proxy_server.proxy_logging_obj.proxy_hook_mapping + hooks_without_managed_files: Final = { + name: installed for name, installed in proxy_hook_mapping.items() if name != "managed_files" + } + return patch.dict( # test-quality-ok: the proxy reads its hooks from this process-wide mapping, no injection seam + proxy_hook_mapping, + {**hooks_without_managed_files, **({} if hook is None else {"managed_files": hook})}, + clear=True, + ) + + +def _prisma_client_saving_a_new_store() -> tuple[MagicMock, MagicMock]: + prisma_client: Final = MagicMock() + managed_vector_store_table: Final = prisma_client.db.litellm_managedvectorstorestable + managed_vector_store_table.find_unique = AsyncMock(return_value=None) + created_row: Final = MagicMock() + created_row.model_dump.return_value = { + "vector_store_id": "vs_new", + "custom_llm_provider": "openai", + "litellm_params": "{}", + "vector_store_metadata": "{}", + } + managed_vector_store_table.create = AsyncMock(return_value=created_row) + return prisma_client, managed_vector_store_table + + +def _router_without_team_deployment() -> Router: + return Router(model_list=[_completion_deployment()]) + + +@pytest.mark.parametrize("ingest_request", [_ingest_form, _ingest_json]) +def test_rag_ingest_returns_a_managed_file_id_for_an_upload_through_the_team_deployment(client_team_a, ingest_request): + router: Final = _team_provider_router() + team_deployment_id: Final = router.registry().all()[0]["model_info"]["id"] + managed_file_id: Final = _managed_file_id("file_123") + hook: Final = _ManagedFilesHookStub(managed_file_id=managed_file_id) + prisma_client, managed_vector_store_table = _prisma_client_saving_a_new_store() + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) + + with ( + aingest_patch, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", router), + _patched_prisma_client(prisma_client), + _installed_managed_files_hook(hook), + ): + response: Final = client_team_a.post("/v1/rag/ingest", **ingest_request({"custom_llm_provider": "openai"})) + + assert response.status_code == 200, response.json() + assert response.json()["file_id"] == managed_file_id + (registration,) = hook.registrations + assert registration.provider_file_id == "file_123" + assert registration.content_type == "text/plain" + assert registration.deployment.deployment_id == team_deployment_id + assert registration.deployment.public_model_name == "openai/team-provider" + assert registration.deployment.credentials["api_key"] == _TEAM_PROVIDER_KEY + assert registration.owner.team_id == _TEAM_ID + saved_data: Final = managed_vector_store_table.create.await_args.kwargs["data"] + saved_metadata: Final = json.loads(saved_data["vector_store_metadata"]) + assert [entry["file_id"] for entry in saved_metadata["ingested_files"]] == [managed_file_id] + + +def test_rag_ingest_registers_a_url_upload_with_the_type_its_path_names(client_team_a): + hook: Final = _ManagedFilesHookStub(managed_file_id=_managed_file_id("file_123")) + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) + + with ( + aingest_patch, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + _installed_managed_files_hook(hook), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + **_ingest_json( + {"custom_llm_provider": "openai"}, file_url="https://example.com/files/handbook.pdf?sig=abc#page=2" + ), + ) + + assert response.status_code == 200, response.json() + (registration,) = hook.registrations + assert registration.content_type == detect_content_type_from_filename("handbook.pdf") + + +@pytest.mark.parametrize( + ("vector_store", "registry_store", "headers", "general_settings"), + [ + pytest.param( + {"custom_llm_provider": "openai", "litellm_credential_name": "caller-openai"}, + None, + {}, + {}, + id="caller-credential-name", + ), + pytest.param( + {"vector_store_id": "registry-store-with-key"}, + _managed_openai_store("registry-store-with-key", {"api_key": _REGISTRY_PROVIDER_KEY}), + {}, + {}, + id="registered-store-key", + ), + pytest.param( + {"custom_llm_provider": "openai"}, + None, + {"Authorization": "Bearer sk-litellm-virtual-key", "x-api-key": "sk-forwarded-provider-key"}, + {"forward_llm_provider_auth_headers": True}, + id="forwarded-provider-key", + ), + pytest.param({"custom_llm_provider": "bedrock"}, None, {}, {}, id="non-openai-provider"), + ], +) +def test_rag_ingest_keeps_the_provider_file_id_when_the_upload_used_other_credentials( + client_team_a, vector_store, registry_store, headers, general_settings +): + hook: Final = _ManagedFilesHookStub(managed_file_id=_managed_file_id("file_123")) + aingest_patch, registry_patch = _patched_ingest_boundary( + registry_store, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) + + with ( + aingest_patch, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + patch( + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), + ), + patch.dict("litellm.proxy.proxy_server.general_settings", general_settings), + _patched_prisma_client(None), + _installed_managed_files_hook(hook), + ): + response: Final = client_team_a.post("/v1/rag/ingest", headers=headers, **_ingest_form(vector_store)) + + assert response.status_code == 200, response.json() + assert response.json()["file_id"] == "file_123" + assert hook.registrations == () + + +def _team_wildcard_router() -> Router: + return Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": _TEAM_PROVIDER_KEY, "api_base": _TEAM_OPENAI_BASE}, + "model_info": {"team_id": _TEAM_ID}, + } + ] + ) + + +_ONE_MODEL_TEAM_KEY: Final = UserAPIKeyAuth(user_id="test_team_user", team_id=_TEAM_ID, models=["openai/gpt-5.4-mini"]) + + +def test_rag_ingest_keeps_the_provider_file_id_for_a_key_the_team_deployment_checks_reject(client_team_a): + hook: Final = _ManagedFilesHookStub(managed_file_id=_managed_file_id("file_123")) + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) + + with ( + aingest_patch as aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_wildcard_router()), + patch.dict(app.dependency_overrides, {user_api_key_auth: lambda: _ONE_MODEL_TEAM_KEY}), + _patched_prisma_client(None), + _installed_managed_files_hook(hook), + ): + response: Final = client_team_a.post("/v1/rag/ingest", **_ingest_form({"custom_llm_provider": "openai"})) + + assert response.status_code == 200, response.json() + assert response.json()["file_id"] == "file_123" + assert hook.registrations == () + assert aingest.await_args.kwargs["ingest_options"]["vector_store"]["api_key"] == _TEAM_PROVIDER_KEY + + +@pytest.mark.parametrize( + ("router_factory", "hook"), + [ + pytest.param( + _router_without_team_deployment, + _ManagedFilesHookStub(managed_file_id=_managed_file_id("file_123")), + id="no-team-deployment", + ), + pytest.param(_team_provider_router, None, id="no-managed-files-hook"), + ], +) +def test_rag_ingest_keeps_the_provider_file_id_without_a_team_deployment_or_managed_files_hook( + client_team_a, router_factory, hook +): + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) + + with ( + aingest_patch, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", router_factory()), + _patched_prisma_client(None), + _installed_managed_files_hook(hook), + ): + response: Final = client_team_a.post("/v1/rag/ingest", **_ingest_form({"custom_llm_provider": "openai"})) + + assert response.status_code == 200, response.json() + assert response.json()["file_id"] == "file_123" + assert hook is None or hook.registrations == () + + +def test_rag_ingest_keeps_the_provider_file_id_when_managed_file_registration_fails(client_team_a): + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) + + with ( + aingest_patch, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + _patched_prisma_client(None), + _installed_managed_files_hook(_ManagedFilesHookStub(managed_file_id=None)), + ): + response: Final = client_team_a.post("/v1/rag/ingest", **_ingest_form({"custom_llm_provider": "openai"})) + + assert response.status_code == 200, response.json() + assert response.json()["file_id"] == "file_123" + + +def test_rag_ingest_with_a_managed_file_id_attaches_the_provider_file_and_returns_the_managed_id(client_team_a): + managed_file_id: Final = _managed_file_id("file-raw") + hook: Final = _ManagedFilesHookStub(managed_file_id=None) + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file-raw"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + _patched_prisma_client(None), + _installed_managed_files_hook(hook), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + json={"file_id": managed_file_id, "ingest_options": {"vector_store": {"custom_llm_provider": "openai"}}}, + ) + + assert response.status_code == 200, response.json() + assert response.json()["file_id"] == managed_file_id + assert mock_aingest.await_args.kwargs["file_id"] == "file-raw" + assert hook.registrations == () + + +def test_rag_ingest_with_a_managed_file_id_attaches_the_team_deployment_copy(client_team_a): + router: Final = _team_provider_router() + team_deployment_id: Final = router.registry().all()[0]["model_info"]["id"] + managed_file_id: Final = _managed_file_id("file-other-copy") + hook: Final = _ManagedFilesHookStub( + managed_file_id=None, + model_mappings={managed_file_id: {"team-deployment": "file-other-copy", team_deployment_id: "file-team-copy"}}, + ) + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file-team-copy"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", router), + _patched_prisma_client(None), + _installed_managed_files_hook(hook), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + json={"file_id": managed_file_id, "ingest_options": {"vector_store": {"custom_llm_provider": "openai"}}}, + ) + + assert response.status_code == 200, response.json() + assert response.json()["file_id"] == managed_file_id + assert mock_aingest.await_args.kwargs["file_id"] == "file-team-copy" + + +def test_rag_ingest_with_a_managed_file_id_attaches_the_copy_matching_the_credentials_for_a_key_limited_to_one_model( + client_team_a, +): + router: Final = _team_wildcard_router() + team_deployment_id: Final = router.registry().all()[0]["model_info"]["id"] + managed_file_id: Final = _managed_file_id("file-other-copy") + hook: Final = _ManagedFilesHookStub( + managed_file_id=None, + model_mappings={managed_file_id: {"team-deployment": "file-other-copy", team_deployment_id: "file-team-copy"}}, + ) + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file-team-copy"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", router), + patch.dict(app.dependency_overrides, {user_api_key_auth: lambda: _ONE_MODEL_TEAM_KEY}), + _patched_prisma_client(None), + _installed_managed_files_hook(hook), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + json={"file_id": managed_file_id, "ingest_options": {"vector_store": {"custom_llm_provider": "openai"}}}, + ) + + assert response.status_code == 200, response.json() + assert response.json()["file_id"] == managed_file_id + assert mock_aingest.await_args.kwargs["file_id"] == "file-team-copy" + assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"]["api_key"] == _TEAM_PROVIDER_KEY + + +def test_rag_ingest_rejects_a_managed_file_id_the_caller_cannot_access(client_team_a): + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file-raw"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + _patched_prisma_client(None), + _installed_managed_files_hook(_ManagedFilesHookStub(managed_file_id=None, can_access=False)), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + json={ + "file_id": _managed_file_id("file-raw"), + "ingest_options": {"vector_store": {"custom_llm_provider": "openai"}}, + }, + ) + + assert response.status_code == 403, response.json() + assert mock_aingest.await_count == 0 + + +def test_rag_ingest_with_an_empty_file_url_never_claims_a_raw_file_id(client_team_a): + hook: Final = _ManagedFilesHookStub(managed_file_id=_managed_file_id("file-raw")) + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file-raw"}, + ) + + with ( + aingest_patch, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + _patched_prisma_client(None), + _installed_managed_files_hook(hook), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + json={ + "file_id": "file-raw", + "file_url": "", + "ingest_options": {"vector_store": {"custom_llm_provider": "openai"}}, + }, + ) + + assert response.status_code == 200, response.json() + assert response.json()["file_id"] == "file-raw" + assert hook.registrations == () + + +def test_rag_ingest_answers_400_for_a_managed_file_id_without_the_managed_files_hook(client_team_a): + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file-raw"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + _patched_prisma_client(None), + _installed_managed_files_hook(None), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + json={ + "file_id": _managed_file_id("file-raw"), + "ingest_options": {"vector_store": {"custom_llm_provider": "openai"}}, + }, + ) + + assert response.status_code == 400, response.json() + assert mock_aingest.await_count == 0 + + def test_rag_ingest_failed_upstream_status_is_returned_without_persisting_store(client_internal_user): prisma_client: Final = MagicMock() managed_vector_store_table: Final = prisma_client.db.litellm_managedvectorstorestable @@ -697,6 +1173,41 @@ def test_rag_query_forwards_team_provider_credentials_to_unmanaged_search( assert search.calls.last.request.headers["authorization"] == f"Bearer {_TEAM_PROVIDER_KEY}" +@pytest.mark.usefixtures("httpx_transport") +def test_rag_query_search_results_carry_the_managed_file_id_ingest_returned( + client_team_a: TestClient, respx_mock: respx.MockRouter +) -> None: + managed_file_id: Final = _managed_file_id("file-ingested") + respx_mock.post(f"{_TEAM_OPENAI_BASE}/vector_stores/provider-native-store/search").respond( + json={ + **_OPENAI_SEARCH_PAGE, + "data": [ + { + "file_id": "file-ingested", + "filename": "handbook.txt", + "score": 0.9, + "attributes": {}, + "content": [{"type": "text", "text": "hello"}], + } + ], + } + ) + + with _installed_managed_files_hook( + _ManagedFilesHookStub(managed_file_id=None, unified_file_ids={"file-ingested": managed_file_id}) + ): + response: Final = _post_rag_query( + client_team_a, + retrieval_config={"vector_store_id": "provider-native-store", "custom_llm_provider": "openai"}, + managed_stores=(), + router=_team_provider_router(), + ) + + assert response.status_code == 200, response.json() + search_results: Final = response.json()["choices"][0]["message"]["provider_specific_fields"]["search_results"] + assert [hit["file_id"] for hit in search_results["data"]] == [managed_file_id] + + @pytest.mark.usefixtures("httpx_transport") def test_rag_query_managed_store_credentials_override_team_provider_credentials( client_team_a: TestClient, respx_mock: respx.MockRouter diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 7f384d8e6bf..75a17d49d3a 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -20,9 +20,11 @@ from fastapi import HTTPException import litellm from litellm.caching.caching import DualCache +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.openai_files_endpoints.common_utils import get_team_provider_deployment from litellm.proxy.route_llm_request import ( JSON_OBJECT_ADAPTER, ProxyModelNotFoundError, @@ -2436,19 +2438,44 @@ def _registered_bedrock_kb_search(team_id: str) -> dict[str, object]: } -def _team_a_router(*team_models: tuple[str, dict[str, object]]) -> litellm.Router: +def _team_a_deployment(public_name: str, litellm_params: Mapping[str, object]) -> dict[str, object]: + return { + "model_name": f"internal-team-a-{public_name}", + "litellm_params": dict(litellm_params), + "model_info": { + "id": f"team-a-{public_name}", + "team_id": "team-a", + "team_public_model_name": public_name, + }, + } + + +def _team_a_router(*team_models: tuple[str, Mapping[str, object]]) -> litellm.Router: + return litellm.Router( + model_list=[_team_a_deployment(public_name, litellm_params) for public_name, litellm_params in team_models] + ) + + +_PROXY_OPENAI_BASE: Final = "https://proxy-openai.example/v1" + + +def _team_a_openai_router_beside_a_proxy_openai_wildcard() -> litellm.Router: return litellm.Router( model_list=[ { - "model_name": f"internal-team-a-{public_name}", - "litellm_params": litellm_params, + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "proxy-openai-key", "api_base": _PROXY_OPENAI_BASE}, + }, + { + "model_name": "model_name_team-a_6f0c8a52", + "litellm_params": dict(_TEAM_OPENAI_PARAMS), "model_info": { - "id": f"team-a-{public_name}", + "id": "team-a-openai", "team_id": "team-a", - "team_public_model_name": public_name, + "team_public_model_name": "openai/*", + "db_model": True, }, - } - for public_name, litellm_params in team_models + }, ] ) @@ -2511,6 +2538,414 @@ async def test_team_key_search_of_registered_bedrock_store_reaches_bedrock_with_ assert "Credential=AKIAREGISTRYKEY00000/" in signed_auth, signed_auth +_TEAM_OPENAI_BASE: Final = "https://team-openai.example/v1" +_TEAM_OPENAI_PARAMS: Final[Mapping[str, object]] = { + "model": "openai/*", + "api_key": "team-openai-key", + "api_base": _TEAM_OPENAI_BASE, +} +_VECTOR_STORE_FILE: Final[Mapping[str, object]] = { + "id": "file-ingested", + "object": "vector_store.file", + "vector_store_id": "vs_team", + "status": "completed", + "created_at": 1, + "usage_bytes": 1, + "last_error": None, +} +_VECTOR_STORE_FILE_CALLS: Final[Mapping[str, tuple[str, str, Mapping[str, object]]]] = { + "avector_store_file_list": ( + "GET", + "/vector_stores/vs_team/files", + {"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False}, + ), + "avector_store_file_retrieve": ("GET", "/vector_stores/vs_team/files/file-ingested", _VECTOR_STORE_FILE), + "avector_store_file_content": ( + "GET", + "/vector_stores/vs_team/files/file-ingested/content", + {"object": "vector_store.file_content.page", "data": [], "has_more": False, "next_page": None}, + ), + "avector_store_file_delete": ( + "DELETE", + "/vector_stores/vs_team/files/file-ingested", + {"id": "file-ingested", "object": "vector_store.file.deleted", "deleted": True}, + ), + "avector_store_file_create": ("POST", "/vector_stores/vs_team/files", _VECTOR_STORE_FILE), + "avector_store_file_update": ("POST", "/vector_stores/vs_team/files/file-ingested", _VECTOR_STORE_FILE), +} +_TEAM_A_MEMBER: Final = UserAPIKeyAuth(api_key="team-a-member-key", team_id="team-a") + + +def _model_less_vector_store_file_request( + route_type: str, litellm_credential_name: str | None = None +) -> Mapping[str, object]: + return { + "model": None, + "vector_store_id": "vs_team", + **({} if route_type == "avector_store_file_list" else {"file_id": "file-ingested"}), + **({"attributes": {"source": "rag-ingest"}} if route_type == "avector_store_file_update" else {}), + "custom_llm_provider": "openai", + "metadata": {"user_api_key_team_id": "team-a"}, + "litellm_credential_name": litellm_credential_name, + } + + +def _team_a_openai_db_deployment(deployment_id: str, api_base: str, created_at: datetime) -> Mapping[str, object]: + return { + "model_name": f"internal-team-a-{deployment_id}", + "litellm_params": {"model": "openai/*", "api_key": f"{deployment_id}-key", "api_base": api_base}, + "model_info": { + "id": deployment_id, + "team_id": "team-a", + "team_public_model_name": f"openai-{deployment_id}", + "db_model": True, + "created_at": created_at, + }, + } + + +@pytest.mark.parametrize("route_type", sorted(_VECTOR_STORE_FILE_CALLS)) +@pytest.mark.parametrize( + "team_models", + [ + pytest.param((("openai/*", _TEAM_OPENAI_PARAMS),), id="openai-only"), + pytest.param( + (("a-langgraph-agent", _LANGGRAPH_AGENT_PARAMS), ("openai/*", _TEAM_OPENAI_PARAMS)), + id="langgraph-first", + ), + ], +) +@pytest.mark.asyncio +async def test_team_key_model_less_vector_store_file_call_reaches_openai_with_the_team_deployment_key( + monkeypatch: pytest.MonkeyPatch, + route_type: str, + team_models: tuple[tuple[str, Mapping[str, object]], ...], +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + method, path, body = _VECTOR_STORE_FILE_CALLS[route_type] + + with respx.mock(assert_all_called=True) as respx_mock: + route: Final = respx_mock.route(method=method, url=f"{_TEAM_OPENAI_BASE}{path}").mock( + return_value=httpx.Response(200, json=body) + ) + await ( + await route_request( + dict(_model_less_vector_store_file_request(route_type)), + _team_a_router(*team_models), + None, + route_type, + user_api_key_dict=_TEAM_A_MEMBER, + ) + ) + + assert route.call_count == 1 + assert route.calls.last.request.headers["Authorization"] == "Bearer team-openai-key" + + +@pytest.mark.parametrize("route_type", sorted(_VECTOR_STORE_FILE_CALLS)) +@pytest.mark.asyncio +async def test_missing_vector_store_files_do_not_cool_down_the_team_deployment( + monkeypatch: pytest.MonkeyPatch, + route_type: str, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + method, path, _ = _VECTOR_STORE_FILE_CALLS[route_type] + router: Final = _team_a_openai_router_beside_a_proxy_openai_wildcard() + + with respx.mock(assert_all_called=True) as respx_mock: + file_route: Final = respx_mock.route(method=method, url=f"{_TEAM_OPENAI_BASE}{path}").mock( + return_value=httpx.Response(404, json={"error": {"message": "No file found", "type": "invalid_request_error"}}) + ) + chat_route: Final = respx_mock.post(f"{_TEAM_OPENAI_BASE}/chat/completions").mock( + return_value=_chat_completion_response("gpt-5.4-mini") + ) + for _ in range(5): + with pytest.raises(litellm.NotFoundError): + await ( + await route_request( + dict(_model_less_vector_store_file_request(route_type)), + router, + None, + route_type, + user_api_key_dict=_TEAM_A_MEMBER, + ) + ) + response: Final = await ( + await route_request( + { + "model": "openai/gpt-5.4-mini", + "messages": list(_HELLO_MESSAGES), + "metadata": {"user_api_key_team_id": "team-a"}, + }, + router, + None, + "acompletion", + user_api_key_dict=_TEAM_A_MEMBER, + ) + ) + + assert file_route.call_count == 5 + assert response.choices[0].message.content == "hi", response + assert chat_route.calls.last.request.headers["Authorization"] == "Bearer team-openai-key" + + +_DIRECT_VECTOR_STORE_FILE_ROUTES: Final = ( + "avector_store_file_content", + "avector_store_file_delete", + "avector_store_file_list", + "avector_store_file_retrieve", +) + + +@pytest.mark.parametrize("route_type", _DIRECT_VECTOR_STORE_FILE_ROUTES) +@pytest.mark.parametrize( + "restricted_member", + [ + pytest.param( + UserAPIKeyAuth(api_key="restricted-member-key", team_id="team-a", models=["a-langgraph-agent"]), + id="key-allowlist", + ), + pytest.param( + UserAPIKeyAuth(api_key="team-blocked-member-key", team_id="team-a", team_models=["a-langgraph-agent"]), + id="team-models", + ), + ], +) +@pytest.mark.asyncio +async def test_a_team_key_restricted_away_from_the_openai_deployment_keeps_proxy_credentials_on_a_file_call( + monkeypatch: pytest.MonkeyPatch, + route_type: str, + restricted_member: UserAPIKeyAuth, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "proxy-openai-key") + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + method, path, body = _VECTOR_STORE_FILE_CALLS[route_type] + + with respx.mock(assert_all_called=True) as respx_mock: + route: Final = respx_mock.route(method=method, url=f"https://api.openai.com/v1{path}").mock( + return_value=httpx.Response(200, json=body) + ) + await ( + await route_request( + dict(_model_less_vector_store_file_request(route_type)), + _team_a_router(("openai/*", _TEAM_OPENAI_PARAMS)), + None, + route_type, + user_api_key_dict=restricted_member, + ) + ) + + assert route.calls.last.request.headers["Authorization"] == "Bearer proxy-openai-key" + + +@pytest.mark.parametrize("route_type", ["avector_store_file_create", "avector_store_file_update"]) +@pytest.mark.asyncio +async def test_a_team_without_an_openai_deployment_never_sends_the_proxy_key_on_a_model_less_file_write( + monkeypatch: pytest.MonkeyPatch, + route_type: str, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "proxy-openai-key") + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + method, path, body = _VECTOR_STORE_FILE_CALLS[route_type] + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.route(method=method, url=f"https://api.openai.com/v1{path}").mock( + return_value=httpx.Response(200, json=body) + ) + with pytest.raises(litellm.APIConnectionError, match="not supported for anthropic"): + await ( + await route_request( + dict(_model_less_vector_store_file_request(route_type)), + _team_a_router(("claude", {"model": "anthropic/claude-sonnet-4-5", "api_key": "team-anthropic"})), + None, + route_type, + user_api_key_dict=_TEAM_A_MEMBER, + ) + ) + + assert route.call_count == 0 + + +class _SuccessRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: Final = asyncio.Event() + self.standard_logging_objects: Final[list[object]] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.standard_logging_objects.append(kwargs.get("standard_logging_object")) + self.logged.set() + + +@pytest.mark.parametrize("route_type", sorted(_VECTOR_STORE_FILE_CALLS)) +@pytest.mark.asyncio +async def test_a_model_less_file_call_logs_the_team_deployment_it_used( + monkeypatch: pytest.MonkeyPatch, + route_type: str, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _SuccessRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + method, path, body = _VECTOR_STORE_FILE_CALLS[route_type] + + with respx.mock(assert_all_called=True) as respx_mock: + respx_mock.route(method=method, url=f"{_TEAM_OPENAI_BASE}{path}").mock( + return_value=httpx.Response(200, json=body) + ) + await ( + await route_request( + dict(_model_less_vector_store_file_request(route_type)), + _team_a_router(("openai/*", _TEAM_OPENAI_PARAMS)), + None, + route_type, + user_api_key_dict=_TEAM_A_MEMBER, + ) + ) + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + + logged: Final = recorder.standard_logging_objects[-1] + assert isinstance(logged, dict), logged + assert (logged["model_id"], logged["model_group"], logged["model"]) == ("team-a-openai/*", "openai/*", "openai/*") + + +@pytest.mark.parametrize("newer_first", [True, False]) +@pytest.mark.asyncio +async def test_a_model_less_file_call_uses_the_team_openai_deployment_ingest_uploads_through( + monkeypatch: pytest.MonkeyPatch, + newer_first: bool, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + older: Final = _team_a_openai_db_deployment("deployment-z", "https://older-openai.example/v1", datetime(2026, 1, 1)) + newer: Final = _team_a_openai_db_deployment("deployment-a", "https://newer-openai.example/v1", datetime(2026, 6, 1)) + router: Final = litellm.Router(model_list=[dict(newer), dict(older)] if newer_first else [dict(older), dict(newer)]) + ingest_deployment: Final = get_team_provider_deployment(router, _TEAM_A_MEMBER, "openai") + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.get(url__regex=r"/v1/vector_stores/vs_team/files/file-ingested$").mock( + return_value=httpx.Response(200, json=_VECTOR_STORE_FILE) + ) + await ( + await route_request( + dict(_model_less_vector_store_file_request("avector_store_file_retrieve")), + router, + None, + "avector_store_file_retrieve", + user_api_key_dict=_TEAM_A_MEMBER, + ) + ) + + assert ingest_deployment is not None + assert route.call_count == 1 + assert route.calls.last.request.headers["Authorization"] == f"Bearer {ingest_deployment.deployment_id}-key" + + +@pytest.mark.asyncio +async def test_registered_store_credential_wins_over_the_team_deployment_on_a_model_less_file_call( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.types.utils import CredentialItem + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="store-credential", + credential_info={"custom_llm_provider": "openai"}, + credential_values={"api_key": "store-credential-key", "api_base": "https://store-openai.example/v1"}, + ) + ], + ) + + with respx.mock(assert_all_called=True) as respx_mock: + route: Final = respx_mock.get("https://store-openai.example/v1/vector_stores/vs_team/files/file-ingested").mock( + return_value=httpx.Response(200, json=_VECTOR_STORE_FILE) + ) + await ( + await route_request( + dict( + _model_less_vector_store_file_request( + "avector_store_file_retrieve", litellm_credential_name="store-credential" + ) + ), + _team_a_router(("openai/*", _TEAM_OPENAI_PARAMS)), + None, + "avector_store_file_retrieve", + user_api_key_dict=_TEAM_A_MEMBER, + ) + ) + + assert route.calls.last.request.headers["Authorization"] == "Bearer store-credential-key" + + + +@pytest.mark.asyncio +async def test_a_model_less_file_call_leaves_the_team_credential_headers_intact_for_later_file_uploads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.openai_files_endpoints.common_utils import get_credentials_for_model + from litellm.types.utils import CredentialItem + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="team-a-openai", + credential_info={"custom_llm_provider": "openai"}, + credential_values={ + "api_key": "team-openai-key", + "api_base": _TEAM_OPENAI_BASE, + "extra_headers": {"X-Team": "team-a"}, + }, + ) + ], + ) + router: Final = _team_a_router(("openai/*", {"model": "openai/*", "litellm_credential_name": "team-a-openai"})) + method, path, body = _VECTOR_STORE_FILE_CALLS["avector_store_file_retrieve"] + + with respx.mock(assert_all_called=True) as respx_mock: + respx_mock.route(method=method, url=f"{_TEAM_OPENAI_BASE}{path}").mock( + return_value=httpx.Response(200, json=body) + ) + upload: Final = respx_mock.post(f"{_TEAM_OPENAI_BASE}/files").mock( + return_value=httpx.Response( + 200, + json={ + "id": "file-uploaded", + "object": "file", + "bytes": 5, + "created_at": 1, + "filename": "notes.txt", + "purpose": "assistants", + "status": "processed", + }, + ) + ) + await ( + await route_request( + dict(_model_less_vector_store_file_request("avector_store_file_retrieve")), + router, + None, + "avector_store_file_retrieve", + user_api_key_dict=_TEAM_A_MEMBER, + ) + ) + await litellm.acreate_file( + file=("notes.txt", b"hello", "text/plain"), + purpose="assistants", + **get_credentials_for_model(llm_router=router, model_id="team-a-openai/*"), + ) + + assert upload.calls.last.request.headers["Content-Type"].startswith("multipart/form-data") + assert upload.calls.last.request.headers["X-Team"] == "team-a" + @pytest.fixture(autouse=True) def discard_cached_user_config_routers() -> Iterator[None]: yield diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_utils.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_utils.py index 5b71fcb8037..8f2f53cba34 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_utils.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_utils.py @@ -1,9 +1,15 @@ +import base64 import json +from collections.abc import Mapping, Sequence +from typing import Final +import httpx import pytest import litellm -from litellm.proxy.vector_store_endpoints.utils import get_litellm_managed_vector_store +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.vector_store_endpoints.utils import get_litellm_managed_vector_store, with_managed_search_file_ids +from litellm.types.utils import SpecialEnums from litellm.vector_stores.vector_store_registry import VectorStoreRegistry @@ -92,3 +98,79 @@ async def test_registry_store_with_non_string_params_is_returned_as_is( assert resolved is store assert resolved["litellm_params"] is stored_params + + +def _managed_file_id(provider_file_id: str) -> str: + unified_file_id: Final = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "unified-1", "openai/team-provider", provider_file_id, "team-deployment" + ) + return base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") + + +class _ManagedFileIds: + def __init__(self, unified_file_ids: Mapping[str, str]) -> None: + self.unified_file_ids: Final = unified_file_ids + + async def get_unified_file_ids_for_provider_file_ids( + self, provider_file_ids: Sequence[str], user_api_key_dict: UserAPIKeyAuth + ) -> Mapping[str, str]: + return {file_id: self.unified_file_ids[file_id] for file_id in provider_file_ids if file_id in self.unified_file_ids} + + +class _UnreachableManagedFilesTable: + async def get_unified_file_ids_for_provider_file_ids( + self, provider_file_ids: Sequence[str], user_api_key_dict: UserAPIKeyAuth + ) -> Mapping[str, str]: + raise httpx.ConnectError("database is unreachable") + + +def _search_hit(file_id: str) -> Mapping[str, object]: + return {"file_id": file_id, "filename": f"{file_id}.txt", "score": 0.5, "content": []} + + +def _search_page(*hits: Mapping[str, object]) -> Mapping[str, object]: + return { + "object": "vector_store.search_results.page", + "search_query": "q", + "data": list(hits), + "has_more": False, + "next_page": None, + } + + +@pytest.mark.asyncio +async def test_search_hits_carry_the_managed_file_id_ingest_returned() -> None: + managed_ids: Final = _ManagedFileIds( + {"file-ingested": _managed_file_id("file-ingested"), "file-mismatched": _managed_file_id("file-other")} + ) + + mapped: Final = await with_managed_search_file_ids( + _search_page(_search_hit("file-ingested"), _search_hit("file-unregistered"), _search_hit("file-mismatched")), + managed_ids, + UserAPIKeyAuth(api_key="team-a-member-key", team_id="team-a"), + ) + + assert mapped == _search_page( + {**_search_hit("file-ingested"), "file_id": _managed_file_id("file-ingested")}, + _search_hit("file-unregistered"), + _search_hit("file-mismatched"), + ) + + +@pytest.mark.asyncio +async def test_search_hits_keep_provider_file_ids_without_the_managed_files_hook() -> None: + page: Final = _search_page(_search_hit("file-ingested")) + + assert await with_managed_search_file_ids(page, None, UserAPIKeyAuth(api_key="k", team_id="team-a")) == page + + +@pytest.mark.asyncio +async def test_search_hits_keep_provider_file_ids_when_the_managed_file_lookup_fails() -> None: + page: Final = _search_page(_search_hit("file-ingested")) + + assert ( + await with_managed_search_file_ids( + page, _UnreachableManagedFilesTable(), UserAPIKeyAuth(api_key="k", team_id="team-a") + ) + == page + ) diff --git a/tests/unit/proxy/vector_store_files_endpoints/test_endpoints.py b/tests/unit/proxy/vector_store_files_endpoints/test_endpoints.py index 271deab7b36..2479c53576f 100644 --- a/tests/unit/proxy/vector_store_files_endpoints/test_endpoints.py +++ b/tests/unit/proxy/vector_store_files_endpoints/test_endpoints.py @@ -9,26 +9,33 @@ A raw or forged managed-looking file id has no ownership row, so without the gua is attached to a vector store or read back under shared provider credentials. """ +import asyncio import base64 +import json from collections.abc import Mapping, Sequence from copy import deepcopy from dataclasses import dataclass +from types import MappingProxyType from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest - - +import respx from fastapi import HTTPException import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.route_llm_request import route_request from litellm.proxy.vector_store_files_endpoints.endpoints import ( + _update_request_data_with_litellm_managed_vector_store_registry, _update_request_data_with_managed_file_id, _with_managed_file_list_ids, _with_provider_file_id_cursors, ) -from litellm.types.utils import SpecialEnums +from litellm.types.utils import CredentialItem, SpecialEnums +from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.types.vector_store_files import ( VectorStoreFileListResponse, VectorStoreFileObject, @@ -290,3 +297,828 @@ async def test_raw_file_id_allowed_when_managed_files_not_required(): assert original is None assert data["file_id"] == RAW_FILE_ID + + +_TEAM_OPENAI_BASE: Final = "https://team-openai.example/v1" +_PROXY_OPENAI_BASE: Final = "https://proxy-openai.example/v1" + + +def _proxy_and_team_b_openai_router( + team_litellm_params: Mapping[str, object] = MappingProxyType( + {"api_key": "team-openai-key", "api_base": _TEAM_OPENAI_BASE} + ), + proxy_litellm_params: Mapping[str, object] = MappingProxyType( + {"api_key": "proxy-openai-key", "api_base": _PROXY_OPENAI_BASE} + ), +) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", **proxy_litellm_params}, + "model_info": {"id": "proxy-openai"}, + }, + { + "model_name": "model_name_team-b_6f0c8a52", + "litellm_params": {"model": "openai/*", **team_litellm_params}, + "model_info": { + "id": "team-b-openai", + "team_id": "team-b", + "team_public_model_name": "openai/*", + "db_model": True, + }, + }, + ] + ) + + +@dataclass(frozen=True) +class ManagedFileCopiesStub: + copies: Mapping[str, str] + + async def can_user_call_unified_file_id(self, unified_file_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + return True + + async def can_user_call_unified_object_id(self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + return False + + async def get_model_file_id_mapping( + self, file_ids: Sequence[str], litellm_parent_otel_span: object + ) -> Mapping[str, Mapping[str, str]]: + return {file_id: self.copies for file_id in file_ids} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller_team_id", "expected_base", "expected_authorization"), + [ + ("team-b", _TEAM_OPENAI_BASE, "Bearer team-openai-key"), + ("team-c", _PROXY_OPENAI_BASE, "Bearer proxy-openai-key"), + ], +) +async def test_managed_file_id_reaches_its_upload_deployment_only_for_that_team( + monkeypatch: pytest.MonkeyPatch, caller_team_id: str, expected_base: str, expected_authorization: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + caller: Final = UserAPIKeyAuth(api_key="sk-test", user_id="caller-user", team_id=caller_team_id) + router: Final = _proxy_and_team_b_openai_router() + managed_file_id: Final = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "ingested-unified-id", "openai/*", RAW_FILE_ID, "team-b-openai" + ).encode() + ) + .decode() + .rstrip("=") + ) + + data, _ = await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": managed_file_id}, + file_id=managed_file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=caller, + managed_files_obj=ManagedResourceAccessCheckerStub(file_access="allow"), + llm_router=router, + ) + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.get(f"{expected_base}/vector_stores/vs-test/files/{RAW_FILE_ID}").mock( + return_value=httpx.Response(200, json=dict(_vector_store_file_row(RAW_FILE_ID))) + ) + await ( + await route_request( + {**data, "custom_llm_provider": "openai", "metadata": {"user_api_key_team_id": caller.team_id}}, + router, + None, + "avector_store_file_retrieve", + user_api_key_dict=caller, + ) + ) + + assert upstream.calls.last.request.headers["Authorization"] == expected_authorization + + +_CALLER_OPENAI_BASE: Final = "https://caller-openai.example/v1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller_credentials", "expected_base", "expected_authorization"), + [ + ({"api_base": _CALLER_OPENAI_BASE}, _TEAM_OPENAI_BASE, "Bearer team-openai-key"), + ({"api_base": _CALLER_OPENAI_BASE, "api_key": "caller-key"}, _CALLER_OPENAI_BASE, "Bearer caller-key"), + ], +) +async def test_managed_file_id_sends_its_deployment_key_only_to_that_deployment_base( + monkeypatch: pytest.MonkeyPatch, + caller_credentials: Mapping[str, str], + expected_base: str, + expected_authorization: str, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + caller: Final = UserAPIKeyAuth(api_key="sk-test", user_id="caller-user", team_id="team-b") + router: Final = _proxy_and_team_b_openai_router() + managed_file_id: Final = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "ingested-unified-id", "openai/*", RAW_FILE_ID, "team-b-openai" + ).encode() + ) + .decode() + .rstrip("=") + ) + + data, _ = await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": managed_file_id, **caller_credentials}, + file_id=managed_file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=caller, + managed_files_obj=ManagedResourceAccessCheckerStub(file_access="allow"), + llm_router=router, + ) + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.get(f"{expected_base}/vector_stores/vs-test/files/{RAW_FILE_ID}").mock( + return_value=httpx.Response(200, json=dict(_vector_store_file_row(RAW_FILE_ID))) + ) + await ( + await route_request( + {**data, "custom_llm_provider": "openai", "metadata": {"user_api_key_team_id": caller.team_id}}, + router, + None, + "avector_store_file_retrieve", + user_api_key_dict=caller, + ) + ) + + assert upstream.calls.last.request.headers["Authorization"] == expected_authorization + + +class _SuccessRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: Final = asyncio.Event() + self.standard_logging_objects: Final[list[object]] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.standard_logging_objects.append(kwargs.get("standard_logging_object")) + self.logged.set() + + +def _gpt_4o_mini_deployment( + deployment_id: str, api_base: str, order: int, team_id: str | None = None +) -> Mapping[str, object]: + return { + "model_name": "gpt-4o-mini" if team_id is None else f"model_name_{team_id}_{deployment_id}", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": f"{deployment_id}-key", + "api_base": api_base, + "order": order, + }, + "model_info": {"id": deployment_id} + if team_id is None + else {"id": deployment_id, "team_id": team_id, "team_public_model_name": "gpt-4o-mini", "db_model": True}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller_team_id", "upload_deployment_id", "upload_base"), + [ + (None, "upload-openai", "https://upload-openai.example/v1"), + ("team-b", "team-b-upload", "https://team-b-upload.example/v1"), + ], +) +async def test_managed_file_id_on_its_upload_deployment_logs_that_deployment_under_its_public_model_name( + monkeypatch: pytest.MonkeyPatch, caller_team_id: str | None, upload_deployment_id: str, upload_base: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _SuccessRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + caller: Final = UserAPIKeyAuth(api_key="sk-test", user_id="caller-user", team_id=caller_team_id) + router: Final = litellm.Router( + model_list=[ + dict(_gpt_4o_mini_deployment("sibling-openai", "https://sibling-openai.example/v1", order=1)), + dict(_gpt_4o_mini_deployment("upload-openai", "https://upload-openai.example/v1", order=2)), + dict(_gpt_4o_mini_deployment("team-b-sibling", "https://team-b-sibling.example/v1", 1, "team-b")), + dict(_gpt_4o_mini_deployment("team-b-upload", "https://team-b-upload.example/v1", 2, "team-b")), + ] + ) + managed_file_id: Final = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "ingested-unified-id", "gpt-4o-mini", RAW_FILE_ID, upload_deployment_id + ).encode() + ) + .decode() + .rstrip("=") + ) + + data, _ = await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": managed_file_id}, + file_id=managed_file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=caller, + managed_files_obj=ManagedResourceAccessCheckerStub(file_access="allow"), + llm_router=router, + ) + with respx.mock(assert_all_called=False) as respx_mock: + upstream: Final = respx_mock.get(f"{upload_base}/vector_stores/vs-test/files/{RAW_FILE_ID}").mock( + return_value=httpx.Response(200, json=dict(_vector_store_file_row(RAW_FILE_ID))) + ) + sibling: Final = respx_mock.route().mock( + return_value=httpx.Response(200, json=dict(_vector_store_file_row(RAW_FILE_ID))) + ) + await ( + await route_request( + {**data, "custom_llm_provider": "openai", "metadata": {"user_api_key_team_id": caller.team_id}}, + router, + None, + "avector_store_file_retrieve", + user_api_key_dict=caller, + ) + ) + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + + assert (upstream.call_count, sibling.call_count) == (1, 0) + assert upstream.calls.last.request.headers["Authorization"] == f"Bearer {upload_deployment_id}-key" + logged: Final = recorder.standard_logging_objects[-1] + assert isinstance(logged, dict), logged + assert (logged["model_group"], logged["model_id"]) == ("gpt-4o-mini", upload_deployment_id) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("key_models", "expected_base", "expected_authorization"), + [ + (["gpt-4o-mini"], "https://allowed-openai.example/v1", "Bearer allowed-key"), + ([], "https://restricted-openai.example/v1", "Bearer restricted-key"), + ], +) +async def test_managed_file_id_reaches_its_upload_deployment_only_when_the_key_may_use_it( + monkeypatch: pytest.MonkeyPatch, key_models: list[str], expected_base: str, expected_authorization: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + caller: Final = UserAPIKeyAuth(api_key="sk-test", user_id="caller-user", models=key_models) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "allowed-key", + "api_base": "https://allowed-openai.example/v1", + }, + "model_info": {"id": "gpt-4o-mini-id"}, + }, + { + "model_name": "openai/*", + "litellm_params": { + "model": "openai/*", + "api_key": "restricted-key", + "api_base": "https://restricted-openai.example/v1", + }, + "model_info": {"id": "restricted-openai"}, + }, + ] + ) + managed_file_id: Final = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "forged-unified-id", "gpt-4o-mini", RAW_FILE_ID, "restricted-openai" + ).encode() + ) + .decode() + .rstrip("=") + ) + + data, _ = await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": managed_file_id}, + file_id=managed_file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=caller, + managed_files_obj=ManagedResourceAccessCheckerStub(file_access="allow"), + llm_router=router, + ) + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.get(f"{expected_base}/vector_stores/vs-test/files/{RAW_FILE_ID}").mock( + return_value=httpx.Response(200, json=dict(_vector_store_file_row(RAW_FILE_ID))) + ) + await ( + await route_request( + {**data, "custom_llm_provider": "openai"}, + router, + None, + "avector_store_file_retrieve", + user_api_key_dict=caller, + ) + ) + + assert upstream.calls.last.request.headers["Authorization"] == expected_authorization + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller_team_id", "upload_copy_status", "expected_url", "expected_authorization"), + [ + ("team-b", 404, f"{_TEAM_OPENAI_BASE}/vector_stores/vs-test/files/file-team-copy", "Bearer team-openai-key"), + ("team-c", 200, f"{_PROXY_OPENAI_BASE}/vector_stores/vs-test/files/file-proxy-copy", "Bearer proxy-openai-key"), + ], +) +async def test_managed_file_id_with_a_copy_on_the_callers_team_deployment_reaches_that_copy( + monkeypatch: pytest.MonkeyPatch, + caller_team_id: str, + upload_copy_status: int, + expected_url: str, + expected_authorization: str, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + caller: Final = UserAPIKeyAuth(api_key="sk-test", user_id="caller-user", team_id=caller_team_id) + router: Final = _proxy_and_team_b_openai_router() + managed_file_id: Final = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "two-copy-unified-id", "openai/*", "file-proxy-copy", "proxy-openai" + ).encode() + ) + .decode() + .rstrip("=") + ) + + with respx.mock(assert_all_called=True) as respx_mock: + respx_mock.get(f"{_PROXY_OPENAI_BASE}/vector_stores/vs-test/files/file-proxy-copy").mock( + return_value=httpx.Response(upload_copy_status, json=dict(_vector_store_file_row("file-proxy-copy"))) + ) + upstream: Final = respx_mock.get(expected_url).mock( + return_value=httpx.Response(200, json=dict(_vector_store_file_row(expected_url.rsplit("/", 1)[-1]))) + ) + data, _ = await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": managed_file_id}, + file_id=managed_file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=caller, + managed_files_obj=ManagedFileCopiesStub( + copies={"proxy-openai": "file-proxy-copy", "team-b-openai": "file-team-copy"} + ), + llm_router=router, + ) + await ( + await route_request( + {**data, "custom_llm_provider": "openai", "metadata": {"user_api_key_team_id": caller.team_id}}, + router, + None, + "avector_store_file_retrieve", + user_api_key_dict=caller, + ) + ) + + assert upstream.calls.last.request.headers["Authorization"] == expected_authorization + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "registered_store", + [ + LiteLLM_ManagedVectorStore( + vector_store_id="vs-test", custom_llm_provider="openai", litellm_params={"api_key": "store-openai-key"} + ), + LiteLLM_ManagedVectorStore( + vector_store_id="vs-test", custom_llm_provider="openai", litellm_credential_name="store-openai" + ), + ], +) +async def test_managed_file_id_on_a_store_with_its_own_credentials_keeps_the_upload_copy( + monkeypatch: pytest.MonkeyPatch, registered_store: LiteLLM_ManagedVectorStore +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="store-openai", credential_values={"api_key": "store-openai-key"}, credential_info={} + ) + ], + ) + caller: Final = UserAPIKeyAuth(api_key="sk-test", user_id="caller-user", team_id="team-b") + router: Final = _proxy_and_team_b_openai_router() + managed_file_id: Final = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "two-copy-unified-id", "openai/*", "file-proxy-copy", "proxy-openai" + ).encode() + ) + .decode() + .rstrip("=") + ) + + data, _ = await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": managed_file_id}, + file_id=managed_file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=caller, + managed_files_obj=ManagedFileCopiesStub( + copies={"proxy-openai": "file-proxy-copy", "team-b-openai": "file-team-copy"} + ), + llm_router=router, + managed_vector_store=registered_store, + ) + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.get(f"{_PROXY_OPENAI_BASE}/vector_stores/vs-test/files/file-proxy-copy").mock( + return_value=httpx.Response(200, json=dict(_vector_store_file_row("file-proxy-copy"))) + ) + await ( + await route_request( + _update_request_data_with_litellm_managed_vector_store_registry( + data=data, + vector_store_id="vs-test", + llm_router=router, + managed_vector_store=registered_store, + should_lookup_registry=False, + ), + router, + None, + "avector_store_file_retrieve", + user_api_key_dict=caller, + ) + ) + + assert upstream.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route_type", "method", "reply"), + [ + ("avector_store_file_retrieve", "GET", dict(_vector_store_file_row("file-proxy-copy"))), + ( + "avector_store_file_delete", + "DELETE", + {"id": "file-proxy-copy", "object": "vector_store.file.deleted", "deleted": True}, + ), + ], +) +async def test_managed_file_id_on_a_store_holding_only_the_upload_copy_reaches_that_copy( + monkeypatch: pytest.MonkeyPatch, route_type: str, method: str, reply: Mapping[str, object] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + caller: Final = UserAPIKeyAuth(api_key="sk-test", user_id="caller-user", team_id="team-b") + router: Final = _proxy_and_team_b_openai_router() + managed_file_id: Final = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "two-copy-unified-id", "openai/*", "file-proxy-copy", "proxy-openai" + ).encode() + ) + .decode() + .rstrip("=") + ) + + with respx.mock(assert_all_called=False) as respx_mock: + team_deployment: Final = respx_mock.route(host="team-openai.example") + upstream: Final = respx_mock.route(url=f"{_PROXY_OPENAI_BASE}/vector_stores/vs-test/files/file-proxy-copy").mock( + side_effect=lambda request: httpx.Response( + 200, json=dict(reply) if request.method == method else dict(_vector_store_file_row("file-proxy-copy")) + ) + ) + data, _ = await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": managed_file_id}, + file_id=managed_file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=caller, + managed_files_obj=ManagedFileCopiesStub( + copies={"proxy-openai": "file-proxy-copy", "team-b-openai": "file-team-copy"} + ), + llm_router=router, + ) + await (await route_request({**data, "custom_llm_provider": "openai", "metadata": {"user_api_key_team_id": caller.team_id}}, router, None, route_type, user_api_key_dict=caller)) + + assert team_deployment.call_count == 0 + assert upstream.calls.last.request.method == method + assert upstream.calls.last.request.headers["Authorization"] == "Bearer proxy-openai-key" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("store_holds_upload_copy", "upstream_base", "attached_copy", "upstream_key"), + [ + (False, _TEAM_OPENAI_BASE, "file-team-copy", "team-openai-key"), + (True, _PROXY_OPENAI_BASE, "file-proxy-copy", "proxy-openai-key"), + ], +) +async def test_managed_file_id_attached_by_a_team_key_attaches_the_copy_the_store_can_hold( + monkeypatch: pytest.MonkeyPatch, + store_holds_upload_copy: bool, + upstream_base: str, + attached_copy: str, + upstream_key: str, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + caller: Final = UserAPIKeyAuth(api_key="sk-test", user_id="caller-user", team_id="team-b") + router: Final = _proxy_and_team_b_openai_router() + managed_file_id: Final = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "two-copy-unified-id", "openai/*", "file-proxy-copy", "proxy-openai" + ).encode() + ) + .decode() + .rstrip("=") + ) + + with respx.mock(assert_all_called=True) as respx_mock: + respx_mock.get(f"{_PROXY_OPENAI_BASE}/vector_stores/vs-test/files/file-proxy-copy").mock( + return_value=( + httpx.Response(200, json=dict(_vector_store_file_row("file-proxy-copy"))) + if store_holds_upload_copy + else httpx.Response(404, json={"error": {"message": "No file found", "type": "invalid_request_error"}}) + ) + ) + upstream: Final = respx_mock.post(f"{upstream_base}/vector_stores/vs-test/files").mock( + return_value=httpx.Response(200, json=dict(_vector_store_file_row(attached_copy))) + ) + data, _ = await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": managed_file_id}, + file_id=managed_file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=caller, + managed_files_obj=ManagedFileCopiesStub( + copies={"proxy-openai": "file-proxy-copy", "team-b-openai": "file-team-copy"} + ), + llm_router=router, + ) + await ( + await route_request( + {**data, "custom_llm_provider": "openai", "metadata": {"user_api_key_team_id": caller.team_id}}, + router, + None, + "avector_store_file_create", + user_api_key_dict=caller, + ) + ) + + assert upstream.call_count == 1 + assert json.loads(upstream.calls.last.request.content)["file_id"] == attached_copy + assert upstream.calls.last.request.headers["Authorization"] == f"Bearer {upstream_key}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("check_status", [429, 500]) +async def test_managed_file_id_whose_store_check_fails_without_a_404_still_reaches_the_upload_copy( + monkeypatch: pytest.MonkeyPatch, check_status: int +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + caller: Final = UserAPIKeyAuth(api_key="sk-test", user_id="caller-user", team_id="team-b") + router: Final = _proxy_and_team_b_openai_router() + managed_file_id: Final = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "two-copy-unified-id", "openai/*", "file-proxy-copy", "proxy-openai" + ).encode() + ) + .decode() + .rstrip("=") + ) + + with respx.mock(assert_all_called=False) as respx_mock: + team_deployment: Final = respx_mock.route(host="team-openai.example") + upstream: Final = respx_mock.get(f"{_PROXY_OPENAI_BASE}/vector_stores/vs-test/files/file-proxy-copy").mock( + side_effect=[ + httpx.Response(check_status, json={"error": {"message": "try again", "type": "server_error"}}), + httpx.Response(200, json=dict(_vector_store_file_row("file-proxy-copy"))), + ] + ) + data, _ = await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": managed_file_id}, + file_id=managed_file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=caller, + managed_files_obj=ManagedFileCopiesStub( + copies={"proxy-openai": "file-proxy-copy", "team-b-openai": "file-team-copy"} + ), + llm_router=router, + ) + response: Final = await ( + await route_request( + {**data, "custom_llm_provider": "openai", "metadata": {"user_api_key_team_id": caller.team_id}}, + router, + None, + "avector_store_file_retrieve", + user_api_key_dict=caller, + ) + ) + + assert response["id"] == "file-proxy-copy" + assert upstream.call_count == 2 + assert upstream.calls.last.request.headers["Authorization"] == "Bearer proxy-openai-key" + assert team_deployment.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route_type", "method", "path"), + [ + ("avector_store_file_create", "POST", "/vector_stores/vs-managed/files"), + ("avector_store_file_retrieve", "GET", "/vector_stores/vs-managed/files/file-team-copy"), + ("avector_store_file_delete", "DELETE", "/vector_stores/vs-managed/files/file-team-copy"), + ], +) +async def test_managed_file_id_on_a_managed_store_reaches_the_store_deployment( + monkeypatch: pytest.MonkeyPatch, route_type: str, method: str, path: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + caller: Final = UserAPIKeyAuth(api_key="sk-test", user_id="caller-user", team_id="team-b") + router: Final = _proxy_and_team_b_openai_router() + managed_file_id: Final = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "team-upload-unified-id", "openai/*", "file-team-copy", "team-b-openai" + ).encode() + ) + .decode() + .rstrip("=") + ) + managed_store_id: Final = ( + base64.urlsafe_b64encode( + b"litellm_proxy:vector_store;unified_id,store-uuid;target_model_names,openai/*;" + b"resource_id,vs-managed;model_id,proxy-openai" + ) + .decode() + .rstrip("=") + ) + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.request(method, f"{_PROXY_OPENAI_BASE}{path}").mock( + return_value=httpx.Response(200, json=dict(_vector_store_file_row("file-team-copy"))) + ) + data, _ = await _update_request_data_with_managed_file_id( + data={"vector_store_id": managed_store_id, "file_id": managed_file_id}, + file_id=managed_file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=caller, + managed_files_obj=ManagedFileCopiesStub( + copies={"proxy-openai": "file-proxy-copy", "team-b-openai": "file-team-copy"} + ), + llm_router=router, + ) + await ( + await route_request( + { + **_update_request_data_with_litellm_managed_vector_store_registry( + data=data, vector_store_id=managed_store_id, llm_router=router + ), + "custom_llm_provider": "openai", + "metadata": {"user_api_key_team_id": "team-b"}, + }, + router, + None, + route_type, + user_api_key_dict=caller, + ) + ) + + assert upstream.call_count == 1 + assert upstream.calls.last.request.headers["Authorization"] == "Bearer proxy-openai-key" + + +@pytest.mark.asyncio +async def test_managed_file_id_on_its_upload_deployment_carries_no_options_from_another_deployment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="team-b-openai", + credential_values={ + "api_key": "team-openai-key", + "api_base": _TEAM_OPENAI_BASE, + "extra_headers": {"X-Team": "team-b"}, + }, + credential_info={}, + ) + ], + ) + caller: Final = UserAPIKeyAuth(api_key="sk-test", user_id="caller-user", team_id="team-b") + upload_base: Final = "https://team-upload.example/v1" + router: Final = litellm.Router( + model_list=[ + { + "model_name": "model_name_team-b_6f0c8a52", + "litellm_params": {"model": "openai/*", "litellm_credential_name": "team-b-openai"}, + "model_info": { + "id": "team-b-openai", + "team_id": "team-b", + "team_public_model_name": "openai/*", + "db_model": True, + }, + }, + { + "model_name": "model_name_team-b_upload", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "upload-key", "api_base": upload_base}, + "model_info": { + "id": "team-b-upload", + "team_id": "team-b", + "team_public_model_name": "team-upload", + "db_model": True, + }, + }, + ] + ) + managed_file_id: Final = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "upload-unified-id", "team-upload", "file-upload-copy", "team-b-upload" + ).encode() + ) + .decode() + .rstrip("=") + ) + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.get(f"{upload_base}/vector_stores/vs-test/files/file-upload-copy").mock( + return_value=httpx.Response(200, json=dict(_vector_store_file_row("file-upload-copy"))) + ) + data, _ = await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": managed_file_id}, + file_id=managed_file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=caller, + managed_files_obj=ManagedFileCopiesStub(copies={"team-b-upload": "file-upload-copy"}), + llm_router=router, + ) + await ( + await route_request( + {**data, "custom_llm_provider": "openai", "metadata": {"user_api_key_team_id": caller.team_id}}, + router, + None, + "avector_store_file_retrieve", + user_api_key_dict=caller, + ) + ) + + assert upstream.calls.last.request.headers["Authorization"] == "Bearer upload-key" + assert "X-Team" not in upstream.calls.last.request.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "credential_setting", [{"timeout": 30}, {"extra_headers": {"X-Proxy": "on"}}, {"no-log": True}] +) +async def test_store_check_sends_the_upload_deployment_named_credential_options( + monkeypatch: pytest.MonkeyPatch, credential_setting: Mapping[str, object] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="proxy-openai", + credential_values={"api_key": "proxy-openai-key", "api_base": _PROXY_OPENAI_BASE, **credential_setting}, + credential_info={}, + ) + ], + ) + caller: Final = UserAPIKeyAuth(api_key="sk-test", user_id="caller-user", team_id="team-b") + router: Final = _proxy_and_team_b_openai_router( + proxy_litellm_params=MappingProxyType({"litellm_credential_name": "proxy-openai"}) + ) + managed_file_id: Final = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "text/plain", "two-copy-unified-id", "openai/*", "file-proxy-copy", "proxy-openai" + ).encode() + ) + .decode() + .rstrip("=") + ) + + with respx.mock(assert_all_called=True) as respx_mock: + upload_copy: Final = respx_mock.get(f"{_PROXY_OPENAI_BASE}/vector_stores/vs-test/files/file-proxy-copy").mock( + return_value=httpx.Response(200, json=dict(_vector_store_file_row("file-proxy-copy"))) + ) + data, _ = await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": managed_file_id}, + file_id=managed_file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=caller, + managed_files_obj=ManagedFileCopiesStub( + copies={"proxy-openai": "file-proxy-copy", "team-b-openai": "file-team-copy"} + ), + llm_router=router, + ) + await ( + await route_request( + {**data, "custom_llm_provider": "openai", "metadata": {"user_api_key_team_id": caller.team_id}}, + router, + None, + "avector_store_file_retrieve", + user_api_key_dict=caller, + ) + ) + + store_check, call = upload_copy.calls + assert call.request.headers["Authorization"] == "Bearer proxy-openai-key" + assert store_check.request.headers.get("X-Proxy") == call.request.headers.get("X-Proxy") diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 13a45ceeaf3..eaa0b3970bd 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -7317,6 +7317,71 @@ def test_load_credentials_from_list_fills_kwargs_from_the_loaded_credential_with assert _credential_warnings(caplog) == [] +@pytest.mark.asyncio +async def test_a_call_by_credential_name_leaves_its_headers_intact_for_later_uploads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="team-openai", + credential_values={ + "api_key": "team-openai-key", + "api_base": "https://team-openai.example/v1", + "extra_headers": {"X-Team": "team-a"}, + }, + credential_info={}, + ) + ], + ) + + with respx.mock(assert_all_called=True) as upstream: + upstream.get("https://team-openai.example/v1/vector_stores/vs-1/files/file-1").respond( + 200, + json={ + "id": "file-1", + "object": "vector_store.file", + "created_at": 1, + "usage_bytes": 1, + "vector_store_id": "vs-1", + "status": "completed", + "last_error": None, + "chunking_strategy": {"type": "auto"}, + "attributes": {}, + }, + ) + upload: Final = upstream.post("https://team-openai.example/v1/files").respond( + 200, + json={ + "id": "file-uploaded", + "object": "file", + "bytes": 5, + "created_at": 1, + "filename": "notes.txt", + "purpose": "assistants", + "status": "processed", + }, + ) + await litellm.avector_store_file_retrieve( + vector_store_id="vs-1", + file_id="file-1", + custom_llm_provider="openai", + litellm_credential_name="team-openai", + ) + await litellm.acreate_file( + file=("notes.txt", b"hello", "text/plain"), + purpose="assistants", + custom_llm_provider="openai", + litellm_credential_name="team-openai", + ) + + assert upload.calls.last.request.headers["Content-Type"].startswith("multipart/form-data") + assert upload.calls.last.request.headers["X-Team"] == "team-a" + + _MISTRAL_CHAT_URL: Final = "https://api.mistral.ai/v1/chat/completions" _MISTRAL_REPLY: Final = { "id": "chatcmpl-oidc",