From 255dad6716b229394c1515691014bb543330f89d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 16 Aug 2026 03:56:02 +0000 Subject: [PATCH] chore(typing): drop 1.3k basedpyright errors across 30 Any hotspot files Replace Any-typed seams with real types in files carrying the highest remaining reportAny/reportExplicitAny density after #34745: the proxy server and its utils, the router, the streaming handler and chunk builder, litellm_logging, the redis cache, the MCP db/tool-registry/spend-writer layer, the anthropic pass-through adapters and guardrail translation, the lasso and presidio guardrail hooks, the azure_ai agents handler, the management endpoints (keys, users, ui_sso, model access groups, config override, MCP, projects), the responses MCP handlers, response polling background streaming, and the containers and vector stores mains No casts, no type: ignore, no noqa, no new suppressions, and no Any annotations that were not already at base. Whole-tree basedpyright: reportAny 14,610 -> 14,009, reportExplicitAny 5,100 -> 4,780, total 144,743 -> 143,471, with no rule increasing repo-wide or in any file. Budgets ratcheted: basedpyright -1,272 across 48 rules, ruff-strict -85, type-discipline -37 --- basedpyright-code-budget.json | 26 +-- .../management_endpoints/project_endpoints.py | 84 ++++++--- litellm/caching/redis_cache.py | 36 ++-- litellm/containers/main.py | 158 ++++++++-------- litellm/litellm_core_utils/litellm_logging.py | 42 +++-- .../streaming_chunk_builder_utils.py | 172 ++++++++++++------ .../litellm_core_utils/streaming_handler.py | 88 ++++++--- .../chat/guardrail_translation/handler.py | 65 +++---- .../adapters/transformation.py | 35 ++-- litellm/llms/azure_ai/agents/handler.py | 156 ++++++++++++---- litellm/llms/custom_httpx/llm_http_handler.py | 15 +- litellm/proxy/_experimental/mcp_server/db.py | 148 +++++++++++---- litellm/proxy/db/db_spend_update_writer.py | 48 ++++- litellm/proxy/db/tool_registry_writer.py | 68 +++++-- .../guardrails/guardrail_hooks/lasso/lasso.py | 42 ++--- .../guardrails/guardrail_hooks/presidio.py | 50 +++-- litellm/proxy/guardrails/usage_endpoints.py | 111 +++++++---- .../config_override_endpoints.py | 77 +++++--- .../internal_user_endpoints.py | 51 ++++-- .../key_management_endpoints.py | 10 +- .../mcp_management_endpoints.py | 68 ++++--- ...model_access_group_management_endpoints.py | 39 +++- litellm/proxy/management_endpoints/ui_sso.py | 45 ++--- litellm/proxy/proxy_server.py | 64 +++---- .../response_polling/background_streaming.py | 68 +++++-- litellm/proxy/utils.py | 69 +++++-- .../responses/mcp/chat_completions_handler.py | 15 +- .../mcp/litellm_proxy_mcp_handler.py | 36 ++-- .../responses/mcp/mcp_streaming_iterator.py | 4 +- litellm/router.py | 2 +- litellm/vector_stores/main.py | 92 +++++----- ruff-strict-budget.json | 16 +- type-discipline-budget.json | 10 +- 33 files changed, 1301 insertions(+), 709 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 06010c706e3..c09369b9b2c 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -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 diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index 66fac8d76ee..579f203554e 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -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}, ) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index a3936fd17e2..934ba500ef9 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -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: diff --git a/litellm/containers/main.py b/litellm/containers/main.py index 69bd48fbb6d..97ca11872c1 100644 --- a/litellm/containers/main.py +++ b/litellm/containers/main.py @@ -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, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a72d46e3fe8..4a4a97b1d85 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index ab4017b144b..67287c903be 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -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 diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 99b1c1a2ab7..51974e85d9c 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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: diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index e4a4d23b438..f2b24984ccf 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -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"): diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 667f9dcaab0..f691a1a4454 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -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, diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index 5f804e901cd..a13b1300e55 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -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 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 721b9545ac1..e22ec89847e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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. diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 08a8b1bc7b3..2f07a8b716c 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -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}} diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index fe68e837a8e..b6dddbb029d 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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::::user_id::" 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()): diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index b2fa538b1cc..187a18be845 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -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}, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index 725c06b8618..ea022510309 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -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") diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 7b2f06e4bfb..c3b7498d9ec 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -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: diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 029a26e84f8..f48e20257db 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -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") diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py index 06184cb40fa..12c99477d3c 100644 --- a/litellm/proxy/management_endpoints/config_override_endpoints.py +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -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, diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 6e1e6d22cb1..2b88658e1b4 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -17,7 +17,7 @@ import json import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final, 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, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7e190e8b19d..ca2607653a1 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 997012dbc65..06c32af2dc2 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -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( diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index 7051f705a03..49c0135ff10 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -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( diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index b87ad8597dc..46af5dd80e1 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -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. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bda6fc25499..e07870c6867 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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() diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 31ab3596418..020698dabd9 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -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": diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 498b6d7ee3d..a972c65c889 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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. diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 38e6d07c626..2a0406f9a4d 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -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) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 56818717c09..b3725d86ec3 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -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, ) diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 022b9ece32e..c7471518398 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -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) diff --git a/litellm/router.py b/litellm/router.py index 0fd3cf6af1b..9d65a72b59a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -254,7 +254,7 @@ if TYPE_CHECKING: ResponsesAPIResponse, ) - Span = _Span | Any + Span = _Span else: Span = Any AutoRouter = Any diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index c8ed6de23b3..9b0ff71730a 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -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, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index bd585bb2719..22d3119d457 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -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 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 8e55b1533ea..ca848190a32 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -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