scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

Merge branch 'litellm_internal_staging' into pre-commit-performance

Browse source 
`make pre-commit` became `make check` on staging, with a working-tree fallback
when nothing is staged, so this branch's three changes move onto that shape:

- `check` depends on `bootstrap-python`, not `bootstrap`
- the script provisions the dashboard off `ui_prettier_changed` /
  `ui_eslint_changed` / `spec_files`, matching the new scope variables, so a
  deleted dashboard file still provisions before the lint that inspects it
- the skip-on-failed-provisioning guard rides on the same variables

Staging's rewrite of `lint_dashboard` dropped the traps that keep the eslint
report from outliving a Ctrl-C, so those come back with the merge.
This commit is contained in:
Claude 2026-08-08 20:30:57 +00:00
commit 340ee50fb2
No known key found for this signature in database
83 changed files with 4912 additions and 1634 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/workflows/test-unit-proxy-endpoints.yml vendored
View file

@ -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

4
CLAUDE.md
View file

@ -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

18
Makefile
View file

@ -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/

28
basedpyright-code-budget.json
View file

@ -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

524
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -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

265
litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py
View file

@ -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)

1
litellm/constants.py
View file

@ -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 = (

20
litellm/integrations/otel/logger.py
View file

@ -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(

67
litellm/integrations/websearch_interception/handler.py
View file

@ -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:

55
litellm/litellm_core_utils/realtime_streaming.py
View file

@ -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):

97
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -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

84
litellm/llms/anthropic/experimental_pass_through/adapters/handler.py
View file

@ -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,

36
litellm/llms/azure/assistants.py
View file

@ -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

2
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -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,

91
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -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)

4
litellm/proxy/_types.py
View file

@ -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(

170
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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

20
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -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)

3
litellm/proxy/common_utils/callback_utils.py
View file

@ -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",

47
litellm/proxy/db/autorouter_session_rollup.py
View file

@ -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:

25
litellm/proxy/fine_tuning_endpoints/endpoints.py
View file

@ -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)
(

173
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -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,

100
litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py
View file

@ -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 []

136
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -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).

60
litellm/proxy/hooks/litellm_skills/main.py
View file

@ -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:

2
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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",

50
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

@ -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(),
)

2
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -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

77
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -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

160
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -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,
)

224
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -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.

187
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -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", []):

76
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -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.

22
litellm/proxy/openai_files_endpoints/files_endpoints.py
View file

@ -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)

123
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -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": {

72
litellm/proxy/vector_store_endpoints/management_endpoints.py
View file

@ -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:

59
litellm/proxy/vector_store_files_endpoints/endpoints.py
View file

@ -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

184
litellm/responses/litellm_completion_transformation/transformation.py
View file

@ -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.

51
litellm/responses/streaming_iterator.py
View file

@ -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

96
litellm/router.py
View file

@ -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(

50
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -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,

30
litellm/router_strategy/complexity_router/config.py
View file

@ -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(

266
litellm/router_utils/pre_call_checks/deployment_affinity_check.py
View file

@ -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",

1
litellm/types/router.py
View file

@ -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)

115
ruff-strict-budget.json
View file

@ -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
}
}

22
ruff.toml
View file

@ -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",
]

51
scripts/budget_ratchet_check.py
View file

@ -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(

2
scripts/install_git_hooks.sh
View file

@ -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

166
scripts/pre_commit_lint.sh
View file

@ -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

2
tests/agent_tests/test_a2a_agent.py
View file

@ -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"):

2
tests/e2e/CONTRIBUTING.md
View file

@ -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

161
tests/e2e/proxy_client.py
View file

@ -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

13
tests/e2e/transport.py
View file

@ -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](

24
tests/litellm_utils_tests/test_proxy_budget_reset.py
View file

@ -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)

4
tests/llm_responses_api_testing/base_responses_api.py
View file

@ -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")

16
tests/proxy_behavior/spend/test_autorouter_session_rollup.py
View file

@ -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(),
)

8
tests/test_litellm/conftest.py
View file

@ -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)

35
tests/test_litellm/integrations/otel/test_otel_v2_logger.py
View file

@ -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():

97
tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py
View file

@ -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

111
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -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."""

173
tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
View file

@ -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

0
tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py Normal file
View file

278
tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py Normal file
View file

@ -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

144
tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py
View file

@ -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()

2
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -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",

0
tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py Normal file
View file

146
tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py Normal file
View file

@ -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

97
tests/test_litellm/router_strategy/test_complexity_router.py
View file

@ -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

9
tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py
View file

@ -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

441
tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py
View file

@ -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)

44
tests/test_litellm/test_budget_ratchet_check.py
View file

@ -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] == ["*"]

23
tests/test_litellm/test_conftest_isolation.py
View file

@ -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

126
tests/test_litellm/test_pre_commit_lint.py
View file

@ -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

65
tests/test_litellm/test_router.py
View file

@ -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)

220
tests/test_litellm/test_ruff_strict_gate.py
View file

@ -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

8
type-discipline-budget.json
View file

@ -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

71
ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx
View file

@ -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[] }) => (
<div data-testid="available-models">{availableModels.join(",")}</div>
),
}));
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 }) => (
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
),
});
};
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(<RouterSettingsAccordion accessToken="test-token" onChange={onChange} />);
renderWithQueryClient(<RouterSettingsAccordion accessToken="test-token" onChange={onChange} />);
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(<RouterSettingsAccordion accessToken="test-token" teamId="team-123" />);
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(<RouterSettingsAccordion accessToken="test-token" />);
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(<RouterSettingsAccordion accessToken="test-token" teamId="team-slow" />);
await waitFor(() => expect(resolvers).toHaveLength(1));
rerender(<RouterSettingsAccordion accessToken="test-token" teamId="team-fast" />);
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(<RouterSettingsAccordion accessToken="test-token" onChange={onChange} />);
const { unmount } = renderWithQueryClient(<RouterSettingsAccordion accessToken="test-token" onChange={onChange} />);
await flushInitialPropagation(onChange);
fireEvent.click(screen.getByText("set-least-busy"));

27
ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx
View file

@ -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<RouterSettingsAccordionRef, RouterSettingsAccordionProps>(
({ accessToken, value, onChange, modelData }, ref) => {
({ accessToken, value, onChange, modelData, teamId }, ref) => {
const [formValue, setFormValue] = useState<RouterSettingsFormValue>({
routerSettings: {},
selectedStrategy: null,
@ -47,7 +49,6 @@ const RouterSettingsAccordion = forwardRef<RouterSettingsAccordionRef, RouterSet
});
const [fallbacks, setFallbacks] = useState<Fallbacks>([]);
const [fallbackGroups, setFallbackGroups] = useState<FallbackGroup[]>([]);
const [modelInfo, setModelInfo] = useState<ModelGroup[]>([]);
const [availableRoutingStrategies, setAvailableRoutingStrategies] = useState<string[]>([]);
const [routerFieldsMetadata, setRouterFieldsMetadata] = useState<{ [key: string]: any }>({});
const [routingStrategyDescriptions, setRoutingStrategyDescriptions] = useState<{ [key: string]: string }>({});
@ -175,21 +176,11 @@ const RouterSettingsAccordion = forwardRef<RouterSettingsAccordionRef, RouterSet
});
}, [accessToken]);
// Fetch available models for fallbacks
useEffect(() => {
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<ModelGroup[]>({
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"] => {

33
ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx Normal file
View file

@ -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([]);
});
});

12
ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx
View file

@ -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<ModelGroup[]> => {
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.
*/

23
ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
View file

@ -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(<TeamInfoView {...defaultProps} />);
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({

1
ui/litellm-dashboard/src/components/team/TeamInfo.tsx
View file

@ -1215,6 +1215,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<RouterSettingsAccordion
ref={routerSettingsRef}
accessToken={accessToken || ""}
teamId={teamId}
value={info.router_settings ? { router_settings: info.router_settings } : undefined}
/>
</Form.Item>

10
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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;