diff --git a/.github/workflows/test-unit-proxy-endpoints.yml b/.github/workflows/test-unit-proxy-endpoints.yml index 645996f779d..2ea3c521e8b 100644 --- a/.github/workflows/test-unit-proxy-endpoints.yml +++ b/.github/workflows/test-unit-proxy-endpoints.yml @@ -38,6 +38,8 @@ jobs: tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/batches_endpoints + tests/test_litellm/proxy/fine_tuning_endpoints + tests/test_litellm/proxy/vector_store_files_endpoints tests/test_litellm/proxy/video_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/image_endpoints diff --git a/CLAUDE.md b/CLAUDE.md index 59929143c46..f1bb46c1fd3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -9,7 +9,7 @@ Don't assume that the existing code is correct or the right way of doing things - easy to maintain/change - modern -In that order of importance +In descending order of importance When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate @@ -41,7 +41,7 @@ Python max line length is 120, not 88 When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing -`make pre-commit` saves its complete output to a log file in .git (overwriting previous pre-commit logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice +`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in diff --git a/Makefile b/Makefile index 39713a9f724..9d7eeb55879 100644 --- a/Makefile +++ b/Makefile @@ -8,7 +8,7 @@ lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \ lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \ install-dev install-proxy-dev install-test-deps install-hooks \ - install-helm-unittest check-circular-imports check-import-safety pre-commit \ + install-helm-unittest check-circular-imports check-import-safety check pre-commit \ lint-install lint-fetch-base bootstrap bootstrap-python bootstrap-dashboard # Default target @@ -24,7 +24,8 @@ help: @echo " make install-test-deps - Install the full local test environment" @echo " make install-helm-unittest - Install helm unittest plugin" @echo " make install-hooks - Install git hooks (Conventional Commits + Branches)" - @echo " make pre-commit - Run CI-equivalent lint on staged files (run before committing)" + @echo " make check - Run CI-equivalent lint on staged files, or on the diff vs the base branch when nothing is staged" + @echo " make pre-commit - Legacy alias for make check" @echo " make format - Apply ruff format code formatting" @echo " make format-check - Check ruff format code formatting (matches CI)" @echo " make lint - Run all linting (Ruff, basedpyright, format check, circular imports, import safety)" @@ -242,13 +243,20 @@ lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline # Faster linting for local development (only checks changed code) lint-dev: lint-format-changed check-circular-imports check-import-safety -# Run the gating CI checks against your staged files right before committing. Mirrors +# Run the gating CI checks against your changes. Scopes to staged files when anything +# is staged (warning about changed files left unstaged); with nothing staged it falls +# back to the working tree's diff against the merge base with the base branch, so a +# fresh merge commit or an unstaged working tree still gets checked. Mirrors # test-linting.yml (Python), test-litellm-ui-build.yml's frontend-lint (dashboard), and -# check-ui-api-types.yml (API-type drift), skipping any whose files you didn't stage. +# check-ui-api-types.yml (API-type drift), skipping any whose files aren't in scope. # Not auto-installed as a git hook so it never slows an unrelated human commit. -pre-commit: bootstrap-python +check: bootstrap-python ./scripts/pre_commit_lint.sh +pre-commit: + @echo "make pre-commit is a legacy alias; use make check" >&2 + @$(MAKE) check + # Testing targets test: install-test-deps $(UV_RUN) pytest tests/ diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 32b8eb3d4d0..0385f7a96e7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 28842 + "limit": 27731 }, "reportArgumentType": { - "limit": 2634 + "limit": 2626 }, "reportAssignmentType": { "limit": 329 @@ -12,7 +12,7 @@ "limit": 514 }, "reportCallIssue": { - "limit": 117 + "limit": 116 }, "reportConstantRedefinition": { "limit": 40 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 9103 + "limit": 8807 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5843 + "limit": 5835 }, "reportMissingTypeArgument": { - "limit": 15816 + "limit": 15790 }, "reportMissingTypeStubs": { "limit": 40 @@ -72,7 +72,7 @@ "limit": 0 }, "reportOptionalMemberAccess": { - "limit": 1078 + "limit": 1077 }, "reportOptionalOperand": { "limit": 0 @@ -90,28 +90,28 @@ "limit": 8 }, "reportReturnType": { - "limit": 218 + "limit": 217 }, "reportTypedDictNotRequiredAccess": { - "limit": 27 + "limit": 26 }, "reportUndefinedVariable": { "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45098 + "limit": 45063 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39826 + "limit": 39773 }, "reportUnknownParameterType": { - "limit": 20237 + "limit": 20207 }, "reportUnknownVariableType": { - "limit": 31371 + "limit": 31281 }, "reportUnnecessaryCast": { "limit": 122 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 864 + "limit": 862 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index d8318962633..f0914240f79 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -3,9 +3,21 @@ import base64 import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Final, + List, + Literal, + Optional, + Protocol, + TypedDict, + Union, + cast, +) from uuid import NAMESPACE_URL, uuid5 from fastapi import HTTPException @@ -98,33 +110,76 @@ def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMB try: batch_obj: Final = LiteLLMBatch.model_validate(_decode_json_blob(row.file_object)) except Exception as e: - verbose_logger.warning( - f"Failed to parse batch object {row.unified_object_id}: {_sanitized_parse_error(e)}" - ) + verbose_logger.warning(f"Failed to parse batch object {row.unified_object_id}: {_sanitized_parse_error(e)}") return None batch_obj.id = row.unified_object_id return batch_obj -def _parse_managed_file_object( - raw_file_object: object, unified_file_id: str -) -> Optional[OpenAIFileObject]: +def _parse_managed_file_object(raw_file_object: object, unified_file_id: str) -> Optional[OpenAIFileObject]: if raw_file_object is None: return None try: return OpenAIFileObject.model_validate(raw_file_object) except Exception as e: - verbose_logger.warning( - f"Failed to parse managed file object {unified_file_id}: {_sanitized_parse_error(e)}" - ) + verbose_logger.warning(f"Failed to parse managed file object {unified_file_id}: {_sanitized_parse_error(e)}") return None +class _ManagedFileRow(Protocol): + unified_file_id: str + file_object: OpenAIFileObject + storage_backend: Optional[str] + storage_url: Optional[str] + created_by: Optional[str] + team_id: Optional[str] + + def model_dump(self) -> Mapping[str, object]: ... + + +class _ManagedFileTableActions(Protocol): + async def find_first(self, where: Mapping[str, object]) -> Optional[_ManagedFileRow]: ... + + async def find_many(self, where: Mapping[str, object]) -> Sequence[_ManagedFileRow]: ... + + async def upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]) -> _ManagedFileRow: ... + + async def delete(self, where: Mapping[str, str]) -> Optional[_ManagedFileRow]: ... + + +class _ManagedObjectTableActions(Protocol): + async def find_first(self, where: Mapping[str, object]) -> "Optional[PrismaManagedObjectRow]": ... + + async def find_many( + self, + where: Mapping[str, object], + take: int, + order: Union[Mapping[str, str], Sequence[Mapping[str, str]]], + cursor: Mapping[str, str] = ..., + skip: int = ..., + ) -> "Sequence[PrismaManagedObjectRow]": ... + + async def upsert( + self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]] + ) -> "PrismaManagedObjectRow": ... + + +class _CursorPageArgs(TypedDict, total=False): + cursor: Mapping[str, str] + skip: int + + +def _managed_file_table(prisma_client: PrismaClient) -> _ManagedFileTableActions: + return prisma_client.db.litellm_managedfiletable + + +def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableActions: + return prisma_client.db.litellm_managedobjecttable + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes - def __init__( - self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient - ): + def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient): self.internal_usage_cache = internal_usage_cache self.prisma_client = prisma_client @@ -143,9 +198,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_mappings: Dict[str, str], user_api_key_dict: UserAPIKeyAuth, ) -> None: - verbose_logger.info( - f"Storing LiteLLM Managed File object with id={file_id} in cache" - ) + verbose_logger.info(f"Storing LiteLLM Managed File object with id={file_id} in cache") if file_object is not None: litellm_managed_file_object = LiteLLM_ManagedFileTable( unified_file_id=file_id, @@ -196,13 +249,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): f"storage_url={db_data.get('storage_url')}" ) - result = await self.prisma_client.db.litellm_managedfiletable.upsert( + result = await _managed_file_table(self.prisma_client).upsert( where={"unified_file_id": file_id}, data={"create": db_data, "update": update_data}, ) - verbose_logger.debug( - f"LiteLLM Managed File object with id={file_id} stored in db: {result}" - ) + verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}") async def store_unified_object_id( self, @@ -213,9 +264,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_purpose: Literal["batch", "fine-tune", "response"], user_api_key_dict: UserAPIKeyAuth, ) -> None: - verbose_logger.info( - f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache" - ) + verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache") litellm_managed_object = LiteLLM_ManagedObjectTable( unified_object_id=unified_object_id, model_object_id=model_object_id, @@ -228,7 +277,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): litellm_parent_otel_span=litellm_parent_otel_span, ) - await self.prisma_client.db.litellm_managedobjecttable.upsert( + await _managed_object_table(self.prisma_client).upsert( where={"unified_object_id": unified_object_id}, data={ "create": { @@ -265,9 +314,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return LiteLLM_ManagedFileTable.model_validate(result) ## CHECK DB - db_object = await self.prisma_client.db.litellm_managedfiletable.find_first( - where={"unified_file_id": file_id} - ) + db_object = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if db_object: return LiteLLM_ManagedFileTable.model_validate(db_object.model_dump()) @@ -277,9 +324,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): self, file_id: str, litellm_parent_otel_span: Optional[Span] = None ) -> OpenAIFileObject: ## get old value - initial_value = await self.prisma_client.db.litellm_managedfiletable.find_first( - where={"unified_file_id": file_id} - ) + initial_value = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if initial_value is None: raise Exception(f"LiteLLM Managed File object with id={file_id} not found") ## delete old value @@ -288,15 +333,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): value=None, litellm_parent_otel_span=litellm_parent_otel_span, ) - await self.prisma_client.db.litellm_managedfiletable.delete( - where={"unified_file_id": file_id} - ) + await _managed_file_table(self.prisma_client).delete(where={"unified_file_id": file_id}) return initial_value.file_object - async def can_user_call_unified_file_id( - self, unified_file_id: str, user_api_key_dict: UserAPIKeyAuth - ) -> bool: - managed_file = await self.prisma_client.db.litellm_managedfiletable.find_first( + async def can_user_call_unified_file_id(self, unified_file_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + managed_file = await _managed_file_table(self.prisma_client).find_first( where={"unified_file_id": unified_file_id} ) @@ -311,13 +352,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"File not found: {unified_file_id}", ) - async def can_user_call_unified_object_id( - self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth - ) -> bool: - managed_object = ( - await self.prisma_client.db.litellm_managedobjecttable.find_first( - where={"unified_object_id": unified_object_id} - ) + async def can_user_call_unified_object_id(self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + managed_object = await _managed_object_table(self.prisma_client).find_first( + where={"unified_object_id": unified_object_id} ) if managed_object: @@ -339,34 +376,28 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): provider: Optional[str] = None, target_model_names: Optional[str] = None, llm_router: Optional[Router] = None, - ) -> Dict[str, Any]: + ) -> Dict[str, object]: # Provider filtering is not supported for managed batches # This is because the encoded object ids stored in the managed objects table do not contain the provider information # To support provider filtering, we would need to store the provider information in the encoded object ids if provider: - raise Exception( - "Filtering by 'provider' is not supported when using managed batches." - ) + raise Exception("Filtering by 'provider' is not supported when using managed batches.") # Model name filtering is not supported for managed batches # This is because the encoded object ids stored in the managed objects table do not contain the model name # A hash of the model name + litellm_params for the model name is encoded as the model id. This is not sufficient to reliably map the target model names to the model ids. if target_model_names: - raise Exception( - "Filtering by 'target_model_names' is not supported when using managed batches." - ) + raise Exception("Filtering by 'target_model_names' is not supported when using managed batches.") owner_filter = build_owner_filter(user_api_key_dict) if owner_filter is None: return build_list_page([]) - where_clause: Dict[str, Any] = {"file_purpose": "batch", **owner_filter} + where_clause: Dict[str, object] = {"file_purpose": "batch", **owner_filter} if after: - cursor_row = ( - await self.prisma_client.db.litellm_managedobjecttable.find_first( - where={**where_clause, "unified_object_id": after} - ) + cursor_row = await _managed_object_table(self.prisma_client).find_first( + where={**where_clause, "unified_object_id": after} ) if cursor_row is None: raise HTTPException( @@ -375,11 +406,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) page_size: Final = min(limit or 20, 100) - cursor_args: Dict[str, Any] = ( - {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} - ) + cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} - batches = await self.prisma_client.db.litellm_managedobjecttable.find_many( + batches = await _managed_object_table(self.prisma_client).find_many( where=where_clause, take=page_size + 1, order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], @@ -389,9 +418,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): has_more = len(batches) > page_size parsed_rows: Final = tuple( - (row, batch_obj) - for row in batches[:page_size] - if (batch_obj := _parse_managed_batch_row(row)) is not None + (row, batch_obj) for row in batches[:page_size] if (batch_obj := _parse_managed_batch_row(row)) is not None ) unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified( raw_file_ids=frozenset( @@ -432,14 +459,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): verbose_proxy_logger=verbose_logger, user_api_key_dict=user_api_key_dict, db_batch_object=row, - unified_batch_id=_is_base64_encoded_unified_file_id( - row.unified_object_id - ), + unified_batch_id=_is_base64_encoded_unified_file_id(row.unified_object_id), ) except Exception as e: - verbose_logger.warning( - f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}" - ) + verbose_logger.warning(f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}") return None return batch_obj @@ -458,7 +481,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if owner_filter is None: return [] - file_ids = await self.prisma_client.db.litellm_managedfiletable.find_many( + file_ids = await _managed_file_table(self.prisma_client).find_many( where={ **owner_filter, "flat_model_file_ids": {"hasSome": model_object_ids}, @@ -467,27 +490,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return [ parsed_file_object.model_copy(update={"id": row.unified_file_id}) for row in file_ids - if ( - parsed_file_object := _parse_managed_file_object( - row.file_object, row.unified_file_id - ) - ) - is not None + if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None ] - async def check_managed_file_id_access( - self, data: Dict, user_api_key_dict: UserAPIKeyAuth - ) -> bool: + async def check_managed_file_id_access(self, data: Dict, user_api_key_dict: UserAPIKeyAuth) -> bool: retrieve_file_id = cast(Optional[str], data.get("file_id")) - potential_file_id = ( - _is_base64_encoded_unified_file_id(retrieve_file_id) - if retrieve_file_id - else False - ) + potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False if potential_file_id and retrieve_file_id: - if await self.can_user_call_unified_file_id( - retrieve_file_id, user_api_key_dict - ): + if await self.can_user_call_unified_file_id(retrieve_file_id, user_api_key_dict): return True else: raise HTTPException( @@ -496,9 +506,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) return False - async def check_file_ids_access( - self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth - ) -> None: + async def check_file_ids_access(self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth) -> None: """ Check if the user has access to a list of file IDs. Only checks managed (unified) file IDs. @@ -513,9 +521,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): for file_id in file_ids: is_unified_file_id = _is_base64_encoded_unified_file_id(file_id) if is_unified_file_id: - if not await self.can_user_call_unified_file_id( - file_id, user_api_key_dict - ): + if not await self.can_user_call_unified_file_id(file_id, user_api_key_dict): raise HTTPException( status_code=403, detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}", @@ -543,10 +549,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ### HANDLE TRANSFORMATIONS ### # Check both completion and acompletion call types - is_completion_call = ( - call_type == CallTypes.completion.value - or call_type == CallTypes.acompletion.value - ) + is_completion_call = call_type == CallTypes.completion.value or call_type == CallTypes.acompletion.value if is_completion_call: messages = data.get("messages") @@ -559,9 +562,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Check if any files are stored in storage backends and need base64 conversion # This is needed for Vertex AI/Gemini which requires base64 content - is_vertex_ai = model and ( - "vertex_ai" in model or "gemini" in model.lower() - ) + is_vertex_ai = model and ("vertex_ai" in model or "gemini" in model.lower()) if is_vertex_ai: await self._convert_storage_files_to_base64( messages=messages, @@ -573,10 +574,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids, user_api_key_dict.parent_otel_span ) data["model_file_id_mapping"] = model_file_id_mapping - elif ( - call_type == CallTypes.aresponses.value - or call_type == CallTypes.responses.value - ): + elif call_type == CallTypes.aresponses.value or call_type == CallTypes.responses.value: # Handle managed files in responses API input and tools file_ids = [] @@ -603,23 +601,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if tools: unified_vs_ids = self.get_vector_store_ids_from_file_search_tools(tools) if unified_vs_ids: - await self.check_vector_store_ids_access( - unified_vs_ids, user_api_key_dict - ) + await self.check_vector_store_ids_access(unified_vs_ids, user_api_key_dict) elif call_type == CallTypes.afile_content.value: retrieve_file_id = cast(Optional[str], data.get("file_id")) - potential_file_id = ( - _is_base64_encoded_unified_file_id(retrieve_file_id) - if retrieve_file_id - else False - ) + potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False if potential_file_id and "llm_output_file_id," in potential_file_id: model_id = self.get_model_id_from_unified_file_id(potential_file_id) if model_id: data["model"] = model_id - data["file_id"] = self.get_output_file_id_from_unified_file_id( - potential_file_id - ) + data["file_id"] = self.get_output_file_id_from_unified_file_id(potential_file_id) elif call_type == CallTypes.acreate_batch.value: input_file_id = cast(Optional[str], data.get("input_file_id")) if input_file_id: @@ -636,10 +626,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ): accessor_key: Optional[str] = None retrieve_object_id: Optional[str] = None - if ( - call_type == CallTypes.aretrieve_batch.value - or call_type == CallTypes.acancel_batch.value - ): + if call_type == CallTypes.aretrieve_batch.value or call_type == CallTypes.acancel_batch.value: accessor_key = "batch_id" elif ( call_type == CallTypes.acancel_fine_tuning_job.value @@ -651,32 +638,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): retrieve_object_id = cast(Optional[str], data.get(accessor_key)) potential_llm_object_id = ( - _is_base64_encoded_unified_file_id(retrieve_object_id) - if retrieve_object_id - else False + _is_base64_encoded_unified_file_id(retrieve_object_id) if retrieve_object_id else False ) if potential_llm_object_id and retrieve_object_id: ## VALIDATE USER HAS ACCESS TO THE OBJECT ## - if not await self.can_user_call_unified_object_id( - retrieve_object_id, user_api_key_dict - ): + if not await self.can_user_call_unified_object_id(retrieve_object_id, user_api_key_dict): raise HTTPException( status_code=403, detail=f"User {user_api_key_dict.user_id} does not have access to the object {retrieve_object_id}", ) ## for managed batch id - get the model id - potential_model_id = get_model_id_from_unified_batch_id( - potential_llm_object_id - ) + potential_model_id = get_model_id_from_unified_batch_id(potential_llm_object_id) if potential_model_id is None: raise Exception( f"LiteLLM Managed {accessor_key} with id={retrieve_object_id} is invalid - does not contain encoded model_id." ) data["model"] = potential_model_id - data[accessor_key] = get_batch_id_from_unified_batch_id( - potential_llm_object_id - ) + data[accessor_key] = get_batch_id_from_unified_batch_id(potential_llm_object_id) elif call_type == CallTypes.acreate_fine_tuning_job.value: input_file_id = cast(Optional[str], data.get("training_file")) if input_file_id: @@ -732,24 +711,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if accessor_key: input_file_id = cast(Optional[str], kwargs.get(accessor_key)) - model_file_id_mapping = cast( - Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping") - ) + model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping")) # model_info may be at top-level or nested under litellm_metadata # (batch/file operations use litellm_metadata) model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None)) if model_id is None: model_id = cast( Optional[str], - kwargs.get("litellm_metadata", {}) - .get("model_info", {}) - .get("id", None), + kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None), ) mapped_file_id: Optional[str] = None if input_file_id and model_file_id_mapping and model_id: - mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get( - model_id, None - ) + mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(model_id, None) if mapped_file_id: kwargs[accessor_key] = mapped_file_id @@ -775,9 +748,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids.append(file_id) return file_ids - def get_file_ids_from_responses_input( - self, input: Union[str, List[Dict[str, Any]]] - ) -> List[str]: + def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, Any]]]) -> List[str]: """ Gets file ids from responses API input. @@ -809,19 +780,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): content = item.get("content") if isinstance(content, list): for content_item in content: - if ( - isinstance(content_item, dict) - and content_item.get("type") == "input_file" - ): + if isinstance(content_item, dict) and content_item.get("type") == "input_file": file_id = content_item.get("file_id") if file_id: file_ids.append(file_id) return file_ids - def get_file_ids_from_responses_tools( - self, tools: List[Dict[str, Any]] - ) -> List[str]: + def get_file_ids_from_responses_tools(self, tools: List[Dict[str, object]]) -> List[str]: """ Gets file ids from responses API tools parameter. @@ -854,9 +820,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return file_ids - def get_vector_store_ids_from_file_search_tools( - self, tools: List[Dict[str, Any]] - ) -> List[str]: + def get_vector_store_ids_from_file_search_tools(self, tools: List[Dict[str, object]]) -> List[str]: """ Extract unified vector_store_ids from file_search tools. @@ -949,9 +913,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ), ) - async def get_model_file_id_mapping( - self, file_ids: List[str], litellm_parent_otel_span: Span - ) -> dict: + async def get_model_file_id_mapping(self, file_ids: List[str], litellm_parent_otel_span: Span) -> dict: """ Get model-specific file IDs for a list of proxy file IDs. Returns a dictionary mapping litellm_proxy/ file_id -> model_id -> model_file_id @@ -981,9 +943,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Get all cache keys matching the pattern file_id:* for file_id in litellm_managed_file_ids: # Search for any cache key starting with this file_id - unified_file_object = await self.get_unified_file_id( - file_id, litellm_parent_otel_span - ) + unified_file_object = await self.get_unified_file_id(file_id, litellm_parent_otel_span) if unified_file_object: file_id_mapping[file_id] = unified_file_object.model_mappings @@ -1001,9 +961,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): raise Exception("LLM Router not initialized. Ensure models added to proxy.") responses = [] for model in target_model_names_list: - individual_response = await llm_router.acreate_file( - model=model, **_create_file_request - ) + individual_response = await llm_router.acreate_file(model=model, **_create_file_request) responses.append(individual_response) return responses @@ -1034,9 +992,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_mappings: Dict[str, str] = {} for file_object in responses: - model_file_id_mapping = file_object._hidden_params.get( - "model_file_id_mapping" - ) + model_file_id_mapping = file_object._hidden_params.get("model_file_id_mapping") if model_file_id_mapping and isinstance(model_file_id_mapping, dict): model_mappings.update(model_file_id_mapping) @@ -1051,17 +1007,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Emit Prometheus metrics for managed file creation prom_logger = self._get_prometheus_logger() if prom_logger: - first_model = ( - target_model_names_list[0] if target_model_names_list else None - ) + first_model = target_model_names_list[0] if target_model_names_list else None first_provider = "" if responses: - first_provider = ( - getattr(responses[0], "_hidden_params", {}).get( - "custom_llm_provider" - ) - or "" - ) + first_provider = getattr(responses[0], "_hidden_params", {}).get("custom_llm_provider") or "" prom_logger.record_managed_file_created( model=first_model or "", api_provider=first_provider, @@ -1104,9 +1053,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) # Convert to URL-safe base64 and strip padding - base64_unified_file_id = ( - base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") - ) + base64_unified_file_id = base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") ## CREATE RESPONSE OBJECT @@ -1123,46 +1070,26 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return response - def get_unified_generic_response_id( - self, model_id: str, generic_response_id: str - ) -> str: - unified_generic_response_id = ( - SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format( - model_id, generic_response_id - ) - ) - return ( - base64.urlsafe_b64encode(unified_generic_response_id.encode()) - .decode() - .rstrip("=") + def get_unified_generic_response_id(self, model_id: str, generic_response_id: str) -> str: + unified_generic_response_id = SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format( + model_id, generic_response_id ) + return base64.urlsafe_b64encode(unified_generic_response_id.encode()).decode().rstrip("=") def get_unified_batch_id(self, batch_id: str, model_id: str) -> str: - unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format( - model_id, batch_id - ) + 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: - 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", - str(deterministic_uuid), - model_name or "", - output_file_id, - model_id, - ) - ) - return ( - base64.urlsafe_b64encode(unified_output_file_id.encode()) - .decode() - .rstrip("=") + def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: Optional[str]) -> 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", + str(deterministic_uuid), + model_name or "", + output_file_id, + model_id, ) + return base64.urlsafe_b64encode(unified_output_file_id.encode()).decode().rstrip("=") def get_model_id_from_unified_file_id(self, file_id: str) -> str: return file_id.split("llm_output_file_model_id,")[1].split(";")[0] @@ -1170,59 +1097,39 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): def get_output_file_id_from_unified_file_id(self, file_id: str) -> str: marker = "llm_output_file_id," if marker not in file_id: - raise ValueError( - f"Unified id does not contain {marker!r}: {file_id[:80]!r}" - ) + raise ValueError(f"Unified id does not contain {marker!r}: {file_id[:80]!r}") return file_id.split(marker, 1)[1].split(";")[0] async def async_post_call_success_hook( self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes - ) -> Any: + ) -> LLMResponseTypes: if isinstance(response, LiteLLMBatch): ## Check if unified_file_id is in the response - unified_file_id = response._hidden_params.get( - "unified_file_id" - ) # managed file id - unified_batch_id = response._hidden_params.get( - "unified_batch_id" - ) # managed batch id + unified_file_id = response._hidden_params.get("unified_file_id") # managed file id + unified_batch_id = response._hidden_params.get("unified_batch_id") # managed batch id model_id = cast(Optional[str], response._hidden_params.get("model_id")) model_name = cast(Optional[str], response._hidden_params.get("model_name")) resolved_model_name = resolve_managed_output_file_model_name( - unified_input_file_id=unified_file_id - if isinstance(unified_file_id, str) - else response.input_file_id, + unified_input_file_id=unified_file_id if isinstance(unified_file_id, str) else response.input_file_id, fallback_model_name=model_name, ) original_response_id = response.id if (unified_batch_id or unified_file_id) and model_id: - response.id = self.get_unified_batch_id( - batch_id=response.id, model_id=model_id - ) + response.id = self.get_unified_batch_id(batch_id=response.id, model_id=model_id) # Handle both output_file_id and error_file_id for file_attr in ["output_file_id", "error_file_id"]: file_id_value = getattr(response, file_attr, None) if file_id_value and model_id: - decoded_output_file_id = _is_base64_encoded_unified_file_id( - file_id_value - ) - if ( - decoded_output_file_id - and "llm_output_file_id," in decoded_output_file_id - ): - provider_file_id = ( - self.get_output_file_id_from_unified_file_id( - decoded_output_file_id - ) - ) + decoded_output_file_id = _is_base64_encoded_unified_file_id(file_id_value) + if decoded_output_file_id and "llm_output_file_id," in decoded_output_file_id: + provider_file_id = self.get_output_file_id_from_unified_file_id(decoded_output_file_id) unified_file_id = file_id_value elif decoded_output_file_id: verbose_logger.warning( - f"Skipping {file_attr}={file_id_value!r}: " - "unified id is not a managed file output id" + f"Skipping {file_attr}={file_id_value!r}: unified id is not a managed file output id" ) continue else: @@ -1241,23 +1148,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Import module and use getattr for better testability with mocks import litellm.proxy.proxy_server as proxy_server_module - _llm_router = getattr( - proxy_server_module, "llm_router", None - ) + _llm_router = getattr(proxy_server_module, "llm_router", None) if _llm_router is not None and model_id: - _creds = ( - _llm_router.get_deployment_credentials_with_provider( - model_id - ) - or {} - ) + _creds = _llm_router.get_deployment_credentials_with_provider(model_id) or {} file_object = await litellm.afile_retrieve( file_id=provider_file_id, **_creds, ) else: file_object = await litellm.afile_retrieve( - custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", # type: ignore[arg-type] + custom_llm_provider=model_name.split("/")[0] + if model_name and "/" in model_name + else "openai", # type: ignore[arg-type] file_id=provider_file_id, ) verbose_logger.debug( @@ -1311,9 +1213,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): elif isinstance(response, LiteLLMFineTuningJob): ## Check if unified_file_id is in the response - unified_file_id = response._hidden_params.get( - "unified_file_id" - ) # managed file id + unified_file_id = response._hidden_params.get("unified_file_id") # managed file id unified_finetuning_job_id = response._hidden_params.get( "unified_finetuning_job_id" ) # managed finetuning job id @@ -1321,9 +1221,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_name = cast(Optional[str], response._hidden_params.get("model_name")) original_response_id = response.id if (unified_file_id or unified_finetuning_job_id) and model_id: - response.id = self.get_unified_generic_response_id( - model_id=model_id, generic_response_id=response.id - ) + response.id = self.get_unified_generic_response_id(model_id=model_id, generic_response_id=response.id) await self.store_unified_object_id( unified_object_id=response.id, file_object=response, @@ -1338,9 +1236,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): """ ## check if file object if hasattr(response, "data") and isinstance(response.data, list): - if all( - isinstance(file_object, FileObject) for file_object in response.data - ): + if all(isinstance(file_object, FileObject) for file_object in response.data): ## Get all file id's ## Check which file id's were created by the user ## Filter the response to only include the files created by the user @@ -1349,9 +1245,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_object.id for file_object in cast(List[FileObject], response.data) # type: ignore ] - user_created_file_ids = await self.get_user_created_file_ids( - user_api_key_dict, file_ids - ) + user_created_file_ids = await self.get_user_created_file_ids(user_api_key_dict, file_ids) ## Filter the response to only include the files created by the user response.data = user_created_file_ids # type: ignore return response @@ -1359,11 +1253,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return response async def afile_retrieve( - self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router=None + self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router: Optional[Router] = None ) -> OpenAIFileObject: - stored_file_object = await self.get_unified_file_id( - file_id, litellm_parent_otel_span - ) + stored_file_object = await self.get_unified_file_id(file_id, litellm_parent_otel_span) # Case 1 : This is not a managed file if not stored_file_object: @@ -1386,21 +1278,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) try: - model_id, model_file_id = next( - iter(stored_file_object.model_mappings.items()) - ) - credentials = ( - llm_router.get_deployment_credentials_with_provider(model_id) or {} - ) - response = await litellm.afile_retrieve( - file_id=model_file_id, **credentials - ) + model_id, model_file_id = next(iter(stored_file_object.model_mappings.items())) + credentials = llm_router.get_deployment_credentials_with_provider(model_id) or {} + response = await litellm.afile_retrieve(file_id=model_file_id, **credentials) response.id = file_id # Replace with unified ID return response except Exception as e: - raise Exception( - f"Failed to retrieve file {file_id} from provider: {str(e)}" - ) from e + raise Exception(f"Failed to retrieve file {file_id} from provider: {str(e)}") from e async def afile_list( self, @@ -1437,12 +1321,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return False except Exception as e: - verbose_logger.warning( - f"Error checking batch polling configuration: {e}. Assuming disabled." - ) + verbose_logger.warning(f"Error checking batch polling configuration: {e}. Assuming disabled.") return False - async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, Any]]: + async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, object]]: """ Find batches that reference this file and still need cost tracking. Find batches that are in non-terminal state and have not yet been processed by CheckBatchCost. @@ -1458,9 +1340,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Get model-specific file IDs for this unified file ID if it's a managed file try: - model_file_id_mapping = await self.get_model_file_id_mapping( - [file_id], litellm_parent_otel_span=None - ) + model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span=None) if model_file_id_mapping and file_id in model_file_id_mapping: # Add all provider file IDs for this unified file @@ -1468,8 +1348,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids_to_check.extend(provider_file_ids) except Exception as e: verbose_logger.debug( - f"Could not get model file ID mapping for {file_id}: {e}. " - f"Will only check unified file ID." + f"Could not get model file ID mapping for {file_id}: {e}. Will only check unified file ID." ) MAX_MATCHES_TO_RETURN = 10 @@ -1487,11 +1366,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): for batch in batches: try: # Parse the batch file_object to check for file references - batch_data = ( - json.loads(batch.file_object) - if isinstance(batch.file_object, str) - else batch.file_object - ) + batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object # Extract file IDs from batch # Batches typically reference the unified file ID in input_file_id @@ -1500,9 +1375,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): output_file_id = batch_data.get("output_file_id") error_file_id = batch_data.get("error_file_id") - referenced_file_ids = [ - fid for fid in [input_file_id, output_file_id, error_file_id] if fid - ] + referenced_file_ids = [fid for fid in [input_file_id, output_file_id, error_file_id] if fid] # Check if any referenced file ID matches the file we're trying to delete if any(ref_id in file_ids_to_check for ref_id in referenced_file_ids): @@ -1514,9 +1387,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): } ) except Exception as e: - verbose_logger.warning( - f"Error parsing batch object {batch.unified_object_id}: {e}" - ) + verbose_logger.warning(f"Error parsing batch object {batch.unified_object_id}: {e}") continue return referencing_batches @@ -1545,21 +1416,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if referencing_batches: # File is referenced by non-terminal batches and polling is enabled - MAX_BATCHES_IN_ERROR = ( - 5 # Limit batches shown in error message for readability - ) + MAX_BATCHES_IN_ERROR = 5 # Limit batches shown in error message for readability # Show up to MAX_BATCHES_IN_ERROR in the error message batches_to_show = referencing_batches[:MAX_BATCHES_IN_ERROR] - batch_statuses = [ - f"{b['batch_id']}: {b['status']}" for b in batches_to_show - ] + batch_statuses = [f"{b['batch_id']}: {b['status']}" for b in batches_to_show] # Determine the count message count_message = f"{len(referencing_batches)}" - if ( - len(referencing_batches) >= 10 - ): # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file + if len(referencing_batches) >= 10: # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file count_message = "10+" error_message = ( @@ -1600,23 +1465,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): await self._check_file_deletion_allowed(file_id) # file_id = convert_b64_uid_to_unified_uid(file_id) - model_file_id_mapping = await self.get_model_file_id_mapping( - [file_id], litellm_parent_otel_span - ) + model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span) delete_response = None specific_model_file_id_mapping = model_file_id_mapping.get(file_id) if specific_model_file_id_mapping: # Remove conflicting keys from data to avoid duplicate keyword arguments - filtered_data = { - k: v for k, v in data.items() if k not in ("model", "file_id") - } + filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")} for model_id, model_file_id in specific_model_file_id_mapping.items(): delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore - stored_file_object = await self.delete_unified_file_id( - file_id, litellm_parent_otel_span - ) + stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span) # Record successful deletion metric only on actual success if stored_file_object or delete_response: @@ -1643,9 +1502,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): Get the content of a file from first model that has it """ model_file_id_mapping = data.pop("model_file_id_mapping", None) - model_file_id_mapping = ( - model_file_id_mapping - or await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span) + model_file_id_mapping = model_file_id_mapping or await self.get_model_file_id_mapping( + [file_id], litellm_parent_otel_span ) specific_model_file_id_mapping = model_file_id_mapping.get(file_id) @@ -1658,13 +1516,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # against the deployment's configured bucket, which they only # trust from this immutable server-side snapshot, never from # request params. - credentials = llm_router.get_deployment_credentials_with_provider( - model_id=model_id - ) + credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id) if credentials is not None: - data["_litellm_internal_model_credentials"] = cast( - Dict, MappingProxyType(dict(credentials)) - ) + data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials))) else: data.pop("_litellm_internal_model_credentials", None) return await llm_router.afile_content(model=model_id, file_id=provider_file_id, **data) # type: ignore @@ -1699,9 +1553,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Check database for storage backend info # IMPORTANT: The database stores the base64 encoded unified_file_id (not the decoded version) # So we query with the original file_id (which is base64 encoded) - db_file = await self.prisma_client.db.litellm_managedfiletable.find_first( - where={"unified_file_id": file_id} - ) + db_file = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if not db_file or not db_file.storage_backend or not db_file.storage_url: continue @@ -1727,22 +1579,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_content = await storage_backend.download_file(storage_url) # Determine content type from file object - content_type = self._get_content_type_from_file_object( - db_file.file_object - ) + content_type = self._get_content_type_from_file_object(db_file.file_object) # Convert to base64 base64_data = base64.b64encode(file_content).decode("utf-8") base64_data_uri = f"data:{content_type};base64,{base64_data}" # Update messages to use base64 instead of file_id - self._update_messages_with_base64_data( - messages, file_id, base64_data_uri, content_type - ) + self._update_messages_with_base64_data(messages, file_id, base64_data_uri, content_type) except Exception as e: - verbose_logger.exception( - f"Error converting file {file_id} from storage backend to base64: {str(e)}" - ) + verbose_logger.exception(f"Error converting file {file_id} from storage backend to base64: {str(e)}") # Continue with other files even if one fails continue diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 339da998d56..024e8c179c2 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -6,16 +6,33 @@ This module provides fake streaming by converting non-streaming responses into s """ import asyncio -from collections.abc import AsyncIterator -from typing import Any, Final, cast +from collections.abc import AsyncIterator, Mapping, Sequence +from typing import Any, Final, Protocol, cast, runtime_checkable from uuid import uuid4 +from pydantic import TypeAdapter + from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) +_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_TEXT_ADAPTER: Final = TypeAdapter(str) + + +@runtime_checkable +class _SupportsModelDump(Protocol): + def model_dump(self, *, mode: str, exclude_none: bool) -> Mapping[str, object]: ... + + +@runtime_checkable +class _SupportsPydanticDict(Protocol): + def dict(self, *, exclude_none: bool) -> Mapping[str, object]: ... + class PydanticAITransformation: """ @@ -28,7 +45,7 @@ class PydanticAITransformation: """ @staticmethod - def _remove_none_values(obj: Any) -> Any: + def _remove_none_values(obj: object) -> object: """ Recursively remove None values from a dict/list structure. @@ -42,14 +59,18 @@ class PydanticAITransformation: Cleaned object with None values removed """ if isinstance(obj, dict): - return {k: PydanticAITransformation._remove_none_values(v) for k, v in obj.items() if v is not None} + typed_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python(obj) + return {k: PydanticAITransformation._remove_none_values(v) for k, v in typed_dict.items() if v is not None} elif isinstance(obj, list): - return [PydanticAITransformation._remove_none_values(item) for item in obj if item is not None] + typed_list: Final = _LIST_ADAPTER.validate_python(obj) + return [PydanticAITransformation._remove_none_values(item) for item in typed_list if item is not None] else: return obj @staticmethod - def _params_to_dict(params: Any) -> dict[str, Any]: + def _params_to_dict( + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", + ) -> Mapping[str, object]: """ Convert params to a dict, handling Pydantic models. @@ -59,10 +80,10 @@ class PydanticAITransformation: Returns: Dict representation of params """ - if hasattr(params, "model_dump"): + if isinstance(params, _SupportsModelDump): # Pydantic v2 model return params.model_dump(mode="python", exclude_none=True) - elif hasattr(params, "dict"): + elif isinstance(params, _SupportsPydanticDict): # Pydantic v1 model return params.dict(exclude_none=True) elif isinstance(params, dict): @@ -75,12 +96,12 @@ class PydanticAITransformation: async def _poll_for_completion( client: AsyncHTTPHandler, endpoint: str, - task_id: str, + task_id: object, request_id: str, max_attempts: int = 30, poll_interval: float = 0.5, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Poll for task completion using tasks/get method. @@ -112,10 +133,10 @@ class PydanticAITransformation: }, ) response.raise_for_status() - poll_data = response.json() + poll_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json()) - result = poll_data.get("result", {}) - status = result.get("status", {}) + result = _STR_KEY_DICT_ADAPTER.validate_python(poll_data.get("result", {})) + status = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {})) state = status.get("state", "") verbose_logger.debug("Pydantic AI: Poll attempt %s/%s, state=%s", attempt + 1, max_attempts, state) @@ -133,10 +154,10 @@ class PydanticAITransformation: async def _send_and_poll_raw( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -153,14 +174,16 @@ class PydanticAITransformation: Raw Pydantic AI task response (with history/artifacts) """ # Convert params to dict if it's a Pydantic model - params_dict = PydanticAITransformation._params_to_dict(params) - # Remove None values - FastA2A doesn't accept null for optional fields - params_dict = PydanticAITransformation._remove_none_values(params_dict) + params_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python( + PydanticAITransformation._remove_none_values(PydanticAITransformation._params_to_dict(params)) + ) # Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI if "message" in params_dict: - params_dict["message"]["kind"] = "message" + message_value: Final = _ANY_KEY_DICT_ADAPTER.validate_python(params_dict["message"]) + message_value["kind"] = "message" + params_dict["message"] = message_value # Build A2A JSON-RPC request using message/send method for FastA2A compatibility a2a_request: Final = { @@ -189,11 +212,11 @@ class PydanticAITransformation: }, ) response.raise_for_status() - response_data = response.json() + response_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json()) # Check if task is already completed - result: Final = response_data.get("result", {}) - status: Final = result.get("status", {}) + result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {})) + status: Final = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {})) state: Final = status.get("state", "") if state != "completed": @@ -217,10 +240,10 @@ class PydanticAITransformation: async def send_non_streaming_request( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a non-streaming A2A request to Pydantic AI agent and wait for completion. @@ -253,10 +276,10 @@ class PydanticAITransformation: async def send_and_get_raw_response( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -282,9 +305,9 @@ class PydanticAITransformation: @staticmethod def _transform_to_a2a_response( - response_data: dict[str, Any], + response_data: Mapping[str, object], request_id: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform Pydantic AI task response to standard A2A non-streaming format. @@ -328,7 +351,7 @@ class PydanticAITransformation: } @staticmethod - def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]: + def _extract_response_text(response_data: Mapping[str, object]) -> tuple[object, object, Sequence[object]]: """ Extract response text from completed task response. @@ -342,52 +365,53 @@ class PydanticAITransformation: Returns: Tuple of (full_text, message_id, parts) """ - result: Final = response_data.get("result", {}) + result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {})) # Try to extract from artifacts first (preferred for results) artifacts: Final = result.get("artifacts", []) if artifacts: - for artifact in artifacts: - parts = artifact.get("parts", []) + for artifact in _LIST_ADAPTER.validate_python(artifacts): + parts = _LIST_ADAPTER.validate_python(_STR_KEY_DICT_ADAPTER.validate_python(artifact).get("parts", [])) for part in parts: - if part.get("kind") == "text": - text = part.get("text", "") + if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text": + text = part_dict.get("text", "") if text: return text, str(uuid4()), parts # Fall back to history - get the last agent message - history: Final = result.get("history", []) + history: Final = _LIST_ADAPTER.validate_python(result.get("history", [])) for msg in reversed(history): - if msg.get("role") == "agent": - parts = msg.get("parts", []) - message_id = msg.get("messageId", str(uuid4())) + if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "agent": + parts = _LIST_ADAPTER.validate_python(msg_dict.get("parts", [])) + message_id = msg_dict.get("messageId", str(uuid4())) full_text = "" for part in parts: - if part.get("kind") == "text": - full_text += part.get("text", "") + if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text": + full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", "")) if full_text: return full_text, message_id, parts # Fall back to message field (original format) message: Final = result.get("message", {}) if message: - parts = message.get("parts", []) - message_id = message.get("messageId", str(uuid4())) + message_dict: Final = _STR_KEY_DICT_ADAPTER.validate_python(message) + parts = _LIST_ADAPTER.validate_python(message_dict.get("parts", [])) + message_id = message_dict.get("messageId", str(uuid4())) full_text = "" for part in parts: - if part.get("kind") == "text": - full_text += part.get("text", "") + if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text": + full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", "")) return full_text, message_id, parts return "", str(uuid4()), [] @staticmethod async def fake_streaming_from_response( - response_data: dict[str, Any], + response_data: Mapping[str, object], request_id: str, chunk_size: int = 50, delay_ms: int = 10, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """ Convert a non-streaming A2A response into fake streaming chunks. @@ -410,12 +434,12 @@ class PydanticAITransformation: full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data) # Extract input message from raw response for history - result: Final = response_data.get("result", {}) - history: Final = result.get("history", []) - input_message = {} + result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {})) + history: Final = _LIST_ADAPTER.validate_python(result.get("history", [])) + input_message = _STR_KEY_DICT_ADAPTER.validate_python({}) for msg in history: - if msg.get("role") == "user": - input_message = msg + if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "user": + input_message = msg_dict break # Generate IDs for streaming events @@ -426,45 +450,49 @@ class PydanticAITransformation: # 1. Emit initial task event (kind: "task", status: "submitted") # Format matches A2ACompletionBridgeTransformation.create_task_event - task_event: Final = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "history": [ - { - "contextId": context_id, - "kind": "message", - "messageId": input_message_id, - "parts": input_message.get("parts", [{"kind": "text", "text": ""}]), - "role": "user", - "taskId": task_id, - } - ], - "id": task_id, - "kind": "task", - "status": { - "state": "submitted", + task_event: Final = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "history": [ + { + "contextId": context_id, + "kind": "message", + "messageId": input_message_id, + "parts": input_message.get("parts", [{"kind": "text", "text": ""}]), + "role": "user", + "taskId": task_id, + } + ], + "id": task_id, + "kind": "task", + "status": { + "state": "submitted", + }, }, - }, - } + } + ) yield task_event # 2. Emit status update (kind: "status-update", status: "working") # Format matches A2ACompletionBridgeTransformation.create_status_update_event - working_event: Final = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "final": False, - "kind": "status-update", - "status": { - "state": "working", + working_event: Final = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "final": False, + "kind": "status-update", + "status": { + "state": "working", + }, + "taskId": task_id, }, - "taskId": task_id, - }, - } + } + ) yield working_event # Small delay to simulate processing @@ -473,29 +501,32 @@ class PydanticAITransformation: # 3. Emit artifact update chunks (kind: "artifact-update") # Format matches A2ACompletionBridgeTransformation.create_artifact_update_event if full_text: + full_text_str: Final = _TEXT_ADAPTER.validate_python(full_text) # Split text into chunks - for i in range(0, len(full_text), chunk_size): - chunk_text = full_text[i : i + chunk_size] - is_last_chunk = (i + chunk_size) >= len(full_text) + for i in range(0, len(full_text_str), chunk_size): + chunk_text = full_text_str[i : i + chunk_size] + is_last_chunk = (i + chunk_size) >= len(full_text_str) - artifact_event = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "kind": "artifact-update", - "taskId": task_id, - "artifact": { - "artifactId": artifact_id, - "parts": [ - { - "kind": "text", - "text": chunk_text, - } - ], + artifact_event = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "kind": "artifact-update", + "taskId": task_id, + "artifact": { + "artifactId": artifact_id, + "parts": [ + { + "kind": "text", + "text": chunk_text, + } + ], + }, }, - }, - } + } + ) yield artifact_event # Add delay between chunks (except for last chunk) @@ -503,19 +534,21 @@ class PydanticAITransformation: await asyncio.sleep(delay_ms / 1000.0) # 4. Emit final status update (kind: "status-update", status: "completed", final: true) - completed_event: Final = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "final": True, - "kind": "status-update", - "status": { - "state": "completed", + completed_event: Final = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "final": True, + "kind": "status-update", + "status": { + "state": "completed", + }, + "taskId": task_id, }, - "taskId": task_id, - }, - } + } + ) yield completed_event verbose_logger.info("Pydantic AI: Fake streaming completed for request_id=%s", request_id) diff --git a/litellm/constants.py b/litellm/constants.py index 30d3bb1f26e..3b91f23fe39 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1322,6 +1322,7 @@ X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks" LITELLM_METADATA_FIELD: Final = "litellm_metadata" OLD_LITELLM_METADATA_FIELD: Final = "metadata" RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name" +SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index e273289168a..2c83406afed 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -33,6 +33,7 @@ from litellm.integrations.otel.model.payloads import ( is_mcp_list_tools, is_mcp_tool_call, ) +from litellm.integrations.otel.model.semconv import Error from litellm.integrations.otel.model.spans import SpanRole, span_role_for_service from litellm.integrations.otel.model.utils import to_ns from litellm.integrations.otel.plumbing.context import ( @@ -634,18 +635,23 @@ class OpenTelemetryV2(CustomLogger): """Stamp the v2 error.* attributes on the FastAPI-owned SERVER span for a failure that dies before any LLM-call span exists (malformed body, auth / validation rejection). Called from the proxy's global exception handler via - ``_close_dangling_otel_server_span``. The instrumentor still owns the span's - status and lifecycle, so this only decorates it — never sets status, never - ends it — and emits no exception event, matching v1's SERVER-span behavior - and avoiding a duplicate of the event ``async_post_call_failure_hook`` or - the ``auth`` phase span already records.""" + ``_close_dangling_otel_server_span``, which swallows the exception into a + ``JSONResponse`` so the instrumentor never sees it and leaves the span + ``UNSET``; the status is set here instead (v1 did the same from the handler) + so a failed request reads as failed and not merely as a span carrying an + error message. The instrumentor still owns the span's lifecycle, so this + never ends it. The exception event is recorded only when nothing stamped + this span already — ``async_post_call_failure_hook`` and the ``auth`` phase + span record their own, and a second event would duplicate it — while the + attributes are always restamped so ``error.code`` stays pinned to the real + response status.""" if span is None or not is_recordable_span(span): return + already_stamped: Final = Error.TYPE in (getattr(span, "attributes", None) or ()) stamp_error( span, _span_error_from_exception(exception, status_code=status_code), - record_event=False, - set_status=False, + record_event=not already_stamped, ) async def async_post_call_failure_hook( diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 9748db2dcd2..2c5f7484dac 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -42,7 +42,12 @@ from litellm.types.integrations.websearch_interception import ( WebSearchInterceptionConfig, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import CallTypes, LlmProviders +from litellm.types.utils import ( + AgenticLoopParams, + CallTypes, + LlmProviders, + StandardLoggingUserAPIKeyMetadata, +) from litellm.utils import ProviderConfigManager if TYPE_CHECKING: @@ -265,7 +270,7 @@ class WebSearchInterceptionLogger(CustomLogger): return None # Check if request has tools with native web_search - tools: Final = kwargs.get("tools") + tools: Final[Sequence[dict[str, object]] | None] = kwargs.get("tools") if not tools: return None @@ -314,7 +319,9 @@ class WebSearchInterceptionLogger(CustomLogger): return kwargs - def _convert_responses_tools(self, kwargs: Mapping[str, object], tools: list[dict[str, object]]) -> dict | None: + def _convert_responses_tools( + self, kwargs: Mapping[str, object], tools: Sequence[dict[str, object]] + ) -> dict[str, object] | None: """Convert Responses API web search tools to the LiteLLM standard function tool.""" if not any(is_web_search_tool_responses(tool) for tool in tools): return None @@ -379,7 +386,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _tool_name(tool: dict[str, Any]) -> str | None: + def _tool_name(tool: Mapping[str, object]) -> object: """Effective tool name, handling OpenAI ``function`` wrapper shape.""" fn: Final = tool.get("function") if tool.get("type") == "function" and isinstance(fn, dict): @@ -1271,7 +1278,7 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs_for_followup: Final = self._prepare_followup_kwargs(kwargs) if logging_obj is not None: - agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {}) + agentic_params: Final[AgenticLoopParams] = logging_obj.model_call_details.get("agentic_loop_params", {}) full_model_name = agentic_params.get("model", model) verbose_logger.debug( "WebSearchInterception: Built anthropic request patch [call_id=%s model=%s messages=%d searches=%d]", @@ -1316,6 +1323,7 @@ class WebSearchInterceptionLogger(CustomLogger): search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router) search_provider: str | None = None search_litellm_params: dict[str, Any] = {} + search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool) if search_tool is not None: await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs) search_litellm_params = dict(search_tool.get("litellm_params", {}) or {}) @@ -1332,12 +1340,30 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.debug( "WebSearchInterception: Executing search for '%s' using provider '%s'", query, search_provider ) + user_api_key_auth: Final = self._get_user_api_key_auth_from_kwargs(kwargs) + search_metadata: Final = ( + None + if user_api_key_auth is None + else self._build_search_request_metadata( + user_api_key_auth=user_api_key_auth, + search_tool_name=search_tool_name, + ) + ) search_kwargs: Final = { key: value for key, value in search_litellm_params.items() if key != "search_provider" and value is not None } - result: Final = await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs) + result: Final = ( + await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs) + if search_metadata is None + else await litellm.asearch( + query=query, + search_provider=search_provider, + litellm_metadata=search_metadata, + **search_kwargs, + ) + ) # Format using transformation function search_result_text: Final = WebSearchTransformation.format_search_response(result) @@ -1394,6 +1420,35 @@ class WebSearchInterceptionLogger(CustomLogger): team_object=team_object, ) + @staticmethod + def _build_search_request_metadata( + user_api_key_auth: "UserAPIKeyAuth", + search_tool_name: str | None, + ) -> Mapping[str, object]: + """ + Spend-tracking metadata for the intercepted search, so its provider cost is logged + and billed against the key/user/team that made the originating LLM request instead + of being dropped by the proxy's spend hook for lack of an owner. + """ + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + user_api_key_metadata: Final[StandardLoggingUserAPIKeyMetadata] = ( + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth) + ) + return { # mutable-ok: litellm's metadata channel is a plain dict its logging path reads and enriches + **user_api_key_metadata, + "model_group": search_tool_name, + "user_api_key": user_api_key_auth.api_key, + "user_api_key_auth": user_api_key_auth, + } + + @staticmethod + def _selected_search_tool_name(search_tool: Mapping[str, object] | None) -> str | None: + if search_tool is None: + return None + search_tool_name: Final = search_tool.get("search_tool_name") + return search_tool_name if isinstance(search_tool_name, str) and search_tool_name else None + @staticmethod def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object] | None) -> "UserAPIKeyAuth | None": if not kwargs: diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 858d10df53b..d68bdc4a250 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,5 +1,6 @@ import asyncio import json +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm @@ -20,11 +21,24 @@ from .litellm_logging import Logging as LiteLLMLogging if TYPE_CHECKING: from websockets.asyncio.client import ClientConnection + from litellm.types.guardrails import GuardrailEventHooks + CLIENT_CONNECTION_CLASS = ClientConnection else: CLIENT_CONNECTION_CLASS = Any +class _ClientWebSocketExceptions(Protocol): + ConnectionClosed: type[Exception] + + +class _ClientWebSocket(Protocol): + exceptions: _ClientWebSocketExceptions + + async def send_text(self, data: str) -> None: ... + async def receive_text(self) -> str: ... + + class RealtimeEventNormalizer(Protocol): def should_drop(self, event: object) -> bool: ... def normalize(self, event: dict) -> dict: ... @@ -48,13 +62,13 @@ class RealTimeStreaming: logging_obj: LiteLLMLogging, provider_config: BaseRealtimeConfig | None = None, model: str = "", - user_api_key_dict: Any | None = None, + user_api_key_dict: object | None = None, request_data: dict | None = None, backend_uses_beta_protocol: bool | None = None, force_transcription_model: str | None = None, event_normalizer: RealtimeEventNormalizer | None = None, ): - self.websocket = websocket + self.websocket: _ClientWebSocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj self.messages: list[OpenAIRealtimeEvents] = [] @@ -127,7 +141,7 @@ class RealTimeStreaming: ] ) _CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(["input_audio_buffer.commit", "input_audio_buffer.end"]) - _AUDIO_FORMAT_MAP: dict[str, dict[str, Any]] = { + _AUDIO_FORMAT_MAP: dict[str, dict[str, str | int]] = { "pcm16": {"type": "audio/pcm", "rate": 24000}, "g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000}, "g711_alaw": {"type": "audio/G711-alaw", "rate": 8000}, @@ -281,6 +295,7 @@ class RealTimeStreaming: if event_obj.get("type") != "response.done": return response: Final = cast(dict[str, Any], event_obj.get("response", {})) + item: Mapping[str, object] for item in response.get("output", []): if item.get("type") == "function_call": self.tool_calls.append( @@ -384,7 +399,7 @@ class RealTimeStreaming: return message try: - message_obj: Final = json.loads(message) + message_obj: Final[Mapping[str, object]] = json.loads(message) except (json.JSONDecodeError, TypeError): return message @@ -487,7 +502,7 @@ class RealTimeStreaming: if self._backend_setup_complete and not self._flushing_pending_messages_until_setup: return False try: - msg_obj: Final = json.loads(message) + msg_obj: Final[Mapping[str, object]] = json.loads(message) except (json.JSONDecodeError, TypeError): return False return msg_obj.get("type") in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES @@ -555,7 +570,7 @@ class RealTimeStreaming: def _event_to_client_json(self, event: dict) -> str: return json.dumps(self._normalize_event_for_ga_client(event)) - async def _send_event_to_client(self, event: Any, event_str: str) -> bool: + async def _send_event_to_client(self, event: object, event_str: str) -> bool: if self._should_drop_event_from_client(event): return False if isinstance(event, dict): @@ -595,12 +610,12 @@ class RealTimeStreaming: def _make_disable_auto_response_message(self) -> str: """Return a session.update that disables VAD auto-response.""" - turn_detection: Final[dict[str, Any]] = { + turn_detection: Final[dict[str, str | bool]] = { "type": "server_vad", "create_response": False, } if self._backend_uses_beta_protocol: - session: dict[str, Any] = {"turn_detection": turn_detection} + session: dict[str, object] = {"turn_detection": turn_detection} else: session = { "type": "realtime", @@ -654,7 +669,7 @@ class RealTimeStreaming: def _has_realtime_guardrails_for_event_hooks( self, - event_hooks: list[Any], + event_hooks: Sequence["GuardrailEventHooks"], ) -> bool: """Return True if any callback would run for one of ``event_hooks``.""" from litellm.integrations.custom_guardrail import CustomGuardrail @@ -699,7 +714,7 @@ class RealTimeStreaming: transcript: str, item_id: str | None = None, pre_block_backend_message: str | None = None, - event_hooks: list[Any] | None = None, + event_hooks: Sequence["GuardrailEventHooks"] | None = None, ) -> bool: """ Run registered guardrails on realtime text (transcript, user message, tool output). @@ -753,7 +768,7 @@ class RealTimeStreaming: raise # Extract the human-readable error from the detail dict (HTTPException) # or fall back to str(e) for plain ValueError. - detail = getattr(e, "detail", None) + detail: object | None = getattr(e, "detail", None) if isinstance(detail, dict): safe_msg = detail.get("error") or str(e) elif detail is not None: @@ -826,7 +841,7 @@ class RealTimeStreaming: return True return False - async def _handle_provider_config_message(self, raw_response) -> None: + async def _handle_provider_config_message(self, raw_response: str) -> None: """Process a backend message when a provider_config is set (transformed path).""" returned_object: Final = self.provider_config.transform_realtime_response( raw_response, @@ -910,7 +925,7 @@ class RealTimeStreaming: await self._send_event_to_client(event, event_str) @staticmethod - def _parse_backend_event(raw_response: str) -> dict | None: + def _parse_backend_event(raw_response: str) -> dict[str, object] | None: """Parse a backend frame once. Returns None for non-JSON or non-object frames.""" try: event: Final = json.loads(raw_response) @@ -1020,7 +1035,7 @@ class RealTimeStreaming: objects and any test doubles that expose a .scope dict. """ try: - headers: Final = websocket.scope.get("headers", []) + headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = websocket.scope.get("headers", []) for name, value in headers: if isinstance(name, bytes): name = name.decode("latin-1") @@ -1071,9 +1086,9 @@ class RealTimeStreaming: session["output_modalities"] = ["text"] # 3-7. Lift flat audio fields into the nested audio object - audio: Final[dict[str, Any]] = {} - inp: Final[dict[str, Any]] = {} - out: Final[dict[str, Any]] = {} + audio: Final[dict[str, object]] = {} + inp: Final[dict[str, object]] = {} + out: Final[dict[str, object]] = {} # voice → audio.output.voice if "voice" in session: @@ -1190,7 +1205,7 @@ class RealTimeStreaming: # model; check them with the same guardrail used for # user text so an attacker cannot smuggle blocked # content into a function_call_output. - output = item.get("output", "") + output: object = item.get("output", "") output_text = output if isinstance(output, str) else json.dumps(output) if output_text: # Build the sanitized function_call_output up @@ -1241,7 +1256,7 @@ class RealTimeStreaming: # interaction turn. continue elif item.get("role") == "user": - content_list = item.get("content", []) + content_list: Sequence[object] = item.get("content", []) texts = [ c.get("text", "") for c in content_list @@ -1280,7 +1295,7 @@ class RealTimeStreaming: and not self._guardrail_turn_detection_update_sent and self._has_audio_transcription_guardrails() ): - session = msg_obj.setdefault("session", {}) + session: object = msg_obj.setdefault("session", {}) if isinstance(session, dict): existing_td = session.get("turn_detection") if not isinstance(existing_td, dict): diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 3f967e29002..886ba6a3a18 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -3,7 +3,7 @@ import time from collections.abc import Iterator, Mapping, Sequence from itertools import groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing import TYPE_CHECKING, Any, Final, TypedDict, Union, cast from litellm._logging import verbose_logger from litellm.types.llms.openai import ( @@ -30,6 +30,7 @@ from litellm.types.utils import ( from litellm.utils import print_verbose, token_counter if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, ) @@ -39,6 +40,60 @@ if TYPE_CHECKING: ) +class _ThinkingBlockFragment(TypedDict, total=False): + type: str | None + data: str | None + thinking: str | None + signature: str | None + + +class _ThinkingDelta(TypedDict, total=False): + thinking_blocks: Sequence[_ThinkingBlockFragment] + + +class _ThinkingChoice(TypedDict, total=False): + delta: _ThinkingDelta + + +class _ThinkingChunk(TypedDict): + choices: Sequence[_ThinkingChoice] + + +class _ContentChoice(TypedDict, total=False): + delta: Mapping[str, str | None] + + +class _ContentChunk(TypedDict): + choices: Sequence[_ContentChoice] + + +class _AudioDelta(TypedDict, total=False): + audio: ChatCompletionAudioDelta | None + + +class _AudioChoice(TypedDict, total=False): + delta: _AudioDelta + + +class _AudioChunk(TypedDict): + choices: Sequence[_AudioChoice] + + +class _UsageBearingChunk(TypedDict, total=False): + usage: Usage | None + _hidden_params: Mapping[str, str] + + +class _UsageSummary(TypedDict): + prompt_tokens: int | None + completion_tokens: int | None + cache_creation_input_tokens: int | None + cache_read_input_tokens: int | None + completion_tokens_details: CompletionTokensDetails | None + prompt_tokens_details: PromptTokensDetailsWrapper | None + cost: float | None + + def capture_cache_creation_token_details( prompt_tokens_details: PromptTokensDetailsWrapper | None, current: CacheCreationTokenDetails | None, @@ -78,7 +133,7 @@ class ChunkProcessor: return [] first_chunk: Final = chunks[0] - first_hidden_params: dict[str, Any] = {} + first_hidden_params: dict[str, object] = {} if isinstance(first_chunk, dict): candidate = first_chunk.get("_hidden_params", {}) if isinstance(candidate, dict): @@ -115,8 +170,8 @@ class ChunkProcessor: @staticmethod def apply_provider_assembled_streaming_metadata( response: ModelResponse, - chunks: list[Any], - logging_obj: Any | None = None, + chunks: list[object], + logging_obj: "Logging | None" = None, ) -> None: if not chunks: return @@ -456,7 +511,7 @@ class ChunkProcessor: ) def get_combined_content( - self, chunks: list[dict[str, Any]], delta_key: str = "content" + self, chunks: Sequence["_ContentChunk"], delta_key: str = "content" ) -> ChatCompletionAssistantContentValue: content_list: Final[list[str]] = [] for chunk in chunks: @@ -475,7 +530,7 @@ class ChunkProcessor: return combined_content def get_combined_thinking_content( - self, chunks: list[dict[str, Any]] + self, chunks: Sequence["_ThinkingChunk"] ) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None: from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, @@ -532,10 +587,10 @@ class ChunkProcessor: return thinking_blocks return None - def get_combined_reasoning_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAssistantContentValue: + def get_combined_reasoning_content(self, chunks: Sequence["_ContentChunk"]) -> ChatCompletionAssistantContentValue: return self.get_combined_content(chunks, delta_key="reasoning_content") - def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAudioResponse: + def get_combined_audio_content(self, chunks: Sequence["_AudioChunk"]) -> ChatCompletionAudioResponse: base64_data_list: Final[list[str]] = [] transcript_list: Final[list[str]] = [] expires_at: int | None = None @@ -544,7 +599,7 @@ class ChunkProcessor: for chunk in chunks: choices = chunk["choices"] for choice in choices: - delta = choice.get("delta") or {} + delta: _AudioDelta = choice.get("delta") or {} audio: ChatCompletionAudioDelta | None = delta.get("audio") if audio is not None: for k, v in audio.items(): @@ -565,7 +620,7 @@ class ChunkProcessor: id=id, ) - def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> dict: + def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> "_UsageSummary": prompt_tokens = 0 completion_tokens = 0 ## anthropic prompt caching information ## @@ -623,8 +678,8 @@ class ChunkProcessor: return reasoning_tokens @staticmethod - def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None: - usage_chunk: Usage | dict[str, Any] | None = None + def _extract_usage_chunk(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Usage | None: + usage_chunk: Usage | None = None if hasattr(chunk, "usage") and chunk.usage is not None: usage_chunk = chunk.usage elif "usage" in chunk: @@ -640,7 +695,7 @@ class ChunkProcessor: def _calculate_usage_per_chunk( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence["_UsageBearingChunk | ModelResponse"], ) -> "UsagePerChunk": from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, @@ -721,13 +776,7 @@ class ChunkProcessor: "web_search_requests", ) - prompt_tokens_details = ( - cast( - PromptTokensDetailsWrapper | None, - usage_chunk_dict["prompt_tokens_details"], - ) - or prompt_tokens_details - ) + prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] or prompt_tokens_details cache_creation_token_details = capture_cache_creation_token_details( prompt_tokens_details, cache_creation_token_details @@ -758,7 +807,7 @@ class ChunkProcessor: @staticmethod def _reset_anthropic_cursor_completion_tokens( - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence["_UsageBearingChunk | ModelResponse"], completion_tokens: int, completion_usage_updates: int, ) -> int: @@ -797,7 +846,7 @@ class ChunkProcessor: def calculate_usage( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence["_UsageBearingChunk | ModelResponse"], model: str, completion_output: str, messages: list | None = None, @@ -851,8 +900,8 @@ class ChunkProcessor: setattr(returned_usage, "cache_read_input_tokens", cache_read_input_tokens) # for anthropic if completion_tokens_details is not None: if isinstance(completion_tokens_details, CompletionTokensDetails): - returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( - **completion_tokens_details.model_dump() + returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper.model_validate( + completion_tokens_details.model_dump() ) else: returned_usage.completion_tokens_details = completion_tokens_details diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index a9751489473..36f3e875a7e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -1,8 +1,9 @@ -from collections.abc import AsyncIterator, Coroutine, Iterator +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping from typing import ( TYPE_CHECKING, Any, Final, + TypeAlias, cast, ) @@ -33,8 +34,12 @@ if TYPE_CHECKING: # Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge. ANTHROPIC_ONLY_REQUEST_KEYS: Final[frozenset[str]] = frozenset({"output_config"}) +_AnthropicMessages: TypeAlias = "list[dict[str, object]]" +_AnthropicSystem: TypeAlias = "str | list[dict[str, object]] | None" +_ContextManagementSpec: TypeAlias = "dict[str, object] | list[dict[str, object]] | None" -def _messages_have_compaction_block(messages: list[dict]) -> bool: + +def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool: """Return True when any message carries a ``compaction`` content block.""" for msg in messages: content = msg.get("content") @@ -54,8 +59,10 @@ def _proxy_router_fallback() -> "Router | None": return _proxy_router -def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None: - """Return ``kwargs["litellm_metadata"]`` when it's a dict; ``None`` otherwise. +def _extract_proxy_litellm_metadata( + kwargs: Mapping[str, object], +) -> "tuple[dict[str, object], UserAPIKeyAuth | None] | tuple[None, None]": + """Return ``(kwargs["litellm_metadata"], its user_api_key_auth)`` when it's a dict; ``(None, None)`` otherwise. The proxy attaches its auth/spend-attribution fields (``user_api_key``, ``user_api_key_team_id``, ``litellm_call_id``, the full ``UserAPIKeyAuth`` @@ -68,18 +75,19 @@ def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | """ litellm_metadata: Final = kwargs.get("litellm_metadata") if not isinstance(litellm_metadata, dict): - return None - return litellm_metadata + return None, None + user_api_key_auth: Final[UserAPIKeyAuth | None] = litellm_metadata.get("user_api_key_auth") + return litellm_metadata, user_api_key_auth async def _prepare_context_managed_request( *, model: str, - messages: list[dict], - tools: list[dict] | None, - system: Any | None, - context_management_spec: Any, - litellm_metadata: dict | None, + messages: _AnthropicMessages, + tools: list[dict[str, object]] | None, + system: _AnthropicSystem, + context_management_spec: _ContextManagementSpec, + litellm_metadata: dict[str, object] | None, additional_drop_params: list[str] | None, llm_router: "Router | None", user_api_key_auth: "UserAPIKeyAuth | None" = None, @@ -102,11 +110,11 @@ async def _prepare_context_managed_request( if polyfill_will_run: history_result: PolyfillResult | None = None - working_messages: list[dict] = messages - working_system: Any | None = system + working_messages: _AnthropicMessages = messages + working_system: _AnthropicSystem = system else: history_result = apply_client_compaction_block_history( - messages=cast(list[dict[str, Any]], messages), + messages=messages, system=system, ) working_messages = history_result.messages if history_result is not None else messages @@ -136,7 +144,7 @@ async def _prepare_context_managed_request( # to non-Anthropic backends that would reject them. if polyfill_will_run and history_result is None: history_result = apply_client_compaction_block_history( - messages=cast(list[dict[str, Any]], messages), + messages=messages, system=system, ) return history_result @@ -144,7 +152,7 @@ async def _prepare_context_managed_request( def _polyfill_will_run( *, - context_management_spec: Any, + context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, ) -> bool: """Return True when ``compact_20260112`` will run via the polyfill dispatcher. @@ -171,7 +179,7 @@ def _polyfill_will_run( def _spec_has_non_compact_edits( *, - context_management_spec: Any, + context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, ) -> bool: """Return True when the spec includes edits other than ``compact_20260112``. @@ -209,9 +217,9 @@ def _context_management_explicitly_dropped(additional_drop_params: list[str] | N def _normalize_spec_edits( *, - context_management_spec: Any, + context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, -) -> list[dict[str, Any]] | None: +) -> list[dict[str, object]] | None: """Return the normalized ``edits`` list, or ``None`` if the polyfill won't run. Delegates spec-shape normalization to the dispatcher's ``_normalize_spec`` @@ -236,11 +244,11 @@ def _normalize_spec_edits( async def _run_polyfill_if_enabled( *, model: str, - messages: list[dict], - tools: list[dict] | None, - system: Any | None, - context_management_spec: Any, - litellm_metadata: dict | None, + messages: _AnthropicMessages, + tools: list[dict[str, object]] | None, + system: _AnthropicSystem, + context_management_spec: _ContextManagementSpec, + litellm_metadata: dict[str, object] | None, additional_drop_params: list[str] | None, llm_router: "Router | None", user_api_key_auth: "UserAPIKeyAuth | None" = None, @@ -306,7 +314,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: def _route_openai_thinking_to_responses_api_if_needed( completion_kwargs: dict[str, Any], *, - thinking: dict[str, Any] | None, + thinking: Mapping[str, object] | None, ) -> None: """ When users call `litellm.anthropic.messages.*` with a non-Anthropic model and @@ -407,12 +415,12 @@ class LiteLLMMessagesToCompletionTransformationHandler: def _prepare_completion_kwargs( *, max_tokens: int, - messages: list[dict], + messages: _AnthropicMessages, model: str, metadata: dict | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, - system: str | list[dict[str, Any]] | None = None, + system: _AnthropicSystem = None, temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, @@ -420,7 +428,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, - extra_kwargs: dict[str, Any] | None = None, + extra_kwargs: Mapping[str, object] | None = None, ) -> tuple[dict[str, Any], dict[str, str]]: """Prepare kwargs for litellm.completion/acompletion. @@ -433,7 +441,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: Logging as LiteLLMLoggingObject, ) - request_data: Final = { + request_data: Final[dict[str, object]] = { "model": model, "messages": messages, "max_tokens": max_tokens, @@ -528,7 +536,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod async def async_anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: _AnthropicMessages, model: str, metadata: dict | None = None, stop_sequences: list[str] | None = None, @@ -537,7 +545,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, - tools: list[dict] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, @@ -551,10 +559,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: requested_router if requested_router is not None else _proxy_router_fallback() ) - proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs) - user_api_key_auth: Final[UserAPIKeyAuth | None] = ( - proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None - ) + proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) polyfill_result: Final = await _prepare_context_managed_request( model=model, @@ -618,7 +623,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod def anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: _AnthropicMessages, model: str, metadata: dict | None = None, stop_sequences: list[str] | None = None, @@ -627,7 +632,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, - tools: list[dict] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, @@ -688,10 +693,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if context_management is None and not _messages_have_compaction_block(messages): polyfill_result: PolyfillResult | None = None else: - proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs) - user_api_key_auth: Final[UserAPIKeyAuth | None] = ( - proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None - ) + proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) polyfill_result = run_async_function( _prepare_context_managed_request, model=model, diff --git a/litellm/llms/azure/assistants.py b/litellm/llms/azure/assistants.py index 671e4633af4..f7b419405ac 100644 --- a/litellm/llms/azure/assistants.py +++ b/litellm/llms/azure/assistants.py @@ -1,8 +1,9 @@ from collections.abc import Coroutine, Iterable -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypedDict import httpx from openai import AsyncAzureOpenAI, AzureOpenAI +from openai.types.shared_params.metadata import Metadata from typing_extensions import overload from ...types.llms.openai import ( @@ -22,6 +23,16 @@ from ...types.llms.openai import ( from .common_utils import BaseAzureLLM +class _RunThreadStreamData(TypedDict): + thread_id: str + assistant_id: str + additional_instructions: str | None + instructions: str | None + metadata: Metadata | None + model: str | None + tools: Iterable[AssistantToolParam] | None + + class AzureAssistantsAPI(BaseAzureLLM): def __init__(self) -> None: super().__init__() @@ -212,9 +223,9 @@ class AzureAssistantsAPI(BaseAzureLLM): response_obj: OpenAIMessage | None = None if getattr(thread_message, "status", None) is None: thread_message.status = "completed" - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) else: - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) return response_obj # fmt: off @@ -301,9 +312,9 @@ class AzureAssistantsAPI(BaseAzureLLM): response_obj: OpenAIMessage | None = None if getattr(thread_message, "status", None) is None: thread_message.status = "completed" - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) else: - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) return response_obj async def async_get_messages( @@ -443,7 +454,7 @@ class AzureAssistantsAPI(BaseAzureLLM): message_thread: Final = await openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread.model_validate(message_thread.dict()) # fmt: off @@ -539,7 +550,7 @@ class AzureAssistantsAPI(BaseAzureLLM): message_thread: Final = azure_openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread.model_validate(message_thread.dict()) async def async_get_thread( self, @@ -566,7 +577,7 @@ class AzureAssistantsAPI(BaseAzureLLM): response: Final = await openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread.model_validate(response.dict()) # fmt: off @@ -642,7 +653,7 @@ class AzureAssistantsAPI(BaseAzureLLM): response: Final = openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread.model_validate(response.dict()) # def delete_thread(self): # pass @@ -730,7 +741,8 @@ class AzureAssistantsAPI(BaseAzureLLM): event_handler: AssistantEventHandler | None, litellm_params: dict | None = None, ) -> AssistantStreamManager[AssistantEventHandler]: - data: Final[dict[str, Any]] = { + stream_fn: Final = client.beta.threads.runs.stream + base_data: Final[_RunThreadStreamData] = { "thread_id": thread_id, "assistant_id": assistant_id, "additional_instructions": additional_instructions, @@ -740,8 +752,8 @@ class AzureAssistantsAPI(BaseAzureLLM): "tools": tools, } if event_handler is not None: - data["event_handler"] = event_handler - return client.beta.threads.runs.stream(**data) + return stream_fn(**base_data, event_handler=event_handler) + return stream_fn(**base_data) # fmt: off diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 76618e0f742..e285feb77ee 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -109,7 +109,7 @@ if MCP_AVAILABLE: ############ MCP Server REST API Routes ################# async def _safe_fire_mcp_tool_call_logging( logging_obj: Any | None, - result: Any, + result: "CallToolResult", start_time: datetime, end_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 1c6ad84ddb4..49a1f1314f0 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -13,9 +13,9 @@ import time import traceback import types import uuid -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol import httpx from fastapi import FastAPI, HTTPException @@ -145,7 +145,7 @@ try: ) # Robust auth lookup keyed by session_object. - _session_obj_auth_storage: "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() + _session_obj_auth_storage: "weakref.WeakKeyDictionary[object, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() except ImportError as e: verbose_logger.debug("MCP module not found: %s", e) MCP_AVAILABLE = False @@ -493,14 +493,14 @@ if MCP_AVAILABLE: def _gateway_create_initialization_options( self, notification_options: NotificationOptions | None = None, - experimental_capabilities: dict[str, dict[str, Any]] | None = None, + experimental_capabilities: dict[str, dict[str, object]] | None = None, ) -> InitializationOptions: opts: Final = Server.create_initialization_options( self, notification_options=notification_options, experimental_capabilities=experimental_capabilities or {}, ) - updates: Final[dict[str, Any]] = {} + updates: Final[dict[str, str]] = {} merged: Final = _mcp_gateway_initialize_instructions.get() if merged is not None: updates["instructions"] = merged @@ -549,6 +549,17 @@ if MCP_AVAILABLE: _stateful_session_locks: Final[dict[str, asyncio.Lock]] = {} _stateful_session_active_request_counts: Final[dict[str, int]] = {} + class _TerminableTransport(Protocol): + async def terminate(self) -> None: ... + + class _TransportRegistry(Protocol): + def __contains__(self, session_id: object, /) -> bool: ... + + def pop(self, session_id: str, default: None, /) -> "_TerminableTransport | None": ... + + def _stateful_server_instances() -> _TransportRegistry: + return getattr(session_manager_stateful, "_server_instances", {}) + def _remove_stateful_session_tracking(session_id: str) -> None: _stateful_session_auth_contexts.pop(session_id, None) _stateful_session_auth_context_last_seen.pop(session_id, None) @@ -578,8 +589,8 @@ if MCP_AVAILABLE: ) -> None: """Terminate expired stateful sessions and drop their auth contexts.""" now = time.monotonic() if now is None else now - server_instances: Final = getattr(session_manager_stateful, "_server_instances", {}) - expired_session_ids: Final = [] + server_instances: Final = _stateful_server_instances() + expired_session_ids: Final[list[str]] = [] for session_id, last_seen in _stateful_session_auth_context_last_seen.items(): if _stateful_session_active_request_counts.get(session_id, 0) > 0: continue @@ -619,7 +630,7 @@ if MCP_AVAILABLE: session may proceed, or ``False`` when the caller is already at the cap with every session in flight (the new ``initialize`` should be rejected). """ - server_instances: Final = getattr(session_manager_stateful, "_server_instances", {}) + server_instances: Final = _stateful_server_instances() def _owned_live_session_ids() -> list[str]: return [ @@ -778,7 +789,7 @@ if MCP_AVAILABLE: get_virtual_tool_definitions, ) - return [Tool(**d) for d in get_virtual_tool_definitions()] + return [Tool.model_validate(d) for d in get_virtual_tool_definitions()] # Get mcp_servers from context variable verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools") @@ -847,7 +858,7 @@ if MCP_AVAILABLE: async def _build_virtual_call_logging_obj( name: str, - arguments: dict[str, Any], + arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth, ) -> LiteLLMLoggingObj | None: """Run the pre-call pipeline (guardrails + logging setup) for a virtual @@ -885,7 +896,7 @@ if MCP_AVAILABLE: async def _dispatch_virtual_mcp_tool( name: str, - arguments: dict[str, Any] | None, + arguments: dict[str, object] | None, user_api_key_auth: UserAPIKeyAuth | None, client_ip: str | None, mcp_servers: list[str] | None = None, @@ -957,7 +968,7 @@ if MCP_AVAILABLE: ) @server.call_tool() - async def mcp_server_tool_call(name: str, arguments: dict[str, Any] | None) -> CallToolResult: + async def mcp_server_tool_call(name: str, arguments: dict[str, object] | None) -> CallToolResult: """ Call a specific tool with the provided arguments Args: @@ -1621,7 +1632,7 @@ if MCP_AVAILABLE: async def _get_user_oauth_extra_headers_from_db( server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None, - prefetched_creds: dict[str, dict[str, Any]] | None = None, + prefetched_creds: 'Mapping[str, "OAuthCredentialPayload"] | None' = None, ) -> dict[str, str] | None: """Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None. @@ -1646,7 +1657,7 @@ if MCP_AVAILABLE: Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. """ - user_id: Final = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None + user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None if not user_id: return {} try: @@ -1871,7 +1882,7 @@ if MCP_AVAILABLE: list_tools_start_time: Final = datetime.now() litellm_logging_obj: LiteLLMLoggingObj | None = None - list_tools_request_data: dict[str, Any] = {} + list_tools_request_data: dict[str, object] = {} if log_list_tools_to_spendlogs: # This is intentionally minimal: only async_success_handler / post_call_failure_hook @@ -1879,7 +1890,7 @@ if MCP_AVAILABLE: list_tools_call_id: Final = str(uuid.uuid4()) # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool) effective_litellm_trace_id: Final = litellm_trace_id or get_chain_id_from_headers(raw_headers) - spend_logs_metadata: Final[dict[str, Any]] = { + spend_logs_metadata: Final[dict[str, object]] = { "mcp_operation": "list_tools", } if isinstance(list_tools_log_source, str): @@ -2615,7 +2626,7 @@ if MCP_AVAILABLE: async def execute_mcp_tool( name: str, - arguments: dict[str, Any], + arguments: dict[str, object], allowed_mcp_servers: list[MCPServer], start_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, @@ -2882,7 +2893,7 @@ if MCP_AVAILABLE: _request_auth_header.reset(_auth_token) _request_extra_headers.reset(_extra_token) _request_resolved_auth_headers.reset(_resolved_token) - response = CallToolResult(content=cast(Any, local_content), isError=False) + response = CallToolResult(content=local_content, isError=False) # Try managed MCP server tool (the name is bare; the prefix boundary was # already resolved above against this server's registered prefixes) @@ -2956,7 +2967,7 @@ if MCP_AVAILABLE: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args local_content = await _handle_local_mcp_tool(original_tool_name, arguments) - response = CallToolResult(content=cast(Any, local_content), isError=False) + response = CallToolResult(content=local_content, isError=False) return await _run_post_mcp_call_guardrails( result=response, @@ -3003,7 +3014,7 @@ if MCP_AVAILABLE: async def _fire_mcp_tool_call_logging( logging_obj: LiteLLMLoggingObj, - result: Any, + result: CallToolResult, start_time: datetime, end_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, @@ -3070,7 +3081,7 @@ if MCP_AVAILABLE: @client async def call_mcp_tool( name: str, - arguments: dict[str, Any] | None = None, + arguments: dict[str, object] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_servers: list[str] | None = None, @@ -3161,7 +3172,7 @@ if MCP_AVAILABLE: async def mcp_get_prompt( name: str, - arguments: dict[str, Any] | None = None, + arguments: dict[str, object] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_servers: list[str] | None = None, @@ -3262,7 +3273,7 @@ if MCP_AVAILABLE: def _get_standard_logging_mcp_tool_call( name: str, - arguments: dict[str, Any], + arguments: dict[str, object], server_name: str | None, session_id: str | None = None, ) -> StandardLoggingMCPToolCall: @@ -3291,13 +3302,13 @@ if MCP_AVAILABLE: async def _handle_managed_mcp_tool( server_name: str, name: str, - arguments: dict[str, Any], + arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, - litellm_logging_obj: Any | None = None, + litellm_logging_obj: LiteLLMLoggingObj | None = None, host_progress_callback: Callable | None = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" @@ -3320,7 +3331,7 @@ if MCP_AVAILABLE: return call_tool_result async def _handle_local_mcp_tool( - name: str, arguments: dict[str, Any] + name: str, arguments: dict[str, object] ) -> list[TextContent | ImageContent | EmbeddedResource]: """ Handle tool execution for local registry tools @@ -3426,7 +3437,8 @@ if MCP_AVAILABLE: Extract mcp-session-id from ASGI scope headers. Returns None if not present. """ - for header_name, header_value in scope.get("headers", []): + scope_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) + for header_name, header_value in scope_headers: name = header_name if isinstance(header_name, bytes) else header_name.encode() if name.lower() == b"mcp-session-id": return header_value.decode() if isinstance(header_value, bytes) else str(header_value) @@ -3528,7 +3540,7 @@ if MCP_AVAILABLE: if message.get("type") != "http.request": break - body = message.get("body", b"") or b"" + body: bytes = message.get("body", b"") or b"" if body: # Only retain up to the remaining peek budget for sniffing. # The full ``message`` is already in memory (delivered by @@ -3571,9 +3583,9 @@ if MCP_AVAILABLE: Fixes https://github.com/BerriAI/litellm/issues/20992 """ _mcp_session_header: Final = b"mcp-session-id" - _headers: Final = scope.get("headers", []) + _headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) - def _normalize_header_name(header_name: Any) -> bytes | None: + def _normalize_header_name(header_name: object) -> bytes | None: if isinstance(header_name, bytes): return header_name.lower() if isinstance(header_name, str): @@ -3902,7 +3914,8 @@ if MCP_AVAILABLE: def _get_authorization_header_from_scope(scope: Scope) -> str | None: """First ``Authorization`` header value in the ASGI scope, or None.""" - for key, value in scope.get("headers", []): + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + for key, value in scope_headers: if key.lower() == b"authorization": return value.decode("latin-1") return None @@ -3921,7 +3934,8 @@ if MCP_AVAILABLE: ``MCPRequestHandler.process_mcp_request``), and forwarding it upstream would leak the proxy key to a third-party MCP server. """ - has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope.get("headers", [])) + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope_headers) if not has_litellm_key_header: return None return _get_authorization_header_from_scope(scope) @@ -4115,7 +4129,7 @@ if MCP_AVAILABLE: async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through StreamableHTTP.""" try: - path: Final = scope.get("path", "") + path: Final[str] = scope.get("path", "") ( user_api_key_auth, mcp_auth_header, @@ -4135,7 +4149,8 @@ if MCP_AVAILABLE: ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. - scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"] + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"] # Apply toolset scope if set server-side via ContextVar (set by # /toolset/{name}/mcp and /{name}/mcp route handlers in proxy_server.py). @@ -4436,7 +4451,7 @@ if MCP_AVAILABLE: async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through SSE.""" try: - path: Final = scope.get("path", "") + path: Final[str] = scope.get("path", "") ( user_api_key_auth, mcp_auth_header, @@ -4456,7 +4471,8 @@ if MCP_AVAILABLE: ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. - scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"] + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"] # Apply toolset scope if set server-side via ContextVar so the # downstream probe list matches the fully-authorized server set @@ -4680,7 +4696,8 @@ if MCP_AVAILABLE: ) -> Send: async def wrapped_send(message: Message) -> None: if message.get("type") == "http.response.start": - for key, value in message.get("headers", []): + response_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = message.get("headers", []) + for key, value in response_headers: header_name = key if isinstance(key, bytes) else str(key).encode() if header_name.lower() == b"mcp-session-id": session_id = value.decode() if isinstance(value, bytes) else str(value) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1fc05ac4653..fa89df39c5f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4557,10 +4557,10 @@ class DefaultInternalUserParams(LiteLLMPydanticObjectBase): user_role: ( Literal[ - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, ] | None ) = Field( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9dc450befc2..4baa7b99a4f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1044,6 +1044,22 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None: request.state.parent_otel_span = parent_otel_span +async def _read_request_body_deferring_parse_failure( + request: Request, +) -> tuple[dict, ProxyException | None]: + """Parse the body, returning a parse failure instead of raising it. + + A body that fails to parse is still a request from a known caller, so auth + must run (resolving identity onto the request's trace) before the 400 goes + out; the caller re-raises the returned exception once identity is seeded. + """ + try: + parsed_body: Final = await _read_request_body(request=request) + except ProxyException as parse_exception: + return {}, parse_exception # mutable-ok: request_data is a plain dict across the whole auth path + return populate_request_with_path_params(request_data=parsed_body, request=request), None + + async def _user_api_key_auth_builder( request: Request, api_key: str, @@ -2516,6 +2532,72 @@ def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) -> ) +async def _authorize_authenticated_request( + user_api_key_auth_obj: UserAPIKeyAuth, + request: Request, + request_data: dict, + route: str, + api_key: str, +) -> UserAPIKeyAuth | None: + """Authorize an already-authenticated request: disabled-route check, the single + ``common_checks`` gate (which also reserves budget), and end-user fallback + resolution. Returns the auth object the exception handler recovered when a check + failed but the request may proceed anyway, else ``None``. + """ + ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## + RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request) + + # Single authorization point. Builder paths MUST NOT call common_checks. + # Route through the same exception handler the builder uses so + # authorization failures (ProxyException, or plain Exception from + # admin-only-route / model-access / budget checks) surface as + # ProxyException consistently with pre-refactor behavior. + try: + await _run_centralized_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request=request, + request_data=request_data, + route=route, + ) + except Exception as e: + return await UserAPIKeyAuthExceptionHandler._handle_authentication_error( + e=e, + request=request, + request_data=request_data, + route=route, + parent_otel_span=user_api_key_auth_obj.parent_otel_span, + api_key=api_key, + resolved_identity=user_api_key_auth_obj, + ) + + # Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return + # paths (no master key, /user/auth route, JWT short-circuits) that bypass + # the end-user resolution block. If those paths produced an auth obj + # without an ``end_user_id`` set, fall back to extracting from the request + # body so spend logs are still attributed correctly. Validation honours + # ``litellm.validate_end_user_id_in_db``. + if user_api_key_auth_obj.end_user_id is None: + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request)) + if raw_end_user_id is not None: + resolved_end_user_id: Final = await resolve_and_validate_end_user_id( + raw_end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth_obj.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if resolved_end_user_id is not None: + user_api_key_auth_obj.end_user_id = resolved_end_user_id + return None + + @tracer.wrap() async def user_api_key_auth( request: Request, @@ -2536,8 +2618,7 @@ async def user_api_key_auth( # close, and the trace never reaches the backend. _ensure_parent_otel_span_on_request_state(request) - request_data = await _read_request_body(request=request) - request_data = populate_request_with_path_params(request_data=request_data, request=request) + request_data, body_parse_exception = await _read_request_body_deferring_parse_failure(request=request) route: Final[str] = get_request_route(request=request) ## CHECK IF ROUTE IS ALLOWED @@ -2545,69 +2626,41 @@ async def user_api_key_auth( # triggers (key/user/team object reads) nest under it instead of flattening # onto the server span. No-op when OTel V2 isn't active. with phase_span(f"auth {route}"): - user_api_key_auth_obj: Final = await _user_api_key_auth_builder( - request=request, - api_key=api_key, - azure_api_key_header=azure_api_key_header, - anthropic_api_key_header=anthropic_api_key_header, - google_ai_studio_api_key_header=google_ai_studio_api_key_header, - azure_apim_header=azure_apim_header, - request_data=request_data, - custom_litellm_key_header=custom_litellm_key_header, - ) + try: + user_api_key_auth_obj: Final = await _user_api_key_auth_builder( + request=request, + api_key=api_key, + azure_api_key_header=azure_api_key_header, + anthropic_api_key_header=anthropic_api_key_header, + google_ai_studio_api_key_header=google_ai_studio_api_key_header, + azure_apim_header=azure_apim_header, + request_data=request_data, + custom_litellm_key_header=custom_litellm_key_header, + ) + except Exception: + # The body was read first, so a caller who sent both a malformed body and + # a rejected key used to get the 400; the response is unchanged, and the + # auth failure is still recorded on the trace by the handler that ran. + if body_parse_exception is not None: + raise body_parse_exception + raise user_api_key_auth_obj.budget_reservation = None - ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## - RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request) - - # Single authorization point. Builder paths MUST NOT call common_checks. - # Route through the same exception handler the builder uses so - # authorization failures (ProxyException, or plain Exception from - # admin-only-route / model-access / budget checks) surface as - # ProxyException consistently with pre-refactor behavior. - try: - await _run_centralized_common_checks( + # A body that never parsed is authenticated (so the trace carries identity + # and this ``auth`` span) but not authorized: there is no model to check it + # against, and budget reservation would increment live spend counters that + # only the endpoint's post-call path releases; the endpoint never runs, since + # the parse failure is raised below. + if body_parse_exception is None: + recovered_auth_obj: Final = await _authorize_authenticated_request( user_api_key_auth_obj=user_api_key_auth_obj, request=request, request_data=request_data, route=route, - ) - except Exception as e: - return await UserAPIKeyAuthExceptionHandler._handle_authentication_error( - e=e, - request=request, - request_data=request_data, - route=route, - parent_otel_span=user_api_key_auth_obj.parent_otel_span, api_key=api_key, - resolved_identity=user_api_key_auth_obj, ) - - # Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return - # paths (no master key, /user/auth route, JWT short-circuits) that bypass - # the end-user resolution block. If those paths produced an auth obj - # without an ``end_user_id`` set, fall back to extracting from the request - # body so spend logs are still attributed correctly. Validation honours - # ``litellm.validate_end_user_id_in_db``. - if user_api_key_auth_obj.end_user_id is None: - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request)) - if raw_end_user_id is not None: - resolved_end_user_id: Final = await resolve_and_validate_end_user_id( - raw_end_user_id=raw_end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth_obj.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, - ) - if resolved_end_user_id is not None: - user_api_key_auth_obj.end_user_id = resolved_end_user_id + if recovered_auth_obj is not None: + return recovered_auth_obj # Identity is now resolved. Seed it AFTER the auth span closes so the Baggage # persists on the request task (detaching the span's context token inside the @@ -2619,6 +2672,9 @@ async def user_api_key_auth( ) user_api_key_auth_obj.request_route = normalize_request_route(route) + if body_parse_exception is not None: + raise body_parse_exception + # Resolve caller identity once, here at the seam, into a single per-request # Principal projected off the key object the builder already fetched (no # second lookup). Downstream consumers read identity off this instead of diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index f7c332f2849..aef1c5ac17e 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_original_file_id, prepare_data_with_credentials, update_batch_in_database, + validate_managed_id_requirement, ) from litellm.proxy.utils import handle_exception_on_proxy, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository @@ -176,6 +177,12 @@ async def create_batch( } input_file_id: Final = _create_batch_data.get("input_file_id", None) + await validate_managed_id_requirement( + resource_id=input_file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) unified_file_id: str | Literal[False] = False model_from_file_id = None @@ -392,6 +399,12 @@ async def retrieve_batch( data: dict = {} try: + await validate_managed_id_requirement( + resource_id=batch_id, + resource_kind="batch", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) model_from_id: Final = decode_model_from_file_id(batch_id) _retrieve_batch_request: Final = RetrieveBatchRequest( batch_id=batch_id, @@ -840,6 +853,13 @@ async def cancel_batch( data: dict = {} try: + await validate_managed_id_requirement( + resource_id=batch_id, + resource_kind="batch", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) + # Check for encoded batch ID with model info model_from_id: Final = decode_model_from_file_id(batch_id) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 22200567012..fbf28e223c1 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional import litellm from litellm import get_secret from litellm._logging import verbose_proxy_logger -from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY +from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, @@ -425,6 +425,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset( "_guardrail_pipelines", "_pipeline_managed_guardrails", PRE_CALL_EXECUTED_GUARDRAILS_KEY, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, "disable_global_guardrails", "disable_global_guardrail", "opted_out_global_guardrails", diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index b1f074c26b7..9732b1d7402 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -33,6 +33,53 @@ if TYPE_CHECKING: CACHE_TTL_5M_SECONDS: Final = 300 CACHE_TTL_1H_SECONDS: Final = 3600 +AUTOROUTER_BENCHMARKS_SQL: Final = """ +WITH windowed AS ( + SELECT * FROM "LiteLLM_AutoRouterSession" + WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp +), +tier_maps AS ( + SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns + FROM ( + SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns + FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv + GROUP BY router_name, router_type, kv.key + ) per_tier + GROUP BY router_name, router_type +) +SELECT + agg.*, + COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns +FROM ( +SELECT + router_name, + router_type, + COUNT(*)::int AS sessions, + COALESCE(SUM(turns), 0)::int AS turns, + COALESCE(SUM(unordered_turns), 0)::int AS unordered_turns, + COALESCE(SUM(covered_turns), 0)::int AS covered_turns, + COALESCE(SUM(cache_hits), 0)::int AS cache_hits, + COALESCE(SUM(same_model_turns), 0)::int AS same_model_turns, + COALESCE(SUM(same_model_hits), 0)::int AS same_model_hits, + COALESCE(SUM(first_visit_turns), 0)::int AS first_visit_turns, + COALESCE(SUM(first_visit_hits), 0)::int AS first_visit_hits, + COALESCE(SUM(return_turns), 0)::int AS return_turns, + COALESCE(SUM(return_hits), 0)::int AS return_hits, + COALESCE(SUM(return_expired_misses), 0)::int AS return_expired_misses, + COALESCE(SUM(return_within_ttl_misses), 0)::int AS return_within_ttl_misses, + COALESCE(SUM(ttl_5m_turns), 0)::int AS ttl_5m_turns, + COALESCE(SUM(ttl_1h_turns), 0)::int AS ttl_1h_turns, + COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens, + COALESCE(SUM(spend), 0)::float8 AS spend, + COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend, + COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds +FROM windowed +GROUP BY router_name, router_type +) agg +LEFT JOIN tier_maps USING (router_name, router_type) +ORDER BY agg.spend DESC +""" + @dataclass(frozen=True, slots=True) class AutoRouterTurnTransaction: diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index f8ffb77edb8..a13ad00713d 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -17,6 +17,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + validate_managed_id_requirement, ) from litellm.proxy.utils import handle_exception_on_proxy from litellm.types.utils import LiteLLMFineTuningJob @@ -134,6 +135,18 @@ async def create_fine_tuning_job( ## CHECK IF MANAGED FILE ID unified_file_id: str | Literal[False] = False training_file: Final = fine_tuning_request.training_file + await validate_managed_id_requirement( + resource_id=training_file, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) + await validate_managed_id_requirement( + resource_id=fine_tuning_request.validation_file, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) response: LiteLLMFineTuningJob | None = None if training_file: unified_file_id = _is_base64_encoded_unified_file_id(training_file) @@ -246,6 +259,12 @@ async def retrieve_fine_tuning_job( try: if premium_user is not True: raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") + await validate_managed_id_requirement( + resource_id=fine_tuning_job_id, + resource_kind="fine-tuning job", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( @@ -513,6 +532,12 @@ async def cancel_fine_tuning_job( try: if premium_user is not True: raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") + await validate_managed_id_requirement( + resource_id=fine_tuning_job_id, + resource_kind="fine-tuning job", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 761d8aabc8a..b68d4a68b79 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -6,10 +6,10 @@ import concurrent.futures import inspect import json import os -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timezone from types import UnionType -from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, Union, cast, get_args, get_origin +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, cast, get_args, get_origin from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request @@ -54,8 +54,8 @@ from litellm.types.guardrails import ( if TYPE_CHECKING: from types import CodeType - from prisma.actions import LiteLLM_GuardrailsTableActions from prisma.models import LiteLLM_GuardrailsTable + from pydantic.fields import FieldInfo from litellm.proxy.utils import PrismaClient @@ -65,24 +65,44 @@ router: Final = APIRouter() GUARDRAIL_REGISTRY: Final = GuardrailRegistry() -def _guardrails_table(prisma_client: "PrismaClient") -> "LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]": - table: Final[LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]] = GuardrailsRepository(prisma_client).table +class _GuardrailsTableActions(Protocol): + async def create(self, data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": ... + + async def delete(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ... + + async def find_unique(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ... + + async def find_many( + self, where: Mapping[str, object], order: Mapping[str, str] + ) -> "Sequence[LiteLLM_GuardrailsTable]": ... + + async def update( + self, where: Mapping[str, object], data: Mapping[str, object] + ) -> "LiteLLM_GuardrailsTable | None": ... + + +def _as_str_object_mapping(mapping: Mapping[str, object]) -> Mapping[str, object]: + return mapping + + +def _guardrails_table(prisma_client: "PrismaClient") -> _GuardrailsTableActions: + table: Final[_GuardrailsTableActions] = GuardrailsRepository(prisma_client).table return table async def _create_guardrail_row(prisma_client: "PrismaClient", data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": - row: Final[LiteLLM_GuardrailsTable] = await GuardrailsRepository(prisma_client).table.create(data=data) + row: Final = await _guardrails_table(prisma_client).create(data=data) return row async def _delete_guardrail_row(prisma_client: "PrismaClient", where: Mapping[str, object]) -> None: - await GuardrailsRepository(prisma_client).table.delete(where=where) + await _guardrails_table(prisma_client).delete(where=where) async def _find_team_guardrail_rows( prisma_client: "PrismaClient", where: Mapping[str, object] ) -> "Sequence[LiteLLM_GuardrailsTable]": - rows: Final[Sequence[LiteLLM_GuardrailsTable]] = await GuardrailsRepository(prisma_client).table.find_many( + rows: Final = await _guardrails_table(prisma_client).find_many( where=where, order={"created_at": "desc"}, ) @@ -499,10 +519,12 @@ async def update_guardrail( if existing_guardrail is None: raise HTTPException(status_code=404, detail=f"Guardrail with ID {guardrail_id} not found") - result: Final = await GUARDRAIL_REGISTRY.update_guardrail_in_db( - guardrail_id=guardrail_id, - guardrail=request.guardrail, - prisma_client=prisma_client, + result: Final = _as_str_object_mapping( + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=request.guardrail, + prisma_client=prisma_client, + ) ) guardrail_name: Final = result.get("guardrail_name", "Unknown") @@ -613,7 +635,7 @@ class RegisterGuardrailRequest(BaseModel): """Request body for POST /guardrails/register. Follows Generic Guardrail API config.""" guardrail_name: str - litellm_params: dict[str, Any] # guardrail, mode, api_base required; api_key, headers, etc. optional + litellm_params: dict[str, object] # guardrail, mode, api_base required; api_key, headers, etc. optional guardrail_info: dict[str, object] | None = None team_id: str | None = None @@ -1172,12 +1194,14 @@ async def patch_guardrail( ) # Update litellm_params if default_on is provided or pii_entities_config is provided - litellm_params = LitellmParams(**dict(existing_guardrail.get("litellm_params", {}))) + existing_litellm_params: Final = _as_str_object_mapping(dict(existing_guardrail.get("litellm_params", {}))) + litellm_params = LitellmParams(**existing_litellm_params) if request.litellm_params is not None: requested_litellm_params: Final = request.litellm_params.model_dump(exclude_unset=True) litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True) litellm_params_dict.update(requested_litellm_params) - litellm_params = LitellmParams(**litellm_params_dict) + merged_litellm_params: Final = _as_str_object_mapping(litellm_params_dict) + litellm_params = LitellmParams(**merged_litellm_params) # Update guardrail_info if provided guardrail_info: Final = ( @@ -1193,10 +1217,12 @@ async def patch_guardrail( litellm_params=litellm_params, guardrail_info=guardrail_info, ) - result: Final = await GUARDRAIL_REGISTRY.update_guardrail_in_db( - guardrail_id=guardrail_id, - guardrail=guardrail, - prisma_client=prisma_client, + result: Final = _as_str_object_mapping( + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=guardrail, + prisma_client=prisma_client, + ) ) guardrail_name = result.get("guardrail_name", "Unknown") @@ -1552,31 +1578,46 @@ async def validate_blocked_words_file(request: dict[str, str]): return {"valid": False, "error": f"Validation error: {e}"} -def _get_field_type_from_annotation(field_annotation: Any) -> str: +def _dunder_origin(annotation: object) -> object: + origin: Final[object] = getattr(annotation, "__origin__", None) + return origin + + +def _dunder_name(annotation: object) -> object: + name: Final[object] = getattr(annotation, "__name__", None) + return name + + +def _dunder_args(annotation: object) -> tuple[object, ...]: + args: Final[tuple[object, ...]] = getattr(annotation, "__args__", ()) + return args + + +def _get_field_type_from_annotation(field_annotation: object) -> str: """ Convert a Python type annotation to a UI-friendly type string """ # Handle Union types (like Optional[T]) if get_origin(field_annotation) is Union or get_origin(field_annotation) is UnionType: # For Optional[T], get the non-None type - args: Final = get_args(field_annotation) + args: Final[tuple[object, ...]] = get_args(field_annotation) non_none_args: Final = [arg for arg in args if arg is not type(None)] if non_none_args: field_annotation = non_none_args[0] # Handle List types - if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is list: + if hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is list: return "array" # Handle Dict types - if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is dict: + if hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is dict: return "dict" # Handle Literal types if hasattr(field_annotation, "__origin__") and hasattr(field_annotation, "__args__"): # Check for Literal types (Python 3.8+) - origin: Final = field_annotation.__origin__ - if hasattr(origin, "__name__") and origin.__name__ == "Literal": + origin: Final = _dunder_origin(field_annotation) + if hasattr(origin, "__name__") and _dunder_name(origin) == "Literal": return "select" # For dropdown/select inputs # Handle basic types @@ -1595,66 +1636,66 @@ def _get_field_type_from_annotation(field_annotation: Any) -> str: return "string" -def _extract_literal_values(annotation: Any) -> list[str]: +def _extract_literal_values(annotation: object) -> Sequence[object]: """ Extract literal values from a Literal type annotation """ if hasattr(annotation, "__origin__") and hasattr(annotation, "__args__"): - origin: Final = annotation.__origin__ - if hasattr(origin, "__name__") and origin.__name__ == "Literal": - return list(annotation.__args__) + origin: Final = _dunder_origin(annotation) + if hasattr(origin, "__name__") and _dunder_name(origin) == "Literal": + return list(_dunder_args(annotation)) return [] -def _get_dict_key_options(field_annotation: Any) -> list[str] | None: +def _get_dict_key_options(field_annotation: object) -> Sequence[object] | None: """ Extract key options from Dict[Literal[...], T] types """ if ( hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is dict + and _dunder_origin(field_annotation) is dict and hasattr(field_annotation, "__args__") ): - args: Final = field_annotation.__args__ + args: Final = _dunder_args(field_annotation) if len(args) >= 2: key_type: Final = args[0] return _extract_literal_values(key_type) return None -def _get_dict_value_type(field_annotation: Any) -> str: +def _get_dict_value_type(field_annotation: object) -> str: """ Get the value type from Dict[K, V] types """ if ( hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is dict + and _dunder_origin(field_annotation) is dict and hasattr(field_annotation, "__args__") ): - args: Final = field_annotation.__args__ + args: Final = _dunder_args(field_annotation) if len(args) >= 2: value_type: Final = args[1] return _get_field_type_from_annotation(value_type) return "string" -def _get_list_element_options(field_annotation: Any) -> list[str] | None: +def _get_list_element_options(field_annotation: object) -> Sequence[object] | None: """ Extract element options from List[Literal[...]] types """ if ( hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is list + and _dunder_origin(field_annotation) is list and hasattr(field_annotation, "__args__") ): - args: Final = field_annotation.__args__ + args: Final = _dunder_args(field_annotation) if len(args) >= 1: element_type: Final = args[0] return _extract_literal_values(element_type) return None -def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool: +def _should_skip_optional_params(field_name: str, field_annotation: object) -> bool: """Check if optional_params field should be skipped (not meaningfully overridden).""" if field_name != "optional_params": return False @@ -1664,12 +1705,12 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool # Check if the annotation is still a generic TypeVar (not specialized) if isinstance(field_annotation, TypeVar) or ( - hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is TypeVar + hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is TypeVar ): return True # Also skip if it's a generic type that wasn't specialized - if hasattr(field_annotation, "__name__") and field_annotation.__name__ in ( + if hasattr(field_annotation, "__name__") and _dunder_name(field_annotation) in ( "T", "TypeVar", ): @@ -1677,18 +1718,18 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool # Handle Optional[T] where T is still a TypeVar if hasattr(field_annotation, "__args__"): - non_none_args: Final = [arg for arg in field_annotation.__args__ if arg is not type(None)] + non_none_args: Final = [arg for arg in _dunder_args(field_annotation) if arg is not type(None)] if non_none_args and isinstance(non_none_args[0], TypeVar): return True return False -def _unwrap_optional_type(field_annotation: Any) -> Any: +def _unwrap_optional_type(field_annotation: object) -> object: """Unwrap Optional types to get the actual type.""" if get_origin(field_annotation) is Union or get_origin(field_annotation) is UnionType: # For Optional[BaseModel], get the non-None type - args: Final = get_args(field_annotation) + args: Final[tuple[object, ...]] = get_args(field_annotation) non_none_args: Final = [arg for arg in args if arg is not type(None)] if non_none_args: return non_none_args[0] @@ -1696,20 +1737,20 @@ def _unwrap_optional_type(field_annotation: Any) -> Any: def _build_field_dict( - field: Any, - field_annotation: Any, + field: "FieldInfo", + field_annotation: object, description: str, required: bool, -) -> dict[str, Any]: +) -> dict[str, object]: """Build field dictionary for non-nested fields.""" # Determine the field type from annotation field_type = _get_field_type_from_annotation(field_annotation) # Check for custom UI type override - field_json_schema_extra: Final = getattr(field, "json_schema_extra", {}) + field_json_schema_extra: Final[Mapping[str, object]] = getattr(field, "json_schema_extra", {}) if field_json_schema_extra and "ui_type" in field_json_schema_extra: ui_type: Final = field_json_schema_extra["ui_type"] - field_type = ui_type.value if hasattr(ui_type, "value") else ui_type + field_type = getattr(ui_type, "value", ui_type) elif field_json_schema_extra and "type" in field_json_schema_extra: field_type = field_json_schema_extra["type"] @@ -1748,8 +1789,9 @@ def _build_field_dict( field_dict["options"] = literal_options # Add default value if it exists - if field.default is not None and field.default is not ...: - field_dict["default_value"] = field.default + field_default: Final[object] = getattr(field, "default", None) + if field_default is not None and field_default is not ...: + field_dict["default_value"] = field_default # Copy min, max, step from json_schema_extra for number/percentage inputs if field_json_schema_extra: @@ -1763,7 +1805,7 @@ def _build_field_dict( def _extract_fields_recursive( model: type[BaseModel], depth: int = 0, -) -> dict[str, Any]: +) -> dict[str, object]: # Check if we've exceeded the maximum recursion depth if depth > DEFAULT_MAX_RECURSE_DEPTH: raise HTTPException( @@ -1817,7 +1859,7 @@ def _extract_fields_recursive( return fields -def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, Any]: +def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, object]: """ Get the fields from a Pydantic model as a nested dictionary structure """ @@ -2141,7 +2183,26 @@ def _resolve_guardrail_input_type(active_guardrail: CustomGuardrail, input_type: return "response" if input_type == "response" else "request" -def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGuardrailRequest) -> None: +class _GuardrailLoggingObj(Protocol): + call_type: str + model_call_details: dict[str, object] + + @property + def update_messages(self) -> "Callable[..., object]": ... + + @property + def async_success_handler(self) -> "Callable[..., Awaitable[object]]": ... + + @property + def success_handler(self) -> "Callable[..., object]": ... + + +class _GuardrailProxyLogging(Protocol): + @property + def post_call_success_hook(self) -> "Callable[..., Awaitable[object]]": ... + + +def _patch_logging_obj_for_guardrail(litellm_logging_obj: _GuardrailLoggingObj, request: ApplyGuardrailRequest) -> None: """Configure the logging object so Langfuse/OTEL extract input and output correctly.""" litellm_logging_obj.call_type = "pass_through_endpoint" litellm_logging_obj.model_call_details["call_type"] = "pass_through_endpoint" @@ -2151,8 +2212,8 @@ def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGua async def _emit_guardrail_success_logs( - proxy_logging_obj: Any, - litellm_logging_obj: Any, + proxy_logging_obj: _GuardrailProxyLogging, + litellm_logging_obj: _GuardrailLoggingObj | None, data: dict, user_api_key_dict: UserAPIKeyAuth, response: ApplyGuardrailResponse, diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index 068a3ecf31b..facb822d00d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -19,7 +19,7 @@ request is sent with the ``X-Cisco-AI-Defense-API-Key`` header. import json import os -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import dataclass, replace from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal @@ -94,13 +94,13 @@ class _CiscoVerdict: is_safe: bool | None classifications: list[str] severity: str | None - rules: list[dict[str, Any]] + rules: list[dict[str, object]] explanation: str | None event_id: str | None action: str | None = None sanitized_text: str | None = None - sanitized_messages: list[dict[str, Any]] | None = None - sanitized_mcp_arguments: dict[str, Any] | None = None + sanitized_messages: list[dict[str, object]] | None = None + sanitized_mcp_arguments: dict[str, object] | None = None class CiscoAIDefenseGuardrailMissingSecrets(Exception): @@ -136,7 +136,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): api_base: str | None = None, inspection_type: str | None = None, inspect_path: str | None = None, - enabled_rules: list[dict[str, Any]] | None = None, + enabled_rules: Sequence[object] | None = None, integration_profile_id: str | None = None, integration_profile_version: str | None = None, integration_tenant_id: str | None = None, @@ -415,7 +415,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: AsyncIterator[Any], + response: AsyncIterator[object], request_data: dict, ): """Buffer and inspect streaming chat output before delivery.""" @@ -437,7 +437,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self.guardrail_name, ) - all_chunks: Final[list[Any]] = [] + all_chunks: Final[list[object]] = [] try: async for chunk in response: all_chunks.append(chunk) @@ -497,7 +497,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): response_obj=assembled, ) except HTTPException as exc: - error_obj: dict[str, Any] = self._http_exception_to_error_obj(exc) + error_obj: dict[str, object] = self._http_exception_to_error_obj(exc) verbose_proxy_logger.warning( "Cisco AI Defense guardrail (%s): streaming response " "blocked — emitting SSE error event instead of " @@ -531,7 +531,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): for chunk in all_chunks: yield chunk - def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> dict[str, Any]: + def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> dict[str, object]: """Canonical block payload used across all four block paths. Same dict is the ``HTTPException.detail`` for chat / MCP request @@ -555,34 +555,34 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): "event_id": verdict.event_id, } - def _http_exception_to_error_obj(self, exc: HTTPException) -> dict[str, Any]: + def _http_exception_to_error_obj(self, exc: HTTPException) -> dict[str, object]: """Wrap an ``HTTPException`` detail into the SSE ``error`` payload. For Cisco's own blocks the detail is already the canonical block payload, so this is a near-passthrough that just adds ``code`` / ``guardrail`` defaults for non-Cisco / unstructured details. """ - error_obj: dict[str, Any] = dict(exc.detail) if isinstance(exc.detail, dict) else {"message": str(exc.detail)} + error_obj: dict[str, object] = {**exc.detail} if isinstance(exc.detail, dict) else {"message": str(exc.detail)} error_obj.setdefault("message", error_obj.get("error", "Guardrail block")) error_obj.setdefault("code", exc.status_code) error_obj.setdefault("guardrail", self.guardrail_name) return error_obj @classmethod - def _streaming_content_was_modified(cls, original_chunks: list[Any], assembled: ModelResponse) -> bool: + def _streaming_content_was_modified(cls, original_chunks: Sequence[object], assembled: ModelResponse) -> bool: """Decide whether redact changed content or tool/function arguments.""" original_text: Final = cls._extract_streaming_chunk_scan_text(original_chunks) assembled_text: Final = " ".join(m.get("content", "") for m in cls._extract_response_messages(assembled)) return original_text != assembled_text @classmethod - def _extract_streaming_chunk_scan_text(cls, chunks: list[Any]) -> str: + def _extract_streaming_chunk_scan_text(cls, chunks: Sequence[object]) -> str: original_text = "" argument_text = "" for chunk in chunks: choices = getattr(chunk, "choices", None) or [] for c in choices: - delta = getattr(c, "delta", None) + delta: object | None = getattr(c, "delta", None) if delta is None: continue text = getattr(delta, "content", None) @@ -595,7 +595,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): args = cls._extract_tool_call_arguments(tc) if args: argument_text += args - fc = getattr(delta, "function_call", None) + fc: object | None = getattr(delta, "function_call", None) if fc is not None: args = cls._extract_function_call_arguments(fc) if args: @@ -673,7 +673,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): allow, WARNING for intervened/redacted, ERROR is left for upstream API failures. """ - fields: Final[dict[str, Any]] = { + fields: Final[dict[str, object]] = { "guardrail": self.guardrail_name, "surface": context.surface, "direction": context.direction, @@ -752,7 +752,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, direction: str = "input", response_obj: object = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: url: Final = f"{self.api_base}{self.inspect_path}" payload: Final = self._build_chat_payload(messages, request_data, user_api_key_dict) start_time: Final = datetime.now() @@ -784,7 +784,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): messages: list[dict[str, str]], request_data: dict, user_api_key_dict: UserAPIKeyAuth, - ) -> dict[str, Any]: + ) -> dict[str, object]: return { "messages": messages, "metadata": self._build_metadata(request_data, user_api_key_dict), @@ -798,9 +798,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): async def _post_inspection( self, url: str, - payload: dict[str, Any], + payload: dict[str, object], surface: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: headers: Final = self._build_headers() verbose_proxy_logger.debug( "Cisco AI Defense guardrail: posting %s inspection to %s", @@ -856,8 +856,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, request_data: dict, user_api_key_dict: UserAPIKeyAuth, - ) -> dict[str, Any]: - metadata: Final[dict[str, Any]] = {} + ) -> dict[str, object]: + metadata: Final[dict[str, object]] = {} user: Final = request_data.get("user") or getattr(user_api_key_dict, "user_id", None) if user: @@ -884,8 +884,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return metadata - def _build_config(self) -> dict[str, Any]: - config: Final[dict[str, Any]] = {} + def _build_config(self) -> dict[str, object]: + config: Final[dict[str, object]] = {} if self.enabled_rules: config["enabled_rules"] = self.enabled_rules if self.integration_profile_id: @@ -899,7 +899,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return config @staticmethod - def _normalize_rule(rule: object) -> dict[str, Any]: + def _normalize_rule(rule: object) -> dict[str, object]: """Coerce a user-supplied rule into the wire-shape dict Cisco expects. Accepts ``str``, ``dict``, and Pydantic model inputs. @@ -922,7 +922,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): rule = dumped if isinstance(rule, dict): - normalized: Final[dict[str, Any]] = {} + normalized: Final[dict[str, object]] = {} rule_name: Final = rule.get("rule_name") if rule_name: normalized["rule_name"] = rule_name @@ -950,7 +950,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): context: _ScanContext, start_time: datetime, response_obj: object = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Parse, log, and (optionally) raise/redact on the Cisco verdict. ``context.direction`` is ``"input"`` for request scans and ``"output"`` @@ -1119,10 +1119,10 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @classmethod def _sanitize_response_for_logging( cls, - inspect_response: dict[str, Any], + inspect_response: Mapping[str, object], surface: str, action: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Drop bulky / privacy-sensitive fields, recursing into nested dicts. MCP verdicts are commonly nested under ``result``, so a @@ -1138,9 +1138,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return sanitized @classmethod - def _strip_sensitive_keys(cls, d: dict[str, Any]) -> dict[str, Any]: + def _strip_sensitive_keys(cls, d: Mapping[str, object]) -> dict[str, object]: """Recursively strip privacy-sensitive keys from a verdict dict.""" - out: Final[dict[str, Any]] = {} + out: Final[dict[str, object]] = {} for key, value in d.items(): if key.startswith("_") or key in cls._REDACTED_LOG_KEYS: continue @@ -1222,8 +1222,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_jsonrpc_error( - inspect_response: dict[str, Any], - ) -> dict[str, Any] | None: + inspect_response: Mapping[str, object], + ) -> dict[str, object] | None: """Detect a JSON-RPC error envelope inside an HTTP 200 response. The Cisco Inspect API can return ``{"error": {...}}`` (or nest one @@ -1270,7 +1270,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_sanitized_text( - inspect_response: dict[str, Any], + inspect_response: Mapping[str, object], ) -> str | None: """Pull ``sanitized_text`` (or camelCase variant) off the verdict.""" for key in ("sanitized_text", "sanitizedText"): @@ -1287,8 +1287,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_sanitized_messages( - inspect_response: dict[str, Any], - ) -> list[dict[str, Any]] | None: + inspect_response: Mapping[str, object], + ) -> list[dict[str, object]] | None: """Pull a sanitized OpenAI-format messages array off the verdict. Cisco can return the rewrite under several keys; we accept any of @@ -1354,7 +1354,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_mcp_input( request_data: dict, sanitized_text: str | None, - sanitized_mcp_arguments: dict[str, Any] | None, + sanitized_mcp_arguments: dict[str, object] | None, ) -> bool: """Rewrite MCP request arguments in all locations the proxy reads.""" if sanitized_mcp_arguments is not None: @@ -1388,7 +1388,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, request_data: dict, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Rewrite chat request input (``messages`` or ``input``).""" if sanitized_messages and self._extract_tool_definition_text(request_data): @@ -1444,7 +1444,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): cls, request_data: dict, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: if sanitized_messages: instruction_text: Final = cls._instruction_text_from_messages(sanitized_messages) @@ -1457,7 +1457,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return False @classmethod - def _instruction_text_from_messages(cls, messages: list[dict[str, Any]]) -> str | None: + def _instruction_text_from_messages(cls, messages: list[dict[str, object]]) -> str | None: for message in messages: if not isinstance(message, dict): continue @@ -1468,7 +1468,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return None @classmethod - def _non_instruction_messages(cls, messages: list[dict[str, Any]] | None) -> list[dict[str, Any]] | None: + def _non_instruction_messages(cls, messages: list[dict[str, object]] | None) -> list[dict[str, object]] | None: if messages is None: return None return [ @@ -1499,7 +1499,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, response_obj: object, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Rewrite chat response (``ModelResponse`` or ``ResponsesAPIResponse``).""" if response_obj is None: @@ -1526,7 +1526,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_model_response_choices( choices: list, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Redact every returned choice, including tool-call/reasoning fields.""" if sanitized_messages: @@ -1570,7 +1570,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_text_completion_choices( choices: list, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Rewrite ``/v1/completions`` text choices after Cisco redaction.""" replacement = sanitized_text @@ -1638,7 +1638,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, output_items: list, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: replacement_text: str | None = sanitized_text if not replacement_text and sanitized_messages: @@ -1672,14 +1672,14 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _sanitized_messages_to_responses_input( - sanitized_messages: list[dict[str, Any]], - ) -> list[dict[str, Any]] | None: + sanitized_messages: list[dict[str, object]], + ) -> list[dict[str, object]] | None: """Convert chat-shape sanitized_messages to Responses API ``input``. Returns ``None`` if nothing usable could be converted, so the caller falls back to ``on_flagged_action``. """ - out: Final[list[dict[str, Any]]] = [] + out: Final[list[dict[str, object]]] = [] for m in sanitized_messages: if not isinstance(m, dict): continue @@ -1764,7 +1764,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): start_time: datetime | None = None, surface: str = "chat", direction: str = "input", - ) -> dict[str, Any]: + ) -> dict[str, object]: verbose_proxy_logger.error( "Cisco AI Defense guardrail (%s): API communication failed: %s", surface, @@ -2060,7 +2060,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return getattr(obj, key, None) @classmethod - def _field_list(cls, obj: object, key: str) -> list[Any]: + def _field_list(cls, obj: object, key: str) -> list[object]: value: Final = cls._field(obj, key) return value if isinstance(value, list) else [] diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index c29da89b15f..5bbb01c6c8e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -8,8 +8,8 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint import copy import json -from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING, Any, Final +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import HTTPException @@ -34,6 +34,9 @@ if TYPE_CHECKING: # Imported lazily at runtime (inside the streaming hook) to avoid a # module-level cyclic import with litellm.integrations.custom_guardrail. from litellm.integrations.custom_guardrail import ModifyResponseException + from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + ) # Call types that use NDJSON streaming (A2A); guardrail HTTPException is emitted as in-stream error A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message) @@ -41,12 +44,35 @@ A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message) GUARDRAIL_NAME: Final = "unified_llm_guardrails" +class _EndpointTranslation(Protocol): + @property + def process_input_messages(self) -> "Callable[..., Awaitable[dict[str, object]]]": ... + + @property + def process_output_response(self) -> "Callable[..., Awaitable[object]]": ... + + @property + def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ... + + @property + def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... + + +def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation: + return translation + + +def _chunk_choices(item: object) -> Sequence[object]: + choices: Final[Sequence[object]] = getattr(item, "choices", None) or [] + return choices + + class _StreamTerminated(Exception): """Internal signal that the incremental transform stream has already emitted its terminal chunks (block message or in-stream error) and must stop.""" -def _get_a2a_request_id(responses_so_far: list[Any], request_data: dict) -> str | None: +def _get_a2a_request_id(responses_so_far: Sequence[object], request_data: dict) -> str | None: """Get JSON-RPC request id from first A2A chunk or request body for in-stream error reporting.""" for item in responses_so_far: if isinstance(item, dict) and "id" in item: @@ -138,7 +164,9 @@ class UnifiedLLMGuardrails(CustomLogger): except ValueError: return data # handle unmapped call types - endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation( + endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + ) _ensure_litellm_metadata(data, user_api_key_dict) @@ -156,7 +184,7 @@ class UnifiedLLMGuardrails(CustomLogger): async def async_moderation_hook( self, data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral - ) -> Any: + ) -> object: """ Runs in parallel to LLM API call Runs on only Input @@ -187,7 +215,9 @@ class UnifiedLLMGuardrails(CustomLogger): if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings: return data - endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation( + endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + ) _ensure_litellm_metadata(data, user_api_key_dict) @@ -202,7 +232,7 @@ class UnifiedLLMGuardrails(CustomLogger): data: dict, user_api_key_dict: UserAPIKeyAuth, response, - ) -> Any: + ) -> object: """ Runs on response from LLM API call @@ -271,7 +301,9 @@ class UnifiedLLMGuardrails(CustomLogger): ) return response - endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation( + endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + ) try: response = await endpoint_translation.process_output_response( @@ -299,10 +331,10 @@ class UnifiedLLMGuardrails(CustomLogger): async def _handle_streaming_block( self, exc: "ModifyResponseException", - endpoint_translation: Any, + endpoint_translation: _EndpointTranslation, stream_started: bool, - responses_so_far: list[Any], - ) -> AsyncGenerator[Any, None]: + responses_so_far: Sequence[object], + ) -> AsyncGenerator[object, None]: """ Terminate a streamed response cleanly when a guardrail blocks it. @@ -323,7 +355,7 @@ class UnifiedLLMGuardrails(CustomLogger): @staticmethod def _resolve_transform_call_type( user_api_key_dict: UserAPIKeyAuth, - mappings: dict, + mappings: Mapping[CallTypes, type["BaseTranslation"]], ) -> str | None: """Resolve the call type for the incremental_diff path, or None if the route is unresolvable / unsupported. @@ -356,9 +388,9 @@ class UnifiedLLMGuardrails(CustomLogger): self, exc: HTTPException, call_type: str | None, - responses_so_far: list[Any], + responses_so_far: Sequence[object], request_data: dict, - ) -> AsyncGenerator[Any, None]: + ) -> AsyncGenerator[object, None]: """Surface a mid-stream HTTPException. For A2A (NDJSON) call types the response has already started, so emit an in-stream JSON-RPC error chunk; otherwise re-raise so the proxy can report it. @@ -387,7 +419,7 @@ class UnifiedLLMGuardrails(CustomLogger): def _build_transform_chunk( self, *, - reference_chunk: Any, + reference_chunk: object, mutated_text_per_choice: dict[int, str], emitted_text_per_choice: dict[int, str], holdback_per_choice: dict[int, int], @@ -500,18 +532,18 @@ class UnifiedLLMGuardrails(CustomLogger): async def _emit_transform_round( self, *, - endpoint_translation: Any, + endpoint_translation: _EndpointTranslation, guardrail_to_apply: CustomGuardrail, request_data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: str, - reference_chunk: Any, - responses_so_far: list[Any], - responses_yielded: list[Any], + reference_chunk: object, + responses_so_far: Sequence[object], + responses_yielded: list[object], emitted_text_per_choice: dict[int, str], finish_reason_per_choice: dict[int, str | None], is_final: bool, - ) -> AsyncGenerator[Any, None]: + ) -> AsyncGenerator[object, None]: """Run one guardrail processing round and emit the resulting diff chunk. Raises ``_StreamTerminated`` (after emitting the terminal block message or @@ -564,14 +596,14 @@ class UnifiedLLMGuardrails(CustomLogger): self, *, guardrail_to_apply: CustomGuardrail, - response: Any, + response: AsyncIterable[object], request_data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: str, sampling_rate: int, end_of_stream_only: bool, - mappings: dict, - ) -> AsyncGenerator[Any, None]: + mappings: Mapping[CallTypes, type["BaseTranslation"]], + ) -> AsyncGenerator[object, None]: """Emit guardrail text transformations as new deltas on the stream. Raw chunks are withheld and accumulated; on each sampled processing round @@ -580,15 +612,15 @@ class UnifiedLLMGuardrails(CustomLogger): synthetic chunk. A BLOCK terminates the stream via the shared block handler; an underflow surfaces as an HTTPException. """ - endpoint_translation: Final = mappings[CallTypes(call_type)]() - responses_so_far: Final[list[Any]] = [] - responses_yielded: Final[list[Any]] = [] + endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + responses_so_far: Final[list[object]] = [] + responses_yielded: Final[list[object]] = [] emitted_text_per_choice: Final[dict[int, str]] = {} finish_reason_per_choice: Final[dict[int, str | None]] = {} chunk_counter = 0 - last_chunk: Any | None = None + last_chunk: object | None = None - def _round(reference_chunk: Any, is_final: bool) -> AsyncGenerator[Any, None]: + def _round(reference_chunk: object, is_final: bool) -> AsyncGenerator[object, None]: return self._emit_transform_round( endpoint_translation=endpoint_translation, guardrail_to_apply=guardrail_to_apply, @@ -694,13 +726,13 @@ class UnifiedLLMGuardrails(CustomLogger): async def _inspect_full_response_for_block( self, *, - endpoint_translation: Any, + endpoint_translation: _EndpointTranslation, guardrail_to_apply: CustomGuardrail, request_data: dict, user_api_key_dict: UserAPIKeyAuth, - responses_so_far: list[Any], - responses_yielded: list[Any], - ) -> AsyncGenerator[Any, None]: + responses_so_far: Sequence[object], + responses_yielded: Sequence[object], + ) -> AsyncGenerator[object, None]: """Run the block-only guardrail inspection over the full assembled response (text + tool calls) so nothing bypasses the block decision. @@ -734,17 +766,17 @@ class UnifiedLLMGuardrails(CustomLogger): raise _StreamTerminated() @staticmethod - def _chunk_has_tool_calls(item: Any) -> bool: - for choice in getattr(item, "choices", None) or []: + def _chunk_has_tool_calls(item: object) -> bool: + for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) if getattr(delta, "tool_calls", None): return True return False @staticmethod - def _chunk_carries_text(item: Any) -> bool: + def _chunk_carries_text(item: object) -> bool: """True if any choice in this chunk has non-empty string ``delta.content``.""" - for choice in getattr(item, "choices", None) or []: + for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) content = getattr(delta, "content", None) if isinstance(content, str) and content != "": @@ -753,7 +785,7 @@ class UnifiedLLMGuardrails(CustomLogger): @staticmethod def _tool_call_passthrough_chunk( - item: Any, + item: object, finish_reason_per_choice: "dict[int, str | None] | None" = None, ) -> ModelResponseStream: """Copy of a chunk carrying tool calls with all text content stripped. @@ -772,7 +804,7 @@ class UnifiedLLMGuardrails(CustomLogger): redaction purpose. """ synthetic_choices: Final[list[StreamingChoices]] = [] - for choice in getattr(item, "choices", None) or []: + for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) idx = getattr(choice, "index", 0) or 0 original_finish = getattr(choice, "finish_reason", None) @@ -801,15 +833,15 @@ class UnifiedLLMGuardrails(CustomLogger): ) @staticmethod - def _record_finish_reasons(item: Any, finish_reason_per_choice: dict[int, str | None]) -> None: - for choice in getattr(item, "choices", None) or []: + def _record_finish_reasons(item: object, finish_reason_per_choice: dict[int, str | None]) -> None: + for choice in _chunk_choices(item): finish_reason = getattr(choice, "finish_reason", None) if finish_reason is not None: finish_reason_per_choice[getattr(choice, "index", 0) or 0] = finish_reason @staticmethod - def _chunk_has_finish_reason(item: Any) -> bool: - choices: Final = getattr(item, "choices", None) or [] + def _chunk_has_finish_reason(item: object) -> bool: + choices: Final = _chunk_choices(item) return any(getattr(choice, "finish_reason", None) is not None for choice in choices) async def async_post_call_streaming_iterator_hook( @@ -845,22 +877,22 @@ class UnifiedLLMGuardrails(CustomLogger): # Get streaming configuration. Resolution order (later wins): default # < guardrail attribute < guardrail_config dict < this callback's # optional_params. - def _streaming_flag(name: str, default: Any) -> Any: + def _streaming_flag(name: str, default: object) -> Any: value = default if guardrail_to_apply is not None: value = getattr(guardrail_to_apply, name, value) - config: Final = getattr(guardrail_to_apply, "guardrail_config", {}) + config: Final[Mapping[str, object]] = getattr(guardrail_to_apply, "guardrail_config", {}) if isinstance(config, dict): value = config.get(name, value) return self.optional_params.get(name, value) - sampling_rate: Final = _streaming_flag("streaming_sampling_rate", 5) + sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5) # Only apply the guardrail at end of stream (not per chunk). - end_of_stream_only = _streaming_flag("streaming_end_of_stream_only", False) + end_of_stream_only: bool = _streaming_flag("streaming_end_of_stream_only", False) # "block_only" (default) drops guardrail text rewrites on the streaming # path; "incremental_diff" emits them as synthetic deltas (see # _run_incremental_transform_stream). - streaming_transform_mode: Final = _streaming_flag("streaming_transform_mode", "block_only") + streaming_transform_mode: Final[str] = _streaming_flag("streaming_transform_mode", "block_only") # Withhold every chunk until end-of-stream moderation passes, then # release the original chunks (clean) or only the block message # (blocked) -- moderating the whole response *before* any content @@ -868,7 +900,9 @@ class UnifiedLLMGuardrails(CustomLogger): # release the original chunks are replayed as-is, so a # content-rewriting guardrail (e.g. PII masking) would leak # unredacted content. Guarded below via mask_response_content. - buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", buffer_until_moderated_default) + buffer_until_moderated: bool = _streaming_flag( + "streaming_buffer_until_moderated", buffer_until_moderated_default + ) if ( buffer_until_moderated @@ -939,9 +973,9 @@ class UnifiedLLMGuardrails(CustomLogger): # Infer call type from first chunk call_type = None chunk_counter = 0 - responses_so_far: Final[list[Any]] = [] - responses_yielded: Final[list[Any]] = [] - pending_end_of_stream_items: Final[list[Any]] = [] + responses_so_far: Final[list[object]] = [] + responses_yielded: Final[list[object]] = [] + pending_end_of_stream_items: Final[list[object]] = [] # Whether any real response chunk has been forwarded to the client. # Drives how a block terminates the stream: continue the in-progress # message (True) vs emit a standalone block message (False, buffered). diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index dd61cad15a1..a64ed764a67 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -26,7 +26,8 @@ Usage: import base64 import json -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -36,7 +37,10 @@ from litellm.llms.litellm_proxy.skills.prompt_injection import ( SkillPromptInjectionHandler, ) from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth -from litellm.types.utils import CallTypes, CallTypesLiteral +from litellm.types.utils import CallTypes, CallTypesLiteral, LLMResponseTypes + +if TYPE_CHECKING: + from litellm.llms.litellm_proxy.skills.sandbox_executor import SkillsSandboxExecutor class SkillsInjectionHook(CustomLogger): @@ -99,7 +103,7 @@ class SkillsInjectionHook(CustomLogger): verbose_proxy_logger.debug("SkillsInjectionHook: Processing %s skills", len(skills)) litellm_skills: Final[list[LiteLLM_SkillsTable]] = [] - anthropic_skills: Final[list[dict[str, Any]]] = [] + anthropic_skills: Final[list[dict[str, object]]] = [] # Separate skills by prefix for skill in skills: @@ -324,9 +328,9 @@ class SkillsInjectionHook(CustomLogger): async def async_post_call_success_deployment_hook( self, request_data: dict, - response: Any, + response: LLMResponseTypes, call_type: CallTypes | None, - ) -> Any | None: + ) -> LLMResponseTypes | None: """ Post-call hook to handle automatic code execution. @@ -372,7 +376,7 @@ class SkillsInjectionHook(CustomLogger): # Check if any tool call needs execution (litellm_code_execution or skill tool) has_executable_tool = False for tc in tool_calls: - tool_name = tc.get("name", "") + tool_name: str = tc.get("name", "") # Execute if it's litellm_code_execution OR a skill tool (litellm_skill_xxx) if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value or tool_name.startswith(LITELLM_SKILL_ID_PREFIX): has_executable_tool = True @@ -441,7 +445,7 @@ class SkillsInjectionHook(CustomLogger): data: dict, response: Any, skill_files: dict[str, bytes], - ) -> Any: + ) -> LLMResponseTypes | None: """ Execute the code execution loop for messages API (Anthropic format). @@ -466,7 +470,7 @@ class SkillsInjectionHook(CustomLogger): max_tokens: Final = data.get("max_tokens", 4096) executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) - generated_files: Final[list[dict[str, Any]]] = [] + generated_files: Final[list[dict[str, object]]] = [] current_response = response for iteration in range(self.max_iterations): @@ -511,9 +515,9 @@ class SkillsInjectionHook(CustomLogger): # Process tool calls tool_results = [] for tc in tool_calls: - tool_name = tc.get("name", "") + tool_name: str = tc.get("name", "") tool_id = tc.get("id", "") - tool_input = tc.get("input", {}) + tool_input: Mapping[str, str] = tc.get("input", {}) # Execute if it's litellm_code_execution OR a skill tool if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: @@ -561,8 +565,8 @@ class SkillsInjectionHook(CustomLogger): self, code: str, skill_files: dict[str, bytes], - executor: Any, - generated_files: list[dict[str, Any]], + executor: "SkillsSandboxExecutor", + generated_files: list[dict[str, object]], ) -> str: """Execute code in sandbox and return result string.""" try: @@ -574,7 +578,8 @@ class SkillsInjectionHook(CustomLogger): # Collect generated files if exec_result.get("files"): - for f in exec_result["files"]: + files: Final[Sequence[Mapping[str, str]]] = exec_result["files"] + for f in files: generated_files.append( { "name": f["name"], @@ -595,10 +600,10 @@ class SkillsInjectionHook(CustomLogger): async def _execute_skill_tool( self, tool_name: str, - tool_input: dict[str, Any], + tool_input: Mapping[str, str], skill_files: dict[str, bytes], - executor: Any, - generated_files: list[dict[str, Any]], + executor: "SkillsSandboxExecutor", + generated_files: list[dict[str, object]], ) -> str: """Execute a skill tool by generating and running code based on skill content.""" # Generate code based on available skill modules @@ -670,7 +675,7 @@ print('No executable skill module found') data: dict, response: Any, skill_files: dict[str, bytes], - ) -> Any: + ) -> LLMResponseTypes: """ Execute the code execution loop until model gives final response. @@ -704,7 +709,7 @@ print('No executable skill module found') kwargs: Final = {k: v for k, v in data.items() if k not in _EXCLUDED_ACOMPLETION_KEYS} executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) - generated_files: Final[list[dict[str, Any]]] = [] + generated_files: Final[list[dict[str, object]]] = [] current_response: Any = response for iteration in range(self.max_iterations): @@ -713,7 +718,7 @@ print('No executable skill module found') stop_reason = current_response.choices[0].finish_reason # Build assistant message for conversation history - assistant_msg_dict: dict[str, Any] = { + assistant_msg_dict: dict[str, object] = { "role": "assistant", "content": assistant_message.content, } @@ -781,13 +786,13 @@ print('No executable skill module found') self, tool_call: Any, skill_files: dict[str, bytes], - executor: Any, - generated_files: list[dict[str, Any]], + executor: "SkillsSandboxExecutor", + generated_files: list[dict[str, object]], ) -> str: """Execute a litellm_code_execution tool call and return result string.""" try: args: Final = json.loads(tool_call.function.arguments) - code: Final = args.get("code", "") + code: Final[str] = args.get("code", "") verbose_proxy_logger.debug("SkillsInjectionHook: Executing code (%s chars)", len(code)) @@ -802,7 +807,8 @@ print('No executable skill module found') # Collect generated files if exec_result.get("files"): tool_result += "\n\nGenerated files:" - for f in exec_result["files"]: + files: Final[Sequence[Mapping[str, str]]] = exec_result["files"] + for f in files: file_content = base64.b64decode(f["content_base64"]) generated_files.append( { @@ -830,8 +836,8 @@ print('No executable skill module found') def _attach_files_to_response( self, response: Any, - generated_files: list[dict[str, Any]], - ) -> Any: + generated_files: list[dict[str, object]], + ) -> LLMResponseTypes: """ Attach generated files to the response object. @@ -841,11 +847,13 @@ print('No executable skill module found') if not generated_files: return response + raw_response: Final = response + # Handle dict response (Anthropic/messages API format) if isinstance(response, dict): response["_litellm_generated_files"] = generated_files verbose_proxy_logger.debug("SkillsInjectionHook: Attached %s files to dict response", len(generated_files)) - return response + return raw_response # Handle object response (OpenAI format) try: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 83ae59ef050..c48fee96646 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -18,6 +18,7 @@ from litellm.constants import ( INTERNAL_CALL_ORIGIN_METADATA_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -226,6 +227,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( "applied_policies", "policy_sources", "routing_decision", + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, "standard_logging_object", "proxy_server_request", diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index d4221845b0c..8b6aafea751 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -26,6 +26,7 @@ from litellm.proxy.auth.auth_checks import ( can_key_call_resolved_model, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.autorouter_session_rollup import AUTOROUTER_BENCHMARKS_SQL from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.repositories.team_repository import TeamRepository from litellm.router_strategy.complexity_router import ComplexityRouter @@ -285,53 +286,6 @@ class _SessionAggRow(BaseModel): _SESSION_AGG_ROWS: Final = TypeAdapter(list[_SessionAggRow]) -_BENCHMARKS_SQL: Final = """ -WITH windowed AS ( - SELECT * FROM "LiteLLM_AutoRouterSession" - WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp -), -tier_maps AS ( - SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns - FROM ( - SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns - FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv - GROUP BY router_name, router_type, kv.key - ) per_tier - GROUP BY router_name, router_type -) -SELECT - agg.*, - COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns -FROM ( -SELECT - router_name, - router_type, - COUNT(*)::int AS sessions, - COALESCE(SUM(turns), 0)::int AS turns, - COALESCE(SUM(unordered_turns), 0)::int AS unordered_turns, - COALESCE(SUM(covered_turns), 0)::int AS covered_turns, - COALESCE(SUM(cache_hits), 0)::int AS cache_hits, - COALESCE(SUM(same_model_turns), 0)::int AS same_model_turns, - COALESCE(SUM(same_model_hits), 0)::int AS same_model_hits, - COALESCE(SUM(first_visit_turns), 0)::int AS first_visit_turns, - COALESCE(SUM(first_visit_hits), 0)::int AS first_visit_hits, - COALESCE(SUM(return_turns), 0)::int AS return_turns, - COALESCE(SUM(return_hits), 0)::int AS return_hits, - COALESCE(SUM(return_expired_misses), 0)::int AS return_expired_misses, - COALESCE(SUM(return_within_ttl_misses), 0)::int AS return_within_ttl_misses, - COALESCE(SUM(ttl_5m_turns), 0)::int AS ttl_5m_turns, - COALESCE(SUM(ttl_1h_turns), 0)::int AS ttl_1h_turns, - COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens, - COALESCE(SUM(spend), 0)::float8 AS spend, - COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend, - COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds -FROM windowed -GROUP BY router_name, router_type -) agg -LEFT JOIN tier_maps USING (router_name, router_type) -ORDER BY agg.spend DESC -""" - def _parse_benchmark_day(value: str) -> datetime: try: @@ -455,7 +409,7 @@ async def get_auto_router_benchmarks( raise HTTPException(status_code=400, detail="end_date must not be earlier than start_date") raw_rows: Final = await prisma_client.db.query_raw( - _BENCHMARKS_SQL, + AUTOROUTER_BENCHMARKS_SQL, start_day.isoformat(), (end_day + timedelta(days=1)).isoformat(), ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 640a735c916..abc5d3e53ff 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -17,7 +17,7 @@ import json import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final, cast +from typing import Any, Final, Literal, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f2cb1124fa0..2de1d177b33 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -18,7 +18,7 @@ import os import re import secrets import traceback -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import Any, Final, Literal, Optional, Protocol, TypeVar, cast @@ -171,8 +171,12 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... + async def create(self, *, data: Mapping[str, object]) -> _PrismaRowT: ... + async def create_many(self, *, data: Sequence[Mapping[str, object]]) -> int: ... + async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ... + async def update( self, *, @@ -181,6 +185,10 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): ) -> _PrismaRowT | None: ... +class _TxTables(Protocol): + litellm_proxymodeltable: _PrismaTableActions[object] + + def _prisma_table( repository: BaseRepository[_RepositoryModelT], ) -> _PrismaTableActions[_RepositoryModelT]: @@ -1650,9 +1658,12 @@ async def generate_key_fn( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - if user_custom_key_generate is not None: - if inspect.iscoroutinefunction(user_custom_key_generate): - result: Final = await user_custom_key_generate(data) + custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = ( + user_custom_key_generate + ) + if custom_key_generate_hook is not None: + if inspect.iscoroutinefunction(custom_key_generate_hook): + result: Final = await custom_key_generate_hook(data) else: raise ValueError("user_custom_key_generate must be a coroutine") decision: Final = result.get("decision", True) @@ -1847,9 +1858,10 @@ async def generate_service_account_key_fn( verbose_proxy_logger.debug("entered /key/generate") - if user_custom_key_generate is not None: - if inspect.iscoroutinefunction(user_custom_key_generate): - result: Final = await user_custom_key_generate(data) + custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_generate + if custom_key_generate_hook is not None: + if inspect.iscoroutinefunction(custom_key_generate_hook): + result: Final = await custom_key_generate_hook(data) else: raise ValueError("user_custom_key_generate must be a coroutine") decision: Final = result.get("decision", True) @@ -1918,7 +1930,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ ) casted_metadata[reserved_field] = existing_value - data_json: Final = data.model_dump(exclude_unset=True, exclude_none=True) + data_json: Final[Mapping[str, object]] = data.model_dump(exclude_unset=True, exclude_none=True) try: for k, v in data_json.items(): @@ -2179,7 +2191,7 @@ async def _process_single_key_update( llm_router: Router | None, user_custom_key_update: Callable | None = None, existing_key_row: LiteLLM_VerificationToken | None = None, -) -> dict[str, Any]: +) -> dict[str, object]: """ Process a single key update with all validations and checks. @@ -2722,9 +2734,10 @@ async def update_key_fn( ) # Custom key update hook - if user_custom_key_update is not None: - if inspect.iscoroutinefunction(user_custom_key_update): - result: Final = await user_custom_key_update(data) + custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_update + if custom_key_update_hook is not None: + if inspect.iscoroutinefunction(custom_key_update_hook): + result: Final = await custom_key_update_hook(data) else: raise ValueError("user_custom_key_update must be a coroutine") decision: Final = result.get("decision", True) @@ -4089,10 +4102,11 @@ async def delete_verification_tokens( failed_tokens: list = [] try: if prisma_client: - tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository( - prisma_client - ).table.find_many(where={"token": {"in": tokens}}) + hashed_tokens: Final[list[str]] = [_hash_token_if_needed(token=key) for key in tokens] + tokens = hashed_tokens + _keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await _prisma_table( + VerificationTokenRepository(prisma_client) + ).find_many(where={"token": {"in": hashed_tokens}}) if len(_keys_being_deleted) == 0: raise HTTPException( @@ -4291,7 +4305,7 @@ async def _rotate_master_key( if models: decrypted_models: Final = proxy_config.decrypt_model_list_from_db(new_models=models) verbose_proxy_logger.debug("ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models)) - new_models: Final = [] + new_models: Final[list[dict[str, object]]] = [] for model in decrypted_models: new_model = await _add_model_to_db( model_params=Deployment(**model), @@ -4306,7 +4320,8 @@ async def _rotate_master_key( _dumped["model_info"] = prisma.Json(_dumped["model_info"]) new_models.append(_dumped) verbose_proxy_logger.debug("Resetting proxy model table") - async with prisma_client.db.tx() as tx: + async with prisma_client.db.tx() as tx_ctx: + tx: Final[_TxTables] = tx_ctx await tx.litellm_proxymodeltable.delete_many() verbose_proxy_logger.debug("Creating %s models", len(new_models)) await tx.litellm_proxymodeltable.create_many( @@ -4630,7 +4645,7 @@ async def _execute_virtual_key_regeneration( _validate_key_alias_format(key_alias=new_key_alias) verbose_proxy_logger.debug("non_default_values: %s", non_default_values) update_data.update(non_default_values) - update_data = prisma_client.jsonify_object(data=update_data) + jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data) # If grace period set, insert deprecated key so old key remains valid await _insert_deprecated_key( @@ -4642,9 +4657,9 @@ async def _execute_virtual_key_regeneration( updated_token: Final = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, - data=update_data, + data=jsonified_update_data, ) - updated_token_dict: Final = dict(updated_token) if updated_token is not None else {} + updated_token_dict: Final[dict[str, object]] = dict(updated_token) if updated_token is not None else {} updated_token_dict["key"] = new_token updated_token_dict["token_id"] = updated_token_dict.pop("token") @@ -5589,7 +5604,7 @@ async def key_aliases( where_sql: Final = " AND ".join(where_parts) count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' - count_rows: Final = await prisma_client.db.query_raw(count_sql, *query_params) + count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params) total_count: Final = int(count_rows[0]["count"]) if count_rows else 0 aliases_params: Final = query_params + [size, (page - 1) * size] @@ -5602,7 +5617,7 @@ async def key_aliases( f" ORDER BY key_alias ASC" f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) - alias_rows: Final = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params) aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages: Final = -(-total_count // size) if total_count > 0 else 0 @@ -5695,7 +5710,7 @@ def _build_key_filter_conditions( agent_id: str | None = None, use_substring_matching: bool = False, expires_filter: str | None = None, -) -> dict[str, str | dict[str, Any] | list[dict[str, Any]]]: +) -> Mapping[str, object]: """Build filter conditions for key listing. Visibility rules: @@ -5707,14 +5722,14 @@ def _build_key_filter_conditions( so former members cannot see service accounts they created after leaving. """ # Prepare filter conditions - where: dict[str, str | dict[str, Any] | list[dict[str, Any]]] = {} + where: dict[str, object] = {} where.update(_get_condition_to_filter_out_ui_session_tokens()) # Build the OR conditions for user's keys and admin team keys - or_conditions: Final[list[dict[str, Any]]] = [] + or_conditions: Final[list[dict[str, object]]] = [] # Base conditions for user's own keys - user_condition: Final[dict[str, Any]] = {} + user_condition: Final[dict[str, object]] = {} if user_id and isinstance(user_id, str): if use_substring_matching: user_condition["user_id"] = { @@ -5784,7 +5799,7 @@ def _build_key_filter_conditions( # Apply team_id, project_id and access_group_id as global AND filters so they # narrow results across all visibility conditions (own keys, team keys, etc.) - global_filters: tuple[dict[str, Any], ...] = ( + global_filters: Final[tuple[dict[str, object], ...]] = ( *( ( {"key_alias": {"contains": key_alias, "mode": "insensitive"}} @@ -5805,7 +5820,7 @@ def _build_key_filter_conditions( else () ), ) - combined_where = {"AND": [where, *global_filters]} if global_filters else where + combined_where: Final[Mapping[str, object]] = {"AND": [where, *global_filters]} if global_filters else where verbose_proxy_logger.debug("Filter conditions: %s", combined_where) return combined_where @@ -5986,7 +6001,7 @@ async def _list_key_helper( ) -def _get_condition_to_filter_out_ui_session_tokens() -> dict[str, Any]: +def _get_condition_to_filter_out_ui_session_tokens() -> Mapping[str, object]: """ Condition to filter out UI session tokens """ @@ -6395,7 +6410,7 @@ async def _can_user_query_key_info( async def test_key_logging( user_api_key_dict: UserAPIKeyAuth, request: Request, - key_logging: list[dict[str, Any]], + key_logging: Sequence[Mapping[str, str]], ) -> LoggingCallbackStatus: """ Test the key-based logging diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 71407c89813..c2087005863 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,9 +13,9 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime import json -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Mapping, Sequence from json import JSONDecodeError -from typing import Any, Final, Literal, cast +from typing import Final, Literal, Protocol, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, ValidationError @@ -78,9 +78,7 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, - DeploymentTypedDict, GenericLiteLLMParams, - LiteLLMParamsTypedDict, updateDeployment, ) from litellm.utils import get_utc_datetime @@ -104,10 +102,80 @@ class UpdatePublicModelGroupsRequest(BaseModel): model_config = ConfigDict(extra="forbid") +class _ProxyModelRow(Protocol): + model_id: str + model_name: str + model_info: Mapping[str, object] | None + + def model_dump_json(self, *, exclude_none: bool = False) -> str: ... + + +class _ProxyModelTable(Protocol): + def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ... + + def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ... + + def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ... + + def delete(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ... + + def delete_many(self, *, where: Mapping[str, object]) -> Awaitable[int]: ... + + +class _TxModelTables(Protocol): + litellm_proxymodeltable: _ProxyModelTable + + +class _TeamRow(Protocol): + models: Sequence[str] + + def model_dump(self) -> Mapping[str, object]: ... + + +class _TeamTable(Protocol): + def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_TeamRow | None]: ... + + def update( + self, *, where: Mapping[str, object], data: Mapping[str, object], include: Mapping[str, bool] + ) -> Awaitable[LiteLLM_TeamTable]: ... + + +class _TeamIdRef(Protocol): + team_id: str + + +class _ModelAliasRow(Protocol): + id: int + model_aliases: dict[str, str] + team: _TeamIdRef | None + + +class _ModelAliasTable(Protocol): + def find_many(self, *, include: Mapping[str, bool]) -> Awaitable[Sequence[_ModelAliasRow]]: ... + + def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[object]: ... + + +def _proxy_model_table(prisma_client: PrismaClient) -> _ProxyModelTable: + return ModelRepository(prisma_client).table + + +def _repo_team_table(prisma_client: PrismaClient) -> _TeamTable: + return TeamRepository(prisma_client).table + + +def _db_team_table(prisma_client: PrismaClient) -> _TeamTable: + return prisma_client.db.litellm_teamtable + + +def _model_alias_table(prisma_client: PrismaClient) -> _ModelAliasTable: + return ModelTableRepository(prisma_client).table + + async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment | None: db_model: Final = cast( BaseModel | None, - await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}), + await _proxy_model_table(prisma_client).find_unique(where={"model_id": model_id}), ) if not db_model: @@ -166,14 +234,9 @@ def _raise_on_strategy_router_write_violation( def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: - merged_deployment_dict: Final = DeploymentTypedDict( - model_name=db_model.model_name, - litellm_params=LiteLLMParamsTypedDict(**db_model.litellm_params.model_dump(exclude_none=True)), - model_info=db_model.model_info.model_dump(exclude_none=True), - ) - # update model name - if updated_patch.model_name: - merged_deployment_dict["model_name"] = updated_patch.model_name + merged_model_name: Final = updated_patch.model_name or db_model.model_name + merged_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True) + merged_model_info: Final = db_model.model_info.model_dump(exclude_none=True) # update litellm params if updated_patch.litellm_params: @@ -182,13 +245,11 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr k: encrypt_value_helper(v) for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items() } - merged_deployment_dict["litellm_params"].update(encrypted_params) + merged_litellm_params.update(encrypted_params) # update model info if updated_patch.model_info: - if "model_info" not in merged_deployment_dict: - merged_deployment_dict["model_info"] = {} - merged_deployment_dict["model_info"].update(updated_patch.model_info.model_dump(exclude_none=True)) + merged_model_info.update(updated_patch.model_info.model_dump(exclude_none=True)) # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI # passes through (which today re-sends the OLD pricing on every save) cannot @@ -202,29 +263,25 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if updated_patch.litellm_params: for field in updated_patch.litellm_params.model_fields_set: if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None: - merged_deployment_dict["litellm_params"].pop(field, None) - merged_deployment_dict.get("model_info", {}).pop(field, None) + merged_litellm_params.pop(field, None) + merged_model_info.pop(field, None) if updated_patch.model_info: for field in updated_patch.model_info.model_fields_set: if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: - merged_deployment_dict["model_info"].pop(field, None) - merged_deployment_dict.get("litellm_params", {}).pop(field, None) + merged_model_info.pop(field, None) + merged_litellm_params.pop(field, None) # convert to prisma compatible format - prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel() - if "model_name" in merged_deployment_dict: - prisma_compatible_model_dict["model_name"] = merged_deployment_dict["model_name"] + for key, value in merged_model_info.items(): + if isinstance(value, datetime.datetime): + merged_model_info[key] = value.isoformat() - if "litellm_params" in merged_deployment_dict: - prisma_compatible_model_dict["litellm_params"] = json.dumps(merged_deployment_dict["litellm_params"]) - - if "model_info" in merged_deployment_dict: - model_info: Final = merged_deployment_dict["model_info"] - for key, value in model_info.items(): - if isinstance(value, datetime.datetime): - model_info[key] = value.isoformat() - prisma_compatible_model_dict["model_info"] = json.dumps(model_info) + prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel( + model_name=merged_model_name, + litellm_params=json.dumps(merged_litellm_params), + model_info=json.dumps(merged_model_info), + ) if updated_patch.blocked is not None: prisma_compatible_model_dict["blocked"] = updated_patch.blocked @@ -338,7 +395,7 @@ async def patch_model( update_data["updated_at"] = cast(str, get_utc_datetime()) # Perform partial update - updated_model: Final = await ModelRepository(prisma_client).table.update( + updated_model: Final = await _proxy_model_table(prisma_client).update( where={"model_id": model_id}, data=update_data, ) @@ -769,8 +826,8 @@ async def _setup_new_team_model_assignment( async def _get_team_deployments( - team_id: str, prisma_client: PrismaClient, table: Any | None = None -) -> list[LiteLLM_ProxyModelTable]: + team_id: str, prisma_client: PrismaClient, table: _ProxyModelTable | None = None +) -> Sequence[_ProxyModelRow]: """ Fetch all deployments for a given team_id from the database. @@ -785,7 +842,7 @@ async def _get_team_deployments( existing transaction. """ prefix: Final = f"model_name_{team_id}_" - table = table or ModelRepository(prisma_client).table + table = table or _proxy_model_table(prisma_client) response: Final = await table.find_many( where={ "model_name": {"startswith": prefix}, @@ -806,7 +863,7 @@ async def _get_team_deployments( async def delete_team_models( team_ids: list[str], prisma_client: PrismaClient, - llm_router: Any | None, + llm_router: Router | None, ) -> list[str]: """ Delete every BYOK model owned by the given teams, from the DB and the router. @@ -820,7 +877,8 @@ async def delete_team_models( Returns the model_ids that were deleted. """ deleted_model_ids: Final[list[str]] = [] - async with prisma_client.db.tx() as tx: + async with prisma_client.db.tx() as tx_ctx: + tx: Final[_TxModelTables] = tx_ctx for team_id in team_ids: rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable) model_ids = [row.model_id for row in rows] @@ -920,11 +978,11 @@ async def _remove_unbacked_team_models( if not names_to_remove: return - existing_team_row: Final = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + existing_team_row: Final = await _db_team_table(prisma_client).find_unique(where={"team_id": team_id}) if existing_team_row is None: return - updated_team_row: Final[LiteLLM_TeamTable] = await prisma_client.db.litellm_teamtable.update( + updated_team_row: Final[LiteLLM_TeamTable] = await _db_team_table(prisma_client).update( where={"team_id": team_id}, data={"models": [model for model in existing_team_row.models if model not in names_to_remove]}, include={"object_permission": True}, @@ -953,7 +1011,7 @@ async def _update_existing_team_model_assignment( """ def _get_team_public_model_name( - model_info: dict | str | None, + model_info: object, ) -> str | None: parsed: Final = model_info_as_mapping(model_info) if parsed is None: @@ -1062,7 +1120,7 @@ class ModelManagementAuthChecks: detail={"error": CommonProxyErrors.not_premium_user.value}, ) - _existing_team_row: Final = await TeamRepository(prisma_client).table.find_unique( + _existing_team_row: Final = await _repo_team_table(prisma_client).find_unique( where={"team_id": model_params.model_info.team_id} ) @@ -1091,7 +1149,7 @@ class ModelManagementAuthChecks: ) -> Literal[True]: ## Check team model auth if model_params.model_info is not None and model_params.model_info.team_id is not None: - team_obj_row: Final = await TeamRepository(prisma_client).table.find_unique( + team_obj_row: Final = await _repo_team_table(prisma_client).find_unique( where={"team_id": model_params.model_info.team_id} ) if team_obj_row is None: @@ -1192,7 +1250,7 @@ async def delete_model( - store keys separately """ # encrypt litellm params # - result: Final = await ModelRepository(prisma_client).table.delete(where={"model_id": model_info.id}) + result: Final = await _proxy_model_table(prisma_client).delete(where={"model_id": model_info.id}) if result is None: raise HTTPException( @@ -1265,9 +1323,9 @@ async def delete_team_model_alias( Returns: - List of team id + model alias pairs that were removed """ - team_model_aliases: Final = await ModelTableRepository(prisma_client).table.find_many(include={"team": True}) + team_model_aliases: Final = await _model_alias_table(prisma_client).find_many(include={"team": True}) tasks: Final = [] - removed_model_aliases: Final = [] + removed_model_aliases: Final[list[tuple[str, str]]] = [] for team_model_alias in team_model_aliases: model_aliases = team_model_alias.model_aliases # {"alias": "public model name"} id = team_model_alias.id @@ -1278,7 +1336,7 @@ async def delete_team_model_alias( removed_model_aliases.append((team_model_alias.team.team_id, key)) del model_aliases[key] tasks.append( - ModelTableRepository(prisma_client).table.update( + _model_alias_table(prisma_client).update( where={"id": id}, data={"model_aliases": json.dumps(model_aliases)}, ) @@ -1492,7 +1550,7 @@ async def update_model( }, ) - _model_id = None + _model_id: str | None = None _model_info: Final = getattr(model_params, "model_info", None) if _model_info is None: raise Exception("model_info not provided") @@ -1551,11 +1609,11 @@ async def update_model( else: pass - _data: Final[dict] = { + _data: Final[dict[str, str]] = { "litellm_params": json.dumps(merged_dictionary), "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, } - model_response: Final = await ModelRepository(prisma_client).table.update( + model_response: Final = await _proxy_model_table(prisma_client).update( where={"model_id": _model_id}, data=_data, ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0b99879f9fe..97f494c51de 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -15,11 +15,11 @@ import math import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Annotated, Final, Protocol, TypeVar, cast +from typing import Annotated, Final, Protocol, TypedDict, TypeVar, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue import litellm from litellm._logging import verbose_proxy_logger @@ -29,6 +29,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( UI_TEAM_ID, BlockTeamRequest, + BudgetNewRequest, CommonProxyErrors, DeleteTeamRequest, LiteLLM_AccessGroupTable, @@ -156,6 +157,15 @@ router: Final = APIRouter() _DbRecordT = TypeVar("_DbRecordT") +class _TeamIdKeyCount(TypedDict): + team_id: int + + +class _TeamIdGroupRow(TypedDict): + team_id: str + _count: _TeamIdKeyCount + + class _PrismaTableActions(Protocol[_DbRecordT]): async def find_unique( self, @@ -220,59 +230,127 @@ class _PrismaTableActions(Protocol[_DbRecordT]): where: Mapping[str, object] | None = None, ) -> int: ... + async def group_by( + self, + by: Sequence[str], + where: Mapping[str, object] | None = None, + count: Mapping[str, bool] | None = None, + ) -> Sequence[_TeamIdGroupRow]: ... + + +class _HasTableActions(Protocol[_DbRecordT]): + @property + def table(self) -> "_PrismaTableActions[_DbRecordT]": ... + + +def _typed_table( + repo: "_HasTableActions[_DbRecordT]", record_type: type[_DbRecordT] +) -> "_PrismaTableActions[_DbRecordT]": + return repo.table + + +def _as_object(value: object) -> object: + return value + + +def _nullable(value: _DbRecordT | None) -> _DbRecordT | None: + return value + + +class _UserIdRow(Protocol): + @property + def user_id(self) -> str | None: ... + + +class _HasUserIdTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_UserIdRow]": ... + + +def _user_id_rows_db(repo: "_HasUserIdTable") -> "_PrismaTableActions[_UserIdRow]": + return repo.table + + +class _RawTeamRow(Protocol): + @property + def members_with_roles(self) -> Sequence[Mapping[str, object]] | None: ... + + +class _HasRawTeamTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_RawTeamRow]": ... + + +def _raw_team_db(repo: "_HasRawTeamTable") -> "_PrismaTableActions[_RawTeamRow]": + return repo.table + + +class _BudgetWriteCall(Protocol): + async def __call__( + self, budget_obj: BudgetNewRequest, user_api_key_dict: UserAPIKeyAuth + ) -> LiteLLM_BudgetTableFull: ... + + +def _as_budget_write(fn: "_BudgetWriteCall") -> "_BudgetWriteCall": + return fn + + +class _TeamFindManyArgs(TypedDict, total=False): + take: int + skip: int + order: Mapping[str, str] + cursor: Mapping[str, object] + + +class _TeamUiViewFilters(TypedDict, total=False): + team_id: Mapping[str, str] + team_alias: Mapping[str, str] + + +class _TeamIdInFilter(TypedDict, total=False): + team_id: Mapping[str, Sequence[str]] + def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]": - team_table: Final[_PrismaTableActions[LiteLLM_TeamTable]] = TeamRepository(prisma_client).table - return team_table + return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable) def _team_membership_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamMembership]": - membership_table: Final[_PrismaTableActions[LiteLLM_TeamMembership]] = TeamMembershipRepository(prisma_client).table - return membership_table + return _typed_table(TeamMembershipRepository(prisma_client), LiteLLM_TeamMembership) def _user_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_UserTable]": - user_table: Final[_PrismaTableActions[LiteLLM_UserTable]] = UserRepository(prisma_client).table - return user_table + return _typed_table(UserRepository(prisma_client), LiteLLM_UserTable) def _model_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_ModelTable]": - model_table: Final[_PrismaTableActions[LiteLLM_ModelTable]] = ModelTableRepository(prisma_client).table - return model_table + return _typed_table(ModelTableRepository(prisma_client), LiteLLM_ModelTable) def _org_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_OrganizationTable]": - org_table: Final[_PrismaTableActions[LiteLLM_OrganizationTable]] = OrganizationRepository(prisma_client).table - return org_table + return _typed_table(OrganizationRepository(prisma_client), LiteLLM_OrganizationTable) def _org_membership_db( prisma_client: PrismaClient | None, ) -> "_PrismaTableActions[LiteLLM_OrganizationMembershipTable]": - org_membership_table: _PrismaTableActions[LiteLLM_OrganizationMembershipTable] = OrganizationMembershipRepository( - prisma_client - ).table - return org_membership_table + return _typed_table(OrganizationMembershipRepository(prisma_client), LiteLLM_OrganizationMembershipTable) def _budget_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_BudgetTableFull]": - budget_table: Final[_PrismaTableActions[LiteLLM_BudgetTableFull]] = BudgetRepository(prisma_client).table - return budget_table + return _typed_table(BudgetRepository(prisma_client), LiteLLM_BudgetTableFull) def _deleted_team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_DeletedTeamTable]": - deleted_team_table: _PrismaTableActions[LiteLLM_DeletedTeamTable] = DeletedTeamRepository(prisma_client).table - return deleted_team_table + return _typed_table(DeletedTeamRepository(prisma_client), LiteLLM_DeletedTeamTable) def _access_group_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_AccessGroupTable]": - access_group_table: _PrismaTableActions[LiteLLM_AccessGroupTable] = AccessGroupRepository(prisma_client).table - return access_group_table + return _typed_table(AccessGroupRepository(prisma_client), LiteLLM_AccessGroupTable) def _tokens_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_VerificationToken]": - tokens_table: _PrismaTableActions[LiteLLM_VerificationToken] = VerificationTokenRepository(prisma_client).table - return tokens_table + return _typed_table(VerificationTokenRepository(prisma_client), LiteLLM_VerificationToken) def _sanitize_for_log(value: object) -> str: @@ -408,7 +486,7 @@ class TeamMemberBudgetHandler: if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - team_member_budget_table: Final = await new_budget( + team_member_budget_table: Final = await _as_budget_write(new_budget)( budget_obj=budget_request, user_api_key_dict=user_api_key_dict, ) @@ -456,7 +534,7 @@ class TeamMemberBudgetHandler: if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - budget_row: Final = await update_budget( + budget_row: Final = await _as_budget_write(update_budget)( budget_obj=budget_request, user_api_key_dict=user_api_key_dict, ) @@ -571,7 +649,7 @@ class TeamMemberBudgetHandler: ) if missing: - await TeamMembershipRepository(prisma_client).table.create_many( + await _team_membership_db(prisma_client).create_many( data=missing, skip_duplicates=True, # safety net against concurrent races ) @@ -1407,9 +1485,10 @@ async def new_team( complete_team_data_dict["metadata"] = encrypt_callback_vars(complete_team_data_dict["metadata"]) complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict) + team_creation_data: Final[Mapping[str, object]] = complete_team_data_dict - team_row: Final[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.create( - data=complete_team_data_dict, + team_row: Final[LiteLLM_TeamTable] = await _team_db(prisma_client).create( + data=team_creation_data, include={"litellm_model_table": True}, ) @@ -1856,7 +1935,7 @@ async def update_team( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team_row = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id}) if existing_team_row is None: raise HTTPException( @@ -1884,7 +1963,7 @@ async def update_team( ) if data.max_budget is not None: - existing_soft_budget: Final = getattr(existing_team_row, "soft_budget", None) + existing_soft_budget: Final[object] = _as_object(getattr(existing_team_row, "soft_budget", None)) soft_budget_to_check: Final = data.soft_budget if data.soft_budget is not None else existing_soft_budget if soft_budget_to_check is not None and isinstance(soft_budget_to_check, (int, float)): if data.max_budget <= soft_budget_to_check: @@ -1943,7 +2022,7 @@ async def update_team( data.organization_id = None # check org team limits - if updating team that belongs to an org - org_id_to_check: Final = ( + org_id_to_check: Final[object] = _as_object( data.organization_id if data.organization_id is not None else existing_team_row.organization_id ) if org_id_to_check is not None and isinstance(org_id_to_check, str) and prisma_client is not None: @@ -1976,7 +2055,7 @@ async def update_team( TeamMemberBudgetHandler.strip_system_managed_metadata_keys(updated_kv["metadata"]) if "metadata" in updated_kv: - stored_metadata: Final = ( + stored_metadata: Final[Mapping[str, JsonValue] | None] = ( { # mutable-ok: the validator payload's isinstance guard requires a plain dict key: value for key, value in existing_team_row.metadata.items() @@ -2079,16 +2158,19 @@ async def update_team( updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Final[LiteLLM_TeamTable | None] = await TeamRepository(prisma_client).table.update( - where={"team_id": data.team_id}, - data=updated_kv, - # `object_permission` is included so `_refresh_cached_team` - # doesn't write a cached team with the relation nulled out — - # see team_model_add for the full rationale. - include={ - "litellm_model_table": True, - "object_permission": True, - }, + team_update_data: Final[Mapping[str, object]] = updated_kv + team_row: Final[LiteLLM_TeamTable | None] = _nullable( + await _team_db(prisma_client).update( + where={"team_id": data.team_id}, + data=team_update_data, + # `object_permission` is included so `_refresh_cached_team` + # doesn't write a cached team with the relation nulled out — + # see team_model_add for the full rationale. + include={ + "litellm_model_table": True, + "object_permission": True, + }, + ) ) if team_row is None or team_row.team_id is None: @@ -2603,7 +2685,7 @@ async def _resolve_existing_member_user_ids( if not requested_user_ids: return frozenset() - found: Final = await UserRepository(prisma_client).table.find_many( + found: Final = await _user_id_rows_db(UserRepository(prisma_client)).find_many( where={ # mutable-ok: Prisma query filters are dict-shaped "user_id": { # mutable-ok: Prisma query filters are dict-shaped "in": sorted(requested_user_ids) @@ -3098,7 +3180,9 @@ async def team_member_delete( key_val["user_id"] = data.user_id elif data.user_email is not None: key_val["user_email"] = data.user_email - existing_user_rows: Final = await UserRepository(prisma_client).table.find_many(where=key_val) + existing_user_rows: Final[Sequence[LiteLLM_UserTable] | None] = await UserRepository(prisma_client).table.find_many( + where=key_val + ) if existing_user_rows is not None and (isinstance(existing_user_rows, list) and len(existing_user_rows) > 0): for existing_user in existing_user_rows: @@ -3106,7 +3190,7 @@ async def team_member_delete( if data.team_id in existing_user.teams: team_list = existing_user.teams team_list.remove(data.team_id) - await UserRepository(prisma_client).table.update( + await _user_db(prisma_client).update( where={ "user_id": existing_user.user_id, }, @@ -3114,7 +3198,7 @@ async def team_member_delete( ) # Also clean up any existing team membership rows for this user and team - user_ids_to_delete: Final = set() + user_ids_to_delete: Final = set[str]() if data.user_id is not None: user_ids_to_delete.add(data.user_id) if existing_user_rows is not None and isinstance(existing_user_rows, list): @@ -3123,9 +3207,7 @@ async def team_member_delete( user_ids_to_delete.add(existing_user.user_id) for _uid in user_ids_to_delete: - await TeamMembershipRepository(prisma_client).table.delete_many( - where={"team_id": data.team_id, "user_id": _uid} - ) + await _team_membership_db(prisma_client).delete_many(where={"team_id": data.team_id, "user_id": _uid}) ## DELETE KEYS CREATED BY USER FOR THIS TEAM if user_ids_to_delete: @@ -3134,9 +3216,7 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository( - prisma_client - ).table.find_many( + keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -3151,7 +3231,7 @@ async def team_member_delete( litellm_changed_by=None, ) - await VerificationTokenRepository(prisma_client).table.delete_many( + await _tokens_db(prisma_client).delete_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -3311,7 +3391,7 @@ async def team_member_update( ### upsert new budget budget_patch: Final = _build_member_budget_patch(data) - async with prisma_client.db.tx() as tx: + async with prisma_client.tx() as tx: await _upsert_budget_and_membership( tx=tx, team_id=data.team_id, @@ -3654,7 +3734,7 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: list[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( + keys_to_delete: list[LiteLLM_VerificationToken] = await _tokens_db(prisma_client).find_many( where={"team_id": {"in": data.team_ids}} ) @@ -4469,7 +4549,7 @@ async def _get_keys_count_by_team( if not page_team_ids: return {} - grouped: Final = await VerificationTokenRepository(prisma_client).table.group_by( + grouped: Final = await _tokens_db(prisma_client).group_by( by=["team_id"], where={"team_id": {"in": page_team_ids}}, count={"team_id": True}, @@ -4786,7 +4866,7 @@ async def _authorize_and_filter_teams( if allowed_org_ids is not None: # Org admin: query DB for teams in their orgs - org_teams: Final = await TeamRepository(prisma_client).table.find_many( + org_teams: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many( where={"organization_id": {"in": allowed_org_ids}}, include={"litellm_model_table": True}, ) @@ -4800,7 +4880,9 @@ async def _authorize_and_filter_teams( ] elif user_id: # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays) - response: Final = await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True}) + response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many( + include={"litellm_model_table": True} + ) return [ team for team in response @@ -4808,7 +4890,7 @@ async def _authorize_and_filter_teams( ] else: # Proxy admin: all teams - return list(await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True})) + return list(await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True})) @router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @@ -4860,7 +4942,7 @@ async def list_team( _team_memberships.append(tm) # add all keys that belong to the team - keys = await VerificationTokenRepository(prisma_client).table.find_many(where={"team_id": team.team_id}) + keys = await _tokens_db(prisma_client).find_many(where={"team_id": team.team_id}) try: returned_responses.append( @@ -4911,7 +4993,7 @@ async def get_paginated_teams( total_count: Final = await _team_db(prisma_client).count() # Get paginated teams - teams: Final = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _team_db(prisma_client).find_many( skip=skip, take=page_size, order={"team_alias": "asc"}, # Sort by team_alias @@ -4961,7 +5043,7 @@ async def ui_view_teams( skip: Final = (page - 1) * page_size # Build where conditions based on provided parameters - where_conditions: Final = {} + where_conditions: Final[_TeamUiViewFilters] = {} if team_id: where_conditions["team_id"] = { @@ -4976,7 +5058,7 @@ async def ui_view_teams( } # Query users with pagination and filters - teams: Final = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _team_db(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -5166,13 +5248,13 @@ async def team_model_delete( ) # Get current models list - current_models: Final = team_obj.models or [] + current_models: Final[Sequence[str]] = team_obj.models or [] # Remove specified models updated_models: Final = [m for m in current_models if m not in data.models] # Update team. See team_model_add for the rationale on `include`. - updated_team: Final = await TeamRepository(prisma_client).table.update( + updated_team: Final = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"models": updated_models}, include={"object_permission": True}, @@ -5425,7 +5507,7 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi BATCH_SIZE: Final = 500 while True: - find_args: dict = { + find_args: _TeamFindManyArgs = { "take": BATCH_SIZE, "order": {"team_id": "asc"}, } @@ -5433,7 +5515,7 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi find_args["cursor"] = {"team_id": cursor} find_args["skip"] = 1 - teams = await TeamRepository(prisma_client).table.find_many(**find_args) + teams = await _team_db(prisma_client).find_many(**find_args) if not teams: break @@ -5528,11 +5610,11 @@ async def get_team_daily_activity( ) ## Fetch team aliases and check team admin status - where_condition: Final = {} + where_condition: Final[_TeamIdInFilter] = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} - team_aliases: Final = await TeamRepository(prisma_client).table.find_many(where=where_condition) - team_alias_metadata: Final = {t.team_id: {"team_alias": t.team_alias} for t in team_aliases} + team_aliases: Final = await _team_db(prisma_client).find_many(where=where_condition) + team_alias_metadata: Final = {t.team_id: {"team_alias": _as_object(t.team_alias)} for t in team_aliases} # Check if user is team admin or has /team/daily/activity permission # If not, filter by user's API keys. diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 44abc56713f..a2c50590dd5 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -16,9 +16,22 @@ import json import os import re import secrets +from collections.abc import Mapping, Sequence from copy import deepcopy from html import escape -from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + NoReturn, + Optional, + Protocol, + TypeVar, + Union, + cast, + overload, +) from urllib.parse import parse_qs, urlencode, urlparse if TYPE_CHECKING: @@ -155,6 +168,102 @@ _CLI_SSO_SECRET_KEY_FRAGMENTS: Final = frozenset( } ) +_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True) + + +class _PrismaTableActions(Protocol[_DbRecordT]): + async def find_unique( + self, + where: Mapping[str, object], + ) -> _DbRecordT | None: ... + + async def find_first( + self, + where: Mapping[str, object] | None = None, + ) -> _DbRecordT | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + ) -> Sequence[_DbRecordT]: ... + + async def update( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> _DbRecordT: ... + + async def update_many( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> int: ... + + +class _UserMetadataRow(Protocol): + @property + def metadata(self) -> Mapping[str, object] | None: ... + + +class _HasUserMetadataTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_UserMetadataRow]": ... + + +def _user_meta_db(repo: "_HasUserMetadataTable") -> "_PrismaTableActions[_UserMetadataRow]": + return repo.table + + +class _SsoConfigRow(Protocol): + @property + def sso_settings(self) -> Mapping[str, object] | None: ... + + +class _HasSsoConfigTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_SsoConfigRow]": ... + + +def _sso_config_db(repo: "_HasSsoConfigTable") -> "_PrismaTableActions[_SsoConfigRow]": + return repo.table + + +class _TeamDetailRow(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +class _HasTeamDetailTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_TeamDetailRow]": ... + + +def _team_detail_db(repo: "_HasTeamDetailTable") -> "_PrismaTableActions[_TeamDetailRow]": + return repo.table + + +class _CustomSsoCall(Protocol): + async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ... + + +class _ServicePrincipalAssignment(Protocol): + def get(self, key: str) -> str: ... + + +class _ServicePrincipalPage(Protocol): + @overload + def get( + self, + key: Literal["value"], + default: Sequence["_ServicePrincipalAssignment"], + ) -> Sequence["_ServicePrincipalAssignment"]: ... + + @overload + def get(self, key: Literal["@odata.nextLink"]) -> str | None: ... + + +def _as_object(value: object) -> object: + return value + def _hash_cli_sso_secret(secret: str) -> str: return hashlib.sha256(secret.encode("utf-8")).hexdigest() @@ -256,7 +365,7 @@ def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict: flow = cache.get_cache(key=cache_key) if isinstance(flow, str): try: - flow = json.loads(flow) + flow = _as_object(json.loads(flow)) except ValueError: flow = None if not isinstance(flow, dict) or "poll_secret_hash" not in flow: @@ -421,7 +530,7 @@ def _flatten_cli_sso_metadata_for_poll( def build_cli_sso_attribution_metadata( result: CustomOpenID | OpenID | dict, -) -> dict[str, Any]: +) -> dict[str, object]: """ Build allowlisted, non-secret scalar attribution metadata from an SSO result. @@ -432,7 +541,7 @@ def build_cli_sso_attribution_metadata( if not claim_map: return {} - metadata: Final[dict[str, Any]] = {} + metadata: Final[dict[str, object]] = {} for source_claim, dest_key in claim_map: if not _is_safe_cli_sso_metadata_dest_key(dest_key): verbose_proxy_logger.debug("Skipping unsafe CLI SSO metadata destination key: %s", dest_key) @@ -474,14 +583,14 @@ def _merge_cli_sso_attribution_metadata( async def _persist_cli_sso_user_metadata( prisma_client: PrismaClient, user_id: str, - attribution_metadata: dict[str, Any], + attribution_metadata: dict[str, object], ) -> None: if not attribution_metadata: return try: - user_row: Final = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) - existing_metadata: dict[str, Any] = {} + user_row: Final = await _user_meta_db(UserRepository(prisma_client)).find_unique(where={"user_id": user_id}) + existing_metadata: dict[str, object] = {} if user_row is not None: row_metadata: Final = user_row.metadata if isinstance(row_metadata, dict): @@ -491,7 +600,7 @@ async def _persist_cli_sso_user_metadata( existing_metadata=existing_metadata, attribution_metadata=attribution_metadata, ) - await UserRepository(prisma_client).table.update_many( + await _user_meta_db(UserRepository(prisma_client)).update_many( where={"user_id": user_id}, data={"metadata": merged_metadata}, ) @@ -1104,7 +1213,7 @@ def generic_response_convertor( ) # Build extra_fields dict from GENERIC_USER_EXTRA_ATTRIBUTES if specified - extra_fields: dict[str, Any] | None = None + extra_fields: dict[str, object] | None = None if generic_user_extra_attributes: extra_fields = {} for attr_name in generic_user_extra_attributes.split(","): @@ -1193,7 +1302,9 @@ async def _setup_team_mappings() -> Optional["TeamMappings"]: prisma_client: Final = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) + sso_db_record: Final = await _sso_config_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) if sso_db_record and sso_db_record.sso_settings: sso_settings_dict: Final = dict(sso_db_record.sso_settings) @@ -1225,7 +1336,9 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: prisma_client: Final = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) + sso_db_record: Final = await _sso_config_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) if sso_db_record and sso_db_record.sso_settings: sso_settings_dict: Final = dict(sso_db_record.sso_settings) @@ -1273,7 +1386,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: return role_mappings -def _parse_generic_sso_headers() -> dict: +def _parse_generic_sso_headers() -> dict[str, str]: """Parse comma-separated GENERIC_SSO_HEADERS env var into a dict.""" raw: Final = os.getenv("GENERIC_SSO_HEADERS", None) if raw is None: @@ -1677,7 +1790,7 @@ def _build_sso_user_update_data( result: Union["CustomOpenID", OpenID, dict] | None, user_email: str | None, user_id: str | None, -) -> dict: +) -> dict[str, object]: """ Build the update data dictionary for SSO user upsert. @@ -1689,7 +1802,7 @@ def _build_sso_user_update_data( Returns: dict: Update data containing user_email and optionally user_role if valid """ - update_data: Final[dict] = {"user_email": normalize_email(user_email)} + update_data: Final[dict[str, object]] = {"user_email": normalize_email(user_email)} # Get SSO role from result and include if valid sso_role: Final = getattr(result, "user_role", None) @@ -1740,7 +1853,7 @@ async def _sync_user_role_from_jwt_role_map( # Update existing DB record if role differs if user_info is not None and user_info.user_role != mapped_role.value: - await UserRepository(prisma_client).table.update( + await _user_meta_db(UserRepository(prisma_client)).update( where={"user_id": user_info.user_id}, data={"user_role": mapped_role.value}, ) @@ -1796,7 +1909,7 @@ async def check_and_update_if_proxy_admin_id(user_role: str, user_id: str, prism return user_role if prisma_client: - await UserRepository(prisma_client).table.update( + await _user_meta_db(UserRepository(prisma_client)).update( where={"user_id": user_id}, data={"user_role": LitellmUserRoles.PROXY_ADMIN.value}, ) @@ -2016,10 +2129,11 @@ async def _build_cli_sso_user_defined_values( ) -> SSOUserDefinedValues | None: from litellm.proxy.proxy_server import user_custom_sso + custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso user_id: Final = parsed_openid_result.get("user_id") - if user_custom_sso is not None: - if inspect.iscoroutinefunction(user_custom_sso): - return await user_custom_sso(result) + if custom_sso_handler is not None: + if inspect.iscoroutinefunction(custom_sso_handler): + return await custom_sso_handler(result) raise ValueError("user_custom_sso must be a coroutine function") if user_id is None: return None @@ -2035,12 +2149,14 @@ async def _build_cli_sso_user_defined_values( async def _fetch_cli_sso_team_details( prisma_client: PrismaClient, - teams: list[str], -) -> list[dict[str, Any]]: - team_details: Final[list[dict[str, Any]]] = [] + teams: Sequence[str], +) -> list[dict[str, object]]: + team_details: Final[list[dict[str, object]]] = [] try: if teams: - prisma_teams: Final = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": teams}}) + prisma_teams: Final = await _team_detail_db(TeamRepository(prisma_client)).find_many( + where={"team_id": {"in": teams}} + ) for team_row in prisma_teams: team_dict = team_row.model_dump() team_details.append( @@ -2257,12 +2373,12 @@ async def cli_poll_key( verbose_proxy_logger.info("Returning teams list for user %s to select from: %s", user_id, user_teams) # Best-effort construction of team_details if it wasn't # already cached for some reason. - team_details_response: list[dict[str, Any]] | None = None + team_details_response: list[dict[str, object]] | None = None if isinstance(user_team_details, list) and user_team_details: team_details_response = user_team_details elif user_teams: team_details_response = [{"team_id": t, "team_alias": None} for t in user_teams] - poll_response: dict[str, Any] = { + poll_response: dict[str, object] = { "status": "ready", "user_id": user_id, "teams": user_teams, @@ -2997,7 +3113,9 @@ class SSOAuthenticationHandler: user_id=user_id, ) - await UserRepository(prisma_client).table.update_many(where={"user_id": user_id}, data=update_data) + await _user_meta_db(UserRepository(prisma_client)).update_many( + where={"user_id": user_id}, data=update_data + ) else: verbose_proxy_logger.info("user not in DB, inserting user into LiteLLM DB") # user not in DB, insert User into LiteLLM DB @@ -3089,7 +3207,9 @@ class SSOAuthenticationHandler: code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) try: - team_obj: Final = await TeamRepository(prisma_client).table.find_first(where={"team_id": litellm_team_id}) + team_obj: Final = await _team_detail_db(TeamRepository(prisma_client)).find_first( + where={"team_id": litellm_team_id} + ) verbose_proxy_logger.debug("Team object: %s", team_obj) # only create a new team if it doesn't exist @@ -3278,9 +3398,10 @@ class SSOAuthenticationHandler: # But if it is, we want their models preferences user_defined_values: SSOUserDefinedValues | None = None - if user_custom_sso is not None: - if inspect.iscoroutinefunction(user_custom_sso): - user_defined_values = await user_custom_sso(result) + custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso + if custom_sso_handler is not None: + if inspect.iscoroutinefunction(custom_sso_handler): + user_defined_values = await custom_sso_handler(result) else: raise ValueError("user_custom_sso must be a coroutine function") elif user_id is not None: @@ -3448,7 +3569,7 @@ class SSOAuthenticationHandler: dict: Token exchange parameters """ # Prepare token exchange parameters (may add code_verifier: str later) - token_params: Final[dict[str, Any]] = {"include_client_id": generic_include_client_id} + token_params: Final[dict[str, object]] = {"include_client_id": generic_include_client_id} # Retrieve PKCE code_verifier if PKCE was used in authorization. # Gate on GENERIC_CLIENT_USE_PKCE to avoid an unnecessary Redis round-trip @@ -3663,7 +3784,7 @@ class SSOAuthenticationHandler: access_token string. Raises ProxyException on any validation failure. """ try: - token_response_raw: Final = response.json() + token_response_raw: Final[object] = _as_object(response.json()) except Exception as json_err: verbose_proxy_logger.error( "Failed to parse token response as JSON: %s. Body: %s", @@ -4253,7 +4374,7 @@ class MicrosoftSSOHandler: while next_link is not None and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES: response = await async_client.get(next_link, headers=headers) - response_json = response.json() + response_json: _ServicePrincipalPage = response.json() verbose_proxy_logger.debug("Response from service principal app role assigned to: %s", response_json) for _object in response_json.get("value", []): diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 9f3c9d9db45..56e986c89cf 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -4,7 +4,7 @@ import re from collections.abc import Mapping from dataclasses import dataclass, field from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, runtime_checkable from litellm.repositories.table_repositories import ( ManagedFileRepository, @@ -22,6 +22,21 @@ if TYPE_CHECKING: from litellm.types.utils import LiteLLMBatch +@runtime_checkable +class ManagedResourceAccessChecker(Protocol): + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: "UserAPIKeyAuth", + ) -> bool: ... + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: "UserAPIKeyAuth", + ) -> bool: ... + + def _is_base64_encoded_unified_file_id(b64_uid: str) -> str | Literal[False]: # Ensure b64_uid is a string and not a mock object if not isinstance(b64_uid, str): @@ -881,6 +896,65 @@ def validate_managed_files_requirement( ) +async def validate_managed_id_requirement( + resource_id: str | None, + resource_kind: Literal["file", "batch", "fine-tuning job"], + user_api_key_dict: "UserAPIKeyAuth", + managed_files_obj: object | None, +) -> None: + """ + Enforce proxy-level managed resources on every route that accepts a provider-issued id + when ``litellm.require_managed_files`` is enabled, and authenticate managed ids against + the caller's stored ownership record. + + Ownership is only recorded for LiteLLM managed ids, so a raw provider id is forwarded to the + provider under shared credentials without any tenant check; knowing another tenant's provider + id would be enough to read, reuse, or destroy the object behind it. + + Raises: + HTTPException: 400 for a raw id, 403 for an inaccessible managed id, or 500 when + ownership validation is unavailable. + """ + from fastapi import HTTPException + + import litellm + + if litellm.require_managed_files is not True: + return + + if not resource_id: + return + + if not _is_base64_encoded_unified_file_id(resource_id): + raise HTTPException( + status_code=400, + detail=( + f"Raw provider {resource_kind} ids cannot be used when require_managed_files is enabled in " + f"litellm_settings. Use the LiteLLM managed {resource_kind} id returned when the " + f"{resource_kind} was created." + ), + ) + + if not isinstance(managed_files_obj, ManagedResourceAccessChecker): + raise HTTPException( + status_code=500, + detail="Managed resource ownership validation is unavailable.", + ) + + can_access: Final = ( + await managed_files_obj.can_user_call_unified_file_id(resource_id, user_api_key_dict) + if resource_kind == "file" + else await managed_files_obj.can_user_call_unified_object_id(resource_id, user_api_key_dict) + ) + if can_access: + return + + raise HTTPException( + status_code=403, + detail=f"The caller does not have access to this managed {resource_kind} id.", + ) + + def _extract_model_param(request: "Request", request_body: dict) -> str | None: """ Extract model parameter from request. diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index bf7aa96121a..0acaac3bf5d 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -50,6 +50,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( handle_model_based_routing, prepare_data_with_credentials, validate_managed_files_requirement, + validate_managed_id_requirement, ) from litellm.proxy.utils import ProxyLogging, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository @@ -612,6 +613,13 @@ async def get_file_content( data: dict = {"file_id": file_id} try: + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) + # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( @@ -908,6 +916,13 @@ async def get_file( data: dict = {"file_id": file_id} try: + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) + custom_llm_provider: Final = ( provider or get_custom_llm_provider_from_request_headers(request=request) @@ -1098,6 +1113,13 @@ async def delete_file( data: dict = {"file_id": file_id} try: + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) + custom_llm_provider: Final = ( provider or get_custom_llm_provider_from_request_headers(request=request) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 382df608a0c..08bb8698cac 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -4,11 +4,11 @@ import json import os from collections import Counter from collections.abc import Mapping -from typing import Any, Final +from typing import Any, Final, Protocol, TypeVar from urllib.parse import urlparse from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile -from pydantic import ConfigDict, ValidationError, create_model +from pydantic import ConfigDict, JsonValue, ValidationError, create_model from pydantic.fields import FieldInfo import litellm @@ -36,6 +36,73 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( router: Final = APIRouter() +_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True) + + +class _PrismaTableActions(Protocol[_DbRecordT]): + async def find_unique(self, where: Mapping[str, object]) -> _DbRecordT | None: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ... + + async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ... + + +class _SsoSettingsMappingRow(Protocol): + @property + def sso_settings(self) -> Mapping[str, object] | None: ... + + +class _HasSsoSettingsMappingTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_SsoSettingsMappingRow]: ... + + +def _sso_settings_mapping_db(repo: _HasSsoSettingsMappingTable) -> _PrismaTableActions[_SsoSettingsMappingRow]: + return repo.table + + +class _StoredSsoSettingsRow(Protocol): + @property + def sso_settings(self) -> object: ... + + +class _HasStoredSsoSettingsTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_StoredSsoSettingsRow]: ... + + +def _stored_sso_settings_db(repo: _HasStoredSsoSettingsTable) -> _PrismaTableActions[_StoredSsoSettingsRow]: + return repo.table + + +class _UiSettingsRow(Protocol): + @property + def ui_settings(self) -> str | Mapping[str, JsonValue] | None: ... + + +class _HasUiSettingsTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_UiSettingsRow]: ... + + +def _ui_settings_db(repo: _HasUiSettingsTable) -> _PrismaTableActions[_UiSettingsRow]: + return repo.table + + +class _ConfigParamRow(Protocol): + @property + def param_value(self) -> str | Mapping[str, object] | None: ... + + +class _HasConfigParamTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_ConfigParamRow]: ... + + +def _config_param_db(repo: _HasConfigParamTable) -> _PrismaTableActions[_ConfigParamRow]: + return repo.table + + # Maps each UIThemeConfig field to the env var the UI branding path reads it # from. /update/ui_theme_settings writes both the stored ui_theme_config and # these env vars, so /get/ui_theme_settings resolves the same env vars to @@ -54,7 +121,7 @@ def _is_public_http_url(value: str | None) -> bool: return parsed.scheme in ("http", "https") and bool(parsed.netloc) -def _resolve_ui_theme_field(stored_values: Mapping[str, Any], field_name: str) -> str | None: +def _resolve_ui_theme_field(stored_values: Mapping[str, object], field_name: str) -> str | None: """Resolve one UI theme field to the value the branding path actually uses. The stored ui_theme_config wins; a field absent or blank there falls back to @@ -263,7 +330,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [ # include generics like ``Optional[int]`` / ``List[str]`` that are not # instances of ``type`` — so tightening this to ``type`` would reject # valid inputs. -_EXTRA_UI_SETTINGS_FIELDS: Final[dict[str, tuple[Any, FieldInfo]]] = {} +_EXTRA_UI_SETTINGS_FIELDS: Final[dict[str, tuple[object, FieldInfo]]] = {} # Settings OSS knows about as enterprise-gated. If a caller sends one of # these keys and no extension package has registered it, the PATCH @@ -275,7 +342,7 @@ _ENTERPRISE_ONLY_UI_SETTINGS: Final[set[str]] = {"enable_projects_ui"} _EFFECTIVE_UI_SETTINGS_CLASS: type[UISettings] | None = None -def register_extra_ui_setting(name: str, annotation: Any, field: FieldInfo) -> None: +def register_extra_ui_setting(name: str, annotation: object, field: FieldInfo) -> None: """Register an additional UI settings field contributed by an extension package. ``field`` must be a ``FieldInfo`` instance — construct it directly @@ -470,7 +537,7 @@ async def delete_allowed_ip( async def _get_settings_with_schema( settings_key: str, - settings_class: Any, + settings_class: type[BaseModel], config: dict, ) -> dict: """ @@ -842,7 +909,9 @@ async def get_sso_settings(): # Resolve the effective SSO config: the stored row wins, else the process # environment, else each field's default. Unlike the legacy read path this # does not write os.environ; a GET has no business mutating the environment. - sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) + sso_db_record: Final = await _sso_settings_mapping_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) sso_db_settings: Final = dict(sso_db_record.sso_settings) if sso_db_record and sso_db_record.sso_settings else None resolved: Final = resolve_sso_config(sso_db_settings, os.environ) @@ -914,8 +983,10 @@ async def update_sso_settings( # before-snapshot has the same shape as after_value, and rely on # create_config_audit_log's secret-name redaction to mask the # *_client_secret fields before the audit row is written. - existing_sso_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) - before_sso_data: dict[str, Any] | None = None + existing_sso_record: Final = await _stored_sso_settings_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) + before_sso_data: dict[str, JsonValue] | None = None if existing_sso_record and existing_sso_record.sso_settings: stored = existing_sso_record.sso_settings if isinstance(stored, str): @@ -948,7 +1019,7 @@ async def update_sso_settings( encrypted_sso_data: Final = proxy_config._encrypt_env_variables(environment_variables=sso_data) # Save to dedicated SSO table - await SSOConfigRepository(prisma_client).table.upsert( + await _stored_sso_settings_db(SSOConfigRepository(prisma_client)).upsert( where={"id": "sso_config"}, data={ "create": { @@ -974,7 +1045,7 @@ async def update_sso_settings( # Remove SSO-related env vars from config.environment_variables try: - env_var_entry: Final = await ConfigRepository(prisma_client).table.find_unique( + env_var_entry: Final = await _config_param_db(ConfigRepository(prisma_client)).find_unique( where={"param_name": "environment_variables"} ) @@ -982,7 +1053,7 @@ async def update_sso_settings( if env_var_entry is not None: if env_var_entry.param_value is not None: if isinstance(env_var_entry.param_value, str): - environment_variables = json.loads(env_var_entry.param_value) + environment_variables: Mapping[str, object] = json.loads(env_var_entry.param_value) else: environment_variables = dict(env_var_entry.param_value) else: @@ -993,7 +1064,7 @@ async def update_sso_settings( key: value for key, value in environment_variables.items() if key not in env_vars_to_remove } - await ConfigRepository(prisma_client).table.update( + await _config_param_db(ConfigRepository(prisma_client)).update( where={"param_name": "environment_variables"}, data={ "param_value": json.dumps(filtered_env_vars, default=str), @@ -1239,8 +1310,10 @@ async def get_ui_settings_cached() -> dict[str, Any]: if prisma_client is None: return {} - db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) - ui_settings: dict[str, Any] = {} + db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( + where={"id": "ui_settings"} + ) + ui_settings: dict[str, JsonValue] = {} if db_record and db_record.ui_settings: raw: Final = db_record.ui_settings ui_settings = json.loads(raw) if isinstance(raw, str) else dict(raw) @@ -1272,9 +1345,11 @@ async def get_ui_settings(): detail={"error": "Database not connected. Please connect a database."}, ) - ui_settings: dict[str, Any] = {} + ui_settings: Mapping[str, JsonValue] = {} - db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) + db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( + where={"id": "ui_settings"} + ) if db_record and db_record.ui_settings: ui_settings_json: Final = db_record.ui_settings @@ -1300,7 +1375,7 @@ async def get_ui_settings(): await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL) # Build config-like object for schema helper - config: Final[dict[str, Any]] = {"litellm_settings": {"ui_settings": ui_settings}} + config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": ui_settings}} return await _get_settings_with_schema( settings_key="ui_settings", @@ -1315,7 +1390,7 @@ async def get_ui_settings(): dependencies=[Depends(user_api_key_auth)], ) async def update_ui_settings( - settings_body: dict[str, Any] = Body(...), + settings_body: dict[str, object] = Body(...), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -1352,7 +1427,7 @@ async def update_ui_settings( raise HTTPException(status_code=422, detail=e.errors()) # Only include fields the caller actually sent (not Pydantic defaults). - settings_dict: Final = settings.model_dump(exclude_unset=True) + settings_dict: Final[Mapping[str, JsonValue]] = settings.model_dump(exclude_unset=True) # Reject enterprise-only settings up front so the caller gets a clear # signal instead of a silent drop. @@ -1373,15 +1448,17 @@ async def update_ui_settings( # Merge with existing persisted settings so a partial PATCH doesn't # overwrite fields the caller didn't send. - existing: dict = {} - db_existing: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) + existing: dict[str, JsonValue] = {} + db_existing: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( + where={"id": "ui_settings"} + ) if db_existing and db_existing.ui_settings: raw: Final = db_existing.ui_settings existing = json.loads(raw) if isinstance(raw, str) else dict(raw) ui_settings: Final = {**existing, **incoming} - await UISettingsRepository(prisma_client).table.upsert( + await _ui_settings_db(UISettingsRepository(prisma_client)).upsert( where={"id": "ui_settings"}, data={ "create": { diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 6b40281198c..2b037bef795 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -10,10 +10,17 @@ All /vector_store management endpoints import copy import json -from typing import Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import APIRouter, Depends, HTTPException +if TYPE_CHECKING: + from prisma.models import LiteLLM_ManagedVectorStoresTable as _VectorStoreRow + + from litellm.proxy.utils import PrismaClient + from litellm.router import Router + import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -43,6 +50,25 @@ from litellm.vector_stores.vector_store_registry import VectorStoreRegistry router: Final = APIRouter() + +class _VectorStoreTableActions(Protocol): + async def find_unique(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ... + + async def create(self, data: Mapping[str, object]) -> "_VectorStoreRow": ... + + async def update(self, where: Mapping[str, str], data: Mapping[str, object]) -> "_VectorStoreRow": ... + + async def delete(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ... + + +def _vector_store_table(prisma_client: "PrismaClient") -> _VectorStoreTableActions: + return ManagedVectorStoresRepository(prisma_client).table + + +def _row_to_vector_store(row: "_VectorStoreRow") -> LiteLLM_ManagedVectorStore: + return LiteLLM_ManagedVectorStore(**row.model_dump()) + + _LITELLM_PARAMS_MASKER: Final = SensitiveDataMasker() @@ -117,22 +143,20 @@ def _redact_sensitive_litellm_params(litellm_params: Any, _depth: int = 0) -> An async def _fetch_and_authorize_vector_store( vector_store_id: str, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, + prisma_client: "PrismaClient", ) -> "LiteLLM_ManagedVectorStore": """ Look up a vector store by id and confirm the caller can access it. Raises HTTPException(404) on miss and HTTPException(403) on access denial. """ - row: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique( - where={"vector_store_id": vector_store_id} - ) + row: Final = await _vector_store_table(prisma_client).find_unique(where={"vector_store_id": vector_store_id}) if row is None: raise HTTPException( status_code=404, detail=f"Vector store with ID {vector_store_id} not found", ) - typed: Final = LiteLLM_ManagedVectorStore(**row.model_dump()) + typed: Final = _row_to_vector_store(row) if not await _check_vector_store_access(typed, user_api_key_dict): raise HTTPException( status_code=403, @@ -141,7 +165,7 @@ async def _fetch_and_authorize_vector_store( return typed -def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, Any] | None: +def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, object] | None: """ Resolve embedding config from router's config-defined models. @@ -177,7 +201,7 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d litellm_params = deployment.litellm_params # Build embedding config from model params - embedding_config: dict[str, Any] = {} + embedding_config: dict[str, object] = {} # Extract api_key api_key = getattr(litellm_params, "api_key", None) @@ -217,7 +241,9 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d return None -async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client) -> dict[str, Any] | None: +async def _resolve_embedding_config_from_db( + embedding_model: str, prisma_client: "PrismaClient" +) -> dict[str, object] | None: """ Resolve embedding config from database model configuration. @@ -307,7 +333,9 @@ async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client) return None -async def _resolve_embedding_config(embedding_model: str, prisma_client, llm_router=None) -> dict[str, Any] | None: +async def _resolve_embedding_config( + embedding_model: str, prisma_client: "PrismaClient | None", llm_router: "Router | None" = None +) -> dict[str, object] | None: """ Resolve embedding config from either router (config-defined) or database models. @@ -388,7 +416,7 @@ async def _check_vector_store_access( async def create_vector_store_in_db( vector_store_id: str, custom_llm_provider: str, - prisma_client, + prisma_client: "PrismaClient | None", vector_store_name: str | None = None, vector_store_description: str | None = None, vector_store_metadata: dict | None = None, @@ -417,7 +445,7 @@ async def create_vector_store_in_db( raise HTTPException(status_code=500, detail="Database not connected") # Check if vector store already exists - existing_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique( + existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique( where={"vector_store_id": vector_store_id} ) if existing_vector_store is not None: @@ -427,7 +455,7 @@ async def create_vector_store_in_db( ) # Prepare data for database - data_to_create: Final[dict[str, Any]] = { + data_to_create: Final[dict[str, object]] = { "vector_store_id": vector_store_id, "custom_llm_provider": custom_llm_provider, } @@ -463,9 +491,9 @@ async def create_vector_store_in_db( data_to_create["litellm_params"] = safe_dumps({}) # Create in database - _new_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.create(data=data_to_create) + _new_vector_store: Final = await _vector_store_table(prisma_client).create(data=data_to_create) - new_vector_store: Final[LiteLLM_ManagedVectorStore] = LiteLLM_ManagedVectorStore(**_new_vector_store.model_dump()) + new_vector_store: Final[LiteLLM_ManagedVectorStore] = _row_to_vector_store(_new_vector_store) # Add vector store to registry if litellm.vector_store_registry is not None: @@ -682,12 +710,12 @@ async def delete_vector_store( memory_vector_store_exists = False vector_store_to_check = None - existing_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique( + existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique( where={"vector_store_id": data.vector_store_id} ) if existing_vector_store is not None: db_vector_store_exists = True - vector_store_to_check = LiteLLM_ManagedVectorStore(**existing_vector_store.model_dump()) + vector_store_to_check = _row_to_vector_store(existing_vector_store) # Check in-memory registry if litellm.vector_store_registry is not None: @@ -715,9 +743,7 @@ async def delete_vector_store( # Delete from database if exists if db_vector_store_exists: - await ManagedVectorStoresRepository(prisma_client).table.delete( - where={"vector_store_id": data.vector_store_id} - ) + await _vector_store_table(prisma_client).delete(where={"vector_store_id": data.vector_store_id}) # Delete from in-memory registry if exists if memory_vector_store_exists and litellm.vector_store_registry is not None: @@ -829,7 +855,7 @@ async def update_vector_store( try: update_data: Final = data.model_dump(exclude_unset=True) - vector_store_id: Final = update_data.pop("vector_store_id") + vector_store_id: Final[str] = update_data.pop("vector_store_id") # Per-store access control: anyone authenticated who passes the # premium-feature gate could otherwise update *any* vector store — @@ -857,12 +883,12 @@ async def update_vector_store( update_data["litellm_params"] = safe_dumps(litellm_params_dict) # Update in database - updated: Final = await ManagedVectorStoresRepository(prisma_client).table.update( + updated: Final = await _vector_store_table(prisma_client).update( where={"vector_store_id": vector_store_id}, data=update_data, ) - updated_vs: Final = LiteLLM_ManagedVectorStore(**updated.model_dump()) + updated_vs: Final = _row_to_vector_store(updated) # Immediately update in-memory registry to keep it in sync if litellm.vector_store_registry is not None: diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index 896b7ca33d7..c9b89bcd390 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -30,10 +30,12 @@ if TYPE_CHECKING: router: Final = APIRouter() -def _update_request_data_with_managed_file_id( +async def _update_request_data_with_managed_file_id( data: dict, file_id: str, request: Request, + user_api_key_dict: UserAPIKeyAuth, + managed_files_obj: object | None, llm_router: Optional["Router"] = None, ) -> tuple[dict, str | None]: """ @@ -65,6 +67,16 @@ def _update_request_data_with_managed_file_id( is_base64_encoded_unified_id, parse_unified_id, ) + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_id_requirement, + ) + + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=managed_files_obj, + ) # First, check if this is a unified managed file ID (base64 encoded) decoded_id: Final = is_base64_encoded_unified_id(file_id) @@ -509,8 +521,13 @@ async def vector_store_file_create( # Handle managed file IDs if present in request body original_managed_file_id = None if "file_id" in data: - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=data["file_id"], request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=data["file_id"], + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -707,8 +724,13 @@ async def vector_store_file_retrieve( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -809,8 +831,13 @@ async def vector_store_file_content( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -911,8 +938,13 @@ async def vector_store_file_update( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -1013,8 +1045,13 @@ async def vector_store_file_delete( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index f54023836e5..b2d065ea23b 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -4,8 +4,8 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion import json import re -from collections.abc import Sequence -from typing import Any, Final, Literal, cast +from collections.abc import Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable from openai.types.chat.chat_completion_named_tool_choice_param import ( ChatCompletionNamedToolChoiceParam, @@ -16,6 +16,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_create_params import ResponseInputParam from openai.types.responses.tool_param import FunctionToolParam +from pydantic import TypeAdapter from typing_extensions import TypedDict from litellm._logging import verbose_logger @@ -78,9 +79,35 @@ from .custom_tools import ( unwrap_custom_tool_arguments, ) +if TYPE_CHECKING: + from openai.types.responses.response_apply_patch_tool_call import ( + ResponseApplyPatchToolCall, + ) + ########### Initialize Classes used for Responses API ########### TOOL_CALLS_CACHE: Final = InMemoryCache() +_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) +_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) +_TEXT_ADAPTER: Final = TypeAdapter(str) + + +@runtime_checkable +class _SupportsIter(Protocol): + def __iter__(self) -> Iterator[object]: ... + + +@runtime_checkable +class _HasToolCalls(Protocol): + tool_calls: object + + +@runtime_checkable +class _HasId(Protocol): + id: object + class ChatCompletionSession(TypedDict, total=False): messages: list[ @@ -205,7 +232,7 @@ class LiteLLMCompletionResponsesConfig: responses_api_request: ResponsesAPIOptionalRequestParams, custom_llm_provider: str | None = None, stream: bool | None = None, - extra_headers: dict[str, Any] | None = None, + extra_headers: Mapping[str, object] | None = None, **kwargs, ) -> dict: """ @@ -462,7 +489,9 @@ class LiteLLMCompletionResponsesConfig: if not chat_completion_messages: continue - deduped_in_place: list[Any] = [] + deduped_in_place: list[ + AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage + ] = [] for m in chat_completion_messages: role = "" if isinstance(m, dict): @@ -472,7 +501,7 @@ class LiteLLMCompletionResponsesConfig: # Drop assistant tool_calls wrappers if we already have this call_id if role == "assistant": - tool_calls: Any = ( + tool_calls: object = ( m.get("tool_calls") if isinstance(m, dict) else getattr(m, "tool_calls", None) ) call_id = "" @@ -534,7 +563,7 @@ class LiteLLMCompletionResponsesConfig: call_id = "" if role == "assistant": - tool_calls: Any = None + tool_calls: object = None if isinstance(tool_call_message, dict): tool_calls = tool_call_message.get("tool_calls") else: @@ -578,7 +607,16 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _find_previous_assistant_idx(messages: list[Any], current_idx: int) -> int | None: + def _find_previous_assistant_idx( + messages: Sequence[ + AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message + ], + current_idx: int, + ) -> int | None: """Find the index of the previous assistant message.""" for j in range(current_idx - 1, -1, -1): if messages[j].get("role") == "assistant": @@ -586,7 +624,18 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _recover_tool_call_id_from_assistant(assistant_message: Any, message: Any) -> str: + def _recover_tool_call_id_from_assistant( + assistant_message: AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message, + message: AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message, + ) -> str: """Try to recover empty tool_call_id from assistant message's tool_calls.""" tool_calls_raw: Final = ( assistant_message.get("tool_calls") @@ -594,17 +643,23 @@ class LiteLLMCompletionResponsesConfig: else getattr(assistant_message, "tool_calls", None) ) if tool_calls_raw and isinstance(tool_calls_raw, list) and len(tool_calls_raw) > 0: - first_tool_call: Final = tool_calls_raw[0] + first_tool_call: Final = _OBJECT_LIST_ADAPTER.validate_python(tool_calls_raw)[0] if isinstance(first_tool_call, dict): - tool_call_id_raw = first_tool_call.get("id", "") + tool_call_id_raw = _ANY_KEY_DICT_ADAPTER.validate_python(first_tool_call).get("id", "") return str(tool_call_id_raw) if tool_call_id_raw is not None else "" - elif hasattr(first_tool_call, "id"): - tool_call_id_raw = getattr(first_tool_call, "id", None) + elif isinstance(first_tool_call, _HasId): + tool_call_id_raw = first_tool_call.id return str(tool_call_id_raw) if tool_call_id_raw is not None else "" return "" @staticmethod - def _get_tool_calls_list(assistant_message: Any) -> list[Any]: + def _get_tool_calls_list( + assistant_message: AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message, + ) -> Sequence[object]: """Extract tool_calls as a list from assistant message.""" tool_calls_raw: Final = ( assistant_message.get("tool_calls") @@ -614,18 +669,18 @@ class LiteLLMCompletionResponsesConfig: if tool_calls_raw is None: return [] if isinstance(tool_calls_raw, list): - return tool_calls_raw - if hasattr(tool_calls_raw, "__iter__") and not isinstance(tool_calls_raw, (str, bytes)): + return _OBJECT_LIST_ADAPTER.validate_python(tool_calls_raw) + if isinstance(tool_calls_raw, _SupportsIter) and not isinstance(tool_calls_raw, (str, bytes)): return list(tool_calls_raw) return [] @staticmethod - def _check_tool_call_exists(tool_calls: list[Any], tool_call_id: str) -> bool: + def _check_tool_call_exists(tool_calls: Sequence[object], tool_call_id: str) -> bool: """Check if a tool_call with the given ID exists in the list.""" for tool_call in tool_calls: - tool_call_id_to_check: str | None = None + tool_call_id_to_check: object = None if isinstance(tool_call, dict): - tool_call_id_to_check = tool_call.get("id") + tool_call_id_to_check = _ANY_KEY_DICT_ADAPTER.validate_python(tool_call).get("id") elif hasattr(tool_call, "id"): tool_call_id_to_check = getattr(tool_call, "id", None) if tool_call_id_to_check == tool_call_id: @@ -633,12 +688,13 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: list[Any]) -> dict[str, Any] | None: + def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: Sequence[object]) -> dict[str, object] | None: """Reconstruct a minimal tool_call definition from tools list.""" for tool in tools: if isinstance(tool, dict): - tool_function = tool.get("function") or {} - tool_name = tool_function.get("name") or tool.get("name") or "" + tool_map = _ANY_KEY_DICT_ADAPTER.validate_python(tool) + tool_function = _ANY_KEY_DICT_ADAPTER.validate_python(tool_map.get("function") or {}) + tool_name = tool_function.get("name") or tool_map.get("name") or "" if tool_name: return { "id": tool_call_id, @@ -651,7 +707,7 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _get_mapping_or_attr_value(obj: Any, key: str, default: Any = None) -> Any: + def _get_mapping_or_attr_value(obj: object, key: str, default: object = None) -> object: """ Safely read a field from dict-like or attribute-based objects. """ @@ -659,7 +715,7 @@ class LiteLLMCompletionResponsesConfig: return default if isinstance(obj, dict): - return obj.get(key, default) + return _ANY_KEY_DICT_ADAPTER.validate_python(obj).get(key, default) getter: Final = getattr(obj, "get", None) if callable(getter): @@ -672,13 +728,13 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _create_tool_call_chunk( - tool_use_definition: dict[str, Any], tool_call_id: str, index: int + tool_use_definition: Mapping[object, object], tool_call_id: str, index: int ) -> ChatCompletionToolCallChunk: """Create a ChatCompletionToolCallChunk from tool_use_definition.""" function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "function") function_name_raw: Final = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "name") function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "arguments") - function: Final[dict[str, Any]] = { + function: Final[dict[str, object]] = { "name": function_name_raw or "", "arguments": function_arguments_raw or "{}", } @@ -697,7 +753,7 @@ class LiteLLMCompletionResponsesConfig: ) @staticmethod - def _normalize_tool_use_definition(tool_use_definition: Any, tool_call_id: str) -> dict[str, Any] | None: + def _normalize_tool_use_definition(tool_use_definition: object, tool_call_id: str) -> dict[object, object] | None: """ Normalize cached tool_call definitions to a dict-like shape consumed by _create_tool_call_chunk. """ @@ -705,7 +761,7 @@ class LiteLLMCompletionResponsesConfig: return None if isinstance(tool_use_definition, dict): - normalized_definition: dict[str, Any] = dict(tool_use_definition) + normalized_definition: dict[object, object] = _ANY_KEY_DICT_ADAPTER.validate_python(tool_use_definition) else: tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "id") tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "type") @@ -738,7 +794,7 @@ class LiteLLMCompletionResponsesConfig: return normalized_definition @staticmethod - def _add_tool_call_to_assistant(assistant_message: Any, tool_call_chunk: ChatCompletionToolCallChunk) -> None: + def _add_tool_call_to_assistant(assistant_message: object, tool_call_chunk: ChatCompletionToolCallChunk) -> None: """Add a tool_call to an assistant message.""" if isinstance(assistant_message, dict): prev_assistant_dict: Final = cast(dict[str, Any], assistant_message) @@ -747,7 +803,7 @@ class LiteLLMCompletionResponsesConfig: tool_calls_list: Final = prev_assistant_dict["tool_calls"] if isinstance(tool_calls_list, list): tool_calls_list.append(tool_call_chunk) - elif hasattr(assistant_message, "tool_calls"): + elif isinstance(assistant_message, _HasToolCalls): if assistant_message.tool_calls is None: assistant_message.tool_calls = [] if isinstance(assistant_message.tool_calls, list): @@ -762,7 +818,7 @@ class LiteLLMCompletionResponsesConfig: | ChatCompletionMessageToolCall | Message ], - tools: list[Any] | None = None, + tools: Sequence[object] | None = None, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -851,7 +907,7 @@ class LiteLLMCompletionResponsesConfig: tool_calls = LiteLLMCompletionResponsesConfig._get_tool_calls_list(prev_assistant) if not LiteLLMCompletionResponsesConfig._check_tool_call_exists(tool_calls, tool_call_id): - _tool_use_definition = TOOL_CALLS_CACHE.get_cache(key=tool_call_id) + _tool_use_definition: object = TOOL_CALLS_CACHE.get_cache(key=tool_call_id) if not _tool_use_definition and tools: _tool_use_definition = LiteLLMCompletionResponsesConfig._reconstruct_tool_call_from_tools( @@ -908,7 +964,7 @@ class LiteLLMCompletionResponsesConfig: function_call=input_item ) else: - content: Final = input_item.get("content") + content: Final[object] = input_item.get("content") # Handle None content: Responses API allows None content, but GenericChatCompletionMessage requires content # Since guardrails skip None content anyway, we return empty list to exclude it from structured messages if content is None: @@ -923,7 +979,7 @@ class LiteLLMCompletionResponsesConfig: ] @staticmethod - def _is_input_item_tool_call_output(input_item: Any) -> bool: + def _is_input_item_tool_call_output(input_item: Mapping[str, object]) -> bool: """ Check if the input item is a tool call output """ @@ -936,7 +992,7 @@ class LiteLLMCompletionResponsesConfig: ] @staticmethod - def _is_input_item_function_call(input_item: Any) -> bool: + def _is_input_item_function_call(input_item: Mapping[str, object]) -> bool: """ Check if the input item is a function call or custom tool call. Both need to be reconstructed as assistant tool_calls for Chat @@ -946,7 +1002,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_tool_call_output_to_chat_completion_message( - tool_call_output: dict[str, Any], + tool_call_output: Mapping[str, object], ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ ChatCompletionToolMessage is used to indicate the output from a tool call @@ -958,7 +1014,7 @@ class LiteLLMCompletionResponsesConfig: return [] def _normalize_function_call_output_to_tool_content( - output: Any, + output: object, ) -> Any: """ Normalize Responses API function_call_output.output into a shape that downstream @@ -981,7 +1037,7 @@ class LiteLLMCompletionResponsesConfig: # Some adapters represent tool output as a list of "input_*" parts if isinstance(output, list): - normalized_blocks: Final[list[dict[str, Any]]] = [] + normalized_blocks: Final[list[dict[str, object]]] = [] text_acc: Final[list[str]] = [] for part in output: if not isinstance(part, dict): @@ -1082,7 +1138,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_function_call_to_chat_completion_message( - function_call: dict[str, Any], + function_call: Mapping[str, str], ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ Transform a Responses API function_call into a Chat Completion message with tool calls @@ -1127,7 +1183,7 @@ class LiteLLMCompletionResponsesConfig: return [chat_completion_response_message] @staticmethod - def _resolve_file_id(item: dict[str, Any]) -> str | None: + def _resolve_file_id(item: Mapping[str, object]) -> object: """ Return the effective file_id for a Responses API input_file item. Explicit file_id takes precedence; file_url is used as fallback so @@ -1136,7 +1192,7 @@ class LiteLLMCompletionResponsesConfig: return item.get("file_id") or item.get("file_url") or None @staticmethod - def _transform_input_file_item_to_file_item(item: dict[str, Any]) -> dict[str, Any]: + def _transform_input_file_item_to_file_item(item: Mapping[str, object]) -> dict[str, object]: """ Transform a Responses API input_file item to a Chat Completion file item @@ -1146,21 +1202,21 @@ class LiteLLMCompletionResponsesConfig: Returns: Dictionary with transformed file structure for Chat Completion """ - file_dict: Final[dict[str, Any]] = {} + file_dict: Final[dict[str, object]] = {} file_id: Final = LiteLLMCompletionResponsesConfig._resolve_file_id(item) if file_id: file_dict["file_id"] = file_id if item.get("file_data"): file_dict["file_data"] = item["file_data"] - new_item: Final[dict[str, Any]] = {"type": "file", "file": file_dict} + new_item: Final[dict[str, object]] = {"type": "file", "file": file_dict} if "cache_control" in item: new_item["cache_control"] = item["cache_control"] return new_item @staticmethod def _transform_input_image_item_to_image_item( - item: dict[str, Any], + item: Mapping[str, str], ) -> ChatCompletionImageObject: """ Transform a Responses API input_image item to a Chat Completion image item @@ -1173,8 +1229,8 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_content_to_chat_completion_content( - content: Any, - ) -> str | list[str | dict[str, Any]]: + content: object, + ) -> str | list[str | dict[str, object]]: """ Transform a Responses API content into a Chat Completion content @@ -1188,7 +1244,7 @@ class LiteLLMCompletionResponsesConfig: elif isinstance(content, str): return content elif isinstance(content, list): - content_list: Final[list[str | dict[str, Any]]] = [] + content_list: Final[list[str | dict[str, object]]] = [] for item in content: if isinstance(item, str): content_list.append(item) @@ -1198,8 +1254,8 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item) ) elif item.get("type") == "input_image": - image_block = dict( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item) + image_block = _STR_KEY_DICT_ADAPTER.validate_python( + dict(LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item)) ) if "cache_control" in item: image_block["cache_control"] = item["cache_control"] @@ -1209,7 +1265,7 @@ class LiteLLMCompletionResponsesConfig: text_value = item.get("text") if text_value is None: continue - content_block: dict[str, Any] = { + content_block: dict[str, object] = { "type": LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type( item.get("type") or "text" ), @@ -1299,7 +1355,7 @@ class LiteLLMCompletionResponsesConfig: parameters = dict(typed_tool.get("parameters", {}) or {}) if not parameters or "type" not in parameters: parameters["type"] = "object" - chat_completion_tool: dict[str, Any] = { + chat_completion_tool: dict[str, object] = { "type": "function", "function": { "name": typed_tool.get("name") or "", @@ -1340,7 +1396,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def transform_chat_completion_tool_params_to_responses_api_tools( chat_completion_tools: list[ChatCompletionToolParam | OpenAIMcpServerTool] | None, - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Transform Chat Completion tool params (e.g. from guardrail output) back to Responses API request tool format. Inverse of @@ -1348,7 +1404,7 @@ class LiteLLMCompletionResponsesConfig: """ if chat_completion_tools is None or not chat_completion_tools: return [] - result: Final[list[dict[str, Any]]] = [] + result: Final[list[dict[str, object]]] = [] for tool in chat_completion_tools: if not isinstance(tool, dict): result.append(tool) @@ -1358,7 +1414,7 @@ class LiteLLMCompletionResponsesConfig: parameters = dict(fn.get("parameters", {}) or {}) if not parameters or "type" not in parameters: parameters["type"] = "object" - responses_tool: dict[str, Any] = { + responses_tool: dict[str, object] = { "type": "function", "name": fn.get("name") or "", "description": fn.get("description") or "", @@ -1510,7 +1566,7 @@ class LiteLLMCompletionResponsesConfig: def convert_response_function_tool_call_to_chat_completion_tool_call( tool_call_item: Any, index: int = 0, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Convert ResponseFunctionToolCall to ChatCompletionToolCallChunk format. @@ -1536,7 +1592,7 @@ class LiteLLMCompletionResponsesConfig: else (dict(provider_fields) if hasattr(provider_fields, "__dict__") else {}) ) - function_dict: Final[dict[str, Any]] = { + function_dict: Final[dict[str, object]] = { "name": tool_call_item.name, "arguments": tool_call_item.arguments, } @@ -1544,7 +1600,7 @@ class LiteLLMCompletionResponsesConfig: if provider_specific_fields: function_dict["provider_specific_fields"] = provider_specific_fields - tool_call_dict: Final[dict[str, Any]] = { + tool_call_dict: Final[dict[str, object]] = { "id": LiteLLMCompletionResponsesConfig._tool_call_id_from_responses_item( getattr(tool_call_item, "id", None), getattr(tool_call_item, "call_id", None), @@ -1561,9 +1617,9 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def convert_apply_patch_tool_call_to_chat_completion_tool_call( - tool_call_item: Any, + tool_call_item: "ResponseApplyPatchToolCall", index: int = 0, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Convert ResponseApplyPatchToolCall to ChatCompletionToolCallChunk format. @@ -1581,7 +1637,7 @@ class LiteLLMCompletionResponsesConfig: import json operation_dict: Final = tool_call_item.operation.model_dump() - tool_call_dict: Final[dict[str, Any]] = { + tool_call_dict: Final[dict[str, object]] = { "id": tool_call_item.call_id, "function": { "name": "apply_patch", @@ -1795,9 +1851,11 @@ class LiteLLMCompletionResponsesConfig: if not images: return image_generation_items - for idx, image_item in enumerate(images): + for idx, image_item in enumerate(_DICT_ITEMS_LIST_ADAPTER.validate_python(images)): # Extract base64 from data URL - image_url = image_item.get("image_url", {}).get("url", "") + image_url = _TEXT_ADAPTER.validate_python( + _ANY_KEY_DICT_ADAPTER.validate_python(image_item.get("image_url", {})).get("url", "") + ) base64_data = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(image_url) if base64_data: @@ -2048,8 +2106,8 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_text_format_to_response_format( - text_param: dict[str, Any] | Any, - ) -> dict[str, Any] | None: + text_param: object, + ) -> dict[str, object] | None: """ Transform Responses API text.format parameter to Chat Completion response_format parameter. diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 820839fc6bf..2e1e1a44594 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,7 +5,7 @@ import json import time import traceback import uuid -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -33,6 +33,7 @@ from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfi from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PART_UNION_TYPES, + ResponseAPIUsage, ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, @@ -112,7 +113,7 @@ _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType( def _error_event_fields(error_obj: object) -> tuple[str, str | None, str | None]: - if isinstance(error_obj, dict): + if _is_json_object(error_obj): raw_message = error_obj.get("message") raw_type = error_obj.get("type") raw_code = error_obj.get("code") @@ -243,7 +244,9 @@ class BaseResponsesAPIStreamingIterator: # Using getattr(..., "response") alone is unsafe with Mocks: they synthesize a # truthy child Mock for any attribute, which breaks tests and is wrong on stream. if "response" in parsed_chunk: - response_object: Final = getattr(openai_responses_api_chunk, "response", None) + response_object: Final[ResponsesAPIResponse | None] = getattr( + openai_responses_api_chunk, "response", None + ) if response_object is not None: response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( responses_api_response=response_object, @@ -279,7 +282,9 @@ class BaseResponsesAPIStreamingIterator: model_id=_stream_model_id, ) elif _event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE: - _part: Final = getattr(openai_responses_api_chunk, "part", None) + _part: Final[PART_UNION_TYPES | Mapping[str, object] | None] = getattr( + openai_responses_api_chunk, "part", None + ) if _part is not None: if isinstance(_part, dict): ResponsesAPIRequestUtils._encode_container_ids_in_annotations( @@ -302,7 +307,7 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): - item: Final = getattr(openai_responses_api_chunk, "item", None) + item: Final[object | None] = getattr(openai_responses_api_chunk, "item", None) if item: encrypted_content: Final = getattr(item, "encrypted_content", None) if encrypted_content and isinstance(encrypted_content, str): @@ -324,9 +329,11 @@ class BaseResponsesAPIStreamingIterator: self.completed_response = openai_responses_api_chunk # Add cost to usage object if include_cost_in_streaming_usage is True if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: - response_obj: Final[Any | None] = getattr(openai_responses_api_chunk, "response", None) + response_obj: Final[ResponsesAPIResponse | None] = getattr( + openai_responses_api_chunk, "response", None + ) if response_obj: - usage_obj: Final[Any | None] = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is not None: try: cost: float | None = self.logging_obj._response_cost_calculator(result=response_obj) @@ -414,7 +421,9 @@ class BaseResponsesAPIStreamingIterator: async_failure_handler / failure_handler so logging integrations correctly record the call as failed. """ - response_obj: Final = getattr(self.completed_response, "response", None) if self.completed_response else None + response_obj: Final[ResponsesAPIResponse | None] = ( + getattr(self.completed_response, "response", None) if self.completed_response else None + ) error_info: Final = getattr(response_obj, "error", None) if response_obj else None error_message, error_type, error_code = _error_event_fields(error_info) self._record_failed_response_usage(response_obj) @@ -429,7 +438,7 @@ class BaseResponsesAPIStreamingIterator: def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj: Final = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is None: return try: @@ -506,7 +515,7 @@ class BaseResponsesAPIStreamingIterator: return request_kwargs = getattr(caching_handler, "request_kwargs", None) - if not isinstance(request_kwargs, dict) or request_kwargs.get("stream") is not True: + if not _is_json_object(request_kwargs) or request_kwargs.get("stream") is not True: return request_kwargs = request_kwargs.copy() preset_cache_key = getattr(caching_handler, "preset_cache_key", None) @@ -606,7 +615,7 @@ class BaseResponsesAPIStreamingIterator: if self.completed_response is None: return - request_payload: Final[dict[str, Any]] = {} + request_payload: Final[dict[str, object]] = {} if isinstance(self.request_data, dict): request_payload.update(self.request_data) try: @@ -695,11 +704,15 @@ class BaseResponsesAPIStreamingIterator: pass -async def call_post_streaming_hooks_for_testing(iterator, chunk): +async def call_post_streaming_hooks_for_testing( + iterator: object, chunk: ResponsesAPIStreamingResponse +) -> ResponsesAPIStreamingResponse: """ Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped. """ - hook_fn: Final = getattr(iterator, "_call_post_streaming_deployment_hook", None) + hook_fn: Final[Callable[[ResponsesAPIStreamingResponse], Awaitable[ResponsesAPIStreamingResponse]] | None] = ( + getattr(iterator, "_call_post_streaming_deployment_hook", None) + ) if hook_fn is None: return chunk return await hook_fn(chunk) @@ -1019,7 +1032,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _dump_response_object(obj: Any) -> dict[str, Any]: if hasattr(obj, "model_dump"): return obj.model_dump() - if isinstance(obj, dict): + if _is_json_object(obj): return obj return {} @@ -1684,7 +1697,7 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final = json.loads(response_str) + evt_obj: Final[Mapping[str, object]] = json.loads(response_str) except (json.JSONDecodeError, TypeError): return response_str @@ -1925,7 +1938,7 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _extract_output_messages( - completed_event: dict[str, Any], + completed_event: dict[str, object], ) -> list[dict[str, object]]: """ Convert the output items in a ``response.completed`` event into @@ -2065,7 +2078,7 @@ class ManagedResponsesWebSocketHandler: Flat: {"type": "response.create", "input": [...], "model": "...", ...} """ nested: Final = msg_obj.get("response") - response_params: Final[dict[str, Any]] = ( + response_params: Final[dict[str, object]] = ( nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"} ) return { @@ -2076,7 +2089,7 @@ class ManagedResponsesWebSocketHandler: def _apply_history( self, - call_kwargs: dict[str, Any], + call_kwargs: dict[str, object], previous_response_id: str | None, current_messages: list[dict[str, object]], prior_history: list[dict[str, object]], @@ -2129,7 +2142,7 @@ class ManagedResponsesWebSocketHandler: return False return event_provider == self._connection_provider - def _inject_credentials(self, call_kwargs: dict[str, Any], model: str | None = None) -> None: + def _inject_credentials(self, call_kwargs: dict[str, object], model: str | None = None) -> None: """Inject connection-level credentials and metadata into call_kwargs.""" if self.api_key is not None: call_kwargs["api_key"] = self.api_key diff --git a/litellm/router.py b/litellm/router.py index 1a76eb5d59b..feaf69a44ae 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -46,6 +46,7 @@ from litellm.constants import ( DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, DEFAULT_MAX_LRU_CACHE_SIZE, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.asyncify import run_async_function @@ -135,6 +136,7 @@ from litellm.router_utils.handle_error import ( from litellm.router_utils.health_state_cache import DeploymentHealthCache from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, + warn_on_unknown_model_group_affinity_flags, ) from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( build_io_token_rate_limit_headers, @@ -603,6 +605,10 @@ class Router: # ``litellm.proxy.auth.auth_checks._is_model_cost_zero``. self._zero_cost_cache: dict[str, bool] = {} + self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds + self.model_group_affinity_config = model_group_affinity_config + warn_on_unknown_model_group_affinity_flags(model_group_affinity_config) + if model_list is not None: # set_model_list will build indices automatically self.set_model_list(model_list) @@ -744,7 +750,6 @@ class Router: litellm.failure_callback = [self.deployment_callback_on_failure] self.routing_strategy_args = routing_strategy_args self.provider_budget_config = provider_budget_config - self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds self.router_budget_logger: RouterBudgetLimiting | None = None if RouterBudgetLimiting.should_init_router_budget_limiter( model_list=model_list, provider_budget_config=self.provider_budget_config @@ -766,7 +771,6 @@ class Router: ) self.model_group_retry_policy: dict[str, RetryPolicy] | None = model_group_retry_policy - self.model_group_affinity_config: dict[str, list[str]] | None = model_group_affinity_config self.allowed_fails_policy: AllowedFailsPolicy | None = None if allowed_fails_policy is not None: @@ -789,21 +793,8 @@ class Router: # If model_group_affinity_config is set but no global affinity checks were # enabled, we still need the DeploymentAffinityCheck callback (with global # flags all False) so per-group config can activate affinity per model group. - if self.model_group_affinity_config and not any( - isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or []) - ): - if self.optional_callbacks is None: - self.optional_callbacks = [] - affinity_callback: Final = DeploymentAffinityCheck( - cache=self.cache, - ttl_seconds=self.deployment_affinity_ttl_seconds, - enable_user_key_affinity=False, - enable_responses_api_affinity=False, - enable_session_id_affinity=False, - model_group_affinity_config=self.model_group_affinity_config, - ) - self.optional_callbacks.append(affinity_callback) - litellm.logging_callback_manager.add_litellm_callback(affinity_callback) + if self.model_group_affinity_config: + self._ensure_deployment_affinity_callback() if self.alerting_config is not None: self._initialize_alerting() @@ -1662,6 +1653,28 @@ class Router: _move_before_deployment_affinity(self.optional_callbacks, ec_callback) _move_before_deployment_affinity(litellm.callbacks, ec_callback) + def _ensure_deployment_affinity_callback(self) -> None: + """Register the DeploymentAffinityCheck callback (global flags all False) if absent. + + Needed when nothing enabled a global affinity flag but affinity can still + activate per request: per-group `model_group_affinity_config` entries, or the + session-affinity marker a complexity router stamps at pre-routing time. + """ + if any(isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or [])): + return + if self.optional_callbacks is None: + self.optional_callbacks = [] + affinity_callback: Final = DeploymentAffinityCheck( + cache=self.cache, + ttl_seconds=self.deployment_affinity_ttl_seconds, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=False, + model_group_affinity_config=self.model_group_affinity_config, + ) + self.optional_callbacks.append(affinity_callback) + litellm.logging_callback_manager.add_litellm_callback(affinity_callback) + def add_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None): if optional_pre_call_checks is None: return @@ -7683,6 +7696,8 @@ class Router: strategy=complexity_router, strategy_label="Complexity-router", ) + if complexity_router._uses_deployment_pin: + self._ensure_deployment_affinity_callback() def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: """True when this deployment opts in via the `auto_router/adaptive_router` model prefix.""" @@ -11190,6 +11205,9 @@ class Router: router_strategy: Final = self._select_pre_routing_strategy(model=model, request_kwargs=request_kwargs) if router_strategy is None: self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) + self._stamp_or_clear_metadata_key( + request_kwargs=request_kwargs, key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, value=None + ) return None pre_routing_hook_response: Final = await router_strategy.async_pre_routing_hook( @@ -11203,6 +11221,11 @@ class Router: request_kwargs=request_kwargs, routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None), ) + self._stamp_or_clear_metadata_key( + request_kwargs=request_kwargs, + key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, + value=(pre_routing_hook_response.session_affinity_ttl_seconds if pre_routing_hook_response else None), + ) # `model` (the alias, e.g. "smart-router") is never the deployment actually # called - apply the alias's own litellm_params (besides `model` itself, @@ -11234,21 +11257,40 @@ class Router: to the deployment that actually served the request. Every attempt therefore writes or clears, never just writes. """ - if routing_decision is None: + Router._stamp_or_clear_metadata_key( + request_kwargs=request_kwargs, + key="routing_decision", + value=( + None + if routing_decision is None + else Router._redact_prompt_text_if_needed( + request_kwargs=request_kwargs, routing_decision=routing_decision + ) + ), + ) + + @staticmethod + def _stamp_or_clear_metadata_key(request_kwargs: dict, key: str, value: object | None) -> None: + """Write a proxy-internal metadata key for THIS routing attempt, or clear it. + + Fallbacks and retries re-enter the pre-routing hook with the same + `request_kwargs`, so every attempt must write or clear, never just write; + a value left behind by an earlier attempt would be attributed to this one. + `get_or_create_metadata_bucket` is the single owner of "which dict holds + proxy-internal metadata": it picks `litellm_metadata` when present (so the + value never lands in the `metadata` dict that routes like /v1/messages + forward to the provider) and replaces a non-dict value rather than silently + skipping the write. Clearing pops from BOTH buckets so a request whose + bucket resolution changed between attempts cannot resurrect a stale value. + """ + if value is None: for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")): if isinstance(bucket, dict): - bucket.pop("routing_decision", None) + bucket.pop(key, None) return - # `get_or_create_metadata_bucket` is the single owner of "which dict holds - # proxy-internal metadata": it picks `litellm_metadata` when present (so the - # decision never lands in the `metadata` dict that routes like /v1/messages - # forward to the provider) and replaces a non-dict value rather than silently - # skipping the write. _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) - metadata_bucket["routing_decision"] = Router._redact_prompt_text_if_needed( - request_kwargs=request_kwargs, routing_decision=routing_decision - ) + metadata_bucket[key] = value @staticmethod def _redact_prompt_text_if_needed( diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index f6ced0bb9d2..32d252f3f68 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1651,6 +1651,28 @@ class ComplexityRouter(CustomLogger): caller_scope: Final = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped" return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}" + @property + def _uses_tier_pin(self) -> bool: + return bool(self.config.session_affinity and not self.config.plugins) + + @property + def _uses_deployment_pin(self) -> bool: + """session_affinity implies the deployment pin: a session frozen onto one model + group but load-balanced across its deployments would still go cache-cold, which + is the exact failure both flags exist to prevent.""" + return bool((self.config.deployment_affinity or self.config.session_affinity) and not self.config.plugins) + + def _with_session_deployment_affinity( + self, response: PreRoutingHookResponse | None + ) -> PreRoutingHookResponse | None: + if response is None or not self._uses_deployment_pin: + return response + return response.model_copy( + update={ # mutable-ok: model_copy types update as a plain dict + "session_affinity_ttl_seconds": self.config.session_affinity_ttl_seconds + } + ) + async def async_pre_routing_hook( self, model: str, @@ -1685,7 +1707,7 @@ class ComplexityRouter(CustomLogger): resolved_messages: Final = self._resolve_messages(messages, request_kwargs) conversation_continuing: Final = _conversation_is_continuing(resolved_messages) - use_session_affinity: Final = self.config.session_affinity and not self.config.plugins + use_session_affinity: Final = self._uses_tier_pin session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None @@ -1724,17 +1746,19 @@ class ComplexityRouter(CustomLogger): "ComplexityRouter: routing decision cause=%s, routed_model=%s", cause, routed_model ) has_original_messages: Final = messages is not None and len(messages) > 0 - return PreRoutingHookResponse( - model=routed_model, - messages=messages if has_original_messages else None, - routing_decision=self._build_routing_decision( - routed_model=routed_model, - cause=cause, - tier=self._tier_for_model(routed_model), - escalation_keyword=pin_escalation_keyword, - escalated=escalated, - conversation_continuing=conversation_continuing, - ), + return self._with_session_deployment_affinity( + PreRoutingHookResponse( + model=routed_model, + messages=messages if has_original_messages else None, + routing_decision=self._build_routing_decision( + routed_model=routed_model, + cause=cause, + tier=self._tier_for_model(routed_model), + escalation_keyword=pin_escalation_keyword, + escalated=escalated, + conversation_continuing=conversation_continuing, + ), + ) ) response: Final = await self._classify_and_route( @@ -1752,7 +1776,7 @@ class ComplexityRouter(CustomLogger): value=response.model, ttl=self.config.session_affinity_ttl_seconds, ) - return response + return self._with_session_deployment_affinity(response) async def _classify_and_route( self, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 69609a973b0..0999af66fd8 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -508,13 +508,39 @@ class ComplexityRouterConfig(BaseModel): "session's first turn and reuse it for every later turn, skipping re-classification. " "Off by default so every turn is classified on its own merits and routed to the cheapest " "adequate tier. Set True to keep a multi-turn session on one model, which preserves " - "provider prompt caches and avoids cross-model conversation-history errors." + "provider prompt caches and avoids cross-model conversation-history errors. Always " + "implies the deployment pin regardless of deployment_affinity: the session sticks to " + "one deployment of the pinned model, since freezing the model while re-shuffling its " + "deployments would still go cache-cold." + ), + ) + deployment_affinity: bool = Field( + default=True, + description=( + "When True and a session_id is resolvable on the request, pin the deployment chosen " + "inside each routed model group and reuse it whenever the session returns to that " + "group, without pinning which group the session routes to. Independent of " + "session_affinity, which pins the model group instead (and always carries this " + "deployment pin with it): with session_affinity off, " + "every turn is still classified on its own merits while a session that escalates to a " + "stronger tier and comes back still lands on the deployment it used before, which is " + "what keeps a provider prompt cache warm. Pins are held per model group, so switching " + "tiers does not disturb the pin left behind in the previous group. On by default " + "because re-shuffling a conversation across deployments of the same model discards " + "that cache for no benefit; set False to keep every turn load-balanced across the " + "group, which is what a deployment set with tight per-deployment rate limits wants. " + "Inert when no session_id is resolvable, since there is nothing to key a pin on, and " + "suppressed when plugins are configured, for the same reason session_affinity is." ), ) session_affinity_ttl_seconds: int = Field( default=3600, gt=0, - description="TTL for the session affinity pin; refreshed on every cache hit", + description=( + "TTL for the session affinity pin; refreshed on every cache hit. Bounds both the " + "session_affinity model pin and the deployment_affinity deployment pin, so it measures " + "idle time for the session's routing decisions rather than total session length" + ), ) plugins: list[RoutingPlugin] | None = Field( diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index b4408d17ffb..7fb90ab89de 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -13,12 +13,15 @@ where routing to a consistent deployment is still beneficial. """ import hashlib +import json +from collections.abc import Mapping, Sequence from typing import Any, Final, cast from typing_extensions import TypedDict from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import AllMessageValues @@ -29,6 +32,47 @@ class DeploymentAffinityCacheValue(TypedDict): model_id: str +VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset( + { + "deployment_affinity", + "responses_api_deployment_check", + "session_affinity", + "encrypted_content_affinity", + } +) + + +def warn_on_unknown_model_group_affinity_flags(model_group_affinity_config: Mapping[str, Sequence[str]] | None) -> None: + """`model_group_affinity_config` is one Router-level config consumed by two callbacks: + DeploymentAffinityCheck acts on three of the flags and EncryptedContentAffinityCheck + on the fourth, so typo detection lives here at the schema, not inside either consumer. + """ + if model_group_affinity_config is None: + return + for group, flags in model_group_affinity_config.items(): + unknown = set(flags) - VALID_MODEL_GROUP_AFFINITY_FLAGS + if unknown: + verbose_router_logger.warning( + "model_group_affinity_config: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s", + unknown, + group, + VALID_MODEL_GROUP_AFFINITY_FLAGS, + ) + + +_CLAIM_PIN_SCRIPT: Final = """ +local current = redis.call('GET', KEYS[1]) +if current == false then + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +if current == ARGV[1] then + redis.call('EXPIRE', KEYS[1], ARGV[2]) +end +return current +""" + + class DeploymentAffinityCheck(CustomLogger): """ Router deployment affinity callback. @@ -38,14 +82,6 @@ class DeploymentAffinityCheck(CustomLogger): """ CACHE_KEY_PREFIX = "deployment_affinity:v1" - VALID_FLAGS = frozenset( - { - "deployment_affinity", - "responses_api_deployment_check", - "session_affinity", - "encrypted_content_affinity", - } - ) def __init__( self, @@ -63,15 +99,6 @@ class DeploymentAffinityCheck(CustomLogger): self.enable_responses_api_affinity = enable_responses_api_affinity self.enable_session_id_affinity = enable_session_id_affinity self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} - for group, flags in self.model_group_affinity_config.items(): - unknown = set(flags) - self.VALID_FLAGS - if unknown: - verbose_router_logger.warning( - "DeploymentAffinityCheck: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s", - unknown, - group, - self.VALID_FLAGS, - ) def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]: """ @@ -218,8 +245,13 @@ class DeploymentAffinityCheck(CustomLogger): return f"{cls.CACHE_KEY_PREFIX}:{model_group}:{hashed_user_key}" @classmethod - def get_session_affinity_cache_key(cls, model_group: str, session_id: str) -> str: - return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{session_id}" + def get_session_affinity_cache_key(cls, model_group: str, session_id: str, user_key: str | None) -> str: + """Session pins are scoped by the caller's hashed API key so two callers reusing + the same client-supplied session_id cannot read or steer each other's pin. + `"unscoped"` covers direct Router usage with no authenticated caller, matching + the complexity router's own session pin key.""" + hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped" + return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}" @staticmethod def _get_user_key_from_metadata_dict(metadata: dict) -> str | None: @@ -278,6 +310,97 @@ class DeploymentAffinityCheck(CustomLogger): return session_id return None + @staticmethod + def _get_marker_session_affinity_ttl(request_kwargs: dict) -> int | None: + """TTL from the session-affinity marker the Router stamps at pre-routing time + when an auto-router routed this request with session_affinity enabled. + Marker presence enables session pinning for this request only; anything that + is not a positive int is treated as absent.""" + for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs): + ttl = metadata.get(SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY) + if isinstance(ttl, int) and not isinstance(ttl, bool) and ttl > 0: + return ttl + return None + + @staticmethod + def _pinned_model_id(stored: object) -> str | None: + """Deployment id held by a stored pin, for both the dict shape this writes and the + bare string older writers left behind. None when the value is neither.""" + if isinstance(stored, dict): + model_id: Final = stored.get("model_id") + return str(model_id) if model_id is not None else None + if isinstance(stored, str): + return stored + return None + + def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: + """The one owner of authoritative local pin writes: a plain set keeps a live + key's original expiry (`allow_ttl_override`), so the entry is replaced to make + the TTL real. Every local pin write goes through here so the redis-winner sync + and the pod-local claim can never disagree about expiry again.""" + self.cache.in_memory_cache.delete_cache(cache_key) + self.cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds) + + async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None: + """First-writer-wins pin write: store `pin_value` only when the key is absent and + return the deployment id the key holds afterwards, so a caller learns whether it won + by comparing against its own id, and None when the stored value is one no reader can + interpret. Concurrent claimers converge on the + first write instead of the last. Re-claiming with the stored value refreshes its + TTL, the same keepalive the complexity router's model pin documents: an active + session must not lose its pin mid-conversation just because it outlives the + original write, so `session_affinity_ttl_seconds` bounds idle time, not total + session length. On Redis one Lua script does the get-or-set-or-refresh + atomically (same registration seam the rate limiters use) and the in-memory + tier is synchronized to the winner; without Redis, and whenever Redis is + unreachable, the pod-local check-and-set below stands in and is atomic because it + runs synchronously on the event loop. Degrading to a pod-local claim rather than + propagating the fault is what keeps same-pod stickiness through a Redis blip: the + caller only logs this result, so an escaping error would leave the session with no + pin at all and reshuffle every turn for the outage, which is worse than losing + cross-pod agreement. The redis tier is + resolved per call because the proxy attaches it after Router construction + (`Router._update_redis_cache`); the compiled script is cached per event loop + underneath the registration seam. + """ + redis_cache: Final = self.cache.redis_cache + if redis_cache is not None: + try: + claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT) + raw: Final = await claim_script(keys=(cache_key,), args=(json.dumps(pin_value), int(ttl_seconds))) + decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw + if not isinstance(decoded, str): + return pin_value["model_id"] + try: + winner: object = json.loads(decoded) + except json.JSONDecodeError: + winner = decoded + self._set_local_pin(cache_key=cache_key, value=winner, ttl_seconds=ttl_seconds) + return self._pinned_model_id(winner) + except Exception as e: # noqa: BLE001 # any Redis/Lua failure degrades to the pod-local claim, never unpins + verbose_router_logger.debug( + "DeploymentAffinityCheck: redis pin claim failed, falling back to pod-local claim. error=%s", e + ) + + return self._claim_pin_in_memory(cache_key=cache_key, pin_value=pin_value, ttl_seconds=ttl_seconds) + + def _claim_pin_in_memory( + self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int + ) -> str | None: + """Pod-local half of the claim, used when no Redis tier is attached and as the + fallback when the Redis claim fails. Mirrors the Lua script exactly, including + the keepalive: re-claiming with the stored value slides the idle window through + `_set_local_pin`. Both branches stay synchronous, hence atomic on the event + loop.""" + existing: Final = self.cache.in_memory_cache.get_cache(cache_key) + if existing is not None: + existing_model_id: Final = self._pinned_model_id(existing) + if existing_model_id == pin_value["model_id"]: + self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds) + return existing_model_id + self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds) + return pin_value["model_id"] + @staticmethod def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None: for deployment in healthy_deployments: @@ -334,12 +457,21 @@ class DeploymentAffinityCheck(CustomLogger): if stable_model_map_key is None: return typed_healthy_deployments + session_affinity_active: Final = ( + enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None + ) + user_key: Final = ( + self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs) + if (session_affinity_active or enable_user_key) + else None + ) + # 2) Session-id -> deployment affinity - if enable_session_id: + if session_affinity_active: session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs=request_kwargs) if session_id is not None: session_cache_key: Final = self.get_session_affinity_cache_key( - model_group=stable_model_map_key, session_id=session_id + model_group=stable_model_map_key, session_id=session_id, user_key=user_key ) session_cache_result: Final = await self.cache.async_get_cache(key=session_cache_key) @@ -371,7 +503,6 @@ class DeploymentAffinityCheck(CustomLogger): if not enable_user_key: return typed_healthy_deployments - user_key: Final = self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs) if user_key is None: return typed_healthy_deployments @@ -438,18 +569,22 @@ class DeploymentAffinityCheck(CustomLogger): enable_session_id, ) = self._get_effective_flags(deployment_model_name) - if not enable_user_key and not enable_session_id: + marker_session_ttl: Final = self._get_marker_session_affinity_ttl(request_kwargs=kwargs) + session_affinity_active: Final = enable_session_id or marker_session_ttl is not None + + if not enable_user_key and not session_affinity_active: return None - user_key = None - if enable_user_key: - user_key = self._get_user_key_from_request_kwargs(request_kwargs=kwargs) + user_key: Final = ( + self._get_user_key_from_request_kwargs(request_kwargs=kwargs) + if (enable_user_key or session_affinity_active) + else None + ) + session_id: Final = ( + self._get_session_id_from_request_kwargs(request_kwargs=kwargs) if session_affinity_active else None + ) - session_id = None - if enable_session_id: - session_id = self._get_session_id_from_request_kwargs(request_kwargs=kwargs) - - if user_key is None and session_id is None: + if not ((enable_user_key and user_key is not None) or session_id is not None): return None model_info = kwargs.get("model_info") @@ -473,22 +608,31 @@ class DeploymentAffinityCheck(CustomLogger): verbose_router_logger.warning("DeploymentAffinityCheck: model_id missing; skipping affinity cache update.") return None - if user_key is not None: + pin_value: Final = DeploymentAffinityCacheValue(model_id=str(model_id)) + + if enable_user_key and user_key is not None: try: cache_key: Final = self.get_affinity_cache_key(model_group=deployment_model_name, user_key=user_key) - await self.cache.async_set_cache( - cache_key, - DeploymentAffinityCacheValue(model_id=str(model_id)), - ttl=self.ttl_seconds, - ) - - verbose_router_logger.debug( - "DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s", - deployment_model_name, - model_id, - self.ttl_seconds, - self._shorten_for_logs(user_key), + claimed_user_pin: Final = await self._claim_pin( + cache_key=cache_key, + pin_value=pin_value, + ttl_seconds=self.ttl_seconds, ) + if claimed_user_pin == pin_value["model_id"]: + verbose_router_logger.debug( + "DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s", + deployment_model_name, + model_id, + self.ttl_seconds, + self._shorten_for_logs(user_key), + ) + else: + verbose_router_logger.debug( + "DeploymentAffinityCheck: affinity pin already claimed model_map_key=%s existing=%s ours=%s", + deployment_model_name, + claimed_user_pin, + model_id, + ) except Exception as e: # Non-blocking: affinity is a best-effort optimization. verbose_router_logger.debug( @@ -500,21 +644,31 @@ class DeploymentAffinityCheck(CustomLogger): # Also persist Session-ID affinity if enabled and session-id is provided if session_id is not None: try: + session_affinity_ttl: Final = marker_session_ttl if marker_session_ttl is not None else self.ttl_seconds session_cache_key: Final = self.get_session_affinity_cache_key( - model_group=deployment_model_name, session_id=session_id + model_group=deployment_model_name, session_id=session_id, user_key=user_key ) - await self.cache.async_set_cache( - session_cache_key, - DeploymentAffinityCacheValue(model_id=str(model_id)), - ttl=self.ttl_seconds, - ) - verbose_router_logger.debug( - "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s", - deployment_model_name, - model_id, - self.ttl_seconds, - session_id, + claimed_session_pin: Final = await self._claim_pin( + cache_key=session_cache_key, + pin_value=pin_value, + ttl_seconds=session_affinity_ttl, ) + if claimed_session_pin == pin_value["model_id"]: + verbose_router_logger.debug( + "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s", + deployment_model_name, + model_id, + session_affinity_ttl, + session_id, + ) + else: + verbose_router_logger.debug( + "DeploymentAffinityCheck: session pin already claimed model_map_key=%s existing=%s ours=%s session_id=%s", + deployment_model_name, + claimed_session_pin, + model_id, + session_id, + ) except Exception as e: verbose_router_logger.debug( "DeploymentAffinityCheck: failed to set session affinity cache. model_map_key=%s error=%s", diff --git a/litellm/types/router.py b/litellm/types/router.py index 4280da08cbb..e166d844735 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -816,6 +816,7 @@ class PreRoutingHookResponse(BaseModel): model: str messages: list[dict[str, Any]] | None routing_decision: StandardLoggingRoutingDecision | None = None + session_affinity_ttl_seconds: int | None = None _PreRoutingStrategyT_co = TypeVar("_PreRoutingStrategyT_co", covariant=True) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index a8af4eabb3f..fdc81fac196 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3121 + "limit": 3114 }, "ANN002": { "limit": 71 @@ -9,7 +9,7 @@ "limit": 834 }, "ANN201": { - "limit": 2032 + "limit": 2031 }, "ANN202": { "limit": 865 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1630 + "limit": 1555 }, "ASYNC230": { "limit": 11 @@ -56,9 +56,6 @@ "B026": { "limit": 3 }, - "B033": { - "limit": 0 - }, "BLE001": { "limit": 2924 }, @@ -81,7 +78,7 @@ "limit": 1 }, "C901": { - "limit": 315 + "limit": 314 }, "D419": { "limit": 6 @@ -113,18 +110,6 @@ "F401": { "limit": 17 }, - "FURB136": { - "limit": 0 - }, - "FURB168": { - "limit": 0 - }, - "FURB188": { - "limit": 0 - }, - "I001": { - "limit": 0 - }, "LOG015": { "limit": 5 }, @@ -137,18 +122,9 @@ "PERF401": { "limit": 12 }, - "PERF402": { - "limit": 0 - }, "PERF403": { "limit": 34 }, - "PIE790": { - "limit": 0 - }, - "PIE800": { - "limit": 0 - }, "PIE804": { "limit": 18 }, @@ -158,9 +134,6 @@ "PLC0206": { "limit": 26 }, - "PLC0208": { - "limit": 0 - }, "PLC0414": { "limit": 46 }, @@ -170,24 +143,12 @@ "PLR0206": { "limit": 1 }, - "PLR0402": { - "limit": 0 - }, "PLR1704": { "limit": 3 }, - "PLR1711": { - "limit": 0 - }, "PLR1714": { "limit": 257 }, - "PLR1730": { - "limit": 0 - }, - "PLR2044": { - "limit": 0 - }, "PLW0127": { "limit": 57 }, @@ -206,27 +167,12 @@ "PLW1510": { "limit": 2 }, - "PYI030": { - "limit": 0 - }, "PYI036": { "limit": 3 }, - "PYI041": { - "limit": 0 - }, - "PYI064": { - "limit": 0 - }, - "RET501": { - "limit": 0 - }, "RET504": { "limit": 177 }, - "RUF010": { - "limit": 0 - }, "RUF012": { "limit": 241 }, @@ -236,18 +182,9 @@ "RUF019": { "limit": 38 }, - "RUF022": { - "limit": 0 - }, - "RUF023": { - "limit": 0 - }, "RUF046": { "limit": 4 }, - "RUF051": { - "limit": 0 - }, "RUF059": { "limit": 67 }, @@ -272,18 +209,12 @@ "SIM113": { "limit": 3 }, - "SIM114": { - "limit": 0 - }, "SIM115": { "limit": 2 }, "SIM117": { "limit": 7 }, - "SIM118": { - "limit": 0 - }, "SIM201": { "limit": 1 }, @@ -302,11 +233,8 @@ "TC004": { "limit": 5 }, - "TC005": { - "limit": 0 - }, "TID251": { - "limit": 1240 + "limit": 1238 }, "TRY002": { "limit": 528 @@ -323,46 +251,13 @@ "TRY300": { "limit": 860 }, - "UP006": { - "limit": 0 - }, - "UP007": { - "limit": 0 - }, - "UP008": { - "limit": 0 - }, - "UP012": { - "limit": 0 - }, - "UP018": { - "limit": 0 - }, - "UP024": { - "limit": 0 - }, "UP028": { "limit": 2 }, "UP031": { "limit": 2 }, - "UP032": { - "limit": 0 - }, - "UP034": { - "limit": 0 - }, - "UP035": { - "limit": 0 - }, "UP036": { "limit": 1 - }, - "UP037": { - "limit": 0 - }, - "UP045": { - "limit": 0 } } diff --git a/ruff.toml b/ruff.toml index 00743e0f38a..9b90910b355 100644 --- a/ruff.toml +++ b/ruff.toml @@ -1,14 +1,26 @@ lint.ignore = ["F405", "E402", "F403"] -lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"] +# The second group is the strict gate's graduates: rules the codebase already has zero +# violations of, so they hard-fail here instead of being ratcheted in ruff-strict-budget.json. +# That gives editors and `ruff check --fix` the diagnostic, which the gate script cannot. +lint.extend-select = [ + "T20", "PGH004", "RUF008", "RUF009", "RUF100", + "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208", + "PLR0402", "PLR1711", "PLR1730", "PLR2044", "PYI030", "PYI041", "PYI064", "RET501", "RUF010", + "RUF022", "RUF023", "RUF051", "SIM114", "SIM118", "TC005", "UP006", "UP007", "UP008", "UP012", + "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045", +] # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip # `# noqa` directives that protect rules enforced elsewhere. List those codes as external # so RUF100 leaves their directives alone: the strict gate (ruff-strict.toml) and upstream # litellm's own ruff config both rely on suppressions this config can't see. lint.external = [ - # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml) - "ANN", "ASYNC230", "B", "C", "D419", "DTZ", "EXE", "FURB", "I001", "LOG015", "N999", "PERF", - "PIE", "PL", "PYI", "RET", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", - "RUF046", "RUF051", "RUF059", "S110", "S112", "SIM", "TC", "TID251", "TRY", "UP", + # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml). + # Family entries whose every strict rule graduated into extend-select above (FURB), and + # standalone graduated codes (I001, RUF010, RUF022, RUF023, RUF051), are dropped so this + # config's RUF100 polices their directives itself. + "ANN", "ASYNC230", "B", "C", "D419", "DTZ", "EXE", "LOG015", "N999", "PERF", + "PIE", "PL", "PYI", "RET", "RUF012", "RUF015", "RUF019", + "RUF046", "RUF059", "S110", "S112", "SIM", "TC", "TID251", "TRY", "UP", # Enforced by upstream litellm's ruff config, but not run in this repo's CI "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index 10a78483643..e97cd1bca00 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -10,7 +10,10 @@ content at the merge-base with the target branch and fails (exits 1, red) if: * a rule was dropped from a budget (its ceiling effectively became infinite), or * an entire budget file was deleted. -New rules and lowered/equal limits are fine. +New rules and lowered/equal limits are fine. So is a rule that graduated: once a +paired config (ruff.toml for the ruff-strict budget) selects the rule outright it +hard-fails at the first violation, which is stricter than any ceiling the budget +could hold, so dropping its entry tightens the guard rather than removing it. This is deliberately NOT a gating check. It should turn the run red so that a loosening is impossible to miss in review, but it must stay OUT of the @@ -30,7 +33,9 @@ import argparse import json import subprocess import sys +import tomllib from pathlib import Path +from types import MappingProxyType from typing import NamedTuple REPO_ROOT = Path(__file__).resolve().parent.parent @@ -40,6 +45,7 @@ DEFAULT_BUDGETS: tuple[str, ...] = ( "type-discipline-budget.json", "basedpyright-code-budget.json", ) +GRADUATION_CONFIGS = MappingProxyType({"ruff-strict-budget.json": "ruff.toml"}) class Regression(NamedTuple): @@ -106,24 +112,57 @@ def _limits(budget: dict) -> dict[str, int]: } +def selectors_hard_failed_by(lint: dict) -> tuple[str, ...]: + """A ruff `[lint]` table's selected codes, minus anything `ignore` turns back off. + + `lint.ignore` wins over `lint.extend-select` in ruff, so an ignored code is not + actually enforced and must not count as a graduation. + """ + ignored = tuple(lint.get("ignore", ())) + return tuple( + selector + for selector in lint.get("extend-select", ()) + if not (ignored and selector.startswith(ignored)) + ) + + +def graduated_selectors(rel: str) -> tuple[str, ...]: + """Selectors the budget's paired ruff config hard-fails, so its ceiling is moot.""" + config = GRADUATION_CONFIGS.get(rel) + if config is None or not (REPO_ROOT / config).exists(): + return () + return selectors_hard_failed_by( + tomllib.loads((REPO_ROOT / config).read_text()).get("lint", {}) + ) + + def _regression_detail( rule: str, base_limits: dict[str, int], head_limits: dict[str, int], + graduated: tuple[str, ...], ) -> str | None: - """Why `rule` regressed vs base, or None when it held flat or fell. + """Why `rule` regressed vs base, or None when it held flat, fell, or graduated. - A dropped rule is terminal; otherwise the only loosening left is a raised limit. + A dropped rule is terminal unless it graduated; otherwise the only loosening + left is a raised limit. """ base_limit = base_limits[rule] if rule not in head_limits: + if graduated and rule.startswith(graduated): + return None return f"rule dropped (limit {base_limit} -> removed)" if head_limits[rule] > base_limit: return f"limit raised {base_limit} -> {head_limits[rule]}" return None -def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regression]: +def regressions_for( + rel: str, + base: dict | None, + head: dict | None, + graduated: tuple[str, ...] = (), +) -> list[Regression]: if base is None: return [] # new budget file: nothing to ratchet against yet if head is None: @@ -133,7 +172,7 @@ def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regr return [ Regression(rel, rule, detail) for rule in sorted(base_limits) - if (detail := _regression_detail(rule, base_limits, head_limits)) is not None + if (detail := _regression_detail(rule, base_limits, head_limits, graduated)) is not None ] @@ -164,7 +203,7 @@ def main() -> int: print(f"skip {rel}: new file (no base at {args.base} to ratchet against)") continue checked.append(rel) - regressions.extend(regressions_for(rel, base, head)) + regressions.extend(regressions_for(rel, base, head, graduated_selectors(rel))) if regressions: print( diff --git a/scripts/install_git_hooks.sh b/scripts/install_git_hooks.sh index 7ea8c3ff2e9..8f4e79e4ddd 100755 --- a/scripts/install_git_hooks.sh +++ b/scripts/install_git_hooks.sh @@ -35,7 +35,7 @@ These hooks enforce Conventional Commits and Conventional Branches. Bypass with --no-verify when you need to (e.g. for emergency hotfixes). The CI-equivalent lint is deliberately not installed as an auto-firing hook -(it can take minutes); run it on demand with 'make pre-commit' before committing. +(it can take minutes); run it on demand with 'make check' before committing. To uninstall: git config --unset core.hooksPath EOF diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index fbe409e29ee..8da9f236d00 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -1,18 +1,27 @@ #!/usr/bin/env bash # -# pre_commit_lint.sh — shift CI lint left. Run it (via `make pre-commit`) right -# before `git commit`; it inspects your staged files and runs only the matching -# gating CI checks, so a clean run means a green CI lint: -# - litellm/ Python staged -> `make lint` (test-linting.yml's lint job) -# - tests/e2e Python staged -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) -# + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) -# - dashboard staged -> `make bootstrap-dashboard` + prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) -# - proxy/types staged -> regenerate dashboard API types and fail on drift (check-ui-api-types.yml) +# pre_commit_lint.sh — shift CI lint left. Run it (via `make check`, formerly +# `make pre-commit`) before `git commit`, or after committing (e.g. a merge +# commit) to predict CI for the branch. It picks the files in scope and runs +# only the matching gating CI checks, so a clean run means a green CI lint: +# - anything staged -> scope is the staged files; changed-but-unstaged files +# whose checks were skipped are called out +# - nothing staged -> scope is the working tree's diff against the merge base +# with origin/litellm_internal_staging, untracked files included +# The per-area checks: +# - litellm/ Python -> `make lint` (test-linting.yml's lint job) +# - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) +# + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) +# - dashboard -> `make bootstrap-dashboard` + prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) +# - proxy/types -> `make bootstrap-dashboard` + regenerate dashboard API types and fail on drift (check-ui-api-types.yml) +# `make check` provisions only the Python half, so the dashboard's npm install runs +# here, for the scopes that reach a node block, and not for a Python-only change. # -# Each block is skipped when no matching files are staged, so unrelated commits stay -# fast. This is intentionally not auto-installed as a git hook (see scripts/install_git_hooks.sh): -# the dashboard and basedpyright passes can take minutes, so it's run on demand rather -# than firing on every human commit. It is hook-compatible if you want that anyway: +# Each block is skipped when no matching files are in scope, so unrelated commits +# stay fast. This is intentionally not auto-installed as a git hook (see +# scripts/install_git_hooks.sh): the dashboard and basedpyright passes can take +# minutes, so it's run on demand rather than firing on every human commit. It is +# hook-compatible if you want that anyway: # `ln -s ../../scripts/pre_commit_lint.sh .git/hooks/pre-commit`. set -eu @@ -20,56 +29,105 @@ set -eu if [ -z "${PRE_COMMIT_LINT_INNER:-}" ]; then log_file=$(git rev-parse --path-format=absolute --git-path pre_commit_lint.log) if : > "$log_file" 2>/dev/null; then - echo "pre-commit: logging full output to $log_file" + echo "check: logging full output to $log_file" PRE_COMMIT_LINT_INNER=1 "$0" "$@" 2>&1 | tee "$log_file" pipe_status=("${PIPESTATUS[@]}") if [ "${pipe_status[1]}" -eq 0 ]; then - echo "pre-commit: full log: $log_file" + echo "check: full log: $log_file" else - echo "pre-commit: WARNING - writing $log_file failed; the log may be incomplete" >&2 + echo "check: WARNING - writing $log_file failed; the log may be incomplete" >&2 fi exit "${pipe_status[0]}" fi - echo "pre-commit: WARNING - cannot write $log_file; output will not be saved" >&2 + echo "check: WARNING - cannot write $log_file; output will not be saved" >&2 PRE_COMMIT_LINT_INNER=1 exec "$0" "$@" fi repo_root=$(git rev-parse --show-toplevel) cd "$repo_root" -staged=$(git diff --cached --name-only --diff-filter=ACMR) -staged_match() { printf '%s\n' "$staged" | grep -E "$1" || true; } +staged=$(git diff --cached --name-only --diff-filter=ACMRD) +unstaged=$(git diff --name-only) +untracked=$(git ls-files --others --exclude-standard) + +if [ -n "$staged" ]; then + scope=$staged +else + git fetch --quiet origin litellm_internal_staging 2>/dev/null || true + merge_base=$(git merge-base origin/litellm_internal_staging HEAD 2>/dev/null) || { + echo "check: cannot resolve the merge base with origin/litellm_internal_staging." >&2 + echo " Fix: git fetch origin litellm_internal_staging" >&2 + exit 1 + } + scope=$(printf '%s\n' "$(git diff --name-only --diff-filter=ACMRD "$merge_base")" "$untracked" | sed '/^$/d' | sort -u) + if [ -z "$scope" ]; then + echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs origin/litellm_internal_staging)" + exit 0 + fi + echo "check: nothing staged; scoping to the working tree's diff against the merge base with origin/litellm_internal_staging:" + printf '%s\n' "$scope" | sed 's/^/ /' +fi + +scope_match() { printf '%s\n' "$scope" | grep -E "$1" || true; } + +existing_files() { + while IFS= read -r f; do + if [ -f "$f" ]; then printf '%s\n' "$f"; fi + done +} + +litellm_py_pattern='^litellm/.*\.py$' +e2e_py_pattern='^tests/e2e/.*\.py$' +spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$' +ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$' +ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$' # CI's lint job (test-linting.yml) only inspects litellm/, so a tests-only or # scripts-only commit can't turn it red; scope the trigger there to skip the slow # make lint when it couldn't catch anything. -litellm_py_files=$(staged_match '^litellm/.*\.py$') -e2e_py_files=$(staged_match '^tests/e2e/.*\.py$') +litellm_py_files=$(scope_match "$litellm_py_pattern") +e2e_py_files=$(scope_match "$e2e_py_pattern") # ruff format (and CI's format step) skip enterprise; the rest of make lint covers it. -fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' || true) +fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files) # check-ui-api-types.yml triggers on any file under litellm/proxy or litellm/types # (Prisma schema and configs included, not just Python) plus the generator and its # lockfiles, so match that whole trigger set rather than a Python subset. -spec_files=$(staged_match '^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$') +spec_files=$(scope_match "$spec_pattern") # CI's frontend-lint runs prettier over a wider extension set than eslint; keep that # split so this flags exactly what the job would. -ui_prettier_files=$(staged_match '^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$') -ui_eslint_files=$(staged_match '^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$') +ui_prettier_changed=$(scope_match "$ui_prettier_pattern") +ui_eslint_changed=$(scope_match "$ui_eslint_pattern") +ui_prettier_files=$(printf '%s\n' "$ui_prettier_changed" | existing_files) +ui_eslint_files=$(printf '%s\n' "$ui_eslint_changed" | existing_files) -# CI lints the committed tree, so this script predicts CI for what you have STAGED -# (every trigger above reads `git diff --cached`). The tools it runs, though, read -# the working tree, so unstaged edits to tracked files and untracked files fold -# into the result and a green/red here won't match a commit of just the staged -# changes. There's no safe way to lint the index in place, so surface the gap -# instead of hiding it: stage everything you intend to commit before trusting a -# pass. This only warns; it never blocks or touches your changes. -unstaged=$(git diff --name-only) -untracked=$(git ls-files --others --exclude-standard) -if [ -n "$unstaged" ] || [ -n "$untracked" ]; then - echo "pre-commit: NOTE - unstaged/untracked changes are included in these checks but" >&2 - echo " won't be in a commit of only your staged changes, so this result may differ from" >&2 - echo " CI. Stage everything you intend to commit (git add) for an accurate prediction:" >&2 - printf '%s\n' "$unstaged" "$untracked" | sed '/^$/d' | sed 's/^/ /' >&2 +# CI lints the committed tree, so with staged files this script predicts CI for +# what you have STAGED (every trigger above reads `git diff --cached`). The tools +# it runs, though, read the working tree, so unstaged edits to tracked files and +# untracked files fold into the result and a green/red here won't match a commit +# of just the staged changes. There's no safe way to lint the index in place, so +# surface the gap instead of hiding it: stage everything you intend to commit +# before trusting a pass. This only warns; it never blocks or touches your changes. +if [ -n "$staged" ]; then + not_staged=$(printf '%s\n' "$unstaged" "$untracked" | sed '/^$/d' | sort -u) + if [ -n "$not_staged" ]; then + echo "check: NOTE - unstaged/untracked changes are included in these checks but" >&2 + echo " won't be in a commit of only your staged changes, so this result may differ from" >&2 + echo " CI. Stage everything you intend to commit (git add) for an accurate prediction:" >&2 + printf '%s\n' "$not_staged" | sed 's/^/ /' >&2 + fi + warn_skipped() { + local check_name=$1 pattern=$2 triggered=$3 + [ -n "$triggered" ] && return 0 + local missed + missed=$(printf '%s\n' "$not_staged" | grep -E "$pattern" || true) + [ -z "$missed" ] && return 0 + echo "check: SKIPPED $check_name because these changed files are not staged:" >&2 + printf '%s\n' "$missed" | sed 's/^/ /' >&2 + } + warn_skipped "Python lint (make lint)" "$litellm_py_pattern" "$litellm_py_files" + warn_skipped "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_pattern" "$e2e_py_files" + warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_changed" + warn_skipped "dashboard API-type sync (npm run gen:api)" "$spec_pattern" "$spec_files" fi lint_dashboard() { @@ -115,15 +173,15 @@ bootstrap_hint() { python_checks() { local rc=0 - echo "pre-commit: linting Python (make lint)" - make lint || { echo "✗ Python lint failed. Fix the reds above, then re-run make pre-commit." >&2; rc=1; } + echo "check: linting Python (make lint)" + make lint || { echo "✗ Python lint failed. Fix the reds above, then re-run make check." >&2; rc=1; } # `make lint` format-checks files in origin/base...HEAD, which at pre-commit time - # predates the staged change, so format-check the staged litellm files directly to + # predates the staged change, so format-check the scoped litellm files directly to # cover a brand-new commit before it lands. if [ -n "$fmt_files" ]; then - echo "pre-commit: ruff format --check (staged litellm files)" + echo "check: ruff format --check (scoped litellm files)" printf '%s\n' "$fmt_files" | xargs uv run --no-sync ruff format --check --exclude '/enterprise/' \ - || { echo "✗ Unformatted staged files. Fix with: make format, then re-stage." >&2; rc=1; } + || { echo "✗ Unformatted files in scope. Fix with: make format, then re-stage." >&2; rc=1; } fi return $rc } @@ -147,28 +205,28 @@ if [ -n "$litellm_py_files" ]; then fi if [ -n "$e2e_py_files" ] && [ -z "$litellm_py_files" ]; then - echo "pre-commit: type-checking tests/e2e (make lint-e2e-basedpyright)" - make lint-e2e-basedpyright || { echo "✗ tests/e2e basedpyright failed. Fix the errors above, then re-run make pre-commit." >&2; status=1; } + echo "check: type-checking tests/e2e (make lint-e2e-basedpyright)" + make lint-e2e-basedpyright || { echo "✗ tests/e2e basedpyright failed. Fix the errors above, then re-run make check." >&2; status=1; } fi if [ -n "$e2e_py_files" ]; then - echo "pre-commit: checking tests/e2e raw HTTP client ban (check_e2e_no_raw_requests)" + echo "check: checking tests/e2e raw HTTP client ban (check_e2e_no_raw_requests)" uv run --no-sync python tests/code_coverage_tests/check_e2e_no_raw_requests.py \ - || { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make pre-commit." >&2; status=1; } + || { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make check." >&2; status=1; } fi dashboard_ready=1 -if [ -n "$ui_prettier_files" ] || [ -n "$ui_eslint_files" ] || [ -n "$spec_files" ]; then - echo "pre-commit: provisioning the dashboard toolchain (make bootstrap-dashboard)" +if [ -n "$ui_prettier_changed" ] || [ -n "$ui_eslint_changed" ] || [ -n "$spec_files" ]; then + echo "check: provisioning the dashboard toolchain (make bootstrap-dashboard)" if ! make bootstrap-dashboard; then - echo "✗ make bootstrap-dashboard failed, so the dashboard lint and API-type checks are skipped rather than run against an unprovisioned toolchain. Fix the error above, then re-run make pre-commit." >&2 + echo "✗ make bootstrap-dashboard failed, so the dashboard lint and API-type checks are skipped rather than run against an unprovisioned toolchain. Fix the error above, then re-run make check." >&2 status=1 dashboard_ready="" fi fi dashboard_checks() { - echo "pre-commit: linting dashboard (prettier + eslint + lint budgets)" + echo "check: linting dashboard (prettier + eslint + lint budgets)" if [ ! -d ui/litellm-dashboard/node_modules ]; then echo "✗ ui/litellm-dashboard/node_modules is missing; dashboard lint cannot run." >&2 bootstrap_hint @@ -177,7 +235,7 @@ dashboard_checks() { lint_dashboard || { echo "✗ Dashboard lint failed. See above; format with: (cd ui/litellm-dashboard && npm run format)." >&2; return 1; } } -if [ -n "$dashboard_ready" ] && { [ -n "$ui_prettier_files" ] || [ -n "$ui_eslint_files" ]; }; then +if [ -n "$dashboard_ready" ] && { [ -n "$ui_prettier_changed" ] || [ -n "$ui_eslint_changed" ]; }; then dash_log=$(mktemp) set -m dashboard_checks > "$dash_log" 2>&1 & @@ -187,7 +245,7 @@ fi genapi_checks() { local status=0 - echo "pre-commit: checking dashboard API types are in sync (npm run gen:api)" + echo "check: checking dashboard API types are in sync (npm run gen:api)" # gen-api-types.mjs imports litellm.proxy.proxy_server, which needs the proxy deps # and an up-to-date Prisma client; check-ui-api-types.yml installs those and runs # prisma generate before gen:api, so mirror that here or a stale client can mask @@ -205,7 +263,7 @@ genapi_checks() { status=1 elif ( cd ui/litellm-dashboard && LITELLM_PYTHON="uv run --no-sync python" npm run gen:api ); then if ! git diff --quiet -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then - echo "✗ Dashboard API types are stale; regenerated src/lib/http/schema.d.ts. Stage it and commit; re-run make pre-commit only if other checks failed too." >&2 + echo "✗ Dashboard API types are stale; regenerated src/lib/http/schema.d.ts. Stage it and commit; re-run make check only if other checks failed too." >&2 status=1 fi else diff --git a/tests/agent_tests/test_a2a_agent.py b/tests/agent_tests/test_a2a_agent.py index f5ad9601369..1f72ced64f1 100644 --- a/tests/agent_tests/test_a2a_agent.py +++ b/tests/agent_tests/test_a2a_agent.py @@ -40,7 +40,7 @@ class MockA2AClient: name="mock-agent", url="http://mock-agent.local" ) - async def send_message(self, request): + async def send_message(self, request, *, context=None): from a2a.compat.v0_3.conversions import pb2_v10 for text in ("hel", "hello"): diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 5999e5772d1..dc69bd42171 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -138,7 +138,7 @@ Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure cover Before you push -1. Run `make lint-e2e-basedpyright` (or `make pre-commit` with your changes staged); the harness is fully typed and the gate allows zero basedpyright errors, enforced in CI on any PR touching `tests/e2e/**/*.py` +1. Run `make lint-e2e-basedpyright` (or `make check` with your changes staged); the harness is fully typed and the gate allows zero basedpyright errors, enforced in CI on any PR touching `tests/e2e/**/*.py` 2. Add the models your test needs to the config your local proxy loads diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 2627bdb8038..5050b6fce68 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -76,12 +76,127 @@ from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] +# After /model/new, poll data-plane /v1/models until the model is listed (or fail). +# Bound by MODEL_SERVABLE_TIMEOUT so a stuck reload does not burn the spend +# poll_timeout (120s). Return on first listing; settle_propagation owns the separate +# wait that lets every worker and replica reload before the caller uses the model. +MODEL_SERVABLE_TIMEOUT = 40.0 +MODEL_SERVABLE_DB_SYNC_SECONDS = 0.0 +MODEL_SERVABLE_INTERVAL = 2.0 +# Cap each /v1/models poll so one slow request cannot outlast the remaining budget. +MODEL_SERVABLE_REQUEST_TIMEOUT = 5.0 + + +@dataclass(frozen=True, slots=True) +class Servable: + """The data plane listed the model within the deadline.""" + + +@dataclass(frozen=True, slots=True) +class NotServable: + """The deadline passed without the data plane listing the model. + + `last_result` is the final /v1/models read, so the caller can tell "the proxy + answered but omitted the model" (propagation) from "the read itself failed" + (network/auth) when reporting.""" + + last_result: Result[ModelsListResponse] | None + + +ServableOutcome = Servable | NotServable + + +def await_servable( + list_models: Callable[[float], Result[ModelsListResponse]], + *, + model_name: str, + timeout: float, + interval: float, + request_timeout: float, + db_sync_seconds: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> ServableOutcome: + """Poll until `model_name` is listed long enough for every worker to DB-sync. + + First listing must happen within `timeout`. After that, the model must stay + listed continuously for `db_sync_seconds` (any miss resets the continuous + window). `db_sync_seconds=0` returns on the first listing. Each poll's request + timeout is clamped to the remaining budget. Sleeps only min(interval, time left) + so a final deadline-clamped poll is never skipped just because a full interval + does not fit. Clock and sleep are injected.""" + started = now() + first_seen_at: float | None = None + last_result: Result[ModelsListResponse] | None = None + while True: + t = now() + phase_deadline = ( + started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds + ) + remaining = phase_deadline - t + if remaining <= 0: + if ( + last_result is not None + and first_seen_at is not None + and (db_sync_seconds <= 0 or t - first_seen_at >= db_sync_seconds) + ): + return Servable() + return NotServable(last_result=last_result) + + poll_timeout = min(request_timeout, remaining) + last_result = list_models(poll_timeout) + listed = isinstance(last_result, Success) and any( + entry.id == model_name for entry in last_result.data.data + ) + t = now() + if not listed: + first_seen_at = None + elif first_seen_at is None: + if t > started + timeout: + return NotServable(last_result=last_result) + first_seen_at = t + if db_sync_seconds <= 0: + return Servable() + elif t - first_seen_at >= db_sync_seconds: + return Servable() + + phase_deadline = ( + started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds + ) + wait = min(interval, phase_deadline - now()) + if wait > 0: + sleep(wait) + + +def servable_timeout_message( + *, + model_name: str, + timeout: float, + db_sync_seconds: float, + last_result: Result[ModelsListResponse] | None, +) -> str: + last_error = ( + f"; last /v1/models poll did not succeed: {last_result}" + if last_result is not None and not isinstance(last_result, Success) + else "" + ) + return ( + f"model {model_name!r} was created but never became servable on the data " + f"plane within {timeout}s of first listing (plus {db_sync_seconds}s continuous " + f"DB sync) after /model/new (control/data-plane propagation or " + f"STORE_MODEL_IN_DB reload issue){last_error}" + ) + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport poll_timeout: float = 120.0 poll_interval: float = 5.0 + model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT + model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS + model_servable_interval: float = MODEL_SERVABLE_INTERVAL + model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT # ---- keys / customers (satisfies lifecycle.ResourceClient) ---------- @@ -192,33 +307,35 @@ class ProxyClient: return model_id def _await_model_servable(self, model_name: str) -> None: - """Block until the data plane lists `model_name`, or fail loudly if it does - not within poll_timeout (a real propagation/config problem, surfaced here - instead of as a downstream "Invalid model name passed").""" - deadline = time.monotonic() + self.poll_timeout - last_result: Result[ModelsListResponse] | None = None - while time.monotonic() < deadline: - last_result = self.transport.get( + """Block until the data plane lists `model_name`, or fail at model_servable_timeout.""" + outcome = await_servable( + lambda poll_timeout: self.transport.get( "/v1/models", headers=self.transport.master, params=NoBody(), response_type=ModelsListResponse, - ) - if isinstance(last_result, Success) and any( - entry.id == model_name for entry in last_result.data.data - ): + timeout=poll_timeout, + ), + model_name=model_name, + timeout=self.model_servable_timeout, + interval=self.model_servable_interval, + request_timeout=self.model_servable_request_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case Servable(): return - time.sleep(self.poll_interval) - last_error = ( - f"; last /v1/models poll did not succeed: {last_result}" - if last_result is not None and not isinstance(last_result, Success) - else "" - ) - raise AssertionError( - f"model {model_name!r} was created but never became servable on the data " - f"plane within {self.poll_timeout}s of /model/new (control/data-plane " - f"propagation or STORE_MODEL_IN_DB reload issue){last_error}" - ) + case NotServable(last_result=last_result): + raise AssertionError( + servable_timeout_message( + model_name=model_name, + timeout=self.model_servable_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, + last_result=last_result, + ) + ) def update_model(self, model_id: str, litellm_params: LiteLLMParamsBody) -> None: """Merge `litellm_params` over the deployment `model_id`'s stored params via diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index a6adf83ed1f..27b11befc8e 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -58,6 +58,7 @@ class Transport(Protocol): headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: ... def delete[R: BaseModel]( @@ -136,13 +137,16 @@ class HttpTransport: headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: + """`timeout` overrides the transport-wide request_timeout for this call, for + pollers whose own deadline is shorter than it.""" return e2e_http.get( self._url(path), headers=headers, params=params, response_type=response_type, - timeout=self.request_timeout, + timeout=self.request_timeout if timeout is None else timeout, ) def delete[R: BaseModel]( @@ -336,9 +340,14 @@ class SplitTransport: headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: return self._route(path).get( - path, headers=headers, params=params, response_type=response_type + path, + headers=headers, + params=params, + response_type=response_type, + timeout=timeout, ) def delete[R: BaseModel]( diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index 44da3ea06a0..00d5380b2f4 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -622,8 +622,12 @@ async def test_service_logger_keys_success(): logger success hook is called with the correct event metadata and no exception is logged. """ keys = [ - {"id": "key1", "spend": 10.0, "budget_duration": 60, "token": "key1"}, - {"id": "key2", "spend": 15.0, "budget_duration": 60, "token": "key2"}, + _attrify( + {"id": "key1", "spend": 10.0, "budget_duration": 60, "token": "key1"} + ), + _attrify( + {"id": "key2", "spend": 15.0, "budget_duration": 60, "token": "key2"} + ), ] prisma_client = MagicMock() prisma_client.get_data = AsyncMock(return_value=keys) @@ -740,8 +744,12 @@ async def test_service_logger_users_success(): the correct metadata and no exception is logged. """ users = [ - {"id": "user1", "spend": 20.0, "budget_duration": 120, "user_id": "user1"}, - {"id": "user2", "spend": 25.0, "budget_duration": 120, "user_id": "user2"}, + _attrify( + {"id": "user1", "spend": 20.0, "budget_duration": 120, "user_id": "user1"} + ), + _attrify( + {"id": "user2", "spend": 25.0, "budget_duration": 120, "user_id": "user2"} + ), ] prisma_client = MagicMock() prisma_client.get_data = AsyncMock(return_value=users) @@ -853,8 +861,12 @@ async def test_service_logger_teams_success(): the proper metadata and nothing is logged as an exception. """ teams = [ - {"id": "team1", "spend": 30.0, "budget_duration": 180, "team_id": "team1"}, - {"id": "team2", "spend": 35.0, "budget_duration": 180, "team_id": "team2"}, + _attrify( + {"id": "team1", "spend": 30.0, "budget_duration": 180, "team_id": "team1"} + ), + _attrify( + {"id": "team2", "spend": 35.0, "budget_duration": 180, "team_id": "team2"} + ), ] prisma_client = MagicMock() prisma_client.get_data = AsyncMock(return_value=teams) diff --git a/tests/llm_responses_api_testing/base_responses_api.py b/tests/llm_responses_api_testing/base_responses_api.py index 407091a65b3..f5751aa79e8 100644 --- a/tests/llm_responses_api_testing/base_responses_api.py +++ b/tests/llm_responses_api_testing/base_responses_api.py @@ -338,7 +338,7 @@ class BaseResponsesAPITest(ABC): ) assert result is not None assert result.id == response.id - assert result.output == response.output + assert result.output_text == response.output_text else: raise ValueError("response is not a ResponsesAPIResponse") else: @@ -352,7 +352,7 @@ class BaseResponsesAPITest(ABC): ) assert result is not None assert result.id == response.id - assert result.output == response.output + assert result.output_text == response.output_text else: raise ValueError("response is not a ResponsesAPIResponse") diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index 65b70f13a3b..7bc61c40ea0 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -12,8 +12,10 @@ from typing import Final import pytest -from litellm.proxy.db.autorouter_session_rollup import UPSERT_AUTOROUTER_SESSION_SQL -from litellm.proxy.management_endpoints.auto_router_endpoints import _BENCHMARKS_SQL +from litellm.proxy.db.autorouter_session_rollup import ( + AUTOROUTER_BENCHMARKS_SQL, + UPSERT_AUTOROUTER_SESSION_SQL, +) pytestmark = pytest.mark.asyncio(loop_scope="session") @@ -164,7 +166,7 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): await _turn(db, key, "A", T0 - timedelta(days=40), session_id=out_of_window, router=router) rows = await db.query_raw( - _BENCHMARKS_SQL, + AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), ) @@ -186,7 +188,7 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, router_type="quality") rows = await db.query_raw( - _BENCHMARKS_SQL, + AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), ) @@ -248,7 +250,7 @@ async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db): await _turn(db, key, "C", T0 + timedelta(seconds=30), session_id=f"s-{uuid.uuid4()}", router=router, tier=None) rows = await db.query_raw( - _BENCHMARKS_SQL, + AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), ) @@ -275,7 +277,7 @@ async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(d ) rows = await db.query_raw( - _BENCHMARKS_SQL, + AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), ) @@ -289,7 +291,7 @@ async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db): await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier=None) rows = await db.query_raw( - _BENCHMARKS_SQL, + AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), ) diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index c0993051d33..0dc8f56f3ce 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -19,6 +19,7 @@ sys.path.insert( import asyncio import litellm +from litellm import router as litellm_router_module from litellm import utils as litellm_utils_module from litellm._logging import ALL_LOGGERS from litellm.litellm_core_utils.prompt_templates import ( @@ -244,6 +245,8 @@ def isolate_litellm_state(): for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() } + original_live_routers = set(litellm_router_module._live_routers) + # Store LiteLLM logger state. Some tests reconfigure handlers/propagation for # JSON logging and do not restore them, which breaks later caplog-based tests. logger_state = {} @@ -313,6 +316,11 @@ def isolate_litellm_state(): litellm_utils_module._runtime_registered_model_cost.clear() litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost) + for _router in tuple(litellm_router_module._live_routers): + litellm_router_module._live_routers.discard(_router) + for _router in original_live_routers: + litellm_router_module._live_routers.add(_router) + # Restore logger configuration mutated by logging-focused tests. for logger in ALL_LOGGERS: original_logger_state = logger_state.get(logger.name) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 2573ad5a375..82b074220fa 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -1195,8 +1195,13 @@ def test_async_post_call_failure_hook_skips_a_transport_that_already_answered(): def test_record_error_attributes_on_span_decorates_without_ending(): """PATH A: a failure that dies before any LLM-call span (malformed body, validation) is stamped onto the instrumentor-owned SERVER span. The method must - not end the span or emit a duplicate exception event, and must pin error.code - to the real response status (not the exception's own code).""" + not end the span, and must pin error.code to the real response status (not the + exception's own code). + + LIT-4780: the instrumentor never sees the exception (the proxy handler turns it + into a JSONResponse), so nothing else marks the span as failed; the status and + the exception event have to come from here or the trace shows the error message + on an otherwise successful-looking request.""" logger, exporter = _logger() server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) logger.record_error_attributes_on_span(server, _proxy_exc("Invalid JSON body", 400), 422) @@ -1206,7 +1211,31 @@ def test_record_error_attributes_on_span_decorates_without_ending(): assert span.attributes["error.type"] == "ProxyException" assert span.attributes["error.message"] == "Invalid JSON body" assert span.attributes["litellm.provider.error.code"] == "422" - assert all(e.name != "exception" for e in span.events) + assert span.status.status_code is StatusCode.ERROR + assert [e.name for e in span.events] == ["exception"] + + +def test_record_error_attributes_on_span_does_not_duplicate_an_already_stamped_error(): + """A failure that already went through ``async_post_call_failure_hook`` reaches + the exception handler too; the second stamp must keep one exception event while + still repinning error.code to the real response status.""" + from litellm.proxy._types import UserAPIKeyAuth + + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + set_request_root_span(server) + exc = _proxy_exc("Authentication Error, invalid key", 401) + asyncio.run( + logger.async_post_call_failure_hook( + request_data={}, original_exception=exc, user_api_key_dict=UserAPIKeyAuth() + ) + ) + logger.record_error_attributes_on_span(server, exc, 400) + server.end() + (span,) = exporter.get_finished_spans() + assert [e.name for e in span.events] == ["exception"] + assert span.attributes["litellm.provider.error.code"] == "400" + assert span.status.status_code is StatusCode.ERROR def test_record_error_attributes_on_span_ignores_below_400_and_missing_span(): diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index b6ff3b70a4d..f39f41a6d12 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -221,14 +221,97 @@ async def test_execute_search_passes_selected_search_tool_litellm_params(monkeyp kwargs={"litellm_params": {"metadata": {"user_api_key_auth": user_api_key_auth}}}, ) - mock_asearch.assert_awaited_once_with( - query="what is litellm", - search_provider="tavily", - api_key="fake-ui-key", - api_base="https://api.tavily.com", - timeout=10.0, - max_retries=2, + forwarded_kwargs = mock_asearch.await_args.kwargs + assert forwarded_kwargs["query"] == "what is litellm" + assert forwarded_kwargs["search_provider"] == "tavily" + assert forwarded_kwargs["api_key"] == "fake-ui-key" + assert forwarded_kwargs["api_base"] == "https://api.tavily.com" + assert forwarded_kwargs["timeout"] == 10.0 + assert forwarded_kwargs["max_retries"] == 2 + + +@pytest.mark.asyncio +async def test_execute_search_attributes_spend_to_the_calling_key(monkeypatch): + """An intercepted search is billed and logged against the key that made the LLM request. + + Without the forwarded attribution metadata the proxy's spend hook skips the search + entirely, so its provider cost never reaches SpendLogs or any budget. + """ + import litellm + from litellm.proxy import proxy_server + from litellm.proxy.hooks.proxy_track_cost_callback import _should_track_cost_callback + + logger = WebSearchInterceptionLogger( + enabled_providers=["bedrock"], + search_tool_name="perplexity-sonar-pro", ) + router = MagicMock() + router.search_tools = [ + { + "search_tool_name": "perplexity-sonar-pro", + "litellm_params": {"search_provider": "perplexity", "api_key": "fake-key"}, + } + ] + mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[])) + user_api_key_auth = UserAPIKeyAuth( + api_key="hashed-sk-1234", + key_alias="alice-key", + user_id="user-alice", + org_id="org-1", + ) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search( + "what is litellm", + kwargs={"litellm_params": {"metadata": {"user_api_key_auth": user_api_key_auth}}}, + ) + + forwarded_metadata = mock_asearch.await_args.kwargs["litellm_metadata"] + assert forwarded_metadata["user_api_key"] == "hashed-sk-1234" + assert forwarded_metadata["user_api_key_hash"] == "hashed-sk-1234" + assert forwarded_metadata["user_api_key_alias"] == "alice-key" + assert forwarded_metadata["user_api_key_user_id"] == "user-alice" + assert forwarded_metadata["user_api_key_org_id"] == "org-1" + assert forwarded_metadata["model_group"] == "perplexity-sonar-pro" + assert ( + _should_track_cost_callback( + user_api_key=forwarded_metadata["user_api_key"], + user_id=forwarded_metadata["user_api_key_user_id"], + team_id=forwarded_metadata["user_api_key_team_id"], + end_user_id=None, + call_type="asearch", + ) + is True + ) + + +@pytest.mark.asyncio +async def test_execute_search_without_proxy_auth_context_stays_sdk_only(monkeypatch): + """SDK callers have no key to attribute the search to, so no proxy metadata is invented.""" + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger( + enabled_providers=["bedrock"], + search_tool_name="perplexity-sonar-pro", + ) + router = MagicMock() + router.search_tools = [ + { + "search_tool_name": "perplexity-sonar-pro", + "litellm_params": {"search_provider": "perplexity", "api_key": "fake-key"}, + } + ] + mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[])) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search("what is litellm", kwargs={"litellm_params": {}}) + + assert "litellm_metadata" not in mock_asearch.await_args.kwargs @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 6fc44ef7519..60d9689dc0b 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -4786,6 +4786,117 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder() setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +async def test_user_api_key_auth_authenticates_before_raising_malformed_body_error(): + """Regression (LIT-4780): a body that fails to parse must still be authenticated + first, so the rejected request's trace carries the caller's key / team / user + identity instead of an anonymous root span. The parse error is re-raised + unchanged once identity is seeded.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + builder_token = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="team-1") + + request = Request( + scope={ + "type": "http", + "headers": [(b"content-type", b"application/json")], + "method": "POST", + } + ) + request._url = URL(url="/chat/completions") + request._body = b'{}{"model": "gpt-4o"}' + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth._user_api_key_auth_builder", + new_callable=AsyncMock, + return_value=builder_token, + ) as mock_builder, + patch( + "litellm.proxy.auth.user_api_key_auth._run_centralized_common_checks", + new_callable=AsyncMock, + ) as mock_common_checks, + patch( + "litellm.proxy.auth.user_api_key_auth.RouteChecks.should_call_route", + ), + patch( + "litellm.proxy.auth.user_api_key_auth.seed_request_identity", + ) as mock_seed, + ): + with pytest.raises(ProxyException) as exc_info: + await user_api_key_auth(request=request, api_key="Bearer sk-test") + + assert "Invalid JSON payload" in str(exc_info.value.message) + assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST) + mock_builder.assert_awaited_once() + assert mock_seed.call_args.args[0] is builder_token + # authorization must not run for a request that is about to be rejected: + # ``common_checks`` reserves budget against live spend counters that only the + # endpoint's post-call path releases, and the endpoint never runs here + mock_common_checks.assert_not_awaited() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_user_api_key_auth_malformed_body_with_rejected_key_still_returns_the_parse_error(): + """The body is read before the key is authenticated, so a caller who sends both a + malformed body and a key that fails auth gets the 400. Authenticating the request + first (LIT-4780) must not turn that into the auth status code.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + request = Request( + scope={ + "type": "http", + "headers": [(b"content-type", b"application/json")], + "method": "POST", + } + ) + request._url = URL(url="/chat/completions") + request._body = b'{}{"model": "gpt-4o"}' + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth._user_api_key_auth_builder", + new_callable=AsyncMock, + side_effect=ProxyException( + message="Authentication Error, invalid key", + type="auth_error", + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.RouteChecks.should_call_route", + ), + ): + with pytest.raises(ProxyException) as exc_info: + await user_api_key_auth(request=request, api_key="Bearer sk-bad") + + assert "Invalid JSON payload" in str(exc_info.value.message) + assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + def _proxy_attrs_for_db_lookup(): """Minimal proxy_server attributes for driving the real ``_user_api_key_auth_builder`` down to the DB key lookup.""" diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index e758aa5ca7f..f9193db143e 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -2253,3 +2253,176 @@ async def test_cancel__provider_only_resolves_named_vertex_credentials(cancel_ha "vertex_location": "us-central1", "vertex_credentials": "/creds/customer-sa.json", } + + +# =========================================================================== # +# require_managed_files - raw provider ids must not reach the provider. # +# # +# Ownership rows only exist for LiteLLM managed ids. A raw provider id sent to # +# these routes is forwarded under the shared provider credentials with no # +# tenant check, so any caller who learns another tenant's id can read its # +# batch, reuse its file as batch input, or cancel its job. These lock the # +# guard on every batches route that accepts a caller-supplied id. # +# =========================================================================== # + + +def _unified_batch_id(model_id: str = "azure/gpt-4o", batch_id: str = "batch-provider-id") -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +def _unified_file_id() -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "managed-id", "gpt-4o-mini", "file-provider-id", "gpt-4o-mini-id" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +@dataclass(frozen=True) +class ManagedResourceAccessCheckerStub: + file_access: bool = True + object_access: bool = True + + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.file_access + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.object_access + + +@pytest.mark.asyncio +async def test_create__raw_input_file_id_rejected_when_managed_files_required(harness): + set_body( + harness, + { + "input_file_id": "file-victim-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "400" + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__model_encoded_input_file_id_rejected_when_managed_files_required(harness): + """A model-encoded id is client-forgeable and has no ownership row, so it is + not a managed file id and must be rejected like any other raw id.""" + set_body( + harness, + { + "input_file_id": AZURE_FILE_ID, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "400" + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__raw_input_file_id_allowed_when_managed_files_not_required(harness): + set_body( + harness, + { + "input_file_id": "file-victim-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with patch.object(litellm, "require_managed_files", False): + await call_create(harness) + + assert harness.acreate_kwargs()["input_file_id"] == "file-victim-abc123" + + +@pytest.mark.asyncio +async def test_create__other_teams_unified_input_file_id_rejected(harness): + set_body( + harness, + { + "input_file_id": _unified_file_id(), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "403" + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_retrieve__raw_batch_id_rejected_when_managed_files_required(retrieve_harness): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_retrieve(retrieve_harness, "batch-victim-abc123") + + assert exc.value.code == "400" + retrieve_harness.litellm_aretrieve.assert_not_called() + retrieve_harness.router_aretrieve.assert_not_called() + + +@pytest.mark.asyncio +async def test_retrieve__unified_batch_id_allowed_when_managed_files_required(retrieve_harness): + retrieve_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + + with patch.object(litellm, "require_managed_files", True): + await call_retrieve(retrieve_harness, _unified_batch_id()) + + assert retrieve_harness.router_aretrieve.call_count == 1 + + +@pytest.mark.asyncio +async def test_cancel__raw_batch_id_rejected_when_managed_files_required(cancel_harness): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_cancel(cancel_harness, "batch-victim-abc123") + + assert exc.value.code == "400" + cancel_harness.litellm_acancel.assert_not_called() + cancel_harness.router_acancel.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancel__unified_batch_id_allowed_when_managed_files_required(cancel_harness): + cancel_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + + with patch.object(litellm, "require_managed_files", True): + await call_cancel(cancel_harness, _unified_batch_id()) + + assert cancel_harness.router_acancel.call_count == 1 diff --git a/tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py b/tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py b/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py new file mode 100644 index 00000000000..b54787bf428 --- /dev/null +++ b/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py @@ -0,0 +1,278 @@ +""" +require_managed_files enforcement for litellm/proxy/fine_tuning_endpoints/endpoints.py + +Ownership rows only exist for LiteLLM managed ids. A raw provider id sent to these +routes is forwarded to the provider under the shared proxy credentials with no tenant +check, so any caller who learns another tenant's file id can train on it, and any +caller who learns another tenant's job id can read or cancel it. + +Each test asserts BOTH that the request is rejected AND that every downstream provider +seam stayed untouched, so a guard that raises after the provider call would still fail. +""" + +import base64 +import os +import sys +from contextlib import ExitStack +from dataclasses import dataclass +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from fastapi import Response + +import litellm +import litellm.proxy.fine_tuning_endpoints.endpoints as endpoints +import litellm.proxy.proxy_server as proxy_server +from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.utils import ProxyLogging +from litellm.router import Router +from litellm.types.llms.openai import LiteLLMFineTuningJobCreate +from litellm.types.utils import LiteLLMFineTuningJob, SpecialEnums + +RAW_FILE_ID = "file-victim-abc123" +RAW_JOB_ID = "ftjob-victim-abc123" + + +def _unified_file_id() -> str: + unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-4o-mini", RAW_FILE_ID, "gpt-4o-mini-id" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +def _unified_job_id() -> str: + unified = SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format("gpt-4o-mini-id", RAW_JOB_ID) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +def _job() -> LiteLLMFineTuningJob: + job = LiteLLMFineTuningJob( + id=RAW_JOB_ID, + created_at=1234567890, + fine_tuned_model=None, + finished_at=None, + hyperparameters={"n_epochs": 1}, + model="gpt-4o-mini", + object="fine_tuning.job", + organization_id="org-test", + result_files=[], + seed=0, + status="running", + trained_tokens=None, + training_file=RAW_FILE_ID, + validation_file=None, + ) + job._hidden_params = {} + return job + + +class FakeRequest: + def __init__(self): + self.headers = {} + self.query_params = {} + + async def json(self): + return {} + + +@dataclass(frozen=True) +class ManagedResourceAccessCheckerStub: + file_access: bool = True + object_access: bool = True + + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.file_access + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.object_access + + +class Seams: + def __init__(self, router: MagicMock, litellm_calls: dict[str, AsyncMock], logging: MagicMock): + self.router = router + self.litellm_calls = litellm_calls + self.logging = logging + + def assert_no_provider_call(self) -> None: + for name, mock in self.litellm_calls.items(): + assert mock.call_count == 0, f"litellm.{name} was called" + for name in ("acreate_fine_tuning_job", "aretrieve_fine_tuning_job", "acancel_fine_tuning_job"): + assert getattr(self.router, name).call_count == 0, f"router.{name} was called" + + +@pytest.fixture +def seams(): + logging = MagicMock(spec=ProxyLogging) + logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"]) + logging.post_call_failure_hook = AsyncMock() + logging.update_request_status = AsyncMock() + logging.get_proxy_hook = MagicMock(return_value=None) + + router = MagicMock(spec=Router) + router.acreate_fine_tuning_job = AsyncMock(return_value=_job()) + router.aretrieve_fine_tuning_job = AsyncMock(return_value=_job()) + router.acancel_fine_tuning_job = AsyncMock(return_value=_job()) + + litellm_calls = { + name: AsyncMock(return_value=_job()) + for name in ("acreate_fine_tuning_job", "aretrieve_fine_tuning_job", "acancel_fine_tuning_job") + } + + with ExitStack() as stack: + stack.enter_context( + patch.object( + ProxyBaseLLMRequestProcessing, + "common_processing_pre_call_logic", + AsyncMock(side_effect=lambda self=None, **kw: (self.data if self else {}, MagicMock())), + ) + ) + stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", MagicMock(return_value={}))) + for name, mock in litellm_calls.items(): + stack.enter_context(patch.object(litellm, name, mock)) + stack.enter_context(patch.object(proxy_server, "llm_router", router)) + stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging)) + stack.enter_context(patch.object(proxy_server, "premium_user", True)) + stack.enter_context(patch.object(proxy_server, "general_settings", {})) + stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock())) + stack.enter_context(patch.object(proxy_server, "version", "test-version")) + stack.enter_context(patch.object(endpoints, "fine_tuning_config", [{"custom_llm_provider": "openai"}])) + yield Seams(router=router, litellm_calls=litellm_calls, logging=logging) + + +async def _create(training_file: str, validation_file: str | None = None): + return await endpoints.create_fine_tuning_job( + request=FakeRequest(), + fastapi_response=Response(), + fine_tuning_request=LiteLLMFineTuningJobCreate( + model="gpt-4o-mini", + training_file=training_file, + validation_file=validation_file, + custom_llm_provider="openai", + ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +async def _retrieve(job_id: str): + return await endpoints.retrieve_fine_tuning_job( + request=FakeRequest(), + fastapi_response=Response(), + fine_tuning_job_id=job_id, + custom_llm_provider="openai", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +async def _cancel(job_id: str): + return await endpoints.cancel_fine_tuning_job( + request=FakeRequest(), + fastapi_response=Response(), + fine_tuning_job_id=job_id, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +@pytest.mark.asyncio +async def test_create__raw_training_file_rejected_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _create(RAW_FILE_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_create__raw_validation_file_rejected_when_managed_files_required(seams): + """The validation file is uploaded and readable exactly like the training file, + so a managed training_file must not smuggle a raw validation_file past the guard.""" + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _create(_unified_file_id(), validation_file=RAW_FILE_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_create__unified_training_file_allowed_when_managed_files_required(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + + with patch.object(litellm, "require_managed_files", True): + await _create(_unified_file_id()) + + assert seams.router.acreate_fine_tuning_job.call_count == 1 + + +@pytest.mark.asyncio +async def test_create__other_teams_unified_training_file_rejected(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _create(_unified_file_id()) + + assert exc.value.code == "403" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_create__raw_training_file_allowed_when_managed_files_not_required(seams): + with patch.object(litellm, "require_managed_files", False): + await _create(RAW_FILE_ID) + + assert seams.litellm_calls["acreate_fine_tuning_job"].call_count == 1 + + +@pytest.mark.asyncio +async def test_retrieve__raw_job_id_rejected_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _retrieve(RAW_JOB_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_retrieve__unified_job_id_allowed_when_managed_files_required(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + + with patch.object(litellm, "require_managed_files", True): + await _retrieve(_unified_job_id()) + + assert seams.router.aretrieve_fine_tuning_job.call_count == 1 + + +@pytest.mark.asyncio +async def test_cancel__raw_job_id_rejected_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _cancel(RAW_JOB_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_cancel__unified_job_id_allowed_when_managed_files_required(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + + with patch.object(litellm, "require_managed_files", True): + await _cancel(_unified_job_id()) + + assert seams.router.acancel_fine_tuning_job.call_count == 1 diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index ac01c6ae1d1..e68e7102fce 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3051,3 +3051,147 @@ def test_list_files_key_allowed_openai_model_still_resolves_team_credentials( mocker, monkeypatch, _team_openai_plus_global_anthropic_router(), ["team-gpt"] ) assert captured_kwargs.get("api_key") == "team-openai-key" + + +@pytest.mark.parametrize( + "http_method, url, patched_litellm_call", + [ + ("get", "/v1/files/file-victim-abc123", "litellm.afile_retrieve"), + ("get", "/v1/files/file-victim-abc123/content", "litellm.afile_content"), + ("delete", "/v1/files/file-victim-abc123", "litellm.afile_delete"), + ], +) +def test_require_managed_files_rejects_raw_provider_file_id( + mocker: MockerFixture, + monkeypatch, + llm_router: Router, + http_method: str, + url: str, + patched_litellm_call: str, +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr("litellm.require_managed_files", True) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + mock_call = mocker.patch(patched_litellm_call, new=mocker.AsyncMock()) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="attacker-user" + ) + + try: + response = getattr(client, http_method)( + url, headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + monkeypatch.setattr("litellm.require_managed_files", False) + + assert response.status_code == 400, response.text + mock_call.assert_not_called() + + +def _unified_managed_file_id() -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-3.5-turbo", "file-victim-abc123", "gpt-3.5-turbo-id" + ) + return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") + + +class _ManagedResourceAccessCheckerStub: + 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 True + + +@pytest.mark.asyncio +async def test_require_managed_files_allows_owned_unified_managed_file_id(monkeypatch): + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_id_requirement, + ) + + monkeypatch.setattr("litellm.require_managed_files", True) + + await validate_managed_id_requirement( + resource_id=_unified_managed_file_id(), + resource_kind="file", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="owner-user"), + managed_files_obj=_ManagedResourceAccessCheckerStub(), + ) + + +@pytest.mark.asyncio +async def test_managed_file_id_requirement_is_opt_in(monkeypatch): + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_id_requirement, + ) + + monkeypatch.setattr("litellm.require_managed_files", False) + + await validate_managed_id_requirement( + resource_id="file-victim-abc123", + resource_kind="file", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + managed_files_obj=None, + ) + + +def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr("litellm.require_managed_files", False) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + mock_retrieve = mocker.patch( + "litellm.afile_retrieve", + new=mocker.AsyncMock( + return_value=OpenAIFileObject( + id="file-victim-abc123", + object="file", + bytes=3, + created_at=1234567890, + filename="test.txt", + purpose="user_data", + status="uploaded", + ) + ), + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="some-user" + ) + + try: + response = client.get( + "/v1/files/file-victim-abc123", headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + mock_retrieve.assert_called_once() diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 6d6fd2e5507..22f9e6bb67a 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -676,6 +676,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "applied_policies": ["spoofed-policy"], "policy_sources": {"spoofed-policy": "request"}, "routing_decision": {"cause": "forged", "routed_model": "spoofed"}, + "_session_deployment_affinity_ttl": 999999, "internal_call_origin": "autorouter_classifier", "_guardrail_pipelines": [{"name": "spoofed"}], "_pipeline_managed_guardrails": ["evaded"], @@ -719,6 +720,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "applied_policies", "policy_sources", "routing_decision", + "_session_deployment_affinity_ttl", "internal_call_origin", "_guardrail_pipelines", "_pipeline_managed_guardrails", diff --git a/tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py b/tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py b/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py new file mode 100644 index 00000000000..da5dd1934e4 --- /dev/null +++ b/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py @@ -0,0 +1,146 @@ +""" +require_managed_files enforcement for litellm/proxy/vector_store_files_endpoints/endpoints.py + +Every vector-store file route (create, retrieve, content, update, delete) resolves its +caller-supplied file id through _update_request_data_with_managed_file_id before the +provider call, so the guard lives there once and covers all five. + +A raw or forged managed-looking file id has no ownership row, so without the guard it +is attached to a vector store or read back under shared provider credentials. +""" + +import base64 +import os +import sys +from dataclasses import dataclass +from typing import Literal +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from fastapi import HTTPException + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.vector_store_files_endpoints.endpoints import ( + _update_request_data_with_managed_file_id, +) +from litellm.types.utils import SpecialEnums + +RAW_FILE_ID = "file-victim-abc123" +CALLER = UserAPIKeyAuth(api_key="sk-test", user_id="attacker-user", team_id="team-b") + + +@dataclass(frozen=True) +class ManagedResourceAccessCheckerStub: + file_access: Literal["allow", "deny", "missing"] + + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + if self.file_access == "missing": + raise HTTPException(status_code=404, detail=f"File not found: {unified_file_id}") + return self.file_access == "allow" + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return False + + +def _unified_file_id() -> str: + unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-4o-mini", RAW_FILE_ID, "gpt-4o-mini-id" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +async def _resolve( + file_id: str, + file_access: Literal["allow", "deny", "missing"] = "allow", +): + return await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": file_id}, + file_id=file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=CALLER, + managed_files_obj=ManagedResourceAccessCheckerStub(file_access=file_access), + llm_router=None, + ) + + +@pytest.mark.asyncio +async def test_raw_file_id_rejected_when_managed_files_required(): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + await _resolve(RAW_FILE_ID) + + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_model_encoded_file_id_rejected_when_managed_files_required(): + """encode_file_id_with_model output is client-forgeable and carries no ownership + row, so it is not a managed file id.""" + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + encoded = encode_file_id_with_model(RAW_FILE_ID, "gpt-4o-mini", id_type="file") + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + await _resolve(encoded) + + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_forged_unified_file_id_rejected_without_ownership_record(): + forged_id = _unified_file_id() + data = {"vector_store_id": "vs-test", "file_id": forged_id} + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + await _update_request_data_with_managed_file_id( + data=data, + file_id=forged_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=CALLER, + managed_files_obj=ManagedResourceAccessCheckerStub(file_access="missing"), + llm_router=None, + ) + + assert exc.value.status_code == 404 + assert data["file_id"] == forged_id + + +@pytest.mark.asyncio +async def test_other_teams_unified_file_id_rejected(): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + await _resolve(_unified_file_id(), file_access="deny") + + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_owned_unified_file_id_allowed_when_managed_files_required(): + with patch.object(litellm, "require_managed_files", True): + data, original = await _resolve(_unified_file_id()) + + assert original == _unified_file_id() + assert data["file_id"] == RAW_FILE_ID + + +@pytest.mark.asyncio +async def test_raw_file_id_allowed_when_managed_files_not_required(): + with patch.object(litellm, "require_managed_files", False): + data, original = await _resolve(RAW_FILE_ID) + + assert original is None + assert data["file_id"] == RAW_FILE_ID diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 356556f3563..94b6b68855b 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -3414,6 +3414,103 @@ class TestSessionAffinity: def _request_kwargs(session_id: str) -> Dict: return {"metadata": {"session_id": session_id}} + @pytest.mark.asyncio + async def test_hook_response_carries_session_affinity_ttl_on_classify_and_pin_paths( + self, mock_router_instance, session_affinity_config + ): + """The hook response's session_affinity_ttl_seconds is what the Router stamps as + the deployment-affinity marker, so both the classify path (turn 1) and the + session-pin path (turn 2) must carry the configured TTL.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**session_affinity_config, "session_affinity_ttl_seconds": 321}, + ) + request_kwargs = self._request_kwargs("marker-session") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.session_affinity_ttl_seconds == 321 + assert second.session_affinity_ttl_seconds == 321 + + @pytest.mark.parametrize( + "session_affinity,deployment_affinity,plugins,tier_pinned,deployment_pinned", + [ + (False, False, False, False, False), + (False, True, False, False, True), + (True, False, False, True, True), + (True, True, False, True, True), + (False, True, True, False, False), + (True, True, True, False, False), + ], + ) + @pytest.mark.asyncio + async def test_tier_pin_and_deployment_pin_are_independently_gated( + self, + mock_router_instance, + basic_config, + session_affinity, + deployment_affinity, + plugins, + tier_pinned, + deployment_pinned, + ): + """deployment_affinity pins the deployment inside each routed group without pinning which + group the session routes to, so with session_affinity off the tier must still reclassify + on every turn while the marker the Router stamps is still emitted. Turn 1 classifies + REASONING and turn 2 SIMPLE, so a reclassified turn 2 moves model while a tier-pinned one + does not. plugins suppress both pins, since a stale pin would bypass the plugin pipeline.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "session_affinity": session_affinity, + "deployment_affinity": deployment_affinity, + **({"plugins": [_DummyPlugin()]} if plugins else {}), + }, + ) + request_kwargs = self._request_kwargs("matrix-session") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.model == "o1-preview" + assert second.model == ("o1-preview" if tier_pinned else "gpt-4o-mini") + assert (first.session_affinity_ttl_seconds is not None) is deployment_pinned + assert (second.session_affinity_ttl_seconds is not None) is deployment_pinned + + @pytest.mark.asyncio + async def test_hook_response_has_no_session_affinity_ttl_when_disabled_or_plugins( + self, mock_router_instance, basic_config, session_affinity_config + ): + mock_router_instance.cache = DualCache() + disabled_router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "deployment_affinity": False}, + ) + plugin_router = ComplexityRouter( + model_name="test-router-plugins", + litellm_router_instance=mock_router_instance, + complexity_router_config={**session_affinity_config, "plugins": [_DummyPlugin()]}, + ) + disabled = await disabled_router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("s-off"), messages=self.SIMPLE_MESSAGE + ) + with_plugins = await plugin_router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("s-plugins"), messages=self.SIMPLE_MESSAGE + ) + assert disabled.session_affinity_ttl_seconds is None + assert with_plugins.session_affinity_ttl_seconds is None + @pytest.mark.asyncio async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config): """Regression: session_affinity defaults to False, so a shared session_id must NOT diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index e29adda3328..428eb0ceafd 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -465,8 +465,7 @@ async def test_async_pre_call_hook_uses_model_map_key_scope(): Deployment affinity caching uses (user_api_key_hash, model_map_key) -> model_id. """ - cache = AsyncMock() - cache.async_set_cache = AsyncMock() + cache = DualCache() callback = DeploymentAffinityCheck( cache=cache, @@ -489,11 +488,7 @@ async def test_async_pre_call_hook_uses_model_map_key_scope(): model_group="claude-sonnet-4-5@20250929", user_key="user-key-abc", ) - cache.async_set_cache.assert_called_once_with( - expected_cache_key, - {"model_id": "model-id-123"}, - ttl=123, - ) + assert await cache.async_get_cache(key=expected_cache_key) == {"model_id": "model-id-123"} @pytest.mark.asyncio diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py index 0bcb0247aad..4053e6d118b 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,6 +1,6 @@ import os import sys -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -10,6 +10,7 @@ import json import litellm from litellm.caching.dual_cache import DualCache +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) @@ -163,7 +164,7 @@ async def test_async_session_id_affinity_priority_over_user_key(): await callback.cache.async_set_cache( DeploymentAffinityCheck.get_session_affinity_cache_key( - "model_group", "session1" + "model_group", "session1", user_key="user1" ), {"model_id": "deployment-2"}, ) @@ -180,3 +181,439 @@ async def test_async_session_id_affinity_priority_over_user_key(): assert len(filtered) == 1 assert filtered[0]["model_info"]["id"] == "deployment-2" + + +MOCK_RESPONSES_API_RESPONSE = { + "id": "resp_mock-resp-456", + "object": "response", + "created_at": 1741476542, + "status": "completed", + "model": "azure/computer-use-preview", + "output": [], + "usage": { + "input_tokens": 5, + "output_tokens": 10, + "total_tokens": 15, + "output_tokens_details": {"reasoning_tokens": 0}, + }, +} + + +def _smart_router(session_affinity=True, ttl_seconds=777, deployment_affinity=True): + return litellm.Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "target-group", + "complexity_router_config": { + "session_affinity": session_affinity, + "deployment_affinity": deployment_affinity, + "session_affinity_ttl_seconds": ttl_seconds, + "tiers": { + "SIMPLE": "target-group", + "MEDIUM": "target-group", + "COMPLEX": "target-group", + "REASONING": "target-group", + }, + }, + }, + }, + { + "model_name": "target-group", + "litellm_params": { + "model": "azure/computer-use-preview-1", + "api_key": "mock-api-key-1", + "api_version": "mock-api-version", + "api_base": "https://mock-endpoint-1.openai.azure.com", + }, + "model_info": {"id": "deployment-1", "base_model": "computer-use-preview"}, + }, + { + "model_name": "target-group", + "litellm_params": { + "model": "azure/computer-use-preview-2", + "api_key": "mock-api-key-2", + "api_version": "mock-api-version-2", + "api_base": "https://mock-endpoint-2.openai.azure.com", + }, + "model_info": {"id": "deployment-2", "base_model": "computer-use-preview"}, + }, + ], + ) + + +def _session_pin_key(session_id, user_key): + return DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="target-group", session_id=session_id, user_key=user_key + ) + + +def _cleanup_router_callbacks(router): + for callback in router.optional_callbacks or []: + litellm.logging_callback_manager.remove_callback_from_all_lists(callback) + + +async def _one_turn(router, model, session_id, key_hash): + """One request with the shuffle forced to deployment-1, so any other landing + deployment can only come from a pin read.""" + with ( + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post, + patch( + "litellm.router_strategy.simple_shuffle.random.choice", + side_effect=lambda seq: seq[0], + ), + ): + mock_post.return_value = MockResponse(MOCK_RESPONSES_API_RESPONSE, 200) + response = await router.aresponses( + model=model, + input=f"turn for {session_id} {key_hash}", + litellm_metadata={"session_id": session_id, "user_api_key_hash": key_hash}, + ) + return response._hidden_params["model_id"] + + +@pytest.mark.asyncio +async def test_auto_router_session_affinity_writes_scoped_pin_and_follows_it(): + """Turn 1 persists a key-scoped deployment pin; a pin seeded to the deployment + the shuffle would never pick is then followed, proving the read path.""" + router = _smart_router() + try: + served = await _one_turn(router, "smart-router", "write-session", "key-1") + assert await router.cache.async_get_cache(key=_session_pin_key("write-session", "key-1")) == { + "model_id": served + } + assert await router.cache.async_get_cache(key=_session_pin_key("write-session", None)) is None + + await router.cache.async_set_cache( + key=_session_pin_key("read-session", "key-1"), value={"model_id": "deployment-2"} + ) + assert await _one_turn(router, "smart-router", "read-session", "key-1") == "deployment-2" + finally: + _cleanup_router_callbacks(router) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,key_hash", + [ + ("target-group", "key-1"), + ("smart-router", "key-2"), + ], + ids=["direct-group-call", "different-api-key"], +) +async def test_seeded_session_pin_is_invisible_outside_its_scope(model, key_hash): + """The pin binds (auto-routed request, api key, session): a direct call to the + group and a different key reusing the session id must both ignore it.""" + router = _smart_router() + try: + await router.cache.async_set_cache( + key=_session_pin_key("scoped-session", "key-1"), value={"model_id": "deployment-2"} + ) + assert await _one_turn(router, model, "scoped-session", key_hash) == "deployment-1" + finally: + _cleanup_router_callbacks(router) + + +@pytest.mark.asyncio +async def test_marker_write_uses_marker_ttl_and_writes_only_the_session_pin(): + """The write hook honors the marker's TTL over the callback default and writes + no user-key entry when only session affinity is active.""" + import time as time_module + + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + + await callback.async_pre_call_deployment_hook( + kwargs={ + "model_info": {"id": "deployment-1"}, + "metadata": { + "deployment_model_name": "target-group", + "session_id": "ttl-session", + "user_api_key_hash": "key-1", + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 777, + }, + }, + call_type=None, + ) + + session_key = _session_pin_key("ttl-session", "key-1") + assert cache.in_memory_cache.cache_dict == {session_key: {"model_id": "deployment-1"}} + assert cache.in_memory_cache.ttl_dict[session_key] == pytest.approx(time_module.time() + 777, abs=5) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_marker", ["777", True, -5, 0, None]) +async def test_malformed_marker_values_do_not_enable_session_affinity(bad_marker): + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + + await callback.async_pre_call_deployment_hook( + kwargs={ + "model_info": {"id": "deployment-1"}, + "metadata": { + "deployment_model_name": "target-group", + "session_id": "bad-marker-session", + "user_api_key_hash": "key-1", + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: bad_marker, + }, + }, + call_type=None, + ) + + assert cache.in_memory_cache.cache_dict == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enable_user_key", [False, True], ids=["session-pin", "user-key-pin"]) +async def test_concurrent_first_requests_never_flip_a_claimed_pin(enable_user_key): + """Two overlapping first requests select different deployments before either + write lands. Pins are first-writer-wins claims, so the second write must leave + the stored pin unchanged instead of flipping it.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=enable_user_key, + enable_responses_api_affinity=False, + ) + + def racing_kwargs(deployment_id): + metadata = {"deployment_model_name": "target-group", "user_api_key_hash": "key-1"} + if not enable_user_key: + metadata["session_id"] = "racing-session" + metadata[SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY] = 777 + return {"model_info": {"id": deployment_id}, "metadata": metadata} + + await callback.async_pre_call_deployment_hook(kwargs=racing_kwargs("deployment-1"), call_type=None) + await callback.async_pre_call_deployment_hook(kwargs=racing_kwargs("deployment-2"), call_type=None) + + pinned_key = ( + DeploymentAffinityCheck.get_affinity_cache_key(model_group="target-group", user_key="key-1") + if enable_user_key + else _session_pin_key("racing-session", "key-1") + ) + assert await cache.async_get_cache(key=pinned_key) == {"model_id": "deployment-1"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "stored_pin", + [{"model_id": "deployment-1"}, "deployment-1"], + ids=["dict-pin", "legacy-string-pin"], +) +async def test_in_memory_reclaim_slides_idle_window_only_for_the_stored_deployment(stored_pin): + """The pod-local claim mirrors the Lua keepalive: the winning deployment's + re-claim extends the pin's expiry, a losing deployment's claim touches neither + the value nor the expiry, so no-Redis setups keep stickiness across an active + session and ttl bounds idle time there too. Sameness is judged on the pinned + model id, so a legacy string pin written by the Redis branch slides the same.""" + import time as time_module + + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + pin_key = _session_pin_key("slide-session", "key-1") + cache.in_memory_cache.set_cache(pin_key, stored_pin, ttl=10) + first_expiry = cache.in_memory_cache.ttl_dict[pin_key] + + reclaimed = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-1"}, ttl_seconds=777) + assert reclaimed == "deployment-1" + assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) + assert cache.in_memory_cache.ttl_dict[pin_key] > first_expiry + + lost = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-2"}, ttl_seconds=10) + assert lost == "deployment-1" + assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) + + +@pytest.mark.asyncio +async def test_claim_pin_uses_redis_attached_after_construction(): + """The proxy attaches Redis via Router._update_redis_cache after the Router (and + this callback) are built. The claim must resolve the redis tier per call, or pins + silently stay pod-local and cross-pod first-writer-wins is lost.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + + captured = {} + + async def fake_runner(keys, args, client=None): + captured["keys"] = keys + captured["args"] = args + return b'{"model_id": "other-pod-winner"}' + + late_redis = MagicMock() + late_redis.async_register_script = MagicMock(return_value=fake_runner) + cache.redis_cache = late_redis + + import time as time_module + + pin_key = _session_pin_key("late-redis-session", "key-1") + cache.in_memory_cache.set_cache(pin_key, {"model_id": "other-pod-winner"}, ttl=10) + + claimed = await callback._claim_pin( + cache_key=pin_key, + pin_value={"model_id": "our-deployment"}, + ttl_seconds=777, + ) + + assert claimed == "other-pod-winner" + assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) + assert captured["keys"] == (pin_key,) + assert captured["args"] == ('{"model_id": "our-deployment"}', 777) + assert cache.in_memory_cache.get_cache(_session_pin_key("late-redis-session", "key-1")) == { + "model_id": "other-pod-winner" + } + + +@pytest.mark.asyncio +async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): + """A Redis outage must cost cross-pod agreement, never same-pod stickiness. The write + hook only logs this result, so an escaping error would leave the session unpinned and + reshuffle every turn for the whole outage. DualCache's write path, which this claim + replaced, wrote the in-memory tier before ever touching Redis.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + + async def exploding_runner(keys, args, client=None): + raise ConnectionError("redis is down") + + down_redis = MagicMock() + down_redis.async_register_script = MagicMock(return_value=exploding_runner) + cache.redis_cache = down_redis + + key = _session_pin_key("outage-session", "key-1") + claimed = await callback._claim_pin(cache_key=key, pin_value={"model_id": "our-deployment"}, ttl_seconds=777) + + assert claimed == "our-deployment" + assert cache.in_memory_cache.get_cache(key) == {"model_id": "our-deployment"} + + second = await callback._claim_pin(cache_key=key, pin_value={"model_id": "another-deployment"}, ttl_seconds=777) + assert second == "our-deployment" + + +@pytest.mark.asyncio +async def test_marker_session_affinity_read_and_write_agree_for_wildcard_groups(): + """Wildcard deployments keep the literal pattern as model_name on both the read + path and the write path, so the marker-gated pin round-trips through one key.""" + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + request_kwargs = { + "model_info": {"id": "wild-deployment-2"}, + "metadata": { + "deployment_model_name": "openai/*", + "session_id": "wild-session", + "user_api_key_hash": "key-1", + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 777, + }, + } + + await callback.async_pre_call_deployment_hook(kwargs=request_kwargs, call_type=None) + filtered = await callback.async_filter_deployments( + model="openai/gpt-4o", + healthy_deployments=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": f"wild-deployment-{i}"}, + } + for i in (1, 2) + ], + messages=[], + request_kwargs=request_kwargs, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["wild-deployment-2"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,session_affinity,deployment_affinity,expect_marker", + [ + ("smart-router", False, True, True), + ("smart-router", True, False, True), + ("smart-router", False, False, False), + ("target-group", False, True, False), + ], + ids=[ + "deployment-affinity-stamps", + "session-affinity-implies-deployment-pin", + "both-off-no-stamp", + "non-auto-routed-clears", + ], +) +async def test_pre_routing_hook_stamps_or_clears_the_marker_per_attempt( + model, session_affinity, deployment_affinity, expect_marker +): + """Every routing attempt writes or clears the marker, so a fallback from an + auto-routed group to a plain group cannot carry a stale marker. session_affinity + implies the deployment pin: a session frozen onto one group must not re-shuffle + across that group's deployments.""" + router = _smart_router(session_affinity=session_affinity, deployment_affinity=deployment_affinity) + try: + request_kwargs = { + "metadata": {SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 111}, + "litellm_metadata": {SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 111}, + } + await router.async_pre_routing_hook( + model=model, + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "Hello"}], + ) + if expect_marker: + assert request_kwargs["litellm_metadata"][SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY] == 777 + else: + assert SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY not in request_kwargs["metadata"] + assert SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY not in request_kwargs["litellm_metadata"] + finally: + _cleanup_router_callbacks(router) + + +def test_complexity_router_with_deployment_affinity_registers_affinity_callback(): + enabled = _smart_router() + session_only = _smart_router(session_affinity=True, deployment_affinity=False) + disabled = _smart_router(session_affinity=False, deployment_affinity=False) + try: + assert [ + (cb.enable_user_key_affinity, cb.enable_responses_api_affinity, cb.enable_session_id_affinity) + for cb in enabled.optional_callbacks or [] + if isinstance(cb, DeploymentAffinityCheck) + ] == [(False, False, False)] + assert any(isinstance(cb, DeploymentAffinityCheck) for cb in session_only.optional_callbacks or []) + assert not any(isinstance(cb, DeploymentAffinityCheck) for cb in disabled.optional_callbacks or []) + finally: + _cleanup_router_callbacks(enabled) + _cleanup_router_callbacks(session_only) + _cleanup_router_callbacks(disabled) diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/test_litellm/test_budget_ratchet_check.py index 1972c1b6386..22d05f4d00d 100644 --- a/tests/test_litellm/test_budget_ratchet_check.py +++ b/tests/test_litellm/test_budget_ratchet_check.py @@ -68,6 +68,50 @@ def test_new_rule_in_head_is_clean(): assert ratchet.regressions_for("b.json", {}, {"new-rule": _spec_of(5)}) == [] +def test_dropped_rule_that_graduated_to_a_hard_failing_config_is_clean(): + base = {"UP006": _spec_of(0)} + assert ratchet.regressions_for("b.json", base, {}, graduated=("UP006",)) == [] + + +def test_graduation_matches_by_prefix_like_ruff_selectors_do(): + base = {"ANN202": _spec_of(865)} + assert ratchet.regressions_for("b.json", base, {}, graduated=("ANN",)) == [] + + +def test_an_unrelated_graduation_does_not_excuse_a_dropped_rule(): + base = {"C901": _spec_of(3)} + regs = ratchet.regressions_for("b.json", base, {}, graduated=("UP006", "SIM118")) + assert [r.rule for r in regs] == ["C901"] + assert "dropped" in regs[0].detail + + +def test_graduation_never_excuses_a_raised_limit(): + base = {"UP006": _spec_of(0)} + regs = ratchet.regressions_for("b.json", base, {"UP006": _spec_of(7)}, graduated=("UP006",)) + assert [r.rule for r in regs] == ["UP006"] + assert "0 -> 7" in regs[0].detail + + +def test_graduated_selectors_come_from_the_paired_ruff_config(): + selectors = ratchet.graduated_selectors("ruff-strict-budget.json") + assert "UP006" in selectors + assert "ANN" not in selectors + + +def test_budgets_without_a_paired_config_can_never_graduate(): + assert ratchet.graduated_selectors("type-discipline-budget.json") == () + assert ratchet.graduated_selectors("basedpyright-code-budget.json") == () + + +def test_a_selector_the_config_also_ignores_does_not_count_as_graduated(): + lint = {"ignore": ["UP006"], "extend-select": ["UP006", "SIM118"]} + assert ratchet.selectors_hard_failed_by(lint) == ("SIM118",) + + +def test_selectors_hard_failed_by_reads_a_config_with_no_ignore_list(): + assert ratchet.selectors_hard_failed_by({"extend-select": ["UP006"]}) == ("UP006",) + + def test_deleted_budget_file_is_a_regression(): regs = ratchet.regressions_for("b.json", {"LIT006": _spec_of(1)}, None) assert [r.rule for r in regs] == ["*"] diff --git a/tests/test_litellm/test_conftest_isolation.py b/tests/test_litellm/test_conftest_isolation.py index 88889ad7740..15183e68f66 100644 --- a/tests/test_litellm/test_conftest_isolation.py +++ b/tests/test_litellm/test_conftest_isolation.py @@ -1,9 +1,15 @@ import litellm +from litellm import Router +from litellm import router as litellm_router_module from litellm import utils as litellm_utils_module CANARY_MODEL = "conftest-isolation-canary-model" +class _CanaryRouterHolder: + router: Router | None = None + + def test_register_model_ledger_entry_is_scoped_to_this_test(): litellm.register_model({CANARY_MODEL: {"litellm_provider": "openai", "input_cost_per_token": 0.001}}) assert CANARY_MODEL in litellm_utils_module._runtime_registered_model_cost @@ -11,3 +17,20 @@ def test_register_model_ledger_entry_is_scoped_to_this_test(): def test_register_model_ledger_entry_was_rolled_back(): assert CANARY_MODEL not in litellm_utils_module._runtime_registered_model_cost + + +def test_live_router_membership_is_scoped_to_this_test(): + _CanaryRouterHolder.router = Router( + model_list=[ + { + "model_name": "conftest-isolation-canary-router", + "litellm_params": {"model": "openai/conftest-isolation-canary-backend", "api_key": "sk-canary"}, + } + ] + ) + assert _CanaryRouterHolder.router in litellm_router_module._live_routers + + +def test_live_router_membership_was_rolled_back(): + assert _CanaryRouterHolder.router is not None + assert _CanaryRouterHolder.router not in litellm_router_module._live_routers diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index f3f7ab51465..6e8fad9072f 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -133,6 +133,126 @@ def _run(repo: Path, bin_dir: Path, extra_env: dict[str, str]) -> subprocess.Com ) +def _commit_all(repo: Path, message: str) -> None: + subprocess.run(["git", "add", "."], cwd=repo, check=True) + subprocess.run( + ["git", "-c", "user.email=t@t", "-c", "user.name=t", "commit", "-qm", message], + cwd=repo, + check=True, + ) + + +def _set_base_ref(repo: Path) -> None: + subprocess.run( + ["git", "update-ref", "refs/remotes/origin/litellm_internal_staging", "HEAD"], + cwd=repo, + check=True, + ) + + +def test_nothing_staged_scopes_to_working_tree_diff_and_runs_checks(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "litellm" / "foo.py").write_text("x = 2\n") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing staged; scoping to the working tree's diff" in proc.stdout + assert "litellm/foo.py" in proc.stdout + assert "linting Python" in proc.stdout + + +def test_nothing_staged_checks_committed_branch_changes(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "litellm" / "foo.py").write_text("x = 2\n") + _commit_all(repo, "branch change") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing staged; scoping to the working tree's diff" in proc.stdout + assert "linting Python" in proc.stdout + + +def test_nothing_staged_includes_untracked_files_in_scope(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "litellm" / "brand_new.py").write_text("z = 3\n") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "litellm/brand_new.py" in proc.stdout + assert "linting Python" in proc.stdout + + +def test_nothing_staged_deletion_only_branch_triggers_checks(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "litellm" / "foo.py").unlink() + _commit_all(repo, "delete module") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing to check" not in proc.stdout + assert "litellm/foo.py" in proc.stdout + assert "linting Python" in proc.stdout + assert "ruff format --check" not in proc.stdout + + +def test_staged_deletion_triggers_checks_without_feeding_missing_files_to_tools(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + subprocess.run(["git", "rm", "-q", "litellm/foo.py"], cwd=repo, check=True) + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing staged" not in proc.stdout + assert "linting Python" in proc.stdout + assert "ruff format --check" not in proc.stdout + + +def test_deleted_dashboard_file_still_triggers_dashboard_lint(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "ui" / "litellm-dashboard" / "src" / "app.ts").unlink() + _commit_all(repo, "delete dashboard file") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "linting dashboard" in proc.stdout + + +def test_nothing_staged_and_no_changes_is_an_explicit_no_op(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing to check" in proc.stdout + assert "linting Python" not in proc.stdout + + +def test_nothing_staged_without_a_base_ref_fails_with_a_fetch_hint(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 1 + assert "cannot resolve the merge base" in proc.stdout + assert "git fetch origin litellm_internal_staging" in proc.stdout + + +def test_partial_staging_warns_which_checks_were_skipped(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + (repo / "notes.md").write_text("hi\n") + subprocess.run(["git", "add", "notes.md"], cwd=repo, check=True) + (repo / "litellm" / "foo.py").write_text("x = 4\n") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "SKIPPED Python lint (make lint)" in proc.stdout + assert "litellm/foo.py" in proc.stdout + assert "linting Python" not in proc.stdout + + def test_python_dashboard_and_gen_api_blocks_run_concurrently_with_grouped_output(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) barrier_dir = tmp_path / "barrier" @@ -207,8 +327,8 @@ def test_full_output_is_saved_to_a_log_file_in_the_git_dir(tmp_path: Path) -> No assert "linting dashboard" in log assert "API types" in log assert "unstaged/untracked changes" in log - assert f"pre-commit: full log: {log_file}" in proc.stdout - assert "pre-commit: full log:" not in log + assert f"check: full log: {log_file}" in proc.stdout + assert "check: full log:" not in log def test_unwritable_log_warns_and_falls_back_to_running_without_one(tmp_path: Path) -> None: @@ -218,7 +338,7 @@ def test_unwritable_log_warns_and_falls_back_to_running_without_one(tmp_path: Pa assert proc.returncode == 0, proc.stdout + proc.stderr assert "linting Python" in proc.stdout assert "output will not be saved" in proc.stderr - assert "pre-commit: full log:" not in proc.stdout + assert "check: full log:" not in proc.stdout failing = _run(repo, bin_dir, {"STUB_FAIL": "make-lint"}) assert failing.returncode == 1 diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index da2d78edb73..67fa827a8e4 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7550,3 +7550,68 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): assert capture.messages, "the fallback failure path did not log at ERROR" assert huge_message not in "".join(capture.messages) assert max(len(message) for message in capture.messages) < 5_000 +def test_stamp_or_clear_metadata_key_writes_and_clears_both_buckets(): + request_kwargs = {"metadata": {}} + litellm.Router._stamp_or_clear_metadata_key(request_kwargs=request_kwargs, key="probe", value=7) + assert request_kwargs["metadata"]["probe"] == 7 + + stale_kwargs = {"metadata": {"probe": 7}, "litellm_metadata": {"probe": 7}} + litellm.Router._stamp_or_clear_metadata_key(request_kwargs=stale_kwargs, key="probe", value=None) + assert "probe" not in stale_kwargs["metadata"] + assert "probe" not in stale_kwargs["litellm_metadata"] + + +@pytest.mark.parametrize( + "complexity_router_config,expect_callback", + [ + ({"tiers": {"SIMPLE": "gpt-4o"}}, True), + ({"tiers": {"SIMPLE": "gpt-4o"}, "deployment_affinity": False}, False), + ({"tiers": {"SIMPLE": "gpt-4o"}, "deployment_affinity": False, "session_affinity": True}, True), + ], +) +def test_complexity_router_registers_affinity_callback_for_deployment_pin(complexity_router_config, expect_callback): + """The marker the complexity router stamps is inert unless a DeploymentAffinityCheck is + registered to read it, so deployment_affinity has to pull the callback in, and its default-on + means a bare config registers one. Opting out must skip the callback entirely rather than + register a filter that can never fire, including when session_affinity is on, since the two + pins are independent.""" + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + { + "model_name": "my-complexity-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": complexity_router_config, + }, + }, + ] + ) + try: + registered = any(isinstance(cb, DeploymentAffinityCheck) for cb in router.optional_callbacks or []) + assert registered is expect_callback + finally: + for cb in router.optional_callbacks or []: + litellm.logging_callback_manager.remove_callback_from_all_lists(cb) + + +def test_ensure_deployment_affinity_callback_is_idempotent(): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + + router = litellm.Router(model_list=[]) + try: + router._ensure_deployment_affinity_callback() + router._ensure_deployment_affinity_callback() + affinity_callbacks = [ + cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) + ] + assert len(affinity_callbacks) == 1 + finally: + for cb in router.optional_callbacks or []: + litellm.logging_callback_manager.remove_callback_from_all_lists(cb) diff --git a/tests/test_litellm/test_ruff_strict_gate.py b/tests/test_litellm/test_ruff_strict_gate.py index 5ca9613e650..8eed90c121f 100644 --- a/tests/test_litellm/test_ruff_strict_gate.py +++ b/tests/test_litellm/test_ruff_strict_gate.py @@ -1,16 +1,24 @@ import importlib.util +import json +import re +import shutil import subprocess +import sys +import tomllib from pathlib import Path import pytest -_MODULE_PATH = Path(__file__).resolve().parents[2] / "scripts" / "ruff_strict_gate.py" +_REPO_ROOT = Path(__file__).resolve().parents[2] +_MODULE_PATH = _REPO_ROOT / "scripts" / "ruff_strict_gate.py" _spec = importlib.util.spec_from_file_location("ruff_strict_gate", _MODULE_PATH) gate = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(gate) Violation = gate.Violation +_ENABLED_BY_RUFF_DEFAULTS = frozenset({"F401"}) + def rule(name, limit): return {name: {"limit": limit}} @@ -158,3 +166,213 @@ def test_base_point_mid_merge_advances_to_the_merged_in_base_tip(tmp_path): repo, _, base_tip = _branched_repo(tmp_path) _git(repo, "merge", "--no-commit", "--no-ff", "main") assert gate.resolve_base_point("main", cwd=repo) == base_tip + + +def _lint_section(config_name: str) -> dict: + return tomllib.loads((_REPO_ROOT / config_name).read_text())["lint"] + + +def _base_external() -> tuple[str, ...]: + return tuple(_lint_section("ruff.toml")["external"]) + + +def _strict_external() -> tuple[str, ...]: + return tuple(_lint_section("ruff-strict.toml")["external"]) + + +def _strict_selected() -> frozenset: + return frozenset(_lint_section("ruff-strict.toml")["select"]) + + +def _prefix_covered(code: str, prefixes: tuple[str, ...]) -> bool: + return any(code.startswith(prefix) for prefix in prefixes) + + +def _selected_by_the_normal_config() -> frozenset: + return frozenset(_lint_section("ruff.toml")["extend-select"]) | _ENABLED_BY_RUFF_DEFAULTS + + +def _budgeted_rules() -> frozenset: + return frozenset(json.loads((_REPO_ROOT / "ruff-strict-budget.json").read_text())) + + +def _ruff_binary() -> str | None: + beside_interpreter = Path(sys.executable).with_name("ruff") + return str(beside_interpreter) if beside_interpreter.exists() else shutil.which("ruff") + + +_RUFF = _ruff_binary() +_needs_ruff = pytest.mark.skipif(_RUFF is None, reason="ruff is not installed in this environment") + + +def _ruff_output_for_noqa(code: str, *extra_args: str) -> str: + proc = subprocess.run( + [ + _RUFF, + "check", + "--no-cache", + "--stdin-filename", + "litellm/types/_external_probe.py", + *extra_args, + "-", + ], + cwd=_REPO_ROOT, + input=f"def _probe(x: int): # noqa: {code}\n return x\n", + capture_output=True, + text=True, + ) + return proc.stdout + + +def test_every_strict_gate_rule_is_protected_from_base_ruf100(): + unprotected = frozenset( + selector + for selector in _strict_selected() + if not _prefix_covered(selector, _base_external()) + and selector not in _selected_by_the_normal_config() + ) + assert unprotected == frozenset(), ( + f"`ruff check` deletes any `# noqa` naming {sorted(unprotected)} as unused, so suppressing " + "one of those strict-gate rules breaks lint. Cover them in ruff.toml's lint.external or " + "enable them in its lint.extend-select." + ) + + +def test_every_selected_rule_keeps_stale_noqa_detection_somewhere(): + policed_by_strict = frozenset( + selector + for selector in _strict_selected() + if not _prefix_covered(selector, _strict_external()) + ) + policed_by_base = frozenset( + selector + for selector in _selected_by_the_normal_config() + if not _prefix_covered(selector, _base_external()) + ) + shadowed = ( + _strict_selected() | _selected_by_the_normal_config() + ) - policed_by_strict - policed_by_base + assert shadowed == frozenset(), ( + f"no config's RUF100 can ever report a stale `# noqa` for {sorted(shadowed)}: every config " + "that selects each of them also shadows it with an external entry. Narrow the external " + "entry in ruff.toml or ruff-strict.toml." + ) + + +_BASE_OWNED_FAMILY = re.compile(r"E[479]\d+|F\d+|T20\d+") +_BASE_OWNED_SINGLES = frozenset({"PGH004", "RUF008", "RUF009", "RUF100"}) + + +@pytest.fixture(scope="module") +def all_ruff_rule_codes() -> frozenset: + listing = subprocess.run( + [_RUFF, "rule", "--all", "--output-format", "json"], + capture_output=True, + text=True, + ) + assert listing.returncode == 0, listing.stderr + return frozenset( + entry["code"] for entry in json.loads(listing.stdout) if "Removed" not in entry["status"] + ) + + +@_needs_ruff +def test_every_base_owned_rule_is_external_or_selected_in_the_strict_config(all_ruff_rule_codes): + base_owned = frozenset( + code + for code in all_ruff_rule_codes + if _BASE_OWNED_FAMILY.fullmatch(code) or code in _BASE_OWNED_SINGLES + ) + stranded = frozenset( + code + for code in base_owned + if code not in _strict_selected() and not _prefix_covered(code, _strict_external()) + ) + assert stranded == frozenset(), ( + f"the strict gate's RUF100 reads a valid `# noqa` for {sorted(stranded)} as unused, the " + "spurious-breach trap ruff-strict.toml's external override exists to prevent. Cover them " + "there." + ) + double_booked = frozenset( + code + for code in base_owned + if code in _strict_selected() and _prefix_covered(code, _strict_external()) + ) + assert double_booked == frozenset(), ( + f"{sorted(double_booked)} are selected by the strict config yet shadowed by its external " + "list, so their stale suppressions can never be reported. Narrow the external entry in " + "ruff-strict.toml." + ) + + +def test_every_budgeted_rule_is_one_the_gate_actually_measures(): + selectors = tuple(_lint_section("ruff-strict.toml")["select"]) + unmeasured = frozenset(code for code in _budgeted_rules() if not code.startswith(selectors)) + assert unmeasured == frozenset(), ( + f"the gate never counts {sorted(unmeasured)}, so their ceilings are dead config that reads " + "as coverage. Either select them in ruff-strict.toml or drop them from the budget." + ) + + +@_needs_ruff +def test_every_strict_selected_rule_is_budgeted_or_hard_failed_by_the_base_config(all_ruff_rule_codes): + strict_enabled = frozenset( + code + for code in all_ruff_rule_codes + if code.startswith(tuple(_lint_section("ruff-strict.toml")["select"])) + ) + base_hard_failed = tuple(_lint_section("ruff.toml")["extend-select"]) + unpoliced = frozenset( + code + for code in strict_enabled + if code not in _budgeted_rules() + and not code.startswith(base_hard_failed) + and code not in _ENABLED_BY_RUFF_DEFAULTS + ) + assert unpoliced == frozenset(), ( + f"nothing enforces {sorted(unpoliced)}: the gate skips rules missing from the budget, and " + "the base config does not hard-fail them. Re-add a budget ceiling or graduate them into " + "ruff.toml's lint.extend-select." + ) + + +@_needs_ruff +def test_a_noqa_for_a_strict_gate_rule_survives_the_normal_ruff_run(): + assert "RUF100" not in _ruff_output_for_noqa("ANN202") + + +@_needs_ruff +def test_the_external_list_is_what_saves_that_noqa(): + assert "RUF100" in _ruff_output_for_noqa("ANN202", "--config", "lint.external=[]") + + +@_needs_ruff +def test_a_stale_noqa_for_a_locally_enabled_rule_is_still_reported(): + assert "RUF100" in _ruff_output_for_noqa("F401") + + +def _ruff_output_for_source(source: str) -> str: + proc = subprocess.run( + [_RUFF, "check", "--no-cache", "--stdin-filename", "litellm/types/_graduate_probe.py", "-"], + cwd=_REPO_ROOT, + input=source, + capture_output=True, + text=True, + ) + return proc.stdout + + +_DEPRECATED_TYPING_ALIAS = "from typing import List # noqa: UP035\n\n\ndef _probe(x: List[int]) -> None: ...\n" + + +@_needs_ruff +def test_a_graduated_rule_now_fails_the_normal_ruff_run_instead_of_waiting_for_the_gate(): + assert "UP006" in _ruff_output_for_source(_DEPRECATED_TYPING_ALIAS) + + +@_needs_ruff +def test_a_graduated_rule_can_still_be_suppressed_without_tripping_unused_noqa(): + suppressed = _DEPRECATED_TYPING_ALIAS.replace("...\n", "... # noqa: UP006\n") + output = _ruff_output_for_source(suppressed) + assert "UP006" not in output + assert "RUF100" not in output diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 0a0cfe9a617..d621e85f09b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23235 + "limit": 23149 }, "LIT002": { - "limit": 27176 + "limit": 27166 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1091 + "limit": 1086 }, "LIT007": { "limit": 0 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16769 + "limit": 16758 }, "LIT011": { "limit": 5598 diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx index a70b7602e5b..5ac3b8b2b64 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx @@ -1,7 +1,9 @@ -import { act, fireEvent, render, screen } from "@testing-library/react"; -import type { ReactNode } from "react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import type { ReactElement, ReactNode } from "react"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm"; +import { fetchAvailableModels, fetchAvailableModelsForTeam } from "@/components/llm_calls/fetch_models"; import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "./RouterSettingsAccordion"; vi.mock("../networking", () => ({ @@ -9,11 +11,14 @@ vi.mock("../networking", () => ({ })); vi.mock("@/components/llm_calls/fetch_models", () => ({ - fetchAvailableModels: vi.fn().mockResolvedValue([]), + fetchAvailableModels: vi.fn().mockResolvedValue([{ model_group: "global-model" }]), + fetchAvailableModelsForTeam: vi.fn().mockResolvedValue([{ model_group: "openai/*" }, { model_group: "gpt-5" }]), })); vi.mock("../Settings/RouterSettings/Fallbacks/FallbackSelectionForm", () => ({ - FallbackSelectionForm: () => null, + FallbackSelectionForm: ({ availableModels }: { availableModels: string[] }) => ( +
{availableModels.join(",")}
+ ), })); vi.mock("@tremor/react", () => ({ @@ -39,9 +44,19 @@ vi.mock("../router_settings/RouterSettingsForm", () => ({ ), })); +const renderWithQueryClient = (ui: ReactElement) => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return render(ui, { + wrapper: ({ children }: { children: ReactNode }) => ( + {children} + ), + }); +}; + describe("RouterSettingsAccordion", () => { beforeEach(() => { - vi.useFakeTimers(); + vi.clearAllMocks(); + vi.useFakeTimers({ shouldAdvanceTime: true }); }); afterEach(() => { @@ -58,7 +73,7 @@ describe("RouterSettingsAccordion", () => { it("debounces propagation and calls onChange once with the last value", async () => { const onChange = vi.fn<(value: RouterSettingsAccordionValue) => void>(); - render(); + renderWithQueryClient(); await flushInitialPropagation(onChange); fireEvent.click(screen.getByText("set-least-busy")); @@ -81,9 +96,51 @@ describe("RouterSettingsAccordion", () => { expect(onChange.mock.calls[0][0].router_settings.routing_strategy).toBe("usage-based-routing"); }); + it("offers the team's own models, including team-scoped BYOK ones, when a teamId is given", async () => { + renderWithQueryClient(); + + await waitFor(() => { + expect(screen.getByTestId("available-models")).toHaveTextContent("gpt-5,openai/*"); + }); + expect(fetchAvailableModelsForTeam).toHaveBeenCalledWith("test-token", "team-123"); + expect(fetchAvailableModels).not.toHaveBeenCalled(); + }); + + it("falls back to the proxy-wide model listing when no teamId is given", async () => { + renderWithQueryClient(); + + await waitFor(() => { + expect(screen.getByTestId("available-models")).toHaveTextContent("global-model"); + }); + expect(fetchAvailableModelsForTeam).not.toHaveBeenCalled(); + }); + + it("ignores a stale team's model response that resolves after a newer team was selected", async () => { + const resolvers: ((models: { model_group: string }[]) => void)[] = []; + vi.mocked(fetchAvailableModelsForTeam).mockImplementation( + () => new Promise((resolve) => resolvers.push(resolve)) as Promise<{ model_group: string }[]>, + ); + + const { rerender } = renderWithQueryClient(); + await waitFor(() => expect(resolvers).toHaveLength(1)); + + rerender(); + await waitFor(() => expect(resolvers).toHaveLength(2)); + + await act(async () => { + resolvers[1]([{ model_group: "fast-team-model" }]); + resolvers[0]([{ model_group: "slow-team-model" }]); + }); + + await waitFor(() => { + expect(screen.getByTestId("available-models")).toHaveTextContent("fast-team-model"); + }); + expect(screen.getByTestId("available-models")).not.toHaveTextContent("slow-team-model"); + }); + it("does not call onChange when unmounted mid-wait", async () => { const onChange = vi.fn<(value: RouterSettingsAccordionValue) => void>(); - const { unmount } = render(); + const { unmount } = renderWithQueryClient(); await flushInitialPropagation(onChange); fireEvent.click(screen.getByText("set-least-busy")); diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx index 08b917e302f..56227abe9ea 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx @@ -1,12 +1,13 @@ import React, { useEffect, useState, useImperativeHandle, forwardRef, useRef } from "react"; import { TabPanel, TabPanels, TabGroup, TabList, Tab } from "@tremor/react"; +import { useQuery } from "@tanstack/react-query"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { getRouterSettingsCall } from "../networking"; import RouterSettingsForm, { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm"; import { Fallbacks } from "../Settings/RouterSettings/Fallbacks/AddFallbacks"; import { FallbackSelectionForm } from "../Settings/RouterSettings/Fallbacks/FallbackSelectionForm"; import { FallbackGroup } from "../Settings/RouterSettings/Fallbacks/FallbackGroupConfig"; -import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; +import { fetchAvailableModels, fetchAvailableModelsForTeam, ModelGroup } from "@/components/llm_calls/fetch_models"; export interface RouterSettingsAccordionValue { router_settings: { @@ -30,6 +31,7 @@ interface RouterSettingsAccordionProps { value?: RouterSettingsAccordionValue; onChange?: (value: RouterSettingsAccordionValue) => void; modelData?: any; + teamId?: string | null; } export interface RouterSettingsAccordionRef { @@ -39,7 +41,7 @@ export interface RouterSettingsAccordionRef { const PROPAGATE_WAIT_MS = 100; const RouterSettingsAccordion = forwardRef( - ({ accessToken, value, onChange, modelData }, ref) => { + ({ accessToken, value, onChange, modelData, teamId }, ref) => { const [formValue, setFormValue] = useState({ routerSettings: {}, selectedStrategy: null, @@ -47,7 +49,6 @@ const RouterSettingsAccordion = forwardRef([]); const [fallbackGroups, setFallbackGroups] = useState([]); - const [modelInfo, setModelInfo] = useState([]); const [availableRoutingStrategies, setAvailableRoutingStrategies] = useState([]); const [routerFieldsMetadata, setRouterFieldsMetadata] = useState<{ [key: string]: any }>({}); const [routingStrategyDescriptions, setRoutingStrategyDescriptions] = useState<{ [key: string]: string }>({}); @@ -175,21 +176,11 @@ const RouterSettingsAccordion = forwardRef { - if (!accessToken) { - return; - } - const loadModels = async () => { - try { - const uniqueModels = await fetchAvailableModels(accessToken); - setModelInfo(uniqueModels); - } catch (error) { - console.error("Error fetching model info for fallbacks:", error); - } - }; - loadModels(); - }, [accessToken]); + const { data: modelInfo = [] } = useQuery({ + queryKey: ["fallbackAvailableModels", accessToken, teamId ?? null], + queryFn: () => (teamId ? fetchAvailableModelsForTeam(accessToken, teamId) : fetchAvailableModels(accessToken)), + enabled: Boolean(accessToken), + }); // Helper function to build router_settings from current state const buildRouterSettings = (): RouterSettingsAccordionValue["router_settings"] => { diff --git a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx new file mode 100644 index 00000000000..bd691c7f629 --- /dev/null +++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx @@ -0,0 +1,33 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { modelAvailableCall } from "@/components/networking"; +import { fetchAvailableModelsForTeam } from "./fetch_models"; + +vi.mock("@/components/networking", () => ({ + modelAvailableCall: vi.fn(), + modelHubCall: vi.fn(), +})); + +const modelAvailableCallMock = vi.mocked(modelAvailableCall); + +describe("fetchAvailableModelsForTeam", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("requests the models scoped to the team so team-only BYOK models are included", async () => { + modelAvailableCallMock.mockResolvedValue({ + data: [{ id: "all-proxy-models" }, { id: "openai/*" }, { id: "gpt-5-mini" }, { id: "openai/*" }], + }); + + const models = await fetchAvailableModelsForTeam("token", "team-123"); + + expect(modelAvailableCallMock).toHaveBeenCalledWith("token", "", "", false, "team-123"); + expect(models).toEqual([{ model_group: "gpt-5-mini" }, { model_group: "openai/*" }]); + }); + + it("returns an empty list when the team has no models", async () => { + modelAvailableCallMock.mockResolvedValue({ data: [] }); + + expect(await fetchAvailableModelsForTeam("token", "team-123")).toEqual([]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx index 0de98330c2e..a1690b1307e 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx @@ -1,12 +1,22 @@ // fetch_models.ts -import { modelHubCall } from "@/components/networking"; +import { excludeProxyWideSentinel } from "@/components/key_team_helpers/fetch_available_models_team_key"; +import { modelAvailableCall, modelHubCall } from "@/components/networking"; export interface ModelGroup { model_group: string; mode?: string; } +export const fetchAvailableModelsForTeam = async (accessToken: string, teamId: string): Promise => { + const response = await modelAvailableCall(accessToken, "", "", false, teamId); + const modelNames: string[] = (response?.data ?? []).map((model: { id: string }) => model.id); + + return excludeProxyWideSentinel(Array.from(new Set(modelNames))) + .sort((a, b) => a.localeCompare(b)) + .map((model) => ({ model_group: model })); +}; + /** * Fetches available models using modelHubCall and formats them for the selection dropdown. */ diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 513719a2ad9..50e10285148 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -919,7 +919,7 @@ describe("TeamInfoView", () => { }); }; - it("prefills pairs from team metadata, hides UI-managed keys, and round-trips typed values on save", async () => { + it("should preserve metadata types and hide managed keys", async () => { const user = userEvent.setup({ delay: null }); vi.mocked(networking.teamInfoCall).mockResolvedValue( createMockTeamData({ @@ -964,27 +964,6 @@ describe("TeamInfoView", () => { expect(updateArg.model_tpm_limit).toEqual({ "gpt-4": 100 }); }); - it("includes a newly added pair in the team update", async () => { - const user = userEvent.setup({ delay: null }); - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["gpt-4"] })); - vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); - - renderWithProviders(); - await openSettingsEditor(user); - - await user.click(screen.getByRole("button", { name: /add key-value pair/i })); - await user.type(screen.getByPlaceholderText("Key"), "cost_center"); - await user.type(screen.getByPlaceholderText("Value"), "eng-1"); - - await user.click(screen.getByRole("button", { name: /save changes/i })); - - await waitFor(() => { - expect(networking.teamUpdateCall).toHaveBeenCalled(); - }); - - expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].metadata).toMatchObject({ cost_center: "eng-1" }); - }); - it("should keep declared keys as ordinary prefilled rows and submit the edited value", async () => { const user = userEvent.setup({ delay: null }); vi.mocked(useTeamMetadataSchema).mockReturnValue({ diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index e57bf59b2da..6c6fdcaedd6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1215,6 +1215,7 @@ const TeamInfoView: React.FC = ({ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a75c23da1cf..b361a16a0e6 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -31695,6 +31695,12 @@ export interface components { * @description Default model to use if tier cannot be determined */ default_model?: string | null; + /** + * Deployment Affinity + * @description When True and a session_id is resolvable on the request, pin the deployment chosen inside each routed model group and reuse it whenever the session returns to that group, without pinning which group the session routes to. Independent of session_affinity, which pins the model group instead (and always carries this deployment pin with it): with session_affinity off, every turn is still classified on its own merits while a session that escalates to a stronger tier and comes back still lands on the deployment it used before, which is what keeps a provider prompt cache warm. Pins are held per model group, so switching tiers does not disturb the pin left behind in the previous group. On by default because re-shuffling a conversation across deployments of the same model discards that cache for no benefit; set False to keep every turn load-balanced across the group, which is what a deployment set with tight per-deployment rate limits wants. Inert when no session_id is resolvable, since there is nothing to key a pin on, and suppressed when plugins are configured, for the same reason session_affinity is. + * @default true + */ + deployment_affinity: boolean; /** * Dimension Weights * @description Weights for each scoring dimension @@ -31752,13 +31758,13 @@ export interface components { semantic_keyword_matching: boolean; /** * Session Affinity - * @description When True and a session_id is resolvable on the request, pin the model chosen on the session's first turn and reuse it for every later turn, skipping re-classification. Off by default so every turn is classified on its own merits and routed to the cheapest adequate tier. Set True to keep a multi-turn session on one model, which preserves provider prompt caches and avoids cross-model conversation-history errors. + * @description When True and a session_id is resolvable on the request, pin the model chosen on the session's first turn and reuse it for every later turn, skipping re-classification. Off by default so every turn is classified on its own merits and routed to the cheapest adequate tier. Set True to keep a multi-turn session on one model, which preserves provider prompt caches and avoids cross-model conversation-history errors. Always implies the deployment pin regardless of deployment_affinity: the session sticks to one deployment of the pinned model, since freezing the model while re-shuffling its deployments would still go cache-cold. * @default false */ session_affinity: boolean; /** * Session Affinity Ttl Seconds - * @description TTL for the session affinity pin; refreshed on every cache hit + * @description TTL for the session affinity pin; refreshed on every cache hit. Bounds both the session_affinity model pin and the deployment_affinity deployment pin, so it measures idle time for the session's routing decisions rather than total session length * @default 3600 */ session_affinity_ttl_seconds: number;