From 64e993773da57f5bd39dc60b1d45b2eaafc5acd6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 05:25:22 +0000 Subject: [PATCH 01/28] chore(typing): clear 1.5k basedpyright Any errors across 54 files Retypes the 54 highest-density reportAny/reportExplicitAny sources with real types instead of shuffling the ceilings around: typed prisma table Protocols so the untyped client surface stops at the query, local TypedDicts for JSON and dict payloads, concrete chunk and logging types on the streaming and callback surfaces, and 3-argument getattr with a Callable annotation where an SDK object is genuinely duck-typed No cast(), no type: ignore, no noqa, no suppression comments, and no new Any annotations. Whole-tree basedpyright drops 1,941 errors with no rule rising anywhere, and all three budget files are ratcheted so the cleared headroom cannot silently grow back Adds a GDC regression test pinning the named AttributeError that the typed credential accessor now raises when with_gdch_audience is missing --- basedpyright-code-budget.json | 26 +-- .../proxy/audit_logging_endpoints.py | 108 +++++++---- litellm/_lazy_imports.py | 79 ++++---- .../transformation.py | 46 ++--- litellm/caching/valkey_semantic_cache.py | 71 ++++--- .../bitbucket/bitbucket_client.py | 49 +++-- .../compression_interception/handler.py | 84 +++++---- litellm/integrations/custom_logger.py | 28 +-- .../opik_payload_builder/payload_builders.py | 12 +- .../prometheus_helpers/prometheus_api.py | 54 +++++- .../websearch_interception/handler.py | 16 +- .../llm_response_utils/response_metadata.py | 31 ++- .../model_response_utils.py | 66 +++++-- .../litellm_core_utils/streaming_handler.py | 8 +- litellm/litellm_core_utils/url_utils.py | 25 +-- .../llms/anthropic/batches/transformation.py | 59 ++++-- .../context_management/dispatcher.py | 81 +++++--- .../responses_adapters/handler.py | 40 ++-- .../llms/anthropic/skills/transformation.py | 30 ++- litellm/llms/azure/files/handler.py | 13 +- litellm/llms/bedrock/realtime/handler.py | 104 ++++++++-- litellm/llms/chatgpt/chat/streaming_utils.py | 41 +++- .../llms/compactifai/chat/transformation.py | 33 +++- .../llms/custom_httpx/container_handler.py | 178 ++++++++++++------ litellm/llms/gdc/chat/transformation.py | 45 ++++- .../llms/infinity/rerank/transformation.py | 29 ++- .../litellm_proxy/skills/code_execution.py | 76 ++++++-- litellm/llms/oci/chat/cohere.py | 38 +++- litellm/llms/ollama/completion/handler.py | 36 +++- .../llms/openai/containers/transformation.py | 72 +++++-- .../vector_stores/rag_api/transformation.py | 127 ++++++++++--- litellm/models/base.py | 2 +- .../mcp_server/oauth2_flow_backfill.py | 63 ++++++- .../_experimental/mcp_server/toolset_db.py | 93 +++++++-- .../proxy/agent_endpoints/agent_registry.py | 81 +++++--- .../proxy/client/cli/commands/credentials.py | 51 ++++- litellm/proxy/client/cli/commands/teams.py | 52 +++-- litellm/proxy/common_utils/callback_utils.py | 14 +- litellm/proxy/common_utils/get_routes.py | 65 ++++--- .../proxy/common_utils/user_api_key_cache.py | 48 ++--- litellm/proxy/db/routing_prisma_wrapper.py | 36 ++-- .../guardrail_hooks/custom_code/sandbox.py | 14 +- .../hiddenlayer/hiddenlayer.py | 33 ++-- .../llm_as_a_judge/__init__.py | 130 ++++++++++--- .../guardrail_hooks/noma/noma_v2.py | 26 +-- .../promptguard/promptguard.py | 20 +- .../jwt_key_mapping_endpoints.py | 76 ++++++-- .../response_polling/background_streaming.py | 47 ++++- .../search_endpoints/search_tool_registry.py | 57 ++++-- litellm/rag/rag_query.py | 65 +++++-- litellm/repositories/base_repository.py | 51 +++-- .../repositories/credentials_repository.py | 58 ++++-- litellm/repositories/team_repository.py | 76 ++++++-- litellm/rust_bridge/responses_websocket.py | 23 ++- .../secret_managers/secret_manager_handler.py | 60 +++++- ruff-strict-budget.json | 18 +- .../gdc/chat/test_gdc_chat_transformation.py | 19 ++ type-discipline-budget.json | 12 +- 58 files changed, 2172 insertions(+), 823 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index b4c324a2c4c..58b8ae01cb0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 19955 + "limit": 18773 }, "reportArgumentType": { "limit": 2566 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 6049 + "limit": 5774 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5663 + "limit": 5640 }, "reportMissingTypeArgument": { - "limit": 15555 + "limit": 15498 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,31 +99,31 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44655 + "limit": 44589 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 39017 + "limit": 38882 }, "reportUnknownParameterType": { - "limit": 19885 + "limit": 19806 }, "reportUnknownVariableType": { - "limit": 30572 + "limit": 30456 }, "reportUnnecessaryCast": { - "limit": 117 + "limit": 116 }, "reportUnnecessaryComparison": { - "limit": 699 + "limit": 698 }, "reportUnnecessaryContains": { "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 836 + "limit": 835 }, "reportUntypedBaseClass": { "limit": 0 @@ -138,9 +138,9 @@ "limit": 139 }, "reportUnusedImport": { - "limit": 545 + "limit": 544 }, "reportUnusedVariable": { - "limit": 146 + "limit": 142 } } diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py index 18ac29b9781..21bde06e86a 100644 --- a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py @@ -7,7 +7,9 @@ GET - /audit/{id} - Get audit log by id GET - /audit - Get all audit logs """ -from typing import Any, Dict, List, Optional +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import Final, Protocol #### AUDIT LOGGING #### from fastapi import APIRouter, Depends, HTTPException, Query @@ -15,6 +17,7 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import ( AuditLogResponse, PaginatedAuditLogResponse, ) +from typing_extensions import ReadOnly, TypedDict from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -22,7 +25,44 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router = APIRouter() -def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]: +class _AuditLogFields(TypedDict): + """Columns of the `LiteLLM_AuditLog` table, as returned by `model_dump()`.""" + + id: ReadOnly[str] + updated_at: ReadOnly[datetime] + changed_by: ReadOnly[str] + changed_by_api_key: ReadOnly[str] + action: ReadOnly[str] + table_name: ReadOnly[str] + object_id: ReadOnly[str] + before_value: ReadOnly[dict[str, object] | None] + updated_values: ReadOnly[dict[str, object] | None] + + +class _AuditLogRecord(Protocol): + """Row of the `LiteLLM_AuditLog` table as materialised by the Prisma client.""" + + def model_dump(self) -> _AuditLogFields: ... + + +class _AuditLogTable(Protocol): + """The `litellm_auditlog` accessor of the Prisma client.""" + + async def find_many( + self, + *, + where: Mapping[str, object], + order: Mapping[str, str], + skip: int, + take: int, + ) -> Sequence[_AuditLogRecord]: ... + + async def count(self, *, where: Mapping[str, object]) -> int: ... + + async def find_unique(self, *, where: Mapping[str, str]) -> _AuditLogRecord | None: ... + + +def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, object]: """ Build an OR condition that matches a value inside a JSON column at the given key, checking both before_value and updated_values. @@ -53,33 +93,33 @@ async def get_audit_logs( page: int = Query(1, ge=1), page_size: int = Query(10, ge=1, le=100), # Filter parameters - changed_by: Optional[str] = Query( + changed_by: str | None = Query( None, description="Filter by user or system that performed the action" ), - changed_by_api_key: Optional[str] = Query( + changed_by_api_key: str | None = Query( None, description="Filter by API key hash that performed the action" ), - action: Optional[str] = Query( + action: str | None = Query( None, description="Filter by action type (create, update, delete)" ), - table_name: Optional[str] = Query( + table_name: str | None = Query( None, description="Filter by table name that was modified" ), - object_id: Optional[str] = Query( + object_id: str | None = Query( None, description="Filter by ID of the object that was modified" ), - start_date: Optional[str] = Query(None, description="Filter logs after this date"), - end_date: Optional[str] = Query(None, description="Filter logs before this date"), - object_team_id: Optional[str] = Query( + start_date: str | None = Query(None, description="Filter logs after this date"), + end_date: str | None = Query(None, description="Filter logs before this date"), + object_team_id: str | None = Query( None, description="Filter by team_id present in before_value or updated_values JSON (PostgreSQL only)", ), - object_key_hash: Optional[str] = Query( + object_key_hash: str | None = Query( None, description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)", ), # Sorting parameters - sort_by: Optional[str] = Query( + sort_by: str | None = Query( None, description="Column to sort by (e.g. 'updated_at', 'action', 'table_name')", ), @@ -102,7 +142,7 @@ async def get_audit_logs( ) # Build filter conditions - where_conditions: Dict[str, Any] = {} + where_conditions: Final[dict[str, object]] = {} if changed_by: where_conditions["changed_by"] = changed_by if changed_by_api_key: @@ -114,33 +154,31 @@ async def get_audit_logs( if object_id: where_conditions["object_id"] = object_id if start_date or end_date: - date_filter: Dict[str, Any] = {} - if start_date: - date_filter["gte"] = start_date - if end_date: - date_filter["lte"] = end_date + date_filter: Final[Mapping[str, str]] = { + bound: bound_value for bound, bound_value in (("gte", start_date), ("lte", end_date)) if bound_value + } where_conditions["updated_at"] = date_filter # JSON field filters (PostgreSQL only) — each filter is AND'd with the # others, but checks both before_value and updated_values internally (OR). - if object_team_id: - where_conditions["AND"] = where_conditions.get("AND", []) + [ - _build_json_field_or_condition("team_id", object_team_id) - ] - if object_key_hash: - where_conditions["AND"] = where_conditions.get("AND", []) + [ - _build_json_field_or_condition("token", object_key_hash) + if object_team_id or object_key_hash: + where_conditions["AND"] = [ + _build_json_field_or_condition(json_key, json_value) + for json_key, json_value in ( + ("team_id", object_team_id), + ("token", object_key_hash), + ) + if json_value ] # Build sort conditions - order_by: Dict[str, Any] = {} - if sort_by and isinstance(sort_by, str): - order_by[sort_by] = sort_order - else: - order_by["updated_at"] = sort_order # Default sort by updated_at + sort_column: Final[str] = sort_by if sort_by and isinstance(sort_by, str) else "updated_at" + order_by: Final[Mapping[str, str]] = {sort_column: sort_order} + + audit_log_table: Final[_AuditLogTable] = prisma_client.db.litellm_auditlog # Get paginated results - audit_logs = await prisma_client.db.litellm_auditlog.find_many( + audit_logs: Final = await audit_log_table.find_many( where=where_conditions, order=order_by, skip=(page - 1) * page_size, @@ -148,8 +186,8 @@ async def get_audit_logs( ) # Get total count for pagination - total_count = await prisma_client.db.litellm_auditlog.count(where=where_conditions) - total_pages = -(-total_count // page_size) # Ceiling division + total_count: Final = await audit_log_table.count(where=where_conditions) + total_pages: Final = -(-total_count // page_size) # Ceiling division # Return paginated response return PaginatedAuditLogResponse( @@ -198,8 +236,10 @@ async def get_audit_log_by_id( detail={"message": CommonProxyErrors.db_not_connected_error.value}, ) + audit_log_table: Final[_AuditLogTable] = prisma_client.db.litellm_auditlog + # Get the audit log by ID - audit_log = await prisma_client.db.litellm_auditlog.find_unique(where={"id": id}) + audit_log: Final = await audit_log_table.find_unique(where={"id": id}) if audit_log is None: raise HTTPException( diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 933464d3f23..dff3b11e353 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -17,8 +17,14 @@ until they're actually needed. import importlib import sys -from collections.abc import Callable -from typing import Any, Final, cast +from collections.abc import Callable, Mapping +from typing import TYPE_CHECKING, Any, Final, cast + +if TYPE_CHECKING: + import httpx + import tiktoken + + from .caching.llm_caching_handler import LLMClientCache as LLMClientCacheType # Import all the data structures that define what can be lazy-loaded # These are just lists of names and maps of where to find them @@ -54,7 +60,7 @@ from ._lazy_imports_registry import ( ) -def get_litellm_globals() -> dict: +def get_litellm_globals() -> dict[str, object]: """ Get the globals dictionary of the litellm module. @@ -64,7 +70,7 @@ def get_litellm_globals() -> dict: return sys.modules["litellm"].__dict__ -def _get_utils_globals() -> dict: +def _get_utils_globals() -> dict[str, object]: """ Get the globals dictionary of the utils module. @@ -74,14 +80,19 @@ def _get_utils_globals() -> dict: return sys.modules["litellm.utils"].__dict__ +def _get_module_level_client_timeout(litellm_globals: Mapping[str, Any]) -> "float | httpx.Timeout | None": + """Read the configured `litellm.request_timeout` used for the module level http clients.""" + return litellm_globals.get("request_timeout") + + # These are special lazy loaders for things that are used internally # They're separate from the main lazy import system because they have specific use cases # Lazy loader for default encoding - avoids importing heavy tiktoken library at startup -_default_encoding: Any | None = None +_default_encoding: "tiktoken.Encoding | None" = None -def _get_default_encoding() -> Any: +def _get_default_encoding() -> "tiktoken.Encoding": """ Lazily load and cache the default OpenAI encoding. @@ -100,10 +111,10 @@ def _get_default_encoding() -> Any: # Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time -_get_modified_max_tokens_func: Any | None = None +_get_modified_max_tokens_func: Callable[..., int | None] | None = None -def _get_modified_max_tokens() -> Any: +def _get_modified_max_tokens() -> Callable[..., int | None]: """ Lazily load and cache the get_modified_max_tokens function. @@ -124,10 +135,10 @@ def _get_modified_max_tokens() -> Any: # Lazy loader for token_counter to avoid importing token_counter module at module import time -_token_counter_new_func: Any | None = None +_token_counter_new_func: Callable[..., int] | None = None -def _get_token_counter_new() -> Any: +def _get_token_counter_new() -> Callable[..., int]: """ Lazily load and cache the token_counter function (aliased as token_counter_new). @@ -154,10 +165,10 @@ def _get_token_counter_new() -> Any: # This registry maps attribute names (like "ModelResponse") to handler functions # It's built once the first time someone accesses a lazy-loaded attribute # Example: {"ModelResponse": _lazy_import_utils, "Cache": _lazy_import_caching, ...} -_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], Any]] | None = None +_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], object]] | None = None -def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: +def _get_lazy_import_registry() -> dict[str, Callable[[str], object]]: """ Build the registry that maps attribute names to their handler functions. @@ -206,7 +217,7 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: return _LAZY_IMPORT_REGISTRY -def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> Any: +def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> object: """ Generic function that handles lazy importing for most attributes. @@ -255,7 +266,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate # Step 6: Get the actual attribute from the module # Example: getattr(utils_module, "ModelResponse") returns the ModelResponse class - value: Final = getattr(module, attr_name) + value: Final[object] = getattr(module, attr_name) # Step 7: Cache it so we don't have to import again next time _globals[name] = value @@ -272,62 +283,62 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate # The registry (above) maps attribute names to these handler functions. -def _lazy_import_utils(name: str) -> Any: +def _lazy_import_utils(name: str) -> object: """Handler for utils module attributes (ModelResponse, token_counter, etc.)""" return _generic_lazy_import(name, _UTILS_IMPORT_MAP, "Utils") -def _lazy_import_cost_calculator(name: str) -> Any: +def _lazy_import_cost_calculator(name: str) -> object: """Handler for cost calculator functions (completion_cost, cost_per_token, etc.)""" return _generic_lazy_import(name, _COST_CALCULATOR_IMPORT_MAP, "Cost calculator") -def _lazy_import_token_counter(name: str) -> Any: +def _lazy_import_token_counter(name: str) -> object: """Handler for token counter utilities""" return _generic_lazy_import(name, _TOKEN_COUNTER_IMPORT_MAP, "Token counter") -def _lazy_import_bedrock_types(name: str) -> Any: +def _lazy_import_bedrock_types(name: str) -> object: """Handler for Bedrock type aliases""" return _generic_lazy_import(name, _BEDROCK_TYPES_IMPORT_MAP, "Bedrock types") -def _lazy_import_types_utils(name: str) -> Any: +def _lazy_import_types_utils(name: str) -> object: """Handler for types from litellm.types.utils (BudgetConfig, ImageObject, etc.)""" return _generic_lazy_import(name, _TYPES_UTILS_IMPORT_MAP, "Types utils") -def _lazy_import_caching(name: str) -> Any: +def _lazy_import_caching(name: str) -> object: """Handler for caching classes (Cache, DualCache, RedisCache, etc.)""" return _generic_lazy_import(name, _CACHING_IMPORT_MAP, "Caching") -def _lazy_import_dotprompt(name: str) -> Any: +def _lazy_import_dotprompt(name: str) -> object: """Handler for dotprompt integration globals""" return _generic_lazy_import(name, _DOTPROMPT_IMPORT_MAP, "Dotprompt") -def _lazy_import_types(name: str) -> Any: +def _lazy_import_types(name: str) -> object: """Handler for type classes (GuardrailItem, etc.)""" return _generic_lazy_import(name, _TYPES_IMPORT_MAP, "Types") -def _lazy_import_llm_configs(name: str) -> Any: +def _lazy_import_llm_configs(name: str) -> object: """Handler for LLM config classes (AnthropicConfig, OpenAILikeChatConfig, etc.)""" return _generic_lazy_import(name, _LLM_CONFIGS_IMPORT_MAP, "LLM config") -def _lazy_import_litellm_logging(name: str) -> Any: +def _lazy_import_litellm_logging(name: str) -> object: """Handler for litellm_logging module (Logging, modify_integration)""" return _generic_lazy_import(name, _LITELLM_LOGGING_IMPORT_MAP, "Litellm logging") -def _lazy_import_llm_provider_logic(name: str) -> Any: +def _lazy_import_llm_provider_logic(name: str) -> object: """Handler for LLM provider logic functions (get_llm_provider, etc.)""" return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic") -def _lazy_import_utils_module(name: str) -> Any: +def _lazy_import_utils_module(name: str) -> object: """ Handler for utils module lazy imports. @@ -355,7 +366,7 @@ def _lazy_import_utils_module(name: str) -> Any: module = importlib.import_module(module_path) # Get the actual attribute from the module - value: Final = getattr(module, attr_name) + value: Final[object] = getattr(module, attr_name) # Cache it so we don't have to import again next time _globals[name] = value @@ -370,7 +381,7 @@ def _lazy_import_utils_module(name: str) -> Any: # These handlers have custom logic that doesn't fit the generic pattern -def _lazy_import_llm_client_cache(name: str) -> Any: +def _lazy_import_llm_client_cache(name: str) -> object: """ Handler for LLM client cache - has special logic for singleton instance. @@ -387,7 +398,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any: # Import the class module: Final = importlib.import_module("litellm.caching.llm_caching_handler") - LLMClientCache: Final = getattr(module, "LLMClientCache") + LLMClientCache: Final[type[LLMClientCacheType]] = getattr(module, "LLMClientCache") # If they want the class itself, return it if name == "LLMClientCache": @@ -403,7 +414,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any: raise AttributeError(f"LLM client cache lazy import: unknown attribute {name!r}") -def _lazy_import_http_handlers(name: str) -> Any: +def _lazy_import_http_handlers(name: str) -> object: """ Handler for HTTP clients - has special logic for creating client instances. @@ -419,8 +430,8 @@ def _lazy_import_http_handlers(name: str) -> Any: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client # Get timeout from module config (if set) - timeout = _globals.get("request_timeout") - params: Final = {"timeout": timeout, "client_alias": "module level aclient"} + async_timeout: Final = _get_module_level_client_timeout(_globals) + params: Final = {"timeout": async_timeout, "client_alias": "module level aclient"} # Create the client instance provider_id: Final = cast(Any, "litellm_module_level_client") @@ -437,8 +448,8 @@ def _lazy_import_http_handlers(name: str) -> Any: # Create a sync HTTP client from litellm.llms.custom_httpx.http_handler import HTTPHandler - timeout = _globals.get("request_timeout") - sync_client: Final = HTTPHandler(timeout=timeout) + sync_timeout: Final = _get_module_level_client_timeout(_globals) + sync_client: Final = HTTPHandler(timeout=sync_timeout) # Cache it _globals["module_level_client"] = sync_client diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index 15cf77708f9..90afa5adf9e 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -17,12 +17,16 @@ A2A Streaming Events: - Artifact update (kind: "artifact-update") - Content/artifact delivery """ +from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final from uuid import uuid4 from litellm._logging import verbose_logger +if TYPE_CHECKING: + from litellm.types.utils import Choices + class A2AStreamingContext: """ @@ -30,7 +34,7 @@ class A2AStreamingContext: Tracks task_id, context_id, and message accumulation. """ - def __init__(self, request_id: str, input_message: dict[str, Any]): + def __init__(self, request_id: str, input_message: Mapping[str, object]): self.request_id = request_id self.task_id = str(uuid4()) self.context_id = str(uuid4()) @@ -46,7 +50,7 @@ class A2ACompletionBridgeTransformation: """ @staticmethod - def _extract_text_from_a2a_parts(parts: list[dict[str, Any]]) -> str: + def _extract_text_from_a2a_parts(parts: Sequence[Mapping[str, object]]) -> str: """Extract text from A2A parts (with or without explicit ``kind``).""" content_parts: Final[list[str]] = [] for part in parts: @@ -62,16 +66,16 @@ class A2ACompletionBridgeTransformation: @staticmethod def get_forward_metadata( - a2a_message: dict[str, Any], + a2a_message: Mapping[str, object], params: dict[str, Any] | None = None, - ) -> dict[str, Any] | None: + ) -> dict[str, object] | None: """ Merge A2A metadata from MessageSendParams and the message for downstream providers. Forwarded once on the LangGraph run payload (``metadata``), not duplicated on each input message — see ``apply_forward_metadata_to_completion_params``. """ - merged: Final[dict[str, Any]] = {} + merged: Final[dict[str, object]] = {} if params and isinstance(params.get("metadata"), dict): merged.update(params["metadata"]) message_metadata: Final = a2a_message.get("metadata") @@ -81,8 +85,8 @@ class A2ACompletionBridgeTransformation: @staticmethod def apply_forward_metadata_to_completion_params( - completion_params: dict[str, Any], - a2a_message: dict[str, Any], + completion_params: dict[str, object], + a2a_message: Mapping[str, object], params: dict[str, Any] | None = None, ) -> None: """ @@ -104,8 +108,8 @@ class A2ACompletionBridgeTransformation: # ``extra_body.metadata`` so the configured keys remain authoritative # and an A2A caller cannot overwrite server-set run metadata. existing_metadata: Final = extra_body.get("metadata") - existing_dict: Final[dict[str, Any]] = existing_metadata if isinstance(existing_metadata, dict) else {} - merged_metadata: Final[dict[str, Any]] = {**forward_metadata, **existing_dict} + existing_dict: Final[dict[str, object]] = existing_metadata if isinstance(existing_metadata, dict) else {} + merged_metadata: Final[dict[str, object]] = {**forward_metadata, **existing_dict} extra_body = {**extra_body, "metadata": merged_metadata} completion_params["extra_body"] = extra_body @@ -114,7 +118,7 @@ class A2ACompletionBridgeTransformation: @staticmethod def a2a_message_to_openai_messages( a2a_message: dict[str, Any], - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Transform an A2A message to OpenAI message format. @@ -124,8 +128,8 @@ class A2ACompletionBridgeTransformation: Returns: List of OpenAI-format messages """ - role: Final = a2a_message.get("role", "user") - parts = a2a_message.get("parts", []) + role: Final[object] = a2a_message.get("role", "user") + parts: Sequence[Mapping[str, object]] = a2a_message.get("parts", []) # Map A2A roles to OpenAI roles openai_role = role @@ -143,7 +147,7 @@ class A2ACompletionBridgeTransformation: # Do not attach A2A message.metadata here — the completion bridge forwards it # once at run level via extra_body.metadata (LangGraph POST /runs/wait shape). - openai_message: Final[dict[str, Any]] = {"role": openai_role, "content": content} + openai_message: Final[dict[str, object]] = {"role": openai_role, "content": content} verbose_logger.debug( "A2A -> OpenAI transform: role=%s -> %s, content_length=%s", role, openai_role, len(content) @@ -155,7 +159,7 @@ class A2ACompletionBridgeTransformation: def openai_response_to_a2a_response( response: Any, request_id: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform a LiteLLM ModelResponse to A2A SendMessageResponse format. @@ -169,7 +173,7 @@ class A2ACompletionBridgeTransformation: # Extract content from response content = "" if hasattr(response, "choices") and response.choices: - choice: Final = response.choices[0] + choice: Final[Choices] = response.choices[0] if hasattr(choice, "message") and choice.message: content = choice.message.content or "" @@ -182,7 +186,7 @@ class A2ACompletionBridgeTransformation: } # Build A2A response - a2a_response: Final = { + a2a_response: Final[dict[str, object]] = { "jsonrpc": "2.0", "id": request_id, "result": a2a_message, @@ -200,7 +204,7 @@ class A2ACompletionBridgeTransformation: @staticmethod def create_task_event( ctx: A2AStreamingContext, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Create the initial task event with status 'submitted'. @@ -235,7 +239,7 @@ class A2ACompletionBridgeTransformation: state: str, final: bool = False, message_text: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Create a status update event. @@ -245,7 +249,7 @@ class A2ACompletionBridgeTransformation: final: Whether this is the final event message_text: Optional message text for 'working' status """ - status: Final[dict[str, Any]] = { + status: Final[dict[str, object]] = { "state": state, "timestamp": A2ACompletionBridgeTransformation._get_timestamp(), } @@ -277,7 +281,7 @@ class A2ACompletionBridgeTransformation: def create_artifact_update_event( ctx: A2AStreamingContext, text: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Create an artifact update event with content. diff --git a/litellm/caching/valkey_semantic_cache.py b/litellm/caching/valkey_semantic_cache.py index 737d212a89d..ac4d4033546 100644 --- a/litellm/caching/valkey_semantic_cache.py +++ b/litellm/caching/valkey_semantic_cache.py @@ -17,6 +17,7 @@ RedisSemanticCache since those are backend agnostic. import asyncio import hashlib import os +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Any, Final @@ -62,7 +63,7 @@ class ValkeySemanticCache(RedisSemanticCache): sync_client: Redis | None = None, async_client: AsyncRedis | None = None, embedding_max_input_tokens: int | None = None, - **kwargs: Any, + **kwargs: object, ): if similarity_threshold is None: raise ValueError("similarity_threshold must be provided, passed None") @@ -115,7 +116,7 @@ class ValkeySemanticCache(RedisSemanticCache): return hashlib.sha256(str(key).encode("utf-8")).hexdigest() @staticmethod - def _embedding_to_bytes(embedding: list[float]) -> bytes: + def _embedding_to_bytes(embedding: Sequence[float]) -> bytes: return pack_vector(embedding) def _index_schema(self, dim: int) -> tuple[TagField, VectorField]: @@ -189,7 +190,9 @@ class ValkeySemanticCache(RedisSemanticCache): def _doc_key(self, key: str) -> str: return f"{self.key_prefix}{self._scope_tag(key)}:{uuid.uuid4()}" - def _doc_mapping(self, key: str, prompt: str, value_str: str, embedding: list[float]) -> dict: + def _doc_mapping( + self, key: str, prompt: str, value_str: str, embedding: Sequence[float] + ) -> Mapping[str | bytes, str | bytes]: return { self.CACHE_KEY_FIELD_NAME: self._scope_tag(key), self.PROMPT_FIELD_NAME: prompt, @@ -205,30 +208,49 @@ class ValkeySemanticCache(RedisSemanticCache): ) return Query(query_string).return_fields(self.RESPONSE_FIELD_NAME, self.DISTANCE_FIELD_NAME).dialect(2) + async def _async_search(self, key: str, embedding: Sequence[float]) -> object: + """Run the KNN query on the async client, stopping the untyped search surface here.""" + return await self.async_client.ft(self.index_name).search( + self._knn_query(key), + query_params={"vec": self._embedding_to_bytes(embedding)}, + ) + @classmethod - def _first_hit(cls, search_result: Any) -> _ValkeyCacheHit | None: - docs: Final = getattr(search_result, "docs", []) + def _first_hit(cls, search_result: object) -> _ValkeyCacheHit | None: + docs: Final[Sequence[object]] = getattr(search_result, "docs", []) if not docs: return None doc: Final = docs[0] + response_field: Final[object] = getattr(doc, cls.RESPONSE_FIELD_NAME) + distance_field: Final[str | bytes | float] = getattr(doc, cls.DISTANCE_FIELD_NAME) return _ValkeyCacheHit( - response=str(getattr(doc, cls.RESPONSE_FIELD_NAME)), - distance=float(getattr(doc, cls.DISTANCE_FIELD_NAME)), + response=str(response_field), + distance=float(distance_field), ) - def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: Any) -> Any: + @staticmethod + def _record_similarity(kwargs: dict[str, Any], similarity: float) -> None: + """Stamp the semantic-similarity score onto the request metadata carried in ``kwargs``.""" + kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity + + @staticmethod + def _embedding_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None: + """The request metadata forwarded to the embedding call.""" + return kwargs.get("metadata") + + def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: object) -> object: if hit is None: - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) return None similarity: Final = 1 - hit.distance - kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity + self._record_similarity(kwargs, similarity) if similarity < self.similarity_threshold: return None return self._get_cache_logic(cached_response=hit.response) - def set_cache(self, key: str, value: Any, **kwargs: Any) -> None: + def set_cache(self, key: str, value: object, **kwargs: object) -> None: print_verbose(f"Valkey semantic-cache set_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) @@ -247,12 +269,12 @@ class ValkeySemanticCache(RedisSemanticCache): except Exception as e: print_verbose(f"Error in Valkey semantic-cache set_cache: {e}") - def get_cache(self, key: str, **kwargs: Any) -> Any: + def get_cache(self, key: str, **kwargs: object) -> object: print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) return None embedding: Final = self._get_embedding(prompt) @@ -265,9 +287,9 @@ class ValkeySemanticCache(RedisSemanticCache): return self._resolve_hit(self._first_hit(search_result), key, **kwargs) except Exception as e: print_verbose(f"Error in Valkey semantic-cache get_cache: {e}") - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) - async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None: + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: print_verbose(f"Async Valkey semantic-cache set_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) @@ -275,7 +297,7 @@ class ValkeySemanticCache(RedisSemanticCache): print_verbose("No prompt provided for semantic caching") return - embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs)) await self._ensure_index_async(len(embedding)) doc_key: Final = self._doc_key(key) @@ -286,31 +308,28 @@ class ValkeySemanticCache(RedisSemanticCache): except Exception as e: print_verbose(f"Error in async Valkey semantic-cache set_cache: {e}") - async def async_get_cache(self, key: str, **kwargs: Any) -> Any: + async def async_get_cache(self, key: str, **kwargs: object) -> object: print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) return None - embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs)) await self._ensure_index_async(len(embedding)) - search_result: Final = await self.async_client.ft(self.index_name).search( - self._knn_query(key), - query_params={"vec": self._embedding_to_bytes(embedding)}, - ) + search_result: Final[object] = await self._async_search(key, embedding) return self._resolve_hit(self._first_hit(search_result), key, **kwargs) except Exception as e: print_verbose(f"Error in async Valkey semantic-cache get_cache: {e}") - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) - async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: Any) -> None: + async def async_set_cache_pipeline(self, cache_list: list[tuple[str, object]], **kwargs: object) -> None: try: await asyncio.gather(*[self.async_set_cache(key, value, **kwargs) for key, value in cache_list]) except Exception as e: print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e}") - async def _index_info(self) -> dict: + async def _index_info(self) -> Mapping[str, object]: return await self.async_client.ft(self.index_name).info() diff --git a/litellm/integrations/bitbucket/bitbucket_client.py b/litellm/integrations/bitbucket/bitbucket_client.py index e06e5ab358f..c6967ef3340 100644 --- a/litellm/integrations/bitbucket/bitbucket_client.py +++ b/litellm/integrations/bitbucket/bitbucket_client.py @@ -4,11 +4,33 @@ BitBucket API client for fetching .prompt files from BitBucket repositories. import base64 import urllib.parse -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final + +from typing_extensions import ReadOnly, TypedDict from litellm.llms.custom_httpx.http_handler import HTTPHandler +class BitBucketSrcEntry(TypedDict, total=False): + """One entry of a BitBucket ``src`` directory listing.""" + + type: ReadOnly[str] + path: ReadOnly[str] + + +class BitBucketSrcListing(TypedDict, total=False): + """A page of a BitBucket ``src`` directory listing.""" + + values: ReadOnly[Sequence[BitBucketSrcEntry]] + + +class BitBucketBranchListing(TypedDict, total=False): + """A page of a BitBucket ``refs/branches`` listing.""" + + values: ReadOnly[Sequence[Mapping[str, object]]] + + def _sanitize_file_path(file_path: str) -> str: """Reject path traversal and URL-encode each path segment.""" if "#" in file_path or "?" in file_path: @@ -31,7 +53,7 @@ class BitBucketClient: - Branch-specific file fetching """ - def __init__(self, config: dict[str, Any]): + def __init__(self, config: Mapping[str, object]): """ Initialize the BitBucket client. @@ -135,16 +157,13 @@ class BitBucketClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final = response.json() - files: Final = [] + data: Final[BitBucketSrcListing] = response.json() - for item in data.get("values", []): - if item.get("type") == "commit_file": - file_path = item.get("path", "") - if file_path.endswith(file_extension): - files.append(file_path) - - return files + return [ + file_path + for item in data.get("values", []) + if item.get("type") == "commit_file" and (file_path := item.get("path", "")).endswith(file_extension) + ] except Exception as e: # Check if it's an HTTP error @@ -162,7 +181,7 @@ class BitBucketClient: else: raise Exception(f"Error listing files in '{directory_path}': {e}") - def get_repository_info(self) -> dict[str, Any]: + def get_repository_info(self) -> Mapping[str, object]: """ Get information about the repository. @@ -191,7 +210,7 @@ class BitBucketClient: except Exception: return False - def get_branches(self) -> list[dict[str, Any]]: + def get_branches(self) -> Sequence[Mapping[str, object]]: """ Get list of branches in the repository. @@ -204,12 +223,12 @@ class BitBucketClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final = response.json() + data: Final[BitBucketBranchListing] = response.json() return data.get("values", []) except Exception as e: raise Exception(f"Failed to get branches: {e}") - def get_file_metadata(self, file_path: str) -> dict[str, Any] | None: + def get_file_metadata(self, file_path: str) -> Mapping[str, object] | None: """ Get metadata about a file (size, last modified, etc.). diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 7ea60053e6f..219718f6771 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -7,7 +7,10 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan. import time import uuid -from typing import Any, Final, cast +from collections.abc import Mapping, Sequence +from typing import Any, Final, Protocol, cast + +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.compression import compress @@ -26,6 +29,19 @@ LITELLM_CONTENT_RETRIEVE_TOOL_NAME: Final = "litellm_content_retrieve" _CACHE_TTL_SECONDS: Final = 15 * 60 +class _AgenticLoopParams(TypedDict, total=False): + """The ``agentic_loop_params`` entry the agentic loop driver records on the logging object.""" + + model: ReadOnly[str] + + +class _AgenticLoopLoggingObj(Protocol): + """Logging object view exposing the untyped call details this handler reads.""" + + @property + def model_call_details(self) -> Mapping[str, _AgenticLoopParams]: ... + + def _compression_savings_from_counts( original_tokens: object, compressed_tokens: object ) -> CompressionSavingsMetadata | None: @@ -78,7 +94,7 @@ class CompressionInterceptionLogger(CustomLogger): compression_trigger: int = 200_000, compression_target: int | None = None, embedding_model: str | None = None, - embedding_model_params: dict[str, Any] | None = None, + embedding_model_params: dict[str, object] | None = None, ): super().__init__() self.enabled = enabled @@ -101,7 +117,7 @@ class CompressionInterceptionLogger(CustomLogger): @staticmethod def initialize_from_proxy_config( litellm_settings: dict[str, Any], - callback_specific_params: dict[str, Any], + callback_specific_params: Mapping[str, object], ) -> "CompressionInterceptionLogger": compression_params: CompressionInterceptionConfig = {} if "compression_interception_params" in litellm_settings: @@ -115,7 +131,9 @@ class CompressionInterceptionLogger(CustomLogger): ) return CompressionInterceptionLogger.from_config_yaml(compression_params) - async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None: + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, Any], call_type: CallTypes | None + ) -> dict[str, object] | None: if not self.enabled: return None if call_type is not None and call_type != CallTypes.anthropic_messages: @@ -145,7 +163,7 @@ class CompressionInterceptionLogger(CustomLogger): cache: Final = cast(dict[str, str], compressed.get("cache", {})) skip_reason: Final = cast(str | None, compressed.get("compression_skipped_reason")) - compressed_tools: Final = cast(list[dict[str, Any]], compressed.get("tools", [])) + compressed_tools: Final = cast(list[dict[str, object]], compressed.get("tools", [])) # Only mutate kwargs when compression actually produced a result. # If compression was a no-op (below trigger, invalid tool sequence, etc.), @@ -156,7 +174,7 @@ class CompressionInterceptionLogger(CustomLogger): kwargs["messages"] = compressed["messages"] if compressed_tools: kwargs["tools"] = self._merge_tools( - existing_tools=cast(list[dict[str, Any]] | None, kwargs.get("tools")), + existing_tools=cast(list[dict[str, object]] | None, kwargs.get("tools")), compressed_tools=compressed_tools, ) call_id = cast(str | None, kwargs.get("litellm_call_id")) @@ -189,14 +207,14 @@ class CompressionInterceptionLogger(CustomLogger): async def async_should_run_agentic_loop( self, - response: Any, + response: object, model: str, - messages: list[dict], - tools: list[dict] | None, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, stream: bool, custom_llm_provider: str, - kwargs: dict, - ) -> tuple[bool, dict]: + kwargs: Mapping[str, object], + ) -> tuple[bool, dict[str, object]]: if not self.enabled: return False, {} if not self._has_retrieval_tool(tools): @@ -214,19 +232,19 @@ class CompressionInterceptionLogger(CustomLogger): async def async_build_agentic_loop_plan( self, - tools: dict, + tools: Mapping[str, object], model: str, - messages: list[dict], - response: Any, - anthropic_messages_provider_config: Any, - anthropic_messages_optional_request_params: dict, - logging_obj: Any, + messages: list[dict[str, object]], + response: object, + anthropic_messages_provider_config: object, + anthropic_messages_optional_request_params: Mapping[str, object], + logging_obj: _AgenticLoopLoggingObj | None, stream: bool, - kwargs: dict, + kwargs: Mapping[str, object], ) -> AgenticLoopPlan: self._prune_expired_cache() - tool_calls: Final = cast(list[dict[str, Any]], tools.get("tool_calls", [])) - thinking_blocks: Final = cast(list[dict[str, Any]], tools.get("thinking_blocks", [])) + tool_calls: Final = cast(list[dict[str, object]], tools.get("tool_calls", [])) + thinking_blocks: Final = cast(list[dict[str, object]], tools.get("thinking_blocks", [])) call_id: Final = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs) cache: Final = self._get_cache(call_id=call_id) @@ -269,7 +287,7 @@ class CompressionInterceptionLogger(CustomLogger): full_model_name = model if logging_obj is not None: agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {}) - full_model_name = cast(str, agentic_params.get("model", model)) + full_model_name = agentic_params.get("model", model) request_patch: Final = AgenticLoopRequestPatch( model=full_model_name, @@ -304,15 +322,15 @@ class CompressionInterceptionLogger(CustomLogger): return {} return cache_entry[0] - def _resolve_call_id(self, logging_obj: Any, kwargs: dict[str, Any]) -> str | None: + def _resolve_call_id(self, logging_obj: _AgenticLoopLoggingObj | None, kwargs: Mapping[str, object]) -> str | None: if logging_obj is not None: logging_call_id: Final = getattr(logging_obj, "litellm_call_id", None) if isinstance(logging_call_id, str) and logging_call_id: return logging_call_id kwargs_call_id: Final = kwargs.get("litellm_call_id") - return cast(str | None, kwargs_call_id if isinstance(kwargs_call_id, str) else None) + return kwargs_call_id if isinstance(kwargs_call_id, str) else None - def _resolve_retrieval_content(self, tool_call: dict[str, Any], cache: dict[str, str]) -> str: + def _resolve_retrieval_content(self, tool_call: Mapping[str, object], cache: Mapping[str, str]) -> str: raw_input: Final = tool_call.get("input", {}) key = "" if isinstance(raw_input, dict): @@ -323,7 +341,9 @@ class CompressionInterceptionLogger(CustomLogger): return cache[key] return f"[compressed content key '{key}' not found]" - def _extract_retrieval_tool_calls(self, response: Any) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + def _extract_retrieval_tool_calls( + self, response: object + ) -> tuple[list[dict[str, object]], list[dict[str, object]]]: if isinstance(response, dict): content = response.get("content", []) else: @@ -332,8 +352,8 @@ class CompressionInterceptionLogger(CustomLogger): if not isinstance(content, list): return [], [] - tool_calls: Final[list[dict[str, Any]]] = [] - thinking_blocks: Final[list[dict[str, Any]]] = [] + tool_calls: Final[list[dict[str, object]]] = [] + thinking_blocks: Final[list[dict[str, object]]] = [] for block in content: if isinstance(block, dict): @@ -380,13 +400,13 @@ class CompressionInterceptionLogger(CustomLogger): return tool_calls, thinking_blocks - def _prepare_followup_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: + def _prepare_followup_kwargs(self, kwargs: Mapping[str, object]) -> dict[str, object]: internal_keys: Final = {"litellm_logging_obj"} return { k: v for k, v in kwargs.items() if not k.startswith("_compression_interception") and k not in internal_keys } - def _has_retrieval_tool(self, tools: Any) -> bool: + def _has_retrieval_tool(self, tools: object) -> bool: if not isinstance(tools, list): return False for tool in tools: @@ -402,9 +422,9 @@ class CompressionInterceptionLogger(CustomLogger): def _merge_tools( self, - existing_tools: list[dict[str, Any]] | None, - compressed_tools: list[dict[str, Any]], - ) -> list[dict[str, Any]]: + existing_tools: Sequence[Mapping[str, object]] | None, + compressed_tools: Sequence[Mapping[str, object]], + ) -> list[Mapping[str, object]]: merged: Final = list(existing_tools or []) if self._has_retrieval_tool(merged): return merged diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a0c78674ac8..417edc77e9c 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -2,7 +2,7 @@ # On success, logs events to Promptlayer import re import traceback -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Sequence from typing import TYPE_CHECKING, Any, Final, Optional from pydantic import BaseModel @@ -103,11 +103,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return [] callbacks: Final = AllCallbacks() - callback_info: Final = getattr(callbacks, lookup_name, None) + callback_info: Final[object] = getattr(callbacks, lookup_name, None) if callback_info is None: return [] - params: Final = getattr(callback_info, "litellm_callback_params", None) + params: Final[Sequence[str] | None] = getattr(callback_info, "litellm_callback_params", None) if not params: return [] @@ -783,7 +783,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac - Converting to string and then truncating the logged content catches this 2. We want to avoid modifying the original `messages`, `response`, and `error_str` in the logging payload since these are in kwargs and could be returned to the user """ - field_value: Final = standard_logging_object.get(field_name) + field_value: Final[object] = standard_logging_object.get(field_name) if field_value: str_value: Final = str(field_value) if len(str_value) > max_length: @@ -937,8 +937,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac • Keep untyped or text content. • Recursively redact inline base64 blobs in *any* string field, at any depth. """ - raw_messages: Final[Any] = payload.get("messages", []) - messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else [] + raw_messages: Final[object] = payload.get("messages", []) + messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else [] verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages)) if messages: @@ -969,8 +969,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac • Keep untyped or text content. • Recursively redact inline base64 blobs in *any* string field, at any depth. """ - raw_messages: Final[Any] = payload.get("messages", []) - messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else [] + raw_messages: Final[object] = payload.get("messages", []) + messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else [] verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages)) if messages: @@ -991,7 +991,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac value: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, - ) -> Any: + ) -> object: """Recursively redact inline base64 from any nested structure with a max recursion depth limit.""" if depth > max_depth: verbose_logger.warning("[CustomLogger] Max recursion depth %s reached while redacting base64", max_depth) @@ -1022,16 +1022,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac def _process_messages( self, - messages: list[Any], + messages: list[object], max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, - ) -> list[dict[str, Any]]: - filtered_messages: Final[list[dict[str, Any]]] = [] + ) -> list[dict[str, object]]: + filtered_messages: Final[list[dict[str, object]]] = [] for msg in messages: if not isinstance(msg, dict): continue - contents: Any = msg.get("content") + contents: object = msg.get("content") if isinstance(contents, list): - cleaned: list[Any] = [] + cleaned: list[object] = [] for c in contents: if self._should_keep_content(content=c): cleaned.append(self._redact_base64(value=c, max_depth=max_depth)) diff --git a/litellm/integrations/opik/opik_payload_builder/payload_builders.py b/litellm/integrations/opik/opik_payload_builder/payload_builders.py index e40d72ea542..855b84ba4c8 100644 --- a/litellm/integrations/opik/opik_payload_builder/payload_builders.py +++ b/litellm/integrations/opik/opik_payload_builder/payload_builders.py @@ -17,12 +17,12 @@ def build_trace_payload( end_time: datetime, input_data: Any, output_data: Any, - metadata: dict[str, Any], + metadata: dict[str, object], tags: list[str], thread_id: str | None, ) -> types.TracePayload: """Build a complete trace payload.""" - trace_name: Final = response_obj.get("object", "unknown type") + trace_name: Final[str] = response_obj.get("object", "unknown type") return types.TracePayload( project_name=project_name, @@ -47,7 +47,7 @@ def build_span_payload( end_time: datetime, input_data: Any, output_data: Any, - metadata: dict[str, Any], + metadata: dict[str, object], tags: list[str], usage: dict[str, int], provider: str | None = None, @@ -56,9 +56,9 @@ def build_span_payload( """Build a complete span payload.""" span_id: Final = utils.create_uuid7() - model: Final = response_obj.get("model", "unknown-model") - obj_type: Final = response_obj.get("object", "unknown-object") - created: Final = response_obj.get("created", 0) + model: Final[str] = response_obj.get("model", "unknown-model") + obj_type: Final[str] = response_obj.get("object", "unknown-object") + created: Final[int] = response_obj.get("created", 0) span_name: Final = f"{model}_{obj_type}_{created}" _logging.verbose_logger.debug("OpikLogger creating span with id %s for trace %s", span_id, trace_id) diff --git a/litellm/integrations/prometheus_helpers/prometheus_api.py b/litellm/integrations/prometheus_helpers/prometheus_api.py index 9f77f87a670..f677db17648 100644 --- a/litellm/integrations/prometheus_helpers/prometheus_api.py +++ b/litellm/integrations/prometheus_helpers/prometheus_api.py @@ -4,9 +4,13 @@ Helper functions to query prometheus API import json import time +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta from typing import Final +from httpx import Response +from typing_extensions import ReadOnly, TypedDict + from litellm import get_secret from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -19,9 +23,45 @@ PROMETHEUS_SELECTED_INSTANCE: Final[str | None] = get_secret("PROMETHEUS_SELECTE async_http_handler: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) +class _PrometheusSeries(TypedDict): + """One time series in a Prometheus query result, with its labels and its ``[timestamp, value]`` samples.""" + + metric: ReadOnly[Mapping[str, str]] + values: ReadOnly[Sequence[Sequence[float | str]]] + + +class _PrometheusResultData(TypedDict): + """The ``data`` envelope of a Prometheus query response.""" + + result: ReadOnly[Sequence[_PrometheusSeries]] + + +class _PrometheusQueryResponse(TypedDict): + """The JSON body returned by the Prometheus ``/api/v1/query`` and ``/api/v1/query_range`` endpoints.""" + + data: ReadOnly[_PrometheusResultData] + + +class _DailySpend(TypedDict): + """One day of spend, in the shape ``get_daily_spend_from_prometheus`` returns.""" + + date: ReadOnly[str] + spend: ReadOnly[float] + + +def _query_body(response: Response) -> _PrometheusQueryResponse: + """Read the untyped JSON body of a Prometheus query response.""" + return response.json() + + +def _query_series(response: Response) -> Sequence[_PrometheusSeries]: + """Read the time series list out of the untyped JSON body of a Prometheus query response.""" + return response.json()["data"]["result"] + + async def get_metric_from_prometheus( metric_name: str, -): +) -> Sequence[_PrometheusSeries]: # Get the start of the current day in Unix timestamp if PROMETHEUS_URL is None: raise ValueError("PROMETHEUS_URL not set please set 'PROMETHEUS_URL=<>' in .env") @@ -31,13 +71,13 @@ async def get_metric_from_prometheus( response: Final = await async_http_handler.get( f"{PROMETHEUS_URL}/api/v1/query", params={"query": query, "time": now} ) # End of the day - _json_response: Final = response.json() + _json_response: Final = _query_body(response) verbose_logger.debug("json response from prometheus /query api %s", _json_response) - results: Final = response.json()["data"]["result"] + results: Final = _query_series(response) return results -async def get_fallback_metric_from_prometheus(): +async def get_fallback_metric_from_prometheus() -> str: """ Gets fallback metrics from prometheus for the last 24 hours """ @@ -96,7 +136,7 @@ def _quote_promql_string_literal(value: str) -> str: return json.dumps(value, ensure_ascii=False) -async def get_daily_spend_from_prometheus(api_key: str | None): +async def get_daily_spend_from_prometheus(api_key: str | None) -> Sequence[_DailySpend]: """ Expected Response Format: [ @@ -133,9 +173,9 @@ async def get_daily_spend_from_prometheus(api_key: str | None): } response: Final = await async_http_handler.get(url, params=params) - _json_response: Final = response.json() + _json_response: Final = _query_body(response) verbose_logger.debug("json response from prometheus /query api %s", _json_response) - results: Final = response.json()["data"]["result"] + results: Final = _query_series(response) formatted_results: Final = [] for result in results: diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index e59ef0449d0..e1ec3360351 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -10,7 +10,7 @@ import asyncio import math import uuid from collections.abc import AsyncIterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, TypedDict, TypeVar, cast from typing_extensions import ReadOnly @@ -74,6 +74,10 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b # ``web_search_tool_result`` blocks to inject into the final response. WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks" +_RESPONSE_CONTENT_FIELD: Final = "content" + +_ResponseT: Final = TypeVar("_ResponseT") + class _PlanMetadataView(TypedDict): websearch_native_blocks: Sequence[Mapping[str, object]] | None @@ -926,17 +930,17 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any: + def _inject_native_blocks(response: _ResponseT, native_blocks: Sequence[Mapping[str, object]]) -> _ResponseT: """Prepend native blocks to response content, dict or object form.""" if not native_blocks: return response if isinstance(response, dict): - existing = response.get("content") or [] - response["content"] = list(native_blocks) + list(existing) + existing = response.get(_RESPONSE_CONTENT_FIELD) or [] + response[_RESPONSE_CONTENT_FIELD] = list(native_blocks) + list(existing) return response - existing = getattr(response, "content", None) or [] + existing = getattr(response, _RESPONSE_CONTENT_FIELD, None) or [] try: - response.content = list(native_blocks) + list(existing) + setattr(response, _RESPONSE_CONTENT_FIELD, list(native_blocks) + list(existing)) except (AttributeError, TypeError): # Object refused write — fall through and leave the response # untouched rather than crash the request. diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index 44fed944d2a..1737e2b8cb0 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -1,6 +1,9 @@ import datetime +from collections.abc import Mapping from typing import Any, Final +import httpx + from litellm.constants import LITELLM_DETAILED_TIMING from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base @@ -25,11 +28,7 @@ class ResponseMetadata: @property def supports_response_time(self) -> bool: """Check if response type supports timing metrics""" - return ( - isinstance(self.result, ModelResponse) - or isinstance(self.result, EmbeddingResponse) - or isinstance(self.result, TranscriptionResponse) - ) + return isinstance(self.result, (ModelResponse, EmbeddingResponse, TranscriptionResponse)) def set_hidden_params(self, logging_obj: LiteLLMLoggingObject, model: str | None, kwargs: dict) -> None: """Set hidden parameters on the response""" @@ -45,14 +44,14 @@ class ResponseMetadata: result=self.result, litellm_model_name=model, router_model_id=model_id ), "additional_headers": process_response_headers( - self._get_value_from_hidden_params("additional_headers") or {}, + self._get_additional_headers_from_hidden_params() or {}, preserve_litellm_internal_headers=True, ), "litellm_model_name": model, } self._update_hidden_params(new_params) - def _update_hidden_params(self, new_params: dict) -> None: + def _update_hidden_params(self, new_params: Mapping[str, object]) -> None: """ Update hidden params - handles when self._hidden_params is a dict or HiddenParams object """ @@ -64,12 +63,12 @@ class ResponseMetadata: for key, value in new_params.items(): setattr(self._hidden_params, key, value) - def _get_value_from_hidden_params(self, key: str) -> Any | None: - """Get value from hidden params - handles when self._hidden_params is a dict or HiddenParams object""" + def _get_additional_headers_from_hidden_params(self) -> httpx.Headers | dict[str, str] | None: + """Get `additional_headers` from hidden params - handles when self._hidden_params is a dict or HiddenParams object""" if isinstance(self._hidden_params, dict): - return self._hidden_params.get(key, None) + return self._hidden_params.get("additional_headers", None) elif isinstance(self._hidden_params, HiddenParams): - return getattr(self._hidden_params, key, None) + return getattr(self._hidden_params, "additional_headers", None) def set_timing_metrics( self, @@ -96,7 +95,7 @@ class ResponseMetadata: ######################################################### # 2. Add LiteLLM overhead duration ######################################################### - llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms") + llm_api_duration_ms: Final[float | None] = logging_obj.model_call_details.get("llm_api_duration_ms") if llm_api_duration_ms is not None: overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4) self._update_hidden_params( @@ -108,7 +107,7 @@ class ResponseMetadata: ######################################################### # 3. Add callback processing duration ######################################################### - callback_duration_ms: Final = getattr(logging_obj, "callback_duration_ms", None) + callback_duration_ms: Final[float | None] = getattr(logging_obj, "callback_duration_ms", None) if callback_duration_ms is not None: self._update_hidden_params( { @@ -136,17 +135,17 @@ class ResponseMetadata: # 5. Detailed per-phase timing (opt-in via env var) ######################################################### if LITELLM_DETAILED_TIMING and llm_api_duration_ms is not None: - detailed: Final[dict] = { + detailed: Final[dict[str, float]] = { "timing_llm_api_ms": round(llm_api_duration_ms, 4), } # message copy time from Logging.__init__() - msg_copy_ms: Final = getattr(logging_obj, "message_copy_duration_ms", None) + msg_copy_ms: Final[float | None] = getattr(logging_obj, "message_copy_duration_ms", None) if msg_copy_ms is not None: detailed["timing_message_copy_ms"] = round(msg_copy_ms, 4) # pre-processing = time from request start to LLM API call start - api_call_start: Final = logging_obj.model_call_details.get("api_call_start_time") + api_call_start: Final[datetime.datetime | None] = logging_obj.model_call_details.get("api_call_start_time") if api_call_start is not None and start_time is not None: pre_ms: Final = (api_call_start - start_time).total_seconds() * 1000 detailed["timing_pre_processing_ms"] = round(pre_ms, 4) diff --git a/litellm/litellm_core_utils/model_response_utils.py b/litellm/litellm_core_utils/model_response_utils.py index ea4be1c856f..7d412229a96 100644 --- a/litellm/litellm_core_utils/model_response_utils.py +++ b/litellm/litellm_core_utils/model_response_utils.py @@ -2,10 +2,54 @@ Utility functions for ModelResponse and ModelResponseStream objects. """ -from typing import Any, Final +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final, Protocol from litellm.types.utils import Delta, ModelResponseBase, ModelResponseStream +_NO_EXTRA_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class _HasModelExtra(Protocol): + """A Pydantic model, seen through the extra fields it collected.""" + + @property + def model_extra(self) -> Mapping[str, object] | None: ... + + +class _StreamingChoice(_HasModelExtra, Protocol): + """The streaming-choice fields this emptiness check reads.""" + + @property + def finish_reason(self) -> object: ... + + @property + def logprobs(self) -> object: ... + + @property + def enhancements(self) -> object: ... + + @property + def delta(self) -> Delta | None: ... + + +def _extra_fields(model: _HasModelExtra) -> Mapping[str, object]: + """The dynamically added fields Pydantic stored on ``model``.""" + return model.model_extra or _NO_EXTRA_FIELDS + + +def _attribute(obj: object, name: str) -> object: + """The named attribute of ``obj``, or ``None`` when it is absent.""" + attribute: Final[object] = getattr(obj, name, None) + return attribute + + +def _has_callable_attribute(obj: object, name: str) -> bool: + """Whether the named attribute of ``obj`` is callable.""" + attribute: Final[object] = getattr(obj, name) + return callable(attribute) + def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: """ @@ -41,7 +85,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: # Check model_extra for dynamically added fields (this is where Pydantic stores them) if hasattr(model_response, "model_extra") and model_response.model_extra: - for extra_field_name, extra_field_value in model_response.model_extra.items(): + for extra_field_name, extra_field_value in _extra_fields(model_response).items(): if _has_meaningful_content(extra_field_value): return False @@ -57,7 +101,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: continue # Check if any other field has meaningful content - model_response_value = getattr(model_response, model_response_field, None) + model_response_value = _attribute(model_response, model_response_field) if _has_meaningful_content(model_response_value): return False @@ -71,7 +115,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: return True -def _has_meaningful_content(value: Any) -> bool: +def _has_meaningful_content(value: object) -> bool: """ Check if a value contains meaningful content. @@ -102,7 +146,7 @@ def _has_meaningful_content(value: Any) -> bool: return True -def _is_choice_non_empty(choice: Any) -> bool: +def _is_choice_non_empty(choice: _StreamingChoice) -> bool: """ Deep check if a choice contains any meaningful content. @@ -131,7 +175,7 @@ def _is_choice_non_empty(choice: Any) -> bool: # Check model_extra for dynamically added fields on the choice if hasattr(choice, "model_extra") and choice.model_extra: - for extra_field_name, extra_field_value in choice.model_extra.items(): + for extra_field_name, extra_field_value in _extra_fields(choice).items(): # Skip certain structural fields that are just default/None placeholders if extra_field_name == "index" and extra_field_value == 0: continue @@ -147,7 +191,7 @@ def _is_choice_non_empty(choice: Any) -> bool: # Skip private attributes, methods, and known empty fields if ( attr_name.startswith("_") - or callable(getattr(choice, attr_name)) + or _has_callable_attribute(choice, attr_name) or attr_name.startswith("model_") or attr_name in { @@ -160,7 +204,7 @@ def _is_choice_non_empty(choice: Any) -> bool: ): continue - attr_value = getattr(choice, attr_name, None) + attr_value = _attribute(choice, attr_name) if _has_meaningful_content(attr_value): return True @@ -179,7 +223,7 @@ def _is_delta_non_empty(delta: Delta) -> bool: """ # Check model_extra for dynamically added fields (this is where Pydantic stores them) if hasattr(delta, "model_extra") and delta.model_extra: - for extra_field_name, extra_field_value in delta.model_extra.items(): + for extra_field_name, extra_field_value in _extra_fields(delta).items(): # Even structural fields are meaningful if they have actual content if _has_meaningful_content(extra_field_value): return True @@ -187,10 +231,10 @@ def _is_delta_non_empty(delta: Delta) -> bool: # Check all regular attributes of the delta object for attr_name in dir(delta): # Skip private attributes, methods, and Pydantic-specific fields - if attr_name.startswith("_") or callable(getattr(delta, attr_name)) or attr_name.startswith("model_"): + if attr_name.startswith("_") or _has_callable_attribute(delta, attr_name) or attr_name.startswith("model_"): continue - attr_value = getattr(delta, attr_name, None) + attr_value = _attribute(delta, attr_name) if _has_meaningful_content(attr_value): return True diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 485091bccd0..300ba427cda 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -92,7 +92,7 @@ def print_verbose(print_statement: object): @dataclass(frozen=True, slots=True) class _ProviderChunkParsed: - response_obj: dict[str, Any] + response_obj: dict[str, object] @dataclass(frozen=True, slots=True) @@ -1242,7 +1242,7 @@ class CustomStreamWrapper: for key, value in anthropic_response_obj["provider_specific_fields"].items(): setattr(model_response, key, value) - response_obj = cast(dict[str, Any], anthropic_response_obj) + response_obj = cast(dict[str, object], anthropic_response_obj) elif self.model == "replicate" or self.custom_llm_provider == "replicate": response_obj = self.handle_replicate_chunk(chunk) completion_obj["content"] = response_obj["text"] @@ -1398,7 +1398,7 @@ class CustomStreamWrapper: if not isinstance(chunk, str): raise ValueError(f"chunk is not a string: {chunk}") response_obj = cast( - dict[str, Any], + dict[str, object], litellm.CodestralTextCompletionConfig()._chunk_parser(chunk), ) completion_obj["content"] = response_obj["text"] @@ -2462,7 +2462,7 @@ def calculate_total_usage(chunks: list[ModelResponse]) -> Usage: prompt_tokens: int = 0 completion_tokens: int = 0 - latest_usage_chunk = None + latest_usage_chunk: Usage | Mapping[str, int] | None = None prompt_tokens_details: PromptTokensDetailsWrapper | None = None completion_tokens_details: CompletionTokensDetailsWrapper | None = None cache_creation_token_details: CacheCreationTokenDetails | None = None diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 0a59eaa75d3..53ae25b928c 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -20,6 +20,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config): """ import socket +from collections.abc import Sequence from ipaddress import ip_address, ip_network from typing import Any, Final from urllib.parse import quote, urlparse, urlunparse @@ -44,7 +45,7 @@ class SSRFError(ValueError): """Raised when a URL targets a blocked network.""" -def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -> str: +def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str: """Percent-encode one user-controlled URL path segment. ``urllib.parse.quote(..., safe="")`` intentionally leaves RFC 3986 @@ -64,7 +65,7 @@ def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") - return quote(value_str, safe="") -def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str: +def encode_url_path_segments(value: object, *, field_name: str = "path") -> str: """Percent-encode a user-controlled URL path made of multiple segments. Empty segments are rejected, so leading, trailing, or consecutive slashes @@ -77,9 +78,9 @@ def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str: if value_str == "": raise ValueError(f"{field_name} is required") - encoded_segments: Final = [] - for segment in value_str.split("/"): - encoded_segments.append(encode_url_path_segment(segment, field_name=field_name)) + encoded_segments: Final = tuple( + encode_url_path_segment(segment, field_name=field_name) for segment in value_str.split("/") + ) return "/".join(encoded_segments) @@ -202,7 +203,7 @@ def _format_host_header(hostname: str, port: int, default_port: int) -> str: return f"{bracketed}:{port}" -def _sockaddr_host(sockaddr: Any) -> str: +def _sockaddr_host(sockaddr: Sequence[object]) -> str: """Return the host element of a ``getaddrinfo`` sockaddr as ``str``. ``getaddrinfo`` with ``IPPROTO_TCP`` returns AF_INET / AF_INET6 sockaddrs @@ -285,7 +286,7 @@ def validate_url(url: str) -> tuple[str, str]: raise SSRFError(f"No addresses found for '{hostname}'") if not is_allowlisted: - for family, type_, proto, canonname, sockaddr in addrinfo: + for _family, _type, _proto, _canonname, sockaddr in addrinfo: resolved_ip = _sockaddr_host(sockaddr) if _is_blocked_ip(resolved_ip): raise SSRFError( @@ -363,7 +364,7 @@ def assert_same_origin(candidate_url: str, expected_url: str) -> None: _MAX_REDIRECTS: Final = 10 -def _extract_redirect_url(response: Any, request_url: str) -> str: +def _extract_redirect_url(response: httpx.Response, request_url: str) -> str: """Extract and resolve the redirect target from a response's Location header.""" location: Final = response.headers.get("location") if not isinstance(location, str) or not location: @@ -372,7 +373,7 @@ def _extract_redirect_url(response: Any, request_url: str) -> str: return str(httpx.URL(request_url).join(location)) -def safe_get(client: Any, url: str, **kwargs: Any) -> Any: +def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: """ Fetch a user-supplied URL with SSRF protection on every redirect hop. @@ -398,7 +399,7 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any: caller_headers: Final = kwargs.pop("headers", {}) for _ in range(_MAX_REDIRECTS): validated_url, original_host = validate_url(url) - response = client.get( + response: httpx.Response = client.get( validated_url, headers={**caller_headers, "Host": original_host}, follow_redirects=False, @@ -412,7 +413,7 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any: raise SSRFError("Too many redirects") -async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any: +async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: """Async version of safe_get.""" if not getattr(litellm, "user_url_validation", True): kwargs.setdefault("follow_redirects", True) @@ -421,7 +422,7 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any: caller_headers: Final = kwargs.pop("headers", {}) for _ in range(_MAX_REDIRECTS): validated_url, original_host = validate_url(url) - response = await client.get( + response: httpx.Response = await client.get( validated_url, headers={**caller_headers, "Host": original_host}, follow_redirects=False, diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index 3f8fd2c27f4..3071c3ddd58 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -1,9 +1,11 @@ import json import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast import httpx from httpx import Headers, Response +from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig @@ -19,6 +21,29 @@ else: LoggingClass = Any +class AnthropicBatchRequestCounts(TypedDict, total=False): + """The ``request_counts`` object of an Anthropic Message Batch.""" + + processing: ReadOnly[int] + succeeded: ReadOnly[int] + errored: ReadOnly[int] + canceled: ReadOnly[int] + expired: ReadOnly[int] + + +class AnthropicMessageBatch(TypedDict, total=False): + """The fields of an Anthropic Message Batch that map onto an OpenAI Batch.""" + + id: ReadOnly[str] + processing_status: ReadOnly[str] + created_at: ReadOnly[str | None] + ended_at: ReadOnly[str | None] + expires_at: ReadOnly[str | None] + cancel_initiated_at: ReadOnly[str | None] + archived_at: ReadOnly[str | None] + request_counts: ReadOnly[AnthropicBatchRequestCounts] + + class AnthropicBatchesConfig(BaseBatchesConfig): def __init__(self): from ..chat.transformation import AnthropicConfig @@ -83,7 +108,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): create_batch_data: CreateBatchRequest, optional_params: dict, litellm_params: dict, - ) -> bytes | str | dict[str, Any]: + ) -> bytes | str | dict[str, object]: """ Transform the batch creation request to Anthropic format. @@ -133,7 +158,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): batch_id: str, optional_params: dict, litellm_params: dict, - ) -> bytes | str | dict[str, Any]: + ) -> bytes | str | dict[str, object]: """ Transform batch retrieval request for Anthropic. @@ -152,7 +177,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): ) -> LiteLLMBatch: """Transform Anthropic MessageBatch retrieval response to LiteLLM format.""" try: - response_data: Final = raw_response.json() + response_data: Final[AnthropicMessageBatch] = raw_response.json() except Exception as e: raise ValueError(f"Failed to parse Anthropic batch response: {e}") @@ -161,18 +186,20 @@ class AnthropicBatchesConfig(BaseBatchesConfig): processing_status: Final = response_data.get("processing_status", "in_progress") # Map Anthropic processing_status to OpenAI status - status_mapping: dict[ - str, - Literal[ - "validating", - "failed", - "in_progress", - "finalizing", - "completed", - "expired", - "cancelling", - "cancelled", - ], + status_mapping: Final[ + Mapping[ + str, + Literal[ + "validating", + "failed", + "in_progress", + "finalizing", + "completed", + "expired", + "cancelling", + "cancelled", + ], + ] ] = { "in_progress": "in_progress", "canceling": "cancelling", @@ -279,7 +306,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): if not line: continue try: - response_json = json.loads(line) + response_json: Mapping[str, Mapping[str, dict[str, object]]] = json.loads(line) # Update model_response with the parsed JSON completion_response = response_json["result"]["message"] transformed_response = self.anthropic_chat_config.transform_parsed_response( diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py index 41795fa0f32..8e06b73a92c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py @@ -1,8 +1,10 @@ """Dispatch ``context_management`` edits to registered polyfill editors.""" import inspect -from collections.abc import Awaitable, Callable -from typing import Any, Final, cast +from collections.abc import Awaitable, Callable, Mapping +from typing import TYPE_CHECKING, Final, TypeAlias, TypedDict, cast + +from typing_extensions import ReadOnly from litellm._logging import verbose_logger from litellm.types.llms.anthropic import AppliedEdit @@ -11,31 +13,54 @@ from .constants import CLEAR_TOOL_USES_EDIT_TYPE, COMPACT_EDIT_TYPE from .editors import apply_clear_tool_uses_20250919, apply_compact_20260112 from .result import PolyfillResult -EditorFn = Callable[..., Any] +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.router import Router -_EDITOR_REGISTRY: Final[dict[str, EditorFn]] = { +AnthropicMessages: TypeAlias = list[dict[str, object]] +AnthropicSystem: TypeAlias = str | list[dict[str, object]] | None +AnthropicTools: TypeAlias = list[dict[str, object]] | None +EditSpec: TypeAlias = dict[str, object] +ContextManagementSpec: TypeAlias = EditSpec | list[EditSpec] | None +SyncEditorReturn: TypeAlias = tuple[AnthropicMessages, AppliedEdit | None] +EditorFn: TypeAlias = Callable[..., object] + + +class EditorKwargs(TypedDict): + """The keyword payload every registered editor accepts.""" + + model: ReadOnly[str] + messages: ReadOnly[AnthropicMessages] + tools: ReadOnly[AnthropicTools] + system: ReadOnly[AnthropicSystem] + edit_spec: ReadOnly[EditSpec] + + +_EDITOR_REGISTRY: Final[Mapping[str, EditorFn]] = { CLEAR_TOOL_USES_EDIT_TYPE: apply_clear_tool_uses_20250919, COMPACT_EDIT_TYPE: apply_compact_20260112, } -def _normalize_spec( - spec: dict[str, Any] | list[dict[str, Any]] | None, -) -> list[dict[str, Any]] | None: +def _map_openai_spec(spec: list[EditSpec]) -> EditSpec | None: + """Translate the OpenAI list form into the Anthropic-native dict form.""" + # Local import to avoid an import cycle at module load. + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + return AnthropicConfig.map_openai_context_management_to_anthropic(spec) + + +def _normalize_spec(spec: ContextManagementSpec) -> list[EditSpec] | None: """Accept Anthropic-native dict form or OpenAI list form; return edits list.""" - if isinstance(spec, list): - # Local import to avoid an import cycle at module load. - from litellm.llms.anthropic.chat.transformation import AnthropicConfig + normalized: Final = _map_openai_spec(spec) if isinstance(spec, list) else spec - spec = AnthropicConfig.map_openai_context_management_to_anthropic(spec) - - edits: Final = spec.get("edits") if isinstance(spec, dict) else None + edits: Final = normalized.get("edits") if isinstance(normalized, dict) else None if not edits or not isinstance(edits, list): return None return [edit for edit in edits if isinstance(edit, dict)] -def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult: +def _wrap_editor_return(raw: object, *, fallback_system: AnthropicSystem) -> PolyfillResult: """Coerce an editor's native return shape into a ``PolyfillResult``. v0 sync editors (e.g. ``clear_tool_uses_20250919``) return a 2-tuple @@ -46,7 +71,7 @@ def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult: return raw # Legacy 2-tuple return — sync editors don't mutate ``system``, so # carry the caller's value forward. - messages, applied = cast(tuple[list[dict[str, Any]], Any], raw) + messages, applied = cast(SyncEditorReturn, raw) return PolyfillResult( messages=messages, system=fallback_system, @@ -57,13 +82,13 @@ def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult: async def apply_context_management( *, model: str, - messages: list[dict[str, Any]], - tools: list[dict[str, Any]] | None, - system: Any, - context_management_spec: dict[str, Any] | list[dict[str, Any]] | None, - litellm_metadata: dict[str, Any] | None = None, - llm_router: Any = None, - user_api_key_auth: Any = None, + messages: AnthropicMessages, + tools: AnthropicTools, + system: AnthropicSystem, + context_management_spec: ContextManagementSpec, + litellm_metadata: Mapping[str, object] | None = None, + llm_router: "Router | None" = None, + user_api_key_auth: "UserAPIKeyAuth | None" = None, ) -> PolyfillResult: """Run edits in order; return a single ``PolyfillResult``. @@ -92,7 +117,7 @@ async def apply_context_management( ) continue - kwargs: dict[str, Any] = { + kwargs: EditorKwargs = { "model": model, "messages": current_messages, "tools": tools, @@ -102,10 +127,12 @@ async def apply_context_management( # Only async editors accept these — passing them to sync v0 editors # would break their signature. if inspect.iscoroutinefunction(editor): - kwargs["litellm_metadata"] = litellm_metadata - kwargs["llm_router"] = llm_router - kwargs["user_api_key_auth"] = user_api_key_auth - raw_result = await cast(Callable[..., Awaitable[Any]], editor)(**kwargs) + raw_result = await cast(Callable[..., Awaitable[PolyfillResult]], editor)( + **kwargs, + litellm_metadata=litellm_metadata, + llm_router=llm_router, + user_api_key_auth=user_api_key_auth, + ) else: raw_result = editor(**kwargs) diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index e6d8686b466..d99c7f556fb 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -33,18 +33,18 @@ def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str, def _build_responses_kwargs( *, max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - context_management: dict | None = None, - metadata: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[AllAnthropicToolsValues | dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: AnthropicOutputSchema | None = None, @@ -134,22 +134,22 @@ class LiteLLMMessagesToResponsesAPIHandler: @staticmethod async def async_anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - context_management: dict | None = None, - metadata: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[AllAnthropicToolsValues | dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: AnthropicOutputSchema | None = None, - **kwargs, + **kwargs: object, ) -> AnthropicMessagesResponse | AsyncIterator[bytes]: responses_kwargs: Final = _build_responses_kwargs( max_tokens=max_tokens, @@ -185,23 +185,23 @@ class LiteLLMMessagesToResponsesAPIHandler: @staticmethod def anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - context_management: dict | None = None, - metadata: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[AllAnthropicToolsValues | dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: AnthropicOutputSchema | None = None, _is_async: bool = False, - **kwargs, + **kwargs: object, ) -> ( AnthropicMessagesResponse | AsyncIterator[bytes] diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index 566322bbdd6..20f3a3d74f0 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -2,9 +2,10 @@ Anthropic Skills API configuration and transformations """ -from typing import Any, Final +from typing import Final import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -23,6 +24,25 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +class _SkillPayload(TypedDict): + """The JSON body Anthropic returns for a single skill, before ``Skill`` validates it.""" + + id: ReadOnly[str] + created_at: ReadOnly[str] + source: ReadOnly[str] + updated_at: ReadOnly[str] + display_title: NotRequired[ReadOnly[str | None]] + latest_version: NotRequired[ReadOnly[str | None]] + type: NotRequired[ReadOnly[str]] + + +class _DeleteSkillPayload(TypedDict): + """The JSON body Anthropic returns for a skill deletion, before ``DeleteSkillResponse`` validates it.""" + + id: ReadOnly[str] + type: NotRequired[ReadOnly[str]] + + class AnthropicSkillsConfig(BaseSkillsAPIConfig): """Anthropic-specific Skills API configuration""" @@ -104,7 +124,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Skill: """Transform Anthropic response to Skill object""" - response_json: Final = raw_response.json() + response_json: Final[_SkillPayload] = raw_response.json() verbose_logger.debug("Transforming create skill response: %s", response_json) return Skill(**response_json) @@ -122,7 +142,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): url: Final = self.get_complete_url(api_base=api_base, endpoint="skills") # Build query parameters - query_params: Final[dict[str, Any]] = {} + query_params: Final[dict[str, int | str]] = {} if "limit" in list_params and list_params["limit"]: query_params["limit"] = list_params["limit"] if "page" in list_params and list_params["page"]: @@ -168,7 +188,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Skill: """Transform Anthropic response to Skill object""" - response_json: Final = raw_response.json() + response_json: Final[_SkillPayload] = raw_response.json() verbose_logger.debug("Transforming get skill response: %s", response_json) return Skill(**response_json) @@ -193,7 +213,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> DeleteSkillResponse: """Transform Anthropic response to DeleteSkillResponse""" - response_json: Final = raw_response.json() + response_json: Final[_DeleteSkillPayload] = raw_response.json() verbose_logger.debug("Transforming delete skill response: %s", response_json) return DeleteSkillResponse(**response_json) diff --git a/litellm/llms/azure/files/handler.py b/litellm/llms/azure/files/handler.py index 4f93896699f..f3ed2e6d6ed 100644 --- a/litellm/llms/azure/files/handler.py +++ b/litellm/llms/azure/files/handler.py @@ -40,6 +40,11 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): data.pop("expires_after", None) return data + @staticmethod + def _to_openai_file_object(response: FileObject) -> OpenAIFileObject: + """Re-wrap the SDK's file object as litellm's, carrying every field across.""" + return OpenAIFileObject(**response.model_dump()) + async def acreate_file( self, create_file_data: CreateFileRequest, @@ -48,7 +53,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): verbose_logger.debug("create_file_data=%s", create_file_data) response = await openai_client.files.create(**self._prepare_create_file_data(create_file_data)) verbose_logger.debug("create_file_response=%s", response) - return OpenAIFileObject(**response.model_dump()) + return self._to_openai_file_object(response) def create_file( self, @@ -61,7 +66,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): max_retries: int | None, client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, litellm_params: dict | None = None, - ) -> OpenAIFileObject | Coroutine[Any, Any, OpenAIFileObject]: + ) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]: openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, @@ -84,7 +89,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): response: Final = cast(AzureOpenAI | OpenAI, openai_client).files.create( **self._prepare_create_file_data(create_file_data) ) - return OpenAIFileObject(**response.model_dump()) + return self._to_openai_file_object(response) async def afile_content( self, @@ -105,7 +110,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): api_version: str | None = None, client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, litellm_params: dict | None = None, - ) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: + ) -> HttpxBinaryResponseContent | Coroutine[None, None, HttpxBinaryResponseContent]: openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 3eeb3cb9fc6..e795e907eb4 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -7,9 +7,10 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic. import asyncio import contextlib import json -from typing import Any, Final +from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict from litellm._logging import _redact_string, verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -18,9 +19,74 @@ from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig +if TYPE_CHECKING: + from litellm.types.realtime import RealtimeResponseTransformInput + _CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) +class _ClientWebSocket(Protocol): + """The client-facing websocket surface used by the Bedrock realtime bridge.""" + + async def send_text(self, data: str) -> None: ... + + async def receive_text(self) -> str: ... + + async def close(self, code: int = ..., reason: str | None = ...) -> None: ... + + +class _BedrockInputStream(Protocol): + """The write half of a Bedrock bidirectional stream.""" + + async def send(self, chunk: object) -> None: ... + + async def close(self) -> None: ... + + +class _BedrockPayloadPart(Protocol): + """A single Bedrock bidirectional output payload.""" + + bytes_: bytes | None + + +class _BedrockOutputChunk(Protocol): + """A chunk read off the read half of a Bedrock bidirectional stream.""" + + value: _BedrockPayloadPart | None + + +class _BedrockOutputStream(Protocol): + """The read half of a Bedrock bidirectional stream.""" + + async def receive(self) -> _BedrockOutputChunk | None: ... + + +class _BedrockBidirectionalStream(Protocol): + """The bidirectional stream returned by ``invoke_model_with_bidirectional_stream``.""" + + input_stream: _BedrockInputStream + + async def await_output(self) -> tuple[object, _BedrockOutputStream]: ... + + +class _ClientSessionPayload(TypedDict, total=False): + """The ``session`` body of a client ``session.update`` frame.""" + + modalities: ReadOnly[object] + + +class _ClientRealtimeFrame(TypedDict, total=False): + """The fields read off a client realtime frame.""" + + type: ReadOnly[str] + session: ReadOnly[_ClientSessionPayload] + + +def _decode_client_frame(payload: str) -> _ClientRealtimeFrame: + """Decode a client realtime frame into the fields this bridge reads.""" + return json.loads(payload) + + class BedrockRealtime(BaseAWSLLM): """Handler for Bedrock Nova Sonic realtime speech-to-speech API.""" @@ -30,7 +96,7 @@ class BedrockRealtime(BaseAWSLLM): async def async_realtime( self, model: str, - websocket: Any, + websocket: _ClientWebSocket, logging_obj: LiteLLMLogging, api_base: str | None = None, api_key: str | None = None, @@ -46,7 +112,7 @@ class BedrockRealtime(BaseAWSLLM): aws_sts_endpoint: str | None = None, aws_bedrock_runtime_endpoint: str | None = None, aws_external_id: str | None = None, - **kwargs, + **kwargs: object, ): """ Establish bidirectional streaming connection with Bedrock Nova Sonic. @@ -118,13 +184,16 @@ class BedrockRealtime(BaseAWSLLM): ) bedrock_client: Final = BedrockRuntimeClient(config=config) + async def open_bidirectional_stream() -> _BedrockBidirectionalStream: + return await bedrock_client.invoke_model_with_bidirectional_stream( + InvokeModelWithBidirectionalStreamOperationInput(model_id=model) + ) + transformation_config: Final = BedrockRealtimeConfig() try: # Initialize the bidirectional stream - bedrock_stream: Final = await bedrock_client.invoke_model_with_bidirectional_stream( - InvokeModelWithBidirectionalStreamOperationInput(model_id=model) - ) + bedrock_stream: Final = await open_bidirectional_stream() verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") @@ -132,7 +201,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") # Track state for transformation - session_state: Final = { + session_state: Final[RealtimeResponseTransformInput] = { "current_output_item_id": None, "current_response_id": None, "current_conversation_id": None, @@ -182,11 +251,11 @@ class BedrockRealtime(BaseAWSLLM): async def _forward_client_to_bedrock( self, - client_ws: Any, - bedrock_stream: Any, + client_ws: _ClientWebSocket, + bedrock_stream: _BedrockBidirectionalStream, transformation_config: BedrockRealtimeConfig, model: str, - session_state: dict, + session_state: "RealtimeResponseTransformInput", logging_obj: LiteLLMLogging | None = None, ): """Forward messages from client WebSocket to Bedrock stream.""" @@ -195,10 +264,11 @@ class BedrockRealtime(BaseAWSLLM): InvokeModelWithBidirectionalStreamInputChunk, ) + def build_input_chunk(payload: bytes) -> object: + return InvokeModelWithBidirectionalStreamInputChunk(value=BidirectionalInputPayloadPart(bytes_=payload)) + async def send_to_bedrock(bedrock_message: str) -> None: - event: Final = InvokeModelWithBidirectionalStreamInputChunk( - value=BidirectionalInputPayloadPart(bytes_=bedrock_message.encode("utf-8")) - ) + event: Final = build_input_chunk(bedrock_message.encode("utf-8")) await bedrock_stream.input_stream.send(event) verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200]) @@ -223,7 +293,7 @@ class BedrockRealtime(BaseAWSLLM): client_message_type: str | None = None requested_modalities: list[str] | None = None with contextlib.suppress(Exception): - parsed_client_message = json.loads(message) + parsed_client_message = _decode_client_frame(message) client_message_type = parsed_client_message.get("type") if client_message_type == "session.update": requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python( @@ -246,12 +316,12 @@ class BedrockRealtime(BaseAWSLLM): async def _forward_bedrock_to_client( self, - bedrock_stream: Any, - client_ws: Any, + bedrock_stream: _BedrockBidirectionalStream, + client_ws: _ClientWebSocket, transformation_config: BedrockRealtimeConfig, model: str, logging_obj: LiteLLMLogging, - session_state: dict, + session_state: "RealtimeResponseTransformInput", ): """Forward messages from Bedrock stream to client WebSocket.""" try: diff --git a/litellm/llms/chatgpt/chat/streaming_utils.py b/litellm/llms/chatgpt/chat/streaming_utils.py index 57f679947f6..f8c54c0a1d8 100644 --- a/litellm/llms/chatgpt/chat/streaming_utils.py +++ b/litellm/llms/chatgpt/chat/streaming_utils.py @@ -4,7 +4,32 @@ Streaming utilities for ChatGPT provider. Normalizes non-spec-compliant tool_call chunks from the ChatGPT backend API. """ -from typing import Any, Final +from collections.abc import Sequence +from typing import Final, Protocol + +from litellm.types.utils import Delta + + +class ChatGPTStreamChoice(Protocol): + """Streaming choice as read by :class:`ChatGPTToolCallNormalizer`.""" + + @property + def delta(self) -> Delta | None: ... + + +class ChatGPTStreamChunk(Protocol): + """Streaming chunk as read by :class:`ChatGPTToolCallNormalizer`.""" + + @property + def choices(self) -> Sequence[ChatGPTStreamChoice]: ... + + +class ChatGPTChunkStream(Protocol): + """Sync/async chunk source wrapped by :class:`ChatGPTToolCallNormalizer`.""" + + def __next__(self) -> ChatGPTStreamChunk: ... + + async def __anext__(self) -> ChatGPTStreamChunk: ... class ChatGPTToolCallNormalizer: @@ -20,36 +45,36 @@ class ChatGPTToolCallNormalizer: chunks to the consumer. """ - def __init__(self, stream: Any): + def __init__(self, stream: ChatGPTChunkStream): self._stream = stream self._seen_ids: dict[str, int] = {} # tool_call_id -> assigned_index self._next_index: int = 0 self._last_id: str | None = None # tracks which tool call the next delta belongs to - def __getattr__(self, name: str) -> Any: + def __getattr__(self, name: str) -> object: return getattr(self._stream, name) - def __iter__(self): + def __iter__(self) -> "ChatGPTToolCallNormalizer": return self - def __aiter__(self): + def __aiter__(self) -> "ChatGPTToolCallNormalizer": return self - def __next__(self): + def __next__(self) -> ChatGPTStreamChunk: while True: chunk = next(self._stream) result = self._normalize(chunk) if result is not None: return result - async def __anext__(self): + async def __anext__(self) -> ChatGPTStreamChunk: while True: chunk = await self._stream.__anext__() result = self._normalize(chunk) if result is not None: return result - def _normalize(self, chunk: Any) -> Any: + def _normalize(self, chunk: ChatGPTStreamChunk) -> ChatGPTStreamChunk | None: """Fix tool_calls in the chunk. Returns None to skip duplicate chunks.""" if not chunk.choices: return chunk diff --git a/litellm/llms/compactifai/chat/transformation.py b/litellm/llms/compactifai/chat/transformation.py index 44e1ab15801..3d189911e9e 100644 --- a/litellm/llms/compactifai/chat/transformation.py +++ b/litellm/llms/compactifai/chat/transformation.py @@ -2,13 +2,16 @@ CompactifAI chat completion transformation """ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.common_utils import OpenAIError from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse from ...openai.chat.gpt_transformation import OpenAIGPTConfig @@ -21,6 +24,18 @@ else: LiteLLMLoggingObj = Any +class CompactifAIResponseFields(TypedDict, total=False): + """The chat completion fields of a CompactifAI response body.""" + + id: ReadOnly[str] + choices: ReadOnly[Sequence[Mapping[str, object]]] + created: ReadOnly[int] + model: ReadOnly[str] + system_fingerprint: ReadOnly[str | None] + usage: ReadOnly[Mapping[str, object]] + object: ReadOnly[str] + + class CompactifAIChatConfig(OpenAIGPTConfig): """ Configuration class for CompactifAI chat completions. @@ -45,11 +60,11 @@ class CompactifAIChatConfig(OpenAIGPTConfig): raw_response: httpx.Response, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, - request_data: dict, - messages: list, - optional_params: dict, - litellm_params: dict, - encoding: Any, + request_data: Mapping[str, object], + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + encoding: object, api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -79,14 +94,18 @@ class CompactifAIChatConfig(OpenAIGPTConfig): message["content"] = tool_calls[0]["function"].get("arguments", "") message["tool_calls"] = None - returned_response: Final = ModelResponse(**response_json) + response_fields: Final[CompactifAIResponseFields] = response_json + + returned_response: Final = ModelResponse(**response_fields) # Set model name with provider prefix returned_response.model = f"compactifai/{model}" return returned_response - def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, str] | httpx.Headers + ) -> BaseLLMException: """ Get the appropriate error class for CompactifAI errors. Since CompactifAI is OpenAI-compatible, we use OpenAI error handling. diff --git a/litellm/llms/custom_httpx/container_handler.py b/litellm/llms/custom_httpx/container_handler.py index 7690351e3b2..343cf043187 100644 --- a/litellm/llms/custom_httpx/container_handler.py +++ b/litellm/llms/custom_httpx/container_handler.py @@ -6,11 +6,13 @@ endpoint defined in endpoints.json, eliminating the need for individual handler """ import json -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping, Sequence from pathlib import Path from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import BaseModel +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -32,22 +34,54 @@ if TYPE_CHECKING: from litellm.llms.base_llm.containers.transformation import BaseContainerConfig +class EndpointConfig(TypedDict): + """One endpoint entry of ``litellm/containers/endpoints.json``.""" + + name: ReadOnly[str] + async_name: ReadOnly[str] + path: ReadOnly[str] + method: ReadOnly[str] + path_params: ReadOnly[Sequence[str]] + query_params: ReadOnly[Sequence[str]] + response_type: ReadOnly[str] + is_multipart: NotRequired[ReadOnly[bool]] + returns_binary: NotRequired[ReadOnly[bool]] + + +class EndpointsConfig(TypedDict): + """The parsed ``litellm/containers/endpoints.json`` document.""" + + endpoints: ReadOnly[Sequence[EndpointConfig]] + + +class ContainerErrorDetail(TypedDict, total=False): + """The ``error`` object of a container API error body.""" + + message: ReadOnly[str] + + +class ContainerResponseBody(TypedDict, total=False): + """The fields this handler reads off a container API JSON body.""" + + error: ReadOnly[ContainerErrorDetail] + + # Response type mapping -RESPONSE_TYPES: Final[dict[str, type]] = { +RESPONSE_TYPES: Final[Mapping[str, type[BaseModel]]] = { "ContainerFileListResponse": ContainerFileListResponse, "ContainerFileObject": ContainerFileObject, "DeleteContainerFileResponse": DeleteContainerFileResponse, } -def _load_endpoints_config() -> dict: +def _load_endpoints_config() -> EndpointsConfig: """Load the endpoints configuration from JSON file.""" config_path: Final = Path(__file__).parent.parent.parent / "containers" / "endpoints.json" with open(config_path) as f: return json.load(f) -def _get_endpoint_config(endpoint_name: str) -> dict | None: +def _get_endpoint_config(endpoint_name: str) -> EndpointConfig | None: """Get config for a specific endpoint by name.""" config: Final = _load_endpoints_config() for endpoint in config["endpoints"]: @@ -56,10 +90,15 @@ def _get_endpoint_config(endpoint_name: str) -> dict | None: return None +def _response_model(response_type_name: str) -> type[BaseModel] | None: + """The pydantic model a container endpoint's ``response_type`` names.""" + return RESPONSE_TYPES.get(response_type_name) + + def _build_url( api_base: str, path_template: str, - path_params: dict[str, str], + path_params: Mapping[str, object], ) -> str: """Build the full URL by substituting path parameters. @@ -89,22 +128,18 @@ def _build_url( def _build_query_params( - query_param_names: list, - kwargs: dict[str, Any], -) -> dict[str, str]: + query_param_names: Sequence[str], + kwargs: Mapping[str, object], +) -> dict[str, object]: """Build query parameters from kwargs.""" - params: Final = {} - for param_name in query_param_names: - value = kwargs.get(param_name) - if value is not None: - params[param_name] = str(value) if not isinstance(value, str) else value - return params + supplied: Final = ((param_name, kwargs.get(param_name)) for param_name in query_param_names) + return {name: value if isinstance(value, str) else str(value) for name, value in supplied if value is not None} def _prepare_multipart_file_upload( file: Any, - headers: dict[str, Any], -) -> tuple: + headers: dict[str, object], +) -> tuple[dict[str, tuple[str, bytes, str]], dict[str, object]]: """ Prepare file and headers for multipart upload. @@ -129,6 +164,52 @@ def _prepare_multipart_file_upload( return files, headers_copy +def _request_headers( + container_provider_config: "BaseContainerConfig", + extra_headers: dict[str, object] | None, + litellm_params: GenericLiteLLMParams, +) -> dict[str, object]: + """The provider auth headers for a container request.""" + return container_provider_config.validate_environment( + headers=extra_headers or {}, + api_key=litellm_params.get("api_key", None), + ) + + +def _request_api_base( + container_provider_config: "BaseContainerConfig", + litellm_params: GenericLiteLLMParams, +) -> str: + """The provider base URL for a container request.""" + return container_provider_config.get_complete_url( + api_base=litellm_params.get("api_base", None), + litellm_params=dict(litellm_params), + ) + + +def _sync_http_client( + client: HTTPHandler | AsyncHTTPHandler | None, + litellm_params: GenericLiteLLMParams, +) -> HTTPHandler: + """The sync HTTP client for a container request, reusing the caller's when usable.""" + if client is None or not isinstance(client, HTTPHandler): + return _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) + return client + + +def _async_http_client( + client: HTTPHandler | AsyncHTTPHandler | None, + litellm_params: GenericLiteLLMParams, +) -> AsyncHTTPHandler: + """The async HTTP client for a container request, reusing the caller's when usable.""" + if client is None or not isinstance(client, AsyncHTTPHandler): + return get_async_httpx_client( + llm_provider=litellm.LlmProviders.OPENAI, + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + return client + + class GenericContainerHandler: """ Generic handler for container file API endpoints. @@ -143,13 +224,13 @@ class GenericContainerHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout = 600, _is_async: bool = False, client: HTTPHandler | AsyncHTTPHandler | None = None, - **kwargs, - ) -> Any | Coroutine[Any, Any, Any]: + **kwargs: object, + ) -> Any | Coroutine[object, object, Any]: """ Generic handler for any container file endpoint. @@ -196,11 +277,11 @@ class GenericContainerHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout = 600, client: HTTPHandler | AsyncHTTPHandler | None = None, - **kwargs, + **kwargs: object, ) -> Any: """Synchronous request handler.""" endpoint_config: Final = _get_endpoint_config(endpoint_name) @@ -208,23 +289,14 @@ class GenericContainerHandler: raise ValueError(f"Unknown endpoint: {endpoint_name}") # Get HTTP client - if client is None or not isinstance(client, HTTPHandler): - http_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) - else: - http_client = client + http_client: Final = _sync_http_client(client, litellm_params) # Build request - headers = container_provider_config.validate_environment( - headers=extra_headers or {}, - api_key=litellm_params.get("api_key", None), - ) + headers = _request_headers(container_provider_config, extra_headers, litellm_params) if extra_headers: headers.update(extra_headers) - api_base: Final = container_provider_config.get_complete_url( - api_base=litellm_params.get("api_base", None), - litellm_params=dict(litellm_params), - ) + api_base: Final = _request_api_base(container_provider_config, litellm_params) # Build URL with path params path_params: Final = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])} @@ -275,11 +347,11 @@ class GenericContainerHandler: return response.content # Check for error response - response_json: Final = response.json() + response_json: Final[ContainerResponseBody] = response.json() if "error" in response_json: from litellm.llms.base_llm.chat.transformation import BaseLLMException - error_msg: Final = response_json.get("error", {}).get("message", str(response_json)) + error_msg: Final = response_json["error"].get("message", str(response_json)) raise BaseLLMException( status_code=response.status_code, message=error_msg, @@ -287,7 +359,7 @@ class GenericContainerHandler: ) # Parse response - response_type: Final = RESPONSE_TYPES.get(endpoint_config["response_type"]) + response_type: Final = _response_model(endpoint_config["response_type"]) if response_type: return response_type(**response_json) return response_json @@ -301,11 +373,11 @@ class GenericContainerHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout = 600, client: HTTPHandler | AsyncHTTPHandler | None = None, - **kwargs, + **kwargs: object, ) -> Any: """Asynchronous request handler.""" endpoint_config: Final = _get_endpoint_config(endpoint_name) @@ -313,26 +385,14 @@ class GenericContainerHandler: raise ValueError(f"Unknown endpoint: {endpoint_name}") # Get HTTP client - if client is None or not isinstance(client, AsyncHTTPHandler): - http_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders.OPENAI, - params={"ssl_verify": litellm_params.get("ssl_verify", None)}, - ) - else: - http_client = client + http_client: Final = _async_http_client(client, litellm_params) # Build request - headers = container_provider_config.validate_environment( - headers=extra_headers or {}, - api_key=litellm_params.get("api_key", None), - ) + headers = _request_headers(container_provider_config, extra_headers, litellm_params) if extra_headers: headers.update(extra_headers) - api_base: Final = container_provider_config.get_complete_url( - api_base=litellm_params.get("api_base", None), - litellm_params=dict(litellm_params), - ) + api_base: Final = _request_api_base(container_provider_config, litellm_params) # Build URL with path params path_params: Final = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])} @@ -383,11 +443,11 @@ class GenericContainerHandler: return response.content # Check for error response - response_json: Final = response.json() + response_json: Final[ContainerResponseBody] = response.json() if "error" in response_json: from litellm.llms.base_llm.chat.transformation import BaseLLMException - error_msg: Final = response_json.get("error", {}).get("message", str(response_json)) + error_msg: Final = response_json["error"].get("message", str(response_json)) raise BaseLLMException( status_code=response.status_code, message=error_msg, @@ -395,7 +455,7 @@ class GenericContainerHandler: ) # Parse response - response_type: Final = RESPONSE_TYPES.get(endpoint_config["response_type"]) + response_type: Final = _response_model(endpoint_config["response_type"]) if response_type: return response_type(**response_json) return response_json diff --git a/litellm/llms/gdc/chat/transformation.py b/litellm/llms/gdc/chat/transformation.py index 2d0322bf10f..03037512551 100644 --- a/litellm/llms/gdc/chat/transformation.py +++ b/litellm/llms/gdc/chat/transformation.py @@ -6,11 +6,25 @@ import json import os import re import threading -from typing import Any, Final +from collections.abc import Callable +from typing import Any, Final, Protocol from urllib.parse import urlsplit import litellm from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig +from litellm.types.llms.openai import AllMessageValues + + +class _GDCHAudienceCredentials(Protocol): + """A GDCH service account credential already bound to an audience, ready to mint a bearer token.""" + + @property + def valid(self) -> bool: ... + + @property + def token(self) -> str: ... + + def refresh(self, request: object) -> None: ... class GDCGeminiConfig(OpenAILikeChatConfig): @@ -21,7 +35,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._creds_lock = threading.Lock() - self._gdch_creds_cache: dict = {} + self._gdch_creds_cache: dict[tuple[str, str], _GDCHAudienceCredentials] = {} def get_supported_openai_params(self, model: str) -> list: return [ @@ -110,7 +124,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): return f"{api_base}/v1/projects/{project}/locations/{location}/chat/completions" - def _read_env_bool(self, val: Any, env_var: str, default: bool = True) -> bool | str: + def _read_env_bool(self, val: bool | str | None, env_var: str, default: bool = True) -> bool | str: def _parse(s: str) -> bool | str: cleaned: Final = s.strip().lower() if cleaned in ("false", "0", "no", "off"): @@ -129,7 +143,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): return default return _parse(_env_val) - def _fetch_auth(self, gdch_creds: Any, ssl_verify: bool | str) -> None: + def _fetch_auth(self, gdch_creds: _GDCHAudienceCredentials, ssl_verify: bool | str) -> None: import requests from google.auth.transport import requests as auth_requests @@ -138,13 +152,24 @@ class GDCGeminiConfig(OpenAILikeChatConfig): auth_request: Final = auth_requests.Request(session=auth_session) gdch_creds.refresh(auth_request) - def _cached_fetch_token(self, creds: Any, audience: str, ssl_verify: bool | str, api_key: str | None = None) -> str: + def _with_gdch_audience(self, creds: object, audience: str) -> _GDCHAudienceCredentials: + """The credential rebound to ``audience``, which GDCH requires before a token refresh.""" + bind_audience: Final[Callable[[str], _GDCHAudienceCredentials] | None] = getattr( + creds, "with_gdch_audience", None + ) + if bind_audience is None: + raise AttributeError("GDC credentials must expose with_gdch_audience to be bound to a request audience") + return bind_audience(audience) + + def _cached_fetch_token( + self, creds: object, audience: str, ssl_verify: bool | str, api_key: str | None = None + ) -> str: # Key cache by both audience and credential identity to prevent cross-caller contamination cache_key: Final = (audience.rstrip("/"), api_key or str(id(creds))) with self._creds_lock: if cache_key not in self._gdch_creds_cache: - self._gdch_creds_cache[cache_key] = creds.with_gdch_audience(audience.rstrip("/")) + self._gdch_creds_cache[cache_key] = self._with_gdch_audience(creds, audience.rstrip("/")) gdch_creds: Final = self._gdch_creds_cache[cache_key] @@ -155,7 +180,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): return token - def _load_creds_from_key(self, api_key: str) -> tuple[Any, bool]: + def _load_creds_from_key(self, api_key: str) -> tuple[object | None, bool]: import google.auth try: @@ -175,7 +200,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): self, headers: dict, model: str, - messages: list[Any], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: str | None = None, @@ -230,7 +255,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): if self._read_env_bool(litellm_params.get("gdc_token_caching"), "GDC_TOKEN_CACHING", default=False): token = self._cached_fetch_token(creds, audience, ssl_verify, api_key) else: - gdch_creds: Final = creds.with_gdch_audience(audience) + gdch_creds: Final = self._with_gdch_audience(creds, audience) self._fetch_auth(gdch_creds, ssl_verify) token = gdch_creds.token headers["Authorization"] = f"Bearer {token}" @@ -252,7 +277,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): def transform_request( self, model: str, - messages: list[Any], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, headers: dict, diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index ebcbf1b5a07..df4adf0c9a2 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -4,9 +4,11 @@ Transformation logic from Cohere's /v1/rerank format to Infinity's `/v1/rerank` Why separate file? Make it easy to see how transformation works """ +from collections.abc import Sequence from typing import Final import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._uuid import uuid @@ -25,6 +27,31 @@ from litellm.types.rerank import ( from ..common_utils import InfinityError +class _InfinityRerankUsage(TypedDict, extra_items=ReadOnly[int]): + """The token counters Infinity reports in the ``usage`` block of a rerank response.""" + + +class _InfinityRerankResult(TypedDict): + """One scored document in an Infinity ``/v1/rerank`` response.""" + + index: ReadOnly[int] + relevance_score: ReadOnly[float] + document: ReadOnly[str] + + +class _InfinityRerankResponse(TypedDict): + """The JSON body returned by Infinity's ``/v1/rerank`` endpoint.""" + + id: ReadOnly[NotRequired[str]] + usage: ReadOnly[NotRequired[_InfinityRerankUsage]] + results: ReadOnly[Sequence[_InfinityRerankResult]] + + +def _parse_rerank_response(raw_response: httpx.Response) -> _InfinityRerankResponse: + """Read the untyped JSON body of an Infinity rerank response.""" + return raw_response.json() + + class InfinityRerankConfig(CohereRerankConfig): def get_complete_url( self, @@ -80,7 +107,7 @@ class InfinityRerankConfig(CohereRerankConfig): No transformation required, Infinity follows Cohere API response format """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _parse_rerank_response(raw_response) except Exception: raise InfinityError(message=raw_response.text, status_code=raw_response.status_code) diff --git a/litellm/llms/litellm_proxy/skills/code_execution.py b/litellm/llms/litellm_proxy/skills/code_execution.py index d435994ce20..f51213ca1c3 100644 --- a/litellm/llms/litellm_proxy/skills/code_execution.py +++ b/litellm/llms/litellm_proxy/skills/code_execution.py @@ -13,12 +13,57 @@ Generated files are returned directly in the response - no separate storage need import base64 import json +from collections.abc import Sequence from enum import Enum -from typing import Any, Final +from typing import Any, Final, Protocol + +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger +class _ToolCallFunction(Protocol): + """Function payload of an assistant tool call.""" + + name: str | None + arguments: str + + +class _ToolCall(Protocol): + """Tool call requested by the assistant on a chat completion choice.""" + + id: str + function: _ToolCallFunction + + +class _AssistantMessage(Protocol): + """Assistant message carried by a chat completion choice.""" + + content: str | None + tool_calls: Sequence[_ToolCall] | None + + +class _CompletionChoice(Protocol): + """Single choice of a chat completion response.""" + + finish_reason: str + message: _AssistantMessage + + +class _SandboxFile(TypedDict): + """File generated inside the sandbox during a code execution run.""" + + name: ReadOnly[str] + mime_type: ReadOnly[str] + content_base64: ReadOnly[str] + + +class _CodeExecutionArguments(TypedDict): + """Arguments the model passes to the `litellm_code_execution` tool.""" + + code: NotRequired[ReadOnly[str]] + + class LiteLLMInternalTools(str, Enum): """ Enum for internal LiteLLM tools that are injected into requests. @@ -30,7 +75,7 @@ class LiteLLMInternalTools(str, Enum): CODE_EXECUTION = "litellm_code_execution" -def get_litellm_code_execution_tool() -> dict[str, Any]: +def get_litellm_code_execution_tool() -> dict[str, object]: """ Returns the litellm_code_execution tool definition in OpenAI format. @@ -51,7 +96,7 @@ def get_litellm_code_execution_tool() -> dict[str, Any]: } -def get_litellm_code_execution_tool_anthropic() -> dict[str, Any]: +def get_litellm_code_execution_tool_anthropic() -> dict[str, object]: """ Returns the litellm_code_execution tool definition in Anthropic/messages API format. @@ -103,7 +148,7 @@ class CodeExecutionHandler: skill_files: dict[str, bytes], skill_id: str | None = None, **kwargs, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Execute an LLM call with automatic code execution handling. @@ -134,8 +179,8 @@ class CodeExecutionHandler: ) current_messages: Final = list(messages) - generated_files: Final[list[dict[str, Any]]] = [] # Files returned directly - execution_results: Final[list[dict]] = [] + generated_files: Final[list[dict[str, object]]] = [] # Files returned directly + execution_results: Final[list[dict[str, object]]] = [] executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) response: Any = None # Initialize to avoid possibly unbound error @@ -151,11 +196,12 @@ class CodeExecutionHandler: **kwargs, ) - assistant_message = response.choices[0].message - stop_reason = response.choices[0].finish_reason + choice: _CompletionChoice = response.choices[0] + assistant_message = choice.message + stop_reason: str = choice.finish_reason # Build assistant message for conversation history - assistant_msg_dict: dict[str, Any] = { + assistant_msg_dict: dict[str, object] = { "role": "assistant", "content": assistant_message.content, } @@ -190,8 +236,8 @@ class CodeExecutionHandler: if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: # Execute code in sandbox try: - args = json.loads(tool_call.function.arguments) - code = args.get("code", "") + args: _CodeExecutionArguments = json.loads(tool_call.function.arguments) + code: str = args.get("code", "") verbose_logger.debug("CodeExecutionHandler: Executing code (%s chars)", len(code)) @@ -202,13 +248,15 @@ class CodeExecutionHandler: verbose_logger.debug("CodeExecutionHandler: Execution result: %s", exec_result) + sandbox_files: Sequence[_SandboxFile] = exec_result["files"] + execution_results.append( { "iteration": iteration, "success": exec_result["success"], "output": exec_result["output"], "error": exec_result["error"], - "files": [f["name"] for f in exec_result["files"]], + "files": [f["name"] for f in sandbox_files], } ) @@ -216,9 +264,9 @@ class CodeExecutionHandler: tool_result = exec_result["output"] or "" # Collect generated files (returned directly, no storage) - if exec_result["files"]: + if sandbox_files: tool_result += "\n\nGenerated files:" - for f in exec_result["files"]: + for f in sandbox_files: file_content = base64.b64decode(f["content_base64"]) # Add to generated files list (returned in response) generated_files.append( diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 7ae438fd4cd..18f75f0e45c 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -8,10 +8,12 @@ response parsing, and streaming chunk parsing for models served with import datetime import json -from typing import Any, Final +from collections.abc import Sequence +from typing import Any, Final, TypedDict import httpx from pydantic import ValidationError +from typing_extensions import ReadOnly from litellm.llms.oci.chat.generic import ( _normalize_oci_finish_reason, @@ -46,6 +48,20 @@ from litellm.types.utils import ( ) +class _OpenAIToolCallFunction(TypedDict, total=False): + """The ``function`` block of an OpenAI-format assistant tool call.""" + + name: ReadOnly[str | None] + arguments: ReadOnly[str | dict[str, object]] + + +class _OpenAIToolCall(TypedDict, total=False): + """An entry of an OpenAI-format assistant message's ``tool_calls``.""" + + id: ReadOnly[str | None] + function: ReadOnly[_OpenAIToolCallFunction] + + def _extract_text_content(content: Any) -> str: """Return the plain-text representation of a message content value.""" if content is None: @@ -78,10 +94,10 @@ def adapt_messages_to_cohere_standard( """ # First pass: build tool_call_id → CohereToolCall so tool-result messages can # reference the originating call by name and parameters. - tool_call_lookup: Final[dict[str, CohereToolCall]] = {} + tool_call_lookup: Final[dict[str | None, CohereToolCall]] = {} for msg in messages: if msg.get("role") == "assistant": - tool_calls_raw: Any = msg.get("tool_calls") or [] + tool_calls_raw: Sequence[_OpenAIToolCall] = msg.get("tool_calls") or [] for tc in tool_calls_raw: tc_id = tc.get("id", "") raw_args = tc.get("function", {}).get("arguments", "{}") @@ -150,8 +166,22 @@ def adapt_messages_to_cohere_standard( return chat_history +class _OpenAIToolDefinitionFunction(TypedDict, total=False): + """The ``function`` block of an OpenAI-format tool definition.""" + + name: ReadOnly[str] + description: ReadOnly[str] + parameters: ReadOnly[dict[str, object]] + + +class _OpenAIToolDefinition(TypedDict, total=False): + """An entry of an OpenAI-format ``tools`` array.""" + + function: ReadOnly[_OpenAIToolDefinitionFunction] + + def adapt_tool_definitions_to_cohere_standard( - tools: list[dict[str, Any]], + tools: Sequence[_OpenAIToolDefinition], ) -> list[CohereTool]: """Adapt OpenAI-format tool definitions to the OCI Cohere format. diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 6e490f3ff15..449952217b5 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -4,16 +4,32 @@ Ollama /chat/completion calls handled in llm_http_handler.py [TODO]: migrate embeddings to a base handler as well. """ -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Any, Final, Protocol, TypedDict + +from typing_extensions import NotRequired, ReadOnly import litellm from litellm.types.utils import EmbeddingResponse +class TokenEncoder(Protocol): + """The tokenizer surface used to estimate prompt tokens.""" + + def encode(self, text: str, /) -> Sequence[int]: ... + + +class OllamaEmbeddingResponse(TypedDict): + """Body of an Ollama ``/api/embed`` response.""" + + embeddings: ReadOnly[list[list[float]]] + prompt_eval_count: ReadOnly[NotRequired[int]] + + def _prepare_ollama_embedding_payload( - model: str, prompts: list[str], optional_params: dict[str, Any] -) -> dict[str, Any]: - data: Final[dict[str, Any]] = {"model": model, "input": prompts} + model: str, prompts: list[str], optional_params: Mapping[str, object] +) -> dict[str, object]: + data: Final[dict[str, object]] = {"model": model, "input": prompts} special_optional_params: Final = ["truncate", "options", "keep_alive", "dimensions"] for k, v in optional_params.items(): @@ -27,12 +43,12 @@ def _prepare_ollama_embedding_payload( def _process_ollama_embedding_response( - response_json: dict, + response_json: OllamaEmbeddingResponse, prompts: list[str], model: str, model_response: EmbeddingResponse, logging_obj: Any, - encoding: Any, + encoding: TokenEncoder | None, ) -> EmbeddingResponse: output_data: Final = [] embeddings: Final[list[list[float]]] = response_json["embeddings"] @@ -72,7 +88,7 @@ async def ollama_aembeddings( model_response: EmbeddingResponse, optional_params: dict, logging_obj: Any, - encoding: Any, + encoding: TokenEncoder | None, ): if not api_base.endswith("/api/embed"): api_base += "/api/embed" @@ -80,7 +96,7 @@ async def ollama_aembeddings( data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) response: Final = await litellm.module_level_aclient.post(url=api_base, json=data) - response_json: Final = response.json() + response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( response_json=response_json, @@ -99,7 +115,7 @@ def ollama_embeddings( optional_params: dict, model_response: EmbeddingResponse, logging_obj: Any, - encoding: Any = None, + encoding: TokenEncoder | None = None, ): if not api_base.endswith("/api/embed"): api_base += "/api/embed" @@ -107,7 +123,7 @@ def ollama_embeddings( data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) response: Final = litellm.module_level_client.post(url=api_base, json=data) - response_json: Final = response.json() + response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( response_json=response_json, diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index 6fc50458aa3..44bd401c115 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -1,6 +1,8 @@ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, Literal import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( @@ -11,9 +13,11 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.containers.main import ( ContainerCreateOptionalRequestParams, ContainerFileListResponse, + ContainerFileObject, ContainerListResponse, ContainerObject, DeleteContainerResult, + ExpiresAfter, ) from litellm.types.router import GenericLiteLLMParams @@ -32,6 +36,46 @@ else: BaseLLMException = Any +class OpenAIContainerPayload(TypedDict): + """The JSON body OpenAI returns for a single container.""" + + id: ReadOnly[str] + object: ReadOnly[Literal["container"]] + created_at: ReadOnly[int] + status: ReadOnly[str] + expires_after: ReadOnly[ExpiresAfter | None] + last_active_at: ReadOnly[int | None] + name: ReadOnly[str | None] + + +class OpenAIContainerListPayload(TypedDict): + """The JSON body OpenAI returns for a page of containers.""" + + object: ReadOnly[Literal["list"]] + data: ReadOnly[list[ContainerObject]] + first_id: ReadOnly[str | None] + last_id: ReadOnly[str | None] + has_more: ReadOnly[bool] + + +class OpenAIContainerDeletedPayload(TypedDict): + """The JSON body OpenAI returns for a deleted container.""" + + id: ReadOnly[str] + object: ReadOnly[Literal["container.deleted"]] + deleted: ReadOnly[bool] + + +class OpenAIContainerFileListPayload(TypedDict): + """The JSON body OpenAI returns for a page of container files.""" + + object: ReadOnly[Literal["list"]] + data: ReadOnly[list[ContainerFileObject]] + first_id: ReadOnly[str | None] + last_id: ReadOnly[str | None] + has_more: ReadOnly[bool] + + class OpenAIContainerConfig(BaseContainerConfig): """Configuration class for OpenAI container API.""" @@ -87,7 +131,7 @@ class OpenAIContainerConfig(BaseContainerConfig): def transform_container_create_request( self, name: str, - container_create_optional_request_params: dict, + container_create_optional_request_params: Mapping[str, object], litellm_params: GenericLiteLLMParams, headers: dict, ) -> dict: @@ -111,7 +155,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerObject: """Transform the OpenAI container creation response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerPayload] = raw_response.json() # Transform the response data container_obj: Final = ContainerObject(**response_data) @@ -140,7 +184,7 @@ class OpenAIContainerConfig(BaseContainerConfig): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """Transform the container list request for OpenAI API. @@ -151,7 +195,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = api_base # Prepare query parameters - params: Final = {} + params: Final[dict[str, object]] = {} if after is not None: params["after"] = after if limit is not None: @@ -171,7 +215,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerListResponse: """Transform the OpenAI container list response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerListPayload] = raw_response.json() # Transform the response data container_list: Final = ContainerListResponse(**response_data) @@ -191,7 +235,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No additional data needed for GET request - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} return url, data @@ -201,7 +245,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerObject: """Transform the OpenAI container retrieve response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerPayload] = raw_response.json() # Transform the response data container_obj: Final = ContainerObject(**response_data) @@ -224,7 +268,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No data needed for DELETE request - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} return url, data @@ -234,7 +278,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> DeleteContainerResult: """Transform the OpenAI container delete response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerDeletedPayload] = raw_response.json() # Transform the response data delete_result: Final = DeleteContainerResult(**response_data) @@ -250,7 +294,7 @@ class OpenAIContainerConfig(BaseContainerConfig): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """Transform the container file list request for OpenAI API. @@ -262,7 +306,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files") # Prepare query parameters - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} if after is not None: params["after"] = after if limit is not None: @@ -282,7 +326,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerFileListResponse: """Transform the OpenAI container file list response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerFileListPayload] = raw_response.json() # Transform the response data file_list: Final = ContainerFileListResponse(**response_data) @@ -308,7 +352,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files/{encoded_file_id}/content") # No query parameters needed - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} return url, params diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 962dfe52c0a..deedf237a50 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.vertex_ai.common_utils import get_vertex_base_url @@ -25,6 +27,68 @@ else: LiteLLMLoggingObj = Any +class _AuthHeadersView(TypedDict): + """Auth headers read out of an untyped ``BaseVectorStoreAuthCredentials``.""" + + headers: ReadOnly[Mapping[str, str]] + + +class _RagContext(TypedDict, total=False): + """One context entry of a Vertex RAG ``:retrieveContexts`` response.""" + + text: ReadOnly[str] + sourceUri: ReadOnly[str] + sourceDisplayName: ReadOnly[str] + pageSpan: ReadOnly[Mapping[str, object]] + score: ReadOnly[float] + + +class _RagContexts(TypedDict, total=False): + """The ``contexts`` envelope wrapping the context list.""" + + contexts: ReadOnly[Sequence[_RagContext]] + + +class _RetrieveContextsBody(TypedDict, total=False): + """Body of a Vertex RAG ``:retrieveContexts`` response.""" + + contexts: ReadOnly[_RagContexts] + + +class _RetrieveContextsView(TypedDict): + """Typed view over the untyped ``:retrieveContexts`` JSON payload.""" + + body: ReadOnly[_RetrieveContextsBody] + + +class _RagCorpusBody(TypedDict, total=False): + """Body of a Vertex RAG ``ragCorpora`` create response.""" + + name: ReadOnly[str] + display_name: ReadOnly[str] + createTime: ReadOnly[str | int | float] + labels: ReadOnly[Mapping[str, str]] + + +class _RagCorpusView(TypedDict): + """Typed view over the untyped ``ragCorpora`` create JSON payload.""" + + body: ReadOnly[_RagCorpusBody] + + +class _RagSearchQuery(TypedDict, total=False): + """The ``query`` block of a Vertex RAG ``:retrieveContexts`` request.""" + + text: ReadOnly[str] + rag_retrieval_config: ReadOnly[Mapping[str, object]] + + +class _LoggedQueryView(TypedDict): + """The search query recovered from the logging object's call details.""" + + search_query: ReadOnly[str] + + class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): """ Configuration for Vertex AI Vector Store RAG API @@ -35,7 +99,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): def __init__(self): super().__init__() - def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials: # Get credentials and project info vertex_credentials: Final = self.get_vertex_ai_credentials(dict(litellm_params)) vertex_project: Final = self.get_vertex_ai_project(dict(litellm_params)) @@ -60,20 +124,23 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): "write": [("POST", "/ragCorpora")], } - def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict: + def validate_environment( + self, headers: dict[str, str], litellm_params: GenericLiteLLMParams | None + ) -> dict[str, str]: """ Validate and set up authentication for Vertex AI RAG API """ litellm_params = litellm_params or GenericLiteLLMParams() auth_headers: Final = self.get_auth_credentials(litellm_params.model_dump()) - headers.update(auth_headers.get("headers", {})) + auth_view: Final[_AuthHeadersView] = {"headers": auth_headers.get("headers", {})} + headers.update(auth_view["headers"]) return headers def get_complete_url( self, api_base: str | None, - litellm_params: dict, + litellm_params: dict[str, object], ) -> str: """ Get the Base endpoint for Vertex AI RAG API @@ -95,9 +162,9 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, - litellm_params: dict, - extra_body: dict[str, Any] | None = None, - ) -> tuple[str, dict[str, Any]]: + litellm_params: dict[str, object], + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: """ Transform search request for Vertex AI RAG API """ @@ -120,35 +187,34 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): # Just the corpus ID, construct full path full_rag_corpus = f"projects/{vertex_project}/locations/{vertex_location}/ragCorpora/{vector_store_id}" - # Build the request body for Vertex AI RAG API - request_body: Final[dict[str, Any]] = { - "vertex_rag_store": {"rag_resources": [{"rag_corpus": full_rag_corpus}]}, - "query": {"text": query}, - } - ######################################################### # Update logging object with details of the request ######################################################### litellm_logging_obj.model_call_details["query"] = query # Add optional parameters + rag_retrieval_config: Final[dict[str, object]] = {} max_num_results: Final = vector_store_search_optional_params.get("max_num_results") if max_num_results is not None: - request_body["query"]["rag_retrieval_config"] = {"top_k": max_num_results} + rag_retrieval_config["top_k"] = max_num_results # Add filters if provided - filters: Final = vector_store_search_optional_params.get("filters") + filters: Final[object] = vector_store_search_optional_params.get("filters") if filters is not None: - if "rag_retrieval_config" not in request_body["query"]: - request_body["query"]["rag_retrieval_config"] = {} - request_body["query"]["rag_retrieval_config"]["filter"] = filters + rag_retrieval_config["filter"] = filters # Add ranking options if provided - ranking_options: Final = vector_store_search_optional_params.get("ranking_options") + ranking_options: Final[object] = vector_store_search_optional_params.get("ranking_options") if ranking_options is not None: - if "rag_retrieval_config" not in request_body["query"]: - request_body["query"]["rag_retrieval_config"] = {} - request_body["query"]["rag_retrieval_config"]["ranking"] = ranking_options + rag_retrieval_config["ranking"] = ranking_options + + query_body: Final[_RagSearchQuery] = ( + {"text": query, "rag_retrieval_config": rag_retrieval_config} if rag_retrieval_config else {"text": query} + ) + request_body: Final[dict[str, object]] = { + "vertex_rag_store": {"rag_resources": [{"rag_corpus": full_rag_corpus}]}, + "query": query_body, + } return url, request_body @@ -159,12 +225,13 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): Transform Vertex AI RAG API response to standard vector store search response """ try: - response_json: Final = response.json() + response_view: Final[_RetrieveContextsView] = {"body": response.json()} + response_json: Final = response_view["body"] # Extract contexts from Vertex AI response - handle nested structure contexts: Final = response_json.get("contexts", {}).get("contexts", []) # Transform contexts to standard format - search_results: Final = [] + search_results: Final[list[VectorStoreSearchResult]] = [] for context in contexts: content = [ VectorStoreResultContent( @@ -202,9 +269,12 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): ) search_results.append(result) + query_view: Final[_LoggedQueryView] = { + "search_query": litellm_logging_obj.model_call_details.get("query", "") + } return VectorStoreSearchResponse( object="vector_store.search_results.page", - search_query=litellm_logging_obj.model_call_details.get("query", ""), + search_query=query_view["search_query"], data=search_results, ) @@ -219,14 +289,14 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): self, vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, api_base: str, - ) -> tuple[str, dict[str, Any]]: + ) -> tuple[str, dict[str, object]]: """ Transform create request for Vertex AI RAG Corpus """ url: Final = f"{api_base}/ragCorpora" # Base URL for creating RAG corpus # Build the request body for Vertex AI RAG Corpus creation - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "display_name": vector_store_create_optional_params.get("name", "litellm-vector-store"), "description": "Vector store created via LiteLLM", } @@ -243,7 +313,8 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): Transform Vertex AI RAG Corpus creation response to standard vector store response """ try: - response_json: Final = response.json() + response_view: Final[_RagCorpusView] = {"body": response.json()} + response_json: Final = response_view["body"] # Extract the corpus ID from the response name corpus_name: Final = response_json.get("name", "") diff --git a/litellm/models/base.py b/litellm/models/base.py index 7eedf10212e..8125bfd0205 100644 --- a/litellm/models/base.py +++ b/litellm/models/base.py @@ -33,6 +33,6 @@ class DomainModel(BaseModel): return cls(**record.dict()) return cls(**dict(record)) - def to_db_dict(self, exclude_unset: bool = False) -> dict[str, Any]: + def to_db_dict(self, exclude_unset: bool = False) -> dict[str, object]: """Convert domain model to a dictionary for database operations.""" return self.model_dump(exclude_none=True, exclude_unset=exclude_unset) diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py b/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py index c09106273e1..150900e7ff2 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py @@ -36,7 +36,10 @@ a healed fleet has no null rows and the backfill exits after one query. import json from collections import Counter -from typing import Any, Final, Literal +from collections.abc import Mapping, Sequence +from typing import Final, Literal, Protocol + +from pydantic import JsonValue from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.db import _decode_oauth_payload, decrypt_credentials @@ -55,9 +58,59 @@ BackfillRule = Literal[ _BACKFILL_AUDIT_ACTOR: Final = "oauth2_flow_backfill" -def _decrypted_credentials(raw_credentials: Any) -> MCPCredentials | None: +class _MCPServerRow(Protocol): + """The ``LiteLLM_MCPServerTable`` columns this backfill reads.""" + + @property + def server_id(self) -> str: ... + + @property + def authorization_url(self) -> str | None: ... + + @property + def registration_url(self) -> str | None: ... + + @property + def token_url(self) -> str | None: ... + + @property + def credentials(self) -> str | Mapping[str, JsonValue] | None: ... + + +class _MCPUserCredentialRow(Protocol): + """The ``LiteLLM_MCPUserCredentials`` columns this backfill reads.""" + + @property + def server_id(self) -> str: ... + + @property + def credential_b64(self) -> str: ... + + +class _MCPServerTable(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_MCPServerRow]: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, str]) -> object: ... + + +class _MCPUserCredentialsTable(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_MCPUserCredentialRow]: ... + + +def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable: + """The MCP server table, typed so the untyped prisma client surface stops here.""" + return prisma_client.db.litellm_mcpservertable + + +def _mcp_user_credentials_table(prisma_client: PrismaClient) -> _MCPUserCredentialsTable: + """The per-user MCP credential table, typed so the untyped prisma client surface stops here.""" + return prisma_client.db.litellm_mcpusercredentials + + +def _decrypted_credentials(raw_credentials: str | Mapping[str, JsonValue] | None) -> MCPCredentials | None: if raw_credentials is None: return None + parsed: JsonValue | Mapping[str, JsonValue] if isinstance(raw_credentials, str): try: parsed = json.loads(raw_credentials) @@ -92,14 +145,14 @@ def classify_null_flow_row( async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[BackfillRule, int]: """Classify every ``auth_type=oauth2`` row whose ``oauth2_flow`` is null; stamp the provable ones, warn on the ambiguous ones, and return counts per rule.""" - null_rows: Final[list[Any]] = await prisma_client.db.litellm_mcpservertable.find_many( + null_rows: Final[Sequence[_MCPServerRow]] = await _mcp_server_table(prisma_client).find_many( where={"auth_type": "oauth2", "oauth2_flow": None}, ) if not null_rows: return {} server_ids: Final = [row.server_id for row in null_rows] - token_rows: Final[list[Any]] = await prisma_client.db.litellm_mcpusercredentials.find_many( + token_rows: Final[Sequence[_MCPUserCredentialRow]] = await _mcp_user_credentials_table(prisma_client).find_many( where={"server_id": {"in": server_ids}}, ) server_ids_with_oauth_tokens: Final[set[str]] = { @@ -141,7 +194,7 @@ async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[Backfi stamped_flows: Final = {flow for _, (flow, _) in classified if flow is not None} for stamped_flow in stamped_flows: server_ids_for_flow = [row.server_id for row, (row_flow, _) in classified if row_flow == stamped_flow] - await prisma_client.db.litellm_mcpservertable.update_many( + await _mcp_server_table(prisma_client).update_many( where={"server_id": {"in": server_ids_for_flow}, "oauth2_flow": None}, data={"oauth2_flow": stamped_flow, "updated_by": _BACKFILL_AUDIT_ACTOR}, ) diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index 9672383a572..ecaaf35e817 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -1,5 +1,9 @@ import json -from typing import Final +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import Final, Protocol + +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -7,18 +11,73 @@ from litellm.proxy.utils import PrismaClient from litellm.repositories.table_repositories import MCPToolsetRepository from litellm.types.mcp_server.mcp_toolset import ( MCPToolset, + MCPToolsetTool, NewMCPToolsetRequest, UpdateMCPToolsetRequest, ) -def _toolset_from_row(row) -> MCPToolset: +class MCPToolsetFields(TypedDict): + """The ``MCPToolset`` constructor keywords a toolset row expands into.""" + + toolset_id: ReadOnly[str] + toolset_name: ReadOnly[str] + description: NotRequired[ReadOnly[str | None]] + tools: NotRequired[ReadOnly[list[MCPToolsetTool]]] + created_at: NotRequired[ReadOnly[datetime | None]] + created_by: NotRequired[ReadOnly[str | None]] + updated_at: NotRequired[ReadOnly[datetime | None]] + updated_by: NotRequired[ReadOnly[str | None]] + + +class MCPToolsetRowData(TypedDict): + """A toolset table row, whose ``tools`` column is stored as JSON.""" + + toolset_id: ReadOnly[str] + toolset_name: ReadOnly[str] + description: NotRequired[ReadOnly[str | None]] + tools: NotRequired[ReadOnly[str | list[MCPToolsetTool]]] + created_at: NotRequired[ReadOnly[datetime | None]] + created_by: NotRequired[ReadOnly[str | None]] + updated_at: NotRequired[ReadOnly[datetime | None]] + updated_by: NotRequired[ReadOnly[str | None]] + + +class MCPToolsetRow(Protocol): + """A row of the toolset table, as the prisma client returns it.""" + + def model_dump(self) -> MCPToolsetRowData: ... + + +class MCPToolsetTable(Protocol): + """The prisma table actions this module runs against the toolset table.""" + + async def create(self, data: Mapping[str, object]) -> MCPToolsetRow: ... + + async def find_unique(self, where: Mapping[str, object]) -> MCPToolsetRow | None: ... + + async def find_first(self, where: Mapping[str, object]) -> MCPToolsetRow | None: ... + + async def find_many(self, where: Mapping[str, object]) -> Sequence[MCPToolsetRow]: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> MCPToolsetRow: ... + + async def delete(self, where: Mapping[str, object]) -> MCPToolsetRow: ... + + +def _toolset_table(prisma_client: PrismaClient) -> MCPToolsetTable: + """The toolset table actions of the prisma client.""" + return MCPToolsetRepository(prisma_client).table + + +def _toolset_from_row(row: MCPToolsetRow) -> MCPToolset: data: Final = row.model_dump() - tools = data.get("tools") or [] - if isinstance(tools, str): - tools = json.loads(tools) - data["tools"] = tools - return MCPToolset(**data) + tools: Final = data.get("tools") or [] + resolved: Final[MCPToolsetFields] = { + **data, + "tools": json.loads(tools) if isinstance(tools, str) else tools, + } + return MCPToolset(**resolved) async def create_mcp_toolset( @@ -31,7 +90,7 @@ async def create_mcp_toolset( data_dict["tools"] = json.dumps(data_dict.get("tools", [])) data_dict["created_by"] = touched_by data_dict["updated_by"] = touched_by - row: Final = await MCPToolsetRepository(prisma_client).table.create(data=data_dict) + row: Final = await _toolset_table(prisma_client).create(data=data_dict) return _toolset_from_row(row) @@ -39,7 +98,7 @@ async def get_mcp_toolset( prisma_client: PrismaClient, toolset_id: str, ) -> MCPToolset | None: - row: Final = await MCPToolsetRepository(prisma_client).table.find_unique(where={"toolset_id": toolset_id}) + row: Final = await _toolset_table(prisma_client).find_unique(where={"toolset_id": toolset_id}) if row is None: return None return _toolset_from_row(row) @@ -47,13 +106,11 @@ async def get_mcp_toolset( async def list_mcp_toolsets( prisma_client: PrismaClient, - toolset_ids: list[str] | None = None, -) -> list[MCPToolset]: + toolset_ids: Sequence[str] | None = None, +) -> Sequence[MCPToolset]: try: - where = {} - if toolset_ids is not None: - where = {"toolset_id": {"in": toolset_ids}} - rows: Final = await MCPToolsetRepository(prisma_client).table.find_many(where=where) + where: Final[Mapping[str, object]] = {} if toolset_ids is None else {"toolset_id": {"in": toolset_ids}} + rows: Final = await _toolset_table(prisma_client).find_many(where=where) return [_toolset_from_row(r) for r in rows] except Exception as e: verbose_proxy_logger.warning("litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - %s", e) @@ -64,7 +121,7 @@ async def get_mcp_toolset_by_name( prisma_client: PrismaClient, toolset_name: str, ) -> MCPToolset | None: - row: Final = await MCPToolsetRepository(prisma_client).table.find_first(where={"toolset_name": toolset_name}) + row: Final = await _toolset_table(prisma_client).find_first(where={"toolset_name": toolset_name}) if row is None: return None return _toolset_from_row(row) @@ -80,7 +137,7 @@ async def update_mcp_toolset( data_dict["tools"] = json.dumps(data_dict["tools"]) data_dict["updated_by"] = touched_by try: - row: Final = await MCPToolsetRepository(prisma_client).table.update( + row: Final = await _toolset_table(prisma_client).update( where={"toolset_id": data.toolset_id}, data=data_dict, ) @@ -98,7 +155,7 @@ async def delete_mcp_toolset( toolset_id: str, ) -> MCPToolset | None: try: - row: Final = await MCPToolsetRepository(prisma_client).table.delete(where={"toolset_id": toolset_id}) + row: Final = await _toolset_table(prisma_client).delete(where={"toolset_id": toolset_id}) except Exception as e: from prisma.errors import RecordNotFoundError diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 64de6827679..edb945fc674 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -12,6 +12,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -83,9 +84,17 @@ class AgentTableClient(Protocol): async def delete(self, where: Mapping[str, object]) -> AgentRecord: ... +class _AgentsRepositoryView(Protocol): + @property + def table(self) -> AgentTableClient: ... + + +def _agents_table_of(repository: _AgentsRepositoryView) -> AgentTableClient: + return repository.table + + def agents_table(prisma_client: PrismaClient) -> AgentTableClient: - table: Final[AgentTableClient] = AgentsRepository(prisma_client).table - return table + return _agents_table_of(AgentsRepository(prisma_client)) class ObjectPermissionGrantRecord(Protocol): @@ -99,9 +108,17 @@ class ObjectPermissionTableClient(Protocol): async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... +class _ObjectPermissionRepositoryView(Protocol): + @property + def table(self) -> ObjectPermissionTableClient: ... + + +def _object_permission_table_of(repository: _ObjectPermissionRepositoryView) -> ObjectPermissionTableClient: + return repository.table + + def object_permission_table(prisma_client: PrismaClient) -> ObjectPermissionTableClient: - table: Final[ObjectPermissionTableClient] = ObjectPermissionRepository(prisma_client).table - return table + return _object_permission_table_of(ObjectPermissionRepository(prisma_client)) class GrantMigrationResult(NamedTuple): @@ -283,19 +300,21 @@ class AgentRegistry: agent_name: Final = agent.get("agent_name") # Serialize litellm_params - litellm_params_obj: Final[Any] = agent.get("litellm_params", {}) - if hasattr(litellm_params_obj, "model_dump"): - litellm_params_dict = litellm_params_obj.model_dump() - else: - litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} + litellm_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("litellm_params", {}) + litellm_params_dict: Final[Mapping[str, object]] = ( + litellm_params_obj.model_dump() + if isinstance(litellm_params_obj, SupportsModelDump) + else (dict(litellm_params_obj) if litellm_params_obj else {}) + ) litellm_params: Final[str] = safe_dumps(litellm_params_dict) # Serialize agent_card_params - agent_card_params_obj: Final[Any] = agent.get("agent_card_params", {}) - if hasattr(agent_card_params_obj, "model_dump"): - agent_card_params_dict = agent_card_params_obj.model_dump() - else: - agent_card_params_dict = dict(agent_card_params_obj) if agent_card_params_obj else {} + agent_card_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("agent_card_params", {}) + agent_card_params_dict: Final[Mapping[str, object]] = ( + agent_card_params_obj.model_dump() + if isinstance(agent_card_params_obj, SupportsModelDump) + else (dict(agent_card_params_obj) if agent_card_params_obj else {}) + ) agent_card_params: Final[str] = safe_dumps(agent_card_params_dict) # Handle object_permission (MCP tool access for agent) @@ -386,15 +405,13 @@ class AgentRegistry: The patched agent """ try: - existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) - if existing_agent is not None: - existing_agent = dict(existing_agent) - - if existing_agent is None: + existing_record: Final = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) + if existing_record is None: raise Exception(f"Agent with ID {agent_id} not found") + existing_agent: Final[Mapping[str, object]] = dict(existing_record) augment_agent: Final = {**existing_agent, **agent} - update_data: Final[dict[str, Any]] = {} + update_data: Final[dict[str, object]] = {} if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if augment_agent.get("litellm_params"): @@ -418,7 +435,7 @@ class AgentRegistry: update_data["extra_headers"] = extra_headers_value if extra_headers_value is not None else [] if agent.get("object_permission") is not None: agent_copy: Final = dict(augment_agent) - existing_object_permission_id: Final = existing_agent.get("object_permission_id") + existing_object_permission_id: Final = existing_record.object_permission_id object_permission_id: Final = await handle_update_object_permission_common( agent_copy, existing_object_permission_id, @@ -460,19 +477,21 @@ class AgentRegistry: agent_name: Final = agent.get("agent_name") # Serialize litellm_params - litellm_params_obj: Final[Any] = agent.get("litellm_params", {}) - if hasattr(litellm_params_obj, "model_dump"): - litellm_params_dict = litellm_params_obj.model_dump() - else: - litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} + litellm_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("litellm_params", {}) + litellm_params_dict: Final[Mapping[str, object]] = ( + litellm_params_obj.model_dump() + if isinstance(litellm_params_obj, SupportsModelDump) + else (dict(litellm_params_obj) if litellm_params_obj else {}) + ) litellm_params: Final[str] = safe_dumps(litellm_params_dict) # Serialize agent_card_params - agent_card_params_obj: Final[Any] = agent.get("agent_card_params", {}) - if hasattr(agent_card_params_obj, "model_dump"): - agent_card_params_dict = agent_card_params_obj.model_dump() - else: - agent_card_params_dict = dict(agent_card_params_obj) if agent_card_params_obj else {} + agent_card_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("agent_card_params", {}) + agent_card_params_dict: Final[Mapping[str, object]] = ( + agent_card_params_obj.model_dump() + if isinstance(agent_card_params_obj, SupportsModelDump) + else (dict(agent_card_params_obj) if agent_card_params_obj else {}) + ) agent_card_params: Final[str] = safe_dumps(agent_card_params_dict) # Serialize static_headers for update diff --git a/litellm/proxy/client/cli/commands/credentials.py b/litellm/proxy/client/cli/commands/credentials.py index c550b39d33f..2beff05b375 100644 --- a/litellm/proxy/client/cli/commands/credentials.py +++ b/litellm/proxy/client/cli/commands/credentials.py @@ -1,14 +1,42 @@ import json +from collections.abc import Mapping, Sequence from typing import Final, Literal import click import requests import rich from rich.table import Table +from typing_extensions import NotRequired, ReadOnly, TypedDict from ...credentials import CredentialsManagementClient +class _CliContext(TypedDict): + """Values the top-level CLI group stores on the click context.""" + + base_url: ReadOnly[str] + api_key: ReadOnly[str | None] + + +class _CliContextView(TypedDict): + obj: ReadOnly[_CliContext] + + +class _CredentialRow(TypedDict): + """Single credential entry as returned by ``GET /credentials``.""" + + credential_name: ReadOnly[NotRequired[str]] + credential_info: ReadOnly[NotRequired[Mapping[str, object]]] + + +class _CredentialRowsView(TypedDict): + rows: ReadOnly[Sequence[_CredentialRow]] + + +class _JsonBodyView(TypedDict): + body: ReadOnly[object] + + @click.group() def credentials(): """Manage credentials for the LiteLLM proxy server""" @@ -25,7 +53,8 @@ def credentials(): @click.pass_context def list(ctx: click.Context, output_format: Literal["table", "json"]): """List all credentials""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final[_CliContextView] = {"obj": ctx.obj} + client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) response: Final = client.list() assert isinstance(response, dict) @@ -39,7 +68,8 @@ def list(ctx: click.Context, output_format: Literal["table", "json"]): table.add_column("Custom LLM Provider", style="green") # Add rows - for cred in response.get("credentials", []): + credential_rows: Final[_CredentialRowsView] = {"rows": response.get("credentials", [])} + for cred in credential_rows["rows"]: info = cred.get("credential_info", {}) table.add_row( str(cred.get("credential_name", "")), @@ -66,7 +96,8 @@ def list(ctx: click.Context, output_format: Literal["table", "json"]): @click.pass_context def create(ctx: click.Context, credential_name: str, info: str, values: str): """Create a new credential""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final[_CliContextView] = {"obj": ctx.obj} + client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) try: credential_info: Final = json.loads(info) credential_values: Final = json.loads(values) @@ -79,8 +110,8 @@ def create(ctx: click.Context, credential_name: str, info: str, values: str): except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) try: - error_body: Final = e.response.json() - rich.print_json(data=error_body) + error_body: Final[_JsonBodyView] = {"body": e.response.json()} + rich.print_json(data=error_body["body"]) except json.JSONDecodeError: click.echo(e.response.text, err=True) raise click.Abort() @@ -91,15 +122,16 @@ def create(ctx: click.Context, credential_name: str, info: str, values: str): @click.pass_context def delete(ctx: click.Context, credential_name: str): """Delete a credential by name""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final[_CliContextView] = {"obj": ctx.obj} + client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) try: response: Final = client.delete(credential_name) rich.print_json(data=response) except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) try: - error_body: Final = e.response.json() - rich.print_json(data=error_body) + error_body: Final[_JsonBodyView] = {"body": e.response.json()} + rich.print_json(data=error_body["body"]) except json.JSONDecodeError: click.echo(e.response.text, err=True) raise click.Abort() @@ -110,6 +142,7 @@ def delete(ctx: click.Context, credential_name: str): @click.pass_context def get(ctx: click.Context, credential_name: str): """Get a credential by name""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final[_CliContextView] = {"obj": ctx.obj} + client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) response: Final = client.get(credential_name) rich.print_json(data=response) diff --git a/litellm/proxy/client/cli/commands/teams.py b/litellm/proxy/client/cli/commands/teams.py index e814ac84ebb..1212e194bde 100644 --- a/litellm/proxy/client/cli/commands/teams.py +++ b/litellm/proxy/client/cli/commands/teams.py @@ -1,21 +1,48 @@ """Team management commands for LiteLLM CLI.""" +from collections.abc import Mapping, Sequence from typing import Any, Final import click import requests from rich.console import Console from rich.table import Table +from typing_extensions import ReadOnly, TypedDict from litellm.proxy.client import Client +class _CliContext(TypedDict): + """The proxy connection settings the CLI group stores on the click context.""" + + base_url: ReadOnly[str] + api_key: ReadOnly[str | None] + + +def _cli_context(ctx: click.Context) -> _CliContext: + """The proxy connection settings the CLI group stored on the click context.""" + ctx_obj: Final[_CliContext] = ctx.obj + return ctx_obj + + +def _proxy_client(ctx: click.Context) -> Client: + """A proxy client for the base URL and API key on the click context.""" + ctx_obj: Final = _cli_context(ctx) + return Client(ctx_obj["base_url"], ctx_obj["api_key"]) + + +def _http_error_detail(error: requests.exceptions.HTTPError) -> object: + """The ``detail`` the proxy reported for a failed request.""" + error_body: Final[Mapping[str, object]] = error.response.json() + return error_body.get("detail", "Unknown error") + + @click.group() def teams(): """Manage teams and team assignments""" -def display_teams_table(teams: list[dict[str, Any]]) -> None: +def display_teams_table(teams: Sequence[dict[str, Any]]) -> None: """Display teams in a formatted table""" console: Final = Console() @@ -33,8 +60,8 @@ def display_teams_table(teams: list[dict[str, Any]]) -> None: for i, team in enumerate(teams): team_alias = team.get("team_alias") or "N/A" - team_id = team.get("team_id", "N/A") - models = team.get("models", []) + team_id: str = team.get("team_id", "N/A") + models: Sequence[str] = team.get("models", []) max_budget = team.get("max_budget") # Format models list @@ -64,7 +91,7 @@ def display_teams_table(teams: list[dict[str, Any]]) -> None: @click.pass_context def list(ctx: click.Context): """List teams that you belong to""" - client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"]) + client: Final = _proxy_client(ctx) try: # Use list() for simpler response structure (returns array directly) @@ -72,8 +99,7 @@ def list(ctx: click.Context): display_teams_table(teams) except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) - error_body: Final = e.response.json() - click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) + click.echo(f"Details: {_http_error_detail(e)}", err=True) raise click.Abort() except Exception as e: click.echo(f"Error: {e}", err=True) @@ -84,7 +110,7 @@ def list(ctx: click.Context): @click.pass_context def available(ctx: click.Context): """List teams that are available to join""" - client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"]) + client: Final = _proxy_client(ctx) try: teams: Final = client.teams.get_available() @@ -96,8 +122,7 @@ def available(ctx: click.Context): click.echo("No available teams to join.") except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) - error_body: Final = e.response.json() - click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) + click.echo(f"Details: {_http_error_detail(e)}", err=True) except Exception as e: click.echo(f"Error: {e}", err=True) raise click.Abort() @@ -108,8 +133,8 @@ def available(ctx: click.Context): @click.pass_context def assign_key(ctx: click.Context, team_id: str | None): """Assign your current CLI key to a team""" - client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"]) - api_key: Final = ctx.obj["api_key"] + client: Final = _proxy_client(ctx) + api_key: Final = _cli_context(ctx)["api_key"] if not api_key: click.echo("No API key found. Please login first using 'litellm login'") @@ -145,7 +170,7 @@ def assign_key(ctx: click.Context, team_id: str | None): teams = client.teams.list() for team in teams: if team.get("team_id") == team_id: - models = team.get("models", []) + models: Sequence[str] = team.get("models", []) if models: click.echo(f"You can now access models: {', '.join(models)}") else: @@ -154,8 +179,7 @@ def assign_key(ctx: click.Context, team_id: str | None): except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) - error_body: Final = e.response.json() - click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) + click.echo(f"Details: {_http_error_detail(e)}", err=True) raise click.Abort() except Exception as e: click.echo(f"Error: {e}", err=True) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 9379a8577a3..a28d03eebf0 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,6 +1,6 @@ import copy import os -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias @@ -525,8 +525,8 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset( def sanitize_openai_provider_metadata( - metadata: dict[str, Any] | None, -) -> dict[str, str] | None: + metadata: Mapping[str, object] | None, +) -> Mapping[str, object] | None: """ Keep only provider-safe OpenAI metadata entries (string keys -> string values). @@ -644,7 +644,7 @@ def process_callback(_callback: str, callback_type: str, environment_variables: return {"name": _callback, "variables": env_vars_dict, "type": callback_type} -def normalize_callback_names(callbacks: Iterable[Any]) -> list[Any]: +def normalize_callback_names(callbacks: Iterable[object] | None) -> list[object]: if callbacks is None: return [] return [c.lower() if isinstance(c, str) else c for c in callbacks] @@ -674,7 +674,7 @@ def decrypt_callback_vars(metadata: Any) -> Any: return _transform_callback_vars(metadata, _decrypt_or_passthrough) -def _transform_callback_vars(metadata: Any, transform: Callable[[str, Any], Any]) -> Any: +def _transform_callback_vars(metadata: object, transform: Callable[[str, Any], Any]) -> object: if not isinstance(metadata, dict): return metadata out: Final = copy.deepcopy(metadata) @@ -704,7 +704,7 @@ def is_sensitive_callback_key( return _CALLBACK_VAR_MASKER.is_sensitive_key(key) -def _encrypt_if_plaintext(key: str, value: Any) -> Any: +def _encrypt_if_plaintext(key: str, value: object) -> object: if not isinstance(value, str) or not value: return value if not is_sensitive_callback_key(key): @@ -725,7 +725,7 @@ def _encrypt_if_plaintext(key: str, value: Any) -> Any: return value -def _decrypt_or_passthrough(key: str, value: Any) -> Any: +def _decrypt_or_passthrough(key: str, value: object) -> object: if not isinstance(value, str) or not value: return value if not value.startswith(_CALLBACK_VAR_ENCRYPTED_PREFIX): diff --git a/litellm/proxy/common_utils/get_routes.py b/litellm/proxy/common_utils/get_routes.py index 2118a6610b4..5aeb755071c 100644 --- a/litellm/proxy/common_utils/get_routes.py +++ b/litellm/proxy/common_utils/get_routes.py @@ -2,63 +2,78 @@ Utility class for getting routes from a FastAPI app. """ +from collections.abc import Sequence from typing import Any, Final from starlette.routing import BaseRoute +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger +class RouteInfo(TypedDict, total=False): + """One entry of the app's route listing.""" + + path: ReadOnly[object] + methods: ReadOnly[object] + name: ReadOnly[object] + endpoint: ReadOnly[str | None] + mounted_app: ReadOnly[bool] + + class GetRoutes: @staticmethod def get_app_routes( route: BaseRoute, endpoint_route: Any, - ) -> list[dict[str, Any]]: + ) -> Sequence[RouteInfo]: """ Get routes for a regular route. """ - routes: Final[list[dict[str, Any]]] = [] - route_info: Final = { + route_info: Final[RouteInfo] = { "path": getattr(route, "path", None), "methods": getattr(route, "methods", None), "name": getattr(route, "name", None), "endpoint": (endpoint_route.__name__ if getattr(route, "endpoint", None) else None), } - routes.append(route_info) - return routes + return [route_info] @staticmethod def get_routes_for_mounted_app( route: BaseRoute, - ) -> list[dict[str, Any]]: + ) -> Sequence[RouteInfo]: """ Get routes for a mounted sub-application. """ - routes: Final[list[dict[str, Any]]] = [] + routes: Final[list[RouteInfo]] = [] mount_path: Final = getattr(route, "path", "") - sub_app: Final = getattr(route, "app", None) - if sub_app and hasattr(sub_app, "routes"): - for sub_route in sub_app.routes: - # Get endpoint - either from endpoint attribute or app attribute - endpoint_func = getattr(sub_route, "endpoint", None) or getattr(sub_route, "app", None) + for sub_route in GetRoutes._mounted_app_routes(route): + endpoint_func: object = getattr(sub_route, "endpoint", None) or getattr(sub_route, "app", None) - if endpoint_func is not None: - sub_route_path = getattr(sub_route, "path", "") - full_path = mount_path.rstrip("/") + sub_route_path + if endpoint_func is not None: + sub_route_path = getattr(sub_route, "path", "") + full_path = mount_path.rstrip("/") + sub_route_path - route_info = { - "path": full_path, - "methods": getattr(sub_route, "methods", ["GET", "POST"]), - "name": getattr(sub_route, "name", None), - "endpoint": GetRoutes._safe_get_endpoint_name(endpoint_func), - "mounted_app": True, - } - routes.append(route_info) + route_info: RouteInfo = { + "path": full_path, + "methods": getattr(sub_route, "methods", ["GET", "POST"]), + "name": getattr(sub_route, "name", None), + "endpoint": GetRoutes._safe_get_endpoint_name(endpoint_func), + "mounted_app": True, + } + routes.append(route_info) return routes @staticmethod - def _safe_get_endpoint_name(endpoint_function: Any) -> str | None: + def _mounted_app_routes(route: BaseRoute) -> Sequence[BaseRoute]: + """The routes of the sub-application mounted at ``route``, if it mounts one.""" + sub_app: Final[object] = getattr(route, "app", None) + if sub_app and hasattr(sub_app, "routes"): + return getattr(sub_app, "routes") + return () + + @staticmethod + def _safe_get_endpoint_name(endpoint_function: object) -> str | None: """ Safely get the name of the endpoint function. """ @@ -66,7 +81,7 @@ class GetRoutes: if hasattr(endpoint_function, "__name__"): return getattr(endpoint_function, "__name__") elif hasattr(endpoint_function, "__class__") and hasattr(endpoint_function.__class__, "__name__"): - return getattr(endpoint_function.__class__, "__name__") + return endpoint_function.__class__.__name__ else: return None except Exception: diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 93d51bdd461..a8a71134034 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -40,30 +40,31 @@ class UserApiKeyCache(DualCache): @overload def get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: object, + parent_otel_span: object = None, local_only: bool = False, *, model_type: type[T], - **kwargs: Any, + **kwargs: object, ) -> T | None: ... @overload def get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: object, + parent_otel_span: object = None, local_only: bool = False, - **kwargs: Any, + model_type: None = None, + **kwargs: object, ) -> Any: ... def get_cache( self, - key, - parent_otel_span=None, + key: object, + parent_otel_span: object = None, local_only: bool = False, model_type: type[BaseModel] | None = None, - **kwargs, + **kwargs: object, ) -> Any | BaseModel | None: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) @@ -85,30 +86,31 @@ class UserApiKeyCache(DualCache): @overload async def async_get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: object, + parent_otel_span: object = None, local_only: bool = False, *, model_type: type[T], - **kwargs: Any, + **kwargs: object, ) -> T | None: ... @overload async def async_get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: object, + parent_otel_span: object = None, local_only: bool = False, - **kwargs: Any, + model_type: None = None, + **kwargs: object, ) -> Any: ... async def async_get_cache( self, - key, - parent_otel_span=None, + key: object, + parent_otel_span: object = None, local_only: bool = False, model_type: type[BaseModel] | None = None, - **kwargs, + **kwargs: object, ) -> Any | BaseModel | None: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) @@ -129,17 +131,17 @@ class UserApiKeyCache(DualCache): return None return decoded - def set_cache(self, key, value, local_only: bool = False, **kwargs): + def set_cache(self, key: object, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) - payload: Final = CacheCodec.serialize(value, model_type=model_type) + payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) - async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): + async def async_set_cache(self, key: object, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) - payload: Final = CacheCodec.serialize(value, model_type=model_type) + payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) - async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs) -> None: + async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs: object) -> None: """ Batch writes with the same Codec boundary as ``async_set_cache`` without ``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged. diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 5aeb52be535..9cc4369fea6 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -6,11 +6,14 @@ otherwise PrismaClient uses the writer-only PrismaWrapper directly. import os from collections.abc import Callable -from typing import Any, Final +from datetime import timedelta +from typing import Any, Final, TypeAlias from litellm._logging import verbose_proxy_logger from litellm.proxy.db.prisma_client import PrismaWrapper +ConnectTimeout: TypeAlias = "int | timedelta | None" + # Per-model action methods that read from the database. These are routed to # the read replica when one is configured. _MODEL_READ_METHODS: Final = frozenset( @@ -31,6 +34,11 @@ _MODEL_READ_METHODS: Final = frozenset( _TOP_LEVEL_READ_METHODS: Final = frozenset({"query_first", "query_raw"}) +def _dynamic_attr(target: object, name: str) -> object: + """Fetch `name` off an unstubbed Prisma object as an opaque value.""" + return getattr(target, name) + + class _RoutedActions: """Per-model accessor that sends reads to the reader and writes to the writer. @@ -43,18 +51,18 @@ class _RoutedActions: def __init__( self, - writer_actions: Any, - reader_actions: Any, + writer_actions: object, + reader_actions: object, should_use_reader: Callable[[], bool], ): self._writer_actions = writer_actions self._reader_actions = reader_actions self._should_use_reader = should_use_reader - def __getattr__(self, name: str) -> Any: + def __getattr__(self, name: str) -> object: if name in _MODEL_READ_METHODS and self._should_use_reader(): - return getattr(self._reader_actions, name) - return getattr(self._writer_actions, name) + return _dynamic_attr(self._reader_actions, name) + return _dynamic_attr(self._writer_actions, name) class RoutingPrismaWrapper: @@ -135,7 +143,7 @@ class RoutingPrismaWrapper: return not self._reader_unavailable @staticmethod - async def _try_connect(client: PrismaWrapper, *args: Any, **kwargs: Any) -> Exception | None: + async def _try_connect(client: PrismaWrapper, *args: ConnectTimeout, **kwargs: ConnectTimeout) -> Exception | None: if client.is_connected() is True: return None try: @@ -144,7 +152,7 @@ class RoutingPrismaWrapper: except Exception as e: return e - async def connect(self, *args: Any, **kwargs: Any) -> None: + async def connect(self, *args: ConnectTimeout, **kwargs: ConnectTimeout) -> None: writer_error: Final = await self._try_connect(self._writer, *args, **kwargs) if writer_error is None: self._writer_unavailable = False @@ -176,7 +184,7 @@ class RoutingPrismaWrapper: writer_error, ) - async def disconnect(self, *args: Any, **kwargs: Any) -> None: + async def disconnect(self, *args: object, **kwargs: object) -> None: first_error: BaseException | None = None for client in (self._writer, self._reader): try: @@ -206,7 +214,7 @@ class RoutingPrismaWrapper: async def recreate_prisma_client( self, new_db_url: str, - http_client: Any | None = None, + http_client: object | None = None, *, expected_generation: int | None = None, ) -> bool: @@ -245,7 +253,7 @@ class RoutingPrismaWrapper: ) return True - async def _recreate_reader(self, http_client: Any | None = None) -> None: + async def _recreate_reader(self, http_client: object | None = None) -> None: """Resolve the reader URL and recreate its Prisma client. IAM-enabled readers regenerate their token (host/port/user came from @@ -265,14 +273,14 @@ class RoutingPrismaWrapper: def __getattr__(self, name: str) -> Any: if name in _TOP_LEVEL_READ_METHODS: - return getattr(self.read_target, name) - writer_attr: Final = getattr(self._writer, name) + return _dynamic_attr(self.read_target, name) + writer_attr: Final = _dynamic_attr(self._writer, name) # Per-model action accessors are non-callable instances that expose # both `find_many` and `create`. Methods like execute_raw / batch_ / # tx are callables and stay on the writer untouched. if not callable(writer_attr) and hasattr(writer_attr, "find_many") and hasattr(writer_attr, "create"): try: - reader_attr: Final = getattr(self._reader, name) + reader_attr: Final = _dynamic_attr(self._reader, name) except AttributeError: return writer_attr return _RoutedActions(writer_attr, reader_attr, self._should_use_reader) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py index da12222f233..707118bf016 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py @@ -15,6 +15,8 @@ restriction intact. """ import operator +from collections.abc import Callable +from types import CodeType from typing import Any, Final from RestrictedPython import ( @@ -58,7 +60,7 @@ class AsyncAwareTransformer(RestrictingNodeTransformer): return self.node_contents_visit(node) -_INPLACE_OPS: Final[dict[str, Any]] = { +_INPLACE_OPS: Final[dict[str, Callable[[object, object], object]]] = { "+=": operator.iadd, "-=": operator.isub, "*=": operator.imul, @@ -75,7 +77,7 @@ _INPLACE_OPS: Final[dict[str, Any]] = { } -def _inplacevar_(op: str, x: Any, y: Any) -> Any: +def _inplacevar_(op: str, x: object, y: object) -> object: # RestrictedPython rewrites ``x += 1`` on a simple name into # ``x = _inplacevar_("+=", x, 1)``. The package deliberately ships no # default, so we dispatch through ``operator``'s in-place helpers, which @@ -86,7 +88,7 @@ def _inplacevar_(op: str, x: Any, y: Any) -> Any: return fn(x, y) -def _build_sandbox_builtins() -> dict[str, Any]: +def _build_sandbox_builtins() -> dict[str, object]: # ``limited_builtins`` overrides ``list``/``tuple``/``range`` from # ``safe_builtins`` with bounds-checking variants (e.g. ``limited_range`` # rejects ``range(10**18)``). ``utility_builtins`` adds ``set``, @@ -98,14 +100,14 @@ def _build_sandbox_builtins() -> dict[str, Any]: } -def build_sandbox_globals() -> dict[str, Any]: +def build_sandbox_globals() -> dict[str, object]: """Assemble the globals dict for executing guardrail code. Includes the LiteLLM-provided primitives (``regex_match``, ``http_get``, ``allow``/``block``/``modify``, etc.) plus the RestrictedPython guards that the compiled bytecode expects to find by name. """ - sandbox: Final[dict[str, Any]] = get_custom_code_primitives().copy() + sandbox: Final[dict[str, object]] = get_custom_code_primitives().copy() sandbox["__builtins__"] = _build_sandbox_builtins() sandbox["_getattr_"] = safer_getattr sandbox["_getitem_"] = default_guarded_getitem @@ -116,7 +118,7 @@ def build_sandbox_globals() -> dict[str, Any]: return sandbox -def compile_sandboxed(source: str, filename: str = "") -> Any: +def compile_sandboxed(source: str, filename: str = "") -> CodeType: """Compile guardrail source with RestrictedPython's AST transformer. Raises ``SyntaxError`` on either a Python syntax error or a restricted diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index 507dd645953..0ae048c040d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -2,7 +2,7 @@ from __future__ import annotations import os from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import urlparse from uuid import uuid4 @@ -11,7 +11,7 @@ import requests from fastapi import HTTPException from httpx import HTTPStatusError from requests.auth import HTTPBasicAuth -from typing_extensions import ReadOnly +from typing_extensions import ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -36,24 +36,31 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options carried by this guardrail's forwarded keyword arguments.""" + + guardrail_name: ReadOnly[str | None] + supported_event_hooks: list[GuardrailEventHooks] | None + + class _HiddenlayerEvaluation(TypedDict, total=False): - action: str - threat_level: str + action: ReadOnly[str] + threat_level: ReadOnly[str] class _HiddenlayerAnalysisEntry(TypedDict, total=False): - name: str - detected: bool + name: ReadOnly[str] + detected: ReadOnly[bool] class _HiddenlayerModifiedSide(TypedDict): - messages: Any + messages: ReadOnly[Any] class _HiddenlayerResponse(TypedDict, total=False): - evaluation: _HiddenlayerEvaluation - analysis: Sequence[_HiddenlayerAnalysisEntry] - modified_data: Mapping[str, _HiddenlayerModifiedSide] + evaluation: ReadOnly[_HiddenlayerEvaluation] + analysis: ReadOnly[Sequence[_HiddenlayerAnalysisEntry]] + modified_data: ReadOnly[Mapping[str, _HiddenlayerModifiedSide]] class _LoggedCallMetadata(TypedDict, total=False): @@ -151,7 +158,7 @@ class HiddenlayerGuardrail(CustomGuardrail): api_key: str | None = None, api_base: str | None = None, auth_url: str | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID") @@ -356,7 +363,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): api_key: str | None = None, api_base: str | None = None, auth_url: str | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID") self.hiddenlayer_client_secret = api_key or os.getenv("HIDDENLAYER_CLIENT_SECRET") @@ -486,7 +493,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): self, payload: Any, input_type: Literal["request", "response"], - hl_headers: dict[str, str], + hl_headers: Mapping[str, str], ) -> httpx.Response: if input_type == "request": path = "detection/v2/request-evaluations" diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py index e3f67f0024b..c5c36ae51c4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py @@ -1,10 +1,11 @@ """LLM-as-a-Judge guardrail: uses an LLM to score responses against weighted criteria.""" -from collections.abc import Callable +from collections.abc import Callable, Mapping, MutableMapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional from fastapi import HTTPException +from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack import litellm from litellm._logging import verbose_logger @@ -41,28 +42,87 @@ _parse_judge_verdict: Final = parse_json_verdict _extract_text_from_content: Final = extract_text_from_content +class _JudgeMessage(TypedDict): + """Chat message, as far as the judge prompt builder reads it.""" + + role: ReadOnly[NotRequired[str]] + content: ReadOnly[NotRequired[object]] + + +class _GuardrailOptions(TypedDict, total=False): + """Base :class:`CustomGuardrail` options forwarded untouched.""" + + mask_request_content: ReadOnly[bool] + mask_response_content: ReadOnly[bool] + violation_message_template: ReadOnly[str | None] + end_session_after_n_fails: ReadOnly[int | None] + on_violation: ReadOnly[str | None] + realtime_violation_message: ReadOnly[str | None] + on_sensitive_data: ReadOnly[str | None] + sensitive_data_route_to_model: ReadOnly[str | None] + sticky_session_routing: ReadOnly[bool] + run_in_parallel: ReadOnly[bool] + only_scan_new_messages: ReadOnly[bool] + + +class _RequestMessagesView(TypedDict): + messages: ReadOnly[Sequence[_JudgeMessage]] + + +class _RequestMetadataView(TypedDict): + metadata: ReadOnly[MutableMapping[str, object]] + + +class _OverallScoreView(TypedDict): + overall_score: ReadOnly[str | float] + + +class _JudgeModelView(TypedDict): + judge_model: ReadOnly[str] + + +class _CriteriaView(TypedDict): + criteria: ReadOnly[Sequence[Mapping[str, str | float]]] + + +class _OnFailureView(TypedDict): + on_failure: ReadOnly[Literal["block", "log"]] + + +class _ThresholdView(TypedDict): + overall_threshold: ReadOnly[str | float] + + +class _ModeView(TypedDict): + mode: ReadOnly[object] + + +class _DefaultOnView(TypedDict): + default_on: ReadOnly[object] + + def _get_litellm_param( litellm_params: "LitellmParams", guardrail: "Guardrail", key: str, - default: Any = None, + default: str | float | bool | None = None, ) -> Any: - val: Final = getattr(litellm_params, key, None) + val: Final[object] = getattr(litellm_params, key, None) if val is not None: return val raw: Final = guardrail.get("litellm_params") if isinstance(raw, dict) and key in raw: return raw[key] if raw is not None and not isinstance(raw, dict): - attr: Final = getattr(raw, key, None) + attr: Final[object] = getattr(raw, key, None) if attr is not None: return attr return default def _build_judge_prompt( - criteria: list[dict[str, Any]], - messages: list[dict[str, Any]], + criteria: Sequence[Mapping[str, object]], + messages: Sequence[_JudgeMessage], response_text: str, ) -> str: criteria_block: Final = "\n".join( @@ -87,13 +147,13 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): self, guardrail_name: str, judge_model: str, - criteria: list[dict[str, Any]], + criteria: Sequence[Mapping[str, object]], overall_threshold: float = 80.0, on_failure: Literal["block", "log"] = "block", event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None, default_on: bool = False, router_provider: "Callable[[], Router | None] | None" = None, - **kwargs: Any, + **kwargs: Unpack[_GuardrailOptions], ) -> None: _event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None if event_hook is not None: @@ -121,9 +181,9 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): async def _run_judge( self, - messages: list[dict[str, Any]], + messages: Sequence[_JudgeMessage], response_text: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: judge_messages: Final = [ {"role": "system", "content": JUDGE_SYSTEM_PROMPT}, { @@ -162,10 +222,10 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): judge_result: dict[str, Any] = {} try: - messages: Final[list[dict[str, Any]]] = request_data.get("messages") or [] + request_messages: Final[_RequestMessagesView] = {"messages": request_data.get("messages") or []} try: - judge_result = await self._run_judge(messages, response_text) + judge_result = await self._run_judge(request_messages["messages"], response_text) except Exception as judge_err: verbose_logger.warning( "llm_as_a_judge guardrail: judge call failed, failing open. Error: %s", judge_err @@ -174,7 +234,8 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): return inputs try: - overall_score: Final = max(0.0, min(100.0, float(judge_result.get("overall_score", 100)))) + raw_score: Final[_OverallScoreView] = {"overall_score": judge_result.get("overall_score", 100)} + overall_score: Final = max(0.0, min(100.0, float(raw_score["overall_score"]))) except (TypeError, ValueError): verbose_logger.warning("llm_as_a_judge: invalid overall_score from judge, failing open") return inputs @@ -189,7 +250,8 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): "threshold": self.overall_threshold, "verdicts": judge_result.get("verdicts", []), } - _metadata: Final = request_data.setdefault("metadata", {}) + request_metadata: Final[_RequestMetadataView] = {"metadata": request_data.setdefault("metadata", {})} + _metadata: Final = request_metadata["metadata"] existing: Final = _metadata.get("eval_information") if isinstance(existing, list): existing.append(eval_info) @@ -238,37 +300,45 @@ def initialize_guardrail( if not guardrail_name: raise ValueError("llm_as_a_judge guardrail requires a guardrail_name") - judge_model: Final = _get_litellm_param(litellm_params, guardrail, "judge_model") - if not judge_model: + judge_model: Final[_JudgeModelView] = {"judge_model": _get_litellm_param(litellm_params, guardrail, "judge_model")} + if not judge_model["judge_model"]: raise ValueError("llm_as_a_judge guardrail requires judge_model in litellm_params") - criteria: Final = _get_litellm_param(litellm_params, guardrail, "criteria") or [] - if not criteria: + criteria: Final[_CriteriaView] = {"criteria": _get_litellm_param(litellm_params, guardrail, "criteria") or []} + if not criteria["criteria"]: raise ValueError("llm_as_a_judge guardrail requires at least one criterion") - weight_total: Final = sum(float(c.get("weight", 0)) for c in criteria) + weight_total: Final = sum(float(c.get("weight", 0)) for c in criteria["criteria"]) if abs(weight_total - 100) > 0.5: raise ValueError(f"llm_as_a_judge criterion weights must sum to 100 (got {weight_total})") - on_failure: Final = _get_litellm_param(litellm_params, guardrail, "on_failure", "block") - if on_failure not in _VALID_ON_FAILURE: - raise ValueError(f"llm_as_a_judge on_failure must be 'block' or 'log', got '{on_failure}'") + on_failure: Final[_OnFailureView] = { + "on_failure": _get_litellm_param(litellm_params, guardrail, "on_failure", "block") + } + if on_failure["on_failure"] not in _VALID_ON_FAILURE: + raise ValueError(f"llm_as_a_judge on_failure must be 'block' or 'log', got '{on_failure['on_failure']}'") - overall_threshold: Final = float(_get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0)) + threshold: Final[_ThresholdView] = { + "overall_threshold": _get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0) + } + overall_threshold: Final = float(threshold["overall_threshold"]) - mode: Final = _get_litellm_param(litellm_params, guardrail, "mode") + mode: Final[_ModeView] = {"mode": _get_litellm_param(litellm_params, guardrail, "mode")} event_hook: GuardrailEventHooks | None = None - if isinstance(mode, str) and mode in {e.value for e in GuardrailEventHooks}: - event_hook = GuardrailEventHooks(mode) + if isinstance(mode["mode"], str) and mode["mode"] in {e.value for e in GuardrailEventHooks}: + event_hook = GuardrailEventHooks(mode["mode"]) + default_on: Final[_DefaultOnView] = { + "default_on": _get_litellm_param(litellm_params, guardrail, "default_on", False) + } instance: Final = LLMAsAJudgeGuardrail( guardrail_name=guardrail_name, - judge_model=judge_model, - criteria=criteria, + judge_model=judge_model["judge_model"], + criteria=criteria["criteria"], overall_threshold=overall_threshold, - on_failure=on_failure, + on_failure=on_failure["on_failure"], event_hook=event_hook, - default_on=bool(_get_litellm_param(litellm_params, guardrail, "default_on", False)), + default_on=bool(default_on["default_on"]), ) litellm.logging_callback_manager.add_litellm_callback(instance) return instance diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index e9cd6addef8..704e2564ef5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -7,8 +7,9 @@ import enum import json import os +from collections.abc import Callable, Mapping from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, cast from urllib.parse import urlparse from litellm._logging import verbose_proxy_logger @@ -36,6 +37,8 @@ _AIDR_SCAN_ENDPOINT: Final = "/litellm/guardrail" _INTERVENED_INPUT_FIELDS: Final = ("texts", "images", "tools", "tool_calls") _DEFAULT_API_BASE_HOSTNAME: Final = urlparse(_DEFAULT_API_BASE).hostname +_GuardrailJsonResponse: TypeAlias = Exception | str | dict[str, object] + _KEYS_DUPLICATING_SCAN_INPUTS: Final = ("messages", "input") _LOGGING_KEYS_DUPLICATING_SCAN_INPUTS: Final = _KEYS_DUPLICATING_SCAN_INPUTS + ( "additional_args", @@ -119,7 +122,7 @@ class NomaV2Guardrail(CustomGuardrail): def _resolve_action_from_response( self, - response_json: dict, + response_json: Mapping[str, object], ) -> _Action: action: Final = response_json.get("action") if isinstance(action, str): @@ -165,10 +168,11 @@ class NomaV2Guardrail(CustomGuardrail): @staticmethod def _sanitize_payload_for_transport(payload: dict) -> dict: - def _default(obj: Any) -> Any: - if hasattr(obj, "model_dump"): + def _default(obj: object) -> object: + model_dump: Final[Callable[[], Mapping[str, object]] | None] = getattr(obj, "model_dump", None) + if model_dump is not None: try: - return obj.model_dump() + return model_dump() except Exception: pass return str(obj) @@ -178,7 +182,7 @@ class NomaV2Guardrail(CustomGuardrail): except (ValueError, TypeError): json_str = safe_dumps(payload) - safe_payload: Final = safe_json_loads(json_str, default={}) + safe_payload: Final[object] = safe_json_loads(json_str, default={}) if safe_payload == {} and payload: verbose_proxy_logger.warning( "Noma v2 guardrail: payload serialization failed, falling back to empty payload" @@ -196,7 +200,7 @@ class NomaV2Guardrail(CustomGuardrail): async def _call_noma_scan( self, payload: dict, - ) -> dict: + ) -> dict[str, object]: headers: Final[dict[str, str]] = {"Content-Type": "application/json"} authorization_header: Final = self._get_authorization_header() if authorization_header: @@ -215,7 +219,7 @@ class NomaV2Guardrail(CustomGuardrail): response.text, ) response.raise_for_status() - response_json: Final = response.json() + response_json: Final[dict[str, object]] = response.json() verbose_proxy_logger.debug( "Noma v2 AIDR response parsed: %s", json.dumps(response_json, default=str), @@ -227,7 +231,7 @@ class NomaV2Guardrail(CustomGuardrail): request_data: dict, start_time: datetime, guardrail_status: GuardrailStatus, - guardrail_json_response: Any, + guardrail_json_response: _GuardrailJsonResponse, ) -> None: end_time: Final = datetime.now() duration: Final = (end_time - start_time).total_seconds() @@ -270,11 +274,11 @@ class NomaV2Guardrail(CustomGuardrail): ) -> GenericGuardrailAPIInputs: start_time: Final = datetime.now() guardrail_status: GuardrailStatus = "success" - guardrail_json_response: Any = {} + guardrail_json_response: _GuardrailJsonResponse = {} dynamic_params = self.get_guardrail_dynamic_request_body_params(request_data) if not isinstance(dynamic_params, dict): dynamic_params = {} - response_json: dict | None = None + response_json: dict[str, object] | None = None # Per-request dynamic params can override configured application context. application_id = self._get_non_empty_str(dynamic_params.get("application_id")) diff --git a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py index 4775a8b3caa..155db816e94 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py @@ -7,8 +7,11 @@ before and after LLM calls. """ import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing_extensions import ReadOnly, TypedDict + from litellm._logging import verbose_proxy_logger from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_guardrail import ( @@ -20,6 +23,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -34,6 +38,16 @@ _DEFAULT_API_BASE: Final = "https://api.promptguard.co" _GUARD_ENDPOINT: Final = "/api/v1/guard" +class PromptGuardResult(TypedDict, total=False): + """The fields this guardrail reads off a PromptGuard Guard API response.""" + + decision: ReadOnly[str] + threat_type: ReadOnly[str] + event_id: ReadOnly[str] + confidence: ReadOnly[float] + redacted_messages: ReadOnly[list[AllMessageValues]] + + class PromptGuardMissingCredentials(Exception): pass @@ -96,7 +110,7 @@ class PromptGuardGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: Mapping[str, object], input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: @@ -114,7 +128,7 @@ class PromptGuardGuardrail(CustomGuardrail): direction: Final = "input" if input_type == "request" else "output" - payload: Final[dict[str, Any]] = { + payload: Final[dict[str, object]] = { "messages": messages, "direction": direction, } @@ -144,7 +158,7 @@ class PromptGuardGuardrail(CustomGuardrail): timeout=10.0, ) response.raise_for_status() - result: Final = response.json() + result: Final[PromptGuardResult] = response.json() except Exception as exc: verbose_proxy_logger.error("PromptGuard API error: %s", str(exc)) if self.block_on_error: diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 41e52f05c01..7b45afa0ad2 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -1,4 +1,6 @@ -from typing import Final +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import Final, Protocol from fastapi import APIRouter, Depends, HTTPException, Query @@ -18,7 +20,59 @@ from litellm.repositories.table_repositories import JWTKeyMappingRepository router: Final = APIRouter() -def _to_response(mapping) -> JWTKeyMappingResponse: +class _JWTKeyMappingRecord(Protocol): + """A ``LiteLLM_JWTKeyMapping`` row, viewed through the columns these endpoints read.""" + + @property + def id(self) -> str: ... + + @property + def jwt_claim_name(self) -> str: ... + + @property + def jwt_claim_value(self) -> str: ... + + @property + def description(self) -> str | None: ... + + @property + def is_active(self) -> bool: ... + + @property + def created_at(self) -> datetime: ... + + @property + def updated_at(self) -> datetime: ... + + @property + def created_by(self) -> str | None: ... + + @property + def updated_by(self) -> str | None: ... + + +class _JWTKeyMappingTable(Protocol): + """The Prisma table actions these endpoints issue against the JWT key mapping table.""" + + async def create(self, *, data: Mapping[str, object]) -> _JWTKeyMappingRecord: ... + + async def find_unique(self, *, where: Mapping[str, object]) -> _JWTKeyMappingRecord | None: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> _JWTKeyMappingRecord: ... + + async def delete(self, *, where: Mapping[str, object]) -> _JWTKeyMappingRecord | None: ... + + async def find_many(self, *, skip: int, take: int, order: Mapping[str, str]) -> Sequence[_JWTKeyMappingRecord]: ... + + async def count(self) -> int: ... + + +def _mapping_table(prisma_client: object) -> _JWTKeyMappingTable: + """View the JWT key mapping repository's untyped Prisma table through the actions used here.""" + return JWTKeyMappingRepository(prisma_client).table + + +def _to_response(mapping: _JWTKeyMappingRecord) -> JWTKeyMappingResponse: """Convert a Prisma mapping object to a safe response (no hashed token).""" return JWTKeyMappingResponse( id=mapping.id, @@ -62,7 +116,7 @@ async def create_jwt_key_mapping( if data.description is not None: create_data["description"] = data.description - new_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.create(data=create_data) + new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data) # Invalidate cache cache_key: Final = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}" @@ -110,7 +164,7 @@ async def update_jwt_key_mapping( try: # Get old mapping for cache invalidation - old_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": data.id}) + old_mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": data.id}) if old_mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") @@ -118,9 +172,7 @@ async def update_jwt_key_mapping( cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" await user_api_key_cache.async_delete_cache(cache_key) - updated_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.update( - where={"id": data.id}, data=update_data - ) + updated_mapping: Final = await _mapping_table(prisma_client).update(where={"id": data.id}, data=update_data) # Invalidate new cache key if claim fields changed cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}" @@ -159,7 +211,7 @@ async def delete_jwt_key_mapping( try: # Get old mapping for cache invalidation - old_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": data.id}) + old_mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": data.id}) if old_mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") @@ -167,7 +219,7 @@ async def delete_jwt_key_mapping( cache_key: Final = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" await user_api_key_cache.async_delete_cache(cache_key) - await JWTKeyMappingRepository(prisma_client).table.delete(where={"id": data.id}) + await _mapping_table(prisma_client).delete(where={"id": data.id}) return {"status": "success"} except HTTPException: raise @@ -195,12 +247,12 @@ async def list_jwt_key_mappings( try: skip: Final = (page - 1) * size - mappings: Final = await JWTKeyMappingRepository(prisma_client).table.find_many( + mappings: Final = await _mapping_table(prisma_client).find_many( skip=skip, take=size, order={"created_at": "desc"}, ) - total_count: Final = await JWTKeyMappingRepository(prisma_client).table.count() + total_count: Final = await _mapping_table(prisma_client).count() return { "mappings": [_to_response(m) for m in mappings], "total_count": total_count, @@ -232,7 +284,7 @@ async def info_jwt_key_mapping( raise HTTPException(status_code=500, detail="Database not connected") try: - mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": id}) + mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": id}) if mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") return _to_response(mapping) diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 020698dabd9..3c5d889d23f 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -10,8 +10,8 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json -from collections.abc import Sequence -from typing import TYPE_CHECKING, Final, TypedDict, cast +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final, TypedDict from fastapi import Request, Response from fastapi.responses import StreamingResponse @@ -38,6 +38,36 @@ class _StreamOutputItem(TypedDict, total=False): content: ReadOnly[Sequence[_StreamContentPart | None]] +class _StreamTerminalResponse(TypedDict, total=False): + """Fields of the ``response`` payload carried by a terminal streaming event.""" + + status: ReadOnly[ResponsesAPIStatus] + tool_choice: ReadOnly[object] + model: ReadOnly[str] + instructions: ReadOnly[str] + temperature: ReadOnly[float] + top_p: ReadOnly[float] + max_output_tokens: ReadOnly[int] + previous_response_id: ReadOnly[str] + truncation: ReadOnly[str] + parallel_tool_calls: ReadOnly[bool] + user: ReadOnly[str] + store: ReadOnly[bool] + output: ReadOnly[Sequence[_StreamOutputItem]] + + +class _StreamEvent(TypedDict, total=False): + """One decoded ``data:`` frame of an OpenAI Responses streaming body.""" + + type: ReadOnly[str] + item: ReadOnly[_StreamOutputItem] + item_id: ReadOnly[str] + part: ReadOnly[_StreamContentPart] + content_index: ReadOnly[int] + delta: ReadOnly[str] + response: ReadOnly[_StreamTerminalResponse] + + async def background_streaming_task( polling_id: str, data, @@ -139,7 +169,7 @@ async def background_streaming_task( None # Will be set by response.completed/failed/incomplete/cancelled ) terminal_error = None - _event_to_status: Final = { + _event_to_status: Final[Mapping[str, ResponsesAPIStatus]] = { "response.completed": "completed", "response.failed": "failed", "response.incomplete": "incomplete", @@ -180,7 +210,7 @@ async def background_streaming_task( break try: - event = json.loads(chunk_data) + event: _StreamEvent = json.loads(chunk_data) event_type = event.get("type", "") # Process different event types based on OpenAI streaming spec @@ -288,12 +318,9 @@ async def background_streaming_task( # Terminal event - extract all ResponsesAPIResponse fields # https://platform.openai.com/docs/api-reference/responses-streaming response_data = event.get("response", {}) - terminal_status = cast( - ResponsesAPIStatus, - response_data.get( - "status", - _event_to_status.get(event_type, "completed"), - ), + terminal_status = response_data.get( + "status", + _event_to_status.get(event_type, "completed"), ) # Extract error for failed and incomplete responses diff --git a/litellm/proxy/search_endpoints/search_tool_registry.py b/litellm/proxy/search_endpoints/search_tool_registry.py index fe4794f3ba1..b25263e4c64 100644 --- a/litellm/proxy/search_endpoints/search_tool_registry.py +++ b/litellm/proxy/search_endpoints/search_tool_registry.py @@ -2,8 +2,9 @@ Search Tool Registry for managing search tool configurations. """ +from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timezone -from typing import Final +from typing import Final, Protocol from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -13,6 +14,40 @@ from litellm.repositories.table_repositories import SearchToolsRepository from litellm.types.search import SearchTool +class SearchToolRecord(Protocol): + search_tool_id: str + search_tool_name: str + created_at: datetime + updated_at: datetime + + def __iter__(self) -> Iterator[tuple[str, object]]: ... + + +class SearchToolTableClient(Protocol): + async def create(self, data: Mapping[str, object]) -> SearchToolRecord: ... + + async def find_unique(self, where: Mapping[str, object]) -> SearchToolRecord | None: ... + + async def find_many(self, order: Mapping[str, str] | None = None) -> Sequence[SearchToolRecord]: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> SearchToolRecord: ... + + async def delete(self, where: Mapping[str, object]) -> SearchToolRecord: ... + + +class _SearchToolsRepositoryView(Protocol): + @property + def table(self) -> SearchToolTableClient: ... + + +def _search_tools_table_of(repository: _SearchToolsRepositoryView) -> SearchToolTableClient: + return repository.table + + +def _search_tools_table(prisma_client: PrismaClient) -> SearchToolTableClient: + return _search_tools_table_of(SearchToolsRepository(prisma_client)) + + class SearchToolRegistry: """ Handles adding, removing, and getting search tools in DB + in memory. @@ -22,7 +57,7 @@ class SearchToolRegistry: pass @staticmethod - def _convert_prisma_to_dict(prisma_obj) -> dict: + def _convert_prisma_to_dict(prisma_obj: SearchToolRecord) -> dict: """ Convert Prisma result to dict with datetime objects as ISO format strings. @@ -35,9 +70,9 @@ class SearchToolRegistry: result: Final = dict(prisma_obj) # Convert datetime objects to ISO format strings if "created_at" in result and result["created_at"]: - result["created_at"] = result["created_at"].isoformat() + result["created_at"] = prisma_obj.created_at.isoformat() if "updated_at" in result and result["updated_at"]: - result["updated_at"] = result["updated_at"].isoformat() + result["updated_at"] = prisma_obj.updated_at.isoformat() return result ########################################################### @@ -61,7 +96,7 @@ class SearchToolRegistry: search_tool_info: Final[str] = safe_dumps(search_tool.get("search_tool_info", {})) # Create search tool in DB - created_search_tool: Final = await SearchToolsRepository(prisma_client).table.create( + created_search_tool: Final = await _search_tools_table(prisma_client).create( data={ "search_tool_name": search_tool_name, "litellm_params": litellm_params, @@ -95,7 +130,7 @@ class SearchToolRegistry: """ try: # Get search tool before deletion for response - existing_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique( + existing_tool: Final = await _search_tools_table(prisma_client).find_unique( where={"search_tool_id": search_tool_id} ) @@ -103,7 +138,7 @@ class SearchToolRegistry: raise Exception(f"Search tool with ID {search_tool_id} not found") # Delete from DB - await SearchToolsRepository(prisma_client).table.delete(where={"search_tool_id": search_tool_id}) + await _search_tools_table(prisma_client).delete(where={"search_tool_id": search_tool_id}) return { "message": f"Search tool {search_tool_id} deleted successfully", @@ -131,7 +166,7 @@ class SearchToolRegistry: search_tool_info: Final[str] = safe_dumps(search_tool.get("search_tool_info", {})) # Update in DB - updated_search_tool: Final = await SearchToolsRepository(prisma_client).table.update( + updated_search_tool: Final = await _search_tools_table(prisma_client).update( where={"search_tool_id": search_tool_id}, data={ "search_tool_name": search_tool_name, @@ -163,7 +198,7 @@ class SearchToolRegistry: try: search_tools_from_db: Final = await call_with_db_reconnect_retry( prisma_client, - lambda: SearchToolsRepository(prisma_client).table.find_many( + lambda: _search_tools_table(prisma_client).find_many( order={"created_at": "desc"}, ), reason="get_all_search_tools_from_db_lookup_failure", @@ -194,7 +229,7 @@ class SearchToolRegistry: Search tool configuration or None if not found """ try: - search_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique( + search_tool: Final = await _search_tools_table(prisma_client).find_unique( where={"search_tool_id": search_tool_id} ) @@ -222,7 +257,7 @@ class SearchToolRegistry: Search tool configuration or None if not found """ try: - search_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique( + search_tool: Final = await _search_tools_table(prisma_client).find_unique( where={"search_tool_name": search_tool_name} ) diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index 16b8f82815c..255faf94402 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -1,11 +1,45 @@ +from collections.abc import Sequence from typing import Any, Final +from typing_extensions import NotRequired, ReadOnly, TypedDict + from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage from litellm.types.utils import ModelResponse -from litellm.types.vector_stores import ( - VectorStoreResultContent, - VectorStoreSearchResponse, -) +from litellm.types.vector_stores import VectorStoreSearchResponse + + +class _ResultContentView(TypedDict): + """Content entry carried by a vector store search result.""" + + type: ReadOnly[NotRequired[str]] + text: ReadOnly[str] + + +class _SearchResultView(TypedDict): + """Vector store search result, as far as :class:`RAGQuery` reads it.""" + + content: ReadOnly[NotRequired[Sequence[_ResultContentView]]] + text: ReadOnly[NotRequired[str]] + + +class _SearchDataView(TypedDict): + results: ReadOnly[Sequence[_SearchResultView]] + + +class _ContextChunksView(TypedDict): + chunks: ReadOnly[Sequence[_SearchResultView | str | None]] + + +class _RerankResultView(TypedDict): + index: ReadOnly[NotRequired[int]] + + +class _RerankResultsView(TypedDict): + results: ReadOnly[Sequence[_RerankResultView]] + + +class _MessageView(TypedDict): + message: ReadOnly[object] class RAGQuery: @@ -42,9 +76,10 @@ class RAGQuery: """ context_content = RAGQuery.CONTENT_PREFIX_STRING - for chunk in context_chunks: + chunks: Final[_ContextChunksView] = {"chunks": context_chunks} + for chunk in chunks["chunks"]: if isinstance(chunk, dict): - result_content: list[VectorStoreResultContent] | None = chunk.get("content") + result_content: Sequence[_ResultContentView] | None = chunk.get("content") if result_content: for content_item in result_content: content_text: str | None = content_item.get("text") @@ -64,14 +99,15 @@ class RAGQuery: def add_search_results_to_response( response: ModelResponse, search_results: VectorStoreSearchResponse, - rerank_results: Any | None = None, + rerank_results: object = None, ) -> ModelResponse: """ Add search results to the response choices. """ if hasattr(response, "choices") and response.choices: for choice in response.choices: - message = getattr(choice, "message", None) + message_view: _MessageView = {"message": getattr(choice, "message", None)} + message = message_view["message"] if message is not None: # Get existing provider_specific_fields or create new dict provider_fields = getattr(message, "provider_specific_fields", None) or {} @@ -91,7 +127,8 @@ class RAGQuery: ) -> list[str | dict[str, Any]]: """Extract text documents from vector store search response.""" documents: Final[list[str | dict[str, Any]]] = [] - for result in search_response.get("data", []): + search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} + for result in search_data["results"]: content_list = result.get("content", []) for content in content_list: if content.get("type") == "text" and content.get("text"): @@ -99,11 +136,13 @@ class RAGQuery: return documents @staticmethod - def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[Any]: + def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[_SearchResultView]: """Get the original search results corresponding to the top reranked results.""" - top_chunks: Final = [] - original_results: Final = search_response.get("data", []) - for result in rerank_response.get("results", []): + top_chunks: Final[list[_SearchResultView]] = [] + search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} + original_results: Final = search_data["results"] + reranked: Final[_RerankResultsView] = {"results": rerank_response.get("results", [])} + for result in reranked["results"]: index = result.get("index") if index is not None and index < len(original_results): top_chunks.append(original_results[index]) diff --git a/litellm/repositories/base_repository.py b/litellm/repositories/base_repository.py index 7008099fe8c..8dfe3ed962e 100644 --- a/litellm/repositories/base_repository.py +++ b/litellm/repositories/base_repository.py @@ -3,7 +3,7 @@ Base repository class with common functionality. """ from abc import ABC, abstractmethod -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from typing import Any, Final, Generic, Protocol, TypeVar, runtime_checkable from pydantic import BaseModel @@ -24,6 +24,22 @@ class SupportsDict(Protocol): DbRecord = Mapping[str, object] | SupportsModelDump | SupportsDict | Sequence[tuple[str, object]] +class PrismaCrudActions(Protocol): + """The Prisma table actions reached by the generic repository CRUD helpers.""" + + async def find_unique(self, *, where: Mapping[str, object]) -> DbRecord | None: ... + + find_many: Callable[..., Awaitable[Sequence[DbRecord]]] + + async def create(self, *, data: Mapping[str, object]) -> DbRecord: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> DbRecord | None: ... + + async def delete(self, *, where: Mapping[str, object]) -> DbRecord | None: ... + + async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... + + def record_to_dict(record: DbRecord) -> Mapping[str, object]: """Project a database record into a mapping of column name to value.""" if isinstance(record, SupportsModelDump): @@ -38,7 +54,7 @@ def record_to_dict(record: DbRecord) -> Mapping[str, object]: class BaseRepository(ABC, Generic[T]): """Abstract base class for all repositories.""" - def __init__(self, prisma_client: Any): # any-ok: PrismaClient is an untyped runtime wrapper + def __init__(self, prisma_client: object): self._prisma_client = prisma_client @property @@ -53,6 +69,11 @@ class BaseRepository(ABC, Generic[T]): """Return the Prisma table for this repository.""" ... + @property + def _crud_actions(self) -> PrismaCrudActions: + """View ``table`` through the action surface the CRUD helpers below use.""" + return self.table + @property @abstractmethod def model_class(self) -> type[T]: @@ -71,18 +92,18 @@ class BaseRepository(ABC, Generic[T]): async def find_by_id(self, id_value: str, id_field: str = "id") -> T | None: """Find a record by its primary key.""" - record: Final = await self.table.find_unique(where={id_field: id_value}) + record: Final = await self._crud_actions.find_unique(where={id_field: id_value}) return self._to_model(record) async def find_many( self, - where: dict[str, Any] | None = None, + where: Mapping[str, object] | None = None, skip: int | None = None, take: int | None = None, - order: dict[str, str] | None = None, + order: Mapping[str, str] | None = None, ) -> list[T]: """Find multiple records matching the criteria.""" - kwargs: Final[dict[str, Any]] = {} + kwargs: Final[dict[str, object]] = {} if where: kwargs["where"] = where if skip is not None: @@ -92,31 +113,31 @@ class BaseRepository(ABC, Generic[T]): if order: kwargs["order"] = order - records: Final = await self.table.find_many(**kwargs) + records: Final = await self._crud_actions.find_many(**kwargs) return self._to_model_list(records) - async def create(self, data: dict[str, Any]) -> T: + async def create(self, data: Mapping[str, object]) -> T: """Create a new record.""" - record: Final = await self.table.create(data=data) + record: Final = await self._crud_actions.create(data=data) model: Final = self._to_model(record) assert model is not None return model - async def update(self, id_value: str, data: dict[str, Any], id_field: str = "id") -> T | None: + async def update(self, id_value: str, data: Mapping[str, object], id_field: str = "id") -> T | None: """Update an existing record.""" - record: Final = await self.table.update(where={id_field: id_value}, data=data) + record: Final = await self._crud_actions.update(where={id_field: id_value}, data=data) return self._to_model(record) async def delete(self, id_value: str, id_field: str = "id") -> T | None: """Delete a record by its primary key.""" - record: Final = await self.table.delete(where={id_field: id_value}) + record: Final = await self._crud_actions.delete(where={id_field: id_value}) return self._to_model(record) - async def count(self, where: dict[str, Any] | None = None) -> int: + async def count(self, where: Mapping[str, object] | None = None) -> int: """Count records matching the criteria.""" - return await self.table.count(where=where) + return await self._crud_actions.count(where=where) async def exists(self, id_value: str, id_field: str = "id") -> bool: """Check if a record exists.""" - record: Final = await self.table.find_unique(where={id_field: id_value}) + record: Final = await self._crud_actions.find_unique(where={id_field: id_value}) return record is not None diff --git a/litellm/repositories/credentials_repository.py b/litellm/repositories/credentials_repository.py index 9fdb6e4aca7..b22f24e0de5 100644 --- a/litellm/repositories/credentials_repository.py +++ b/litellm/repositories/credentials_repository.py @@ -6,11 +6,41 @@ credential values is the caller's responsibility (see ``CredentialHelperUtils``) so reads return the stored values verbatim. """ -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm.models.credentials import CredentialItem from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync +if TYPE_CHECKING: + from prisma.models import LiteLLM_CredentialsTable + + +class _CredentialsDb(Protocol): + @property + def litellm_credentialstable(self) -> object: ... + + +class _PrismaClientView(Protocol): + @property + def db(self) -> _CredentialsDb: ... + + +class _CredentialsActions(Protocol): + """Prisma table actions used by :class:`CredentialsRepository`.""" + + async def find_many(self) -> "Sequence[LiteLLM_CredentialsTable]": ... + + async def create(self, *, data: Mapping[str, object]) -> "LiteLLM_CredentialsTable": ... + + async def find_unique(self, *, where: Mapping[str, object]) -> "LiteLLM_CredentialsTable | None": ... + + async def update( + self, *, where: Mapping[str, object], data: Mapping[str, object] + ) -> "LiteLLM_CredentialsTable | None": ... + + async def delete(self, *, where: Mapping[str, object]) -> "LiteLLM_CredentialsTable | None": ... + class CredentialsRepository: """Repository for credentials database operations, keyed by credential name.""" @@ -19,7 +49,7 @@ class CredentialsRepository: self._prisma_client = prisma_client @property - def prisma_client(self) -> Any: + def prisma_client(self) -> _PrismaClientView: if self._prisma_client is None: raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") return self._prisma_client @@ -31,6 +61,10 @@ class CredentialsRepository: table_name="litellm_credentialstable", ) + @property + def _credentials_table(self) -> _CredentialsActions: + return self.table + @staticmethod def _to_model(record: Any) -> CredentialItem | None: if record is None: @@ -42,18 +76,20 @@ class CredentialsRepository: credential_info=data.get("credential_info") or {}, ) - async def find_all(self) -> Any: - return await self.table.find_many() + async def find_all(self) -> "Sequence[LiteLLM_CredentialsTable]": + return await self._credentials_table.find_many() - async def create(self, data: dict[str, Any]) -> Any: - return await self.table.create(data=data) + async def create(self, data: Mapping[str, object]) -> "LiteLLM_CredentialsTable": + return await self._credentials_table.create(data=data) async def find_by_name(self, credential_name: str) -> CredentialItem | None: - record: Final = await self.table.find_unique(where={"credential_name": credential_name}) + record: Final = await self._credentials_table.find_unique(where={"credential_name": credential_name}) return self._to_model(record) - async def update_by_name(self, credential_name: str, data: dict[str, Any]) -> Any: - return await self.table.update(where={"credential_name": credential_name}, data=data) + async def update_by_name( + self, credential_name: str, data: Mapping[str, object] + ) -> "LiteLLM_CredentialsTable | None": + return await self._credentials_table.update(where={"credential_name": credential_name}, data=data) - async def delete_by_name(self, credential_name: str) -> Any: - return await self.table.delete(where={"credential_name": credential_name}) + async def delete_by_name(self, credential_name: str) -> "LiteLLM_CredentialsTable | None": + return await self._credentials_table.delete(where={"credential_name": credential_name}) diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index 7efd32288e4..68490797348 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -3,9 +3,10 @@ Team repository for database operations on LiteLLM_TeamTable. """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from contextlib import AbstractAsyncContextManager from datetime import datetime -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from pydantic import TypeAdapter @@ -13,12 +14,45 @@ from litellm.models.team import LiteLLM_TeamTable, Member from litellm.repositories.base_repository import ( BaseRepository, DbRecord, + PrismaCrudActions, record_to_dict, ) if TYPE_CHECKING: from prisma import Prisma + +class _TeamTables(Protocol): + """The Prisma tables this repository reaches, on the client or inside a transaction.""" + + litellm_teamtable: PrismaCrudActions + litellm_deletedteamtable: PrismaCrudActions + + +class _TeamArrays(Protocol): + """The string array columns of a team row, which the domain model leaves untyped.""" + + @property + def members(self) -> Sequence[str]: ... + + @property + def admins(self) -> Sequence[str]: ... + + @property + def models(self) -> Sequence[str]: ... + + +def _team_arrays(team: LiteLLM_TeamTable) -> _TeamArrays: + """View a team's untyped list columns as sequences of ids.""" + return team + + +class _TeamDatabase(_TeamTables, Protocol): + """The Prisma client surface used for team reads, writes, and archival transactions.""" + + def tx(self) -> AbstractAsyncContextManager[_TeamTables]: ... + + _MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member]) _JSON_ENCODED_TEAM_FIELDS: Final = ( "metadata", @@ -34,12 +68,16 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): """Repository for team database operations.""" @property - def table(self) -> Any: # any-ok: PrismaClient.db is an untyped runtime wrapper - return self.prisma_client.db.litellm_teamtable + def _db(self) -> _TeamDatabase: + return self.prisma_client.db @property - def deleted_table(self) -> Any: # any-ok: PrismaClient.db is an untyped runtime wrapper - return self.prisma_client.db.litellm_deletedteamtable + def table(self) -> Any: # any-ok: callers reach model-specific actions this repository does not use + return self._db.litellm_teamtable + + @property + def deleted_table(self) -> PrismaCrudActions: + return self._db.litellm_deletedteamtable @property def model_class(self) -> type[LiteLLM_TeamTable]: @@ -75,8 +113,8 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): ) if not rows: return None - raw_value: Final = rows[0]["members_with_roles"] - parsed: Final = json.loads(raw_value) if isinstance(raw_value, str) else raw_value + raw_value: Final[object] = rows[0]["members_with_roles"] + parsed: Final[object] = json.loads(raw_value) if isinstance(raw_value, str) else raw_value if not parsed: return [] return _MEMBERS_WITH_ROLES_ADAPTER.validate_python(parsed) @@ -86,24 +124,24 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): async def find_by_alias(self, team_alias: str) -> LiteLLM_TeamTable | None: """Find a team by alias.""" - records: Final = await self.table.find_many(where={"team_alias": team_alias}) + records: Final = await self._crud_actions.find_many(where={"team_alias": team_alias}) if records: return self._to_model(records[0]) return None async def find_by_organization_id(self, organization_id: str) -> list[LiteLLM_TeamTable]: """Find all teams belonging to an organization.""" - records: Final = await self.table.find_many(where={"organization_id": organization_id}) + records: Final = await self._crud_actions.find_many(where={"organization_id": organization_id}) return self._to_model_list(records) async def find_by_member(self, user_id: str) -> list[LiteLLM_TeamTable]: """Find all teams where user is a member.""" - records: Final = await self.table.find_many(where={"members": {"has": user_id}}) + records: Final = await self._crud_actions.find_many(where={"members": {"has": user_id}}) return self._to_model_list(records) async def find_by_admin(self, user_id: str) -> list[LiteLLM_TeamTable]: """Find all teams where user is an admin.""" - records: Final = await self.table.find_many(where={"admins": {"has": user_id}}) + records: Final = await self._crud_actions.find_many(where={"admins": {"has": user_id}}) return self._to_model_list(records) async def create_team( @@ -232,7 +270,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): archive_data["litellm_changed_by"] = litellm_changed_by archive_data["deleted_at"] = datetime.utcnow() - async with self.prisma_client.db.tx() as tx: + async with self._db.tx() as tx: await tx.litellm_deletedteamtable.create(data=archive_data) await tx.litellm_teamtable.delete(where={"team_id": team_id}) @@ -293,7 +331,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if not await self.exists(team_id, id_field="team_id"): return None - record: Final = await self.table.update( + record: Final = await self._crud_actions.update( where={"team_id": team_id}, data={"members": {"push": user_id}}, ) @@ -310,7 +348,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if team is None: return None - members: Final = [m for m in team.members if m != user_id] + members: Final = [m for m in _team_arrays(team).members if m != user_id] return await self.update(team_id, {"members": members}, id_field="team_id") async def add_admin(self, team_id: str, user_id: str) -> LiteLLM_TeamTable | None: @@ -318,7 +356,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if not await self.exists(team_id, id_field="team_id"): return None - record: Final = await self.table.update( + record: Final = await self._crud_actions.update( where={"team_id": team_id}, data={"admins": {"push": user_id}}, ) @@ -335,7 +373,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if team is None: return None - admins: Final = [a for a in team.admins if a != user_id] + admins: Final = [a for a in _team_arrays(team).admins if a != user_id] return await self.update(team_id, {"admins": admins}, id_field="team_id") async def add_models(self, team_id: str, models: list[str]) -> LiteLLM_TeamTable | None: @@ -343,7 +381,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if not await self.exists(team_id, id_field="team_id"): return None - record: Final = await self.table.update( + record: Final = await self._crud_actions.update( where={"team_id": team_id}, data={"models": {"push": models}}, ) @@ -360,5 +398,5 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if team is None: return None - current_models: Final = [m for m in team.models if m not in models] + current_models: Final = [m for m in _team_arrays(team).models if m not in models] return await self.update(team_id, {"models": current_models}, id_field="team_id") diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py index 20fc2634a8a..c57a2b14dcb 100644 --- a/litellm/rust_bridge/responses_websocket.py +++ b/litellm/rust_bridge/responses_websocket.py @@ -2,8 +2,9 @@ from __future__ import annotations +from collections.abc import Awaitable from dataclasses import dataclass -from typing import Any, Final, Protocol +from typing import Final, Protocol import httpx from websockets.exceptions import ConnectionClosedOK @@ -12,6 +13,16 @@ from litellm.rust_bridge.loader import get_native_bridge from litellm.rust_bridge.timeouts import timeout_to_seconds +class RustResponsesWebSocketSocket(Protocol): + """Open socket handle handed back by the native bridge.""" + + def send_text(self, text: str) -> Awaitable[None]: ... + + def recv_text(self) -> Awaitable[str | None]: ... + + def close(self) -> Awaitable[None]: ... + + class RustResponsesWebSocketConnection(Protocol): @classmethod def connect( @@ -19,7 +30,7 @@ class RustResponsesWebSocketConnection(Protocol): url: str, headers: dict[str, str], timeout_seconds: float | None, - ) -> Any: + ) -> Awaitable[RustResponsesWebSocketSocket]: raise NotImplementedError @@ -32,7 +43,7 @@ _UNSET: Final[_Unset] = _Unset() @dataclass(slots=True) class _RustResponsesWebSocketState: - connection: Any = None + connection: type[RustResponsesWebSocketConnection] | None = None _STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState() @@ -40,13 +51,13 @@ _STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState() def set_rust_responses_websocket( *, - connection: Any = _UNSET, + connection: type[RustResponsesWebSocketConnection] | None | _Unset = _UNSET, ) -> None: if not isinstance(connection, _Unset): _STATE.connection = connection -def load_rust_responses_websocket() -> Any: +def load_rust_responses_websocket() -> type[RustResponsesWebSocketConnection] | None: if _STATE.connection is not None: return _STATE.connection native_bridge: Final = get_native_bridge() @@ -59,7 +70,7 @@ def load_rust_responses_websocket() -> Any: class _ConnectionAdapter: - def __init__(self, connection: Any): + def __init__(self, connection: RustResponsesWebSocketSocket): self._connection = connection async def send(self, text: str) -> None: diff --git a/litellm/secret_managers/secret_manager_handler.py b/litellm/secret_managers/secret_manager_handler.py index c77d2505d0e..a85382f863e 100644 --- a/litellm/secret_managers/secret_manager_handler.py +++ b/litellm/secret_managers/secret_manager_handler.py @@ -6,14 +6,62 @@ Handles retrieving secrets from different secret management systems. import base64 import os -from typing import Any, Final +from collections.abc import Mapping +from typing import Final, Protocol, overload import litellm from litellm._logging import print_verbose -from litellm.types.secret_managers.main import KeyManagementSystem +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem -def _is_base64(s): +class _VaultSecret(Protocol): + """The secret object returned by the Azure Key Vault and Infisical clients.""" + + @property + def value(self) -> str | None: ... + + @property + def secret_value(self) -> str | None: ... + + +class _KmsPlaintext(Protocol): + """The decrypted payload the AWS KMS client exposes under the ``Plaintext`` response key.""" + + def decode(self, encoding: str) -> str | None: ... + + +class _KmsDecryptResponse(Protocol): + """The decrypt response returned by the Google KMS and AWS KMS clients.""" + + @property + def plaintext(self) -> bytes: ... + + def __getitem__(self, key: str) -> _KmsPlaintext: ... + + +class _SecretManagerClient(Protocol): + """The untyped secret manager client surface reached by ``get_secret_from_manager``.""" + + def get_secret(self, secret_name: str, /) -> _VaultSecret: ... + + @overload + def decrypt(self, *, request: Mapping[str, object]) -> _KmsDecryptResponse: ... + + @overload + def decrypt(self, *, CiphertextBlob: bytes) -> _KmsDecryptResponse: ... + + def sync_read_secret( + self, + *, + secret_name: str, + primary_secret_name: str | None = None, + optional_params: Mapping[str, object] | None = None, + ) -> str | None: ... + + def get_secret_from_google_secret_manager(self, secret_name: str, /) -> str | None: ... + + +def _is_base64(s: str | bytes) -> bool: """Check if a string is valid base64.""" import binascii @@ -24,10 +72,10 @@ def _is_base64(s): def get_secret_from_manager( - client: Any, + client: _SecretManagerClient, key_manager: str, secret_name: str, - key_management_settings: Any | None = None, + key_management_settings: KeyManagementSettings | None = None, ) -> str | None: """ Get a secret from the configured secret manager. @@ -56,7 +104,7 @@ def get_secret_from_manager( elif ( key_manager == KeyManagementSystem.GOOGLE_KMS.value or client.__class__.__name__ == "KeyManagementServiceClient" ): - encrypted_secret: Any = os.getenv(secret_name) + encrypted_secret: str | bytes | None = os.getenv(secret_name) if encrypted_secret is None: raise ValueError("Google KMS requires the encrypted secret to be in the environment!") b64_flag: Final = _is_base64(encrypted_secret) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 5c312dcf1c8..ceaf1ec504f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,21 +1,21 @@ { "ANN001": { - "limit": 3020 + "limit": 3012 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 827 + "limit": 818 }, "ANN201": { - "limit": 2017 + "limit": 2014 }, "ANN202": { - "limit": 852 + "limit": 851 }, "ANN204": { - "limit": 711 + "limit": 707 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1188 + "limit": 1061 }, "ASYNC230": { "limit": 11 @@ -198,7 +198,7 @@ "limit": 22 }, "SIM101": { - "limit": 58 + "limit": 57 }, "SIM102": { "limit": 317 @@ -234,7 +234,7 @@ "limit": 5 }, "TID251": { - "limit": 1212 + "limit": 1202 }, "TRY002": { "limit": 524 @@ -249,7 +249,7 @@ "limit": 113 }, "TRY300": { - "limit": 859 + "limit": 858 }, "UP028": { "limit": 2 diff --git a/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py b/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py index d106cf7ea21..a768f8ccd8c 100644 --- a/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py +++ b/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py @@ -313,6 +313,25 @@ class TestGDCGeminiConfig: api_base=TEST_API_BASE, ) + def test_validate_environment_credentials_missing_audience_binding_are_named(self): + config = GDCGeminiConfig() + creds_without_audience_binding = MagicMock(spec=[]) + + with patch( + "google.auth.load_credentials_from_dict", + return_value=(creds_without_audience_binding, None), + ): + with pytest.raises(AttributeError, match="must expose with_gdch_audience"): + config.validate_environment( + headers={}, + model=TEST_MODEL, + messages=[], + optional_params={}, + litellm_params={"vertex_project": TEST_PROJECT}, + api_key=TEST_API_KEY, + api_base=TEST_API_BASE, + ) + def test_validate_environment_string_false_disables_token_caching(self): config = GDCGeminiConfig() mock_creds = MagicMock() diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 20cd1165577..890184fad78 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22805 + "limit": 22728 }, "LIT002": { - "limit": 26878 + "limit": 26860 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1069 + "limit": 1066 }, "LIT007": { "limit": 0 @@ -27,12 +27,12 @@ "limit": 0 }, "LIT010": { - "limit": 16695 + "limit": 16673 }, "LIT011": { - "limit": 5588 + "limit": 5586 }, "LIT012": { - "limit": 4519 + "limit": 4512 } } From 6f24d85490f03c8af4f6b94e93aff0ffa2240c1b Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 25 Aug 2026 18:26:57 +0000 Subject: [PATCH 02/28] fix(proxy): grant users with empty models list direct access in model listing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 4 +- .../test_team_model_name_translation.py | 55 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0abcdeaf3f6..3d3169eeab9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12382,8 +12382,10 @@ def get_direct_access_models( The 'all-proxy-models' sentinel grants direct access to every non-team deployment, mirroring how get_key_models expands it for the key/team path. + An empty models list means unrestricted access at call time (see + can_user_call_model), so it resolves the same way. """ - if SpecialModelNames.all_proxy_models.value in user_db_object.models: + if not user_db_object.models or SpecialModelNames.all_proxy_models.value in user_db_object.models: return llm_router.get_model_ids(exclude_team_models=True) return [ diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 2f4018b55ab..cb1521f10ae 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -1507,6 +1507,61 @@ def test_get_direct_access_models_resolves_explicit_model_names(): router.get_model_list.assert_called_once_with(model_name="gpt-4o") +def test_get_direct_access_models_empty_models_grants_all_non_team_models(): + """An empty user.models list means unrestricted access at call time + (can_user_call_model), so the listing must resolve it like 'all-proxy-models' + instead of returning nothing. Regression for a user with models=[] and no + teams seeing an empty Models+Endpoints page.""" + router = MagicMock() + router.get_model_ids.return_value = ["global-id-1", "global-id-2"] + + user = LiteLLM_UserTable(user_id="u", models=[], teams=[]) + + result = ps.get_direct_access_models(user_db_object=user, llm_router=router) + + assert result == ["global-id-1", "global-id-2"] + router.get_model_ids.assert_called_once_with(exclude_team_models=True) + router.get_model_list.assert_not_called() + + +@pytest.mark.asyncio +async def test_populate_team_access_grants_empty_models_user_direct_access(monkeypatch): + """An internal user with models=[] and no teams can call every non-team model, + so the Models+Endpoints page must list them instead of rendering empty.""" + global_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + + router = MagicMock() + router.get_model_ids.return_value = ["global-id-1"] + + user_row = LiteLLM_UserTable( + user_id="u", + user_role=LitellmUserRoles.INTERNAL_USER.value, + models=[], + teams=[], + ) + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + + monkeypatch.setattr(ps, "get_all_team_models", AsyncMock(return_value={})) + + caller = UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.INTERNAL_USER, team_models=[]) + + populated = await ps._populate_team_access_on_models( + user_api_key_dict=caller, + prisma_client=prisma_client, + llm_router=router, + all_models=[global_row], + ) + visible = ps._filter_models_to_user_accessible(populated) + + assert [m["model_info"]["id"] for m in visible] == ["global-id-1"] + assert visible[0]["model_info"]["direct_access"] is True + + @pytest.mark.asyncio async def test_populate_team_access_grants_all_proxy_models_user_direct_access( monkeypatch, From bd2b93c6889614f04986fcd9d94067e061a20567 Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 25 Aug 2026 20:19:47 +0000 Subject: [PATCH 03/28] chore(proxy): drop redundant docstring note in get_direct_access_models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3d3169eeab9..f536f521d57 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12382,8 +12382,6 @@ def get_direct_access_models( The 'all-proxy-models' sentinel grants direct access to every non-team deployment, mirroring how get_key_models expands it for the key/team path. - An empty models list means unrestricted access at call time (see - can_user_call_model), so it resolves the same way. """ if not user_db_object.models or SpecialModelNames.all_proxy_models.value in user_db_object.models: return llm_router.get_model_ids(exclude_team_models=True) From f81408535447db7cf6be81caa48c1447965f2288 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 27 Aug 2026 01:34:33 +0000 Subject: [PATCH 04/28] feat(mcp): bulk-import Anthropic MCP connectors via API and admin UI Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 1147 ++++++++++++++++- .../mcp_connector_import.py | 168 +++ .../mcp_management_endpoints.py | 117 ++ .../test_mcp_connector_import.py | 157 +++ .../test_mcp_management_endpoints.py | 140 ++ .../_components/ImportMCPServers.tsx | 134 ++ .../_components/importConnectorConfig.test.ts | 74 ++ .../_components/importConnectorConfig.ts | 64 + .../mcp-servers/_components/mcp_servers.tsx | 19 +- .../src/components/networking.tsx | 9 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 486 ++++++- 11 files changed, 2506 insertions(+), 9 deletions(-) create mode 100644 litellm/proxy/management_endpoints/mcp_connector_import.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ImportMCPServers.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/importConnectorConfig.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/importConnectorConfig.ts diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 0e3cbbfd560..5e0b5ba3faf 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -14727,6 +14727,17 @@ "title": "Args", "type": "array" }, + "audience": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Audience" + }, "auth_type": { "anyOf": [ { @@ -14738,7 +14749,11 @@ "authorization", "oauth2", "aws_sigv4", - "token" + "token", + "oauth2_token_exchange", + "oauth2_id_jag", + "true_passthrough", + "oauth_delegate" ], "type": "string" }, @@ -14793,6 +14808,17 @@ ], "title": "Command" }, + "connected_app_reachable": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Connected App Reachable" + }, "created_at": { "anyOf": [ { @@ -14826,6 +14852,22 @@ } ] }, + "dcr_bridge": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Dcr Bridge" + }, + "delegate_auth_to_upstream": { + "default": false, + "title": "Delegate Auth To Upstream", + "type": "boolean" + }, "description": { "anyOf": [ { @@ -14844,6 +14886,20 @@ "title": "Env", "type": "object" }, + "env_vars": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/MCPEnvVar" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Env Vars" + }, "extra_headers": { "items": { "type": "string" @@ -14889,6 +14945,17 @@ "title": "Is Byok", "type": "boolean" }, + "issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Issuer" + }, "last_health_check": { "anyOf": [ { @@ -14901,6 +14968,17 @@ ], "title": "Last Health Check" }, + "max_concurrent_requests": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Max Concurrent Requests" + }, "mcp_access_groups": { "items": { "type": "string" @@ -14920,6 +14998,26 @@ ], "title": "Mcp Info" }, + "oauth2_flow": { + "anyOf": [ + { + "enum": [ + "client_credentials", + "authorization_code" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Oauth2 Flow" + }, + "oauth_passthrough": { + "default": false, + "title": "Oauth Passthrough", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -15023,6 +15121,17 @@ "description": "Health status: 'healthy', 'unhealthy', 'unknown'", "title": "Status" }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" + }, "submitted_at": { "anyOf": [ { @@ -15063,6 +15172,39 @@ "title": "Teams", "type": "array" }, + "timeout": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Timeout" + }, + "token_exchange_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Endpoint" + }, + "token_exchange_profile": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Profile" + }, "token_url": { "anyOf": [ { @@ -15153,8 +15295,236 @@ "title": "LiteLLM_MCPServerTable", "type": "object" }, + "MCPConnectorEntry": { + "properties": { + "args": { + "items": { + "type": "string" + }, + "title": "Args", + "type": "array" + }, + "authorization_token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Authorization Token" + }, + "command": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Command" + }, + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "env": { + "additionalProperties": { + "type": "string" + }, + "title": "Env", + "type": "object" + }, + "headers": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Headers" + }, + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Name" + }, + "type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Type" + }, + "url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Url" + } + }, + "title": "MCPConnectorEntry", + "type": "object" + }, + "MCPConnectorImportFailure": { + "properties": { + "error": { + "title": "Error", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "name", + "error" + ], + "title": "MCPConnectorImportFailure", + "type": "object" + }, + "MCPConnectorImportRequest": { + "properties": { + "mcp_servers": { + "anyOf": [ + { + "additionalProperties": { + "$ref": "#/components/schemas/MCPConnectorEntry" + }, + "type": "object" + }, + { + "items": { + "$ref": "#/components/schemas/MCPConnectorEntry" + }, + "type": "array" + } + ], + "title": "Mcp Servers" + } + }, + "required": [ + "mcp_servers" + ], + "title": "MCPConnectorImportRequest", + "type": "object" + }, + "MCPConnectorImportResponse": { + "properties": { + "errors": { + "items": { + "$ref": "#/components/schemas/MCPConnectorImportFailure" + }, + "title": "Errors", + "type": "array" + }, + "imported": { + "items": { + "$ref": "#/components/schemas/MCPConnectorImportResult" + }, + "title": "Imported", + "type": "array" + }, + "skipped": { + "items": { + "$ref": "#/components/schemas/MCPConnectorImportSkipped" + }, + "title": "Skipped", + "type": "array" + } + }, + "required": [ + "imported", + "skipped", + "errors" + ], + "title": "MCPConnectorImportResponse", + "type": "object" + }, + "MCPConnectorImportResult": { + "properties": { + "alias": { + "title": "Alias", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "server_id": { + "title": "Server Id", + "type": "string" + } + }, + "required": [ + "name", + "server_id", + "alias" + ], + "title": "MCPConnectorImportResult", + "type": "object" + }, + "MCPConnectorImportSkipped": { + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "reason": { + "title": "Reason", + "type": "string" + } + }, + "required": [ + "name", + "reason" + ], + "title": "MCPConnectorImportSkipped", + "type": "object" + }, "MCPCredentials": { "properties": { + "audience": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Audience" + }, "auth_value": { "anyOf": [ { @@ -15243,6 +15613,17 @@ ], "title": "Aws Session Token" }, + "client_assertion_signing_alg": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Assertion Signing Alg" + }, "client_id": { "anyOf": [ { @@ -15254,6 +15635,28 @@ ], "title": "Client Id" }, + "client_private_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Private Key" + }, + "client_private_key_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Private Key Id" + }, "client_secret": { "anyOf": [ { @@ -15265,6 +15668,42 @@ ], "title": "Client Secret" }, + "id_jag_resource": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Resource" + }, + "id_jag_resource_token_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Resource Token Endpoint" + }, + "redirect_uris": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Redirect Uris" + }, "scopes": { "anyOf": [ { @@ -15278,11 +15717,113 @@ } ], "title": "Scopes" + }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" + }, + "token_endpoint_auth_method": { + "anyOf": [ + { + "enum": [ + "client_secret_basic", + "client_secret_post" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Endpoint Auth Method" + }, + "token_exchange_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Endpoint" + }, + "token_exchange_profile": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Profile" + }, + "upstream_resource": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Upstream Resource" } }, "title": "MCPCredentials", "type": "object" }, + "MCPEnvVar": { + "description": "One environment variable for an MCP server.\n\nVariables can be interpolated into ``static_headers`` using ``${NAME}``\nsyntax. ``scope=global`` values are stored on the server. ``scope=user``\nvalues are stored per-user in ``LiteLLM_MCPUserEnvVars`` and supplied by\neach user.", + "properties": { + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "name": { + "title": "Name", + "type": "string" + }, + "scope": { + "$ref": "#/components/schemas/MCPEnvVarScope", + "default": "global" + }, + "value": { + "default": "", + "title": "Value", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "MCPEnvVar", + "type": "object" + }, + "MCPEnvVarScope": { + "description": "Scope for an MCP server environment variable.\n\n- ``global``: value is provided by the admin and used for all users.\n- ``user``: each user must provide their own value via the per-user\n env-var endpoint. The admin-supplied ``value`` is treated as a\n placeholder/hint and is not used at request time.", + "enum": [ + "global", + "user" + ], + "title": "MCPEnvVarScope", + "type": "string" + }, "MCPOAuthUserCredentialRequest": { "description": "Stores a user's OAuth2 token for an OpenAPI MCP server.", "properties": { @@ -15537,6 +16078,112 @@ "title": "MCPUserCredentialResponse", "type": "object" }, + "MCPUserEnvVarSpec": { + "description": "Describes one per-user env var slot for the calling user.\n\nStored values are write-only: the status only reports whether a value\n``is_set`` and never echoes the decrypted secret back to the client.", + "properties": { + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "is_set": { + "default": false, + "title": "Is Set", + "type": "boolean" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "MCPUserEnvVarSpec", + "type": "object" + }, + "MCPUserEnvVarsRequest": { + "description": "Payload for storing the calling user's per-user env var values.", + "properties": { + "values": { + "additionalProperties": { + "type": "string" + }, + "title": "Values", + "type": "object" + } + }, + "required": [ + "values" + ], + "title": "MCPUserEnvVarsRequest", + "type": "object" + }, + "MCPUserEnvVarsStatus": { + "description": "Per-user env var status for a single MCP server.", + "properties": { + "alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Alias" + }, + "missing_count": { + "default": 0, + "title": "Missing Count", + "type": "integer" + }, + "required": { + "items": { + "$ref": "#/components/schemas/MCPUserEnvVarSpec" + }, + "title": "Required", + "type": "array" + }, + "server_id": { + "title": "Server Id", + "type": "string" + }, + "server_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Server Name" + }, + "setup_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Setup Url" + } + }, + "required": [ + "server_id" + ], + "title": "MCPUserEnvVarsStatus", + "type": "object" + }, "MakeMCPServersPublicRequest": { "properties": { "mcp_server_ids": { @@ -15604,6 +16251,17 @@ "title": "Args", "type": "array" }, + "audience": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Audience" + }, "auth_type": { "anyOf": [ { @@ -15615,7 +16273,11 @@ "authorization", "oauth2", "aws_sigv4", - "token" + "token", + "oauth2_token_exchange", + "oauth2_id_jag", + "true_passthrough", + "oauth_delegate" ], "type": "string" }, @@ -15680,6 +16342,22 @@ } ] }, + "dcr_bridge": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Dcr Bridge" + }, + "delegate_auth_to_upstream": { + "default": false, + "title": "Delegate Auth To Upstream", + "type": "boolean" + }, "description": { "anyOf": [ { @@ -15698,6 +16376,20 @@ "title": "Env", "type": "object" }, + "env_vars": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/MCPEnvVar" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Env Vars" + }, "extra_headers": { "anyOf": [ { @@ -15728,6 +16420,28 @@ "title": "Is Byok", "type": "boolean" }, + "issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Issuer" + }, + "max_concurrent_requests": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Max Concurrent Requests" + }, "mcp_access_groups": { "items": { "type": "string" @@ -15762,6 +16476,11 @@ ], "title": "Oauth2 Flow" }, + "oauth_passthrough": { + "default": false, + "title": "Oauth Passthrough", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -15831,6 +16550,17 @@ ], "title": "Static Headers" }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" + }, "submitted_at": { "anyOf": [ { @@ -15856,6 +16586,39 @@ "description": "Server-managed: set by the endpoint; caller values are overridden.", "title": "Submitted By" }, + "timeout": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Timeout" + }, + "token_exchange_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Endpoint" + }, + "token_exchange_profile": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Profile" + }, "token_url": { "anyOf": [ { @@ -16008,6 +16771,17 @@ "title": "Args", "type": "array" }, + "audience": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Audience" + }, "auth_type": { "anyOf": [ { @@ -16019,7 +16793,11 @@ "authorization", "oauth2", "aws_sigv4", - "token" + "token", + "oauth2_token_exchange", + "oauth2_id_jag", + "true_passthrough", + "oauth_delegate" ], "type": "string" }, @@ -16084,6 +16862,22 @@ } ] }, + "dcr_bridge": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Dcr Bridge" + }, + "delegate_auth_to_upstream": { + "default": false, + "title": "Delegate Auth To Upstream", + "type": "boolean" + }, "description": { "anyOf": [ { @@ -16102,6 +16896,20 @@ "title": "Env", "type": "object" }, + "env_vars": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/MCPEnvVar" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Env Vars" + }, "extra_headers": { "anyOf": [ { @@ -16132,6 +16940,28 @@ "title": "Is Byok", "type": "boolean" }, + "issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Issuer" + }, + "max_concurrent_requests": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Max Concurrent Requests" + }, "mcp_access_groups": { "items": { "type": "string" @@ -16151,6 +16981,26 @@ ], "title": "Mcp Info" }, + "oauth2_flow": { + "anyOf": [ + { + "enum": [ + "client_credentials", + "authorization_code" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Oauth2 Flow" + }, + "oauth_passthrough": { + "default": false, + "title": "Oauth Passthrough", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -16213,6 +17063,50 @@ ], "title": "Static Headers" }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" + }, + "timeout": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Timeout" + }, + "token_exchange_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Endpoint" + }, + "token_exchange_profile": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Profile" + }, "token_url": { "anyOf": [ { @@ -16331,6 +17225,13 @@ }, "ValidationError": { "properties": { + "ctx": { + "title": "Context", + "type": "object" + }, + "input": { + "title": "Input" + }, "loc": { "items": { "anyOf": [ @@ -16600,6 +17501,18 @@ "description": "Filter MCP servers by team scope. When provided, returns only servers the team has access to plus globally available (allow_all_keys) servers. Used by the Create Key UI to show team-scoped MCP servers.", "title": "Team Id" } + }, + { + "description": "Annotate each returned server with connected_app_reachable: whether a connected app authorized by the calling user (a gateway OAuth session) is served this server on the aggregate MCP endpoint.", + "in": "query", + "name": "connected_app_view", + "required": false, + "schema": { + "default": false, + "description": "Annotate each returned server with connected_app_reachable: whether a connected app authorized by the calling user (a gateway OAuth session) is served this server on the aggregate MCP endpoint.", + "title": "Connected App View", + "type": "boolean" + } } ], "responses": { @@ -16827,6 +17740,53 @@ ] } }, + "/v1/mcp/server/import": { + "post": { + "description": "Bulk-import MCP connectors from Anthropic mcpServers or mcp_servers JSON", + "operationId": "import_mcp_servers_v1_mcp_server_import_post", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPConnectorImportRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPConnectorImportResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Import Mcp Servers", + "tags": [ + "mcp_management" + ] + } + }, "/v1/mcp/server/oauth/session": { "post": { "description": "Temporarily cache an MCP server in memory without writing to the database", @@ -17438,6 +18398,156 @@ ] } }, + "/v1/mcp/server/{server_id}/user-env-vars": { + "delete": { + "description": "Clear the calling user's per-user MCP env var values for this server.", + "operationId": "clear_mcp_user_env_vars_v1_mcp_server__server_id__user_env_vars_delete", + "parameters": [ + { + "in": "path", + "name": "server_id", + "required": true, + "schema": { + "title": "Server Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPUserEnvVarsStatus" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Clear Mcp User Env Vars", + "tags": [ + "mcp_management" + ] + }, + "get": { + "description": "Return the calling user's per-user MCP env var status for this server.", + "operationId": "get_mcp_user_env_vars_v1_mcp_server__server_id__user_env_vars_get", + "parameters": [ + { + "in": "path", + "name": "server_id", + "required": true, + "schema": { + "title": "Server Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPUserEnvVarsStatus" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Mcp User Env Vars", + "tags": [ + "mcp_management" + ] + }, + "post": { + "description": "Store the calling user's per-user MCP env var values for this server. Submitted values are merged over any previously stored values, so you only send the fields you want to set or change; a variable omitted (or sent empty) keeps its stored value. Use DELETE to clear all stored values.", + "operationId": "store_mcp_user_env_vars_v1_mcp_server__server_id__user_env_vars_post", + "parameters": [ + { + "in": "path", + "name": "server_id", + "required": true, + "schema": { + "title": "Server Id", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPUserEnvVarsRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPUserEnvVarsStatus" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Store Mcp User Env Vars", + "tags": [ + "mcp_management" + ] + } + }, "/v1/mcp/tools": { "get": { "description": "Get all MCP tools available for the current key, including those from access groups", @@ -17746,6 +18856,37 @@ "mcp_management" ] } + }, + "/v1/mcp/user-env-vars/status": { + "get": { + "description": "Per-user MCP env var status across every server the user can access. Used by the dashboard to highlight servers with missing per-user vars.", + "operationId": "list_mcp_user_env_var_status_v1_mcp_user_env_vars_status_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/MCPUserEnvVarsStatus" + }, + "title": "Response List Mcp User Env Var Status V1 Mcp User Env Vars Status Get", + "type": "array" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "List Mcp User Env Var Status", + "tags": [ + "mcp_management" + ] + } } } }, diff --git a/litellm/proxy/management_endpoints/mcp_connector_import.py b/litellm/proxy/management_endpoints/mcp_connector_import.py new file mode 100644 index 00000000000..d576131ae1b --- /dev/null +++ b/litellm/proxy/management_endpoints/mcp_connector_import.py @@ -0,0 +1,168 @@ +""" +Convert Anthropic MCP connector definitions into LiteLLM MCP server create requests. + +Two interchange shapes are accepted: +- the ``mcpServers`` mapping used by Claude Desktop / Claude Code config files +- the ``mcp_servers`` array used by the Anthropic Messages API MCP connector +""" + +import re +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, ValidationError + +from litellm.proxy._types import MCPApprovalStatus, NewMCPServerRequest +from litellm.types.mcp import MCPAuth, MCPCredentials, MCPTransport + + +class MCPConnectorEntry(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + name: str | None = None + type: str | None = None + url: str | None = None + authorization_token: str | None = Field( + default=None, validation_alias=AliasChoices("authorization_token", "authorizationToken") + ) + headers: Mapping[str, str] | None = None + command: str | None = None + args: tuple[str, ...] = Field(default_factory=tuple) + env: Mapping[str, str] = Field(default_factory=dict) + description: str | None = None + + +class MCPConnectorImportRequest(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + mcp_servers: Mapping[str, MCPConnectorEntry] | tuple[MCPConnectorEntry, ...] = Field( + validation_alias=AliasChoices("mcp_servers", "mcpServers") + ) + + +@dataclass(frozen=True, slots=True) +class ConvertedConnector: + name: str + request: NewMCPServerRequest + + +@dataclass(frozen=True, slots=True) +class ConnectorConversionError: + name: str + error: str + + +class MCPConnectorImportResult(BaseModel): + name: str + server_id: str + alias: str + + +class MCPConnectorImportSkipped(BaseModel): + name: str + reason: str + + +class MCPConnectorImportFailure(BaseModel): + name: str + error: str + + +class MCPConnectorImportResponse(BaseModel): + imported: tuple[MCPConnectorImportResult, ...] + skipped: tuple[MCPConnectorImportSkipped, ...] + errors: tuple[MCPConnectorImportFailure, ...] + + +_INVALID_SERVER_NAME_CHARS: Final = re.compile(r"[^A-Za-z0-9_]") + + +def sanitize_connector_name(name: str) -> str: + sanitized: Final = re.sub(r"_+", "_", _INVALID_SERVER_NAME_CHARS.sub("_", name.strip())).strip("_") + return sanitized + + +_SSE_TYPES: Final = frozenset({"sse"}) +_URL_TYPES: Final = frozenset({"url", "http", "streamable_http", "streamable-http", "sse", ""}) + + +def _convert_entry(name: str, entry: MCPConnectorEntry) -> ConvertedConnector | ConnectorConversionError: + sanitized_name: Final = sanitize_connector_name(name) + if not sanitized_name: + return ConnectorConversionError(name=name, error="Connector name is empty after sanitization.") + + if entry.url and entry.command: + return ConnectorConversionError(name=name, error="Connector cannot have both a url and a command.") + + if entry.command: + try: + stdio_request: Final = NewMCPServerRequest( + server_name=sanitized_name, + alias=sanitized_name, + description=entry.description, + approval_status=MCPApprovalStatus.active, + transport=MCPTransport.stdio, + command=entry.command, + args=entry.args, + env=entry.env, + ) + except ValidationError as e: + return ConnectorConversionError(name=name, error=_first_validation_message(e)) + return ConvertedConnector(name=name, request=stdio_request) + + if not entry.url: + return ConnectorConversionError(name=name, error="Connector must have either a url or a command.") + + entry_type: Final = (entry.type or "").lower() + if entry_type not in _URL_TYPES: + return ConnectorConversionError(name=name, error=f"Unsupported connector type '{entry.type}'.") + + transport: Final = MCPTransport.sse if entry_type in _SSE_TYPES else MCPTransport.http + credentials: Final = _bearer_credentials(entry.authorization_token) + try: + remote_request: Final = NewMCPServerRequest( + server_name=sanitized_name, + alias=sanitized_name, + description=entry.description, + approval_status=MCPApprovalStatus.active, + transport=transport, + url=entry.url, + auth_type=MCPAuth.bearer_token if entry.authorization_token else MCPAuth.none, + credentials=credentials, + static_headers=entry.headers, + ) + except ValidationError as e: + return ConnectorConversionError(name=name, error=_first_validation_message(e)) + return ConvertedConnector(name=name, request=remote_request) + + +def _bearer_credentials(token: str | None) -> MCPCredentials | None: + if not token: + return None + credentials: Final[MCPCredentials] = {"auth_value": token} + return credentials + + +def _first_validation_message(error: ValidationError) -> str: + messages: Final = tuple(str(detail.get("msg", "")) for detail in error.errors()) + return messages[0] if messages else str(error) + + +def convert_connector_entries( + payload: MCPConnectorImportRequest, +) -> tuple[ConvertedConnector | ConnectorConversionError, ...]: + servers: Final = payload.mcp_servers + if isinstance(servers, Mapping): + return tuple(_convert_entry(name, entry) for name, entry in servers.items()) + return tuple( + _convert_entry(entry.name or "", entry) if entry.name else _named_entry_error(index, entry) + for index, entry in enumerate(servers) + ) + + +def _named_entry_error(index: int, entry: MCPConnectorEntry) -> ConnectorConversionError: + return ConnectorConversionError( + name=entry.url or f"entry {index}", + error="Connector entries in list form must have a name.", + ) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 54a591a5e1a..6583d5405f8 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -133,6 +133,7 @@ if MCP_AVAILABLE: delete_mcp_server, delete_user_credential, delete_user_env_vars, + get_all_mcp_servers, get_all_mcp_servers_for_user, get_draft_mcp_server, get_mcp_server, @@ -199,6 +200,16 @@ if MCP_AVAILABLE: populate_request_with_path_params, ) from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view + from litellm.proxy.management_endpoints.mcp_connector_import import ( + ConnectorConversionError, + ConvertedConnector, + MCPConnectorImportFailure, + MCPConnectorImportRequest, + MCPConnectorImportResponse, + MCPConnectorImportResult, + MCPConnectorImportSkipped, + convert_connector_entries, + ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.types.mcp import ( MCP_ADMIN_CONFIG_CREDENTIAL_KEYS, @@ -1607,6 +1618,112 @@ if MCP_AVAILABLE: return _redact_mcp_credentials(new_mcp_server) + @router.post( + "/server/import", + description="Bulk-import MCP connectors from Anthropic mcpServers or mcp_servers JSON", + dependencies=(Depends(user_api_key_auth),), + response_model=MCPConnectorImportResponse, + status_code=status.HTTP_200_OK, + ) + @management_endpoint_wrapper + async def import_mcp_servers( + payload: MCPConnectorImportRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection + ): + """ + Bulk-import MCP connectors. Accepts the Claude Desktop / Claude Code + ``mcpServers`` mapping or the Anthropic Messages API ``mcp_servers`` + array, creates each entry as a LiteLLM MCP server, and returns + per-entry results so partial imports are visible to the caller. + """ + prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") + + if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict + "error": "User does not have permission to import mcp servers. You can only import mcp servers if you are a PROXY_ADMIN." + }, + ) + + conversions: Final = convert_connector_entries(payload) + existing_servers: Final = await get_all_mcp_servers(prisma_client) + existing_names: Final = frozenset( + name for server in existing_servers for name in (server.alias, server.server_name) if name + ) + + def _classify( + index: int, conversion: ConvertedConnector | ConnectorConversionError + ) -> ConvertedConnector | ConnectorConversionError | MCPConnectorImportSkipped: + if isinstance(conversion, ConnectorConversionError): + return conversion + alias: Final = conversion.request.alias or "" + if alias in existing_names: + return MCPConnectorImportSkipped( + name=conversion.name, reason=f"An MCP server named '{alias}' already exists." + ) + earlier_aliases: Final = frozenset( + earlier.request.alias or "" + for earlier in conversions[:index] + if isinstance(earlier, ConvertedConnector) + ) + if alias in earlier_aliases: + return MCPConnectorImportSkipped( + name=conversion.name, reason=f"Duplicate connector name '{alias}' in the import payload." + ) + return conversion + + async def _create( + conversion: ConvertedConnector, + ) -> MCPConnectorImportResult | MCPConnectorImportFailure: + try: + validate_and_normalize_mcp_server_payload(conversion.request) + except HTTPException as e: + error_text: Final = ( + str(e.detail.get("error", e.detail)) if isinstance(e.detail, dict) else str(e.detail) + ) + return MCPConnectorImportFailure(name=conversion.name, error=error_text) + try: + created: Final = await create_mcp_server( + prisma_client, + conversion.request, + touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, + ) + except Exception as e: # noqa: BLE001 # any create failure must become a per-entry error, not a 500 + verbose_proxy_logger.exception("Error importing mcp server %s: %s", conversion.name, e) + return MCPConnectorImportFailure(name=conversion.name, error=str(e)) + return MCPConnectorImportResult( + name=conversion.name, server_id=created.server_id, alias=created.alias or "" + ) + + classified: Final = tuple(_classify(index, conversion) for index, conversion in enumerate(conversions)) + outcomes: Final = tuple( + [ + await _create(entry) if isinstance(entry, ConvertedConnector) else entry for entry in classified + ] # mutable-ok: await is illegal in a generator expression here + ) + + imported: Final = tuple(entry for entry in outcomes if isinstance(entry, MCPConnectorImportResult)) + if imported: + # Best-effort registry refresh, mirroring add_mcp_server: rows are + # already committed, so a refresh failure must not surface as a 500. + try: + await global_mcp_server_manager.reload_servers_from_database() + except Exception as e: # noqa: BLE001 # rows are committed; a refresh failure must not surface as a 500 + verbose_proxy_logger.exception("MCP connector import committed but registry refresh failed: %s", e) + + return MCPConnectorImportResponse( + imported=imported, + skipped=tuple(entry for entry in outcomes if isinstance(entry, MCPConnectorImportSkipped)), + errors=tuple( + MCPConnectorImportFailure(name=entry.name, error=entry.error) + if isinstance(entry, ConnectorConversionError) + else entry + for entry in outcomes + if isinstance(entry, (ConnectorConversionError, MCPConnectorImportFailure)) + ), + ) + @router.post( "/server/oauth/session", description="Temporarily cache an MCP server in memory without writing to the database", diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py new file mode 100644 index 00000000000..9f2874bcbf3 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py @@ -0,0 +1,157 @@ +import pytest + +from litellm.proxy.management_endpoints.mcp_connector_import import ( + ConnectorConversionError, + ConvertedConnector, + MCPConnectorImportRequest, + convert_connector_entries, + sanitize_connector_name, +) +from litellm.types.mcp import MCPAuth, MCPTransport + + +def _single(payload: dict) -> ConvertedConnector | ConnectorConversionError: + results = convert_connector_entries(MCPConnectorImportRequest.model_validate(payload)) + assert len(results) == 1 + return results[0] + + +class TestSanitizeConnectorName: + @pytest.mark.parametrize( + "raw,expected", + [ + ("my-server", "my_server"), + (" spaced name ", "spaced_name"), + ("already_ok", "already_ok"), + ("a.b.c", "a_b_c"), + ("---", ""), + ], + ) + def test_sanitizes_to_mcp_safe_names(self, raw, expected): + assert sanitize_connector_name(raw) == expected + + +class TestConvertMcpServersMapping: + def test_url_connector_with_authorization_token(self): + result = _single( + { + "mcpServers": { + "github-mcp": { + "url": "https://api.example.com/mcp", + "authorization_token": "secret-token", + "headers": {"X-Env": "prod"}, + "description": "GitHub connector", + } + } + } + ) + assert isinstance(result, ConvertedConnector) + assert result.request.server_name == "github_mcp" + assert result.request.alias == "github_mcp" + assert result.request.transport == MCPTransport.http + assert result.request.url == "https://api.example.com/mcp" + assert result.request.auth_type == MCPAuth.bearer_token + assert result.request.credentials == {"auth_value": "secret-token"} + assert result.request.static_headers == {"X-Env": "prod"} + assert result.request.description == "GitHub connector" + + def test_camel_case_authorization_token_alias(self): + result = _single( + {"mcpServers": {"srv": {"url": "https://x.example/mcp", "authorizationToken": "tok"}}} + ) + assert isinstance(result, ConvertedConnector) + assert result.request.credentials == {"auth_value": "tok"} + + def test_url_connector_without_token_uses_no_auth(self): + result = _single({"mcpServers": {"open": {"url": "https://open.example/mcp"}}}) + assert isinstance(result, ConvertedConnector) + assert result.request.auth_type == MCPAuth.none + assert result.request.credentials is None + + def test_sse_type_maps_to_sse_transport(self): + result = _single({"mcpServers": {"legacy": {"type": "sse", "url": "https://sse.example/mcp"}}}) + assert isinstance(result, ConvertedConnector) + assert result.request.transport == MCPTransport.sse + + def test_stdio_connector(self): + result = _single( + { + "mcpServers": { + "local": { + "command": "npx", + "args": ["-y", "@example/mcp-server"], + "env": {"API_KEY": "value"}, + } + } + } + ) + assert isinstance(result, ConvertedConnector) + assert result.request.transport == MCPTransport.stdio + assert result.request.command == "npx" + assert result.request.args == ["-y", "@example/mcp-server"] + assert result.request.env == {"API_KEY": "value"} + + def test_disallowed_stdio_command_returns_error(self): + result = _single({"mcpServers": {"evil": {"command": "rm", "args": ["-rf", "/"]}}}) + assert isinstance(result, ConnectorConversionError) + assert "not in the allowed commands list" in result.error + + def test_unsupported_type_returns_error(self): + result = _single({"mcpServers": {"ws": {"type": "websocket", "url": "wss://x.example"}}}) + assert isinstance(result, ConnectorConversionError) + assert "Unsupported connector type" in result.error + + def test_missing_url_and_command_returns_error(self): + result = _single({"mcpServers": {"empty": {}}}) + assert isinstance(result, ConnectorConversionError) + assert "either a url or a command" in result.error + + def test_url_and_command_together_returns_error(self): + result = _single({"mcpServers": {"both": {"url": "https://x.example/mcp", "command": "npx"}}}) + assert isinstance(result, ConnectorConversionError) + assert "both a url and a command" in result.error + + def test_name_empty_after_sanitization_returns_error(self): + result = _single({"mcpServers": {"---": {"url": "https://x.example/mcp"}}}) + assert isinstance(result, ConnectorConversionError) + assert "empty after sanitization" in result.error + + +class TestConvertMcpServersList: + def test_anthropic_messages_api_list_shape(self): + result = _single( + { + "mcp_servers": [ + { + "type": "url", + "url": "https://mcp.example.com/sse", + "name": "deepwiki", + "authorization_token": "tok", + } + ] + } + ) + assert isinstance(result, ConvertedConnector) + assert result.request.server_name == "deepwiki" + assert result.request.transport == MCPTransport.http + assert result.request.credentials == {"auth_value": "tok"} + + def test_list_entry_without_name_returns_error(self): + result = _single({"mcp_servers": [{"type": "url", "url": "https://x.example/mcp"}]}) + assert isinstance(result, ConnectorConversionError) + assert "must have a name" in result.error + + def test_partial_conversion_preserves_per_entry_results(self): + results = convert_connector_entries( + MCPConnectorImportRequest.model_validate( + { + "mcpServers": { + "good": {"url": "https://good.example/mcp"}, + "bad": {"type": "websocket", "url": "wss://bad.example"}, + } + } + ) + ) + assert len(results) == 2 + assert isinstance(results[0], ConvertedConnector) + assert isinstance(results[1], ConnectorConversionError) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 0d639e1cb6a..7b3580f166a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -6766,3 +6766,143 @@ class TestConnectedAppViewAnnotation: assert all(server.connected_app_reachable is None for server in result) reload_mock.assert_not_awaited() + + +class TestImportMCPServers: + """Bulk connector import must be admin-only and report per-entry outcomes.""" + + @staticmethod + def _import_patches(existing_servers, create_mock, mock_manager): + return ( + patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers", + AsyncMock(return_value=existing_servers), + ), + patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", + create_mock, + ), + patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ) + + @pytest.mark.asyncio + async def test_non_admin_is_rejected(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + MCPConnectorImportRequest, + import_mcp_servers, + ) + + payload = MCPConnectorImportRequest.model_validate( + {"mcpServers": {"srv": {"url": "https://x.example/mcp"}}} + ) + caller = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.INTERNAL_USER) + + with patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ): + with pytest.raises(HTTPException) as exc_info: + await import_mcp_servers(payload=payload, user_api_key_dict=caller) + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_import_reports_imported_skipped_and_errors(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + MCPConnectorImportRequest, + import_mcp_servers, + ) + + payload = MCPConnectorImportRequest.model_validate( + { + "mcpServers": { + "new-server": {"url": "https://new.example/mcp", "authorization_token": "tok"}, + "existing": {"url": "https://existing.example/mcp"}, + "broken": {"type": "websocket", "url": "wss://x.example"}, + } + } + ) + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") + existing = generate_mock_mcp_server_db_record(server_id="existing-1", alias="existing") + created = generate_mock_mcp_server_db_record(server_id="created-1", alias="new_server") + create_mock = AsyncMock(return_value=created) + mock_manager = MagicMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ExitStack() as stack: + for p in self._import_patches([existing], create_mock, mock_manager): + stack.enter_context(p) + result = await import_mcp_servers(payload=payload, user_api_key_dict=admin) + + assert [entry.name for entry in result.imported] == ["new-server"] + assert result.imported[0].server_id == "created-1" + assert [entry.name for entry in result.skipped] == ["existing"] + assert "already exists" in result.skipped[0].reason + assert [entry.name for entry in result.errors] == ["broken"] + create_mock.assert_awaited_once() + sent_request = create_mock.await_args[0][1] + assert sent_request.credentials == {"auth_value": "tok"} + mock_manager.reload_servers_from_database.assert_awaited_once() + + @pytest.mark.asyncio + async def test_duplicate_names_within_payload_are_skipped(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + MCPConnectorImportRequest, + import_mcp_servers, + ) + + payload = MCPConnectorImportRequest.model_validate( + { + "mcp_servers": [ + {"type": "url", "url": "https://a.example/mcp", "name": "dup srv"}, + {"type": "url", "url": "https://b.example/mcp", "name": "dup-srv"}, + ] + } + ) + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") + created = generate_mock_mcp_server_db_record(server_id="created-1", alias="dup_srv") + create_mock = AsyncMock(return_value=created) + mock_manager = MagicMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ExitStack() as stack: + for p in self._import_patches([], create_mock, mock_manager): + stack.enter_context(p) + result = await import_mcp_servers(payload=payload, user_api_key_dict=admin) + + assert len(result.imported) == 1 + assert len(result.skipped) == 1 + assert "Duplicate connector name" in result.skipped[0].reason + create_mock.assert_awaited_once() + + @pytest.mark.asyncio + async def test_no_imports_skips_registry_refresh(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + MCPConnectorImportRequest, + import_mcp_servers, + ) + + payload = MCPConnectorImportRequest.model_validate( + {"mcpServers": {"existing": {"url": "https://existing.example/mcp"}}} + ) + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") + existing = generate_mock_mcp_server_db_record(server_id="existing-1", alias="existing") + create_mock = AsyncMock() + mock_manager = MagicMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ExitStack() as stack: + for p in self._import_patches([existing], create_mock, mock_manager): + stack.enter_context(p) + result = await import_mcp_servers(payload=payload, user_api_key_dict=admin) + + assert result.imported == () + create_mock.assert_not_awaited() + mock_manager.reload_servers_from_database.assert_not_awaited() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ImportMCPServers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ImportMCPServers.tsx new file mode 100644 index 00000000000..d04160bca3e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ImportMCPServers.tsx @@ -0,0 +1,134 @@ +import React, { useState } from "react"; +import { Button } from "@/components/ui/button"; +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { Textarea } from "@/components/ui/textarea"; +import { Alert, AlertTitle } from "@/components/shared/Alert"; +import { importMCPServers } from "@/components/networking"; +import { toast } from "@/lib/toast"; +import { MCPConnectorImportResponse, parseConnectorConfig } from "./importConnectorConfig"; + +interface ImportMCPServersProps { + accessToken: string; + open: boolean; + onClose: () => void; + onImported: () => void; +} + +const PLACEHOLDER = `{ + "mcpServers": { + "my_server": { + "url": "https://example.com/mcp", + "authorization_token": "..." + } + } +}`; + +const ImportMCPServers: React.FC = ({ accessToken, open, onClose, onImported }) => { + const [configText, setConfigText] = useState(""); + const [parseError, setParseError] = useState(null); + const [isImporting, setIsImporting] = useState(false); + const [result, setResult] = useState(null); + + const handleClose = () => { + setConfigText(""); + setParseError(null); + setResult(null); + onClose(); + }; + + const handleImport = async () => { + const parsed = parseConnectorConfig(configText); + if (!parsed.ok) { + setParseError(parsed.error); + return; + } + setParseError(null); + setIsImporting(true); + try { + const response = (await importMCPServers(accessToken, parsed.payload)) as MCPConnectorImportResponse; + setResult(response); + if (response.imported.length > 0) { + toast.success(`Imported ${response.imported.length} MCP server${response.imported.length === 1 ? "" : "s"}`); + onImported(); + } + } catch (error) { + console.error("Failed to import MCP servers:", error); + setParseError("Import request failed. Check the proxy logs for details."); + } finally { + setIsImporting(false); + } + }; + + return ( + !isOpen && handleClose()}> + + + Import MCP Connectors + +
+

+ Paste an Anthropic connector configuration: the mcpServers mapping from a Claude Desktop / + Claude Code config file, or the mcp_servers array from the Anthropic Messages API. +

+