scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge pull request #37073 from BerriAI/litellm_decrease_anys_fable_round2

Browse source
This commit is contained in:
Mateo Wang 2026-08-17 09:40:40 -07:00 • committed by GitHub
commit 1dbed6eb60
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
33 changed files with 1301 additions and 709 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
basedpyright-code-budget.json
View file

@ -1,9 +1,9 @@
{
"reportAny": {
"limit": 22945
"limit": 22344
},
"reportArgumentType": {
"limit": 2579
"limit": 2578
},
"reportAssignmentType": {
"limit": 323
@ -24,13 +24,13 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 7311
"limit": 6991
},
"reportFunctionMemberAccess": {
"limit": 7
},
"reportGeneralTypeIssues": {
"limit": 157
"limit": 154
},
"reportIncompatibleMethodOverride": {
"limit": 56
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5707
"limit": 5681
},
"reportMissingTypeArgument": {
"limit": 15640
"limit": 15609
},
"reportMissingTypeStubs": {
"limit": 40
@ -72,7 +72,7 @@
"limit": 0
},
"reportOptionalMemberAccess": {
"limit": 1069
"limit": 1061
},
"reportOptionalOperand": {
"limit": 0
@ -99,19 +99,19 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44776
"limit": 44709
},
"reportUnknownLambdaType": {
"limit": 113
"limit": 112
},
"reportUnknownMemberType": {
"limit": 39237
"limit": 39154
},
"reportUnknownParameterType": {
"limit": 19967
"limit": 19947
},
"reportUnknownVariableType": {
"limit": 30881
"limit": 30772
},
"reportUnnecessaryCast": {
"limit": 117
@ -123,7 +123,7 @@
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 853
"limit": 851
},
"reportUntypedBaseClass": {
"limit": 0

84
enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py
View file

@ -11,7 +11,7 @@ Endpoints for /project operations
#### PROJECT MANAGEMENT ####
import json
from collections.abc import Mapping, Sequence
from collections.abc import Sequence
from typing import TYPE_CHECKING
from fastapi import APIRouter, Depends, HTTPException, Request
@ -29,7 +29,11 @@ from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy
if TYPE_CHECKING:
from prisma import models as prisma_models
from prisma.actions import LiteLLM_TeamTableActions
from prisma.actions import (
LiteLLM_ProjectTableActions,
LiteLLM_TeamTableActions,
LiteLLM_VerificationTokenActions,
)
router = APIRouter()
@ -39,6 +43,27 @@ def _team_table(prisma_client: PrismaClient) -> "LiteLLM_TeamTableActions[prisma
return team_table
def _project_table(prisma_client: PrismaClient) -> "LiteLLM_ProjectTableActions[prisma_models.LiteLLM_ProjectTable]":
project_table: LiteLLM_ProjectTableActions[prisma_models.LiteLLM_ProjectTable] = (
prisma_client.db.litellm_projecttable
)
return project_table
def _verification_token_table(
prisma_client: PrismaClient,
) -> "LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]":
verification_token_table: LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken] = (
prisma_client.db.litellm_verificationtoken
)
return verification_token_table
def _jsonified(prisma_client: PrismaClient, payload: dict[str, object]) -> dict[str, object]:
jsonified: dict[str, object] = prisma_client.jsonify_object(payload)
return jsonified
async def _check_user_permission_for_project(
user_api_key_dict: UserAPIKeyAuth,
team_id: str | None,
@ -137,7 +162,7 @@ def _check_team_project_limits(
# --- Validate project models are a subset of team models ---
project_models = data.models
team_models = team_object.models or []
team_models: list[str] = team_object.models or []
if project_models and len(team_models) > 0:
# If team has 'all-proxy-models', skip validation as it allows all models
if SpecialModelNames.all_proxy_models.value not in team_models:
@ -188,11 +213,11 @@ async def _create_budget_for_project(
) -> str:
"""Create a budget for the project and return budget_id."""
budget_params = LiteLLM_BudgetTable.model_fields.keys()
_json_data: Mapping[str, object] = data.json(exclude_none=True)
_json_data: dict[str, object] = data.model_dump(exclude_none=True)
_budget_data = {k: v for k, v in _json_data.items() if k in budget_params}
budget_row = LiteLLM_BudgetTable.model_validate(_budget_data)
new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True))
new_budget = _jsonified(prisma_client, budget_row.model_dump(exclude_none=True))
_budget: prisma_models.LiteLLM_BudgetTable = await prisma_client.db.litellm_budgettable.create(
data={
@ -227,7 +252,7 @@ async def _set_project_object_permission(
return None
def _remove_budget_fields_from_project_data(project_data: dict) -> dict:
def _remove_budget_fields_from_project_data(project_data: dict[str, object]) -> dict[str, object]:
"""
Remove budget fields from project data.
Budget fields belong to LiteLLM_BudgetTable, not LiteLLM_ProjectTable.
@ -396,9 +421,7 @@ async def new_project(
data.project_id = str(uuid.uuid4())
else:
# Check if project_id already exists
existing_project = await prisma_client.db.litellm_projecttable.find_unique(
where={"project_id": data.project_id}
)
existing_project = await _project_table(prisma_client).find_unique(where={"project_id": data.project_id})
if existing_project is not None:
raise ProxyException(
message=f"Project id = {data.project_id} already exists. Please use a different project id.",
@ -423,11 +446,14 @@ async def new_project(
)
# Create project row (following organization_endpoints.py pattern)
project_row = LiteLLM_ProjectTable(
**data.json(exclude_none=True),
object_permission_id=object_permission_id,
created_by=user_api_key_dict.user_id or litellm_proxy_admin_name,
updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name,
project_row_payload: dict[str, object] = data.model_dump(exclude_none=True)
project_row = LiteLLM_ProjectTable.model_validate(
{
**project_row_payload,
"object_permission_id": object_permission_id,
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
}
)
for field in LiteLLM_ManagementEndpoint_MetadataFields:
@ -438,7 +464,7 @@ async def new_project(
value=getattr(data, field),
)
new_project_row = prisma_client.jsonify_object(project_row.json(exclude_none=True))
new_project_row = _jsonified(prisma_client, project_row.model_dump(exclude_none=True))
# Remove budget fields (following organization_endpoints.py pattern)
new_project_row = _remove_budget_fields_from_project_data(new_project_row)
@ -560,7 +586,7 @@ async def update_project(
# Fetch existing project
existing_project: (
prisma_models.LiteLLM_ProjectTable | None
) = await prisma_client.db.litellm_projecttable.find_unique(where={"project_id": data.project_id})
) = await _project_table(prisma_client).find_unique(where={"project_id": data.project_id})
if existing_project is None:
raise ProxyException(
@ -617,8 +643,7 @@ async def update_project(
)
# Prepare update data
update_data = data.json(exclude_none=True, exclude={"project_id"})
update_data = prisma_client.jsonify_object(update_data)
update_data = _jsonified(prisma_client, data.model_dump(exclude_none=True, exclude={"project_id"}))
update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name
# Handle budget updates
@ -660,9 +685,10 @@ async def update_project(
# Handle metadata fields
for field in LiteLLM_ManagementEndpoint_MetadataFields:
if field in update_data:
if update_data.get("metadata") is None:
update_data["metadata"] = {}
update_data["metadata"][field] = update_data.pop(field)
existing_metadata = update_data.get("metadata")
metadata_dict: dict[str, object] = existing_metadata if isinstance(existing_metadata, dict) else {}
metadata_dict[field] = update_data.pop(field)
update_data["metadata"] = metadata_dict
# Remove budget fields (following organization_endpoints.py pattern)
update_data = _remove_budget_fields_from_project_data(update_data)
@ -748,11 +774,11 @@ async def delete_project(
detail={"error": "Only admins can delete projects"},
)
deleted_projects = []
deleted_projects: list[prisma_models.LiteLLM_ProjectTable | None] = []
for project_id in data.project_ids:
# Check if project exists
existing_project = await prisma_client.db.litellm_projecttable.find_unique(where={"project_id": project_id})
existing_project = await _project_table(prisma_client).find_unique(where={"project_id": project_id})
if existing_project is None:
raise ProxyException(
@ -765,7 +791,7 @@ async def delete_project(
# Check if there are any keys associated with this project
associated_keys: Sequence[
prisma_models.LiteLLM_VerificationToken
] = await prisma_client.db.litellm_verificationtoken.find_many(where={"project_id": project_id})
] = await _verification_token_table(prisma_client).find_many(where={"project_id": project_id})
if len(associated_keys) > 0:
raise ProxyException(
@ -778,7 +804,7 @@ async def delete_project(
# Delete the project
deleted_project: (
prisma_models.LiteLLM_ProjectTable | None
) = await prisma_client.db.litellm_projecttable.delete(where={"project_id": project_id})
) = await _project_table(prisma_client).delete(where={"project_id": project_id})
await delete_cached_project_object(
project_id=project_id,
@ -829,7 +855,7 @@ async def project_info(
)
# Fetch project
project: prisma_models.LiteLLM_ProjectTable | None = await prisma_client.db.litellm_projecttable.find_unique(
project: prisma_models.LiteLLM_ProjectTable | None = await _project_table(prisma_client).find_unique(
where={"project_id": project_id},
include={"litellm_budget_table": True, "object_permission": True},
)
@ -901,7 +927,7 @@ async def list_projects(
if user_api_key_has_admin_view(user_api_key_dict):
projects: Sequence[
prisma_models.LiteLLM_ProjectTable
] = await prisma_client.db.litellm_projecttable.find_many(
] = await _project_table(prisma_client).find_many(
include={"litellm_budget_table": True, "object_permission": True}
)
else:
@ -911,9 +937,9 @@ async def list_projects(
user_record: prisma_models.LiteLLM_UserTable | None = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
)
user_team_ids: Sequence[str] = user_record.teams if user_record is not None and user_record.teams else []
user_team_ids: list[str] = user_record.teams if user_record is not None and user_record.teams else []
projects = await prisma_client.db.litellm_projecttable.find_many(
projects = await _project_table(prisma_client).find_many(
where={"team_id": {"in": user_team_ids}},
include={"litellm_budget_table": True, "object_permission": True},
)

36
litellm/caching/redis_cache.py
View file

@ -49,7 +49,7 @@ if TYPE_CHECKING:
cluster_pipeline = ClusterPipeline
async_redis_client = Redis
async_redis_cluster_client = RedisCluster
Span = _Span | Any
Span = _Span
else:
pipeline = Any
cluster_pipeline = Any
@ -625,7 +625,11 @@ class RedisCache(BaseCache):
f"{self.namespace}-{hashlib.sha256(script.encode()).hexdigest()[:16]}"
)
async def run_script(keys: Sequence[str], args: Sequence[Any], client: Any = None) -> Any:
async def run_script(
keys: Sequence[str],
args: Sequence[str | bytes | int | float],
client: object = None,
) -> object:
async def execute() -> object:
executor: Callable[..., Awaitable[Any]] | None = litellm.in_memory_llm_clients_cache.get_cache(
key=script_cache_key
@ -650,7 +654,11 @@ class RedisCache(BaseCache):
if hasattr(_redis_client, "register_script"):
registered_script: Final = _redis_client.register_script(script)
async def standalone_executor(keys: Sequence[str], args: Sequence[Any], client: Any = None) -> Any:
async def standalone_executor(
keys: Sequence[str],
args: Sequence[str | bytes | int | float],
client: object = None,
) -> object:
namespaced_keys: Final = tuple(self.check_and_fix_namespace(key=key) for key in keys)
return await registered_script(keys=namespaced_keys, args=args, client=client)
@ -659,7 +667,11 @@ class RedisCache(BaseCache):
if hasattr(_redis_client, "script_load"):
script_sha: Final = _redis_client.script_load(script)
async def cluster_executor(keys: Sequence[str], args: Sequence[Any], client: Any = None) -> Any:
async def cluster_executor(
keys: Sequence[str],
args: Sequence[str | bytes | int | float],
client: object = None,
) -> object:
namespaced_keys: Final = tuple(self.check_and_fix_namespace(key=key) for key in keys)
return await _redis_client.evalsha(script_sha, len(namespaced_keys), *namespaced_keys, *args)
@ -757,7 +769,7 @@ class RedisCache(BaseCache):
async def _pipeline_helper(
self,
pipe: pipeline | cluster_pipeline,
cache_list: list[tuple[Any, Any]],
cache_list: Sequence[tuple[str, object]],
ttl: float | None,
) -> list:
"""
@ -783,7 +795,9 @@ class RedisCache(BaseCache):
return results
@_redis_circuit_breaker_guard
async def async_set_cache_pipeline(self, cache_list: list[tuple[Any, Any]], ttl: float | None = None, **kwargs):
async def async_set_cache_pipeline(
self, cache_list: Sequence[tuple[str, object]], ttl: float | None = None, **kwargs
):
"""
Use Redis Pipelines for bulk write operations
"""
@ -795,7 +809,7 @@ class RedisCache(BaseCache):
start_time: Final = time.time()
print_verbose(f"Set Async Redis Cache: key list: {cache_list}\nttl={ttl}, redis_version={self.redis_version}")
cache_value: Final[Any] = None
cache_value: Final = None
try:
async with _redis_client.pipeline(transaction=False) as pipe:
results: Final = await self._pipeline_helper(pipe, cache_list, ttl)
@ -1074,7 +1088,7 @@ class RedisCache(BaseCache):
# NON blocking - notify users Redis is throwing an exception
verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: ", e)
def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]:
"""
Wrapper to call `mget` on the redis client
@ -1082,7 +1096,7 @@ class RedisCache(BaseCache):
"""
return self.redis_client.mget(keys=keys)
async def _async_run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
async def _async_run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]:
"""
Wrapper to call `mget` on the redis client
@ -1115,7 +1129,7 @@ class RedisCache(BaseCache):
cache_key = self.check_and_fix_namespace(key=cache_key or "")
_keys.append(cache_key)
start_time: Final = time.time()
results: Final[list] = self._run_redis_mget_operation(keys=_keys)
results: Final = self._run_redis_mget_operation(keys=_keys)
end_time: Final = time.time()
_duration: Final = end_time - start_time
self.service_logger_obj.service_success_hook(
@ -1522,7 +1536,7 @@ class RedisCache(BaseCache):
async def async_rpush(
self,
key: str,
values: list[Any],
values: Sequence[str | bytes | int | float],
parent_otel_span: Span | None = None,
**kwargs,
) -> int:

158
litellm/containers/main.py
View file

@ -1,9 +1,11 @@
import asyncio
import contextvars
import json
from collections.abc import Coroutine
from collections.abc import Coroutine, Mapping
from functools import partial
from typing import Any, Final, Literal, overload
from typing import Final, Literal, overload
import httpx
import litellm
from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT
@ -48,16 +50,16 @@ __all__ = [
@client
async def acreate_container(
name: str,
expires_after: dict[str, Any] | None = None,
expires_after: Mapping[str, object] | None = None,
file_ids: list[str] | None = None,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
# LiteLLM specific params,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
) -> ContainerObject:
"""Asynchronously calls the `create_container` function with the given arguments and keyword arguments.
@ -120,9 +122,9 @@ async def acreate_container(
@overload
def create_container(
name: str,
expires_after: dict[str, Any] | None = None,
expires_after: Mapping[str, object] | None = None,
file_ids: list[str] | None = None,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
@ -130,16 +132,16 @@ def create_container(
*,
acreate_container: Literal[True],
**kwargs,
) -> Coroutine[Any, Any, ContainerObject]:
) -> Coroutine[object, object, ContainerObject]:
...
@overload
def create_container(
name: str,
expires_after: dict[str, Any] | None = None,
expires_after: Mapping[str, object] | None = None,
file_ids: list[str] | None = None,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
@ -156,20 +158,20 @@ def create_container(
@client
def create_container(
name: str,
expires_after: dict[str, Any] | None = None,
expires_after: Mapping[str, object] | None = None,
file_ids: list[str] | None = None,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
) -> ContainerObject | Coroutine[Any, Any, ContainerObject]:
) -> ContainerObject | Coroutine[object, object, ContainerObject]:
"""Create a container using the OpenAI Container API.
Currently supports OpenAI
@ -281,13 +283,13 @@ async def alist_containers(
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
) -> ContainerListResponse:
"""Asynchronously list containers.
@ -351,7 +353,7 @@ def list_containers(
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
@ -359,7 +361,7 @@ def list_containers(
*,
alist_containers: Literal[True],
**kwargs,
) -> Coroutine[Any, Any, ContainerListResponse]:
) -> Coroutine[object, object, ContainerListResponse]:
...
@ -368,7 +370,7 @@ def list_containers(
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
@ -387,18 +389,18 @@ def list_containers(
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
) -> ContainerListResponse | Coroutine[Any, Any, ContainerListResponse]:
) -> ContainerListResponse | Coroutine[object, object, ContainerListResponse]:
"""List containers using the OpenAI Container API.
Currently supports OpenAI
@ -481,13 +483,13 @@ def list_containers(
@client
async def aretrieve_container(
container_id: str,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
) -> ContainerObject:
"""Asynchronously retrieve a container.
@ -545,7 +547,7 @@ async def aretrieve_container(
@overload
def retrieve_container(
container_id: str,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
@ -553,14 +555,14 @@ def retrieve_container(
*,
aretrieve_container: Literal[True],
**kwargs,
) -> Coroutine[Any, Any, ContainerObject]:
) -> Coroutine[object, object, ContainerObject]:
...
@overload
def retrieve_container(
container_id: str,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
@ -577,18 +579,18 @@ def retrieve_container(
@client
def retrieve_container(
container_id: str,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
) -> ContainerObject | Coroutine[Any, Any, ContainerObject]:
) -> ContainerObject | Coroutine[object, object, ContainerObject]:
"""Retrieve a container using the OpenAI Container API.
Currently supports OpenAI
@ -696,13 +698,13 @@ def retrieve_container(
@client
async def adelete_container(
container_id: str,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
) -> DeleteContainerResult:
"""Asynchronously delete a container.
@ -760,7 +762,7 @@ async def adelete_container(
@overload
def delete_container(
container_id: str,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
@ -768,14 +770,14 @@ def delete_container(
*,
adelete_container: Literal[True],
**kwargs,
) -> Coroutine[Any, Any, DeleteContainerResult]:
) -> Coroutine[object, object, DeleteContainerResult]:
...
@overload
def delete_container(
container_id: str,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
@ -792,18 +794,18 @@ def delete_container(
@client
def delete_container(
container_id: str,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
) -> DeleteContainerResult | Coroutine[Any, Any, DeleteContainerResult]:
) -> DeleteContainerResult | Coroutine[object, object, DeleteContainerResult]:
"""Delete a container using the OpenAI Container API.
Currently supports OpenAI
@ -914,11 +916,11 @@ async def alist_container_files(
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
) -> ContainerFileListResponse:
"""Asynchronously list files in a container.
@ -985,7 +987,7 @@ def list_container_files(
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600,
timeout: float | httpx.Timeout = 600,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
@ -993,7 +995,7 @@ def list_container_files(
*,
alist_container_files: Literal[True],
**kwargs,
) -> Coroutine[Any, Any, ContainerFileListResponse]:
) -> Coroutine[object, object, ContainerFileListResponse]:
...
@ -1003,7 +1005,7 @@ def list_container_files(
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600,
timeout: float | httpx.Timeout = 600,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
@ -1023,16 +1025,16 @@ def list_container_files(
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
) -> ContainerFileListResponse | Coroutine[Any, Any, ContainerFileListResponse]:
) -> ContainerFileListResponse | Coroutine[object, object, ContainerFileListResponse]:
"""List files in a container using the OpenAI Container API.
Currently supports OpenAI
@ -1125,11 +1127,11 @@ def list_container_files(
async def aupload_container_file(
container_id: str,
file: FileTypes,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
) -> ContainerFileObject:
"""Asynchronously upload a file to a container.
@ -1211,7 +1213,7 @@ async def aupload_container_file(
def upload_container_file(
container_id: str,
file: FileTypes,
timeout=600,
timeout: float | httpx.Timeout = 600,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
@ -1219,7 +1221,7 @@ def upload_container_file(
*,
aupload_container_file: Literal[True],
**kwargs,
) -> Coroutine[Any, Any, ContainerFileObject]:
) -> Coroutine[object, object, ContainerFileObject]:
...
@ -1227,7 +1229,7 @@ def upload_container_file(
def upload_container_file(
container_id: str,
file: FileTypes,
timeout=600,
timeout: float | httpx.Timeout = 600,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
@ -1245,16 +1247,16 @@ def upload_container_file(
def upload_container_file(
container_id: str,
file: FileTypes,
timeout=600, # default to 10 minutes
timeout: float | httpx.Timeout = 600, # default to 10 minutes
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
) -> ContainerFileObject | Coroutine[Any, Any, ContainerFileObject]:
) -> ContainerFileObject | Coroutine[object, object, ContainerFileObject]:
"""Upload a file to a container using the OpenAI Container API.
This endpoint allows uploading files directly to a container session,

42
litellm/litellm_core_utils/litellm_logging.py
View file

@ -13,6 +13,7 @@ import traceback
from collections.abc import Callable, Mapping, Sequence
from datetime import datetime as dt_object
from functools import lru_cache
from types import TracebackType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast
from httpx import Response
@ -1189,6 +1190,7 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["additional_args"] = additional_args
self.model_call_details["log_event_type"] = "post_api_call"
attr: Literal["warning", "debug"]
if self.litellm_request_debug:
attr = "warning"
else:
@ -1802,7 +1804,7 @@ class Logging(LiteLLMLoggingBaseClass):
if self.model_call_details.get("litellm_params") is None:
return
metadata_hidden_params: Final = hidden_params.copy()
response_cost: Final = self.model_call_details.get("response_cost")
response_cost: Final[object] = self.model_call_details.get("response_cost")
if metadata_hidden_params.get("response_cost") is None and response_cost is not None:
metadata_hidden_params["response_cost"] = response_cost
@ -1844,7 +1846,10 @@ class Logging(LiteLLMLoggingBaseClass):
logging_result, start_time, end_time
)
if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None:
standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get(
"standard_logging_object"
)
if standard_logging_payload is not None:
emit_standard_logging_payload(standard_logging_payload)
def _build_standard_logging_payload(
@ -2109,7 +2114,7 @@ class Logging(LiteLLMLoggingBaseClass):
def _success_handler_body(
self,
result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml)
result: object = None,
start_time: datetime.datetime | None = None,
end_time: datetime.datetime | None = None,
cache_hit: bool | None = None,
@ -2150,7 +2155,10 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
complete_streaming_response, start_time, end_time
)
if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None:
standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get(
"standard_logging_object"
)
if standard_logging_payload is not None:
# Only emit for sync requests (async_success_handler handles async)
if is_sync_request:
emit_standard_logging_payload(standard_logging_payload)
@ -2981,7 +2989,7 @@ class Logging(LiteLLMLoggingBaseClass):
global_callbacks=litellm.failure_callback,
)
result = None # result sent to all loggers, init this to None incase it's not created
result: object = None # result sent to all loggers, init this to None incase it's not created
result = redact_message_input_output_from_logging(
model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}),
@ -3395,11 +3403,11 @@ class Logging(LiteLLMLoggingBaseClass):
def _get_assembled_streaming_response(
self,
result: ModelResponse | TextCompletionResponse | ModelResponseStream | ResponseCompletedEvent | Any,
result: ModelResponse | TextCompletionResponse | ModelResponseStream | ResponseCompletedEvent | object,
start_time: datetime.datetime,
end_time: datetime.datetime,
is_async: bool,
streaming_chunks: list[Any],
streaming_chunks: list[object],
) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | None:
if self.stream is not True:
return None
@ -3677,9 +3685,7 @@ def set_callbacks(callback_list, function_id=None):
from sentry_sdk.scrubber import EventScrubber
sentry_sdk_instance = sentry_sdk
sentry_trace_rate = (
os.environ.get("SENTRY_API_TRACE_RATE") if "SENTRY_API_TRACE_RATE" in os.environ else "1.0"
)
sentry_trace_rate = os.environ.get("SENTRY_API_TRACE_RATE", "1.0")
sentry_sample_rate = (
os.environ.get("SENTRY_API_SAMPLE_RATE") if "SENTRY_API_SAMPLE_RATE" in os.environ else "1.0"
)
@ -5150,13 +5156,13 @@ class StandardLoggingPayloadSetup:
# ProxyException uses .code, LiteLLM exceptions use .status_code,
# httpx.HTTPStatusError exposes status only as .response.status_code.
# Stringified for Prisma JSON compatibility.
error_code_attr: Final = getattr(original_exception, "code", None)
error_code_attr: Final[object] = getattr(original_exception, "code", None)
if error_code_attr is not None and str(error_code_attr) not in ("", "None"):
error_status: str = str(error_code_attr)
else:
status_code_attr = getattr(original_exception, "status_code", None)
status_code_attr: object = getattr(original_exception, "status_code", None)
if status_code_attr is None:
response_attr: Final = getattr(original_exception, "response", None)
response_attr: Final[object] = getattr(original_exception, "response", None)
status_code_attr = getattr(response_attr, "status_code", None)
error_status = str(status_code_attr) if status_code_attr is not None else ""
error_class: Final[str] = str(original_exception.__class__.__name__) if original_exception else ""
@ -5165,7 +5171,7 @@ class StandardLoggingPayloadSetup:
# Get traceback information (first 100 lines)
traceback_info = traceback_str or ""
if original_exception:
tb: Final = getattr(original_exception, "__traceback__", None)
tb: Final[TracebackType | None] = getattr(original_exception, "__traceback__", None)
if tb:
tb_lines: Final = traceback.format_tb(tb)
traceback_info += "".join(tb_lines[:MAXIMUM_TRACEBACK_LINES_TO_LOG]) # Limit to first 100 lines
@ -5276,11 +5282,11 @@ class StandardLoggingPayloadSetup:
"""
dynamic_litellm_session_id: Final = litellm_params.get("litellm_session_id")
dynamic_litellm_trace_id: Final = litellm_params.get("litellm_trace_id")
metadata: Final = litellm_params.get("metadata")
metadata: Final[Mapping[str, object] | None] = litellm_params.get("metadata")
metadata_session_id: Final = metadata.get("session_id") if metadata else None
metadata_trace_id: Final = metadata.get("trace_id") if metadata else None
ordered_candidates: Final[tuple[Any, Any, Any, Any]] = (
ordered_candidates: Final[tuple[object, object, object, object]] = (
(dynamic_litellm_trace_id, dynamic_litellm_session_id, metadata_trace_id, metadata_session_id)
if litellm.request_correlation_in_logs
else (dynamic_litellm_session_id, dynamic_litellm_trace_id, metadata_session_id, metadata_trace_id)
@ -5305,10 +5311,10 @@ class StandardLoggingPayloadSetup:
"""
if not litellm.request_correlation_in_logs:
return ""
dynamic_litellm_session_id: Final = litellm_params.get("litellm_session_id")
dynamic_litellm_session_id: Final[object] = litellm_params.get("litellm_session_id")
if dynamic_litellm_session_id:
return str(dynamic_litellm_session_id)
metadata: Final = litellm_params.get("metadata")
metadata: Final[Mapping[str, object] | None] = litellm_params.get("metadata")
metadata_session_id: Final = metadata.get("session_id") if metadata else None
if metadata_session_id:
return str(metadata_session_id)

172
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -3,7 +3,9 @@ import time
from collections.abc import Iterator, Mapping, Sequence
from itertools import groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, TypedDict, Union, cast
from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast
from typing_extensions import ReadOnly, Required
from litellm._logging import verbose_logger
from litellm.types.llms.openai import (
@ -14,6 +16,9 @@ from litellm.types.utils import (
CacheCreationTokenDetails,
ChatCompletionAudioResponse,
ChatCompletionCustomToolCallPayload,
ChatCompletionDeltaCustomToolCall,
ChatCompletionDeltaCustomToolCallPayload,
ChatCompletionDeltaToolCall,
ChatCompletionMessageCustomToolCall,
ChatCompletionMessageToolCall,
Choices,
@ -25,6 +30,7 @@ from litellm.types.utils import (
ModelResponseStream,
PromptTokensDetailsWrapper,
ServerToolUse,
StreamingChoices,
Usage,
)
from litellm.utils import print_verbose, token_counter
@ -79,6 +85,51 @@ class _AudioChunk(TypedDict):
choices: Sequence[_AudioChoice]
_ChunkHiddenParams: TypeAlias = dict[str, object]
class _BaseChunk(TypedDict, total=False):
id: ReadOnly[str]
object: ReadOnly[str]
created: ReadOnly[int]
model: ReadOnly[str]
system_fingerprint: ReadOnly[str | None]
choices: ReadOnly[Required[Sequence[StreamingChoices]]]
_hidden_params: ReadOnly[_ChunkHiddenParams]
class _ToolCallFunctionFragment(TypedDict, total=False):
name: ReadOnly[str]
arguments: ReadOnly[str]
provider_specific_fields: ReadOnly[dict[str, object]]
class _ToolCallCustomFragment(TypedDict, total=False):
name: ReadOnly[str]
input: ReadOnly[str]
class _ToolCallFragment(TypedDict, total=False):
index: ReadOnly[int]
id: ReadOnly[str | None]
type: ReadOnly[str | None]
function: ReadOnly[_ToolCallFunctionFragment | Function | None]
custom: ReadOnly[_ToolCallCustomFragment | None]
provider_specific_fields: ReadOnly[dict[str, object] | None]
class _ToolCallDelta(TypedDict, total=False):
tool_calls: ReadOnly[Sequence[_ToolCallFragment | ChatCompletionDeltaToolCall | ChatCompletionDeltaCustomToolCall]]
class _ToolCallChoice(TypedDict, total=False):
delta: ReadOnly[_ToolCallDelta]
class _ToolCallChunk(TypedDict):
choices: ReadOnly[Sequence[_ToolCallChoice]]
class _UsageBearingChunk(TypedDict, total=False):
usage: Usage | None
_hidden_params: Mapping[str, str]
@ -158,7 +209,7 @@ class ChunkProcessor:
return chunks
def update_model_response_with_hidden_params(
self, model_response: ModelResponse, chunk: Mapping[str, dict[str, object]] | None = None
self, model_response: ModelResponse, chunk: "_BaseChunk | None" = None
) -> ModelResponse:
if chunk is None:
return model_response
@ -214,18 +265,18 @@ class ChunkProcessor:
)
@staticmethod
def _get_chunk_id(chunks: Sequence[Mapping[str, str]]) -> str:
def _get_chunk_id(chunks: Sequence["_BaseChunk"]) -> str:
"""
Chunks:
[{"id": ""}, {"id": "1"}, {"id": "1"}]
"""
for chunk in chunks:
if chunk.get("id"):
return chunk["id"]
if chunk_id := chunk.get("id"):
return chunk_id
return ""
@staticmethod
def _get_model_from_chunks(chunks: Sequence[Mapping[str, str]], first_chunk_model: str) -> str:
def _get_model_from_chunks(chunks: Sequence["_BaseChunk"], first_chunk_model: str) -> str:
"""
Get the actual model from chunks, preferring a model that differs from the first chunk.
@ -241,7 +292,7 @@ class ChunkProcessor:
# Fall back to first chunk's model if no different model found
return first_chunk_model
def build_base_response(self, chunks: list[dict[str, Any]]) -> ModelResponse:
def build_base_response(self, chunks: Sequence["_BaseChunk"]) -> ModelResponse:
chunk = self.first_chunk
id: Final = ChunkProcessor._get_chunk_id(chunks)
object: Final = chunk["object"]
@ -292,7 +343,7 @@ class ChunkProcessor:
@staticmethod
def _iter_tool_call_fragments(
tool_call_chunks: Sequence[Mapping[str, Any]],
tool_call_chunks: Sequence["_ToolCallChunk"],
) -> Iterator[tuple[int, str, str]]:
for chunk in tool_call_chunks:
for choice in chunk["choices"]:
@ -306,21 +357,21 @@ class ChunkProcessor:
index = tool_call.get("index", 0)
function = tool_call.get("function")
if isinstance(function, dict):
if function.get("arguments"):
yield index, "arguments", function["arguments"]
elif getattr(function, "arguments", None):
yield index, "arguments", function.arguments
if fragment_arguments := function.get("arguments"):
yield index, "arguments", fragment_arguments
elif function_arguments := getattr(function, "arguments", None):
yield index, "arguments", function_arguments
custom = tool_call.get("custom")
if isinstance(custom, dict) and custom.get("input"):
yield index, "custom_input", custom["input"]
if isinstance(custom, dict) and (custom_input := custom.get("input")):
yield index, "custom_input", custom_input
else:
index = getattr(tool_call, "index", 0)
function = getattr(tool_call, "function", None)
if getattr(function, "arguments", None):
yield index, "arguments", function.arguments
if object_arguments := getattr(function, "arguments", None):
yield index, "arguments", object_arguments
custom = getattr(tool_call, "custom", None)
if getattr(custom, "input", None):
yield index, "custom_input", custom.input
if object_custom_input := getattr(custom, "input", None):
yield index, "custom_input", object_custom_input
@staticmethod
def _join_fragments_by_index_and_field(
@ -337,7 +388,7 @@ class ChunkProcessor:
)
def get_combined_tool_content(
self, tool_call_chunks: Sequence[Mapping[str, Any]]
self, tool_call_chunks: Sequence["_ToolCallChunk"]
) -> list[
ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall
]: # mutable-ok: assigned verbatim to Message.tool_calls, a list field
@ -364,7 +415,7 @@ class ChunkProcessor:
has_function = "function" in tool_call and tool_call["function"] is not None
has_custom = "custom" in tool_call and tool_call["custom"] is not None
else:
has_function = hasattr(tool_call, "function") and tool_call.function is not None
has_function = getattr(tool_call, "function", None) is not None
has_custom = getattr(tool_call, "custom", None) is not None
if not has_function and not has_custom:
@ -387,61 +438,67 @@ class ChunkProcessor:
# Extract id, type, and function data (handle both dict and object)
if isinstance(tool_call, dict):
if tool_call.get("id"):
tool_call_map[index]["id"] = tool_call["id"]
if tool_call.get("type"):
tool_call_map[index]["type"] = tool_call["type"]
if fragment_id := tool_call.get("id"):
tool_call_map[index]["id"] = fragment_id
if fragment_type := tool_call.get("type"):
tool_call_map[index]["type"] = fragment_type
function = tool_call.get("function", {})
if isinstance(function, dict):
if function.get("name"):
tool_call_map[index]["name"] = function["name"]
if fragment_name := function.get("name"):
tool_call_map[index]["name"] = fragment_name
else:
# function is an object
if hasattr(function, "name") and function.name:
tool_call_map[index]["name"] = function.name
if function_name := getattr(function, "name", None):
tool_call_map[index]["name"] = function_name
custom = tool_call.get("custom")
if isinstance(custom, dict):
if custom.get("name"):
tool_call_map[index]["custom_name"] = custom["name"]
if custom_name := custom.get("name"):
tool_call_map[index]["custom_name"] = custom_name
else:
# tool_call is an object
if hasattr(tool_call, "id") and tool_call.id:
tool_call_map[index]["id"] = tool_call.id
if hasattr(tool_call, "type") and tool_call.type:
tool_call_map[index]["type"] = tool_call.type
if hasattr(tool_call, "function"):
if hasattr(tool_call.function, "name") and tool_call.function.name:
tool_call_map[index]["name"] = tool_call.function.name
if object_function_name := getattr(getattr(tool_call, "function", None), "name", None):
tool_call_map[index]["name"] = object_function_name
custom = getattr(tool_call, "custom", None)
if custom is not None:
if getattr(custom, "name", None):
tool_call_map[index]["custom_name"] = custom.name
object_custom: ChatCompletionDeltaCustomToolCallPayload | None = getattr(
tool_call, "custom", None
)
if object_custom is not None:
if getattr(object_custom, "name", None):
tool_call_map[index]["custom_name"] = object_custom.name
# Preserve provider_specific_fields from streaming chunks
provider_fields = None
provider_fields: object = None
if isinstance(tool_call, dict):
provider_fields = tool_call.get("provider_specific_fields")
if not provider_fields and isinstance(tool_call.get("function"), dict):
provider_fields = tool_call["function"].get("provider_specific_fields")
if not provider_fields and isinstance(fragment_function := tool_call.get("function"), dict):
provider_fields = fragment_function.get("provider_specific_fields")
else:
if hasattr(tool_call, "provider_specific_fields") and tool_call.provider_specific_fields:
provider_fields = tool_call.provider_specific_fields
elif (
hasattr(tool_call, "function")
and hasattr(tool_call.function, "provider_specific_fields")
and tool_call.function.provider_specific_fields
):
provider_fields = tool_call.function.provider_specific_fields
object_provider_fields: object = getattr(tool_call, "provider_specific_fields", None)
if object_provider_fields:
provider_fields = object_provider_fields
else:
function_provider_fields: object = getattr(
getattr(tool_call, "function", None),
"provider_specific_fields",
None,
)
if function_provider_fields:
provider_fields = function_provider_fields
if provider_fields:
# Merge provider_specific_fields if multiple chunks have them
if tool_call_map[index]["provider_specific_fields"] is None:
tool_call_map[index]["provider_specific_fields"] = {}
merged_provider_fields = tool_call_map[index]["provider_specific_fields"]
if merged_provider_fields is None:
merged_provider_fields = {}
tool_call_map[index]["provider_specific_fields"] = merged_provider_fields
if isinstance(provider_fields, dict):
tool_call_map[index]["provider_specific_fields"].update(provider_fields)
merged_provider_fields.update(provider_fields)
joined_fragments: Final = self._join_fragments_by_index_and_field(
self._iter_tool_call_fragments(tool_call_chunks)
@ -762,19 +819,14 @@ class ChunkProcessor:
server_tool_use = usage_chunk.server_tool_use
else:
server_tool_use = ServerToolUse.model_validate(usage_chunk.server_tool_use)
if (
usage_chunk_dict["prompt_tokens_details"] is not None
and getattr(
if usage_chunk_dict["prompt_tokens_details"] is not None:
chunk_web_search_requests: int | None = getattr(
usage_chunk_dict["prompt_tokens_details"],
"web_search_requests",
None,
)
is not None
):
web_search_requests = getattr(
usage_chunk_dict["prompt_tokens_details"],
"web_search_requests",
)
if chunk_web_search_requests is not None:
web_search_requests = chunk_web_search_requests
prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] or prompt_tokens_details

88
litellm/litellm_core_utils/streaming_handler.py
View file

@ -6,7 +6,7 @@ import logging
import threading
import time
import traceback
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Final, NoReturn, Protocol, TypeVar, cast
@ -155,6 +155,33 @@ class _TextCompletionChoiceLike(Protocol):
finish_reason: str | None
class _VertexFunctionCallLike(Protocol):
name: str
args: Mapping[str, Iterable[object]]
class _VertexPartLike(Protocol):
function_call: _VertexFunctionCallLike
class _VertexContentLike(Protocol):
parts: Sequence[_VertexPartLike]
class _VertexFinishReasonLike(Protocol):
name: str
class _VertexCandidateLike(Protocol):
content: _VertexContentLike
finish_reason: _VertexFinishReasonLike
class _VertexChunkLike(Protocol):
text: str
candidates: Sequence[_VertexCandidateLike]
class CustomStreamWrapper:
def __init__(
self,
@ -291,13 +318,13 @@ class CustomStreamWrapper:
that has since taken over the same Task/thread's context.
"""
try:
logging_obj: Final = getattr(self, "logging_obj", None)
logging_obj: Final[object | None] = getattr(self, "logging_obj", None)
if logging_obj is None:
return
method_name: Final = (
"_restore_correlation_context_if_unclaimed" if guarded else "_restore_correlation_context"
)
restore: Final = getattr(logging_obj, method_name, None)
restore: Final[Callable[[], object] | None] = getattr(logging_obj, method_name, None)
if restore is not None:
restore()
except Exception as restore_error: # noqa: BLE001 # best-effort cleanup; must not raise into the caller
@ -1261,18 +1288,18 @@ class CustomStreamWrapper:
raise Exception("An unknown error occurred with the stream")
self.received_finish_reason = "stop"
elif self.custom_llm_provider == "vertex_ai" and not isinstance(chunk, ModelResponseStream):
chunk = cast(Any, chunk)
vertex_chunk: Final = cast(_VertexChunkLike, chunk)
import proto
if hasattr(chunk, "candidates") is True:
if hasattr(vertex_chunk, "candidates") is True:
try:
try:
completion_obj["content"] = chunk.text
completion_obj["content"] = vertex_chunk.text
except Exception as e:
original_exception: Final = e
if "Part has no text." in str(e):
## check for function calling
function_call: Final = chunk.candidates[0].content.parts[0].function_call
function_call: Final = vertex_chunk.candidates[0].content.parts[0].function_call
args_dict: Final = {}
@ -1311,15 +1338,15 @@ class CustomStreamWrapper:
else:
raise original_exception
if (
hasattr(chunk.candidates[0], "finish_reason")
and chunk.candidates[0].finish_reason.name != "FINISH_REASON_UNSPECIFIED"
hasattr(vertex_chunk.candidates[0], "finish_reason")
and vertex_chunk.candidates[0].finish_reason.name != "FINISH_REASON_UNSPECIFIED"
): # every non-final chunk in vertex ai has this
self.received_finish_reason = map_finish_reason(chunk.candidates[0].finish_reason.name)
self.received_finish_reason = map_finish_reason(vertex_chunk.candidates[0].finish_reason.name)
except Exception:
if chunk.candidates[0].finish_reason.name == "SAFETY":
raise Exception(f"The response was blocked by VertexAI. {chunk}")
if vertex_chunk.candidates[0].finish_reason.name == "SAFETY":
raise Exception(f"The response was blocked by VertexAI. {vertex_chunk}")
else:
completion_obj["content"] = str(chunk)
completion_obj["content"] = str(vertex_chunk)
elif self.custom_llm_provider == "petals":
if self.completion_stream is None or len(self.completion_stream) == 0:
if self.received_finish_reason is not None:
@ -1357,13 +1384,14 @@ class CustomStreamWrapper:
if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"]
if response_obj["usage"] is not None:
_text_completion_usage: Final[Usage] = response_obj["usage"]
setattr(
model_response,
"usage",
litellm.Usage(
prompt_tokens=response_obj["usage"].prompt_tokens,
completion_tokens=response_obj["usage"].completion_tokens,
total_tokens=response_obj["usage"].total_tokens,
prompt_tokens=_text_completion_usage.prompt_tokens,
completion_tokens=_text_completion_usage.completion_tokens,
total_tokens=_text_completion_usage.total_tokens,
),
)
elif self.custom_llm_provider == "text-completion-codestral":
@ -1395,15 +1423,17 @@ class CustomStreamWrapper:
if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"]
elif self.custom_llm_provider == "cached_response":
chunk = cast(ModelResponseStream, chunk)
chunk_finish_reason: Final = chunk.choices[0].finish_reason
cached_chunk: Final = cast(ModelResponseStream, chunk)
chunk_finish_reason: Final = cached_chunk.choices[0].finish_reason
response_obj = {
"text": chunk.choices[0].delta.content,
"text": cached_chunk.choices[0].delta.content,
"is_finished": chunk_finish_reason is not None,
"finish_reason": chunk_finish_reason,
"original_chunk": chunk,
"original_chunk": cached_chunk,
"tool_calls": (
chunk.choices[0].delta.tool_calls if hasattr(chunk.choices[0].delta, "tool_calls") else None
cached_chunk.choices[0].delta.tool_calls
if hasattr(cached_chunk.choices[0].delta, "tool_calls")
else None
),
}
@ -1411,11 +1441,11 @@ class CustomStreamWrapper:
if response_obj["tool_calls"] is not None:
completion_obj["tool_calls"] = response_obj["tool_calls"]
print_verbose(f"completion obj content: {completion_obj['content']}")
if hasattr(chunk, "id"):
model_response.id = chunk.id
self.response_id = chunk.id
if hasattr(chunk, "system_fingerprint"):
self.system_fingerprint = chunk.system_fingerprint
if hasattr(cached_chunk, "id"):
model_response.id = cached_chunk.id
self.response_id = cached_chunk.id
if hasattr(cached_chunk, "system_fingerprint"):
self.system_fingerprint = cached_chunk.system_fingerprint
if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"]
else: # openai / azure chat model
@ -2310,16 +2340,16 @@ class CustomStreamWrapper:
def _normalize_status_code(exc: Exception) -> int | None:
"""Best-effort status_code extraction."""
try:
code: Final = getattr(exc, "status_code", None)
code: Final[int | str | None] = getattr(exc, "status_code", None)
if code is not None:
return int(code)
except Exception:
pass
response: Final = getattr(exc, "response", None)
response: Final[object | None] = getattr(exc, "response", None)
if response is not None:
try:
status_code: Final = getattr(response, "status_code", None)
status_code: Final[int | str | None] = getattr(response, "status_code", None)
if status_code is not None:
return int(status_code)
except Exception:

65
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -13,7 +13,7 @@ Pattern Overview:
"""
import json
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, cast
@ -61,6 +61,7 @@ if TYPE_CHECKING:
ModifyResponseException,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -123,7 +124,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _build_streaming_usage_response(
responses_so_far: list[Any],
responses_so_far: list[object],
request_data: dict | None,
) -> ModelResponse | None:
chunks: Final = tuple(response for response in responses_so_far if isinstance(response, (str, bytes)))
@ -141,7 +142,7 @@ class AnthropicMessagesHandler(BaseTranslation):
self,
exc: "ModifyResponseException",
stream_started: bool = False,
responses_so_far: list[Any] | None = None,
responses_so_far: list[object] | None = None,
) -> list[bytes]:
"""
Build an Anthropic SSE sequence delivering the guardrail block message
@ -184,7 +185,7 @@ class AnthropicMessagesHandler(BaseTranslation):
)
return list(FakeAnthropicMessagesStreamIterator(response=block_response))
def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[Any]) -> list[bytes]:
def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[object]) -> list[bytes]:
"""Continue an already-started message: close the open content block,
append the block message as a new text block, then end the message --
without a second message_start."""
@ -234,7 +235,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _content_block_state(
responses_so_far: list[Any],
responses_so_far: list[object],
) -> tuple[int | None, int | None]:
"""From the SSE chunks already sent to the client, return (open
content-block index or None, highest content-block index seen or None).
@ -260,7 +261,7 @@ class AnthropicMessagesHandler(BaseTranslation):
return open_index, max_index
@staticmethod
def _iter_sse_events(item: Any) -> list[dict]:
def _iter_sse_events(item: object) -> list[dict[str, object]]:
"""Yield the event-data dicts in one stream chunk.
Handles both formats this stream can carry (see
@ -271,14 +272,16 @@ class AnthropicMessagesHandler(BaseTranslation):
return [item]
if not isinstance(item, (bytes, bytearray)):
return []
events: Final[list[dict]] = []
events: Final[list[dict[str, object]]] = []
for block in item.decode("utf-8", errors="replace").split("\n\n"):
for line in block.split("\n"):
line = line.strip()
if not line.startswith("data:"):
continue
try:
parsed = json.loads(line[len("data:") :].strip())
parsed: str | int | float | bool | None | Sequence[object] | Mapping[str, object] = json.loads(
line[len("data:") :].strip()
)
except json.JSONDecodeError:
continue
if isinstance(parsed, dict):
@ -315,7 +318,7 @@ class AnthropicMessagesHandler(BaseTranslation):
self,
data: dict,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Any | None = None,
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
) -> Any:
"""
Process input messages by applying guardrails to text content.
@ -467,8 +470,8 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _openai_system_message_to_anthropic(
message: dict[str, Any],
) -> dict[str, Any] | None: # mutable-ok: API message payload
message: dict[str, object],
) -> dict[str, object] | None: # mutable-ok: API message payload
"""Convert an OpenAI system message to the client's Anthropic-shaped entry."""
content: Final = message.get("content")
if isinstance(content, str):
@ -477,14 +480,14 @@ class AnthropicMessagesHandler(BaseTranslation):
) # mutable-ok: API message payload
if not isinstance(content, list):
return None
blocks: Final[list[dict[str, Any]]] = [] # mutable-ok: API message payload
blocks: Final[list[dict[str, object]]] = [] # mutable-ok: API message payload
for block in content:
if not isinstance(block, dict) or block.get("type") != "text":
continue
text = block.get("text")
if not isinstance(text, str) or not text:
continue
anthropic_block: dict[str, Any] = { # mutable-ok: API message payload
anthropic_block: dict[str, object] = { # mutable-ok: API message payload
"type": "text",
"text": text,
} # mutable-ok: API message payload
@ -602,7 +605,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _extract_midturn_system_text(
message: dict[str, Any], # mutable-ok: API message payload
message: Mapping[str, object],
msg_idx: int,
) -> ExtractedInput:
"""Match the adapter's filtering so positional guardrail write-back stays aligned."""
@ -636,7 +639,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@classmethod
def _extract_input_text_and_images(
cls,
message: dict[str, Any],
message: Mapping[str, object],
msg_idx: int,
skip_system_message: bool = False,
skip_tool_message: bool = False,
@ -707,7 +710,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@classmethod
def _extract_tool_result(
cls,
content_item: Mapping[str, Any],
content_item: Mapping[str, object],
msg_idx: int,
content_idx: int,
) -> ExtractedInput:
@ -736,7 +739,7 @@ class AnthropicMessagesHandler(BaseTranslation):
)
@staticmethod
def _image_sources(block: Mapping[str, Any]) -> tuple[str, ...]:
def _image_sources(block: Mapping[str, object]) -> tuple[str, ...]:
source: Final = block.get("source")
if not isinstance(source, Mapping):
return ()
@ -746,7 +749,7 @@ class AnthropicMessagesHandler(BaseTranslation):
async def _apply_guardrail_responses_to_input(
self,
messages: list[dict[str, Any]],
messages: list[dict[str, object]],
responses: list[str],
scanned: tuple[ScannedText, ...],
) -> None:
@ -788,10 +791,10 @@ class AnthropicMessagesHandler(BaseTranslation):
self,
response: "AnthropicMessagesResponse",
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Any | None = None,
user_api_key_dict: Any | None = None,
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
) -> Any:
) -> "AnthropicMessagesResponse":
"""
Process output response by applying guardrails to text content and tool calls.
@ -869,8 +872,8 @@ class AnthropicMessagesHandler(BaseTranslation):
self,
responses_so_far: list[Any],
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Any | None = None,
user_api_key_dict: Any | None = None,
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
) -> list[Any]:
"""
@ -950,8 +953,8 @@ class AnthropicMessagesHandler(BaseTranslation):
def _prepare_request_data(
self,
request_data: dict | None,
response: Any,
user_api_key_dict: Any | None,
response: object,
user_api_key_dict: "UserAPIKeyAuth | None",
key: str,
) -> dict:
"""Ensure request_data has the response/responses_so_far key and metadata."""
@ -968,7 +971,7 @@ class AnthropicMessagesHandler(BaseTranslation):
return request_data
@staticmethod
def _get_response_content(response: Any) -> list[Any]:
def _get_response_content(response: object) -> list[Any]:
"""Extract content list from a dict or object response."""
if isinstance(response, dict):
return response.get("content", []) or []
@ -986,10 +989,10 @@ class AnthropicMessagesHandler(BaseTranslation):
) -> None:
"""Extract text, images, and tool calls from content blocks."""
for content_idx, content_block in enumerate(response_content):
block_dict: dict[str, Any] = {}
block_dict: dict[str, object] = {}
if isinstance(content_block, dict):
block_type = content_block.get("type")
block_dict = cast(dict[str, Any], content_block)
block_dict = cast(dict[str, object], content_block)
elif hasattr(content_block, "type"):
block_type = getattr(content_block, "type", None)
if hasattr(content_block, "model_dump"):
@ -1017,7 +1020,7 @@ class AnthropicMessagesHandler(BaseTranslation):
texts_to_check: list[str],
images_to_check: list[str],
tool_calls_to_check: list["ChatCompletionToolCallChunk"],
response: Any,
response: object,
) -> "GenericGuardrailAPIInputs":
"""Build GenericGuardrailAPIInputs with optional images, tool calls, model."""
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
@ -1212,7 +1215,7 @@ class AnthropicMessagesHandler(BaseTranslation):
def _extract_output_text_and_images(
self,
content_block: dict[str, Any],
content_block: dict[str, object],
content_idx: int,
texts_to_check: list[str],
images_to_check: list[str],
@ -1282,7 +1285,7 @@ class AnthropicMessagesHandler(BaseTranslation):
# Handle both dict and Pydantic object content blocks
if isinstance(content_block, dict):
if content_block.get("type") == "text":
cast(dict[str, Any], content_block)["text"] = guardrail_response
cast(dict[str, object], content_block)["text"] = guardrail_response
elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text":
# Update Pydantic object's text attribute
if hasattr(content_block, "text"):

35
litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
View file

@ -84,6 +84,7 @@ from litellm.types.llms.anthropic import (
AnthropicResponseContentBlockText,
AnthropicResponseContentBlockThinking,
AnthropicResponseContentBlockToolUse,
AnthropicThinkingParam,
AppliedEdit,
ContentBlockDelta,
ContentJsonBlockDelta,
@ -305,7 +306,7 @@ class LiteLLMAnthropicMessagesAdapter:
target["cache_control"] = cache_control
else:
# Fallback for non-dict objects (shouldn't happen in practice)
cast(dict[str, Any], target)["cache_control"] = cache_control
cast(dict[str, object], target)["cache_control"] = cache_control
def translatable_anthropic_params(self) -> list[str]:
"""
@ -323,7 +324,7 @@ class LiteLLMAnthropicMessagesAdapter:
"stop_sequences",
]
def _is_web_search_tool(self, tool: dict[str, Any]) -> bool:
def _is_web_search_tool(self, tool: Mapping[str, object]) -> bool:
"""
Check if a tool is an Anthropic web search tool.
@ -498,7 +499,7 @@ class LiteLLMAnthropicMessagesAdapter:
assistant_message_str = str(content)
elif isinstance(content, dict):
if content.get("type") == "text":
text_block: dict[str, Any] = {
text_block: dict[str, object] = {
"type": "text",
"text": content.get("text", ""),
}
@ -513,10 +514,12 @@ class LiteLLMAnthropicMessagesAdapter:
"name": tool_name,
"arguments": json.dumps(content.get("input", {})),
}
signature = self._extract_signature_from_tool_use_content(cast(dict[str, Any], content))
signature = self._extract_signature_from_tool_use_content(
cast(dict[str, object], content)
)
if signature:
provider_specific_fields: dict[str, Any] = (
provider_specific_fields: dict[str, object] = (
function_chunk.get("provider_specific_fields") or {}
)
provider_specific_fields["thought_signature"] = signature
@ -575,7 +578,7 @@ class LiteLLMAnthropicMessagesAdapter:
@staticmethod
def translate_anthropic_thinking_to_reasoning_effort(
thinking: dict[str, Any],
thinking: AnthropicThinkingParam,
) -> str | None:
"""
Translate Anthropic's thinking parameter to OpenAI's reasoning_effort.
@ -632,9 +635,9 @@ class LiteLLMAnthropicMessagesAdapter:
@staticmethod
def translate_thinking_for_model(
thinking: dict[str, Any],
thinking: AnthropicThinkingParam,
model: str,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Translate Anthropic thinking parameter based on the target model.
@ -670,7 +673,7 @@ class LiteLLMAnthropicMessagesAdapter:
@staticmethod
def _apply_reasoning_summary_wrapping(
reasoning_effort: str,
thinking: dict[str, Any],
thinking: Mapping[str, object],
) -> Any:
"""
Apply the reasoning_effort/summary wrapping rules shared by every
@ -770,7 +773,7 @@ class LiteLLMAnthropicMessagesAdapter:
return new_tools, tool_name_mapping
def translate_anthropic_output_format_to_openai(self, output_format: Any) -> dict[str, Any] | None:
def translate_anthropic_output_format_to_openai(self, output_format: Any) -> dict[str, object] | None:
"""
Translate Anthropic's output_format to OpenAI's response_format.
@ -889,7 +892,7 @@ class LiteLLMAnthropicMessagesAdapter:
model_name: Final = anthropic_message_request.get("model", "")
for block in system_content:
if isinstance(block, dict) and block.get("type") == "text":
text_block: dict[str, Any] = {
text_block: dict[str, object] = {
"type": "text",
"text": block.get("text", ""),
}
@ -959,7 +962,7 @@ class LiteLLMAnthropicMessagesAdapter:
web_search_tools: Final[list[AllAnthropicToolsValues]] = []
regular_tools: Final[list[AllAnthropicToolsValues]] = []
for tool in tools:
cast_tool = cast(dict[str, Any], tool)
cast_tool = cast(dict[str, object], tool)
if self._is_web_search_tool(cast_tool):
web_search_tools.append(cast(AllAnthropicToolsValues, tool))
else:
@ -1007,7 +1010,7 @@ class LiteLLMAnthropicMessagesAdapter:
new_kwargs["output_config"] = effort_config # rebind-ok: out-param store like thinking above
return
reasoning_effort = self.translate_anthropic_thinking_to_reasoning_effort(cast(dict[str, Any], thinking))
reasoning_effort = self.translate_anthropic_thinking_to_reasoning_effort(cast(AnthropicThinkingParam, thinking))
if not reasoning_effort:
return
@ -1020,7 +1023,7 @@ class LiteLLMAnthropicMessagesAdapter:
reasoning_effort = output_config["effort"]
new_kwargs["reasoning_effort"] = self._apply_reasoning_summary_wrapping(
reasoning_effort, cast(dict[str, Any], thinking)
reasoning_effort, cast(dict[str, object], thinking)
)
def _translate_output_format_to_openai(
@ -1040,7 +1043,7 @@ class LiteLLMAnthropicMessagesAdapter:
``output_format`` takes precedence when both are provided.
"""
output_format: Any = anthropic_message_request.get("output_format")
output_format: object = anthropic_message_request.get("output_format")
if not output_format:
output_config: Final = anthropic_message_request.get("output_config")
if isinstance(output_config, dict):
@ -1407,7 +1410,7 @@ class LiteLLMAnthropicMessagesAdapter:
if THOUGHT_SIGNATURE_SEPARATOR in raw_id:
parts = raw_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)
thought_sig = parts[1] if len(parts) > 1 else None
tool_block: dict[str, Any] = {
tool_block: dict[str, object] = {
"type": "tool_use",
"id": normalize_anthropic_tool_use_id(raw_id),
"name": tool_name,

156
litellm/llms/azure_ai/agents/handler.py
View file

@ -22,10 +22,11 @@ import asyncio
import json
import time
import uuid
from collections.abc import AsyncIterator, Callable
from typing import TYPE_CHECKING, Any, Final
from collections.abc import AsyncIterator, Awaitable, Mapping
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict
import httpx
from typing_extensions import ReadOnly
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
@ -33,7 +34,11 @@ from litellm.llms.azure_ai.agents.transformation import (
AzureAIAgentsConfig,
AzureAIAgentsError,
)
from litellm.types.utils import ModelResponse
from litellm.types.llms.openai import (
ChatCompletionAnnotation,
ChatCompletionAnnotationURLCitation,
)
from litellm.types.utils import ModelResponse, ModelResponseStream
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
@ -46,6 +51,69 @@ else:
AsyncHTTPHandler = Any
class _AzureRawAnnotation(TypedDict, total=False):
type: ReadOnly[str]
text: ReadOnly[str]
start_index: ReadOnly[int]
end_index: ReadOnly[int]
url_citation: ReadOnly[ChatCompletionAnnotationURLCitation]
_TransformedAnnotation: TypeAlias = ChatCompletionAnnotation | _AzureRawAnnotation
class _AzureText(TypedDict, total=False):
value: ReadOnly[str]
annotations: ReadOnly[list[_AzureRawAnnotation]]
class _AzureContentItem(TypedDict, total=False):
type: ReadOnly[str]
text: ReadOnly[_AzureText]
class _AzureMessage(TypedDict, total=False):
role: ReadOnly[str]
content: ReadOnly[list[_AzureContentItem]]
class _AzureMessagesData(TypedDict, total=False):
data: ReadOnly[list[_AzureMessage]]
class _CreatedObject(TypedDict):
id: ReadOnly[str]
class _RunError(TypedDict, total=False):
message: ReadOnly[str]
class _RunStatus(TypedDict, total=False):
status: ReadOnly[str]
last_error: ReadOnly[_RunError]
class _SSEDelta(TypedDict, total=False):
content: ReadOnly[list[_AzureContentItem]]
class _SSEEventData(TypedDict, total=False):
id: ReadOnly[str]
content: ReadOnly[list[_AzureContentItem]]
delta: ReadOnly[_SSEDelta]
class _SyncAgentRequest(Protocol):
def __call__(self, method: str, url: str, json_data: Mapping[str, object] | None = None) -> httpx.Response: ...
class _AsyncAgentRequest(Protocol):
def __call__(
self, method: str, url: str, json_data: Mapping[str, object] | None = None
) -> Awaitable[httpx.Response]: ...
class AzureAIAgentsHandler:
"""
Handler for Azure AI Agent Service.
@ -89,7 +157,9 @@ class AzureAIAgentsHandler:
# -------------------------------------------------------------------------
# Response Helpers
# -------------------------------------------------------------------------
def _extract_content_from_messages(self, messages_data: dict) -> tuple[str, list[dict[str, Any]] | None]:
def _extract_content_from_messages(
self, messages_data: _AzureMessagesData
) -> tuple[str, list[_TransformedAnnotation] | None]:
"""Extract assistant content and annotations from the messages response.
Returns (content, annotations) where annotations is a list of
@ -108,8 +178,8 @@ class AzureAIAgentsHandler:
def _transform_annotations(
self,
raw_annotations: list[dict[str, Any]] | None,
) -> list[dict[str, Any]] | None:
raw_annotations: list[_AzureRawAnnotation] | None,
) -> list[_TransformedAnnotation] | None:
"""Transform Azure AI Foundry annotations to OpenAI-compatible format.
Azure AI returns annotations like:
@ -123,11 +193,11 @@ class AzureAIAgentsHandler:
if not raw_annotations:
return None
result: Final[list[dict[str, Any]]] = []
result: Final[list[_TransformedAnnotation]] = []
for ann in raw_annotations:
ann_type = ann.get("type")
if ann_type == "url_citation":
url_citation = dict(ann.get("url_citation", {}))
url_citation: ChatCompletionAnnotationURLCitation = {**ann.get("url_citation", {})}
# Azure puts start/end_index at annotation level; OpenAI
# expects them inside url_citation
if "start_index" in ann and "start_index" not in url_citation:
@ -147,8 +217,8 @@ class AzureAIAgentsHandler:
content: str,
model_response: ModelResponse,
thread_id: str,
messages: list[dict[str, Any]],
annotations: list[dict[str, Any]] | None = None,
messages: list[dict[str, object]],
annotations: list[_TransformedAnnotation] | None = None,
) -> ModelResponse:
"""Build the ModelResponse from agent output."""
from litellm.types.utils import Choices, Message, Usage
@ -201,7 +271,7 @@ class AzureAIAgentsHandler:
api_key: str,
optional_params: dict,
headers: dict | None,
) -> tuple:
) -> tuple[dict[str, str], str, str, str | None, str]:
"""Prepare common parameters for completion.
Azure Foundry Agents API uses Bearer token authentication:
@ -241,7 +311,7 @@ class AzureAIAgentsHandler:
def completion(
self,
model: str,
messages: list[dict[str, Any]],
messages: list[dict[str, object]],
api_base: str,
api_key: str,
model_response: ModelResponse,
@ -266,7 +336,7 @@ class AzureAIAgentsHandler:
api_base,
) = self._prepare_completion_params(model, api_base, api_key, optional_params, headers)
def make_request(method: str, url: str, json_data: dict | None = None) -> httpx.Response:
def make_request(method: str, url: str, json_data: Mapping[str, object] | None = None) -> httpx.Response:
if method == "GET":
return client.get(url=url, headers=headers)
return client.post(
@ -290,14 +360,14 @@ class AzureAIAgentsHandler:
def _execute_agent_flow_sync(
self,
make_request: Callable,
make_request: _SyncAgentRequest,
api_base: str,
api_version: str,
agent_id: str,
thread_id: str | None,
messages: list[dict[str, Any]],
messages: list[dict[str, object]],
optional_params: dict,
) -> tuple[str, str, list[dict[str, Any]] | None]:
) -> tuple[str, str, list[_TransformedAnnotation] | None]:
"""Execute the agent flow synchronously. Returns (thread_id, content, annotations)."""
# Step 1: Create thread if not provided
@ -305,7 +375,8 @@ class AzureAIAgentsHandler:
verbose_logger.debug("Creating thread at: %s", self._build_thread_url(api_base, api_version))
response = make_request("POST", self._build_thread_url(api_base, api_version), {})
self._check_response(response, [200, 201], "Failed to create thread")
thread_id = response.json()["id"]
thread_data: Final[_CreatedObject] = response.json()
thread_id = thread_data["id"]
verbose_logger.debug("Created thread: %s", thread_id)
# At this point thread_id is guaranteed to be a string
@ -325,7 +396,8 @@ class AzureAIAgentsHandler:
response = make_request("POST", self._build_runs_url(api_base, thread_id, api_version), run_payload)
self._check_response(response, [200, 201], "Failed to create run")
run_id: Final = response.json()["id"]
run_data: Final[_CreatedObject] = response.json()
run_id: Final = run_data["id"]
verbose_logger.debug("Created run: %s", run_id)
# Step 4: Poll for completion
@ -334,13 +406,15 @@ class AzureAIAgentsHandler:
response = make_request("GET", status_url)
self._check_response(response, [200], "Failed to get run status")
status = response.json().get("status")
status_data: _RunStatus = response.json()
status = status_data.get("status")
verbose_logger.debug("Run status: %s", status)
if status == "completed":
break
elif status in ["failed", "cancelled", "expired"]:
error_msg = response.json().get("last_error", {}).get("message", "Unknown error")
error_data: _RunStatus = response.json()
error_msg = error_data.get("last_error", {}).get("message", "Unknown error")
raise AzureAIAgentsError(status_code=500, message=f"Run {status}: {error_msg}")
time.sleep(self.config.POLL_INTERVAL_SECONDS)
@ -351,7 +425,8 @@ class AzureAIAgentsHandler:
response = make_request("GET", self._build_list_messages_url(api_base, thread_id, api_version))
self._check_response(response, [200], "Failed to get messages")
content, annotations = self._extract_content_from_messages(response.json())
messages_data: Final[_AzureMessagesData] = response.json()
content, annotations = self._extract_content_from_messages(messages_data)
return thread_id, content, annotations
# -------------------------------------------------------------------------
@ -360,7 +435,7 @@ class AzureAIAgentsHandler:
async def acompletion(
self,
model: str,
messages: list[dict[str, Any]],
messages: list[dict[str, object]],
api_base: str,
api_key: str,
model_response: ModelResponse,
@ -389,7 +464,7 @@ class AzureAIAgentsHandler:
api_base,
) = self._prepare_completion_params(model, api_base, api_key, optional_params, headers)
async def make_request(method: str, url: str, json_data: dict | None = None) -> httpx.Response:
async def make_request(method: str, url: str, json_data: Mapping[str, object] | None = None) -> httpx.Response:
if method == "GET":
return await client.get(url=url, headers=headers)
return await client.post(
@ -413,14 +488,14 @@ class AzureAIAgentsHandler:
async def _execute_agent_flow_async(
self,
make_request: Callable,
make_request: _AsyncAgentRequest,
api_base: str,
api_version: str,
agent_id: str,
thread_id: str | None,
messages: list[dict[str, Any]],
messages: list[dict[str, object]],
optional_params: dict,
) -> tuple[str, str, list[dict[str, Any]] | None]:
) -> tuple[str, str, list[_TransformedAnnotation] | None]:
"""Execute the agent flow asynchronously. Returns (thread_id, content, annotations)."""
# Step 1: Create thread if not provided
@ -428,7 +503,8 @@ class AzureAIAgentsHandler:
verbose_logger.debug("Creating thread at: %s", self._build_thread_url(api_base, api_version))
response = await make_request("POST", self._build_thread_url(api_base, api_version), {})
self._check_response(response, [200, 201], "Failed to create thread")
thread_id = response.json()["id"]
thread_data: Final[_CreatedObject] = response.json()
thread_id = thread_data["id"]
verbose_logger.debug("Created thread: %s", thread_id)
# At this point thread_id is guaranteed to be a string
@ -448,7 +524,8 @@ class AzureAIAgentsHandler:
response = await make_request("POST", self._build_runs_url(api_base, thread_id, api_version), run_payload)
self._check_response(response, [200, 201], "Failed to create run")
run_id: Final = response.json()["id"]
run_data: Final[_CreatedObject] = response.json()
run_id: Final = run_data["id"]
verbose_logger.debug("Created run: %s", run_id)
# Step 4: Poll for completion
@ -457,13 +534,15 @@ class AzureAIAgentsHandler:
response = await make_request("GET", status_url)
self._check_response(response, [200], "Failed to get run status")
status = response.json().get("status")
status_data: _RunStatus = response.json()
status = status_data.get("status")
verbose_logger.debug("Run status: %s", status)
if status == "completed":
break
elif status in ["failed", "cancelled", "expired"]:
error_msg = response.json().get("last_error", {}).get("message", "Unknown error")
error_data: _RunStatus = response.json()
error_msg = error_data.get("last_error", {}).get("message", "Unknown error")
raise AzureAIAgentsError(status_code=500, message=f"Run {status}: {error_msg}")
await asyncio.sleep(self.config.POLL_INTERVAL_SECONDS)
@ -474,7 +553,8 @@ class AzureAIAgentsHandler:
response = await make_request("GET", self._build_list_messages_url(api_base, thread_id, api_version))
self._check_response(response, [200], "Failed to get messages")
content, annotations = self._extract_content_from_messages(response.json())
messages_data: Final[_AzureMessagesData] = response.json()
content, annotations = self._extract_content_from_messages(messages_data)
return thread_id, content, annotations
# -------------------------------------------------------------------------
@ -483,7 +563,7 @@ class AzureAIAgentsHandler:
async def acompletion_stream(
self,
model: str,
messages: list[dict[str, Any]],
messages: list[dict[str, object]],
api_base: str,
api_key: str,
logging_obj: LiteLLMLoggingObj,
@ -491,7 +571,7 @@ class AzureAIAgentsHandler:
litellm_params: dict,
timeout: float,
headers: dict | None = None,
) -> AsyncIterator:
) -> AsyncIterator[ModelResponseStream]:
"""Execute async streaming completion using Azure Agent Service with native SSE."""
import litellm
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
@ -505,12 +585,12 @@ class AzureAIAgentsHandler:
) = self._prepare_completion_params(model, api_base, api_key, optional_params, headers)
# Build payload for create-thread-and-run with streaming
thread_messages: Final = []
thread_messages: Final[list[dict[str, object]]] = []
for msg in messages:
if msg.get("role") in ["user", "system"]:
thread_messages.append({"role": "user", "content": msg.get("content", "")})
payload: Final[dict[str, Any]] = {
payload: Final[dict[str, object]] = {
"assistant_id": agent_id,
"stream": True,
}
@ -552,14 +632,14 @@ class AzureAIAgentsHandler:
self,
response: httpx.Response,
model: str,
) -> AsyncIterator:
) -> AsyncIterator[ModelResponseStream]:
"""Process SSE stream and yield OpenAI-compatible streaming chunks."""
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
response_id: Final = f"chatcmpl-{uuid.uuid4().hex[:8]}"
created: Final = int(time.time())
thread_id = None
collected_annotations: list[dict[str, Any]] | None = None
collected_annotations: list[_TransformedAnnotation] | None = None
current_event = None
@ -597,7 +677,7 @@ class AzureAIAgentsHandler:
return
try:
data = json.loads(data_str)
data: _SSEEventData = json.loads(data_str)
except json.JSONDecodeError:
continue

15
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -49,6 +49,7 @@ from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse
from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
@ -5930,10 +5931,10 @@ class BaseLLMHTTPHandler:
self,
api_base: str,
api_key: str,
request_data: dict[str, Any],
request_data: dict[str, object],
logging_obj: LiteLLMLoggingObj,
timeout: float | httpx.Timeout,
provider_config: Any | None = None,
provider_config: BaseRealtimeHTTPConfig | None = None,
model: str | None = None,
extra_headers: dict[str, object] | None = None,
client: HTTPHandler | AsyncHTTPHandler | None = None,
@ -5963,10 +5964,10 @@ class BaseLLMHTTPHandler:
self,
api_base: str,
api_key: str,
request_data: dict[str, Any],
request_data: dict[str, object],
logging_obj: LiteLLMLoggingObj,
timeout: float | httpx.Timeout,
provider_config: Any | None = None,
provider_config: BaseRealtimeHTTPConfig | None = None,
model: str | None = None,
extra_headers: dict[str, object] | None = None,
client: HTTPHandler | AsyncHTTPHandler | None = None,
@ -5992,7 +5993,7 @@ class BaseLLMHTTPHandler:
endpoint: Literal["client_secrets", "transcription_sessions"],
api_base: str,
api_key: str,
request_data: dict[str, Any],
request_data: dict[str, object],
logging_obj: LiteLLMLoggingObj,
timeout: float | httpx.Timeout,
provider_config: Any | None = None,
@ -11077,7 +11078,7 @@ class BaseLLMHTTPHandler:
client: HTTPHandler | AsyncHTTPHandler | None = None,
stream: bool = False,
litellm_metadata: dict[str, object] | None = None,
system_instruction: Any | None = None,
system_instruction: object | None = None,
) -> Any:
"""
Handles Google GenAI generate content requests.
@ -11208,7 +11209,7 @@ class BaseLLMHTTPHandler:
client: AsyncHTTPHandler | None = None,
stream: bool = False,
litellm_metadata: dict[str, object] | None = None,
system_instruction: Any | None = None,
system_instruction: object | None = None,
) -> Any:
"""
Async version of the generate content handler.

148
litellm/proxy/_experimental/mcp_server/db.py
View file

@ -4,7 +4,7 @@ import hashlib
import json
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, TypeVar, cast
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -45,10 +45,47 @@ from litellm.types.mcp import MCPCredentials
if TYPE_CHECKING:
from prisma import models as prisma_db_models
from prisma import types as prisma_db_types
from prisma.actions import LiteLLM_MCPUserCredentialsActions, LiteLLM_MCPUserEnvVarsActions
from litellm.types.mcp_server.mcp_server_manager import MCPServer
_RowT = TypeVar("_RowT")
class _TableActions(Protocol[_RowT]):
async def find_unique(
self, where: Mapping[str, object], include: Mapping[str, object] | None = None
) -> _RowT | None: ...
async def find_many(
self,
take: int | None = None,
where: Mapping[str, object] | None = None,
order: Mapping[str, object] | None = None,
) -> list[_RowT]: ...
async def create(self, data: Mapping[str, object]) -> _RowT: ...
async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT | None: ...
async def delete(self, where: Mapping[str, object]) -> _RowT | None: ...
async def delete_many(self, where: Mapping[str, object] | None = None) -> int: ...
class _UserEnvVarsTransactionClient(Protocol):
litellm_mcpuserenvvars: "_TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]"
async def execute_raw(self, query: str, *args: object) -> int: ...
class _UserEnvVarsTransaction(Protocol):
async def __aenter__(self) -> _UserEnvVarsTransactionClient: ...
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
_AUTH_FLOW_SCOPED_FIELDS: Final["frozenset[str]"] = frozenset(
{
"issuer",
@ -434,23 +471,54 @@ def _credentials_blob_to_mutable_dict(blob: str | Mapping[str, object]) -> dict[
return parsed_blob
def _mcp_server_table_actions(
prisma_client: PrismaClient,
) -> "_TableActions[prisma_db_models.LiteLLM_MCPServerTable]":
table: Final[_TableActions[prisma_db_models.LiteLLM_MCPServerTable]] = MCPServerRepository(prisma_client).table
return table
def _verification_token_table_actions(
prisma_client: PrismaClient,
) -> "_TableActions[prisma_db_models.LiteLLM_VerificationToken]":
table: Final[_TableActions[prisma_db_models.LiteLLM_VerificationToken]] = VerificationTokenRepository(
prisma_client
).table
return table
def _team_table_actions(
prisma_client: PrismaClient,
) -> "_TableActions[prisma_db_models.LiteLLM_TeamTable]":
table: Final[_TableActions[prisma_db_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
return table
def _oauth_client_table_actions(
prisma_client: PrismaClient,
) -> "_TableActions[prisma_db_models.LiteLLM_MCPServerOAuthClient]":
table: Final[_TableActions[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = MCPServerOAuthClientRepository(
prisma_client
).table
return table
def _db_transaction_manager(prisma_client: PrismaClient) -> _UserEnvVarsTransaction:
manager: Final[_UserEnvVarsTransaction] = prisma_client.db.tx()
return manager
async def _db_find_mcp_server_rows(
prisma_client: PrismaClient,
where: "prisma_db_types.LiteLLM_MCPServerTableWhereInput | None" = None,
) -> "list[prisma_db_models.LiteLLM_MCPServerTable]":
rows: list[prisma_db_models.LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_many(
where=where
)
return rows
return await _mcp_server_table_actions(prisma_client).find_many(where=where)
async def _db_find_mcp_server_row(
prisma_client: PrismaClient, server_id: str
) -> "prisma_db_models.LiteLLM_MCPServerTable | None":
row: prisma_db_models.LiteLLM_MCPServerTable | None = await MCPServerRepository(prisma_client).table.find_unique(
where={"server_id": server_id}
)
return row
return await _mcp_server_table_actions(prisma_client).find_unique(where={"server_id": server_id})
async def _db_update_mcp_server_row(
@ -467,19 +535,17 @@ async def _db_update_mcp_server_row(
def _user_credential_actions(
prisma_client: PrismaClient,
) -> "LiteLLM_MCPUserCredentialsActions[prisma_db_models.LiteLLM_MCPUserCredentials]":
table: Final[LiteLLM_MCPUserCredentialsActions[prisma_db_models.LiteLLM_MCPUserCredentials]] = (
MCPUserCredentialsRepository(prisma_client).table
)
) -> "_TableActions[prisma_db_models.LiteLLM_MCPUserCredentials]":
table: Final[_TableActions[prisma_db_models.LiteLLM_MCPUserCredentials]] = MCPUserCredentialsRepository(
prisma_client
).table
return table
def _user_env_var_actions(
prisma_client: PrismaClient,
) -> "LiteLLM_MCPUserEnvVarsActions[prisma_db_models.LiteLLM_MCPUserEnvVars]":
table: Final[LiteLLM_MCPUserEnvVarsActions[prisma_db_models.LiteLLM_MCPUserEnvVars]] = (
prisma_client.db.litellm_mcpuserenvvars
)
) -> "_TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]":
table: Final[_TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]] = prisma_client.db.litellm_mcpuserenvvars
return table
@ -501,7 +567,7 @@ async def _db_find_user_credential_rows(
async def _db_upsert_user_credential_row(
prisma_client: PrismaClient, user_id: str, server_id: str, credential_b64: str
) -> None:
await MCPUserCredentialsRepository(prisma_client).table.upsert(
await _user_credential_actions(prisma_client).upsert(
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}},
data={
"create": {
@ -592,9 +658,9 @@ async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]
"""
Returns the matching mcp servers from the db with the server_ids
"""
_mcp_servers: Final[list[prisma_db_models.LiteLLM_MCPServerTable]] = await MCPServerRepository(
_mcp_servers: Final[list[prisma_db_models.LiteLLM_MCPServerTable]] = await _mcp_server_table_actions(
prisma_client
).table.find_many(
).find_many(
where={
"server_id": {"in": server_ids},
}
@ -612,9 +678,9 @@ async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, toke
"""
Returns the mcp servers from the db for the verification token
"""
verification_token_record: prisma_db_models.LiteLLM_VerificationToken | None = await VerificationTokenRepository(
prisma_client
).table.find_unique(
verification_token_record: (
prisma_db_models.LiteLLM_VerificationToken | None
) = await _verification_token_table_actions(prisma_client).find_unique(
where={
"token": token,
},
@ -633,7 +699,7 @@ async def get_mcp_servers_by_team(prisma_client: PrismaClient, team_id: str) ->
"""
Returns the mcp servers from the db for the team id
"""
team_record: prisma_db_models.LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.find_unique(
team_record: prisma_db_models.LiteLLM_TeamTable | None = await _team_table_actions(prisma_client).find_unique(
where={
"team_id": team_id,
},
@ -760,9 +826,9 @@ async def delete_mcp_server(
if deleted_server is not None:
credential_user_ids: list[str] = []
try:
credential_rows: Sequence[
prisma_db_models.LiteLLM_MCPUserCredentials
] = await prisma_client.db.litellm_mcpusercredentials.find_many(where={"server_id": server_id})
credential_rows: Sequence[prisma_db_models.LiteLLM_MCPUserCredentials] = await _user_credential_actions(
prisma_client
).find_many(where={"server_id": server_id})
credential_user_ids = [row.user_id for row in credential_rows]
except Exception as e: # noqa: BLE001 - enumeration is best-effort; cached tokens expire by TTL
verbose_proxy_logger.warning(
@ -771,9 +837,9 @@ async def delete_mcp_server(
e,
)
for model, label in (
(prisma_client.db.litellm_mcpusercredentials, "credential"),
(prisma_client.db.litellm_mcpuserenvvars, "env var"),
(prisma_client.db.litellm_mcpserveroauthclient, "OAuth client"),
(_user_credential_actions(prisma_client), "credential"),
(_user_env_var_actions(prisma_client), "env var"),
(_oauth_client_table_actions(prisma_client), "OAuth client"),
):
try:
await model.delete_many(where={"server_id": server_id})
@ -1042,9 +1108,9 @@ async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, s
LiteLLM_MCPServerTable row, so their dynamically registered client lives here keyed
by server_id. The returned value is the raw credentials blob for
``_get_persisted_dcr_credentials`` to parse."""
row: Final[prisma_db_models.LiteLLM_MCPServerOAuthClient | None] = await MCPServerOAuthClientRepository(
row: Final[prisma_db_models.LiteLLM_MCPServerOAuthClient | None] = await _oauth_client_table_actions(
prisma_client
).table.find_unique(where={"server_id": server_id})
).find_unique(where={"server_id": server_id})
if row is None:
return None
return row.credentials
@ -1062,7 +1128,7 @@ async def upsert_mcp_server_oauth_client_credentials(
encrypted: Final = encrypt_credentials(credentials=MCPCredentials(**credentials), encryption_key=_get_salt_key())
blob: Final = safe_dumps(encrypted)
await MCPServerOAuthClientRepository(prisma_client).table.upsert(
await _oauth_client_table_actions(prisma_client).upsert(
where={"server_id": server_id},
data={
"create": {"server_id": server_id, "credentials": blob},
@ -1109,21 +1175,21 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient,
continue
update_data["updated_by"] = touched_by
await MCPServerRepository(prisma_client).table.update(
await _mcp_server_table_actions(prisma_client).update(
where={"server_id": mcp_server.server_id},
data=update_data,
)
updated += 1
oauth_clients: Final[list[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = await MCPServerOAuthClientRepository(
oauth_clients: Final[list[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = await _oauth_client_table_actions(
prisma_client
).table.find_many()
).find_many()
oauth_updated = 0
for oauth_client in oauth_clients:
rotated_credentials = _reencrypt_mcp_credentials_blob(oauth_client.credentials, new_master_key)
if rotated_credentials is None:
continue
await MCPServerOAuthClientRepository(prisma_client).table.update(
await _oauth_client_table_actions(prisma_client).update(
where={"server_id": oauth_client.server_id},
data={"credentials": rotated_credentials},
)
@ -1813,7 +1879,9 @@ async def get_mcp_submissions(
along with a summary count breakdown by approval_status.
Mirrors get_guardrail_submissions() from guardrail_endpoints.py.
"""
rows: list[prisma_db_models.LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_many(
rows: Final[list[prisma_db_models.LiteLLM_MCPServerTable]] = await _mcp_server_table_actions(
prisma_client
).find_many(
where={"submitted_at": {"not": None}},
order={"submitted_at": "desc"},
take=500, # safety cap; paginate if needed in a future iteration
@ -1915,7 +1983,7 @@ async def merge_user_env_vars(
"big",
signed=True,
)
async with prisma_client.db.tx() as tx:
async with _db_transaction_manager(prisma_client) as tx:
await tx.execute_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key)
row: Final[prisma_db_models.LiteLLM_MCPUserEnvVars | None] = await tx.litellm_mcpuserenvvars.find_unique(
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}

48
litellm/proxy/db/db_spend_update_writer.py
View file

@ -13,7 +13,7 @@ import random
import time
import traceback
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload
import litellm
from litellm._logging import verbose_proxy_logger
@ -64,6 +64,7 @@ from litellm.proxy.spend_tracking.savings import (
extract_cache_read_tokens,
)
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
from litellm.repositories.prisma_protocols import BatchTable
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -72,6 +73,37 @@ else:
ProxyLogging = Any
class _SpendBatch(Protocol):
litellm_usertable: BatchTable
litellm_verificationtoken: BatchTable
litellm_teamtable: BatchTable
litellm_teammembership: BatchTable
litellm_organizationtable: BatchTable
litellm_tagtable: BatchTable
litellm_agentstable: BatchTable
class _SpendBatchManager(Protocol):
async def __aenter__(self) -> _SpendBatch: ...
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
class _SpendTransaction(Protocol):
def batch_(self) -> _SpendBatchManager: ...
class _SpendTransactionManager(Protocol):
async def __aenter__(self) -> _SpendTransaction: ...
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager:
tx: Final[_SpendTransactionManager] = prisma_client.db.tx(timeout=timedelta(seconds=60))
return tx
def _get_llm_router():
"""The proxy's router, or None outside a running proxy.
@ -1195,7 +1227,7 @@ class DBSpendUpdateWriter:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction:
async with _spend_update_tx(prisma_client) as transaction:
async with transaction.batch_() as batcher:
# Sort by ID for consistent lock ordering across pods to prevent deadlocks.
# batch_() issues statements sequentially within the tx, so iteration
@ -1237,7 +1269,7 @@ class DBSpendUpdateWriter:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction:
async with _spend_update_tx(prisma_client) as transaction:
async with transaction.batch_() as batcher:
# Sort by token for consistent lock ordering across pods to prevent deadlocks.
for token, response_cost in sorted(key_list_transactions.items()):
@ -1270,7 +1302,7 @@ class DBSpendUpdateWriter:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction:
async with _spend_update_tx(prisma_client) as transaction:
async with transaction.batch_() as batcher:
# Sort by team_id for consistent lock ordering across pods to prevent deadlocks.
for team_id, response_cost in sorted(team_list_transactions.items()):
@ -1311,7 +1343,7 @@ class DBSpendUpdateWriter:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction:
async with _spend_update_tx(prisma_client) as transaction:
async with transaction.batch_() as batcher:
# Sort by composite key for consistent lock ordering across pods to prevent deadlocks.
# Key format "team_id::<v>::user_id::<v>" makes the string sort equivalent to sorting by (team_id, user_id).
@ -1362,7 +1394,7 @@ class DBSpendUpdateWriter:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction:
async with _spend_update_tx(prisma_client) as transaction:
async with transaction.batch_() as batcher:
# Sort by org_id for consistent lock ordering across pods to prevent deadlocks.
for org_id, response_cost in sorted(org_list_transactions.items()):
@ -1420,7 +1452,7 @@ class DBSpendUpdateWriter:
async def _update_entity_spend_in_db(
entity_name: str,
transactions: dict[str, float] | None,
table_accessor: Any,
table_accessor: Literal["litellm_tagtable", "litellm_agentstable"],
where_field: str,
n_retry_times: int,
prisma_client: PrismaClient,
@ -1445,7 +1477,7 @@ class DBSpendUpdateWriter:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction:
async with _spend_update_tx(prisma_client) as transaction:
async with transaction.batch_() as batcher:
# Sort by entity_id for consistent lock ordering across pods to prevent deadlocks.
for entity_id, response_cost in sorted(transactions.items()):

68
litellm/proxy/db/tool_registry_writer.py
View file

@ -6,8 +6,9 @@ Admins use the management endpoints to read and update input_policy / output_pol
"""
import uuid
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ToolDiscoveryQueueItem
@ -20,8 +21,41 @@ from litellm.types.tool_management import (
)
if TYPE_CHECKING:
from prisma import models as prisma_db_models
from litellm.proxy.utils import PrismaClient
_RowT_co: Final = TypeVar("_RowT_co", covariant=True)
class _TableActions(Protocol[_RowT_co]):
async def find_unique(self, where: Mapping[str, object]) -> _RowT_co | None: ...
async def find_many(
self,
where: Mapping[str, object] | None = None,
order: Mapping[str, object] | None = None,
include: Mapping[str, object] | None = None,
) -> Sequence[_RowT_co]: ...
async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT_co: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT_co | None: ...
def _tool_table_actions(prisma_client: "PrismaClient") -> "_TableActions[prisma_db_models.LiteLLM_ToolTable]":
table: Final[_TableActions[prisma_db_models.LiteLLM_ToolTable]] = ToolRepository(prisma_client).table
return table
def _object_permission_table_actions(
prisma_client: "PrismaClient",
) -> "_TableActions[prisma_db_models.LiteLLM_ObjectPermissionTable]":
table: Final[_TableActions[prisma_db_models.LiteLLM_ObjectPermissionTable]] = ObjectPermissionRepository(
prisma_client
).table
return table
def _row_to_model(row: dict | Any) -> LiteLLM_ToolTableRow:
"""Convert a Prisma model instance or dict to LiteLLM_ToolTableRow."""
@ -87,7 +121,7 @@ async def batch_upsert_tools(
if not data:
return
now: Final = datetime.now(timezone.utc)
table: Final = ToolRepository(prisma_client).table
table: Final = _tool_table_actions(prisma_client)
for item in data:
tool_name = item.get("tool_name", "")
origin = item.get("origin") or "user_defined"
@ -132,8 +166,8 @@ async def list_tools(
) -> list[LiteLLM_ToolTableRow]:
"""Return all tools, optionally filtered by input_policy."""
try:
where: Final = {"input_policy": input_policy} if input_policy is not None else {}
rows: Final = await ToolRepository(prisma_client).table.find_many(
where: Final[Mapping[str, str]] = {"input_policy": input_policy} if input_policy is not None else {}
rows: Final = await _tool_table_actions(prisma_client).find_many(
where=where,
order={"created_at": "desc"},
)
@ -149,7 +183,7 @@ async def get_tool(
) -> LiteLLM_ToolTableRow | None:
"""Return a single tool row by tool_name."""
try:
row: Final = await ToolRepository(prisma_client).table.find_unique(
row: Final = await _tool_table_actions(prisma_client).find_unique(
where={"tool_name": tool_name},
)
if row is None:
@ -172,7 +206,7 @@ async def update_tool_policy(
_updated_by: Final = updated_by or "system"
now: Final = datetime.now(timezone.utc)
create_data: Final[dict] = {
create_data: Final[dict[str, object]] = {
"tool_id": str(uuid.uuid4()),
"tool_name": tool_name,
"input_policy": input_policy or "untrusted",
@ -182,7 +216,7 @@ async def update_tool_policy(
"created_at": now,
"updated_at": now,
}
update_data: Final[dict] = {
update_data: Final[dict[str, object]] = {
"updated_by": _updated_by,
"updated_at": now,
}
@ -191,7 +225,7 @@ async def update_tool_policy(
if output_policy is not None:
update_data["output_policy"] = output_policy
await ToolRepository(prisma_client).table.upsert(
await _tool_table_actions(prisma_client).upsert(
where={"tool_name": tool_name},
data={
"create": create_data,
@ -214,7 +248,7 @@ async def get_tools_by_names(
if not tool_names:
return {}
try:
rows: Final = await ToolRepository(prisma_client).table.find_many(
rows: Final = await _tool_table_actions(prisma_client).find_many(
where={"tool_name": {"in": tool_names}},
)
return {
@ -239,7 +273,7 @@ async def list_overrides_for_tool(
"""
out: Final[list[ToolPolicyOverrideRow]] = []
try:
perms: Final = await ObjectPermissionRepository(prisma_client).table.find_many(
perms: Final = await _object_permission_table_actions(prisma_client).find_many(
where={"blocked_tools": {"has": tool_name}},
include={
"verification_tokens": True,
@ -302,7 +336,7 @@ class ToolPolicyRegistry:
try:
tools: Final = await call_with_db_reconnect_retry(
prisma_client,
lambda: ToolRepository(prisma_client).table.find_many(),
lambda: _tool_table_actions(prisma_client).find_many(),
reason="sync_tool_policy_from_db_tools_lookup_failure",
)
self._tool_input_policies = {
@ -314,7 +348,7 @@ class ToolPolicyRegistry:
perms: Final = await call_with_db_reconnect_retry(
prisma_client,
lambda: ObjectPermissionRepository(prisma_client).table.find_many(),
lambda: _object_permission_table_actions(prisma_client).find_many(),
reason="sync_tool_policy_from_db_perms_lookup_failure",
)
self._blocked_tools_by_op_id = {}
@ -352,7 +386,7 @@ class ToolPolicyRegistry:
"""
if not tool_names:
return {}
blocked: Final[set] = set()
blocked: Final[set[str]] = set()
for op_id in (object_permission_id, team_object_permission_id):
if op_id and op_id.strip():
blocked.update(self._blocked_tools_by_op_id.get(op_id.strip(), []))
@ -385,7 +419,7 @@ async def add_tool_to_object_permission_blocked(
if not object_permission_id or not tool_name:
return False
try:
row: Final = await ObjectPermissionRepository(prisma_client).table.find_unique(
row: Final = await _object_permission_table_actions(prisma_client).find_unique(
where={"object_permission_id": object_permission_id},
)
if row is None:
@ -394,7 +428,7 @@ async def add_tool_to_object_permission_blocked(
if tool_name in current:
return True
current.append(tool_name)
await ObjectPermissionRepository(prisma_client).table.update(
await _object_permission_table_actions(prisma_client).update(
where={"object_permission_id": object_permission_id},
data={"blocked_tools": current},
)
@ -413,7 +447,7 @@ async def remove_tool_from_object_permission_blocked(
if not object_permission_id or not tool_name:
return False
try:
row: Final = await ObjectPermissionRepository(prisma_client).table.find_unique(
row: Final = await _object_permission_table_actions(prisma_client).find_unique(
where={"object_permission_id": object_permission_id},
)
if row is None:
@ -422,7 +456,7 @@ async def remove_tool_from_object_permission_blocked(
if tool_name not in current:
return False
current = [t for t in current if t != tool_name]
await ObjectPermissionRepository(prisma_client).table.update(
await _object_permission_table_actions(prisma_client).update(
where={"object_permission_id": object_permission_id},
data={"blocked_tools": current},
)

42
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py
View file

@ -53,7 +53,7 @@ class LassoResponse(TypedDict):
violations_detected: bool
deputies: dict[str, bool]
findings: dict[str, list[dict[str, Any]]]
findings: dict[str, list[dict[str, object]]]
messages: list[dict[str, str]] | None
@ -120,7 +120,7 @@ class LassoGuardrail(CustomGuardrail):
super().__init__(**kwargs)
@staticmethod
def _get_field(obj: Any, field: str, default: Any = None) -> Any:
def _get_field(obj: Any, field: str, default: object = None) -> Any:
"""Get a field from either a dict or a Pydantic object."""
if isinstance(obj, dict):
return obj.get(field, default)
@ -129,7 +129,7 @@ class LassoGuardrail(CustomGuardrail):
@staticmethod
def _extract_tool_call_fields(
call: Any,
) -> tuple[str | None, str | None, dict[str, Any] | None]:
) -> tuple[str | None, str | None, dict[str, object] | None]:
"""Extract (call_id, name, parsed_input) from a tool call.
Handles both dict-style and Pydantic object-style tool_calls.
@ -142,7 +142,7 @@ class LassoGuardrail(CustomGuardrail):
return call_id, None, None
name: Final = get(func, "name")
args_str: Final = get(func, "arguments")
input_data: dict[str, Any] | None = None
input_data: dict[str, object] | None = None
if args_str:
try:
parsed = json.loads(args_str)
@ -248,7 +248,7 @@ class LassoGuardrail(CustomGuardrail):
# Extract messages from the response for validation
if isinstance(response, litellm.ModelResponse):
response_messages: Final[list[dict[str, Any]]] = []
response_messages: Final[list[dict[str, object]]] = []
for choice in response.choices:
if not hasattr(choice, "message"):
continue
@ -392,7 +392,7 @@ class LassoGuardrail(CustomGuardrail):
LassoGuardrailAPIError: If the Lasso API call fails
HTTPException: If blocking violations are detected
"""
raw_messages: Final[list[dict[str, Any]]] = data.get("messages") or []
raw_messages: Final[list[dict[str, object]]] = data.get("messages") or []
messages: list[dict[str, Any]] = self._expand_messages_for_classification(raw_messages) if raw_messages else []
messages_count: Final = len(messages)
if data.get("input") is not None:
@ -417,7 +417,7 @@ class LassoGuardrail(CustomGuardrail):
data: dict,
cache: DualCache,
message_type: Literal["PROMPT", "COMPLETION"],
messages: list[dict[str, Any]],
messages: list[dict[str, object]],
) -> dict:
"""Handle classification without masking."""
try:
@ -435,7 +435,7 @@ class LassoGuardrail(CustomGuardrail):
data: dict,
cache: DualCache,
message_type: Literal["PROMPT", "COMPLETION"],
messages: list[dict[str, Any]],
messages: list[dict[str, object]],
messages_count: int,
) -> dict:
"""Handle masking with classifix endpoint.
@ -477,7 +477,7 @@ class LassoGuardrail(CustomGuardrail):
self,
original_messages: list[dict[str, Any]],
masked_messages: list[dict[str, Any]],
) -> list[dict[str, Any]]:
) -> list[dict[str, object]]:
"""Map Lasso-format masked messages back onto the original OpenAI-format messages.
Lasso receives expanded messages (tool_use / tool_result blocks) and returns them
@ -487,7 +487,7 @@ class LassoGuardrail(CustomGuardrail):
while preserving the original structure.
"""
# Index masked content by type so we can look up by id without caring about order.
masked_tool_use: Final[dict[str, dict[str, Any]]] = {}
masked_tool_use: Final[dict[str, dict[str, object]]] = {}
masked_tool_result: Final[dict[str, str]] = {}
masked_text: Final[list[str]] = []
@ -524,7 +524,7 @@ class LassoGuardrail(CustomGuardrail):
},
)
result: Final[list[dict[str, Any]]] = []
result: Final[list[dict[str, object]]] = []
text_cursor = 0
for orig_msg in original_messages:
@ -563,9 +563,9 @@ class LassoGuardrail(CustomGuardrail):
def _update_tool_calls_from_masked(
self,
tool_calls: list[Any],
masked_tool_use: dict[str, dict[str, Any]],
) -> list[Any]:
tool_calls: list[object],
masked_tool_use: dict[str, dict[str, object]],
) -> list[object]:
"""Replace tool_call arguments with masked values returned by Lasso."""
updated: Final = []
for call in tool_calls:
@ -745,11 +745,11 @@ class LassoGuardrail(CustomGuardrail):
def _prepare_payload(
self,
messages: list[dict[str, Any]],
messages: list[dict[str, object]],
data: dict,
cache: DualCache,
message_type: Literal["PROMPT", "COMPLETION"] = "PROMPT",
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Prepare the payload for the Lasso API request.
@ -759,7 +759,7 @@ class LassoGuardrail(CustomGuardrail):
data: Request data (used for conversation_id generation and tools extraction)
cache: Cache instance for storing conversation_id (optional for post-call)
"""
payload: Final[dict[str, Any]] = {
payload: Final[dict[str, object]] = {
"messages": messages,
"messageType": message_type,
# Drives the "Used By" badge on Lasso Application API Keys: every call from this
@ -776,7 +776,7 @@ class LassoGuardrail(CustomGuardrail):
payload["sessionId"] = conversation_id
# Map OpenAI ChatCompletionToolParam array → ToolDefinition array
tools_data: Final[list[dict[str, Any]]] = data.get("tools") or []
tools_data: Final[list[dict[str, object]]] = data.get("tools") or []
if tools_data:
get: Final = self._get_field
tool_definitions: Final = []
@ -787,7 +787,7 @@ class LassoGuardrail(CustomGuardrail):
name = get(func, "name")
if not name:
continue
td: dict[str, Any] = {"name": name}
td: dict[str, object] = {"name": name}
description = get(func, "description")
if description:
td["description"] = description
@ -803,7 +803,7 @@ class LassoGuardrail(CustomGuardrail):
async def _call_lasso_api(
self,
headers: dict[str, str],
payload: dict[str, Any],
payload: dict[str, object],
api_url: str | None = None,
) -> LassoResponse:
"""Call the Lasso API and return the response."""
@ -921,7 +921,7 @@ class LassoGuardrail(CustomGuardrail):
) -> None:
"""Apply masking to the actual model response when mask=True and masked content is available."""
# Index masked tool_use blocks by id for O(1) lookup.
masked_tool_use: Final[dict[str, dict[str, Any]]] = {}
masked_tool_use: Final[dict[str, dict[str, object]]] = {}
masked_text: Final[list[str]] = []
for masked_msg in masked_messages:
content = masked_msg.get("content")

50
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

@ -14,9 +14,10 @@ import threading
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, cast
import aiohttp
from typing_extensions import NotRequired, ReadOnly
import litellm
from litellm import get_secret
@ -53,9 +54,18 @@ from litellm.utils import (
)
class _PresidioAnonymizeItem(TypedDict, total=False):
entity_type: ReadOnly[str | None]
class _PresidioAnonymizeResponse(TypedDict):
text: ReadOnly[str]
items: ReadOnly[NotRequired[list[_PresidioAnonymizeItem]]]
class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
user_api_key_cache = None
ad_hoc_recognizers = None
ad_hoc_recognizers: list[str] | None = None
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@ -72,7 +82,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
def __init__(
self,
mock_testing: bool = False,
mock_redacted_text: dict | None = None,
mock_redacted_text: _PresidioAnonymizeResponse | None = None,
presidio_analyzer_api_base: str | None = None,
presidio_anonymizer_api_base: str | None = None,
output_parse_pii: bool | None = False,
@ -91,7 +101,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
super().__init__(**kwargs)
self.guardrail_provider = "presidio"
self.pii_tokens: dict = {} # mapping of PII token to original text - only used with Presidio `replace` operation
self.pii_tokens: dict[
str, str
] = {} # mapping of PII token to original text - only used with Presidio `replace` operation
self.mock_redacted_text = mock_redacted_text
self.output_parse_pii = output_parse_pii or False
self.apply_to_output = apply_to_output
@ -265,7 +277,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
text: str,
presidio_config: PresidioPerRequestConfig | None,
request_data: dict,
) -> list[PresidioAnalyzeResponseItem] | dict:
) -> list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse:
"""
Send text to the Presidio analyzer endpoint and get analysis results
"""
@ -385,7 +397,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
# contain API keys or other secrets) in error responses.
raise Exception(f"Presidio PII analysis failed: {type(e).__name__}") from e
async def _post_presidio_anonymize(self, text: str, analyze_results: Any) -> Any:
async def _post_presidio_anonymize(
self,
text: str,
analyze_results: list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse,
) -> _PresidioAnonymizeResponse | None:
"""POST to Presidio anonymize; returns parsed JSON body."""
# Use shared session to prevent memory leak (issue #14540)
async with self._get_session_iterator() as session:
@ -417,7 +433,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
def _finalize_presidio_anonymize_simple(
self,
redacted_text: dict[str, Any],
redacted_text: _PresidioAnonymizeResponse,
masked_entity_count: dict[str, int],
) -> str:
# No need to build numbered tokens — just use Presidio's
@ -483,7 +499,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
async def anonymize_text(
self,
text: str,
analyze_results: Any,
analyze_results: list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse,
output_parse_pii: bool,
masked_entity_count: dict[str, int],
request_data: dict | None = None,
@ -517,8 +533,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
raise Exception(f"Presidio PII anonymization failed: {type(e).__name__}") from e
def filter_analyze_results_by_score(
self, analyze_results: list[PresidioAnalyzeResponseItem] | dict
) -> list[PresidioAnalyzeResponseItem] | dict:
self, analyze_results: list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse
) -> list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse:
"""
Drop detections that fall below configured per-entity score thresholds
or match an entity type in the deny list.
@ -556,7 +572,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
return filtered_results
def raise_exception_if_blocked_entities_detected(self, analyze_results: list[PresidioAnalyzeResponseItem] | dict):
def raise_exception_if_blocked_entities_detected(
self, analyze_results: list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse
):
"""
Raise an exception if blocked entities are detected
"""
@ -590,7 +608,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
Calls Presidio Analyze + Anonymize endpoints for PII Analysis + Masking
"""
start_time: Final = datetime.now()
analyze_results: list[PresidioAnalyzeResponseItem] | dict | None = None
analyze_results: list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse | None = None
status: GuardrailStatus = "success"
masked_entity_count: Final[dict[str, int]] = {}
exception_str: str = ""
@ -895,7 +913,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
return text
@staticmethod
def _is_anthropic_message_response(response: Any) -> bool:
def _is_anthropic_message_response(
response: ModelResponse | EmbeddingResponse | ImageResponse | dict[str, object],
) -> bool:
"""Check if the response is an Anthropic native message dict."""
return (
isinstance(response, dict)
@ -1283,8 +1303,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
@staticmethod
def _preserve_usage_from_last_chunk(
assembled_model_response: Any,
chunks: list[Any],
assembled_model_response: ModelResponse,
chunks: list[ModelResponseStream],
) -> None:
"""Copy usage metadata from the last chunk when stream_chunk_builder misses it."""
if not getattr(assembled_model_response, "usage", None) and chunks:

111
litellm/proxy/guardrails/usage_endpoints.py
View file

@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, overload
from fastapi import APIRouter, Depends, Query
from pydantic import BaseModel
from typing_extensions import NotRequired, TypedDict
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -26,7 +26,12 @@ from litellm.repositories.table_repositories import (
if TYPE_CHECKING:
from prisma import models as prisma_models
from prisma import types as prisma_types
from prisma.actions import LiteLLM_GuardrailsTableActions, LiteLLM_PolicyTableActions
from prisma.actions import (
LiteLLM_DailyGuardrailMetricsActions,
LiteLLM_DailyPolicyMetricsActions,
LiteLLM_GuardrailsTableActions,
LiteLLM_PolicyTableActions,
)
from litellm.proxy.utils import PrismaClient
from litellm.types.guardrails import Guardrail
@ -55,9 +60,51 @@ def _policies_table(
return policies_table
def _daily_guardrail_metrics_table(
prisma_client: "PrismaClient",
) -> "LiteLLM_DailyGuardrailMetricsActions[prisma_models.LiteLLM_DailyGuardrailMetrics]":
metrics_table: Final[LiteLLM_DailyGuardrailMetricsActions[prisma_models.LiteLLM_DailyGuardrailMetrics]] = (
DailyGuardrailMetricsRepository(prisma_client).table
)
return metrics_table
def _daily_policy_metrics_table(
prisma_client: "PrismaClient",
) -> "LiteLLM_DailyPolicyMetricsActions[prisma_models.LiteLLM_DailyPolicyMetrics]":
metrics_table: Final[LiteLLM_DailyPolicyMetricsActions[prisma_models.LiteLLM_DailyPolicyMetrics]] = (
DailyPolicyMetricsRepository(prisma_client).table
)
return metrics_table
async def _find_daily_guardrail_metrics(
prisma_client: "PrismaClient",
where: "prisma_types.LiteLLM_DailyGuardrailMetricsWhereInput",
) -> "Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics]":
return await _daily_guardrail_metrics_table(prisma_client).find_many(where=where)
async def _find_daily_policy_metrics(
prisma_client: "PrismaClient",
where: "prisma_types.LiteLLM_DailyPolicyMetricsWhereInput",
) -> "Sequence[prisma_models.LiteLLM_DailyPolicyMetrics]":
return await _daily_policy_metrics_table(prisma_client).find_many(where=where)
# --- Response models ---
class _GuardrailRunInfo(TypedDict, total=False):
guardrail_id: ReadOnly[str | None]
guardrail_name: ReadOnly[str | None]
guardrail_status: ReadOnly[str | None]
duration: ReadOnly[float | None]
confidence_score: ReadOnly[float | None]
risk_score: ReadOnly[float | None]
guardrail_response: ReadOnly[str | Mapping[str, object] | Sequence[Mapping[str, object]] | None]
class UsageChartPoint(TypedDict):
date: str
passed: int
@ -356,15 +403,15 @@ async def guardrails_usage_overview(
guardrails: Final[Sequence[_DbOrConfigGuardrail]] = [*db_guardrails, *config_guardrails]
# Daily metrics in range
metrics: Final[Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics]] = await DailyGuardrailMetricsRepository(
prisma_client
).table.find_many(where={"date": {"gte": start, "lte": end}})
metrics: Final[Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics]] = await _find_daily_guardrail_metrics(
prisma_client, where={"date": {"gte": start, "lte": end}}
)
# Previous period for trend
start_prev: Final = (datetime.strptime(start, "%Y-%m-%d") - timedelta(days=7)).strftime("%Y-%m-%d")
metrics_prev: Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics] = await DailyGuardrailMetricsRepository(
prisma_client
).table.find_many(where={"date": {"gte": start_prev, "lt": start}})
metrics_prev: Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics] = await _find_daily_guardrail_metrics(
prisma_client, where={"date": {"gte": start_prev, "lt": start}}
)
agg: Final = _aggregate_daily_metrics(metrics, "guardrail_id")
prev_agg: Final = _prev_fail_rates(metrics_prev, "guardrail_id")
@ -424,21 +471,19 @@ async def guardrails_usage_detail(
logical_id: Final = _get_guardrail_field(guardrail, "guardrail_name")
metric_ids: Final = [i for i in (logical_id, guardrail_id) if i]
metrics: Final[Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics]] = await DailyGuardrailMetricsRepository(
prisma_client
).table.find_many(
metrics: Final[Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics]] = await _find_daily_guardrail_metrics(
prisma_client,
where={
"guardrail_id": {"in": metric_ids},
"date": {"gte": start, "lte": end},
}
},
)
metrics_prev: Final[Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics]] = await DailyGuardrailMetricsRepository(
prisma_client
).table.find_many(
metrics_prev: Final[Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics]] = await _find_daily_guardrail_metrics(
prisma_client,
where={
"guardrail_id": {"in": metric_ids},
"date": {"lt": start},
}
},
)
requests: Final = sum(int(m.requests_evaluated or 0) for m in metrics)
@ -510,7 +555,9 @@ def _build_usage_logs_where(
def _usage_log_entry_from_row(
r: "prisma_models.LiteLLM_SpendLogGuardrailIndex", sl: Any, action_filter: str | None
r: "prisma_models.LiteLLM_SpendLogGuardrailIndex",
sl: "prisma_models.LiteLLM_SpendLogs",
action_filter: str | None,
) -> UsageLogEntry | None:
meta = sl.metadata
if isinstance(meta, str):
@ -518,8 +565,8 @@ def _usage_log_entry_from_row(
meta = json.loads(meta)
except Exception:
meta = {}
guardrail_info_list: Final = (meta or {}).get("guardrail_information") or []
entry_for_guardrail = None
guardrail_info_list: Final[Sequence[_GuardrailRunInfo]] = (meta or {}).get("guardrail_information") or []
entry_for_guardrail: _GuardrailRunInfo | None = None
for gi in guardrail_info_list:
if (gi.get("guardrail_id") or gi.get("guardrail_name")) == r.guardrail_id:
entry_for_guardrail = gi
@ -567,13 +614,12 @@ def _snippet(text: Any, max_len: int = 200) -> str | None:
if isinstance(text, str):
s = text
elif isinstance(text, list):
parts: Final = []
for item in text:
if isinstance(item, dict) and "content" in item:
c = item["content"]
parts.append(c if isinstance(c, str) else str(c))
else:
parts.append(str(item))
parts: Final[Sequence[str]] = [
(c if isinstance(c := item["content"], str) else str(c))
if isinstance(item, dict) and "content" in item
else str(item)
for item in text
]
s = " ".join(parts)
else:
s = str(text)
@ -705,18 +751,17 @@ async def policies_usage_overview(
try:
policies: Final = await _policies_table(prisma_client).find_many()
metrics: Final[Sequence[prisma_models.LiteLLM_DailyPolicyMetrics]] = await DailyPolicyMetricsRepository(
prisma_client
).table.find_many(where={"date": {"gte": start, "lte": end}})
metrics_prev: Final[Sequence[prisma_models.LiteLLM_DailyPolicyMetrics]] = await DailyPolicyMetricsRepository(
prisma_client
).table.find_many(
metrics: Final[Sequence[prisma_models.LiteLLM_DailyPolicyMetrics]] = await _find_daily_policy_metrics(
prisma_client, where={"date": {"gte": start, "lte": end}}
)
metrics_prev: Final[Sequence[prisma_models.LiteLLM_DailyPolicyMetrics]] = await _find_daily_policy_metrics(
prisma_client,
where={
"date": {
"gte": (datetime.strptime(start, "%Y-%m-%d") - timedelta(days=7)).strftime("%Y-%m-%d"),
"lt": start,
}
}
},
)
agg: Final = _aggregate_daily_metrics(metrics, "policy_id")
prev_agg: Final = _prev_fail_rates(metrics_prev, "policy_id")

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

@ -1,12 +1,13 @@
import asyncio
import json
import os
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final, Protocol
from fastapi import APIRouter, Depends, Header, HTTPException
from pydantic import TypeAdapter
from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm._uuid import uuid
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
@ -38,13 +39,32 @@ from litellm.types.proxy.management_endpoints.config_overrides import (
HashicorpVaultConfig,
)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
router: Final = APIRouter()
class _ConfigOverrideRow(Protocol):
config_value: str | Mapping[str, object] | None
class _ConfigOverridesTableClient(Protocol):
async def find_unique(self, where: Mapping[str, str]) -> _ConfigOverrideRow | None: ...
async def upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]) -> object: ...
async def delete(self, where: Mapping[str, str]) -> object: ...
def _config_overrides_table(prisma_client: "PrismaClient") -> _ConfigOverridesTableClient:
return ConfigOverridesRepository(prisma_client).table
_AUDIT_REDACTED: Final = "***REDACTED***"
def _redact_config(config: Mapping[str, Any] | None) -> dict[str, Any]:
def _redact_config(config: Mapping[str, object] | None) -> dict[str, str]:
"""Strip values from a config snapshot before audit-log emission.
Hashicorp Vault config carries ``vault_token``, ``approle_secret_id``,
@ -68,8 +88,8 @@ def _log_audit_task_exception(task: "asyncio.Task[None]") -> None:
async def _emit_hashicorp_vault_audit_log(
*,
action: AUDIT_ACTIONS,
before_config: Mapping[str, Any] | None,
after_config: Mapping[str, Any] | None,
before_config: Mapping[str, object] | None,
after_config: Mapping[str, object] | None,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
) -> None:
@ -136,9 +156,9 @@ _sensitive_masker: Final = SensitiveDataMasker()
# --- Shared helpers ---
def _mask_sensitive_fields(data: dict[str, Any], sensitive_fields: set[str]) -> dict[str, Any]:
def _mask_sensitive_fields(data: Mapping[str, object], sensitive_fields: set[str]) -> dict[str, object]:
"""Mask sensitive fields for API responses. Non-sensitive fields are left as-is."""
masked: Final = {}
masked: Final[dict[str, object]] = {}
for key, value in data.items():
if value is not None and key in sensitive_fields and isinstance(value, str):
masked[key] = _sensitive_masker._mask_value(value)
@ -147,7 +167,7 @@ def _mask_sensitive_fields(data: dict[str, Any], sensitive_fields: set[str]) ->
return masked
def _get_current_env_values(env_var_mapping: dict[str, str]) -> dict[str, Any]:
def _get_current_env_values(env_var_mapping: dict[str, str]) -> dict[str, str | None]:
"""Read current env var values as fallback when no DB record exists."""
values: Final = {}
for field_name, env_var_name in env_var_mapping.items():
@ -156,7 +176,13 @@ def _get_current_env_values(env_var_mapping: dict[str, str]) -> dict[str, Any]:
return values
def _extract_field_type(field_info: dict[str, Any]) -> str:
class _JsonSchemaField(TypedDict, total=False):
type: ReadOnly[str]
anyOf: ReadOnly[Sequence["_JsonSchemaField"]]
description: ReadOnly[str]
def _extract_field_type(field_info: _JsonSchemaField) -> str:
"""Extract the non-null type from a Pydantic v2 JSON schema field."""
if "type" in field_info:
return field_info["type"]
@ -166,11 +192,12 @@ def _extract_field_type(field_info: dict[str, Any]) -> str:
return "string"
def _build_field_schema(model_class: type) -> dict[str, Any]:
def _build_field_schema(model_class: type[BaseModel]) -> dict[str, object]:
"""Build field_schema dict from a Pydantic model for UI rendering."""
schema: Final = TypeAdapter(model_class).json_schema(by_alias=True)
raw_properties: Final[Mapping[str, _JsonSchemaField]] = schema.get("properties", {})
properties: Final = {}
for field_name, field_info in schema.get("properties", {}).items():
for field_name, field_info in raw_properties.items():
properties[field_name] = {
"description": field_info.get("description", ""),
"type": _extract_field_type(field_info),
@ -181,14 +208,14 @@ def _build_field_schema(model_class: type) -> dict[str, Any]:
}
def _parse_config_value(raw: Any) -> dict[str, Any]:
def _parse_config_value(raw: str | Mapping[str, object]) -> dict[str, object]:
"""Parse a config_value from DB (may be JSON string or dict)."""
if isinstance(raw, str):
return safe_json_loads(raw, default={})
return dict(raw)
def _set_env_vars(config_data: dict[str, Any]) -> None:
def _set_env_vars(config_data: Mapping[str, object]) -> None:
"""Set HCP_VAULT_* env vars from config data. Unsets vars for missing/None/empty fields."""
for field_name, env_var_name in HASHICORP_ENV_VAR_MAPPING.items():
value = config_data.get(field_name)
@ -242,15 +269,15 @@ async def update_hashicorp_vault_config(
detail=CommonProxyErrors.db_not_connected_error.value,
)
config_data = config.model_dump(exclude_none=True)
config_data: dict[str, object] = config.model_dump(exclude_none=True)
# Merge ALL fields the user didn't send: try DB first, fall back to env vars.
# Omitted field = keep existing; empty string = clear/remove the field.
existing_record: Final = await ConfigOverridesRepository(prisma_client).table.find_unique(
existing_record: Final = await _config_overrides_table(prisma_client).find_unique(
where={"config_type": "hashicorp_vault"}
)
existing_decrypted: dict[str, Any] | None = None
env_values: dict[str, Any] = {}
existing_decrypted: dict[str, object] | None = None
env_values: dict[str, str | None] = {}
if existing_record is not None and existing_record.config_value is not None:
existing_data: Final = _parse_config_value(existing_record.config_value)
existing_decrypted = proxy_config._decrypt_db_variables(existing_data)
@ -307,7 +334,7 @@ async def update_hashicorp_vault_config(
# Only persist to DB after successful init
encrypted_data: Final = proxy_config._encrypt_env_variables(config_data)
config_value: Final = safe_dumps(encrypted_data)
await ConfigOverridesRepository(prisma_client).table.upsert(
await _config_overrides_table(prisma_client).upsert(
where={"config_type": "hashicorp_vault"},
data={
"create": {
@ -377,7 +404,7 @@ async def get_hashicorp_vault_config(
field_schema: Final = _build_field_schema(HashicorpVaultConfig)
# Try to load from DB
db_record: Final = await ConfigOverridesRepository(prisma_client).table.find_unique(
db_record: Final = await _config_overrides_table(prisma_client).find_unique(
where={"config_type": "hashicorp_vault"}
)
@ -385,7 +412,7 @@ async def get_hashicorp_vault_config(
config_data: Final = _parse_config_value(db_record.config_value)
# Decrypt then mask sensitive fields so plaintext secrets are never sent to the UI
decrypted_data: Final = proxy_config._decrypt_db_variables(config_data)
decrypted_data: Final[Mapping[str, object]] = proxy_config._decrypt_db_variables(config_data)
masked_data: Final = _mask_sensitive_fields(decrypted_data, HASHICORP_SENSITIVE_FIELDS)
return ConfigOverrideSettingsResponse(
@ -434,10 +461,10 @@ async def delete_hashicorp_vault_config(
# Capture the prior config before delete so the audit-log row can
# show *what* was removed (keys only — values get redacted).
existing_record: Final = await ConfigOverridesRepository(prisma_client).table.find_unique(
existing_record: Final = await _config_overrides_table(prisma_client).find_unique(
where={"config_type": "hashicorp_vault"}
)
before_config: dict[str, Any] | None = None
before_config: dict[str, object] | None = None
if existing_record is not None and existing_record.config_value is not None:
try:
before_config = proxy_config._decrypt_db_variables(_parse_config_value(existing_record.config_value))
@ -447,7 +474,7 @@ async def delete_hashicorp_vault_config(
# Delete DB record if it exists — ignore if not found
deleted = False
try:
await ConfigOverridesRepository(prisma_client).table.delete(where={"config_type": "hashicorp_vault"})
await _config_overrides_table(prisma_client).delete(where={"config_type": "hashicorp_vault"})
deleted = True
except RecordNotFoundError:
verbose_proxy_logger.debug("No existing Hashicorp Vault config record to delete")
@ -502,7 +529,7 @@ async def test_hashicorp_vault_connection(
# Step 1: Authenticate (exercises AppRole login, TLS cert login, or direct token)
try:
headers: Final = await asyncio.to_thread(client._get_request_headers)
headers: Final[dict[str, str]] = await asyncio.to_thread(client._get_request_headers)
except Exception as e:
raise HTTPException(
status_code=502,

51
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, Literal, cast
from typing import Any, Final, Literal, Protocol, cast
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
@ -88,6 +88,7 @@ if TYPE_CHECKING:
from prisma.actions import (
LiteLLM_InvitationLinkActions,
LiteLLM_OrganizationMembershipActions,
LiteLLM_OrganizationTableActions,
LiteLLM_TeamMembershipActions,
LiteLLM_TeamTableActions,
LiteLLM_UserTableActions,
@ -142,6 +143,15 @@ def _invitation_link_table(
return invitation_table
def _organization_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_OrganizationTableActions[prisma_models.LiteLLM_OrganizationTable]":
organization_table: Final[LiteLLM_OrganizationTableActions[prisma_models.LiteLLM_OrganizationTable]] = (
OrganizationRepository(prisma_client).table
)
return organization_table
def _team_membership_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]":
@ -234,7 +244,7 @@ async def _check_duplicate_user_field(
if case_insensitive:
where_clause[field_name]["mode"] = "insensitive"
existing_user: Final = await UserRepository(prisma_client).table.find_first(where=where_clause)
existing_user: Final[object] = await UserRepository(prisma_client).table.find_first(where=where_clause)
if existing_user is not None:
existing_value: Final = getattr(existing_user, field_name, value)
@ -737,11 +747,11 @@ async def _get_user_info_teams(
user_id: str | None,
user_info: Any | None,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[list[Any], list[Any] | None]:
) -> tuple[list[TeamListResponseObject], list[TeamListResponseObject] | None]:
"""Fetch and merge teams from membership + user.teams field."""
from litellm.proxy.management_endpoints.team_endpoints import list_team
team_list: list[Any] = []
team_list: list[TeamListResponseObject] = []
team_id_list: list[str] = []
teams_1: Final = await list_team(
@ -756,7 +766,7 @@ async def _get_user_info_teams(
team_list = teams_1
team_id_list = [team.team_id for team in teams_1]
teams_2: list[Any] | None = None
teams_2: list[TeamListResponseObject] | None = None
target_team_ids: Final = getattr(user_info, "teams", None)
if target_team_ids and isinstance(target_team_ids, list):
@ -766,7 +776,7 @@ async def _get_user_info_teams(
query_type="find_all",
)
elif user_api_key_dict.user_id is not None and user_id is None:
caller_user_info: Final = await prisma_client.get_data(user_id=user_api_key_dict.user_id)
caller_user_info: Final[object] = await prisma_client.get_data(user_id=user_api_key_dict.user_id)
caller_team_ids: Final = getattr(caller_user_info, "teams", None)
if caller_team_ids:
teams_2 = await prisma_client.get_data(
@ -805,8 +815,8 @@ def _build_user_info_response(
user_id: str | None,
user_info: Any | None,
keys: list[LiteLLM_VerificationToken] | None,
team_list: list[Any],
teams_1: list[Any] | None,
team_list: list[TeamListResponseObject],
teams_1: list[TeamListResponseObject] | None,
) -> UserInfoResponse:
"""Create UserInfoResponse while filtering sensitive fields."""
if user_info is None and keys is not None:
@ -1085,7 +1095,7 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
verbose_proxy_logger.debug("results_keys: %s", results)
_keys_in_db: Final[list] = results[0]["keys"] or []
_keys_in_db: Final[Sequence[dict[str, object]]] = results[0]["keys"] or []
# cast all keys to LiteLLM_VerificationToken
keys_in_db: Final = []
for key in _keys_in_db:
@ -1094,7 +1104,7 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
keys_in_db.append(LiteLLM_VerificationToken.model_validate(key))
# cast all teams to LiteLLM_TeamTable
_teams_in_db: list = results[0]["teams"] or []
_teams_in_db: list[LiteLLM_TeamTable] = results[0]["teams"] or []
_teams_in_db = [LiteLLM_TeamTable.model_validate(team) for team in _teams_in_db]
_teams_in_db.sort(key=lambda x: getattr(x, "team_alias", "") or "")
returned_keys: Final = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db)
@ -1885,7 +1895,7 @@ async def get_user_key_counts(
# Get count for each user_id individually
for user_id in user_ids:
count = await VerificationTokenRepository(prisma_client).table.count(
count = await _verification_token_table(prisma_client).count(
where={
"user_id": user_id,
"OR": [
@ -2166,6 +2176,13 @@ async def get_users(
}
class _DeleteTeamRow(Protocol):
team_id: str
members_with_roles: object
def model_dump(self) -> Mapping[str, object]: ...
@router.post(
"/user/delete",
tags=["Internal User management"],
@ -2308,7 +2325,9 @@ async def delete_user(
)
## CLEANUP MEMBERS_WITH_ROLES
fetch_all_teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_row.teams}})
fetch_all_teams: Sequence[_DeleteTeamRow] = await TeamRepository(prisma_client).table.find_many(
where={"team_id": {"in": user_row.teams}}
)
teams_to_update = []
for team in fetch_all_teams:
removed_team_members, new_team_members = _cleanup_members_with_roles(
@ -2363,7 +2382,7 @@ async def add_internal_user_to_organization(
user_id: str,
organization_id: str,
user_role: LitellmUserRoles,
):
) -> "prisma_models.LiteLLM_OrganizationMembership":
"""
Helper function to add an internal user to an organization
@ -2382,14 +2401,16 @@ async def add_internal_user_to_organization(
try:
# Check if organization_id exists
organization_row: Final = await OrganizationRepository(prisma_client).table.find_unique(
organization_row: Final = await _organization_table(prisma_client).find_unique(
where={"organization_id": organization_id}
)
if organization_row is None:
raise Exception(f"Organization not found, passed organization_id={organization_id}")
# Create a new organization membership entry
new_membership: Final = await OrganizationMembershipRepository(prisma_client).table.create(
new_membership: Final[prisma_models.LiteLLM_OrganizationMembership] = await OrganizationMembershipRepository(
prisma_client
).table.create(
data={
"user_id": user_id,
"organization_id": organization_id,

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

@ -998,7 +998,7 @@ async def _common_key_generation_helper(
)
new_budget: Final = prisma_client.jsonify_object(budget_row.json(exclude_none=True))
_budget: Final = await BudgetRepository(prisma_client).table.create(
_budget: Final[LiteLLM_BudgetTable] = await BudgetRepository(prisma_client).table.create(
data={
**new_budget,
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
@ -4755,7 +4755,9 @@ async def _execute_virtual_key_regeneration(
grace_period=data.grace_period if data else None,
)
updated_token: Final[Mapping[str, object] | None] = await VerificationTokenRepository(prisma_client).table.update(
updated_token: Final[LiteLLM_VerificationToken | None] = await _prisma_table(
VerificationTokenRepository(prisma_client)
).update(
where={"token": hashed_api_key},
data=with_settings_updated_at(jsonified_update_data),
)
@ -5307,7 +5309,9 @@ async def validate_key_list_check(
if key_hash:
try:
key_info: Final = await VerificationTokenRepository(prisma_client).table.find_unique(
key_info: Final[LiteLLM_VerificationToken] = await VerificationTokenRepository(
prisma_client
).table.find_unique(
where={"token": key_hash},
)
except Exception:

68
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -19,10 +19,10 @@ import functools
import importlib
import json
import os
from collections.abc import Iterable
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Final, Literal
from typing import TYPE_CHECKING, Final, Literal, Protocol
from fastapi import (
APIRouter,
@ -36,6 +36,7 @@ from fastapi import (
status,
)
from fastapi.responses import JSONResponse
from typing_extensions import ReadOnly, TypedDict
try:
from prisma.errors import RecordNotFoundError, UniqueViolationError
@ -77,7 +78,11 @@ TEMPORARY_MCP_SERVER_TTL_SECONDS: Final = 300
TEMPORARY_MCP_SERVER_REDIS_KEY_PREFIX: Final = "litellm:mcp:temporary_server"
def does_mcp_server_exist(mcp_server_records: Iterable[Any], mcp_server_id: str) -> bool:
class _HasServerId(Protocol):
server_id: str
def does_mcp_server_exist(mcp_server_records: Iterable[_HasServerId], mcp_server_id: str) -> bool:
"""
Check if the mcp server with the given id exists in the iterable of mcp servers.
@ -93,6 +98,8 @@ def does_mcp_server_exist(mcp_server_records: Iterable[Any], mcp_server_id: str)
DEFAULT_MCP_REGISTRY_VERSION: Final = "1.0.0"
if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.proxy.utils import PrismaClient
try:
@ -111,7 +118,7 @@ if MCP_AVAILABLE:
class _ToolNameValidationResult(BaseModel):
is_valid: bool = True
warnings: list = []
warnings: list[str] = []
def validate_tool_name(name: str) -> _ToolNameValidationResult:
return _ToolNameValidationResult()
@ -263,7 +270,7 @@ if MCP_AVAILABLE:
_VALID_MCP_REQUIRED_FIELDS: Final[frozenset] = frozenset(NewMCPServerRequest.model_fields)
def _validate_mcp_required_fields(payload: Any) -> None:
def _validate_mcp_required_fields(payload: NewMCPServerRequest) -> None:
"""Validate submission payload against admin-configured mcp_required_fields."""
from litellm.proxy.proxy_server import (
general_settings as proxy_general_settings,
@ -329,7 +336,18 @@ if MCP_AVAILABLE:
return server.server_name
return server.server_id
def _build_mcp_registry_entry_for_server(server: MCPServer, base_url: str) -> dict[str, Any]:
class _McpRegistryRemote(TypedDict):
type: ReadOnly[str]
url: ReadOnly[str]
class _McpRegistryEntry(TypedDict):
name: ReadOnly[str]
title: ReadOnly[str]
description: ReadOnly[str]
version: ReadOnly[str]
remotes: ReadOnly[Sequence[_McpRegistryRemote]]
def _build_mcp_registry_entry_for_server(server: MCPServer, base_url: str) -> _McpRegistryEntry:
server_name: Final = _build_mcp_registry_server_name(server)
title: Final = server_name
description: Final = server_name
@ -353,7 +371,7 @@ if MCP_AVAILABLE:
],
}
def _build_builtin_registry_entry(base_url: str) -> dict[str, Any]:
def _build_builtin_registry_entry(base_url: str) -> _McpRegistryEntry:
remote_url: Final = _build_registry_remote_url(base_url, "/mcp")
return {
"name": LITELLM_MCP_SERVER_NAME,
@ -400,7 +418,7 @@ if MCP_AVAILABLE:
if cache_backend is None or not hasattr(cache_backend, "async_set_cache"):
return
payload: Final[dict[str, Any]] = server.model_dump(mode="json")
payload: Final[dict[str, object]] = server.model_dump(mode="json")
payload_json: Final = json.dumps(payload)
try:
encrypted_payload: Final = encrypt_value_helper(payload_json)
@ -464,7 +482,7 @@ if MCP_AVAILABLE:
return None
if not isinstance(loaded, dict):
return None
payload_dict: Final[dict[str, Any]] = loaded
payload_dict: Final[dict[str, object]] = loaded
try:
return MCPServer.model_validate(payload_dict)
@ -725,7 +743,7 @@ if MCP_AVAILABLE:
one, so a form that round-trips it must not read as "credentials supplied"."""
if not credentials:
return False
as_dict: Final[dict[str, Any]] = dict(credentials)
as_dict: Final[dict[str, object]] = dict(credentials)
return any(value for key, value in as_dict.items() if key not in MCP_ADMIN_CONFIG_CREDENTIAL_KEYS)
def _inherit_credentials_from_existing_server(
@ -738,7 +756,7 @@ if MCP_AVAILABLE:
if existing_server is None:
return payload
inherited_credentials: dict[str, Any] = {
inherited_credentials: dict[str, object] = {
credential_key: value
for server_attr, credential_key in _INHERITED_CREDENTIAL_FIELDS
if (value := getattr(existing_server, server_attr, None))
@ -755,7 +773,7 @@ if MCP_AVAILABLE:
except AttributeError:
pass
payload_dict: dict[str, Any]
payload_dict: dict[str, object]
try:
payload_dict = payload.model_dump()
except AttributeError:
@ -888,7 +906,9 @@ if MCP_AVAILABLE:
# Get from DB
if prisma_client is not None:
try:
mcp_servers: Final = await MCPServerRepository(prisma_client).table.find_many()
mcp_servers: Final[Sequence[prisma_models.LiteLLM_MCPServerTable]] = await MCPServerRepository(
prisma_client
).table.find_many()
for server in mcp_servers:
if hasattr(server, "mcp_access_groups") and server.mcp_access_groups:
access_groups.update(server.mcp_access_groups)
@ -930,7 +950,7 @@ if MCP_AVAILABLE:
verbose_proxy_logger.debug("MCP registry request from IP=%s", client_ip)
base_url: Final = get_request_base_url(request)
registry_servers: Final[list[dict[str, Any]]] = []
registry_servers: Final[list[dict[str, _McpRegistryEntry]]] = []
registry_servers.append({"server": _build_builtin_registry_entry(base_url)})
# Centralized IP-based filtering: external callers only see public servers
@ -1126,7 +1146,9 @@ if MCP_AVAILABLE:
if user_id and _byok_prisma_client is not None:
byok_server_ids: Final = [s.server_id for s in redacted_mcp_servers if getattr(s, "is_byok", False)]
if byok_server_ids:
cred_rows: Final = await MCPUserCredentialsRepository(_byok_prisma_client).table.find_many(
cred_rows: Final[
Sequence[prisma_models.LiteLLM_MCPUserCredentials]
] = await MCPUserCredentialsRepository(_byok_prisma_client).table.find_many(
where={"user_id": user_id, "server_id": {"in": byok_server_ids}}
)
cred_set: Final = {r.server_id for r in cred_rows}
@ -1680,7 +1702,7 @@ if MCP_AVAILABLE:
options={"verify_exp": False, "verify_aud": False},
)
if decoded.get("login_method") in ("sso", "username_password"):
cookie_key: Final = decoded.get("key", "")
cookie_key: Final[str] = decoded.get("key", "")
if cookie_key:
api_key = f"Bearer {cookie_key}"
except _jwt.InvalidTokenError:
@ -1707,7 +1729,7 @@ if MCP_AVAILABLE:
get_request_route,
)
server_id: Final = request.path_params.get("server_id", "")
server_id: Final[str] = request.path_params.get("server_id", "")
if server_id:
_s = global_mcp_server_manager.get_mcp_server_by_id(server_id)
if not _s:
@ -2324,7 +2346,7 @@ if MCP_AVAILABLE:
required: Final[list[MCPUserEnvVarSpec]] = []
missing_count = 0
for spec in user_specs:
name = spec["name"]
name: str = spec["name"]
if name not in blocking:
continue
value = stored_values.get(name)
@ -2672,16 +2694,16 @@ if MCP_AVAILABLE:
"mcp_registry.json",
)
_mcp_registry_cache: dict[str, Any] | None = None
_mcp_registry_cache: Mapping[str, Sequence[Mapping[str, str]]] | None = None
def _load_mcp_registry() -> dict[str, Any]:
def _load_mcp_registry() -> Mapping[str, Sequence[Mapping[str, str]]]:
"""Load the curated MCP registry from disk. Cached after first read."""
global _mcp_registry_cache
if _mcp_registry_cache is not None:
return _mcp_registry_cache
try:
with open(_MCP_REGISTRY_PATH, "r") as f:
data: dict[str, Any] = json.load(f)
data: Mapping[str, Sequence[Mapping[str, str]]] = json.load(f)
except Exception as e:
verbose_proxy_logger.warning("Failed to load MCP registry from %s: %s", _MCP_REGISTRY_PATH, e)
data = {"servers": []}
@ -2747,9 +2769,9 @@ if MCP_AVAILABLE:
)
@functools.lru_cache(maxsize=1)
def _load_openapi_registry() -> dict[str, Any]:
def _load_openapi_registry() -> dict[str, object]:
with open(_OPENAPI_REGISTRY_PATH, "r") as f:
data: Final[dict[str, Any]] = json.load(f)
data: Final[dict[str, object]] = json.load(f)
return data
@router.get(

39
litellm/proxy/management_endpoints/model_access_group_management_endpoints.py
View file

@ -7,7 +7,7 @@ Endpoints here:
import json
from collections.abc import Mapping, Sequence
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final, Protocol
from fastapi import APIRouter, Depends, HTTPException
@ -33,10 +33,31 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import
UpdateModelGroupRequest,
)
if TYPE_CHECKING:
from litellm import Router
router: Final = APIRouter()
def validate_models_exist(model_names: list[str], llm_router) -> tuple[bool, list[str]]:
class _DeploymentRow(Protocol):
model_id: str
model_name: str
model_info: object
class _ModelTableClient(Protocol):
async def find_many(self, where: Mapping[str, object] | None = None) -> Sequence[_DeploymentRow]: ...
async def find_unique(self, where: Mapping[str, object]) -> _DeploymentRow | None: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
def _model_table(prisma_client: PrismaClient) -> _ModelTableClient:
return ModelRepository(prisma_client).table
def validate_models_exist(model_names: list[str], llm_router: "Router | None") -> tuple[bool, list[str]]:
"""
Validate that all requested model names exist in the router.
Checks only exact model name matches.
@ -117,7 +138,7 @@ async def _tag_deployment_with_access_group(
)
if not was_modified:
return None
await ModelRepository(prisma_client).table.update(
await _model_table(prisma_client).update(
where={"model_id": model_id},
data={"model_info": json.dumps(updated_model_info)},
)
@ -150,7 +171,7 @@ async def _strip_access_group_from_deployment(
)
if not was_modified:
return None
await ModelRepository(prisma_client).table.update(
await _model_table(prisma_client).update(
where={"model_id": model_id},
data={"model_info": json.dumps(updated_model_info)},
)
@ -174,7 +195,7 @@ async def update_deployments_with_access_group(
The (model_id, updated model_info) pair of every deployment actually written,
so callers can verify each one survived the post-write reload
"""
deployments: Final = await ModelRepository(prisma_client).table.find_many(where={"model_name": {"in": model_names}})
deployments: Final = await _model_table(prisma_client).find_many(where={"model_name": {"in": model_names}})
verbose_proxy_logger.debug("Found %s deployments for model_names: %s", len(deployments), model_names)
found_names: Final = {deployment.model_name for deployment in deployments}
@ -225,8 +246,8 @@ async def update_specific_deployments_with_access_group(
return tuple(pair for pair in tagged if pair is not None)
async def _find_deployment_or_400(model_id: str, prisma_client: PrismaClient) -> Mapping[str, object] | None:
deployment: Final = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id})
async def _find_deployment_or_400(model_id: str, prisma_client: PrismaClient) -> object:
deployment: Final = await _model_table(prisma_client).find_unique(where={"model_id": model_id})
if deployment is None:
raise HTTPException(
status_code=400,
@ -646,7 +667,7 @@ async def update_access_group(
try:
# Step 1: Remove access group from ALL DB deployments (skip config models)
all_deployments: Final = await ModelRepository(prisma_client).table.find_many()
all_deployments: Final = await _model_table(prisma_client).find_many()
stripped: Final = [
await _strip_access_group_from_deployment(
@ -764,7 +785,7 @@ async def delete_access_group(
try:
# Remove access group from all DB deployments (skip config models)
all_deployments: Final = await ModelRepository(prisma_client).table.find_many()
all_deployments: Final = await _model_table(prisma_client).find_many()
removed: Final = [
await _strip_access_group_from_deployment(

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

@ -479,7 +479,7 @@ def _is_safe_cli_sso_metadata_dest_key(dest_key: str) -> bool:
return not any(fragment in lowered for fragment in _CLI_SSO_SECRET_KEY_FRAGMENTS)
def _is_safe_cli_sso_scalar_claim_value(value: Any) -> bool:
def _is_safe_cli_sso_scalar_claim_value(value: object) -> bool:
if not isinstance(value, _CLI_SSO_SCALAR_TYPES):
return False
if isinstance(value, str):
@ -490,17 +490,17 @@ def _is_safe_cli_sso_scalar_claim_value(value: Any) -> bool:
return True
def _sso_result_to_dict(result: CustomOpenID | OpenID | dict) -> dict[str, Any]:
def _sso_result_to_dict(result: CustomOpenID | OpenID | dict[str, object]) -> dict[str, object]:
if isinstance(result, dict):
return result
if hasattr(result, "model_dump"):
dumped: Final = result.model_dump()
if isinstance(dumped, dict):
return cast(dict[str, Any], dumped)
return dumped
return {}
def _get_nested_claim_value(data: dict[str, Any], claim_path: str) -> Any:
def _get_nested_claim_value(data: Mapping[str, object], claim_path: str) -> object:
"""Resolve a dot-notation claim path against an SSO result dict.
Unlike ``get_nested_value``, this does not strip a leading ``metadata.``
@ -514,7 +514,7 @@ def _get_nested_claim_value(data: dict[str, Any], claim_path: str) -> Any:
placeholder: Final = "\x00"
parts = claim_path.replace("\\.", placeholder).split(".")
parts = [p.replace(placeholder, ".") for p in parts]
current: Any = data
current: object = data
for part in parts:
if isinstance(current, dict) and part in current:
current = current[part]
@ -523,7 +523,7 @@ def _get_nested_claim_value(data: dict[str, Any], claim_path: str) -> Any:
return current
def _extract_sso_claim_value(result: CustomOpenID | OpenID | dict, claim_path: str) -> Any:
def _extract_sso_claim_value(result: CustomOpenID | OpenID | dict[str, object], claim_path: str) -> object:
extra_fields: Final = getattr(result, "extra_fields", None)
if isinstance(extra_fields, dict):
if claim_path in extra_fields:
@ -539,7 +539,7 @@ def _extract_sso_claim_value(result: CustomOpenID | OpenID | dict, claim_path: s
return _get_nested_claim_value(result_dict, claim_path)
def _set_nested_metadata_value(metadata: dict[str, Any], key_path: str, value: Any) -> None:
def _set_nested_metadata_value(metadata: dict[str, object], key_path: str, value: object) -> None:
placeholder: Final = "\x00"
parts = key_path.replace("\\.", placeholder).split(".")
parts = [p.replace(placeholder, ".") for p in parts]
@ -554,24 +554,25 @@ def _set_nested_metadata_value(metadata: dict[str, Any], key_path: str, value: A
def _flatten_cli_sso_metadata_for_poll(
metadata: dict[str, Any],
metadata: Mapping[str, object],
) -> dict[str, str | int | float | bool]:
"""Expose scalar attribution metadata as a flat dict for CLI poll responses."""
flattened: Final[dict[str, str | int | float | bool]] = {}
stack: Final[list[tuple[str, Any]]] = [("", metadata)]
stack: Final[list[tuple[str, object]]] = [("", metadata)]
while stack:
prefix, value = stack.pop()
if isinstance(value, dict):
for key, nested in value.items():
nested_items: Mapping[str, object] = value
for key, nested in nested_items.items():
nested_prefix = f"{prefix}.{key}" if prefix else key
stack.append((nested_prefix, nested))
elif _is_safe_cli_sso_scalar_claim_value(value):
elif isinstance(value, (str, int, float, bool)) and _is_safe_cli_sso_scalar_claim_value(value):
flattened[prefix] = value
return flattened
def build_cli_sso_attribution_metadata(
result: CustomOpenID | OpenID | dict,
result: CustomOpenID | OpenID | dict[str, object],
) -> dict[str, object]:
"""
Build allowlisted, non-secret scalar attribution metadata from an SSO result.
@ -599,8 +600,8 @@ def build_cli_sso_attribution_metadata(
def _merge_cli_sso_attribution_metadata(
existing_metadata: dict[str, Any], attribution_metadata: dict[str, Any]
) -> dict[str, Any]:
existing_metadata: dict[str, object], attribution_metadata: dict[str, object]
) -> dict[str, object]:
"""Merge attribution metadata into existing user metadata in-place.
Preserves original value types (in particular, string claim values that
@ -608,7 +609,7 @@ def _merge_cli_sso_attribution_metadata(
are merged iteratively so attribution claims do not clobber unrelated keys
under the same parent.
"""
pending: Final[list[tuple[dict[str, Any], dict[str, Any]]]] = [(existing_metadata, attribution_metadata)]
pending: Final[list[tuple[dict[str, object], dict[str, object]]]] = [(existing_metadata, attribution_metadata)]
while pending:
target, source = pending.pop()
for key, value in source.items():
@ -656,7 +657,7 @@ async def _persist_cli_sso_user_metadata(
def _cli_poll_attribution_metadata_from_session(
session_data: dict[str, Any],
session_data: Mapping[str, object],
) -> dict[str, str | int | float | bool]:
stored: Final = session_data.get("attribution_metadata")
if isinstance(stored, dict):
@ -960,11 +961,12 @@ def process_sso_jwt_access_token(
# Try role_mappings first (group-based role determination)
if role_mappings is not None and role_mappings.roles:
group_claim: Final = role_mappings.group_claim
user_groups_raw: Final[Any] = get_nested_value(access_token_payload, group_claim)
user_groups_raw: Final[object] = get_nested_value(access_token_payload, group_claim)
user_groups: list[str] = []
if isinstance(user_groups_raw, list):
user_groups = [str(g) for g in user_groups_raw]
raw_groups: Final[Sequence[object]] = user_groups_raw
user_groups = [str(g) for g in raw_groups]
elif isinstance(user_groups_raw, str):
user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()]
elif user_groups_raw is not None:
@ -1214,12 +1216,13 @@ def generic_response_convertor(
]:
# Use role_mappings to determine role from groups
group_claim: Final = role_mappings.group_claim
user_groups_raw: Final[Any] = get_nested_value(response, group_claim)
user_groups_raw: Final[object] = get_nested_value(response, group_claim)
# Handle different formats: could be a list, string (comma-separated), or single value
user_groups: list[str] = []
if isinstance(user_groups_raw, list):
user_groups = [str(g) for g in user_groups_raw]
raw_groups: Final[Sequence[object]] = user_groups_raw
user_groups = [str(g) for g in raw_groups]
elif isinstance(user_groups_raw, str):
# Handle comma-separated string
user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()]
@ -3093,7 +3096,7 @@ class SSOAuthenticationHandler:
def _get_generic_sso_redirect_params(
state: str | None = None,
generic_authorization_endpoint: str | None = None,
) -> tuple[dict, str | None]:
) -> tuple[dict[str, str], str | None]:
"""
Get redirect parameters for Generic SSO with proper state priority handling.
Optionally generates PKCE parameters if GENERIC_CLIENT_USE_PKCE is enabled.

64
litellm/proxy/proxy_server.py
View file

@ -864,7 +864,7 @@ async def _flush_spend_logs_queue_on_shutdown() -> None:
verbose_proxy_logger.exception("Error flushing spend logs queue on shutdown: %s", e)
async def proxy_shutdown_event():
async def proxy_shutdown_event() -> None:
global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update
verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server")
if prisma_client:
@ -958,7 +958,7 @@ async def _initialize_shared_aiohttp_session():
@asynccontextmanager
async def proxy_startup_event(app: FastAPI):
async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
global \
prisma_client, \
master_key, \
@ -3732,11 +3732,11 @@ _DB_OVERLAY_REMOTE_MODULE_LIST_FIELDS: Final[dict[str, tuple[str, ...]]] = {
}
def _is_remote_module_url(value: Any) -> bool:
def _is_remote_module_url(value: object) -> bool:
return isinstance(value, str) and (value.startswith("s3://") or value.startswith("gcs://"))
def _scrub_guardrail_inner(inner: dict[str, Any]) -> None:
def _scrub_guardrail_inner(inner: dict[str, JsonValue]) -> None:
"""Strip remote-URL entries from a guardrail's ``callbacks`` list
and ``guardrail`` (v2 module-path) field. Mutates in place."""
cbs: Final = inner.get("callbacks")
@ -3756,7 +3756,7 @@ def _scrub_guardrail_inner(inner: dict[str, Any]) -> None:
inner["guardrail"] = None
def _scrub_db_overlay_remote_module_loads(section: str, db_value: Any) -> Any:
def _scrub_db_overlay_remote_module_loads(section: str, db_value: JsonValue) -> JsonValue:
"""Strip ``s3://`` / ``gcs://`` entries from the DB-overlay value for
fields whose contents reach ``get_instance_fn``. The same scheme is
allowed from a YAML config (the documented operator flow) but a
@ -4064,8 +4064,8 @@ class ProxyConfig:
def __init__(self) -> None:
self.config: dict[str, Any] = {}
self._last_semantic_filter_config: dict[str, Any] | None = None
self._last_hashicorp_vault_config: dict[str, Any] | None = None
self._last_semantic_filter_config: dict[str, object] | None = None
self._last_hashicorp_vault_config: dict[str, object] | None = None
self.worker_registry: list[WorkerRegistryEntry] = []
self.config_sync_subscriber: ConfigSyncSubscriber | None = None
self.auth_cache_invalidation_subscriber: AuthCacheInvalidationSubscriber | None = None
@ -5955,7 +5955,7 @@ class ProxyConfig:
)
@staticmethod
def _parse_router_settings_value(value: Any) -> dict | None:
def _parse_router_settings_value(value: object) -> dict | None:
"""
Parse a router_settings value that may be a dict or a JSON/YAML string.
@ -6499,7 +6499,7 @@ class ProxyConfig:
as "all models deleted" and must not evict existing router deployments.
"""
try:
new_models: Final = await ModelRepository(prisma_client).table.find_many()
new_models: Final[list[_ModelTableRow]] = await ModelRepository(prisma_client).table.find_many()
return new_models
except Exception as e:
verbose_proxy_logger.exception(
@ -7563,9 +7563,9 @@ def _get_client_requested_model_for_streaming(request_data: dict) -> str:
return requested_model if isinstance(requested_model, str) else ""
def _is_positive_int_like(value: Any) -> bool:
def _is_positive_int_like(value: str | float | None) -> bool:
try:
return int(value) > 0
return value is not None and int(value) > 0
except (TypeError, ValueError):
return False
@ -7832,7 +7832,7 @@ _STREAM_KEEPALIVE: Final = object()
_KEEPALIVE_MIN_SECONDS: Final = 1.0
_KEEPALIVE_MAX_SECONDS: Final = 300.0
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
_EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({})
async def _iter_with_keepalive(
@ -7887,7 +7887,7 @@ async def _iter_with_keepalive(
class _DeploymentKeepaliveConfig(NamedTuple):
keepalive_seconds: Any
keepalive_seconds: object
allow_client_override: bool
@ -7945,7 +7945,7 @@ def _is_explicit_keepalive_disable(raw: object) -> bool:
return False
def _resolve_keepalive_seconds(request_data: Mapping[str, Any], response: object = None) -> float:
def _resolve_keepalive_seconds(request_data: Mapping[str, object], response: object = None) -> float:
deployment_config: Final = _keepalive_from_deployment_config(request_data, response)
deployment_raw: Final = deployment_config.keepalive_seconds if deployment_config is not None else None
allow_client_override: Final = deployment_config.allow_client_override if deployment_config is not None else False
@ -7992,7 +7992,7 @@ def _resolve_keepalive_seconds(request_data: Mapping[str, Any], response: object
_KEEPALIVE_CACHE_TTL_SECONDS: Final = 5.0
def _make_keepalive_resolver(request_data: Mapping[str, Any]) -> Callable[[object], float]:
def _make_keepalive_resolver(request_data: Mapping[str, object]) -> Callable[[object], float]:
"""Wrap `_resolve_keepalive_seconds` with a memo keyed on the serving
deployment's model_id. The steady-state case (no mid-stream fallback, the
overwhelming majority of streams) sees the same model_id on every chunk, so
@ -9784,7 +9784,7 @@ async def model_info(
)
def _blocked_response_usage(original_response: Any | None) -> "litellm.Usage":
def _blocked_response_usage(original_response: object | None) -> "litellm.Usage":
"""
Token usage for a synthetic guardrail-blocked response.
@ -12303,7 +12303,7 @@ def _enrich_model_info_with_litellm_data(
async def _get_caller_byok_team_scope(
user_api_key_dict: UserAPIKeyAuth | None,
prisma_client: Any | None,
prisma_client: PrismaClient | None,
) -> set[str] | None:
"""
Return the team IDs whose BYOK rows the caller is allowed to see via
@ -12338,7 +12338,7 @@ async def _get_caller_byok_team_scope(
return key_team_scope | set(user_row.teams or [])
def _byok_row_outside_caller_teams(model_info_dict: dict[str, Any], allowed_team_ids: set[str] | None) -> bool:
def _byok_row_outside_caller_teams(model_info_dict: dict[str, JsonValue], allowed_team_ids: set[str] | None) -> bool:
"""Whether a team BYOK row belongs to a team the caller is not a member of.
`team_id` is only set on team BYOK rows; non-team rows fall through
@ -12360,15 +12360,15 @@ _SORTED_SEARCH_DB_FETCH_CAP: Final = 500
async def _fetch_db_models_for_search(
prisma_client: Any,
proxy_config: Any,
prisma_client: PrismaClient,
proxy_config: ProxyConfig,
search_lower: str,
db_model_ids_in_router: set[str],
router_models_count: int,
page: int,
size: int,
sort_by: str | None,
is_byok_outside_caller_teams: Callable[[dict[str, Any]], bool],
is_byok_outside_caller_teams: Callable[[dict[str, JsonValue]], bool],
) -> tuple[list[dict[str, Any]], int]:
"""
Run the bounded DB query that backs `/v2/model/info?search=`. Returns
@ -12414,7 +12414,7 @@ async def _fetch_db_models_for_search(
if not is_byok_outside_caller_teams(m.model_info if isinstance(m.model_info, dict) else {})
]
decrypted: Final[list[dict[str, Any]]] = []
decrypted: Final[list[dict[str, object]]] = []
for db_model in matching_db_rows:
decrypted_models = proxy_config.decrypt_model_list_from_db([db_model])
if decrypted_models:
@ -12426,8 +12426,8 @@ async def _fetch_db_models_for_search(
async def _apply_search_filter_to_models(
all_models: list[dict[str, Any]],
search: str,
prisma_client: Any | None,
proxy_config: Any,
prisma_client: PrismaClient | None,
proxy_config: ProxyConfig,
user_api_key_dict: UserAPIKeyAuth | None = None,
page: int = 1,
size: int = 50,
@ -12466,7 +12466,7 @@ async def _apply_search_filter_to_models(
prisma_client=prisma_client,
)
def _is_byok_outside_caller_teams(model_info_dict: dict[str, Any]) -> bool:
def _is_byok_outside_caller_teams(model_info_dict: dict[str, JsonValue]) -> bool:
return _byok_row_outside_caller_teams(model_info_dict, allowed_team_ids)
def _model_matches_search(m: dict[str, Any]) -> bool:
@ -12532,7 +12532,7 @@ async def _apply_search_filter_to_models(
return filtered_router_models + db_models, search_total_count
def _normalize_datetime_for_sorting(dt: Any) -> datetime | None:
def _normalize_datetime_for_sorting(dt: object) -> datetime | None:
"""
Normalize a datetime value to a timezone-aware UTC datetime for sorting.
@ -12685,7 +12685,7 @@ def _paginate_models_response(
size: int,
total_count: int | None,
search: str | None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Paginate models and return response dictionary.
@ -12724,7 +12724,7 @@ def _paginate_models_response(
}
def _team_models_resolve_to_names(team_models: list[str], access_groups: dict[str, Any]) -> list[str]:
def _team_models_resolve_to_names(team_models: list[str], access_groups: Mapping[str, Sequence[str]]) -> list[str]:
"""Expand team model entries (including access group names) to concrete model names."""
resolved: Final[list[str]] = []
for name in team_models:
@ -13600,7 +13600,7 @@ async def model_metrics_exceptions(
return {"data": response, "exception_types": list(exception_types)}
def _deployment_matches_allowed_model_names(model: dict[str, Any], allowed_model_names: set[str]) -> bool:
def _deployment_matches_allowed_model_names(model: dict[str, JsonValue], allowed_model_names: set[str]) -> bool:
"""Match a router deployment against allowed public model names.
Team-scoped rows store an internal routing key in ``model_name``; callers
@ -14860,7 +14860,7 @@ async def _rollback_onboarding_invite_claim(
verbose_proxy_logger.exception("Failed to roll back onboarding invitation after session key mint failed.")
async def _generate_onboarding_ui_session_token(user_obj: Any) -> str:
async def _generate_onboarding_ui_session_token(user_obj: _UserTableRow) -> str:
global master_key, general_settings
response: Final = await generate_key_helper_fn(
@ -15975,7 +15975,7 @@ def _general_settings_ui_litellm_default(
return False if spec["type"] == "Boolean" else None
def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> GeneralSettingsUILiteLLMValue:
def _validate_general_settings_ui_litellm_value(field_name: str, value: object) -> GeneralSettingsUILiteLLMValue:
spec: Final = _GENERAL_SETTINGS_UI_LITELLM_FIELDS[field_name]
field_type: Final = spec["type"]
if value is None or value == "":
@ -16015,7 +16015,7 @@ def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) ->
async def _persist_general_settings_ui_litellm_field(
field_name: str, value: Any, user_api_key_dict: UserAPIKeyAuth
field_name: str, value: object, user_api_key_dict: UserAPIKeyAuth
) -> dict:
validated: Final = _validate_general_settings_ui_litellm_value(field_name, value)
config: Final = await proxy_config.get_config()

68
litellm/proxy/response_polling/background_streaming.py
View file

@ -10,10 +10,12 @@ https://platform.openai.com/docs/api-reference/responses-streaming
import asyncio
import json
from typing import TYPE_CHECKING, Any, Final, cast
from collections.abc import Sequence
from typing import TYPE_CHECKING, Final, TypedDict, cast
from fastapi import Request, Response
from fastapi.responses import StreamingResponse
from typing_extensions import ReadOnly
from litellm._logging import verbose_proxy_logger
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
@ -27,6 +29,15 @@ if TYPE_CHECKING:
from litellm.router import Router
class _StreamContentPart(TypedDict, total=False):
text: ReadOnly[str]
class _StreamOutputItem(TypedDict, total=False):
id: ReadOnly[str]
content: ReadOnly[Sequence[_StreamContentPart | None]]
async def background_streaming_task(
polling_id: str,
data,
@ -97,8 +108,9 @@ async def background_streaming_task(
# Process streaming response following OpenAI events format
# https://platform.openai.com/docs/api-reference/responses-streaming
output_items: Final[dict[str, dict[str, Any]]] = {} # Track output items by ID
accumulated_text: Final = {} # Track accumulated text deltas by (item_id, content_index)
output_items: Final[dict[str, _StreamOutputItem]] = {} # Track output items by ID
# Track accumulated text deltas by (item_id, content_index)
accumulated_text: Final[dict[tuple[str, int], str]] = {}
# ResponsesAPIResponse fields to extract from response.completed
usage_data = None
@ -187,16 +199,19 @@ async def background_streaming_task(
if item_id and item_id in output_items:
# Update the output item with new content
if "content" not in output_items[item_id]:
output_items[item_id]["content"] = []
output_items[item_id]["content"].append(content_part)
current_item = output_items[item_id]
appended_item: _StreamOutputItem = {
**current_item,
"content": (*current_item.get("content", ()), content_part),
}
output_items[item_id] = appended_item
state_dirty = True
elif event_type == "response.output_text.delta":
# Text delta - accumulate text content
# https://platform.openai.com/docs/api-reference/responses-streaming/response-text-delta
item_id = event.get("item_id")
content_index = event.get("content_index", 0)
content_index: int = event.get("content_index", 0)
delta = event.get("delta", "")
if item_id and item_id in output_items:
@ -207,12 +222,24 @@ async def background_streaming_task(
accumulated_text[key] += delta
# Update the content in output_items
if "content" in output_items[item_id]:
content_list = output_items[item_id]["content"]
if content_index < len(content_list):
# Update existing content part with accumulated text
if isinstance(content_list[content_index], dict):
content_list[content_index]["text"] = accumulated_text[key]
current_item = output_items[item_id]
content_list: Sequence[_StreamContentPart | None] = current_item.get("content", ())
if content_index < len(content_list):
# Update existing content part with accumulated text
content_entry = content_list[content_index]
if isinstance(content_entry, dict):
delta_part: _StreamContentPart = {
**content_entry,
"text": accumulated_text[key],
}
delta_item: _StreamOutputItem = {
**current_item,
"content": tuple(
delta_part if index == content_index else entry
for index, entry in enumerate(content_list)
),
}
output_items[item_id] = delta_item
state_dirty = True
elif event_type == "response.content_part.done":
@ -223,10 +250,17 @@ async def background_streaming_task(
if item_id and item_id in output_items:
# Update with final content from event
if "content" in output_items[item_id]:
content_list = output_items[item_id]["content"]
if content_index < len(content_list):
content_list[content_index] = content_part
current_item = output_items[item_id]
content_list = current_item.get("content", ())
if content_index < len(content_list):
finalized_item: _StreamOutputItem = {
**current_item,
"content": tuple(
content_part if index == content_index else entry
for index, entry in enumerate(content_list)
),
}
output_items[item_id] = finalized_item
state_dirty = True
elif event_type == "response.output_item.done":

69
litellm/proxy/utils.py
View file

@ -11,13 +11,15 @@ import sys
import threading
import time
import traceback
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence
from dataclasses import dataclass, field
from datetime import date, datetime, timedelta, timezone
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, TypeVar, Union, cast, overload
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, TypeVar, Union, cast, overload
from typing_extensions import ReadOnly, TypedDict
from litellm import _custom_logger_compatible_callbacks_literal
from litellm.constants import (
@ -170,7 +172,9 @@ from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams
if TYPE_CHECKING:
from mcp.types import CallToolResult
from opentelemetry.trace import Span as _Span
from prisma.actions import LiteLLM_DeprecatedVerificationTokenActions
from prisma.client import TransactionManager
from prisma.models import LiteLLM_DeprecatedVerificationToken
from prisma.types import HttpConfig
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -185,6 +189,24 @@ else:
_T: Final = TypeVar("_T")
class _ViewCountRow(TypedDict):
view_count: ReadOnly[int]
view_names: ReadOnly[Sequence[str] | None]
class _RelTuplesRow(TypedDict):
reltuples: ReadOnly[int]
class _EndUserBatchTable(Protocol):
def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ...
class _EndUserSpendBatch(Protocol):
@property
def litellm_endusertable(self) -> _EndUserBatchTable: ...
unified_guardrail: Final = UnifiedLLMGuardrails()
NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES: "frozenset[CallTypes]" = frozenset({CallTypes.anthropic_messages})
@ -363,10 +385,10 @@ def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback:
detail: Final = getattr(exc, "detail", None)
if not isinstance(detail, dict):
return
guardrail_name: Final = getattr(callback, "guardrail_name", None)
guardrail_name: Final[object] = getattr(callback, "guardrail_name", None)
if guardrail_name:
detail.setdefault("guardrail_name", guardrail_name)
event_hook: Final = getattr(callback, "event_hook", None)
event_hook: Final[object] = getattr(callback, "event_hook", None)
if event_hook:
detail.setdefault("guardrail_mode", event_hook)
@ -1043,7 +1065,7 @@ class ProxyLogging:
# Select guardrail using router's load balancing
selected_guardrail: Final = llm_router.get_available_guardrail(guardrail_name=guardrail_name)
callback: Final = selected_guardrail.get("callback")
callback: Final[CustomGuardrail | None] = selected_guardrail.get("callback")
if callback is None:
raise ValueError(f"No callback found for guardrail: {guardrail_name}")
@ -2107,7 +2129,7 @@ class ProxyLogging:
Related issue - https://github.com/BerriAI/litellm/issues/3395
"""
litellm_debug_info: Final = getattr(original_exception, "litellm_debug_info", None)
litellm_debug_info: Final[str | None] = getattr(original_exception, "litellm_debug_info", None)
exception_str = str(original_exception)
if litellm_debug_info is not None:
exception_str += litellm_debug_info
@ -2429,7 +2451,7 @@ class ProxyLogging:
#################################################################
for callback in other_callbacks:
callback_response = await callback.async_post_call_success_hook(
callback_response: LLMResponseTypes | None = await callback.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict, data=data, response=response
)
if callback_response is not None:
@ -2707,6 +2729,9 @@ class ProxyLogging:
complete_response = str_so_far + response_str
else:
complete_response = response_str
callback_response: (
ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream | None
)
callback_response = await _callback.async_post_call_streaming_hook(
user_api_key_dict=user_api_key_dict,
response=complete_response,
@ -2813,8 +2838,10 @@ class ProxyLogging:
logging_obj: Final = request_data.get("litellm_logging_obj")
if logging_obj is None:
return
_deferred_cb: Final = getattr(logging_obj, "_on_deferred_stream_complete", None)
_args: Final = getattr(logging_obj, "_deferred_stream_complete_args", None)
_deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr(
logging_obj, "_on_deferred_stream_complete", None
)
_args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None)
if _deferred_cb is not None and _args is not None:
logging_obj._on_deferred_stream_complete = None
logging_obj._deferred_stream_complete_args = None
@ -2908,7 +2935,10 @@ async def _lookup_deprecated_key(
_deprecated_key_cache.pop(hashed_token, None)
try:
deprecated_row: Final = await db.litellm_deprecatedverificationtoken.find_first(
deprecated_keys_table: Final[
LiteLLM_DeprecatedVerificationTokenActions[LiteLLM_DeprecatedVerificationToken]
] = db.litellm_deprecatedverificationtoken
deprecated_row: Final = await deprecated_keys_table.find_first(
where={
"token": hashed_token,
"revoke_at": {"gt": now},
@ -3337,7 +3367,7 @@ class PrismaClient:
required_view: Final = "LiteLLM_VerificationTokenView"
expected_views_str: Final = ", ".join(f"'{view}'" for view in expected_views)
pg_schema: Final = os.getenv("DATABASE_SCHEMA", "public")
ret: Final = await self.db.query_raw(f"""
ret: Final[Sequence[_ViewCountRow]] = await self.db.query_raw(f"""
WITH existing_views AS (
SELECT viewname
FROM pg_views
@ -4345,7 +4375,9 @@ class PrismaClient:
else:
filter_query = {"token": {"in": hashed_tokens}}
deleted_tokens: Final = await VerificationTokenRepository(self).table.delete_many(where=filter_query)
deleted_tokens: Final[int] = await VerificationTokenRepository(self).table.delete_many(
where=filter_query
)
verbose_proxy_logger.debug("deleted_tokens: %s", deleted_tokens)
return {"deleted_keys": deleted_tokens}
elif table_name == "team" and team_id_list is not None and isinstance(team_id_list, list):
@ -4450,7 +4482,7 @@ class PrismaClient:
engine: Final = prisma_obj._engine
process: Final = getattr(engine, "process", None) if engine is not None else None
if process is not None:
pid: Final = process.pid
pid: Final[object] = process.pid
if isinstance(pid, int):
return pid
except (AttributeError, TypeError):
@ -5257,7 +5289,7 @@ class PrismaClient:
about to check, and attribute the failure to the wrong replacement.
"""
sql_query: Final = "SELECT 1"
response: Final = await wrapper.query_raw(sql_query)
response: Final[object] = await wrapper.query_raw(sql_query)
return response
async def _probe_answers_now(self, wrapper: PrismaWrapper) -> bool:
@ -5383,7 +5415,7 @@ class PrismaClient:
FROM pg_class
WHERE oid = '"LiteLLM_SpendLogs"'::regclass;
"""
result: Final = await self.db.query_raw(query=sql_query)
result: Final[Sequence[_RelTuplesRow]] = await self.db.query_raw(query=sql_query)
return result[0]["reltuples"]
try:
@ -5540,7 +5572,7 @@ async def _cache_user_row(user_id: str, cache: DualCache, db: PrismaClient):
if user_row is not None:
print_verbose(f"User Row: {user_row}, type = {type(user_row)}")
if hasattr(user_row, "model_dump_json") and callable(getattr(user_row, "model_dump_json")):
cache_value: Final = user_row.model_dump_json()
cache_value: Final[str] = user_row.model_dump_json()
cache.set_cache(key=cache_key, value=cache_value, ttl=600) # store for 10 minutes
@ -5766,6 +5798,7 @@ class ProxyUpdateSpend:
start_time = time.time()
try:
async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction:
batcher: _EndUserSpendBatch
async with transaction.batch_() as batcher:
# Sort by end_user_id for consistent lock ordering across pods to prevent deadlocks.
for end_user_id, response_cost in sorted(end_user_list_transactions.items()):
@ -6400,7 +6433,7 @@ def _check_and_merge_model_level_guardrails(
# Medium on #29654).
team_id: Final = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id")
model_level_guardrails: list | None = None
model_level_guardrails: list[object] | None = None
if model_id is not None:
deployment: Final = llm_router.get_deployment(model_id=model_id)
if deployment is None:
@ -6449,7 +6482,7 @@ def _check_and_merge_model_level_guardrails(
return _merge_guardrails_with_existing(data, model_level_guardrails)
def _merge_guardrails_with_existing(data: dict, model_level_guardrails: Any) -> dict:
def _merge_guardrails_with_existing(data: dict, model_level_guardrails: object) -> dict:
"""
Merge model-level guardrails with any existing guardrails in the request data.

15
litellm/responses/mcp/chat_completions_handler.py
View file

@ -1,15 +1,18 @@
"""Helpers for handling MCP-aware `/chat/completions` requests."""
import logging
from typing import Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, cast
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
from litellm.responses.mcp.request_context import MCPRequestContext
from litellm.types.utils import ModelResponse
from litellm.types.utils import Message, ModelResponse
from litellm.utils import CustomStreamWrapper
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
def _add_mcp_metadata_to_response(
response: ModelResponse | CustomStreamWrapper,
@ -55,7 +58,7 @@ def _add_mcp_metadata_to_response(
# Add MCP metadata to all choices' messages
for choice in response.choices:
message = getattr(choice, "message", None)
message: Message | None = getattr(choice, "message", None)
if message is not None:
# Get existing provider_specific_fields or create new dict
provider_fields = getattr(message, "provider_specific_fields", None) or {}
@ -109,7 +112,7 @@ async def acompletion_with_mcp(
)
context: Final = MCPRequestContext.resolve(kwargs=kwargs, tools=tools)
user_api_key_auth: Final = context.user_api_key_auth
user_api_key_auth: Final[UserAPIKeyAuth | None] = context.user_api_key_auth
request_tags: Final = list(context.request_tags) if context.request_tags else None
mcp_auth_header: Final = context.mcp_auth_header
mcp_server_auth_headers: Final = context.mcp_server_auth_headers
@ -165,7 +168,7 @@ async def acompletion_with_mcp(
return response
# For auto-execute: handle streaming vs non-streaming differently
stream: Final = kwargs.get("stream", False)
stream: Final[bool] = kwargs.get("stream", False)
mock_tool_calls: Final = base_call_args.pop("mock_tool_calls", None)
if stream:
@ -539,7 +542,7 @@ async def acompletion_with_mcp(
self.__iter__()
return next(self._sync_iterator)
def __getattr__(self, name):
def __getattr__(self, name: str) -> object:
# Delegate all other attributes to original wrapper
return getattr(self._original_wrapper, name)

36
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -25,7 +25,10 @@ from litellm.types.llms.openai import (
from litellm.types.llms.openai import ToolParam as ResponsesToolParam
from litellm.types.utils import (
CallTypes,
ChatCompletionMessageCustomToolCall,
ChatCompletionMessageToolCall,
Choices,
Message,
ModelResponse,
StandardLoggingMCPToolCall,
)
@ -419,12 +422,14 @@ class LiteLLM_Proxy_MCP_Handler:
if not mcp_tools_with_litellm_proxy:
return [], {}
typed_user_api_key_auth: Final[UserAPIKeyAuth | None] = user_api_key_auth
# Step 1: Fetch MCP tools from manager
(
mcp_tools_fetched,
allowed_mcp_servers,
) = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager(
user_api_key_auth=user_api_key_auth,
user_api_key_auth=typed_user_api_key_auth,
mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy,
litellm_trace_id=litellm_trace_id,
mcp_auth_header=mcp_auth_header,
@ -527,10 +532,12 @@ class LiteLLM_Proxy_MCP_Handler:
try:
for choice in response.choices:
message = getattr(choice, "message", None)
message: Message | None = getattr(choice, "message", None)
if message is None:
continue
tool_call_entries = getattr(message, "tool_calls", None)
tool_call_entries: (
Sequence[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None
) = getattr(message, "tool_calls", None)
if tool_call_entries:
for tool_call in tool_call_entries:
if hasattr(tool_call, "model_dump"):
@ -564,7 +571,7 @@ class LiteLLM_Proxy_MCP_Handler:
else:
tool_call_id = getattr(tool_call, "call_id", None) or getattr(tool_call, "id", None)
function_obj: Final = getattr(tool_call, "function", None)
function_obj: Final[object] = getattr(tool_call, "function", None)
if function_obj is not None:
tool_name = getattr(function_obj, "name", None)
tool_arguments = getattr(function_obj, "arguments", None)
@ -655,6 +662,7 @@ class LiteLLM_Proxy_MCP_Handler:
tool_call_id: str | None = None
rules_obj: Final = Rules()
logging_safe_headers: Final = logging_safe_mcp_headers(raw_headers)
typed_user_api_key_auth: Final[UserAPIKeyAuth | None] = user_api_key_auth
for tool_call in tool_calls:
logging_request_data: dict[str, object] = {}
tool_name: str | None = None
@ -722,18 +730,18 @@ class LiteLLM_Proxy_MCP_Handler:
logging_request_data["litellm_trace_id"] = litellm_trace_id
if request_tags:
logging_metadata["tags"] = request_tags
if user_api_key_auth is not None:
if typed_user_api_key_auth is not None:
from litellm.proxy.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
)
LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
data=logging_request_data,
user_api_key_dict=user_api_key_auth,
user_api_key_dict=typed_user_api_key_auth,
_metadata_variable_name="metadata",
)
user_identifier = getattr(user_api_key_auth, "end_user_id", None) or getattr(
user_api_key_auth, "user_id", None
user_identifier = getattr(typed_user_api_key_auth, "end_user_id", None) or getattr(
typed_user_api_key_auth, "user_id", None
)
if user_identifier:
logging_request_data["user"] = user_identifier
@ -792,7 +800,7 @@ class LiteLLM_Proxy_MCP_Handler:
server_name=server_name,
name=sanitized_tool_name,
arguments=parsed_arguments,
user_api_key_auth=user_api_key_auth,
user_api_key_auth=typed_user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
@ -808,7 +816,7 @@ class LiteLLM_Proxy_MCP_Handler:
if litellm_logging_obj
else {"mcp_tool_name": tool_name}
),
user_api_key_dict=user_api_key_auth,
user_api_key_dict=typed_user_api_key_auth,
)
if litellm_logging_obj:
@ -844,7 +852,7 @@ class LiteLLM_Proxy_MCP_Handler:
except BlockedPiiEntityError as e:
await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure(
proxy_logging_obj=proxy_logging_obj,
user_api_key_auth=user_api_key_auth,
user_api_key_auth=typed_user_api_key_auth,
request_data=logging_request_data,
error=e,
)
@ -860,7 +868,7 @@ class LiteLLM_Proxy_MCP_Handler:
except GuardrailRaisedException as e:
await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure(
proxy_logging_obj=proxy_logging_obj,
user_api_key_auth=user_api_key_auth,
user_api_key_auth=typed_user_api_key_auth,
request_data=logging_request_data,
error=e,
)
@ -878,7 +886,7 @@ class LiteLLM_Proxy_MCP_Handler:
except HTTPException as e:
await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure(
proxy_logging_obj=proxy_logging_obj,
user_api_key_auth=user_api_key_auth,
user_api_key_auth=typed_user_api_key_auth,
request_data=logging_request_data,
error=e,
)
@ -894,7 +902,7 @@ class LiteLLM_Proxy_MCP_Handler:
except Exception as e:
await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure(
proxy_logging_obj=proxy_logging_obj,
user_api_key_auth=user_api_key_auth,
user_api_key_auth=typed_user_api_key_auth,
request_data=logging_request_data,
error=e,
)

4
litellm/responses/mcp/mcp_streaming_iterator.py
View file

@ -511,7 +511,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
if self.base_iterator:
if hasattr(self.base_iterator, "__anext__"):
try:
chunk: Final = await cast(Any, self.base_iterator).__anext__()
chunk: Final[ResponsesAPIStreamingResponse] = await cast(Any, self.base_iterator).__anext__()
# Capture the response ID from the first event to ensure consistency
if self._cached_response_id is None and hasattr(chunk, "response"):
@ -569,7 +569,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
if not self.base_iterator or not hasattr(self.base_iterator, "__anext__"):
raise StopAsyncIteration
chunk: Final = await cast(Any, self.base_iterator).__anext__()
chunk: Final[ResponsesAPIStreamingResponse] = await cast(Any, self.base_iterator).__anext__()
if self._cached_response_id is None and hasattr(chunk, "response"):
new_response: Final[ResponsesAPIResponse | None] = getattr(chunk, "response", None)

2
litellm/router.py
View file

@ -254,7 +254,7 @@ if TYPE_CHECKING:
ResponsesAPIResponse,
)
Span = _Span | Any
Span = _Span
else:
Span = Any
AutoRouter = Any

92
litellm/vector_stores/main.py
View file

@ -5,9 +5,9 @@ LiteLLM SDK Functions for Creating and Searching Vector Stores
import asyncio
import builtins
import contextvars
from collections.abc import Coroutine
from collections.abc import Coroutine, Mapping
from functools import partial
from typing import Any, Final
from typing import Final
import httpx
@ -96,9 +96,9 @@ async def acreate(
metadata: dict[str, str] | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: str | None = None,
@ -160,14 +160,14 @@ def create(
metadata: dict[str, str] | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: str | None = None,
**kwargs,
) -> VectorStoreCreateResponse | Coroutine[Any, Any, VectorStoreCreateResponse]:
) -> VectorStoreCreateResponse | Coroutine[object, object, VectorStoreCreateResponse]:
"""
Create a vector store.
@ -274,9 +274,9 @@ async def asearch(
rewrite_query: bool | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: str | None = None,
@ -341,14 +341,14 @@ def search(
rewrite_query: bool | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: str | None = None,
**kwargs,
) -> VectorStoreSearchResponse | Coroutine[Any, Any, VectorStoreSearchResponse]:
) -> VectorStoreSearchResponse | Coroutine[object, object, VectorStoreSearchResponse]:
"""
Search a vector store for relevant chunks based on a query and file attributes filter.
@ -466,9 +466,9 @@ def search(
@client
async def aretrieve(
vector_store_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -518,13 +518,13 @@ async def aretrieve(
@client
def retrieve(
vector_store_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> VectorStoreCreateResponse | Coroutine[Any, Any, VectorStoreCreateResponse]:
) -> VectorStoreCreateResponse | Coroutine[object, object, VectorStoreCreateResponse]:
"""
Retrieve a vector store.
@ -601,13 +601,13 @@ async def alist(
before: str | None = None,
limit: int | None = 20,
order: str | None = "desc",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
):
) -> Mapping[str, object]:
"""
Async: List vector stores.
"""
@ -638,7 +638,7 @@ async def alist(
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
response: Mapping[str, object] = await init_response
else:
response = init_response
@ -659,9 +659,9 @@ def list(
before: str | None = None,
limit: int | None = 20,
order: str | None = "desc",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -753,9 +753,9 @@ async def aupdate(
name: str | None = None,
expires_after: dict | None = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -811,13 +811,13 @@ def update(
name: str | None = None,
expires_after: dict | None = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> VectorStoreCreateResponse | Coroutine[Any, Any, VectorStoreCreateResponse]:
) -> VectorStoreCreateResponse | Coroutine[object, object, VectorStoreCreateResponse]:
"""
Update a vector store.
@ -905,13 +905,13 @@ def update(
@client
async def adelete(
vector_store_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
):
) -> Mapping[str, object]:
"""
Async: Delete a vector store.
"""
@ -939,7 +939,7 @@ async def adelete(
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
response: Mapping[str, object] = await init_response
else:
response = init_response
@ -957,9 +957,9 @@ async def adelete(
@client
def delete(
vector_store_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,

16
ruff-strict-budget.json
View file

@ -1,6 +1,6 @@
{
"ANN001": {
"limit": 3046
"limit": 3026
},
"ANN002": {
"limit": 71
@ -9,13 +9,13 @@
"limit": 827
},
"ANN201": {
"limit": 2022
"limit": 2017
},
"ANN202": {
"limit": 855
},
"ANN204": {
"limit": 712
"limit": 711
},
"ANN205": {
"limit": 114
@ -24,7 +24,7 @@
"limit": 133
},
"ANN401": {
"limit": 1341
"limit": 1290
},
"ASYNC230": {
"limit": 11
@ -39,7 +39,7 @@
"limit": 505
},
"B009": {
"limit": 60
"limit": 59
},
"B010": {
"limit": 190
@ -201,7 +201,7 @@
"limit": 58
},
"SIM102": {
"limit": 321
"limit": 319
},
"SIM103": {
"limit": 119
@ -213,7 +213,7 @@
"limit": 2
},
"SIM117": {
"limit": 7
"limit": 6
},
"SIM201": {
"limit": 1
@ -234,7 +234,7 @@
"limit": 5
},
"TID251": {
"limit": 1220
"limit": 1216
},
"TRY002": {
"limit": 524

10
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 22938
"limit": 22909
},
"LIT002": {
"limit": 26901
"limit": 26898
},
"LIT003": {
"limit": 269
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1072
"limit": 1071
},
"LIT007": {
"limit": 0
@ -27,10 +27,10 @@
"limit": 0
},
"LIT010": {
"limit": 16715
"limit": 16713
},
"LIT011": {
"limit": 5593
"limit": 5591
},
"LIT012": {
"limit": 4519