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