mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(proxy)!: return managed file ids from rag ingest and use team credentials on vector store file routes (#50459)
* fix(proxy): give rag ingest managed file ids and team credentials on vector store file routes RAG ingest through a team's OpenAI deployment now registers the uploaded file as a LiteLLM managed file and returns its managed id, so GET /v1/files lists it for the caller. Re-ingesting a managed id attaches the provider file and returns the caller's managed id Model-less vector store file calls (create, list, retrieve, content, update, delete) from a team key now use the team's OpenAI deployment that ingest uploads through, via the router, and fall back to the direct path when the key may not use one. Explicit api_key, api_base, litellm_credential_name, and forwarded provider keys still win, and blank values count as absent * fix(proxy): route managed vector store file ids to their team deployment and keep missing-file 404s off the router Model-less team vector store file calls now go straight to the provider with the team deployment's credentials, so 404s on missing files no longer put the team deployment in cooldown. A managed file id on the vector store file routes now uses the deployment named by its llm_output_file_model_id when that deployment is proxy-wide or belongs to the caller's team, and falls back to target_model_names otherwise. * fix(rag): keep managed ids and team checks on every file and search path Ingest no longer registers a raw file id when file_url is empty, a managed file id on a proxy without the managed files hook answers 400, and vector store search and /v1/rag/query hits show the managed id ingest returned. Model-less vector store file calls use the team deployment only after the key, team, and member model checks pass, fall through to the router when the team has no OpenAI deployment, and log the deployment they used. * fix(rag): keep provider ids on lookup failures and mint managed ids only for keys that can use them A database error while mapping search hits to managed ids keeps OpenAI's ids instead of failing a finished search. Ingest registers a managed id only when the caller passes the team deployment's model checks, so a model-limited key keeps the raw id it can still use, and the file routes share that check. * fix(rag): attach the managed file copy that matches the ingest credentials * fix(vector_store_files): route a managed file id to the caller team's copy * fix(vector_store_files): keep the upload copy for a store registered with its own credentials * fix(vector_store_files): use a managed file's team copy only when the store holds it * fix(vector_store_files): attach a managed file's team copy and keep named credential options out of the store check * fix(vector_store_files): attach the team copy only to a store that routes to the team deployment * fix(vector_store_files): keep a managed store's routing and check the upload copy through its own deployment * fix(vector_store_files): route a managed file to its upload deployment only when the key may call it A managed file id names its upload deployment, and with require_managed_files off that id is not checked against an ownership record. The routes now use that deployment only when the caller may call its public model, else they keep target_model_names routing as on main. * fix(vector_store_files): log a pinned managed file call under its deployment's public model name * fix(vector_store_files): send a pinned managed file call's deployment key only to that deployment's base * fix(vector_store_files): keep stored credential headers intact and fail the store check on errors other than 404 * fix(vector_store_files): answer a failed store check with the provider status and copy credentials read by name * fix(vector_store_files): send a call to the upload copy when the store check fails without a 404 A store check error other than a 404 now leaves the call on the upload copy, the way main always routed it, so the real call keeps the router's retries, the provider's status, and its failure log instead of failing on one transient error * fix(vector_store_files): keep a credential's no-log key out of the store check A named credential carrying a no-log key made the store check call aretrieve with no-log twice, so every managed id call through that deployment answered 500 where main answered 200 --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tavren27
GitHub
parent
fe6cf07cc8
commit
96b916028a
20 changed files with 2922 additions and 113 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue