From 3f38730d2c3059964a2232856e5fa5e688f58a59 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 13:31:15 -0800 Subject: [PATCH] refactor(typing): clear about 1,100 basedpyright Any errors across 303 files (#49986) * refactor(typing): clear basedpyright Any errors (wave 1) * refactor(typing): clear basedpyright Any errors (wave 2) * refactor(typing): clear basedpyright Any errors (wave 3) * refactor(typing): clear basedpyright Any errors (wave 4) * refactor(typing): clear basedpyright Any errors (wave 5) * refactor(typing): clear basedpyright Any errors (wave 6) * fix(typing): keep the cached-plan fallback's direct query_first calls so span naming resolves * fix(logging): keep callbacks running when litellm_params is a params model should_run_callback and the sync success and failure handlers read litellm_params through any object with a get method again, as they did before this branch, so a logging hook that returns a GenericLiteLLMParams no longer stops sync logging for that request Adds behavior tests for the changed lines codecov listed as not run * fix: read the regional uplift for every data residency again and reuse one dict TypeAdapter The typed lookup only read the eu and us multipliers, so a deployment that set regional_processing_uplift_multiplier_fedramp stopped applying it. /model/update and the model write paths now share one module-level TypeAdapter instead of building one per call. * fix: bill regional uplift multipliers float() accepts, bytes included * fix: raise the in operator's own error from one shared json_contains, and name the bad uplift type --- litellm/a2a_protocol/card_resolver.py | 37 ++++- .../a2a_protocol/exception_mapping_utils.py | 5 +- .../litellm_completion_bridge/handler.py | 7 +- litellm/a2a_protocol/main.py | 15 +- litellm/caching/caching.py | 25 +++- litellm/caching/caching_handler.py | 39 +++--- litellm/caching/redis_cache.py | 27 +++- litellm/caching/redis_semantic_cache.py | 5 +- .../transformation.py | 9 +- litellm/cost_calculator.py | 33 ++++- litellm/exceptions.py | 22 ++- litellm/harness/endpoint.py | 47 +++++-- litellm/images/utils.py | 21 ++- .../SlackAlerting/user_spend_alerts.py | 4 +- litellm/integrations/argilla.py | 23 +++- litellm/integrations/arize/_utils.py | 23 +++- .../azure_sentinel/azure_sentinel.py | 18 ++- .../bitbucket/bitbucket_client.py | 32 ++++- .../bitbucket/bitbucket_prompt_manager.py | 13 +- .../clickhouse/clickhouse_spend_logger.py | 40 ++++-- litellm/integrations/custom_guardrail.py | 13 +- litellm/integrations/email_alerting.py | 25 +++- litellm/integrations/focus/database.py | 4 +- litellm/integrations/galileo.py | 13 +- litellm/integrations/gcs_bucket/gcs_bucket.py | 14 +- .../generic_prompt_manager.py | 2 +- litellm/integrations/humanloop.py | 24 ++-- litellm/integrations/lago.py | 14 +- .../integrations/langfuse/langfuse_otel.py | 5 +- .../langfuse/langfuse_otel_attributes.py | 13 +- litellm/integrations/mlflow.py | 25 +++- litellm/integrations/openmeter.py | 15 +- litellm/integrations/opentelemetry.py | 83 +++++++----- .../base_otel_llm_obs_attributes.py | 3 +- litellm/integrations/otel/plumbing/metrics.py | 20 ++- litellm/integrations/prompt_layer.py | 9 +- litellm/integrations/weave/weave_otel.py | 45 ++++-- .../websearch_interception/tools.py | 22 ++- litellm/integrations/weights_biases.py | 12 +- .../transformation.py | 12 +- litellm/interactions/streaming_iterator.py | 19 ++- .../litellm_core_utils/audio_utils/utils.py | 2 +- litellm/litellm_core_utils/core_helpers.py | 25 +++- litellm/litellm_core_utils/dd_tracing.py | 17 ++- litellm/litellm_core_utils/litellm_logging.py | 93 +++++++++---- .../usage_object_transformation.py | 35 +++-- .../litellm_core_utils/llm_cost_calc/utils.py | 12 +- .../prompt_templates/common_utils.py | 4 +- .../prompt_templates/factory.py | 33 +++-- litellm/litellm_core_utils/redact_messages.py | 15 +- .../litellm_core_utils/streaming_handler.py | 38 ++++-- .../a2a/chat/guardrail_translation/handler.py | 77 ++++++----- litellm/llms/anthropic/chat/handler.py | 2 +- litellm/llms/anthropic/common_utils.py | 4 +- .../native/agentic_streaming_iterator.py | 14 +- .../responses_adapters/streaming_iterator.py | 15 +- .../responses_adapters/transformation.py | 17 ++- .../anthropic/messages/stream_assembly.py | 36 +++-- .../llms/apiserpent/search/transformation.py | 29 ++-- litellm/llms/azure/azure.py | 12 +- .../llms/azure/responses/transformation.py | 5 +- litellm/llms/azure_ai/agents/handler.py | 6 +- .../llms/azure_ai/agents/transformation.py | 2 +- .../anthropic/count_tokens/handler.py | 8 +- .../anthropic/count_tokens/token_counter.py | 6 +- .../anthropic/messages_transformation.py | 10 +- .../image_edit/flux2_transformation.py | 13 +- .../guardrail_translation/base_translation.py | 16 ++- .../base_llm/managed_resources/isolation.py | 14 +- litellm/llms/bedrock/chat/invoke_handler.py | 15 +- litellm/llms/bedrock/common_utils.py | 32 +++-- .../count_tokens/bedrock_token_counter.py | 2 +- litellm/llms/bedrock/count_tokens/handler.py | 10 +- .../bedrock/count_tokens/transformation.py | 4 +- litellm/llms/bedrock/files/transformation.py | 25 +++- ...n_nova_canvas_image_edit_transformation.py | 17 ++- .../image_edit/stability_transformation.py | 34 +++-- .../guardrail_translation/handler.py | 5 +- .../bedrock/vector_stores/transformation.py | 54 ++++++-- .../image_edit/transformation.py | 34 +++-- litellm/llms/bytez/chat/transformation.py | 8 +- litellm/llms/cohere/common_utils.py | 10 +- litellm/llms/custom_httpx/llm_http_handler.py | 38 ++++-- .../llms/dashscope/embed/transformation.py | 22 ++- .../llms/dataforseo/search/transformation.py | 33 ++++- litellm/llms/gemini/common_utils.py | 34 +++-- litellm/llms/gemini/count_tokens/handler.py | 55 ++++++-- .../gemini/interactions/transformation.py | 19 ++- litellm/llms/gigachat/chat/streaming.py | 47 +++++-- litellm/llms/huggingface/embedding/handler.py | 51 +++++-- .../huggingface/embedding/transformation.py | 5 +- .../litellm_proxy/skills/code_execution.py | 19 ++- .../litellm_proxy/skills/transformation.py | 6 +- .../milvus/vector_stores/transformation.py | 9 +- litellm/llms/mistral/chat/transformation.py | 12 +- litellm/llms/oci/chat/cohere.py | 2 +- litellm/llms/oci/chat/generic.py | 3 +- litellm/llms/oci/chat/transformation.py | 4 +- litellm/llms/ollama/common_utils.py | 24 +++- litellm/llms/ollama/completion/handler.py | 21 ++- .../openai/responses/count_tokens/handler.py | 9 +- .../responses/count_tokens/token_counter.py | 2 +- .../llms/recraft/image_edit/transformation.py | 33 +++-- litellm/llms/sap/credentials.py | 25 ++-- .../snowflake/embedding/transformation.py | 32 ++++- .../llms/soniox/realtime/transformation.py | 23 ++-- .../valkey/vector_stores/transformation.py | 6 +- litellm/llms/vertex_ai/common_utils.py | 2 +- .../llms/vertex_ai/files/transformation.py | 2 +- .../llms/vertex_ai/gemini/transformation.py | 24 +++- .../vertex_imagen_transformation.py | 33 +++-- .../count_tokens/handler.py | 5 +- .../llms/vertex_ai/videos/transformation.py | 11 +- litellm/main.py | 15 +- litellm/proxy/_lazy_features.py | 4 +- litellm/proxy/admin_mcp.py | 6 +- litellm/proxy/batches_endpoints/endpoints.py | 23 +++- litellm/proxy/caching_routes.py | 5 +- litellm/proxy/client/chat.py | 6 +- litellm/proxy/client/cli/commands/http.py | 5 +- litellm/proxy/client/credentials.py | 19 +-- litellm/proxy/client/health.py | 10 +- litellm/proxy/client/http_client.py | 11 +- litellm/proxy/client/keys.py | 19 +-- litellm/proxy/client/models.py | 11 +- litellm/proxy/client/teams.py | 7 +- litellm/proxy/common_request_processing.py | 19 ++- .../common_utils/cache_pydantic_utils.py | 4 +- litellm/proxy/common_utils/callback_utils.py | 2 +- .../proxy/common_utils/http_parsing_utils.py | 36 +++-- .../proxy/common_utils/user_api_key_cache.py | 5 +- .../container_endpoints/handler_factory.py | 8 +- litellm/proxy/db/health_check_latest.py | 21 +-- litellm/proxy/db/prisma_client.py | 20 ++- litellm/proxy/db/proxy_worker_heartbeat.py | 4 +- litellm/proxy/db/routing_prisma_wrapper.py | 22 ++- .../guardrail_hooks/aporia_ai/aporia_ai.py | 7 +- .../guardrails/guardrail_hooks/azure/base.py | 8 +- .../guardrail_hooks/bedrock_guardrails.py | 58 ++++---- .../guardrail_hooks/compresr/compresr.py | 16 ++- .../guardrail_hooks/dynamoai/dynamoai.py | 2 +- .../guardrail_hooks/enkryptai/enkryptai.py | 2 +- .../guardrail_hooks/grayswan/grayswan.py | 12 +- .../guardrails_ai/guardrails_ai.py | 6 +- .../guardrail_hooks/headroom/headroom.py | 13 +- .../guardrail_hooks/lakera_ai_v2.py | 3 +- .../litellm_content_filter/content_filter.py | 5 +- .../litellm_content_filter/patterns.py | 24 +++- .../mcp_jwt_signer/__init__.py | 2 +- .../mcp_jwt_signer/mcp_jwt_signer.py | 10 +- .../guardrail_hooks/microsoft_purview/base.py | 27 ++-- .../model_armor/model_armor.py | 17 ++- .../guardrails/guardrail_hooks/onyx/onyx.py | 14 +- .../ovalix/ovalix_extraction.py | 9 +- .../guardrail_hooks/pangea/pangea.py | 14 +- .../panw_prisma_airs/panw_prisma_airs.py | 13 +- .../guardrail_hooks/pillar/pillar.py | 11 +- .../guardrail_hooks/qualifire/qualifire.py | 16 ++- .../guardrail_hooks/singulr/singulr.py | 78 ++++++----- litellm/proxy/hooks/batch_rate_limiter.py | 6 +- .../proxy/hooks/key_management_event_hooks.py | 16 ++- litellm/proxy/hooks/litellm_skills/main.py | 40 ++++-- litellm/proxy/hooks/max_iterations_limiter.py | 4 +- .../proxy/hooks/mcp_semantic_filter/hook.py | 59 ++++---- .../hooks/parallel_request_limiter_v3.py | 3 +- litellm/proxy/hooks/responses_id_security.py | 8 +- litellm/proxy/litellm_pre_call_utils.py | 32 +++-- .../auto_router_endpoints.py | 6 +- .../cache_settings_endpoints.py | 13 +- .../internal_user_endpoints.py | 4 +- .../key_management_endpoints.py | 3 +- .../management_v1/budgets.py | 8 +- .../model_management_endpoints.py | 55 ++++---- .../policy_endpoints/ai_policy_suggester.py | 22 ++- litellm/proxy/management_endpoints/ui_sso.py | 4 +- .../user_agent_analytics_endpoints.py | 4 +- .../object_permission_utils.py | 19 ++- .../management_helpers/team_list_sorting.py | 4 +- .../management_helpers/user_list_sorting.py | 5 +- .../guardrail_translation/handler.py | 14 +- .../proxy/mcp_server/oauth2_flow_backfill.py | 15 +- .../proxy/mcp_server/oauth2_token_cache.py | 5 +- .../mcp_server/openapi_to_mcp_generator.py | 25 +++- litellm/proxy/mcp_server/operations.py | 10 +- .../sso_assertion_store.py | 9 +- litellm/proxy/mcp_server/utils.py | 9 +- .../openai_files_endpoints/common_utils.py | 11 +- .../openai_passthrough_logging_handler.py | 6 +- .../pass_through_endpoints.py | 44 +++--- .../passthrough_guardrails.py | 7 +- litellm/proxy/proxy_cli.py | 6 +- litellm/proxy/proxy_server.py | 58 +++++--- .../spend_tracking/key_metadata_recovery.py | 6 +- .../spend_tracking/ptu_flat_cost_rollup.py | 11 +- .../spend_tracking/spend_capture_rate.py | 10 +- litellm/proxy/types_utils/utils.py | 4 +- litellm/proxy/utils.py | 65 ++++++--- .../proxy/vector_store_endpoints/endpoints.py | 45 ++++-- .../management_endpoints.py | 12 +- litellm/proxy/video_endpoints/utils.py | 8 +- litellm/rag/main.py | 3 +- litellm/rag/rag_query.py | 60 +++++--- .../autorouter_session_repository.py | 16 +-- litellm/repositories/base_repository.py | 7 +- litellm/repositories/budget_repository.py | 20 +-- .../repositories/credentials_repository.py | 20 ++- .../object_permission_repository.py | 20 +-- .../repositories/organization_repository.py | 20 +-- litellm/repositories/project_repository.py | 16 +-- litellm/repositories/table_repositories.py | 30 +++- litellm/repositories/team_repository.py | 11 +- litellm/repositories/user_repository.py | 18 ++- .../verification_token_repository.py | 4 +- .../session_handler.py | 20 ++- .../streaming_iterator.py | 15 +- .../transformation.py | 41 +++--- litellm/responses/main.py | 10 +- .../mcp/litellm_proxy_mcp_handler.py | 2 +- .../responses/mcp/mcp_streaming_iterator.py | 10 +- litellm/responses/utils.py | 5 +- litellm/router.py | 11 +- litellm/router_config.py | 5 +- litellm/router_deprecated.py | 6 +- litellm/router_python_host.py | 5 +- .../router_strategy/adaptive_router/hooks.py | 86 ++++++++---- .../adaptive_router/signals.py | 26 ++-- .../auto_router/auto_router.py | 19 ++- litellm/router_strategy/budget_limiter.py | 92 +++++++------ .../complexity_router/complexity_router.py | 21 ++- .../quality_router/quality_router.py | 4 +- litellm/router_utils/cooldown_callbacks.py | 4 +- .../router_utils/fallback_event_handlers.py | 4 +- litellm/router_utils/handle_error.py | 2 +- litellm/router_utils/health_state_cache.py | 10 +- litellm/skills/main.py | 21 +-- .../text_completions/streaming_iterator.py | 6 +- litellm/types/llms/anthropic_skills.py | 3 +- litellm/utils.py | 44 ++++-- .../vector_stores/vector_store_registry.py | 2 +- litellm/videos/utils.py | 10 +- tests/unit/a2a_protocol/test_card_resolver.py | 32 +++++ tests/unit/caching/test_caching.py | 23 ++++ .../unit/caching/test_redis_semantic_cache.py | 23 ++++ tests/unit/conftest.py | 8 +- .../test_prometheus_unit_tests.py | 22 +-- .../test_bitbucket_prompt_manager.py | 18 +++ .../integrations/focus/test_focus_database.py | 16 +++ .../unit/integrations/test_email_alerting.py | 33 +++++ tests/unit/integrations/test_galileo.py | 29 ++++ tests/unit/integrations/test_humanloop.py | 57 +++++++- tests/unit/integrations/test_lago.py | 55 ++++++++ tests/unit/integrations/test_langfuse_otel.py | 29 ++++ tests/unit/integrations/test_opentelemetry.py | 51 +++++++ tests/unit/integrations/test_prompt_layer.py | 75 ++++++++++ tests/unit/integrations/test_wandb.py | 67 +++++++++ tests/unit/integrations/test_weave_otel.py | 51 +++++++ ...test_gemini_interactions_transformation.py | 11 ++ .../test_litellm_responses_bridge.py | 22 ++- .../llm_cost_calc/test_llm_cost_calc_utils.py | 40 ++++++ .../litellm_core_utils/test_core_helpers.py | 38 +++++- .../litellm_core_utils/test_dd_tracing.py | 34 +++++ .../test_litellm_logging.py | 49 +++++++ .../test_streaming_handler.py | 47 +++++++ .../test_a2a_guardrail_handler.py | 57 ++++++++ .../messages/test_stream_assembly.py | 22 +++ ...test_azure_ai_image_edit_transformation.py | 22 +++ .../test_stability_transformation.py | 66 +++++++++ .../llms/cohere/test_cohere_common_utils.py | 18 +++ .../custom_httpx/test_llm_http_handler.py | 41 ++++++ ...test_dashscope_embedding_transformation.py | 22 ++- .../test_dataforseo_search_transformation.py | 15 ++ .../llms/gemini/test_gemini_common_utils.py | 14 ++ .../chat/test_gigachat_chat_streaming.py | 33 +++++ .../test_huggingface_embedding_handler.py | 47 +++++++ .../skills/test_code_execution.py | 22 +++ .../test_mistral_chat_transformation.py | 19 +++ .../unit/llms/ollama/test_ollama_embedding.py | 40 +++++- .../test_recraft_image_edit_transformation.py | 8 ++ tests/unit/llms/sap/test_sap_fetch_creds.py | 12 ++ .../test_vertex_ai_gemini_transformation.py | 24 ++++ ...est_vertex_ai_image_edit_transformation.py | 27 +++- tests/unit/proxy/client/test_credentials.py | 22 +++ tests/unit/proxy/client/test_keys.py | 24 ++++ tests/unit/proxy/client/test_teams.py | 41 ++++++ .../common_utils/test_http_parsing_utils.py | 24 +++- .../guardrails_ai/test_guardrails_ai.py | 55 ++++++++ .../guardrail_hooks/test_lakera_ai_v2.py | 36 +++++ .../guardrail_hooks/test_microsoft_purview.py | 8 ++ .../guardrails/test_pillar_guardrails.py | 32 +++++ .../hooks/test_key_management_event_hooks.py | 31 +++++ .../hooks/test_max_iterations_limiter.py | 39 ++++++ .../test_user_agent_analytics_endpoints.py | 54 ++++++++ .../test_object_permission_utils.py | 20 +++ .../mcp_server/test_semantic_tool_filter.py | 59 ++++++++ .../test_pass_through_endpoints.py | 104 ++++++++++++++ tests/unit/proxy/test_caching_routes.py | 17 +++ .../test_prisma_client_get_data.py | 11 ++ .../test_vector_store_endpoints.py | 128 ++++++++++++++++++ tests/unit/responses/test_responses_utils.py | 10 +- .../adaptive_router/test_hooks.py | 41 ++++++ .../adaptive_router/test_signals.py | 32 ++++- tests/unit/test_main.py | 35 +++++ tests/unit/test_utils.py | 29 ++++ 303 files changed, 5111 insertions(+), 1411 deletions(-) create mode 100644 tests/unit/integrations/test_email_alerting.py create mode 100644 tests/unit/integrations/test_lago.py create mode 100644 tests/unit/integrations/test_prompt_layer.py create mode 100644 tests/unit/llms/cohere/test_cohere_common_utils.py create mode 100644 tests/unit/proxy/management_endpoints/test_user_agent_analytics_endpoints.py diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index 08e95f99c87..441b326f4d2 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/litellm/a2a_protocol/card_resolver.py @@ -4,13 +4,14 @@ Custom A2A Card Resolver for LiteLLM. Extends the A2A SDK's card resolver to support multiple well-known paths. """ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable from litellm._logging import verbose_logger from litellm.a2a_protocol.exceptions import A2AAgentCardDiscoveryError from litellm.constants import LOCALHOST_URL_PATTERNS +from litellm.litellm_core_utils.core_helpers import has_attributes if TYPE_CHECKING: from a2a.types import AgentCard @@ -130,15 +131,34 @@ def normalize_agent_card_interfaces(agent_card: "AgentCard", protocol_version: s return normalized +class _UrlCarrier(Protocol): + url: str + + +class _InterfacesCarrier(Protocol): + @property + def supported_interfaces(self) -> Sequence[_UrlCarrier]: ... + + +def _url_of(value: object) -> str | None: + return value.url if has_attributes(value, _UrlCarrier, "url") else None + + +def _supported_interfaces(agent_card: object) -> Sequence[_UrlCarrier] | None: + if has_attributes(agent_card, _InterfacesCarrier, "supported_interfaces"): + return agent_card.supported_interfaces + return None + + def get_agent_card_url(agent_card: "AgentCard") -> str | None: """Return the agent endpoint URL from the resolved SDK card.""" - url: Final = getattr(agent_card, "url", None) + url: Final = _url_of(agent_card) if url: return url - interfaces: Final = getattr(agent_card, "supported_interfaces", None) + interfaces: Final = _supported_interfaces(agent_card) if interfaces: - return getattr(interfaces[0], "url", None) + return _url_of(interfaces[0]) return None @@ -168,16 +188,16 @@ def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard": Returns: The agent card with the URL fixed if necessary """ - card_url: Final = getattr(agent_card, "url", None) + card_url: Final = _url_of(agent_card) if card_url and is_localhost_or_internal_url(card_url): # Normalize base_url to ensure it ends with / fixed_url: Final = base_url.rstrip("/") + "/" agent_card.url = fixed_url - interfaces: Final = getattr(agent_card, "supported_interfaces", None) + interfaces: Final = _supported_interfaces(agent_card) if interfaces: - interface_url: Final = getattr(interfaces[0], "url", None) + interface_url: Final = _url_of(interfaces[0]) if interface_url and is_localhost_or_internal_url(interface_url): interfaces[0].url = base_url.rstrip("/") + "/" @@ -194,6 +214,9 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): - /agentCard/v1.0 """ + if TYPE_CHECKING: + base_url: str + async def get_agent_card( self, relative_card_path: str | None = None, diff --git a/litellm/a2a_protocol/exception_mapping_utils.py b/litellm/a2a_protocol/exception_mapping_utils.py index 1abf32915f5..1fe8b2bf0af 100644 --- a/litellm/a2a_protocol/exception_mapping_utils.py +++ b/litellm/a2a_protocol/exception_mapping_utils.py @@ -4,7 +4,7 @@ A2A Protocol Exception Mapping Utils. Maps A2A SDK exceptions to LiteLLM A2A exception types. """ -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger from litellm.a2a_protocol.card_resolver import ( @@ -21,6 +21,7 @@ from litellm.constants import CONNECTION_ERROR_PATTERNS if TYPE_CHECKING: from a2a.client import Client as A2AClientType + from a2a.types import AgentCard _A2A_KIND_DIALECT_HINT: Final = ( @@ -165,7 +166,7 @@ def map_a2a_exception( async def handle_a2a_localhost_retry( error: A2ALocalhostURLError, - agent_card: Any, + agent_card: "AgentCard | None", a2a_client: "A2AClientType", is_streaming: bool = False, ) -> "A2AClientType": diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 1e003a6e4bc..06a910a08fd 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -10,7 +10,7 @@ A2A Streaming Events (in order): 4. Status update (kind: "status-update") - Final status "completed" with final=true """ -from collections.abc import AsyncIterator, Callable, Coroutine, Mapping +from collections.abc import AsyncIterator, Mapping from types import MappingProxyType from typing import Any, Final @@ -126,10 +126,7 @@ class A2ACompletionBridgeHandler: @staticmethod async def _acompletion(completion_params: Mapping[str, object]) -> ModelResponse | CustomStreamWrapper: - acompletion_fn: Final[Callable[..., Coroutine[object, object, ModelResponse | CustomStreamWrapper]]] = vars( - litellm - )["acompletion"] - return await acompletion_fn(**completion_params) + return await litellm.acompletion(**completion_params) @staticmethod async def handle_non_streaming( diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index dcbe7a6506e..7d617613b0c 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -14,7 +14,7 @@ import datetime import uuid from collections.abc import AsyncIterator, Coroutine, Mapping from types import MappingProxyType, ModuleType -from typing import TYPE_CHECKING, Final, Optional, cast +from typing import TYPE_CHECKING, Final, Optional, Protocol, cast from typing_extensions import TypedDict, Unpack @@ -24,6 +24,7 @@ from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator from litellm.a2a_protocol.utils import A2ARequestUtils from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -36,7 +37,6 @@ if TYPE_CHECKING: from a2a.client import Client as A2AClientType from a2a.client import ClientCallContext as A2ACallContextType from a2a.compat.v0_3.types import ( - AgentCard, Message, SendMessageRequest, SendMessageResponse, @@ -45,6 +45,7 @@ if TYPE_CHECKING: SendStreamingMessageSuccessResponse, Task, ) + from a2a.types import AgentCard from a2a.types.a2a_pb2 import SendMessageRequest as CoreSendMessageRequest from a2a.types.a2a_pb2 import StreamResponse as CoreStreamResponse @@ -95,6 +96,12 @@ class _SendMessageOptions(TypedDict, total=False, extra_items=object): pass +class _StreamingRequestConverter(Protocol): + def to_core_send_message_request( + self, compat_req: "SendStreamingMessageRequest", / + ) -> "CoreSendMessageRequest": ... + + def _set_usage_on_logging_obj( kwargs: Mapping[str, object], prompt_tokens: int, @@ -375,7 +382,9 @@ async def _stream_messages( "The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk" ) - pb_request: Final[CoreSendMessageRequest] = _a2a_conversions.to_core_send_message_request(request) + pb_request: Final[CoreSendMessageRequest] = require_attributes( + _a2a_conversions, _StreamingRequestConverter, "to_core_send_message_request" + ).to_core_send_message_request(request) async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)): compat_chunk = _to_compat_stream_response(event, request_id=request.id) yield SendStreamingMessageResponse(root=compat_chunk) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 3b7ea66ad6b..b9cef2ea301 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -13,10 +13,10 @@ import json import logging import time import traceback -from collections.abc import Generator, Mapping +from collections.abc import Awaitable, Callable, Generator, Mapping from contextlib import contextmanager from enum import Enum -from typing import Any, Final, Literal, cast +from typing import Any, Final, Literal, Protocol, cast from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -25,6 +25,7 @@ from litellm._internal_context import current_service_target, service_target from litellm._logging import verbose_logger from litellm.constants import CACHED_STREAMING_CHUNK_DELAY from litellm.integrations.otel.runtime import phase_span +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.model_param_helper import ( ModelParamHelper, embeds_thought_signature_in_tool_call_id, @@ -52,6 +53,16 @@ _REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter( ) +class _SupportsAsyncPing(Protocol): + @property + def ping(self) -> Callable[[], Awaitable[object]] | None: ... + + +class _SupportsAsyncDeleteCacheKeys(Protocol): + @property + def delete_cache_keys(self) -> Callable[[object], Awaitable[object]] | None: ... + + def _keyed_provider_params_in_extra_body(kwargs: Mapping[str, object]) -> str: raw_extra_body: Final = kwargs.get("extra_body") if not isinstance(raw_extra_body, dict): @@ -1032,14 +1043,16 @@ class Cache: cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs) await self.cache.batch_cache_write(cache_key, cached_data, **kwargs) - async def ping(self): - cache_ping: Final = getattr(self.cache, "ping") + async def ping(self) -> object: + cache_ping: Final = require_attributes(self.cache, _SupportsAsyncPing, "ping").ping if cache_ping: return await cache_ping() return None - async def delete_cache_keys(self, keys): - cache_delete_cache_keys: Final = getattr(self.cache, "delete_cache_keys") + async def delete_cache_keys(self, keys: object) -> object: + cache_delete_cache_keys: Final = require_attributes( + self.cache, _SupportsAsyncDeleteCacheKeys, "delete_cache_keys" + ).delete_cache_keys if cache_delete_cache_keys: return await cache_delete_cache_keys(keys) return None diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 1d80d6952a4..82804d2d707 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -22,6 +22,7 @@ from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError +from typing_extensions import TypedDict import litellm from litellm._internal_context import is_internal_call, post_response_phase @@ -90,6 +91,10 @@ class CachingHandlerResponse(LiteLLMBaseModel): in_memory_cache_obj: Final = InMemoryCache() +class _CacheCallKwargs(TypedDict, total=False, extra_items=object): + pass + + def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str, object]: """ The caching handler is stored on the Logging object @@ -455,7 +460,7 @@ class LLMCachingHandler: if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function): args = args or () # Now that we confirmed caching will happen, prepare kwargs - new_kwargs: Final = kwargs.copy() + new_kwargs: Final = dict[str, object](kwargs) new_kwargs.update( convert_args_to_kwargs( self.original_function, @@ -470,7 +475,7 @@ class LLMCachingHandler: print_verbose("Checking Sync Cache") logging_obj.reset_caching_details() with response_cache_phase("get"): - cached_result = litellm.cache.get_cache(**new_kwargs) + cached_result = litellm.cache.get_cache(**_CacheCallKwargs(**new_kwargs)) if is_response_without_output(cached_result): verbose_logger.debug("LiteLLM Cache: cached response has no output, treating it as a miss") return CachingHandlerResponse(cached_result=None) @@ -480,7 +485,7 @@ class LLMCachingHandler: pass else: call_type = original_function.__name__ - cached_result = self._convert_cached_result_to_model_response( + converted_result: Final = self._convert_cached_result_to_model_response( cached_result=cached_result, call_type=call_type, kwargs=kwargs, @@ -508,16 +513,16 @@ class LLMCachingHandler: logging_obj=logging_obj, model=f"{custom_llm_provider}/{model}", kwargs=kwargs, - cached_result=cached_result, + cached_result=converted_result, is_async=False, custom_llm_provider=custom_llm_provider, ) if not is_internal_call.get() and not _should_defer_streaming_cache_hit_callbacks( - cached_result=cached_result + cached_result=converted_result ): logging_obj.handle_sync_cache_hit_success_callbacks( - result=cached_result, + result=converted_result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, @@ -527,8 +532,8 @@ class LLMCachingHandler: or self.request_kwargs.get("cache_key") or litellm.cache.get_cache_key(**self.request_kwargs) ) - _set_cached_hidden_param(cached_result, "cache_key", cache_key) - return CachingHandlerResponse(cached_result=cached_result) + _set_cached_hidden_param(converted_result, "cache_key", cache_key) + return CachingHandlerResponse(cached_result=converted_result) return CachingHandlerResponse(cached_result=cached_result) _sync_get_cache = sync_get_cache @@ -950,22 +955,14 @@ class LLMCachingHandler: def _convert_cached_result_to_model_response( self, - cached_result: Any, + cached_result: object, call_type: str, kwargs: dict[str, object], logging_obj: LiteLLMLoggingObj, model: str, args: tuple[object, ...], custom_llm_provider: str | None = None, - ) -> ( - ModelResponse - | TextCompletionResponse - | EmbeddingResponse - | RerankResponse - | TranscriptionResponse - | CustomStreamWrapper - | None - ): + ) -> object: """ Internal method to process the cached result @@ -1141,7 +1138,7 @@ class LLMCachingHandler: self, result: object, original_function: Callable, - kwargs: dict[str, Any], + kwargs: Mapping[str, object], args: tuple[object, ...] | None = None, ): """ @@ -1169,7 +1166,7 @@ class LLMCachingHandler: return cache: Final = litellm.cache - new_kwargs: Final = kwargs.copy() + new_kwargs: Final = dict(kwargs) new_kwargs.update( convert_args_to_kwargs( original_function, @@ -1206,7 +1203,7 @@ class LLMCachingHandler: ) ) else: - create_cache_write_task(lambda: cache.async_add_cache(result, **new_kwargs)) + create_cache_write_task(lambda: cache.async_add_cache(result, **_CacheCallKwargs(**new_kwargs))) def sync_set_cache( self, diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 62ee5d83634..4d9d26fc70c 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -40,7 +40,7 @@ from litellm.constants import ( REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, REDIS_TIMEOUT_LOG_INTERVAL, ) -from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs +from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, has_attributes from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.redis.store_settings import RedisStoreSettings, breaker_settings from litellm.types.caching import ( @@ -97,6 +97,22 @@ class _AsyncRedisCommands(Protocol): def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... +class _AsyncLuaScript(Protocol): + def __call__( + self, keys: Sequence[str], args: Sequence[str | bytes | int | float], client: object + ) -> Awaitable[object]: ... + + +class _SupportsRegisterScript(Protocol): + def register_script(self, script: str, /) -> _AsyncLuaScript: ... + + +class _SupportsScriptLoad(Protocol): + def script_load(self, script: str, /) -> object: ... + + def evalsha(self, sha: object, numkeys: int, /, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... + + _GENERIC_CALLER_MODULES: Final = frozenset( { __name__, @@ -844,6 +860,9 @@ class RedisCache(BaseCache): def _async_commands(self) -> _AsyncRedisCommands: return self.init_async_client() + def _async_client_object(self) -> object: + return self.init_async_client() + def check_and_fix_namespace(self, key: str) -> str: """ Make sure each key starts with the given namespace @@ -1117,8 +1136,8 @@ class RedisCache(BaseCache): Kept separate from async_register_script so each loop caches its own executor; see that method for why the binding must be per loop. """ - _redis_client: Final[Any] = self.init_async_client() - if hasattr(_redis_client, "register_script"): + _redis_client: Final = self._async_client_object() + if has_attributes(_redis_client, _SupportsRegisterScript, "register_script"): registered_script: Final = _redis_client.register_script(script) async def standalone_executor( @@ -1131,7 +1150,7 @@ class RedisCache(BaseCache): return standalone_executor - if hasattr(_redis_client, "script_load"): + if has_attributes(_redis_client, _SupportsScriptLoad, "script_load"): script_sha: Final = _redis_client.script_load(script) async def cluster_executor( diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index 7268511fe7e..c884eb1b001 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -43,6 +43,7 @@ if TYPE_CHECKING: _CACHE_RESPONSE_ADAPTER: Final[TypeAdapter[str | bytes | None]] = TypeAdapter( str | bytes | None, config=ConfigDict(strict=True, hide_input_in_errors=True) ) +_DECODED_CACHE_VALUE: Final = TypeAdapter(object) class RedisSemanticCache(BaseCache): @@ -407,10 +408,10 @@ class RedisSemanticCache(BaseCache): # Convert string representation to Python object try: - return json.loads(response_text) + return _DECODED_CACHE_VALUE.validate_python(json.loads(response_text)) except json.JSONDecodeError: try: - return ast.literal_eval(response_text) + return _DECODED_CACHE_VALUE.validate_python(ast.literal_eval(response_text)) except (ValueError, SyntaxError) as e: print_verbose(f"Error parsing cached response: {e}") return None diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index b24c75e69bc..fc7a6048538 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -39,7 +39,7 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import ( parse_chat_allowed_tools_tool_choice, responses_allowed_tools_tool_choice, ) -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, iter_items from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params from litellm.litellm_core_utils.prompt_templates.prompt_cache_breakpoint_utils import ( prompt_cache_breakpoint_with_readable_mode, @@ -349,6 +349,9 @@ _TOOL_SEARCH_REPLAY_ITEM_TYPES: Final = frozenset({"reasoning", *_TOOL_SEARCH_IT _TOOL_SEARCH_ITEMS_KEY: Final = "tool_search_items" +_DECODED_JSON: Final = TypeAdapter(object) + + def _json_field(value: object, key: str) -> object: return cast(Mapping[str, object], value).get(key) if isinstance(value, dict) else None # cast-ok: untyped JSON @@ -857,8 +860,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Handle message items with output_text content if item_type == "message": - content_list: Final = item.get("content", []) - for content_item in content_list: + content_list: Final = _DECODED_JSON.validate_python(item.get("content", [])) + for content_item in iter_items(content_list): if isinstance(content_item, dict): content_type = content_item.get("type") if content_type == "output_text": diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 3910a554905..5f7d646e150 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -8,7 +8,7 @@ from datetime import datetime from functools import lru_cache from itertools import accumulate from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, cast from httpx import Response from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError @@ -24,7 +24,12 @@ from litellm.constants import ( REALTIME_METERED_USAGE_KEY, ) from litellm.litellm_core_utils.audio_utils.utils import normalized_audio_duration_seconds -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, billed_usage_cost +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + attribute_or_none, + billed_usage_cost, + require_attributes, +) from litellm.litellm_core_utils.cost_discount import resolve_cost_discount from litellm.litellm_core_utils.decisions_pricing import ( decisions_pricing_fields, @@ -1188,6 +1193,10 @@ def _model_contains_known_llm_provider(model: str) -> bool: return _provider_prefix in LlmProvidersSet +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + def _get_response_model(completion_response: object) -> str | None: """ Extract the model name from a completion response object. @@ -1219,7 +1228,7 @@ _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: Final[dict] = { } -def _map_traffic_type_to_service_tier(traffic_type: str | None) -> str | None: +def _map_traffic_type_to_service_tier(traffic_type: object) -> str | None: """ Map a Gemini usageMetadata.trafficType value to a LiteLLM service_tier string. @@ -1273,11 +1282,13 @@ def _served_service_tier(completion_response: object, usage_object: Usage | None usage_tier: Final = _extract_service_tier(usage_object) if isinstance(usage_tier, str): return usage_tier - hidden_params: Final = getattr(completion_response, "_hidden_params", None) + hidden_params: Final[object] = getattr(completion_response, "_hidden_params", None) if hidden_params is None: return None - provider_specific: Final = hidden_params.get("provider_specific_fields") or {} - raw_traffic_type: Final = provider_specific.get("traffic_type") + provider_specific: Final[object] = ( + require_attributes(hidden_params, _SupportsGet, "get").get("provider_specific_fields") or {} + ) + raw_traffic_type: Final = require_attributes(provider_specific, _SupportsGet, "get").get("traffic_type") if not raw_traffic_type: return None return _map_traffic_type_to_service_tier(raw_traffic_type) or "default" @@ -1341,8 +1352,16 @@ def get_usage_object( return None +class _DumpedUsageFields(TypedDict, total=False, extra_items=object): + pass + + +def _dumped_usage_fields(usage_model: BaseModel) -> Mapping[str, object]: + return usage_model.model_dump() + + def _usage_from_model(usage_model: BaseModel) -> Usage: - return Usage(**usage_model.model_dump()) + return Usage(**_DumpedUsageFields(**_dumped_usage_fields(usage_model))) def _decisions_prompt_tokens_details(usage: DecisionsUsage | OpenAIDecisionUsage) -> PromptTokensDetailsWrapper: diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 6662cf4a2c2..56a5f136b62 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -12,10 +12,11 @@ import datetime import enum from collections.abc import Mapping, Sequence -from typing import Any, Final +from typing import Any, Final, Protocol, SupportsIndex, SupportsInt, runtime_checkable import httpx import openai +from typing_extensions import Buffer import litellm from litellm.types.utils import LiteLLMCommonStrings @@ -1250,6 +1251,19 @@ class BlockedPiiEntityError(Exception): super().__init__(self.message) +@runtime_checkable +class _SupportsTrunc(Protocol): + def __trunc__(self) -> int: ... + + +def _status_code_int(value: object) -> int: + if isinstance(value, (str, Buffer, SupportsInt, SupportsIndex, _SupportsTrunc)): + return int(value) + raise TypeError( + f"int() argument must be a string, a bytes-like object or a real number, not '{type(value).__name__}'" + ) + + class MidStreamFallbackError(ServiceUnavailableError): def __init__( self, @@ -1265,8 +1279,8 @@ class MidStreamFallbackError(ServiceUnavailableError): is_pre_first_chunk: bool = False, emitted_disqualifying_content: bool = False, ): - original_status: Final = getattr(original_exception, "status_code", None) - self.status_code = int(original_status) if original_status is not None else 503 + original_status: Final[object] = getattr(original_exception, "status_code", None) + self.status_code = _status_code_int(original_status) if original_status is not None else 503 self.message = f"litellm.MidStreamFallbackError: {message}" self.model = model self.llm_provider = llm_provider @@ -1309,7 +1323,7 @@ class MidStreamFallbackError(ServiceUnavailableError): ) # Restore the propagated status and original response/request objects - self.status_code = int(original_status) if original_status is not None else 503 + self.status_code = _status_code_int(original_status) if original_status is not None else 503 self.response = _saved_response self.request = _saved_request self.message = _saved_message diff --git a/litellm/harness/endpoint.py b/litellm/harness/endpoint.py index 71e2742fa2e..249bb2cebb7 100644 --- a/litellm/harness/endpoint.py +++ b/litellm/harness/endpoint.py @@ -14,10 +14,11 @@ import contextlib import itertools import json import logging +import math import secrets -from collections.abc import AsyncIterable, AsyncIterator, Mapping +from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping from dataclasses import dataclass -from types import MappingProxyType, ModuleType +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol import httpx @@ -36,6 +37,7 @@ from litellm.constants import ( from litellm.harness.context import GatewayTarget from litellm.harness.errors import HarnessError, HarnessInstallFailed from litellm.harness.types import Harness, Usage +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.custom_http import httpxSpecialProvider @@ -45,7 +47,7 @@ if TYPE_CHECKING: from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route - from uvicorn import Server + from uvicorn import Config, Server verbose_logger: Final = logging.getLogger("LiteLLM") _JSON_VALUE_ADAPTER: Final = TypeAdapter(object, config=ConfigDict(strict=True)) @@ -110,9 +112,22 @@ class _ResponsesModule(Protocol): StreamingResponse: type[StreamingResponse] +class _UvicornModule(Protocol): + Config: type[Config] + Server: type[Server] + + +class _CapturesSignals(Protocol): + capture_signals: Callable[[], contextlib.AbstractContextManager[None]] + + +class _InstallsSignalHandlers(Protocol): + install_signal_handlers: Callable[[], None] + + @dataclass(frozen=True) class _ServerDeps: - uvicorn: ModuleType + uvicorn: _UvicornModule applications: _ApplicationsModule routing: _RoutingModule responses: _ResponsesModule @@ -438,6 +453,14 @@ def _noop() -> None: return None +def _signal_capture(server: object) -> _CapturesSignals | None: + return server if has_attributes(server, _CapturesSignals, "capture_signals") else None + + +def _signal_handler_installer(server: object) -> _InstallsSignalHandlers | None: + return server if has_attributes(server, _InstallsSignalHandlers, "install_signal_handlers") else None + + class ModelEndpoint: """Local HTTP endpoint for one harness session. Use as an async context manager.""" @@ -521,7 +544,7 @@ class ModelEndpoint: await asyncio.sleep(DEFAULT_POLLING_INTERVAL) def _build_server(self, deps: _ServerDeps) -> Server: - config = deps.uvicorn.Config( + config: Final = deps.uvicorn.Config( self._build_app(deps), host=HARNESS_ENDPOINT_HOST, port=0, @@ -529,14 +552,16 @@ class ModelEndpoint: log_level="warning", access_log=False, lifespan="off", - timeout_graceful_shutdown=HARNESS_PROCESS_KILL_GRACE_SECONDS, + timeout_graceful_shutdown=math.ceil(HARNESS_PROCESS_KILL_GRACE_SECONDS), ) - server = deps.uvicorn.Server(config) + server: Final = deps.uvicorn.Server(config) # Never touch the host process's signal handlers. - if hasattr(server, "capture_signals"): - server.capture_signals = contextlib.nullcontext - if hasattr(server, "install_signal_handlers"): - server.install_signal_handlers = _noop + signal_capture: Final = _signal_capture(server) + if signal_capture is not None: + signal_capture.capture_signals = contextlib.nullcontext + handler_installer: Final = _signal_handler_installer(server) + if handler_installer is not None: + handler_installer.install_signal_handlers = _noop return server def _build_app(self, deps: _ServerDeps) -> Starlette: diff --git a/litellm/images/utils.py b/litellm/images/utils.py index d1a2dd76af0..efd37065e33 100644 --- a/litellm/images/utils.py +++ b/litellm/images/utils.py @@ -6,12 +6,25 @@ from pathlib import Path from typing import Any, Final, Protocol, cast, get_type_hints, runtime_checkable import litellm +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.litellm_core_utils.token_counter import get_image_type from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.types.files import FILE_MIME_TYPES, FileType from litellm.types.images.main import ImageEditOptionalRequestParams +class _ReadsBytes(Protocol): + def read(self, size: int, /) -> bytes: ... + + +class _Tells(Protocol): + def tell(self) -> int: ... + + +class _Seeks(Protocol): + def seek(self, offset: int, /) -> object: ... + + class ImageEditRequestUtils: @staticmethod def get_optional_params_image_edit( @@ -114,12 +127,12 @@ class ImageEditRequestUtils: bytes_data = image_data[:100] else: # For other types, try to read if possible - if hasattr(image_data, "read"): - current_pos = getattr(image_data, "tell", lambda: 0)() - if hasattr(image_data, "seek"): + if has_attributes(image_data, _ReadsBytes, "read"): + current_pos = image_data.tell() if has_attributes(image_data, _Tells, "tell") else 0 + if has_attributes(image_data, _Seeks, "seek"): image_data.seek(0) bytes_data = image_data.read(100) - if hasattr(image_data, "seek"): + if has_attributes(image_data, _Seeks, "seek"): image_data.seek(current_pos) else: return FILE_MIME_TYPES[FileType.PNG] # Default fallback diff --git a/litellm/integrations/SlackAlerting/user_spend_alerts.py b/litellm/integrations/SlackAlerting/user_spend_alerts.py index 08799a2f59f..67357cff403 100644 --- a/litellm/integrations/SlackAlerting/user_spend_alerts.py +++ b/litellm/integrations/SlackAlerting/user_spend_alerts.py @@ -53,11 +53,13 @@ async def fetch_user_spend_rows( today: datetime.date, baseline_days: int, ) -> tuple[UserSpendRow, ...]: + from litellm.proxy.utils import query_raw_rows + today_str: Final = today.strftime("%Y-%m-%d") month_start_str: Final = today.replace(day=1).strftime("%Y-%m-%d") baseline_start_str: Final = (today - datetime.timedelta(days=max(baseline_days, 1))).strftime("%Y-%m-%d") return USER_SPEND_ROWS_ADAPTER.validate_python( - await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str) + await query_raw_rows(prisma_client.db, USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str) ) diff --git a/litellm/integrations/argilla.py b/litellm/integrations/argilla.py index f873f274ab6..c20b2c332d3 100644 --- a/litellm/integrations/argilla.py +++ b/litellm/integrations/argilla.py @@ -7,11 +7,12 @@ import json import os import random import types -from collections.abc import Mapping -from typing import Final +from collections.abc import Mapping, Sequence +from typing import Final, Protocol import httpx from pydantic import BaseModel +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -29,6 +30,22 @@ from litellm.types.integrations.argilla import ( from litellm.types.utils import StandardLoggingPayload +class _ArgillaDataset(TypedDict): + id: ReadOnly[str] + + +class _ArgillaDatasetsBody(TypedDict, total=False): + items: ReadOnly[Sequence[_ArgillaDataset]] + + +class _ArgillaDatasetsResponse(Protocol): + def json(self) -> _ArgillaDatasetsBody: ... + + +def _datasets_body(response: _ArgillaDatasetsResponse) -> _ArgillaDatasetsBody: + return response.json() + + def is_serializable(value): non_serializable_types: Final = ( types.CoroutineType, @@ -102,7 +119,7 @@ class ArgillaLogger(CustomBatchLogger): url=f"{_credentials_base_url}/api/v1/me/datasets?name={_credentials_dataset_name}", headers={"X-Argilla-Api-Key": _credentials_api_key}, ) - json_response: Final = dataset_response.json() + json_response: Final = _datasets_body(dataset_response) if ( "items" in json_response and isinstance(json_response["items"], list) diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index fd912839c95..ea71300b8a6 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -1,6 +1,6 @@ import json from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Protocol from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict, override @@ -10,6 +10,7 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor BaseLLMObsOTELAttributes, safe_set_attribute, ) +from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes, subscript from litellm.litellm_core_utils.redact_messages import ( should_redact_message_logging, ) @@ -34,32 +35,40 @@ from litellm.integrations._types.open_inference import ( _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) +class _SupportsGetWithDefault(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + class ArizeOTELAttributes(BaseLLMObsOTELAttributes): @staticmethod @override - def set_messages(span: "Span", kwargs: dict[str, Any]): + def set_messages(span: "Span", kwargs: Mapping[str, object]): messages: Final = kwargs.get("messages") # for /chat/completions # https://docs.arize.com/arize/large-language-models/tracing/semantic-conventions if messages: - last_message: Final = messages[-1] + last_message: Final = subscript(messages, -1) safe_set_attribute( span, SpanAttributes.INPUT_VALUE, - last_message.get("content", ""), + require_attributes(last_message, _SupportsGetWithDefault, "get").get("content", ""), ) # LLM_INPUT_MESSAGES shows up under `input_messages` tab on the span page. - for idx, msg in enumerate(messages): + for idx, msg in enumerate(iter_items(messages)): prefix = f"{SpanAttributes.LLM_INPUT_MESSAGES}.{idx}" # Set the role per message. - safe_set_attribute(span, f"{prefix}.{MessageAttributes.MESSAGE_ROLE}", msg.get("role")) + safe_set_attribute( + span, + f"{prefix}.{MessageAttributes.MESSAGE_ROLE}", + require_attributes(msg, _SupportsGetWithDefault, "get").get("role"), + ) # Set the content per message. safe_set_attribute( span, f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}", - msg.get("content", ""), + require_attributes(msg, _SupportsGetWithDefault, "get").get("content", ""), ) # Additive: emit structured tool_calls / multimodal content diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index 84dfc770e8f..f8047f97e75 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -18,10 +18,11 @@ import time import traceback from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final, TypeVar +from typing import Final, Protocol, TypeVar from urllib.parse import urlparse import httpx +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.integrations.batch_utils import ( @@ -52,6 +53,19 @@ MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType( ) +class _OAuthTokenBody(TypedDict, total=False): + access_token: ReadOnly[str] + expires_in: ReadOnly[int] + + +class _OAuthTokenResponse(Protocol): + def json(self) -> _OAuthTokenBody: ... + + +def _oauth_token_body(response: _OAuthTokenResponse) -> _OAuthTokenBody: + return response.json() + + class AzureSentinelLogger(CustomBatchLogger): """ Logger that sends LiteLLM logs to Azure Sentinel via Azure Monitor Logs Ingestion API @@ -226,7 +240,7 @@ class AzureSentinelLogger(CustomBatchLogger): if response.status_code != 200: raise Exception(f"Failed to get OAuth2 token: {response.status_code} - {response.text}") - token_response: Final = response.json() + token_response: Final = _oauth_token_body(response) self.oauth_token = token_response.get("access_token") expires_in: Final = token_response.get("expires_in", 3600) diff --git a/litellm/integrations/bitbucket/bitbucket_client.py b/litellm/integrations/bitbucket/bitbucket_client.py index e4f5a50550e..5ebde5d4fb6 100644 --- a/litellm/integrations/bitbucket/bitbucket_client.py +++ b/litellm/integrations/bitbucket/bitbucket_client.py @@ -5,7 +5,7 @@ BitBucket API client for fetching .prompt files from BitBucket repositories. import base64 import urllib.parse from collections.abc import Mapping -from typing import Final, TypedDict +from typing import Final, Protocol, TypedDict from typing_extensions import NotRequired, ReadOnly @@ -36,6 +36,30 @@ class BitBucketFileMetadata(TypedDict): last_modified: ReadOnly[str | None] +class _SrcListingResponse(Protocol): + def json(self) -> BitBucketSrcListing: ... + + +class _RepositoryInfoResponse(Protocol): + def json(self) -> Mapping[str, object]: ... + + +class _BranchListingResponse(Protocol): + def json(self) -> BitBucketBranchListing: ... + + +def _src_listing(response: _SrcListingResponse) -> BitBucketSrcListing: + return response.json() + + +def _repository_info(response: _RepositoryInfoResponse) -> Mapping[str, object]: + return response.json() + + +def _branch_listing(response: _BranchListingResponse) -> BitBucketBranchListing: + return response.json() + + 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: @@ -166,7 +190,7 @@ class BitBucketClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final[BitBucketSrcListing] = response.json() + data: Final = _src_listing(response) return [ file_path for item in data.get("values", []) @@ -201,7 +225,7 @@ class BitBucketClient: try: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - return response.json() + return _repository_info(response) except Exception as e: raise Exception(f"Failed to get repository info: {e}") @@ -231,7 +255,7 @@ class BitBucketClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final[BitBucketBranchListing] = response.json() + data: Final = _branch_listing(response) return data.get("values", []) except Exception as e: raise Exception(f"Failed to get branches: {e}") diff --git a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py index fc885b7ee2c..5296f4b710b 100644 --- a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py +++ b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py @@ -4,12 +4,13 @@ Fetches .prompt files from BitBucket repositories and provides team-based access """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment from litellm.integrations.custom_prompt_management import CustomPromptManagement +from litellm.litellm_core_utils.core_helpers import require_attributes if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -26,6 +27,10 @@ from litellm.types.utils import StandardCallbackDynamicParams from .bitbucket_client import BitBucketClient +class _SupportsGetWithDefault(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + class BitBucketPromptTemplate: """ Represents a prompt template loaded from BitBucket. @@ -35,7 +40,7 @@ class BitBucketPromptTemplate: self, template_id: str, content: str, - metadata: dict[str, Any], + metadata: Mapping[str, object], model: str | None = None, ): self.template_id = template_id @@ -44,7 +49,9 @@ class BitBucketPromptTemplate: self.model = model or metadata.get("model") self.temperature = metadata.get("temperature") self.max_tokens = metadata.get("max_tokens") - self.input_schema = metadata.get("input", {}).get("schema", {}) + self.input_schema = require_attributes(metadata.get("input", {}), _SupportsGetWithDefault, "get").get( + "schema", {} + ) self.optional_params = {k: v for k, v in metadata.items() if k not in ["model", "input", "content"]} def __repr__(self): diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index c8453456052..cd1250b17f1 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -7,10 +7,10 @@ so `response_id` is always the raw provider response id (cache-hit suffix stripp import json import re -from collections.abc import Iterator, Mapping +from collections.abc import Iterable, Iterator, Mapping from math import isfinite from types import MappingProxyType -from typing import Any, Final +from typing import Any, Final, Protocol from pydantic import JsonValue, TypeAdapter, ValidationError @@ -19,6 +19,7 @@ from litellm._logging import verbose_logger from litellm.integrations.clickhouse.clickhouse_batch_logger import ClickHouseBatchLogger from litellm.integrations.clickhouse.context import is_lens_analysis from litellm.integrations.clickhouse.schema import SPEND_LOGS_TABLE +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload @@ -40,6 +41,14 @@ _INTERNAL_METADATA_KEYS: Final = frozenset( ) +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +class _SupportsItems(Protocol): + def items(self) -> Iterable[tuple[object, object]]: ... + + def _request_metadata_fields(value: object) -> Iterator[tuple[str, JsonValue]]: if value is None: return @@ -85,22 +94,22 @@ def _json_mapping(value: Mapping[str, object]) -> str: return _json(dict(value)) -def _find_traceparent(metadata: Mapping[str, Any], kwargs: Mapping[str, Any]) -> tuple[str, str]: +def _find_traceparent(metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> tuple[str, str]: custom_headers = metadata.get("requester_custom_headers") or MappingProxyType({}) - proxy_request = (kwargs.get("litellm_params") or MappingProxyType({})).get( + proxy_request = require_attributes(kwargs.get("litellm_params") or MappingProxyType({}), _SupportsGet, "get").get( "proxy_server_request" ) or MappingProxyType({}) - request_headers = proxy_request.get("headers") or MappingProxyType({}) + request_headers = require_attributes(proxy_request, _SupportsGet, "get").get("headers") or MappingProxyType({}) for headers in (custom_headers, request_headers): - for name, value in headers.items(): + for name, value in require_attributes(headers, _SupportsItems, "items").items(): if str(name).lower() == "traceparent": return parse_traceparent(value) return "", "" -def _cache_tokens(usage: Mapping[str, Any]) -> tuple[int, int]: +def _cache_tokens(usage: Mapping[str, object]) -> tuple[int, int]: """(cache_read, cache_write) from a Usage dict: OpenAI prompt_tokens_details first, Anthropic fields as fallback.""" - details = usage.get("prompt_tokens_details") or MappingProxyType({}) + details = require_attributes(usage.get("prompt_tokens_details") or MappingProxyType({}), _SupportsGet, "get") cache_read = _int(details.get("cached_tokens")) or _int(usage.get("cache_read_input_tokens")) cache_write = ( _int(details.get("cache_write_tokens")) @@ -116,10 +125,17 @@ def _request_tags(value: object) -> list[str]: return [str(tag) for tag in value] -def _session_id(payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, Any]) -> str: +def _session_id(payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, object]) -> str: """Mirrors proxy `_get_session_id_for_spend_log`: explicit session id, else the payload trace id.""" - request_metadata = (kwargs.get("litellm_params") or MappingProxyType({})).get("metadata") or MappingProxyType({}) - return str(payload.get("session_id") or request_metadata.get("session_id") or payload.get("trace_id") or "") + request_metadata = require_attributes( + kwargs.get("litellm_params") or MappingProxyType({}), _SupportsGet, "get" + ).get("metadata") or MappingProxyType({}) + return str( + payload.get("session_id") + or require_attributes(request_metadata, _SupportsGet, "get").get("session_id") + or payload.get("trace_id") + or "" + ) def _is_trace_ingest(payload: StandardLoggingPayload) -> bool: @@ -128,7 +144,7 @@ def _is_trace_ingest(payload: StandardLoggingPayload) -> bool: def spend_log_row_from_payload( - payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, Any] + payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, object] ) -> SpendLogRecord: metadata: Mapping[str, Any] = payload.get("metadata") or MappingProxyType({}) hidden_params: Mapping[str, Any] = payload.get("hidden_params") or MappingProxyType({}) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index b7a0ccda731..497d9d89ad4 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -7,7 +7,7 @@ from collections.abc import Iterable, Mapping, Sequence from datetime import datetime from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, cast, get_args import httpx from pydantic import TypeAdapter @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, redact_nested_match_and_regex_keys, + require_attributes, ) from litellm.secret_managers.main import str_to_bool from litellm.types.guardrails import ( @@ -167,18 +168,22 @@ def _strict_guardrail_modes_enabled() -> bool: return True if parsed is None else parsed -def get_session_id_from_request_data(request_data: dict[str, Any]) -> str | None: +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +def get_session_id_from_request_data(request_data: Mapping[str, object]) -> str | None: """Extract session_id from request data (litellm_session_id or metadata).""" session_id = request_data.get("litellm_session_id") if session_id: return str(session_id) - metadata: Final = request_data.get("metadata") or {} + metadata: Final = require_attributes(request_data.get("metadata") or {}, _SupportsGet, "get") session_id = metadata.get("session_id") if session_id: return str(session_id) - litellm_metadata: Final = request_data.get("litellm_metadata") or {} + litellm_metadata: Final = require_attributes(request_data.get("litellm_metadata") or {}, _SupportsGet, "get") session_id = litellm_metadata.get("session_id") if session_id: return str(session_id) diff --git a/litellm/integrations/email_alerting.py b/litellm/integrations/email_alerting.py index 27ee6bb60b4..d7228628d35 100644 --- a/litellm/integrations/email_alerting.py +++ b/litellm/integrations/email_alerting.py @@ -3,17 +3,29 @@ Functions for sending Email Alerts """ import os -from typing import Final +from collections.abc import Awaitable, Mapping, Sequence +from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.repositories.team_repository import TeamRepository from litellm.types.integrations.slack_alerting import WebhookEvent +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + # we use this for the email header, please send a test email if you change this. verify it looks good on email LITELLM_LOGO_URL: Final = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png" LITELLM_SUPPORT_CONTACT: Final = "support@berri.ai" +def _query_raw_rows( + client: "PrismaClient", query: str, *args: object +) -> Awaitable[Sequence[Mapping[str, object]] | None]: + from litellm.proxy.utils import query_raw_rows + + return query_raw_rows(client.db, query, *args) + + async def get_all_team_member_emails(team_id: str | None = None) -> list: verbose_logger.debug("Email Alerting: Getting all team members for team_id=%s", team_id) if team_id is None: @@ -51,17 +63,18 @@ async def get_all_team_member_emails(team_id: str | None = None) -> list: WHERE user_id = ANY($1::TEXT[]); """ - _result: Final = await prisma_client.db.query_raw(sql_query, _team_member_user_ids) + _result: Final = await _query_raw_rows(prisma_client, sql_query, _team_member_user_ids) verbose_logger.debug("Email Alerting: Got all Emails for team, emails=%s", _result) if _result is None: return [] - emails: Final = [] - for user in _result: - if user and isinstance(user, dict) and user.get("user_email", None) is not None: - emails.append(user.get("user_email")) + emails: Final = [ + user.get("user_email") + for user in _result + if user and isinstance(user, dict) and user.get("user_email", None) is not None + ] return emails diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index 891318f1c54..044439aeec2 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -123,7 +123,9 @@ class FocusLiteLLMDatabase: ORDER BY ordinal_position; """ try: - columns_response: Final = await client.db.query_raw(info_query) + from litellm.proxy.utils import query_raw_rows + + columns_response: Final = await query_raw_rows(client.db, info_query) return {"columns": columns_response, "table_name": "LiteLLM_DailyUserSpend"} except Exception as exc: raise RuntimeError(f"Error getting table info: {exc}") from exc diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index 56147cd720d..0f7dcdbcc81 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -15,6 +15,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import require_attributes, subscript from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, get_content_from_model_response, @@ -60,6 +61,10 @@ class _JsonResponse(Protocol): def json(self) -> object: ... +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + def _login_access_token(response: _GalileoLoginResponse) -> str: """Read the bearer token out of a Galileo login response body.""" return response.json()["access_token"] @@ -463,13 +468,13 @@ class GalileoObserve(CustomLogger): pass @staticmethod - def _build_prompt(kwargs: Mapping[str, Any]) -> dict[str, object]: - optional_params: Final[Mapping[str, object]] = kwargs.get("optional_params", {}) or {} + def _build_prompt(kwargs: Mapping[str, object]) -> dict[str, object]: + optional_params: Final = require_attributes(kwargs.get("optional_params", {}) or {}, _SupportsGet, "get") prompt: Final[dict[str, object]] = {"messages": kwargs.get("messages")} if optional_params.get("functions") is not None: - prompt["functions"] = optional_params["functions"] + prompt["functions"] = subscript(optional_params, "functions") if optional_params.get("tools") is not None: - prompt["tools"] = optional_params["tools"] + prompt["tools"] = subscript(optional_params, "tools") return prompt @staticmethod diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 7057e719764..99a8a5979b4 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -5,7 +5,7 @@ import os import time from collections.abc import Mapping from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -16,7 +16,7 @@ from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase from litellm.litellm_core_utils.cloud_storage_security import ( sanitize_cloud_object_component, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, require_attributes from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.gcs_bucket import * @@ -28,6 +28,10 @@ else: VertexBase = Any +class _SupportsGetWithDefault(Protocol): + def get(self, key: str, default: str, /) -> object: ... + + class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): def __init__(self, bucket_name: str | None = None) -> None: from litellm.proxy.proxy_server import premium_user @@ -298,7 +302,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): await self._send_individual_logs(items_to_process) def _get_object_name( - self, kwargs: Mapping[str, object], logging_payload: StandardLoggingPayload, response_obj: Any + self, kwargs: Mapping[str, object], logging_payload: StandardLoggingPayload, response_obj: object ) -> str: """ Get the object name to use for the current payload @@ -313,7 +317,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): return f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}" if logging_payload.get("error_str", None) is not None: return self._generate_failure_object_name(request_date_str=current_date) - raw_response_id: Final[object] = response_obj.get("id", "") + raw_response_id: Final[object] = require_attributes(response_obj, _SupportsGetWithDefault, "get").get("id", "") response_id: Final = raw_response_id if isinstance(raw_response_id, str) else "" return self._generate_success_object_name( request_date_str=current_date, @@ -360,7 +364,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): def _generate_success_object_name( self, request_date_str: str, - response_id: str, + response_id: object, ) -> str: return f"{request_date_str}/{response_id}" diff --git a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py index 2601f44e0bf..fda537c55c6 100644 --- a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py +++ b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py @@ -132,7 +132,7 @@ class GenericPromptManager(CustomPromptManagement): async def async_fetch_prompt_from_api( self, prompt_id: str | None, prompt_spec: PromptSpec | None - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Fetch a prompt from the API asynchronously. """ diff --git a/litellm/integrations/humanloop.py b/litellm/integrations/humanloop.py index 093a68c9284..394f087c132 100644 --- a/litellm/integrations/humanloop.py +++ b/litellm/integrations/humanloop.py @@ -7,10 +7,12 @@ https://humanloop.com/ from typing import Final, cast import httpx +from pydantic import TypeAdapter from typing_extensions import TypedDict import litellm from litellm.caching import DualCache +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, decoded_json_object, subscript from litellm.llms.custom_httpx.http_handler import get_httpx_client from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -19,6 +21,8 @@ from litellm.types.utils import StandardCallbackDynamicParams from .custom_logger import CustomLogger +_DECODED_JSON: Final = TypeAdapter(object) + class PromptManagementClient(TypedDict): prompt_id: str @@ -79,18 +83,20 @@ class HumanLoopPromptManager(DualCache): raise Exception(f"Error getting prompt from Humanloop: {e.response.text}") json_response: Final = response.json() - template_message: Final = json_response["template"] - if isinstance(template_message, dict): - template_messages = [template_message] - elif isinstance(template_message, list): - template_messages = template_message + decoded_response: Final = _DECODED_JSON.validate_python(json_response) + template_message: Final = subscript(decoded_response, "template") + template_mapping: Final = as_str_mapping(template_message) + template_list: Final = as_object_list(template_message) + if template_mapping is not None: + template_messages = [template_mapping] + elif template_list is not None: + template_messages = template_list else: raise ValueError(f"Invalid template message type: {type(template_message)}") template_model: Final = json_response["model"] - optional_params: Final = {} - for k, v in json_response.items(): - if k in litellm.OPENAI_CHAT_COMPLETION_PARAMS: - optional_params[k] = v + optional_params: Final = { + k: v for k, v in decoded_json_object(decoded_response).items() if k in litellm.OPENAI_CHAT_COMPLETION_PARAMS + } return PromptManagementClient( prompt_id=humanloop_prompt_id, prompt_template=cast(list[AllMessageValues], template_messages), diff --git a/litellm/integrations/lago.py b/litellm/integrations/lago.py index ec4388bbb8d..4b79484ee84 100644 --- a/litellm/integrations/lago.py +++ b/litellm/integrations/lago.py @@ -3,7 +3,7 @@ import json import os -from typing import Final, Literal +from typing import Final, Literal, Protocol import httpx @@ -11,6 +11,7 @@ import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, @@ -18,6 +19,10 @@ from litellm.llms.custom_httpx.http_handler import ( ) +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + def get_utc_datetime(): import datetime as dt from datetime import datetime @@ -70,10 +75,11 @@ class LagoLogger(CustomLogger): if ( isinstance(response_obj, litellm.ModelResponse) or isinstance(response_obj, litellm.EmbeddingResponse) ) and hasattr(response_obj, "usage"): + response_usage: Final = require_attributes(response_obj["usage"], _SupportsGet, "get") usage = { - "prompt_tokens": response_obj["usage"].get("prompt_tokens", 0), - "completion_tokens": response_obj["usage"].get("completion_tokens", 0), - "total_tokens": response_obj["usage"].get("total_tokens"), + "prompt_tokens": response_usage.get("prompt_tokens", 0), + "completion_tokens": response_usage.get("completion_tokens", 0), + "total_tokens": response_usage.get("total_tokens"), } litellm_params: Final = kwargs.get("litellm_params", {}) or {} diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index 0117348da66..977fa816eed 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -7,6 +7,7 @@ from itertools import chain from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, runtime_checkable from opentelemetry import trace +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger @@ -33,6 +34,8 @@ from litellm.types.integrations.langfuse_otel import ( ) from litellm.types.utils import StandardCallbackDynamicParams +_DECODED_JSON: Final = TypeAdapter(object) + if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -643,7 +646,7 @@ def _tool_call_arguments(arguments: object) -> object: if not isinstance(arguments, str): return arguments try: - return json.loads(arguments) + return _DECODED_JSON.validate_python(json.loads(arguments)) except json.JSONDecodeError: return {} diff --git a/litellm/integrations/langfuse/langfuse_otel_attributes.py b/litellm/integrations/langfuse/langfuse_otel_attributes.py index 2873dba93ed..28530076c57 100644 --- a/litellm/integrations/langfuse/langfuse_otel_attributes.py +++ b/litellm/integrations/langfuse/langfuse_otel_attributes.py @@ -7,7 +7,7 @@ Relevant Issue: https://github.com/BerriAI/litellm/issues/13764 import json from collections.abc import Mapping from datetime import datetime -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Protocol from pydantic import BaseModel from typing_extensions import override @@ -16,6 +16,7 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor BaseLLMObsOTELAttributes, safe_set_attribute, ) +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.utils import ( EmbeddingResponse, @@ -30,6 +31,10 @@ if TYPE_CHECKING: from opentelemetry.trace import Span +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + def serialize_langfuse_datetime(value: datetime) -> str: return json.dumps(value.astimezone().isoformat()) @@ -87,11 +92,11 @@ def get_output_content_by_type( class LangfuseLLMObsOTELAttributes(BaseLLMObsOTELAttributes): @staticmethod @override - def set_messages(span: "Span", kwargs: dict[str, Any]): + def set_messages(span: "Span", kwargs: Mapping[str, object]): prompt: Final = {"messages": kwargs.get("messages")} optional_params: Final = kwargs.get("optional_params", {}) - functions: Final = optional_params.get("functions") - tools: Final = optional_params.get("tools") + functions: Final = require_attributes(optional_params, _SupportsGet, "get").get("functions") + tools: Final = require_attributes(optional_params, _SupportsGet, "get").get("tools") if functions is not None: prompt["functions"] = functions if tools is not None: diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py index 8731e96440f..9e5b1109a77 100644 --- a/litellm/integrations/mlflow.py +++ b/litellm/integrations/mlflow.py @@ -1,9 +1,11 @@ import json import threading -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import attribute_or_none, iter_items, require_attributes if TYPE_CHECKING: from litellm.types.utils import StandardLoggingPayload @@ -11,6 +13,25 @@ else: StandardLoggingPayload = Any +class _ChoiceWithMessage(Protocol): + @property + def message(self) -> object: ... + + +class _ExcludeNoneDumpable(Protocol): + def model_dump(self, *, exclude_none: bool) -> object: ... + + +def _dumped_choice_message(choice: object) -> object: + message: Final = require_attributes(choice, _ChoiceWithMessage, "message").message + return require_attributes(message, _ExcludeNoneDumpable, "model_dump").model_dump(exclude_none=True) + + +def _output_messages(response_obj: object) -> Sequence[object]: + choices: Final = attribute_or_none(response_obj, "choices") if hasattr(response_obj, "choices") else () + return [_dumped_choice_message(c) for c in iter_items(choices)] + + class MlflowLogger(CustomLogger): def __init__(self): from mlflow.tracking import MlflowClient @@ -62,7 +83,7 @@ class MlflowLogger(CustomLogger): inputs: Final = self._construct_input(kwargs) input_messages: Final = inputs.get("messages", []) - output_messages: Final = [c.message.model_dump(exclude_none=True) for c in getattr(response_obj, "choices", [])] + output_messages: Final = _output_messages(response_obj) if messages := [*input_messages, *output_messages]: set_span_chat_messages(span, messages) if tools := inputs.get("tools"): diff --git a/litellm/integrations/openmeter.py b/litellm/integrations/openmeter.py index 918e849a086..9668fe8fa89 100644 --- a/litellm/integrations/openmeter.py +++ b/litellm/integrations/openmeter.py @@ -4,13 +4,13 @@ import json import os from collections.abc import Mapping -from typing import Final +from typing import Final, Protocol import httpx import litellm from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, require_attributes from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, @@ -19,6 +19,10 @@ from litellm.llms.custom_httpx.http_handler import ( ) +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + def get_utc_datetime(): import datetime as dt from datetime import datetime @@ -60,10 +64,11 @@ class OpenMeterLogger(CustomLogger): if ( isinstance(response_obj, litellm.ModelResponse) or isinstance(response_obj, litellm.EmbeddingResponse) ) and hasattr(response_obj, "usage"): + response_usage: Final = require_attributes(response_obj["usage"], _SupportsGet, "get") usage = { - "prompt_tokens": response_obj["usage"].get("prompt_tokens", 0), - "completion_tokens": response_obj["usage"].get("completion_tokens", 0), - "total_tokens": response_obj["usage"].get("total_tokens"), + "prompt_tokens": response_usage.get("prompt_tokens", 0), + "completion_tokens": response_usage.get("completion_tokens", 0), + "total_tokens": response_usage.get("total_tokens"), } # OPENMETER_TRUST_REQUEST_USER (default "true"): when set to "false", diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 39c485cb31f..5cd201a2f47 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -6,7 +6,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -35,6 +35,7 @@ from litellm.integrations.otel.model.metadata import flatten_metadata from litellm.integrations.otel.model.semconv import LiteLLM, Metric from litellm.integrations.otel.plumbing.otlp_tls import resolve_otlp_http_tls from litellm.integrations.otel.routing import routing_decision_attributes +from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.secret_redaction import redact_string @@ -47,7 +48,6 @@ from litellm.secret_managers.main import get_secret_bool, str_to_bool from litellm.types.integrations.otel_span_attributes import SpanAttributes from litellm.types.services import ServiceLoggerPayload from litellm.types.utils import ( - ChatCompletionMessageToolCall, CostBreakdown, Function, LLMResponseTypes, @@ -130,6 +130,22 @@ def _text_as_stored_on_span(span: Span, key: str, text: str) -> str: _MAX_DYNAMIC_TRACER_PROVIDERS: Final = 256 +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +def _response_field(response_obj: object, key: str) -> object: + if not response_obj: + return None + return require_attributes(response_obj, _SupportsGet, "get").get(key) + + +def _billed_usage(response_obj: object, call_type: str | None, litellm_params: Mapping[str, object] | None) -> object: + if not response_obj or is_unbilled_non_inference_call_from_params(call_type, litellm_params, response_obj): + return None + return _response_field(response_obj, "usage") + + class _CacheHitView(TypedDict, total=False): cache_hit: ReadOnly[object] saved_cache_cost: ReadOnly[object] @@ -2407,17 +2423,17 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): @staticmethod def _tool_calls_kv_pair( - tool_calls: list[ChatCompletionMessageToolCall], + tool_calls: Iterable[object], ) -> dict[str, object]: kv_pairs: Final[dict[str, object]] = {} for idx, tool_call in enumerate(tool_calls): - _function = tool_call.get("function") + _function: object = require_attributes(tool_call, _SupportsGet, "get").get("function") if not _function: continue keys = Function.__annotations__.keys() for key in keys: - _value = _function.get(key) + _value = require_attributes(_function, _SupportsGet, "get").get(key) if _value: kv_pairs[f"{SpanAttributes.LLM_COMPLETIONS.value}.{idx}.function_call.{key}"] = _value @@ -2584,7 +2600,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # but Embeddings and Image-gen responses do not. Fall back to # the litellm call ID so every call type can be correlated # across LiteLLM UI, Phoenix traces, and provider logs (Issue #8). - response_id = (response_obj.get("id") if response_obj else None) or standard_logging_payload.get("id") + response_id = _response_field(response_obj, "id") or standard_logging_payload.get("id") if response_id: self.safe_set_attribute( span=span, @@ -2601,42 +2617,37 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) # The model used to generate the response. - if response_obj and response_obj.get("model"): + response_model: Final = _response_field(response_obj, "model") + if response_model: self.safe_set_attribute( span=span, key=SpanAttributes.LLM_RESPONSE_MODEL.value, - value=response_obj.get("model"), + value=response_model, ) self._set_service_tier_attributes(span=span, standard_logging_payload=standard_logging_payload) - usage: Final = ( - response_obj.get("usage") - if response_obj - and not is_unbilled_non_inference_call_from_params( - kwargs.get("call_type"), litellm_params, response_obj - ) - else None - ) + usage: Final = _billed_usage(response_obj, kwargs.get("call_type"), litellm_params) if usage: + usage_fields: Final = require_attributes(usage, _SupportsGet, "get") self.safe_set_attribute( span=span, key=SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS.value, - value=usage.get("total_tokens"), + value=usage_fields.get("total_tokens"), ) # The number of tokens used in the LLM response (completion). self.safe_set_attribute( span=span, key=SpanAttributes.GEN_AI_USAGE_OUTPUT_TOKENS.value, - value=usage.get("completion_tokens"), + value=usage_fields.get("completion_tokens"), ) # The number of tokens used in the LLM prompt. self.safe_set_attribute( span=span, key=SpanAttributes.GEN_AI_USAGE_INPUT_TOKENS.value, - value=usage.get("prompt_tokens"), + value=usage_fields.get("prompt_tokens"), ) ######################################################################## @@ -2712,7 +2723,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ########## LLM Response Attributes ########## ############################################# if response_obj is not None: - if response_obj.get("choices"): + response_view: Final = require_attributes(response_obj, _SupportsGet, "get") + if response_view.get("choices"): transformed_choices: Final = self._transform_choices_to_otel_semantic_conventions( response_obj.get("choices") ) @@ -2722,10 +2734,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): value=safe_dumps(transformed_choices), ) - finish_reasons: Final = [] - for idx, choice in enumerate(response_obj.get("choices")): - if choice.get("finish_reason"): - finish_reasons.append(choice.get("finish_reason")) + choice_views: Final = [ + require_attributes(choice, _SupportsGet, "get") + for choice in iter_items(response_view.get("choices")) + ] + finish_reasons: Final = [ + choice.get("finish_reason") for choice in choice_views if choice.get("finish_reason") + ] if finish_reasons: self.safe_set_attribute( @@ -2734,11 +2749,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): value=safe_dumps(finish_reasons), ) - for idx, choice in enumerate(response_obj.get("choices")): + for choice in choice_views: if choice.get("finish_reason"): - message = choice.get("message") - if tool_calls := tool_calls_with_server_tool_calls(message): - kv_pairs = OpenTelemetry._tool_calls_kv_pair(list(tool_calls)) + tool_calls = tool_calls_with_server_tool_calls(choice.get("message")) + if tool_calls: + kv_pairs = OpenTelemetry._tool_calls_kv_pair(tool_calls) for key, value in kv_pairs.items(): self.safe_set_attribute( span=span, @@ -2746,13 +2761,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): value=value, ) - elif response_obj.get("output"): + elif response_view.get("output"): # Responses API: ResponsesAPIResponse has an "output" # list instead of "choices". Each item with # type="message" contains a "content" list of # OutputText objects (type="output_text"). - output_items: Final = response_obj.get("output") - output_messages: Final = self._transform_responses_api_output_to_otel(output_items) + output_items: Final = response_view.get("output") + output_messages: Final = self._transform_responses_api_output_to_otel(iter_items(output_items)) if output_messages: self.safe_set_attribute( span=span, @@ -2766,7 +2781,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # ChatCompletionMessageToolCall format expected by # _tool_calls_kv_pair. tool_calls = [] - for out_item in output_items: + for out_item in iter_items(output_items): item_d = self._to_dict(out_item) if item_d and item_d.get("type") == "function_call": tool_calls.append( @@ -2787,7 +2802,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) # Extract finish reason from ResponsesAPIResponse.status - status: Final = response_obj.get("status") + status: Final = response_view.get("status") if status: self.safe_set_attribute( span=span, @@ -2907,7 +2922,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): return obj.model_dump() return None - def _transform_responses_api_output_to_otel(self, output: list) -> list[dict]: + def _transform_responses_api_output_to_otel(self, output: Iterable[object]) -> list[dict]: """ Transform Responses API output items into OTEL GenAI 1.38 format. diff --git a/litellm/integrations/opentelemetry_utils/base_otel_llm_obs_attributes.py b/litellm/integrations/opentelemetry_utils/base_otel_llm_obs_attributes.py index 97a970c8966..350c61d22f6 100644 --- a/litellm/integrations/opentelemetry_utils/base_otel_llm_obs_attributes.py +++ b/litellm/integrations/opentelemetry_utils/base_otel_llm_obs_attributes.py @@ -1,4 +1,5 @@ from abc import ABC +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: @@ -7,7 +8,7 @@ if TYPE_CHECKING: class BaseLLMObsOTELAttributes(ABC): @staticmethod - def set_messages(span: "Span", kwargs: dict[str, Any]): + def set_messages(span: "Span", kwargs: Mapping[str, object]): pass @staticmethod diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index 1197ae055d0..27ee20362a2 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -11,6 +11,7 @@ identical metrics. The attribute cardinality filter is reused from v1 by import from collections.abc import Mapping from dataclasses import dataclass from datetime import datetime +from types import MappingProxyType from typing import Any, Final, Literal, Protocol, TypeAlias from opentelemetry.metrics import Histogram, Meter @@ -34,6 +35,7 @@ from litellm.integrations.otel.model.semconv import ( resolve_provider, ) from litellm.integrations.otel.model.utils import to_seconds +from litellm.litellm_core_utils.core_helpers import attribute_or_none, require_attributes from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -177,7 +179,14 @@ class _MetricKwargs(TypedDict, total=False): end_time: ReadOnly[datetime | float | str | None] -def resolve_error_type(kwargs: Mapping[str, Any]) -> str: +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def resolve_error_type(kwargs: Mapping[str, object]) -> str: """The ``error.type`` value for a failed request. Bounded by construction: the mapped provider exception's class name (the same @@ -188,8 +197,13 @@ def resolve_error_type(kwargs: Mapping[str, Any]) -> str: span and its exception event, where high cardinality is free. """ std_log: Final = kwargs.get("standard_logging_object") - info: Final = getattr(std_log, "error_information", None) or (std_log or {}).get("error_information") or {} - error_class: Final = info.get("error_class") or info.get("error_code") + info: Final = ( + attribute_or_none(std_log, "error_information") + or require_attributes(std_log or _NO_FIELDS, _SupportsGet, "get").get("error_information") + or _NO_FIELDS + ) + info_fields: Final = require_attributes(info, _SupportsGet, "get") + error_class: Final = info_fields.get("error_class") or info_fields.get("error_code") if error_class: return str(error_class) exception: Final = kwargs.get("exception") diff --git a/litellm/integrations/prompt_layer.py b/litellm/integrations/prompt_layer.py index 1efd7560a1d..5bdc56fedad 100644 --- a/litellm/integrations/prompt_layer.py +++ b/litellm/integrations/prompt_layer.py @@ -4,9 +4,12 @@ import os import traceback from typing import Final -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm +from litellm.litellm_core_utils.core_helpers import decoded_json_object + +_DECODED_JSON: Final = TypeAdapter(object) class PromptLayerLogger: @@ -61,8 +64,8 @@ class PromptLayerLogger: }, ) - response_json: Final = request_response.json() - if not request_response.json().get("success", False): + response_json: Final = decoded_json_object(_DECODED_JSON.validate_python(request_response.json())) + if not response_json.get("success", False): raise Exception("Promptlayer did not successfully log the response!") print_verbose(f"Prompt Layer Logging: success - final response object: {request_response.text}") diff --git a/litellm/integrations/weave/weave_otel.py b/litellm/integrations/weave/weave_otel.py index 2bff3682de7..68a9dc3e162 100644 --- a/litellm/integrations/weave/weave_otel.py +++ b/litellm/integrations/weave/weave_otel.py @@ -4,7 +4,8 @@ import base64 import json import os from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Protocol from opentelemetry.trace import Status, StatusCode from typing_extensions import override @@ -19,6 +20,7 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor BaseLLMObsOTELAttributes, safe_set_attribute, ) +from litellm.litellm_core_utils.core_helpers import has_attributes, require_attributes from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.types.integrations.weave_otel import WeaveOtelConfig, WeaveSpanAttributes from litellm.types.utils import StandardCallbackDynamicParams @@ -34,6 +36,21 @@ WEAVE_BASE_URL: Final = "https://trace.wandb.ai" WEAVE_OTEL_ENDPOINT: Final = "/otel/v1/traces" +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +class _SupportsReplace(Protocol): + def replace(self, old: str, new: str, /) -> object: ... + + +class _ModelDumpable(Protocol): + def model_dump(self) -> object: ... + + class WeaveLLMObsOTELAttributes(BaseLLMObsOTELAttributes): """ Weave-specific LLM observability OTEL attributes. @@ -44,15 +61,15 @@ class WeaveLLMObsOTELAttributes(BaseLLMObsOTELAttributes): @staticmethod @override - def set_messages(span: Span, kwargs: dict[str, Any]): + def set_messages(span: Span, kwargs: Mapping[str, object]): """Set input messages as span attributes using OpenInference conventions.""" messages: Final = kwargs.get("messages") or [] - optional_params: Final = kwargs.get("optional_params") or {} + optional_params: Final = kwargs.get("optional_params") or _NO_FIELDS prompt: Final = {"messages": messages} - functions: Final = optional_params.get("functions") - tools: Final = optional_params.get("tools") + functions: Final = require_attributes(optional_params, _SupportsGet, "get").get("functions") + tools: Final = require_attributes(optional_params, _SupportsGet, "get").get("tools") if functions is not None: prompt["functions"] = functions if tools is not None: @@ -60,7 +77,7 @@ class WeaveLLMObsOTELAttributes(BaseLLMObsOTELAttributes): safe_set_attribute(span, OpenInferenceSpanAttributes.INPUT_VALUE, json.dumps(prompt)) -def _set_weave_specific_attributes(span: Span, kwargs: Mapping[str, Any], response_obj: Any): +def _set_weave_specific_attributes(span: Span, kwargs: Mapping[str, object], response_obj: object): """ Sets Weave-specific metadata attributes onto the OTEL span. @@ -69,25 +86,27 @@ def _set_weave_specific_attributes(span: Span, kwargs: Mapping[str, Any], respon """ # Extract all needed data upfront - litellm_params: Final = kwargs.get("litellm_params") or {} + litellm_params: Final = kwargs.get("litellm_params") or _NO_FIELDS # optional_params = kwargs.get("optional_params") or {} - metadata: Final = kwargs.get("metadata") or {} + metadata: Final = kwargs.get("metadata") or _NO_FIELDS model: Final = kwargs.get("model") or "" - custom_llm_provider: Final = litellm_params.get("custom_llm_provider") or "" + custom_llm_provider: Final = ( + require_attributes(litellm_params, _SupportsGet, "get").get("custom_llm_provider") or "" + ) # Weave supports a custom display name and will default to the model name if not provided. - display_name = metadata.get("display_name") + display_name = require_attributes(metadata, _SupportsGet, "get").get("display_name") if not display_name and model: if custom_llm_provider: display_name = f"{custom_llm_provider}/{model}" else: display_name = model if display_name: - display_name = display_name.replace("/", "__") + display_name = require_attributes(display_name, _SupportsReplace, "replace").replace("/", "__") safe_set_attribute(span, WeaveSpanAttributes.DISPLAY_NAME.value, display_name) # Weave threads are OpenInference sessions. - if (session_id := metadata.get("session_id")) is not None: + if (session_id := require_attributes(metadata, _SupportsGet, "get").get("session_id")) is not None: if isinstance(session_id, (list, dict)): session_id = safe_dumps(session_id) safe_set_attribute(span, WeaveSpanAttributes.THREAD_ID.value, session_id) @@ -97,7 +116,7 @@ def _set_weave_specific_attributes(span: Span, kwargs: Mapping[str, Any], respon # but we override them here to better match Weave's expectations if response_obj: output_dict = None - if hasattr(response_obj, "model_dump"): + if has_attributes(response_obj, _ModelDumpable, "model_dump"): output_dict = response_obj.model_dump() elif hasattr(response_obj, "get"): output_dict = response_obj diff --git a/litellm/integrations/websearch_interception/tools.py b/litellm/integrations/websearch_interception/tools.py index 36b02c48dc5..227f9f33baa 100644 --- a/litellm/integrations/websearch_interception/tools.py +++ b/litellm/integrations/websearch_interception/tools.py @@ -7,9 +7,19 @@ to this format for consistent interception and execution. """ from collections.abc import Mapping -from typing import Any, Final +from typing import Final, Protocol from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME +from litellm.litellm_core_utils.core_helpers import require_attributes + + +class _SupportsGetWithDefault(Protocol): + def get(self, key: str, default: str, /) -> object: ... + + +class _SupportsStartswith(Protocol): + def startswith(self, prefix: str, /) -> object: ... + _WEB_SEARCH_TOOL_DESCRIPTION: Final = ( "Search the web for information. Use this when you need current " @@ -164,7 +174,7 @@ def is_chat_completion_request_web_search_tool(tool: Mapping[str, object]) -> bo return tool.get("type") == "web_search" or is_web_search_tool(dict(tool)) -def is_web_search_tool_chat_completion(tool: dict[str, Any]) -> bool: +def is_web_search_tool_chat_completion(tool: Mapping[str, object]) -> bool: """ Check if a tool is a web search tool for Chat Completions API (strict check). @@ -197,7 +207,7 @@ def is_web_search_tool_chat_completion(tool: dict[str, Any]) -> bool: # Check for OpenAI format: {"type": "function", "function": {"name": "litellm_web_search"}} if tool_type == "function" and "function" in tool: function_def: Final = tool.get("function", {}) - function_name: Final = function_def.get("name", "") + function_name: Final = require_attributes(function_def, _SupportsGetWithDefault, "get").get("name", "") if function_name == LITELLM_WEB_SEARCH_TOOL_NAME: return True @@ -229,7 +239,7 @@ def is_anthropic_native_web_search_tool(tool: Mapping[str, object]) -> bool: return tool_type.startswith("web_search_") and tool_type != "function" -def is_web_search_tool(tool: dict[str, Any]) -> bool: +def is_web_search_tool(tool: Mapping[str, object]) -> bool: """ Check if a tool is a web search tool (native or LiteLLM standard). @@ -281,7 +291,7 @@ def is_web_search_tool(tool: dict[str, Any]) -> bool: # Check for OpenAI format: {"type": "function", "function": {"name": "..."}} if tool_type == "function" and "function" in tool: function_def: Final = tool.get("function", {}) - function_name: Final = function_def.get("name", "") + function_name: Final = require_attributes(function_def, _SupportsGetWithDefault, "get").get("name", "") if function_name == LITELLM_WEB_SEARCH_TOOL_NAME: return True @@ -290,7 +300,7 @@ def is_web_search_tool(tool: dict[str, Any]) -> bool: return True # Check for native Anthropic web_search_* types - if tool_type.startswith("web_search_"): + if require_attributes(tool_type, _SupportsStartswith, "startswith").startswith("web_search_"): return True # Check for Claude Code's web_search with a type field diff --git a/litellm/integrations/weights_biases.py b/litellm/integrations/weights_biases.py index 9bc070a1f9a..01670501f22 100644 --- a/litellm/integrations/weights_biases.py +++ b/litellm/integrations/weights_biases.py @@ -1,11 +1,13 @@ from typing import Final +from litellm.litellm_core_utils.core_helpers import iter_items, subscript + imported_openAIResponse = True try: import io import logging from collections.abc import Mapping - from typing import Any, Literal, Protocol, TypeVar + from typing import Literal, Protocol, TypeVar from wandb.sdk.data_types import trace_tree @@ -25,7 +27,7 @@ try: class OpenAIRequestResponseResolver: def __call__( self, - request: dict[str, Any], + request: Mapping[str, object], response: OpenAIResponse, time_elapsed: float, ) -> trace_tree.WBTraceTree | None: @@ -110,14 +112,14 @@ try: def _resolve_chat_completion( self, - request: dict[str, Any], + request: Mapping[str, object], response: OpenAIResponse, time_elapsed: float, ) -> trace_tree.WBTraceTree: """Resolves the request and response objects for `openai.Completion`.""" prompt: Final = io.StringIO() - for message in request["messages"]: - prompt.write(f"\n\n**{message['role']}**: {message['content']}\n") + for message in iter_items(request["messages"]): + prompt.write(f"\n\n**{subscript(message, 'role')}**: {subscript(message, 'content')}\n") request_str: Final = prompt.getvalue() choices: Final = [ diff --git a/litellm/interactions/litellm_responses_transformation/transformation.py b/litellm/interactions/litellm_responses_transformation/transformation.py index 2e9b62dcfe2..6bc7413dfa4 100644 --- a/litellm/interactions/litellm_responses_transformation/transformation.py +++ b/litellm/interactions/litellm_responses_transformation/transformation.py @@ -8,11 +8,11 @@ This module handles transforming between: from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final, cast +from typing import Final, Protocol, cast from pydantic import BaseModel -from litellm.litellm_core_utils.core_helpers import billed_usage_cost +from litellm.litellm_core_utils.core_helpers import attribute_or_none, billed_usage_cost, has_attributes from litellm.types.interactions import ( InteractionInput, InteractionsAPIOptionalRequestParams, @@ -27,6 +27,10 @@ from litellm.types.llms.openai import ( _STEP_TYPE_ROLES: Final = MappingProxyType({"user_input": "user", "model_output": "assistant"}) +class _IsoFormattable(Protocol): + def isoformat(self) -> object: ... + + class LiteLLMResponsesInteractionsConfig: """Configuration class for transforming between Interactions API and Responses API.""" @@ -232,12 +236,12 @@ class LiteLLMResponsesInteractionsConfig: ) # Convert created_at to ISO string - created_at: Final = getattr(responses_response, "created_at", None) + created_at: Final = attribute_or_none(responses_response, "created_at") if isinstance(created_at, int): from datetime import datetime created = datetime.fromtimestamp(created_at).isoformat() - elif created_at is not None and hasattr(created_at, "isoformat"): + elif created_at is not None and has_attributes(created_at, _IsoFormattable, "isoformat"): created = created_at.isoformat() else: created = None diff --git a/litellm/interactions/streaming_iterator.py b/litellm/interactions/streaming_iterator.py index c7faa3fc510..674c05a7545 100644 --- a/litellm/interactions/streaming_iterator.py +++ b/litellm/interactions/streaming_iterator.py @@ -7,8 +7,9 @@ from the Google Interactions API, similar to the responses API streaming iterato import asyncio import json +from collections.abc import Mapping from datetime import datetime -from typing import Any, Final +from typing import Final, Protocol import httpx from pydantic import ConfigDict, TypeAdapter @@ -16,7 +17,7 @@ from pydantic import ConfigDict, TypeAdapter from litellm._logging import verbose_logger from litellm.constants import STREAM_SSE_DONE_STRING from litellm.litellm_core_utils.asyncify import run_async_function -from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.core_helpers import process_response_headers, require_attributes from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.thread_pool_executor import executor @@ -32,6 +33,10 @@ _INTERACTIONS_STREAMING_CHUNK_ADAPTER: Final = TypeAdapter( ) +class _SupportsGet(Protocol): + def get(self, key: str, default: object, /) -> object: ... + + class BaseInteractionsAPIStreamingIterator: """ Base class for streaming iterators that process responses from the Interactions API. @@ -45,7 +50,7 @@ class BaseInteractionsAPIStreamingIterator: model: str | None, interactions_api_config: BaseInteractionsAPIConfig, logging_obj: LiteLLMLoggingObj, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, object] | None = None, custom_llm_provider: str | None = None, ): self.response = response @@ -65,9 +70,9 @@ class BaseInteractionsAPIStreamingIterator: model=model or "", optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), ) - _model_info: Final[dict] = litellm_metadata.get("model_info", {}) if litellm_metadata else {} + _model_info: Final[object] = litellm_metadata.get("model_info", {}) if litellm_metadata else {} self._hidden_params = { - "model_id": _model_info.get("id", None), + "model_id": require_attributes(_model_info, _SupportsGet, "get").get("id", None), "api_base": _api_base, } self._hidden_params["additional_headers"] = process_response_headers(self.response.headers or {}) @@ -135,7 +140,7 @@ class InteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator): model: str | None, interactions_api_config: BaseInteractionsAPIConfig, logging_obj: LiteLLMLoggingObj, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, object] | None = None, custom_llm_provider: str | None = None, ): super().__init__( @@ -202,7 +207,7 @@ class SyncInteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator) model: str | None, interactions_api_config: BaseInteractionsAPIConfig, logging_obj: LiteLLMLoggingObj, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, object] | None = None, custom_llm_provider: str | None = None, ): super().__init__( diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index 29b10ac3c83..808a424815e 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -313,7 +313,7 @@ def calculate_request_duration(file: FileTypes) -> float | None: file_content = content elif hasattr(content, "read") and not isinstance(content, (str, os.PathLike)): # File-like object in tuple - current_pos: Final = getattr(content, "tell", lambda: None)() + current_pos: Final = content.tell() if hasattr(content, "tell") else None # Seek to start to ensure we read the entire content if hasattr(content, "seek"): content.seek(0) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index b28d5a4ea21..f31d4205b82 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -4,6 +4,7 @@ import copy import logging import math import re +import sys from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable @@ -133,7 +134,7 @@ class _SupportsIter(Protocol): @runtime_checkable class _SupportsGetItem(Protocol): - def __getitem__(self, index: int, /) -> object: ... + def __getitem__(self, key: str | int, /) -> object: ... def attribute_or_none(source: object, name: str) -> object: @@ -166,6 +167,28 @@ def iter_items(value: object) -> Iterable[object]: raise TypeError(f"'{type(value).__name__}' object is not iterable") +@runtime_checkable +class _SupportsContains(Protocol): + def __contains__(self, key: object, /) -> bool: ... + + +_NOT_A_CONTAINER: Final = "is not a container or iterable" if sys.version_info >= (3, 14) else "is not iterable" + + +def json_contains(value: object, key: object) -> bool: + if isinstance(value, _SupportsContains): + return key in value + if isinstance(value, (_SupportsIter, _SupportsGetItem)): + return key in tuple(iter_items(value)) + raise TypeError(f"argument of type '{type(value).__name__}' {_NOT_A_CONTAINER}") + + +def subscript(value: object, key: str | int) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + def attribute_items(source: object, name: str) -> Iterable[object]: value: Final = attribute_or_none(source, name) return iter_items(value) if value else () diff --git a/litellm/litellm_core_utils/dd_tracing.py b/litellm/litellm_core_utils/dd_tracing.py index a4db3bcd5c2..e72c7d4980b 100644 --- a/litellm/litellm_core_utils/dd_tracing.py +++ b/litellm/litellm_core_utils/dd_tracing.py @@ -5,8 +5,9 @@ If the ddtrace package is not installed, the tracer will be a no-op. """ from contextlib import contextmanager -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.secret_managers.main import get_secret_bool if TYPE_CHECKING: @@ -15,6 +16,14 @@ else: DD_TRACER = Any +class _TagStrSpan(Protocol): + def set_tag_str(self, key: str, value: str, /) -> object: ... + + +class _TagSpan(Protocol): + def set_tag(self, key: str, value: str, /) -> object: ... + + class NullSpan: """A no-op span implementation.""" @@ -81,7 +90,7 @@ else: tracer = NullTracer() -def get_active_span() -> Any | None: +def get_active_span() -> object: """ Return the active Datadog span, checking current span first and then root span. """ @@ -115,10 +124,10 @@ def set_active_span_tag(tag_key: str, tag_value: str) -> bool: return False try: - if hasattr(span, "set_tag_str"): + if has_attributes(span, _TagStrSpan, "set_tag_str"): span.set_tag_str(tag_key, str(tag_value)) return True - if hasattr(span, "set_tag"): + if has_attributes(span, _TagSpan, "set_tag"): span.set_tag(tag_key, str(tag_value)) return True except Exception: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9b6ec8d28d6..fcc521cf6cf 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -19,7 +19,7 @@ from datetime import datetime as dt_object from enum import Enum from functools import lru_cache from types import MappingProxyType, TracebackType -from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, TypeVar, cast from httpx import Response, Timeout from openai import AsyncStream as OpenAIAsyncStream @@ -90,6 +90,7 @@ from litellm.litellm_core_utils.classifier_logging import ( is_classifier_call, ) from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, finite_cost, get_provider_response_headers_from_hidden_params, is_expected_client_error, @@ -99,6 +100,7 @@ from litellm.litellm_core_utils.core_helpers import ( proxy_stamped_used_client_oauth_token, reconstruct_model_name, reports_native_messages_cost, + require_attributes, set_response_cost_in_hidden_params, single_log_line, ) @@ -666,6 +668,31 @@ class UsageObservation(Enum): ESTIMATED = "estimated" +class _A2AUsageSlot(Protocol): + usage: Mapping[str, object] | None + + +class _A2AResponseCopier(Protocol): + model_copy: Callable[..., _A2AUsageSlot] + + +class _MetadataLookup(Protocol): + get: Callable[[str], object] + + +class _ParamsLookup(Protocol): + def get(self, key: str, default: object, /) -> object: ... + + +def _callback_control_params(litellm_params: object) -> dict[str, object]: + mapping: Final = as_str_mapping(litellm_params) + if isinstance(mapping, dict): + return mapping + if mapping is not None: + return dict(mapping) + return litellm_params.model_dump() if isinstance(litellm_params, BaseModel) else {} + + class Logging(LiteLLMLoggingBaseClass): global \ supabaseClient, \ @@ -929,7 +956,9 @@ class Logging(LiteLLMLoggingBaseClass): new_call.start_time = start_time new_call.litellm_params = litellm_params new_call.model_call_details = { - **{key: value for key, value in self.model_call_details.items() if not key.startswith("_accounting_")}, + **{ + key: value for key, value in self._model_call_details_view.items() if not key.startswith("_accounting_") + }, "litellm_call_id": litellm_call_id, "start_time": start_time, "first_api_call_start_time": None, @@ -2467,7 +2496,7 @@ class Logging(LiteLLMLoggingBaseClass): return self.response_cost_calculator(result=result, cache_hit=cache_hit) @staticmethod - def is_sync_litellm_request(litellm_params: Mapping[str, object]) -> bool: + def is_sync_litellm_request(litellm_params: _ParamsLookup) -> bool: """True for sync SDK entrypoints (``completion``), false for async (``acompletion``, etc.).""" return ( litellm_params.get(CallTypes.acompletion.value, False) is not True @@ -2918,11 +2947,12 @@ class Logging(LiteLLMLoggingBaseClass): """Deprecated alias of mark_logging_complete, kept for callers of the old name""" self.mark_logging_complete(event_type=event_type) - def should_run_callback(self, callback: litellm.CALLBACK_TYPES, litellm_params: dict, event_hook: str) -> bool: + def should_run_callback(self, callback: litellm.CALLBACK_TYPES, litellm_params: object, event_hook: str) -> bool: if litellm.global_disable_no_log_param: return True - if litellm_params.get("no-log", False) is True: + request_params: Final = require_attributes(litellm_params, _ParamsLookup, "get") + if request_params.get("no-log", False) is True: # proxy cost tracking cal backs should run if not (isinstance(callback, CustomLogger) and "_PROXY_" in callback.__class__.__name__): @@ -2932,7 +2962,7 @@ class Logging(LiteLLMLoggingBaseClass): # Check for dynamically disabled callbacks via headers if EnterpriseCallbackControls is not None and EnterpriseCallbackControls.is_callback_disabled_dynamically( callback=callback, - litellm_params=litellm_params, + litellm_params=_callback_control_params(litellm_params), standard_callback_dynamic_params=self.standard_callback_dynamic_params, ): verbose_logger.debug( @@ -3643,8 +3673,8 @@ class Logging(LiteLLMLoggingBaseClass): build_logging_payload=True, publish_logging_payload=not for_async_request, ) - litellm_params = self.model_call_details.get("litellm_params", {}) - is_sync_request: Final = self.is_sync_litellm_request(litellm_params) + litellm_params = self._model_call_details_view.get("litellm_params", {}) + is_sync_request: Final = self.is_sync_litellm_request(require_attributes(litellm_params, _ParamsLookup, "get")) try: ## BUILD COMPLETE STREAMED RESPONSE complete_streaming_response: ( @@ -3800,18 +3830,23 @@ class Logging(LiteLLMLoggingBaseClass): model = call_details_view["model"] messages = call_details_view["messages"] - optional_params = kwargs.get("optional_params", {}) - litellm_params = kwargs.get("litellm_params", {}) + optional_params = call_details_view.get("optional_params", {}) + litellm_params = call_details_view.get("litellm_params", {}) supabaseClient.log_event( model=model, messages=messages, - end_user=optional_params.get("user", "default"), + end_user=require_attributes(optional_params, _ParamsLookup, "get").get("user", "default"), response_obj=result, start_time=start_time, end_time=end_time, litellm_call_id=( current_call_id - if (current_call_id := litellm_params.get("litellm_call_id")) is not None + if ( + current_call_id := require_attributes(litellm_params, _ParamsLookup, "get").get( + "litellm_call_id", None + ) + ) + is not None else str(uuid.uuid4()) ), print_verbose=print_verbose, @@ -4369,7 +4404,7 @@ class Logging(LiteLLMLoggingBaseClass): for callback in delivery_callbacks: # check if callback can run for this request - litellm_params = self.model_call_details.get("litellm_params", {}) + litellm_params = self._model_call_details_view.get("litellm_params", {}) should_run = self.should_run_callback( callback=callback, litellm_params=litellm_params, @@ -4658,8 +4693,8 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug("Logging Details LiteLLM-Failure Call: %s", litellm.failure_callback) if not self.should_run_logging(event_type="sync_failure"): # prevent double logging return - litellm_params: Final = self.model_call_details.get("litellm_params", {}) - is_sync_request: Final = self.is_sync_litellm_request(litellm_params) + litellm_params: Final = self._model_call_details_view.get("litellm_params", {}) + is_sync_request: Final = self.is_sync_litellm_request(require_attributes(litellm_params, _ParamsLookup, "get")) try: start_time, end_time = self._failure_handler_helper_fn( @@ -4879,7 +4914,7 @@ class Logging(LiteLLMLoggingBaseClass): self.mark_logging_complete(event_type="async_failure") for callback in callbacks: try: - litellm_params = self.model_call_details.get("litellm_params", {}) + litellm_params = self._model_call_details_view.get("litellm_params", {}) should_run = self.should_run_callback( callback=callback, litellm_params=litellm_params, @@ -5441,7 +5476,7 @@ class Logging(LiteLLMLoggingBaseClass): httpx_response: Final[Response | None] = self.model_call_details.get("httpx_response", None) if httpx_response is None: raise ValueError("Google GenAI Generate Content: httpx_response is None") - dict_result: Final = httpx_response.json() + dict_result: Final[GenerateContentResponseBody] = httpx_response.json() result = litellm.VertexGeminiConfig().transform_google_generate_content_to_openai_model_response( completion_response=dict_result, model_response=litellm.ModelResponse(), @@ -5454,7 +5489,7 @@ class Logging(LiteLLMLoggingBaseClass): ) return result - def _handle_a2a_response_logging(self, result: Any) -> object: + def _handle_a2a_response_logging(self, result: object) -> object: """ Handles logging for A2A (Agent-to-Agent) responses. @@ -5473,7 +5508,7 @@ class Logging(LiteLLMLoggingBaseClass): return result # Deep copy result and add usage - result_copy: Final = result.model_copy(deep=True) + result_copy: Final = require_attributes(result, _A2AResponseCopier, "model_copy").model_copy(deep=True) result_copy.usage = usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) return result_copy @@ -7475,7 +7510,7 @@ class StandardLoggingPayloadSetup: @staticmethod def get_standard_logging_payload_trace_id( logging_obj: Logging, - litellm_params: Mapping[str, Any], + litellm_params: Mapping[str, object], ) -> str: """ Returns the `litellm_trace_id` for this request @@ -7491,9 +7526,13 @@ class StandardLoggingPayloadSetup: """ dynamic_litellm_session_id: Final = litellm_params.get("litellm_session_id") dynamic_litellm_trace_id: Final = litellm_params.get("litellm_trace_id") - metadata: Final[Mapping[str, object] | None] = litellm_params.get("metadata") - metadata_session_id: Final = metadata.get("session_id") if metadata else None - metadata_trace_id: Final = metadata.get("trace_id") if metadata else None + metadata: Final = litellm_params.get("metadata") + metadata_session_id: Final = ( + require_attributes(metadata, _MetadataLookup, "get").get("session_id") if metadata else None + ) + metadata_trace_id: Final = ( + require_attributes(metadata, _MetadataLookup, "get").get("trace_id") if metadata else None + ) ordered_candidates: Final[tuple[object, object, object, object]] = ( (dynamic_litellm_trace_id, dynamic_litellm_session_id, metadata_trace_id, metadata_session_id) @@ -7508,7 +7547,7 @@ class StandardLoggingPayloadSetup: @staticmethod def get_standard_logging_payload_session_id( logging_obj: Logging, - litellm_params: Mapping[str, Any], + litellm_params: Mapping[str, object], ) -> str: """ Returns the end-user/conversation `litellm_session_id` for this request, independent of trace_id. @@ -7523,8 +7562,10 @@ class StandardLoggingPayloadSetup: dynamic_litellm_session_id: Final[object] = litellm_params.get("litellm_session_id") if dynamic_litellm_session_id: return str(dynamic_litellm_session_id) - metadata: Final[Mapping[str, object] | None] = litellm_params.get("metadata") - metadata_session_id: Final = metadata.get("session_id") if metadata else None + metadata: Final = litellm_params.get("metadata") + metadata_session_id: Final = ( + require_attributes(metadata, _MetadataLookup, "get").get("session_id") if metadata else None + ) if metadata_session_id: return str(metadata_session_id) return logging_obj.litellm_session_id diff --git a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py index 78a8cb51bac..9a3197cc044 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py +++ b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py @@ -1,7 +1,8 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Any, Final +from typing import Any, Final, Protocol +from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes from litellm.types.utils import ( CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, @@ -55,19 +56,29 @@ _INTERACTIONS_MODALITY_FIELDS: Final[Mapping[str, str]] = MappingProxyType( ) -def _modality_field(entry: Mapping[str, object]) -> str | None: - return _INTERACTIONS_MODALITY_FIELDS.get(str(entry.get("modality", "")).lower()) +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +def _modality_field(entry: object) -> str | None: + return _INTERACTIONS_MODALITY_FIELDS.get( + str(require_attributes(entry, _SupportsGet, "get").get("modality", "")).lower() + ) def _token_count(value: object) -> int: return value if isinstance(value, int) else 0 -def _modality_token_sums(entries: Sequence[Mapping[str, object]]) -> Mapping[str, int]: +def _modality_token_sums(entries: Sequence[object]) -> Mapping[str, int]: fields: Final = frozenset(field for entry in entries if (field := _modality_field(entry)) is not None) return MappingProxyType( { - field: sum(_token_count(entry.get("tokens")) for entry in entries if _modality_field(entry) == field) + field: sum( + _token_count(require_attributes(entry, _SupportsGet, "get").get("tokens")) + for entry in entries + if _modality_field(entry) == field + ) for field in fields } ) @@ -119,12 +130,16 @@ class InteractionsUsageObjectTransformation: return "total_input_tokens" in usage_object or "total_output_tokens" in usage_object @staticmethod - def transform_interactions_usage_object(usage_object: Mapping[str, Any]) -> Usage: - input_entries: Final = tuple(usage_object.get("input_tokens_by_modality") or ()) + tuple( - usage_object.get("tool_use_tokens_by_modality") or () + def transform_interactions_usage_object(usage_object: Mapping[str, object]) -> Usage: + input_entries: Final = tuple(iter_items(usage_object.get("input_tokens_by_modality") or ())) + tuple( + iter_items(usage_object.get("tool_use_tokens_by_modality") or ()) + ) + cached_sums: Final = _modality_token_sums( + tuple(iter_items(usage_object.get("cached_tokens_by_modality") or ())) + ) + output_sums: Final = _modality_token_sums( + tuple(iter_items(usage_object.get("output_tokens_by_modality") or ())) ) - cached_sums: Final = _modality_token_sums(tuple(usage_object.get("cached_tokens_by_modality") or ())) - output_sums: Final = _modality_token_sums(tuple(usage_object.get("output_tokens_by_modality") or ())) total_cached_tokens: Final = _token_count(usage_object.get("total_cached_tokens")) input_sums: Final = _subtract_cached_from_input( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 5eb9cdb11ce..3fd5390b427 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -6,10 +6,10 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field, replace from datetime import datetime, timezone, tzinfo from types import MappingProxyType -from typing import Final, Literal, Protocol, TypedDict, cast +from typing import Final, Literal, Protocol, SupportsFloat, SupportsIndex, TypedDict, cast from zoneinfo import ZoneInfo, ZoneInfoNotFoundError -from typing_extensions import ReadOnly +from typing_extensions import Buffer, ReadOnly import litellm from litellm._internal_context import current_billing_time @@ -48,7 +48,6 @@ _IMAGE_RESPONSE_CALL_TYPES: Final = frozenset( } ) -# Pre-resolved DataResidency enum values for fast membership checks _VALID_DATA_RESIDENCIES: Final = frozenset(r.value for r in DataResidency) _DEPLOYMENT_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields) @@ -1428,6 +1427,13 @@ def get_regional_uplift_multiplier(model_info: ModelInfo, data_residency: str | multiplier: Final = model_info.get(f"regional_processing_uplift_multiplier_{residency}") if multiplier is None: return 1.0 + if not isinstance(multiplier, (str, Buffer, SupportsFloat, SupportsIndex)): + verbose_logger.error( + "Invalid regional_processing_uplift_multiplier_%s for model (got %s); defaulting to 1.0", + residency, + type(multiplier).__name__, + ) + return 1.0 try: return float(multiplier) except (TypeError, ValueError): diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index ec857218c79..0572ca66fc9 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -6,7 +6,7 @@ import json from collections.abc import Hashable, Iterable, Mapping, Sequence from dataclasses import dataclass from itertools import chain, groupby -from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast +from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast from pydantic import TypeAdapter @@ -1143,7 +1143,7 @@ def strip_prefill_prefix_field(message: AllMessageValues) -> AllMessageValues: return _message_without_fields(message, frozenset({"prefix"})) -def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Any: +def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Mapping[object, object]: """ Filters a value from a dictionary diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 17a1fa7d81b..ae36484b294 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -10,7 +10,7 @@ from collections.abc import Callable, Container, Iterable, Iterator, Mapping, Se from dataclasses import dataclass, replace from enum import Enum from types import MappingProxyType -from typing import Annotated, Any, Final, TypeAlias, TypedDict, TypeVar, cast, overload +from typing import Annotated, Any, Final, Protocol, TypeAlias, TypedDict, TypeVar, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import BaseModel, BeforeValidator, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator @@ -22,7 +22,7 @@ import litellm.types.llms from litellm import verbose_logger from litellm._uuid import uuid from litellm.constants import REDACTED_BY_LITELLM -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, require_attributes from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( anthropic_system_messages, find_anthropic_server_tool_result, @@ -4150,7 +4150,7 @@ class BedrockImageProcessor: _validate_format = validate_format @staticmethod - def _get_document_format(mime_type: str, supported_doc_formats: list[str]) -> str: + def _get_document_format(mime_type: str, supported_doc_formats: Sequence[str]) -> str: """ Get the document format from the mime type @@ -6571,6 +6571,20 @@ def has_tool_with_name(tools: object, tool_name: str) -> bool: return False +class _SupportsModelDump(Protocol): + def model_dump(self) -> dict[str, object]: ... + + +class _StructuredMessagesSource(Protocol): + def get_structured_messages(self, data: dict[str, object], /) -> Sequence[object] | None: ... + + +def _instantiate_structured_message_sources( + source_classes: Iterable[Callable[[], _StructuredMessagesSource]], +) -> Sequence[_StructuredMessagesSource]: + return [source_cls() for source_cls in source_classes] + + def resolve_structured_messages( messages: list[dict[str, object]] | None, request_kwargs: dict[str, Any], @@ -6608,16 +6622,17 @@ def resolve_structured_messages( break # 2. Fallback: try each mapped handler until one produces messages - handlers_to_try: Final[list[Any]] = [] - if call_type is not None and call_type in mappings: - handlers_to_try.append(mappings[call_type]()) - else: - handlers_to_try.extend(handler_cls() for handler_cls in mappings.values()) + handlers_to_try: Final = _instantiate_structured_message_sources( + (mappings[call_type],) if call_type is not None and call_type in mappings else mappings.values() + ) for handler in handlers_to_try: structured = handler.get_structured_messages(request_kwargs) if structured: - return [msg if isinstance(msg, dict) else msg.model_dump() for msg in structured] + return [ + msg if isinstance(msg, dict) else require_attributes(msg, _SupportsModelDump, "model_dump").model_dump() + for msg in structured + ] return None diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index a5e79d7b2b4..c13c9bc31fb 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -12,7 +12,7 @@ import copy import inspect from collections.abc import Iterable, Mapping, Sequence from dataclasses import replace -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError @@ -21,7 +21,9 @@ from litellm.constants import REDACTED_BY_LITELLM from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.classifier_logging import without_classifier_audit from litellm.litellm_core_utils.core_helpers import ( + attribute_or_none, get_metadata_variable_name_from_kwargs, + has_attributes, ) from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items @@ -42,6 +44,11 @@ if TYPE_CHECKING: else: LiteLLMLoggingObject = Any + +class _HasArguments(Protocol): + arguments: str + + MESSAGE_REDACTION_ENABLE_HEADERS: Final = ("litellm-enable-message-redaction", "x-litellm-enable-message-redaction") @@ -163,7 +170,7 @@ def _redact_bedrock_guardrail_trace(response: object) -> None: } response.update(traces) return - for field, value in ((name, getattr(response, name, None)) for name in BEDROCK_GUARDRAIL_TRACE_FIELDS): + for field, value in ((name, attribute_or_none(response, name)) for name in BEDROCK_GUARDRAIL_TRACE_FIELDS): if value is not None: setattr(response, field, _redacted_bedrock_trace(value)) @@ -191,8 +198,8 @@ def _redact_tool_calls(tool_calls) -> None: if not tool_calls: return for tool_call in tool_calls: - function = getattr(tool_call, "function", None) - if function is not None and hasattr(function, "arguments"): + function = attribute_or_none(tool_call, "function") + if function is not None and has_attributes(function, _HasArguments, "arguments"): function.arguments = REDACTED_BY_LITELLM diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index aa6ce6ea99e..6839e11f8c7 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -7,7 +7,7 @@ import logging import threading import time import traceback -from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence, Sized from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeVar, cast @@ -17,7 +17,7 @@ import httpx from openai.types.chat import ChatCompletionChunk from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice from openai.types.chat.chat_completion_chunk import ChoiceDelta -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import NotRequired, TypedDict import litellm @@ -27,7 +27,9 @@ from litellm.constants import STREAM_CHUNK_CHOICE_FIELDS, STREAM_FINISH_CHOICE_F from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import ( UPSTREAM_USAGE_COST, + attribute_or_none, chunk_usage, + decoded_json_object, provider_reported_usage_cost, stamp_litellm_bill_on_usage, ) @@ -129,6 +131,13 @@ def _apply_provider_fields(response: ModelResponseStream, fields: Mapping[str, o setattr(response, key, value) +def _entry_len(entries: Mapping[str, object], key: str) -> int: + value: Final = entries[key] + if isinstance(value, Sized): + return len(value) + raise TypeError(f"object of type '{type(value).__name__}' has no len()") + + def _prepend_upstream_choice_extras( choice: StreamingChoices, upstream_choice: ChunkChoice, extra_keys: frozenset[str] ) -> None: @@ -346,6 +355,9 @@ class _ParsedChunkHiddenParams(LiteLLMBaseModel): upstream_provider: str | None = None +_DECODED_JSON: Final = TypeAdapter(object) + + def _provider_response_model(chunk: object) -> str | None: model: Final[object] = chunk.get("model") if isinstance(chunk, Mapping) else getattr(chunk, "model", None) return model if isinstance(model, str) and model else None @@ -432,9 +444,8 @@ class CustomStreamWrapper: self._closed = False self._stream_created_time: float = time.time() - litellm_params: Final[GenericLiteLLMParams] = GenericLiteLLMParams.model_validate( - dict(**self.logging_obj.model_call_details.get("litellm_params", {})) - ) + raw_litellm_params: Final = dict[str, object](**self.logging_obj.model_call_details.get("litellm_params", {})) + litellm_params: Final[GenericLiteLLMParams] = GenericLiteLLMParams.model_validate(raw_litellm_params) self.merge_reasoning_content_in_choices: bool = litellm_params.merge_reasoning_content_in_choices or False self.sent_first_thinking_block = False self.sent_last_thinking_block = False @@ -465,7 +476,7 @@ class CustomStreamWrapper: _api_base: Final = get_api_base( model=model or "", - optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), + optional_params=raw_litellm_params, ) self._hidden_params: dict[str, object] = { @@ -953,7 +964,7 @@ class CustomStreamWrapper: if "text_output" in chunk: response = CustomStreamWrapper.strip_sse_data_from_chunk(chunk) or "" response = response.strip() - parsed_response = json.loads(response) + parsed_response = decoded_json_object(_DECODED_JSON.validate_python(json.loads(response))) else: return { "text": "", @@ -1087,7 +1098,7 @@ class CustomStreamWrapper: def is_chunk_non_empty( self, - completion_obj: dict[str, Any], + completion_obj: Mapping[str, object], model_response: ModelResponseStream, response_obj: Mapping[str, object], ) -> bool: @@ -1097,7 +1108,7 @@ class CustomStreamWrapper: or ( "tool_calls" in completion_obj and completion_obj["tool_calls"] is not None - and len(completion_obj["tool_calls"]) > 0 + and _entry_len(completion_obj, "tool_calls") > 0 ) or ("function_call" in completion_obj and completion_obj["function_call"] is not None) or ( @@ -1850,7 +1861,9 @@ class CustomStreamWrapper: elif delta.tool_calls is not None: if isinstance(delta.tool_calls, list): for t in delta.tool_calls: - if hasattr(t, "functions") and hasattr(getattr(t, "functions"), "arguments"): + if hasattr(t, "functions") and hasattr( + attribute_or_none(t, "functions"), "arguments" + ): if t.function.arguments is None: t.function.arguments = "" _json_delta: Final[_DeltaDumpDict] = delta.model_dump() @@ -1888,7 +1901,7 @@ class CustomStreamWrapper: return ## CHECK FOR TOOL USE - if "tool_calls" in completion_obj and len(completion_obj["tool_calls"]) > 0: + if "tool_calls" in completion_obj and _entry_len(completion_obj, "tool_calls") > 0: if self.is_function_call is True: # user passed in 'functions' param completion_obj["function_call"] = completion_obj["tool_calls"][0]["function"] completion_obj["tool_calls"] = None @@ -2817,8 +2830,9 @@ class CustomStreamWrapper: def _record_and_strip_usage_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None: self.chunks.append(chunk.model_copy()) + dumped_chunk: Final = MappingProxyType[str, object](chunk.model_dump()) usage_free_chunk: Final = self.model_response_creator( - chunk={key: value for key, value in chunk.model_dump().items() if key != "usage"}, + chunk={key: value for key, value in dumped_chunk.items() if key != "usage"}, hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage Mapping[str, object], getattr(chunk, HIDDEN_PARAMS_ATTR) ), diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index 5dc9679fe78..48a947c31ca 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -11,13 +11,15 @@ A2A Protocol Format: """ import json -from collections.abc import Sequence -from typing import TYPE_CHECKING, Any, Final, Optional +from collections.abc import Mapping, Sequence +from functools import reduce +from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, runtime_checkable from pydantic import ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import has_attributes, require_attributes, subscript from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -30,6 +32,38 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) +_UNSUPPORTED_RESPONSE: Final = object() + + +class _ModelDumpable(Protocol): + def model_dump(self) -> object: ... + + +class _ResultLookup(Protocol): + def get(self, key: str, default: object, /) -> object: ... + + +@runtime_checkable +class _SupportsSetItem(Protocol): + def __setitem__(self, key: str, value: object, /) -> None: ... + + +def _item_assignment_target(value: object) -> _SupportsSetItem: + if isinstance(value, _SupportsSetItem): + return value + raise TypeError(f"'{type(value).__name__}' object does not support item assignment") + + +def _response_payload(response: object) -> object: + if has_attributes(response, _ModelDumpable, "model_dump"): + return response.model_dump() + if isinstance(response, dict): + return response + return _UNSUPPORTED_RESPONSE + + +def _path_step(current: object, key: str) -> object: + return subscript(current, int(key) if key.isdigit() else key) class _A2ATextPart(TypedDict, total=False): @@ -123,7 +157,7 @@ class A2AGuardrailHandler(BaseTranslation): async def process_output_response( self, - response: Any, + response: object, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, @@ -147,17 +181,12 @@ class A2AGuardrailHandler(BaseTranslation): Returns: Modified response with guardrails applied to text content """ - # Handle both dict and Pydantic model responses - if hasattr(response, "model_dump"): - response_dict = response.model_dump() - is_pydantic = True - elif isinstance(response, dict): - response_dict = response - is_pydantic = False - else: + payload: Final = _response_payload(response) + if payload is _UNSUPPORTED_RESPONSE: verbose_proxy_logger.warning("A2A: Unknown response type %s, skipping guardrail", type(response)) return response + response_dict: Final = require_attributes(payload, _ResultLookup, "get") result: Final = response_dict.get("result", {}) if not result or not isinstance(result, dict): verbose_proxy_logger.debug("A2A: No result in response, skipping guardrail") @@ -218,15 +247,8 @@ class A2AGuardrailHandler(BaseTranslation): verbose_proxy_logger.debug("A2A: Processed output response") - # Update the original response - if is_pydantic: - # For Pydantic models, we need to update the underlying dict - # and the model will reflect the changes - response_dict["result"] = result - return response - else: - response["result"] = result - return response + _item_assignment_target(response_dict)["result"] = result + return response async def process_output_streaming_response( self, @@ -448,20 +470,11 @@ class A2AGuardrailHandler(BaseTranslation): def _apply_text_to_path( self, - result: dict[str | int, Any], + result: Mapping[str, object], path: tuple[str, ...], part_idx: int, text: str, ) -> None: """Apply guardrailed text back to the specified path in the result.""" - # Navigate to the parts list - current = result - for key in path: - if key.isdigit(): - # Array index - current = current[int(key)] - else: - current = current[key] - - # Update the text in the part - current[part_idx]["text"] = text + parts: Final = reduce(_path_step, path, result) + _item_assignment_target(subscript(parts, part_idx))["text"] = text diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 4062a957e00..ed4ed295dd3 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -1081,7 +1081,7 @@ class ModelResponseIterator: if "caller" in content_block_start["content_block"]: caller_data: Final = content_block_start["content_block"]["caller"] if caller_data and tool_use is not None: - tool_use["caller"] = cast(dict[str, Any], caller_data) + tool_use["caller"] = caller_data elif content_block_start["content_block"]["type"] == "redacted_thinking": ( thinking_blocks, diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 5d4957fa166..1c42dd15853 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -2107,7 +2107,7 @@ def _strip_empty_content_blocks_from_message(message: object) -> object: def strip_empty_content_blocks_from_anthropic_messages( messages: Sequence[object], -) -> list[Any]: +) -> Sequence[object]: """ Return a new message list with empty or whitespace-only ``{"type": "text"}`` content blocks and empty *unsigned* ``{"type": "thinking"}`` blocks removed. @@ -2220,7 +2220,7 @@ def _sanitize_tool_use_id_content_block(block: object) -> object: return block -def sanitize_tool_use_ids_in_anthropic_messages(messages: Sequence[object]) -> list[Any]: +def sanitize_tool_use_ids_in_anthropic_messages(messages: Sequence[object]) -> Sequence[object]: """ Return a new message list with ``tool_use`` / ``server_tool_use`` ``id`` and ``tool_result`` ``tool_use_id`` values rewritten to satisfy Anthropic's diff --git a/litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py b/litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py index 0aa30998557..dfda86339e7 100644 --- a/litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py @@ -17,11 +17,14 @@ execute it. import asyncio import contextlib import json -from collections.abc import AsyncIterator +from collections.abc import AsyncIterable, AsyncIterator from typing import TYPE_CHECKING, Any, Final, cast +from pydantic import TypeAdapter + from litellm._logging import verbose_logger from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES +from litellm.litellm_core_utils.core_helpers import decoded_json_object, has_attributes from litellm.types.integrations.custom_logger import AgenticLoopLimitError, AgenticLoopSafetyError if TYPE_CHECKING: @@ -42,6 +45,7 @@ AGENTIC_LOOP_LIMIT_ERROR_SSE_BYTES: Final = ( b'"The agentic loop reached its limit before this turn finished. ' b'Raise max_agentic_loops to allow more loops."}}\n\n' ) +_DECODED_JSON: Final = TypeAdapter(object) def is_server_fulfilled_tool_leak_error(chunk: object) -> bool: @@ -76,10 +80,10 @@ def _parse_sse_events(raw: bytes) -> list[tuple]: continue data_str = stripped[len("data:") :].strip() try: - data = json.loads(data_str) + data = _DECODED_JSON.validate_python(json.loads(data_str)) except (json.JSONDecodeError, ValueError): continue - event_type = current_event_type or data.get("type", "") + event_type = current_event_type or decoded_json_object(data).get("type", "") current_event_type = None events.append((event_type, data)) return events @@ -371,7 +375,7 @@ class AgenticAnthropicStreamingIterator: verbose_logger.debug("AgenticStreamingIterator: Could not rebuild response from SSE bytes") return - result: Final = await self._http_handler._call_agentic_completion_hooks( + result: Final[object] = await self._http_handler._call_agentic_completion_hooks( response=rebuilt, model=self._model, messages=self._messages, @@ -386,7 +390,7 @@ class AgenticAnthropicStreamingIterator: if result is None: return - if hasattr(result, "__aiter__"): + if has_attributes(result, AsyncIterable[object], "__aiter__"): self._follow_up_iterator = result.__aiter__() elif isinstance(result, dict): from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( diff --git a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py index 5c899e8e8bf..20dd13a329e 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py @@ -2,8 +2,8 @@ ## Translates OpenAI call to Anthropic `/v1/messages` format import asyncio from collections import deque -from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Iterable, Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Final import anyio from pydantic import ConfigDict, field_validator @@ -12,6 +12,7 @@ from litellm import verbose_logger from litellm._logging import redact_internal_details_from_client_message from litellm._uuid import uuid from litellm.exceptions import APIError, MidStreamFallbackError +from litellm.litellm_core_utils.core_helpers import has_attributes, iter_items from litellm.litellm_core_utils.prompt_templates.reasoning_utils import ( encrypted_reasoning_signature, ) @@ -126,7 +127,7 @@ class AnthropicResponsesStreamWrapper: def __init__( self, - responses_stream: Any, + responses_stream: object, model: str, litellm_logging_obj: "LiteLLMLoggingObject | None" = None, safeguards_evaluator: "SafeguardsEvaluator | None" = None, @@ -414,7 +415,9 @@ class AnthropicResponsesStreamWrapper: response_obj: Final = getattr(event, "response", None) or ( event.get("response") if isinstance(event, dict) else None ) - output: Final = (getattr(response_obj, "output", None) or ()) if response_obj is not None else () + output: Final[Iterable[object]] = ( + (getattr(response_obj, "output", None) or ()) if response_obj is not None else () + ) refusal_text: Final = responses_output_refusal_text(output) or (self._refusal_text or None) status: Final = getattr(response_obj, "status", None) if response_obj is not None else None has_tool_call: Final = any( @@ -508,14 +511,14 @@ class AnthropicResponsesStreamWrapper: return self._chunk_queue.popleft() try: - if hasattr(self.responses_stream, "__aiter__"): + if has_attributes(self.responses_stream, AsyncIterable[object], "__aiter__"): async for event in self.responses_stream: self._process_event(event) if self._chunk_queue: return self._chunk_queue.popleft() else: if self._sync_responses_iterator is None: - self._sync_responses_iterator = iter(self.responses_stream) + self._sync_responses_iterator = iter(iter_items(self.responses_stream)) sync_iterator: Final = self._sync_responses_iterator missing: Final = object() while (event := await asyncio.to_thread(next, sync_iterator, missing)) is not missing: diff --git a/litellm/llms/anthropic/messages/responses_adapters/transformation.py b/litellm/llms/anthropic/messages/responses_adapters/transformation.py index b4455f1a4e2..d7b8592f45c 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/messages/responses_adapters/transformation.py @@ -10,6 +10,8 @@ from collections.abc import Iterable, Mapping, Sequence from itertools import groupby from typing import Any, Final, cast +from pydantic import TypeAdapter + from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, @@ -58,6 +60,7 @@ from litellm.types.llms.openai import ( REASONING_SUMMARY_PART_SEPARATOR: Final = "\n\n" RESPONSES_INCLUDE_ENCRYPTED_REASONING: Final = "reasoning.encrypted_content" +_DECODED_JSON: Final = TypeAdapter(object) class LiteLLMAnthropicToResponsesAPIAdapter: @@ -805,15 +808,17 @@ class LiteLLMAnthropicToResponsesAPIAdapter: content.append(reasoning_block) elif item_type == "function_call": try: - input_data = json.loads(item.get("arguments", "{}")) + input_data = _DECODED_JSON.validate_python(json.loads(item.get("arguments", "{}"))) except (json.JSONDecodeError, TypeError): input_data = {} content.append( - AnthropicResponseContentBlockToolUse( - type="tool_use", - id=item.get("call_id") or item.get("id", ""), - name=item.get("name", ""), - input=input_data, + AnthropicResponseContentBlockToolUse.model_validate( + { + "type": "tool_use", + "id": item.get("call_id") or item.get("id", ""), + "name": item.get("name", ""), + "input": input_data, + } ).model_dump(exclude_none=True) ) stop_reason = "tool_use" diff --git a/litellm/llms/anthropic/messages/stream_assembly.py b/litellm/llms/anthropic/messages/stream_assembly.py index 843e8a6e495..9a31d171bc8 100644 --- a/litellm/llms/anthropic/messages/stream_assembly.py +++ b/litellm/llms/anthropic/messages/stream_assembly.py @@ -4,10 +4,13 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping, map_finish_reason +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object, map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict from litellm.llms.anthropic.chat.handler import ( ModelResponseIterator as AnthropicModelResponseIterator, ) @@ -21,6 +24,8 @@ from litellm.types.utils import ( TextCompletionResponse, ) +_DECODED_JSON: Final = TypeAdapter(object) + def split_sse_chunk_into_events(chunk: str | bytes) -> list[str]: """ @@ -98,6 +103,14 @@ _FAST_PATH_DISALLOWED_DELTA_TYPES: Final = frozenset( ) +def _joined_text(parts: Sequence[object]) -> str: + strings: Final = [part for part in parts if isinstance(part, str)] + if len(strings) == len(parts): + return "".join(strings) + first_bad: Final = next(position for position, part in enumerate(parts) if not isinstance(part, str)) + raise TypeError(f"sequence item {first_bad}: expected str instance, {type(parts[first_bad]).__name__} found") + + def collapse_pure_text_chunks( all_chunks: Sequence[str | bytes], ) -> list[str] | None: @@ -120,10 +133,10 @@ def collapse_pure_text_chunks( if ev: normalized.append(ev) - text_block_indexes: Final[set] = set() + text_block_indexes: Final[set[object]] = set() out: Final[list[str]] = [] - pending_text: list[str] = [] - pending_index: int | None = None + pending_text: list[object] = [] + pending_index: object = None saw_any_text_delta = False def flush() -> None: @@ -132,7 +145,7 @@ def collapse_pure_text_chunks( merged: Final = { "type": "content_block_delta", "index": pending_index if pending_index is not None else 0, - "delta": {"type": "text_delta", "text": "".join(pending_text)}, + "delta": {"type": "text_delta", "text": _joined_text(pending_text)}, } out.append("data: " + json.dumps(merged)) pending_text = [] @@ -158,20 +171,20 @@ def collapse_pure_text_chunks( if not json_str.startswith("{"): continue try: - data = json.loads(json_str) + data = decoded_json_object(_DECODED_JSON.validate_python(json.loads(json_str))) except (json.JSONDecodeError, ValueError): return None etype = data.get("type") if etype == "content_block_start": - block = data.get("content_block") or {} + block = decoded_json_object(data.get("content_block") or {}) if block.get("type") != "text": return None text_block_indexes.add(data.get("index")) flush() out.append(ev) elif etype == "content_block_delta": - delta = data.get("delta") or {} + delta = decoded_json_object(data.get("delta") or {}) dtype = delta.get("type") if dtype in _FAST_PATH_DISALLOWED_DELTA_TYPES: return None @@ -291,9 +304,10 @@ def extract_sse_data(event_str: str) -> dict | None: if payload is None or not payload.startswith("{"): return None try: - return json.loads(payload) + decoded: Final = _DECODED_JSON.validate_python(json.loads(payload)) except (ValueError, TypeError): return None + return decoded if is_str_keyed_dict(decoded) else None def build_usage_only_response_from_chunks( @@ -492,7 +506,7 @@ def _carries_no_data(fields: Sequence[tuple[str, str]]) -> bool: def _json_mapping(payload: str) -> Mapping[str, object] | None: try: - value: Final[object] = json.loads(payload) + value: Final = _DECODED_JSON.validate_python(json.loads(payload)) except ValueError: return None return as_str_mapping(value) @@ -555,7 +569,7 @@ def is_anthropic_sse_stream(all_chunks: Sequence[object]) -> bool: def _is_json_no_reader_can_take(payload: str) -> bool: try: - value: Final[object] = json.loads(payload) + value: Final = _DECODED_JSON.validate_python(json.loads(payload)) except ValueError: return False if value is None: diff --git a/litellm/llms/apiserpent/search/transformation.py b/litellm/llms/apiserpent/search/transformation.py index faf1903ffe8..2e5fd2032c1 100644 --- a/litellm/llms/apiserpent/search/transformation.py +++ b/litellm/llms/apiserpent/search/transformation.py @@ -12,7 +12,9 @@ from typing import Final, Literal, cast from urllib.parse import urlencode import httpx +from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object, iter_items from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.apiserpent.search.defaults import ( @@ -33,6 +35,7 @@ from litellm.secret_managers.main import get_secret_str DEEP_SEARCH_PARAM: Final = "deep" APISERPENT_BASE: Final = "https://apiserpent.com" APISERPENT_PARAMS_KEY: Final = "_apiserpent_params" +_DECODED_JSON: Final = TypeAdapter(object) class APISerpentSearchConfig(BaseSearchConfig): @@ -150,22 +153,24 @@ class APISerpentSearchConfig(BaseSearchConfig): Full format nests results under ``results.organic[]``; simple format returns a flat ``results[]`` array. Both expose title/url/snippet. """ - response_json: Final = raw_response.json() + response_json: Final = decoded_json_object(_DECODED_JSON.validate_python(raw_response.json())) raw_results: Final = response_json.get("results") or {} - organic: Final = raw_results.get("organic", []) if isinstance(raw_results, dict) else raw_results + nested_results: Final = as_str_mapping(raw_results) + organic: Final = nested_results.get("organic", []) if nested_results is not None else raw_results - results: Final[list[SearchResult]] = [] - for result in organic: - results.append( - SearchResult( - title=result.get("title", ""), - url=result.get("url", ""), - snippet=result.get("snippet", ""), - date=result.get("date"), - last_updated=None, - ) + results: Final = [ + SearchResult.model_validate( + { + "title": result.get("title", ""), + "url": result.get("url", ""), + "snippet": result.get("snippet", ""), + "date": result.get("date"), + "last_updated": None, + } ) + for result in map(decoded_json_object, iter_items(organic)) + ] return SearchResponse( results=results, diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 69ffac5d6f9..ef74042156e 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -2,7 +2,7 @@ import asyncio import json import time from collections.abc import Callable, Coroutine -from typing import Final +from typing import Final, Protocol import httpx from openai import ( @@ -17,6 +17,7 @@ from pydantic import BaseModel import litellm from litellm.constants import DEFAULT_MAX_RETRIES +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing from litellm.llms.custom_httpx.http_handler import ( @@ -60,6 +61,11 @@ from .image_generation import ( from .image_generation.http_utils import azure_deployment_image_generation_json_body +class _HasResponseText(Protocol): + @property + def text(self) -> str: ... + + class AzureOpenAIAssistantsAPIConfig: """ Reference: https://learn.microsoft.com/en-us/azure/ai-services/openai/assistants-reference-messages?tabs=python#create-message @@ -853,11 +859,11 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except Exception as e: status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) error_text = str(e) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - error_text = error_response.text + error_text = require_attributes(error_response, _HasResponseText, "text").text raise AzureOpenAIError( status_code=status_code, message=error_text, diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index a814a9fbac0..6a0740b2f15 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -251,9 +251,8 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): def compact_url(self, api_base: str) -> str: params: Final = httpx.URL(api_base).params - v1_params: Final = ( - params if BaseAzureLLM.is_azure_v1_api_version(params.get("api-version")) else params.remove("api-version") - ) + api_version: Final = next(iter(params.get_list("api-version")), None) + v1_params: Final = params if BaseAzureLLM.is_azure_v1_api_version(api_version) else params.remove("api-version") return super().compact_url(str(_v1_responses_url(api_base).copy_with(params=v1_params))) ######################################################### diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index 423a09c5c92..5dae6cc9fc3 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -308,7 +308,7 @@ class AzureAIAgentsHandler: logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, - timeout: float, + timeout: float | str | httpx.Timeout | None, client: HTTPHandler | None = None, headers: dict | None = None, ) -> ModelResponse: @@ -432,7 +432,7 @@ class AzureAIAgentsHandler: logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, - timeout: float, + timeout: float | str | httpx.Timeout | None, client: AsyncHTTPHandler | None = None, headers: dict | None = None, ) -> ModelResponse: @@ -559,7 +559,7 @@ class AzureAIAgentsHandler: logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, - timeout: float, + timeout: float | str | httpx.Timeout | None, headers: dict | None = None, ) -> AsyncIterator[ModelResponseStream]: """Execute async streaming completion using Azure Agent Service with native SSE.""" diff --git a/litellm/llms/azure_ai/agents/transformation.py b/litellm/llms/azure_ai/agents/transformation.py index 18f8f7916b7..f6e70496247 100644 --- a/litellm/llms/azure_ai/agents/transformation.py +++ b/litellm/llms/azure_ai/agents/transformation.py @@ -323,7 +323,7 @@ class AzureAIAgentsConfig(BaseConfig): logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, - timeout: float | Any, + timeout: float | str | httpx.Timeout | None, acompletion: bool, stream: bool | None = False, headers: dict | None = None, diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 3609871fdf8..5fcb84593d2 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -8,7 +8,7 @@ from collections.abc import Mapping from typing import Final import httpx -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import verbose_logger @@ -22,6 +22,8 @@ from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +_DECODED_JSON: Final = TypeAdapter(object) + class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): """ @@ -42,7 +44,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): system: object = None, optional_params: Mapping[str, JsonValue] | None = None, context_management: JsonValue = None, - ) -> dict[str, object]: + ) -> object: """ Handle a CountTokens request using httpx with Azure authentication. @@ -126,7 +128,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): message=error_text, ) - azure_response: Final = response.json() + azure_response: Final = _DECODED_JSON.validate_python(response.json()) verbose_logger.debug("Azure AI Anthropic response: %s", azure_response) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 14f7d4bfd2b..8657fcddab4 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -11,6 +11,7 @@ from pydantic import JsonValue import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountFields, ChatCountTranslationError, @@ -139,13 +140,14 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): ) if result is not None: + response_fields: Final = decoded_json_object(result) return TokenCountResponse.model_validate( { - "total_tokens": result.get("input_tokens", 0), + "total_tokens": response_fields.get("input_tokens", 0), "request_model": request_model, "model_used": model_to_use, "tokenizer_type": "azure_ai_anthropic_api", - "original_response": result, + "original_response": response_fields, } ) except AnthropicError as e: diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index 1f9ded827b6..63f9f726c71 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -3,7 +3,7 @@ Azure Anthropic messages transformation config - extends AnthropicMessagesConfig """ from collections.abc import Sequence -from typing import Any, Final +from typing import Final from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, @@ -109,14 +109,10 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): Processes both `system` and `messages` content blocks. """ - def _sanitize(cache_control: Any) -> None: - if isinstance(cache_control, dict): - cache_control.pop("scope", None) - def _process_content_list(content: list) -> None: for item in content: - if isinstance(item, dict) and "cache_control" in item: - _sanitize(item["cache_control"]) + if isinstance(item, dict) and "cache_control" in item and isinstance(item["cache_control"], dict): + item["cache_control"].pop("scope", None) if "system" in anthropic_messages_request: system: Final = anthropic_messages_request["system"] diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index 9bd476efdff..00d3ad63fda 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -2,11 +2,12 @@ import base64 from collections.abc import Mapping, Sequence from io import BufferedReader from types import MappingProxyType -from typing import Any, Final +from typing import Final, Protocol from httpx._types import RequestFiles import litellm +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.llms.azure_ai.common_utils import ( AzureFoundryModelInfo, get_azure_ai_auth_headers, @@ -21,6 +22,10 @@ from litellm.types.llms.openai import FileTypes from litellm.types.router import GenericLiteLLMParams +class _BinaryReadable(Protocol): + def read(self) -> bytes: ... + + class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): """ Azure AI Foundry FLUX 2 image edit config @@ -100,7 +105,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): if image is None: raise ValueError("FLUX 2 image edit requires an image.") - images: Final = tuple(image) if isinstance(image, list) else (image,) + images: Final[Sequence[object]] = [*image] if isinstance(image, list) else [image] if not images: raise ValueError("FLUX 2 image edit requires at least one image.") max_reference_images: Final = 10 if "flex" in model.lower() else 8 @@ -121,14 +126,14 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): } return request_body, [] - def _convert_image_to_base64(self, image: Any) -> str: + def _convert_image_to_base64(self, image: object) -> str: """Convert image file to base64 string""" if isinstance(image, BufferedReader): image_bytes = image.read() image.seek(0) # Reset file pointer for potential reuse elif isinstance(image, bytes): image_bytes = image - elif hasattr(image, "read"): + elif has_attributes(image, _BinaryReadable, "read"): image_bytes = image.read() else: raise ValueError(f"Unsupported image type: {type(image)}") diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 71368f4ef61..6794739be24 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,7 +1,9 @@ from abc import ABC, abstractmethod from collections.abc import Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional, Protocol + +from litellm.litellm_core_utils.core_helpers import has_attributes if TYPE_CHECKING: from fastapi import HTTPException @@ -15,6 +17,10 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +class _ModelDumpable(Protocol): + def model_dump(self) -> object: ... + + @dataclass(slots=True) class StreamTransformSink: """Out-parameter used by ``process_output_streaming_response`` to hand the @@ -81,7 +87,7 @@ class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( - user_api_key_dict: Any | None, + user_api_key_dict: object, ) -> dict[str, object]: """ Transform user_api_key_dict to a metadata dict with prefixed keys. @@ -99,7 +105,11 @@ class BaseTranslation(ABC): return {} # Convert to dict if it's a Pydantic object - user_dict = user_api_key_dict.model_dump() if hasattr(user_api_key_dict, "model_dump") else user_api_key_dict + user_dict: Final = ( + user_api_key_dict.model_dump() + if has_attributes(user_api_key_dict, _ModelDumpable, "model_dump") + else user_api_key_dict + ) if not isinstance(user_dict, dict): return {} diff --git a/litellm/llms/base_llm/managed_resources/isolation.py b/litellm/llms/base_llm/managed_resources/isolation.py index 09ca23a8ded..135986a6563 100644 --- a/litellm/llms/base_llm/managed_resources/isolation.py +++ b/litellm/llms/base_llm/managed_resources/isolation.py @@ -11,8 +11,9 @@ are denied so an empty user_id can never select an unscoped query. """ from collections.abc import Mapping, Sequence -from typing import Any, Final +from typing import Final, Protocol +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.types.proxy.auth.user_api_key_auth import ( UserAPIKeyAuth, ) @@ -47,15 +48,20 @@ def resolve_resource_owner_id( return None -def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, object]: +class _Identified(Protocol): + @property + def id(self) -> object: ... + + +def build_list_page(items: Sequence[object], has_more: bool = False) -> dict[str, object]: """Build the OpenAI-style paginated list response shape used by managed file/batch/vector-store listings. ``first_id`` and ``last_id`` are sourced from each item's ``.id`` attribute.""" return { "object": "list", "data": items, - "first_id": items[0].id if items else None, - "last_id": items[-1].id if items else None, + "first_id": require_attributes(items[0], _Identified, "id").id if items else None, + "last_id": require_attributes(items[-1], _Identified, "id").id if items else None, "has_more": has_more, } diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index bc2ef2441d2..57cc39a03f4 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -11,7 +11,7 @@ from litellm import verbose_logger from litellm._uuid import uuid from litellm.caching.caching import InMemoryCache from litellm.constants import RESPONSE_FORMAT_TOOL_NAME -from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.core_helpers import decoded_json_object, map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers @@ -50,6 +50,7 @@ from litellm.types.utils import GenericStreamingChunk as GChunk from ..common_utils import ( BedrockError, BedrockEventStreamResponseDict, + bedrock_event_chunk_text, bedrock_event_stream_header, bedrock_event_stream_response, bedrock_stream_event_error_status, @@ -59,6 +60,7 @@ from ..common_utils import ( get_bedrock_response_stream_shape, get_bedrock_stream_event_statuses, get_bedrock_tool_name, + parse_bedrock_event_stream, ) bedrock_tool_name_mappings: Final[InMemoryCache] = InMemoryCache(max_size_in_memory=50, default_ttl=600) @@ -922,7 +924,7 @@ class AWSEventStreamDecoder: response_stream_shape: "Shape", request_id_headers: httpx.Headers | None = None, ) -> str | None: - parsed_response: Final = self.parser.parse(response_dict, response_stream_shape) + parsed_response: Final = parse_bedrock_event_stream(self.parser, response_dict, response_stream_shape) if response_dict["status_code"] != 200: raise build_bedrock_stream_error(response_dict, response_stream_shape, headers=request_id_headers) @@ -932,11 +934,12 @@ class AWSEventStreamDecoder: raise build_bedrock_stream_event_error( event_type, event_error_status, response_dict["body"], headers=request_id_headers ) - if "chunk" in parsed_response: - chunk = parsed_response.get("chunk") - if not chunk: + parsed_fields: Final = decoded_json_object(parsed_response) + if "chunk" in parsed_fields: + event_chunk: Final = parsed_fields.get("chunk") + if not event_chunk: return None - return chunk.get("bytes").decode() + return bedrock_event_chunk_text(event_chunk) else: chunk = response_dict.get("body") if not chunk: diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 0f8694a7e46..de0132ec98c 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -13,7 +13,7 @@ import warnings from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias from typing_extensions import ( ReadOnly, @@ -24,6 +24,7 @@ from typing_extensions import ( if TYPE_CHECKING: from botocore.eventstream import EventStreamMessage from botocore.model import ServiceModel, Shape + from botocore.parsers import EventStreamJSONParser from litellm.types.llms.bedrock import BedrockCreateBatchRequest @@ -33,7 +34,7 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm import verbose_logger from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix -from litellm.litellm_core_utils.core_helpers import lax_bool +from litellm.litellm_core_utils.core_helpers import decoded_json_object, lax_bool, require_attributes from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.litellm_core_utils.prompt_templates.json_schema_utils import ( schema_property_count, @@ -135,7 +136,7 @@ _BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = (*AWS_AUTH_PARAM_KEYS def merge_bedrock_aws_request_params( litellm_params: Mapping[str, object], optional_params: Mapping[str, object], -) -> dict[str, Any]: +) -> dict[str, object]: """Merge deployment and request parameters without allowing auth escalation. Deployment configuration is authoritative for AWS authentication. When a @@ -2516,15 +2517,16 @@ class BedrockEventStreamDecoderBase: ), ) response_dict: Final = event.to_response_dict() - parsed_response: Final = self.parser.parse(response_dict, response_stream_shape) + parsed_response: Final = parse_bedrock_event_stream(self.parser, response_dict, response_stream_shape) if response_dict["status_code"] != 200: raise build_bedrock_stream_error(response_dict, response_stream_shape) - if "chunk" in parsed_response: - chunk = parsed_response.get("chunk") - if not chunk: + parsed_fields: Final = decoded_json_object(parsed_response) + if "chunk" in parsed_fields: + event_chunk: Final = parsed_fields.get("chunk") + if not event_chunk: return None - return chunk.get("bytes").decode() + return bedrock_event_chunk_text(event_chunk) else: chunk = response_dict.get("body") if not chunk: @@ -2541,6 +2543,20 @@ def _decoded_json_value(raw: str) -> object: return _JSON_VALUE.validate_python(json.loads(raw)) +class _DecodablePayload(Protocol): + def decode(self) -> str: ... + + +def parse_bedrock_event_stream( + parser: EventStreamJSONParser, response_dict: Mapping[str, object], response_stream_shape: Shape +) -> object: + return _JSON_VALUE.validate_python(parser.parse(response_dict, response_stream_shape)) + + +def bedrock_event_chunk_text(chunk: object) -> str: + return require_attributes(decoded_json_object(chunk).get("bytes"), _DecodablePayload, "decode").decode() + + def get_anthropic_beta_from_headers(headers: dict) -> list[str]: """ Extract anthropic-beta header values and convert them to a list. diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index dec099f37c0..6f0a18b0088 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -45,7 +45,7 @@ class _CountTokensReply(LiteLLMPydanticObjectBase): @dataclass(frozen=True, slots=True) class CountedTokens: input_tokens: int - original_response: Mapping[str, JsonValue] + original_response: Mapping[str, object] tokenizer_type: str diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 5f8f478e4cd..84067923788 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -7,14 +7,18 @@ Simplified handler leveraging existing LiteLLM Bedrock infrastructure. from typing import Any, Final import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.bedrock.base_aws_llm import run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client +_DECODED_JSON: Final = TypeAdapter(object) + class BedrockCountTokensHandler(BedrockCountTokensConfig): """ @@ -29,7 +33,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): litellm_params: dict[str, Any], resolved_model: str, client: AsyncHTTPHandler | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Handle a CountTokens request using existing LiteLLM patterns. @@ -111,12 +115,12 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): response=response, ) - bedrock_response: Final = response.json() + bedrock_response: Final = _DECODED_JSON.validate_python(response.json()) verbose_logger.debug("Bedrock response: %s", bedrock_response) # Transform response back to expected format - final_response: Final = self.transform_bedrock_response_to_anthropic(bedrock_response) + final_response: Final = self.transform_bedrock_response_to_anthropic(decoded_json_object(bedrock_response)) verbose_logger.debug("Final response: %s", final_response) diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index df5a3827c1f..cd59e714ba0 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -247,9 +247,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): ) return join_api_base(base_url, f"model/{encoded_model_id}/count-tokens") - def transform_bedrock_response_to_anthropic( - self, bedrock_response: Mapping[str, JsonValue] - ) -> dict[str, JsonValue]: + def transform_bedrock_response_to_anthropic(self, bedrock_response: Mapping[str, object]) -> dict[str, object]: """ Transform Bedrock CountTokens response to Anthropic format. diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index ef54c85c131..e5f470fb87b 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -11,7 +11,7 @@ from datetime import datetime from email.utils import parsedate_to_datetime from itertools import chain from types import MappingProxyType -from typing import Any, Final, Literal, TypeAlias, TypedDict +from typing import Any, Final, Literal, Protocol, TypeAlias, TypedDict from urllib.parse import quote, unquote, urlencode import httpx @@ -19,6 +19,7 @@ from httpx import Headers, Response from openai.types.file_deleted import FileDeleted from pydantic import ConfigDict, Field, TypeAdapter from typing_extensions import ReadOnly +from typing_extensions import TypedDict as ExtraItemsTypedDict from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -38,6 +39,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( split_configured_cloud_bucket_name, validate_managed_cloud_file_id, ) +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.prompt_templates.common_utils import ( text_completion_prompt_to_messages, ) @@ -151,6 +153,10 @@ def _retrieved_s3_file_size(raw_response: Response) -> int: ) +class _SupportsRstrip(Protocol): + def rstrip(self, chars: str, /) -> str: ... + + class _S3DeleteContext(LiteLLMBaseModel): file_id: str = Field(min_length=1) @@ -168,6 +174,10 @@ _LITELLM_PARAMS_THE_MAPPER_TAKES: Final = frozenset({"allowed_openai_params"}) _MAPPED_PARAMS_THE_REQUEST_HANDLER_STRIPS: Final = frozenset({"json_mode"}) +class _ParamsTheMapperOwns(ExtraItemsTypedDict, total=False, extra_items=object): + pass + + _UPLOAD_PARAMS: Final = TypeAdapter(dict[str, object]) @@ -710,9 +720,11 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): encoded_object_name: Final = encode_s3_object_key_for_url(object_name) # S3 endpoint URL format - s3_endpoint_url: Final = ( + s3_endpoint_url: Final = require_attributes( request_params.get("s3_endpoint_url") - or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" + or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}", + _SupportsRstrip, + "rstrip", ).rstrip("/") return f"{s3_endpoint_url}/{bucket_name}/{encoded_object_name}" @@ -1062,7 +1074,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): from litellm.types.utils import LlmProviders messages: Final = map_developer_role_to_system_role(openai_request_body.get("messages", [])) - optional_params: Final = {k: v for k, v in openai_request_body.items() if k not in ["model", "messages"]} + request_params: Final[Mapping[str, object]] = openai_request_body + optional_params: Final = {k: v for k, v in request_params.items() if k not in ["model", "messages"]} # --- Anthropic: use existing AmazonAnthropicClaudeConfig --- if provider == LlmProviders.ANTHROPIC: @@ -1085,8 +1098,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): custom_llm_provider="bedrock", messages=messages, model_info=deployment_model_info(litellm_params), - **MappingProxyType( - { + **_ParamsTheMapperOwns( + **{ k: v for k, v in optional_params.items() if not is_litellm_owned_kwarg(k) or k in _LITELLM_PARAMS_THE_MAPPER_TAKES diff --git a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py index 1cece8743de..def53da4493 100644 --- a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py +++ b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py @@ -18,8 +18,10 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles +from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.images.main import ImageEditOptionalRequestParams @@ -38,6 +40,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + def _nova_canvas_task_body( *, @@ -401,7 +405,7 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): json_mode: bool | None = None, ) -> ImageResponse: try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error parsing Nova Canvas image edit response: {e}", @@ -416,18 +420,19 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): headers=raw_response.headers, ) - images: Final[list[str]] = response_data.get("images") or [] + response_object: Final = decoded_json_object(response_data) + images: Final[object] = response_object.get("images") or [] - if "errors" in response_data and not images: + if "errors" in response_object and not images: raise self.get_error_class( - error_message=f"Nova Canvas image edit error: {response_data['errors']}", + error_message=f"Nova Canvas image edit error: {response_object['errors']}", status_code=raw_response.status_code, headers=raw_response.headers, ) # Nova Canvas InvokeModel success body uses "images" and optional "error" (AWS docs); # it does not use Stability-style "finish_reasons". - error_msg = response_data.get("message") or response_data.get("error") + error_msg = response_object.get("message") or response_object.get("error") if error_msg and not images: if not isinstance(error_msg, str): error_msg = str(error_msg) @@ -439,7 +444,7 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): model_response: Final = ImageResponse() model_response.data = [] - for image_b64 in images: + for image_b64 in iter_items(images): if image_b64: model_response.data.append( ImageObject( diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index d0ce5895970..8605e7e29e9 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -22,11 +22,13 @@ API Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parame """ import base64 -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, runtime_checkable import httpx from httpx._types import RequestFiles +from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items, json_contains from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.images.main import ImageEditOptionalRequestParams @@ -44,6 +46,19 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + + +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, key: str | int, /) -> object: ... + + +def _subscript(value: object, key: str | int) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + class BedrockStabilityImageEditConfig(BaseImageEditConfig): """ @@ -299,7 +314,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): OpenAI expects: {"data": [{"b64_json": "base64..."}], "created": timestamp} """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error parsing Bedrock Stability response: {e}", @@ -308,18 +323,19 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): ) # Check for errors in response - if "errors" in response_data: + if json_contains(response_data, "errors"): raise self.get_error_class( - error_message=f"Bedrock Stability error: {response_data['errors']}", + error_message=f"Bedrock Stability error: {_subscript(response_data, 'errors')}", status_code=raw_response.status_code, headers=raw_response.headers, ) # Check finish_reasons - finish_reasons: Final = response_data.get("finish_reasons", []) - if finish_reasons and finish_reasons[0]: + response_fields: Final = decoded_json_object(response_data) + finish_reasons: Final = response_fields.get("finish_reasons", []) + if finish_reasons and _subscript(finish_reasons, 0): raise self.get_error_class( - error_message=f"Bedrock Stability error: {finish_reasons[0]}", + error_message=f"Bedrock Stability error: {_subscript(finish_reasons, 0)}", status_code=400, headers=raw_response.headers, ) @@ -329,9 +345,9 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): model_response.data = [] # Extract images from response - images: Final = response_data.get("images", []) + images: Final = response_fields.get("images", []) if images: - for image_b64 in images: + for image_b64 in iter_items(images): if image_b64: model_response.data.append( ImageObject( diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 3514a5e6e5c..a20eae0b59d 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -14,7 +14,7 @@ from pydantic import InstanceOf, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, @@ -849,6 +849,7 @@ class InvokeStreamFrame: _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) _Frame: Final = TypeVar("_Frame", ConverseStreamFrame, InvokeStreamFrame) @@ -1298,7 +1299,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): texts: Sequence[tuple[_StreamGroupKey, str]] = () if event_type == "contentBlockDelta": try: - payload_dict: dict[str, object] = json.loads(payload_bytes) + payload_dict = decoded_json_object(_DECODED_JSON.validate_python(json.loads(payload_bytes))) texts = _delta_frame_texts(payload_dict) except Exception as e: verbose_proxy_logger.debug( diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 96d1da19bcc..b2b3a9b0faf 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -1,11 +1,12 @@ from collections.abc import Mapping from copy import deepcopy -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, runtime_checkable from urllib.parse import urlparse import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -34,6 +35,32 @@ else: LiteLLMLoggingObj = Any +@runtime_checkable +class _SupportsSetDefault(Protocol): + def setdefault(self, key: str, default: object, /) -> object: ... + + +@runtime_checkable +class _SupportsSetItem(Protocol): + def __setitem__(self, key: str, value: object, /) -> None: ... + + +def _setdefault_target(value: object) -> _SupportsSetDefault: + if isinstance(value, _SupportsSetDefault): + return value + raise AttributeError(f"{type(value).__name__!r} object has no attribute 'setdefault'", name="setdefault", obj=value) + + +def _item_assignment_target(value: object) -> _SupportsSetItem: + if isinstance(value, _SupportsSetItem): + return value + raise TypeError(f"'{type(value).__name__}' object does not support item assignment") + + +def _vector_search_configuration(retrieval_config: object) -> Mapping[str, object]: + return decoded_json_object(decoded_json_object(retrieval_config).get("vectorSearchConfiguration", {})) + + class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): """Vector store configuration for AWS Bedrock Knowledge Bases.""" @@ -206,7 +233,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: if isinstance(query, list): query = " ".join(query) @@ -218,30 +245,33 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): "retrievalQuery": BedrockKBRetrievalQuery(text=query), } - retrieval_config: dict[str, Any] = {} - - if isinstance(extra_body, dict): - retrieval_config = deepcopy( - extra_body.get("retrievalConfiguration") or extra_body.get("retrieval_configuration") or {} - ) + retrieval_config: Final[object] = ( + deepcopy(extra_body.get("retrievalConfiguration") or extra_body.get("retrieval_configuration") or {}) + if isinstance(extra_body, dict) + else {} + ) max_results: Final = vector_store_search_optional_params.get("max_num_results") if max_results is not None: - existing_number_of_results = retrieval_config.get("vectorSearchConfiguration", {}).get("numberOfResults") + existing_number_of_results: Final = _vector_search_configuration(retrieval_config).get("numberOfResults") if existing_number_of_results is not None and existing_number_of_results != max_results: verbose_logger.debug( "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.numberOfResults (%s) with max_num_results=%s", existing_number_of_results, max_results, ) - retrieval_config.setdefault("vectorSearchConfiguration", {})["numberOfResults"] = max_results + _item_assignment_target(_setdefault_target(retrieval_config).setdefault("vectorSearchConfiguration", {}))[ + "numberOfResults" + ] = max_results filters: Final = vector_store_search_optional_params.get("filters") if filters is not None: - existing_filter: Final = retrieval_config.get("vectorSearchConfiguration", {}).get("filter") + existing_filter: Final = _vector_search_configuration(retrieval_config).get("filter") if existing_filter is not None and existing_filter != filters: verbose_logger.debug( "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.filter with filters from vector_store_search_optional_params" ) - retrieval_config.setdefault("vectorSearchConfiguration", {})["filter"] = filters + _item_assignment_target(_setdefault_target(retrieval_config).setdefault("vectorSearchConfiguration", {}))[ + "filter" + ] = filters if retrieval_config: request_body["retrievalConfiguration"] = cast(BedrockKBRetrievalConfiguration, retrieval_config) user_context: Final = self._user_context(extra_body=extra_body, litellm_params=litellm_params) diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 38bcbb0d905..8eb02def465 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -10,7 +10,7 @@ API Reference: https://docs.bfl.ai/ import base64 import time from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol import httpx from httpx._types import RequestFiles @@ -19,6 +19,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.core_helpers import as_object_list, has_attributes from litellm.litellm_core_utils.url_utils import async_safe_get, join_api_base, safe_get from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str @@ -52,6 +53,19 @@ class _BflPolledResult(TypedDict): _BFL_POLLED_RESULT_ADAPTER: Final = TypeAdapter(_BflPolledResult) _DECODED_JSON: Final = TypeAdapter(object) +_UNTYPED_VALUE: Final = TypeAdapter(object) + + +class _BinaryReadable(Protocol): + def read(self) -> bytes: ... + + +class _Tellable(Protocol): + def tell(self) -> int: ... + + +class _Seekable(Protocol): + def seek(self, offset: int, /) -> object: ... _BFL_REQUEST_PARAMS: Final = ( @@ -197,7 +211,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): def _read_image_bytes( self, - image: Any, + image: object, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, ) -> bytes: @@ -208,9 +222,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): ) if isinstance(image, bytes): return image - elif isinstance(image, list): - # If it's a list, take the first image - return self._read_image_bytes(image[0], depth=depth + 1, max_depth=max_depth) + elif (image_list := as_object_list(image)) is not None: + return self._read_image_bytes(image_list[0], depth=depth + 1, max_depth=max_depth) elif isinstance(image, str): if image.startswith(("http://", "https://")): response: Final = safe_get(litellm.module_level_client, image, timeout=60.0) @@ -221,13 +234,12 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): "Unsupported image input: plain string values that are not URLs are not accepted. " "Provide image bytes or a file-like object." ) - elif hasattr(image, "read"): - # File-like object - pos: Final = getattr(image, "tell", lambda: 0)() - if hasattr(image, "seek"): + elif has_attributes(image, _BinaryReadable, "read"): + pos: Final = image.tell() if has_attributes(image, _Tellable, "tell") else 0 + if has_attributes(image, _Seekable, "seek"): image.seek(0) data: Final = image.read() - if hasattr(image, "seek"): + if has_attributes(image, _Seekable, "seek"): image.seek(pos) return data else: @@ -265,7 +277,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): # Handle mask if provided (for inpainting) if "mask" in image_edit_optional_request_params: - mask: Final = image_edit_optional_request_params["mask"] + mask: Final = _UNTYPED_VALUE.validate_python(image_edit_optional_request_params["mask"]) mask_bytes: Final = self._read_image_bytes(mask) request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8") diff --git a/litellm/llms/bytez/chat/transformation.py b/litellm/llms/bytez/chat/transformation.py index c297eb13db1..465f505c729 100644 --- a/litellm/llms/bytez/chat/transformation.py +++ b/litellm/llms/bytez/chat/transformation.py @@ -6,7 +6,9 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.url_utils import encode_url_path_segments @@ -36,6 +38,8 @@ else: # 5 minute timeout (models may need to load) STREAMING_TIMEOUT: Final = 60 * 5 +_DECODED_JSON: Final = TypeAdapter(object) + class BytezChatConfig(BaseConfig): """ @@ -194,11 +198,11 @@ class BytezChatConfig(BaseConfig): ) -> ModelResponse: json: Final = raw_response.json() - error: Final = json.get("error") + error: Final = decoded_json_object(_DECODED_JSON.validate_python(json)).get("error") if error is not None: raise BytezError( - message=str(json["error"]), + message=str(error), status_code=raw_response.status_code, ) diff --git a/litellm/llms/cohere/common_utils.py b/litellm/llms/cohere/common_utils.py index 65531451d52..4ba5f9bd57b 100644 --- a/litellm/llms/cohere/common_utils.py +++ b/litellm/llms/cohere/common_utils.py @@ -1,6 +1,8 @@ import json from typing import Final, Literal +from pydantic import TypeAdapter + from litellm.llms.base_llm.base_utils import BaseLLMModelInfo from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.cohere import CohereV2StreamedToolCall @@ -12,6 +14,8 @@ from litellm.types.utils import ( ProviderSpecificModelInfo, ) +_DECODED_JSON: Final = TypeAdapter(object) + class CohereError(BaseLLMException): def __init__(self, status_code, message): @@ -190,8 +194,10 @@ class ModelResponseIterator: if index != -1: str_line = str_line[index:] - data_json: Final = json.loads(str_line) - return self.chunk_parser(chunk=data_json) + data_json: Final = _DECODED_JSON.validate_python(json.loads(str_line)) + if isinstance(data_json, dict): + return self.chunk_parser(chunk=data_json) + raise AttributeError(f"{type(data_json).__name__!r} object has no attribute 'get'", name="get", obj=data_json) # Async iterator def __aiter__(self): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index df82891f789..8c9de272c0d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -153,7 +153,7 @@ from litellm.types.containers.main import ( ) from litellm.types.decisions import DecisionsIRRequest, DecisionsIRResponse from litellm.types.extraction import ExtractionResponse -from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig +from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig, TwoStepFileUploadRequest from litellm.types.integrations.custom_logger import ( NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, AgenticLoopLimitError, @@ -188,6 +188,7 @@ from litellm.types.responses.main import DeleteResponseResult from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + AgenticLoopParams, CallTypes, EmbeddingResponse, FileTypes, @@ -239,6 +240,10 @@ def _client_request_headers( return _merge_headers_case_insensitive(optional_headers or {}, extra_headers) +def _requested_stream(response_api_optional_request_params: Mapping[str, object]) -> object: + return response_api_optional_request_params.get("stream", False) + + def _rust_responses_websocket_enabled( custom_llm_provider: str | None, ) -> bool: @@ -386,12 +391,19 @@ class _SignedBodyKwargs(TypedDict, total=False): json: ReadOnly[dict[str, object]] +class _AgenticLoopParamsView(TypedDict): + agentic_loop_params: ReadOnly[AgenticLoopParams] + + _GOAWAY_RETRY_MESSAGE: Final = "Vector store search connection was closed by an HTTP/2 GOAWAY, retrying once" def _closed_by_http2_goaway(error: httpx.RemoteProtocolError) -> bool: cause: Final = error.__cause__ - return cause is not None and any(isinstance(arg, ConnectionTerminated) for arg in cause.args) + if cause is None: + return False + cause_args: Final[tuple[object, ...]] = cause.args + return any(isinstance(arg, ConnectionTerminated) for arg in cause_args) async def _asend_with_goaway_retry(send: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response: @@ -433,8 +445,12 @@ def _signed_body_kwargs(*, signed_body: bytes | None, data: dict[str, object]) - return {"json": data} +def _response_content_type(headers: Mapping[str, str]) -> str: + return headers.get("content-type", "") + + def _answered_json(response: httpx.Response) -> bool: - return "application/json" in response.headers.get("content-type", "").lower() + return "application/json" in _response_content_type(response.headers).lower() def _stream_query_params(stream: bool, starting_after: int | None) -> Mapping[str, object]: @@ -3064,7 +3080,7 @@ class BaseLLMHTTPHandler: } # Check if streaming is requested - stream = response_api_optional_request_params.get("stream", False) + stream = _requested_stream(response_api_optional_request_params) caller_requested_stream: Final = bool(stream or (extra_body or {}).get("stream")) api_base: Final = responses_api_provider_config.get_complete_url( @@ -4136,7 +4152,7 @@ class BaseLLMHTTPHandler: ): # Handle pre-signed requests (e.g., from Bedrock S3 uploads) # Type narrowing: this is a plain dict, not TwoStepFileUploadConfig - presigned_request: Final = cast(dict[str, Any], transformed_request) + presigned_request: Final = cast(TwoStepFileUploadRequest, transformed_request) upload_response = getattr(sync_httpx_client, presigned_request["method"].lower())( url=presigned_request["url"], headers=presigned_request["headers"], @@ -4350,7 +4366,7 @@ class BaseLLMHTTPHandler: ): # Handle pre-signed requests (e.g., from Bedrock S3 uploads) # Type narrowing: this is a plain dict, not TwoStepFileUploadConfig - presigned_request: Final = cast(dict[str, Any], transformed_request) + presigned_request: Final = cast(TwoStepFileUploadRequest, transformed_request) upload_response = await getattr(async_httpx_client, presigned_request["method"].lower())( url=presigned_request["url"], headers=presigned_request["headers"], @@ -6129,8 +6145,10 @@ class BaseLLMHTTPHandler: full_model_name = model if logging_obj is not None: - agentic_params: Final[Mapping[str, object]] = logging_obj.model_call_details.get("agentic_loop_params", {}) - full_model_name = cast(str, agentic_params.get("model", model)) + agentic_view: Final[_AgenticLoopParamsView] = { + "agentic_loop_params": logging_obj.model_call_details.get("agentic_loop_params", {}) + } + full_model_name = agentic_view["agentic_loop_params"].get("model", model) optional_params: Final = dict(anthropic_messages_optional_request_params) optional_params.update(patch.optional_params) @@ -11116,7 +11134,7 @@ class BaseLLMHTTPHandler: extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, - ): + ) -> Mapping[str, object]: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( llm_provider=litellm.LlmProviders(custom_llm_provider), @@ -11407,7 +11425,7 @@ class BaseLLMHTTPHandler: extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, - ): + ) -> Mapping[str, object]: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( llm_provider=litellm.LlmProviders(custom_llm_provider), diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index 0f01a4b3172..e1efa075dc8 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -11,12 +11,13 @@ Endpoint Docs - https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api """ -from typing import Final +from typing import Final, Protocol, runtime_checkable import httpx from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict +from litellm.litellm_core_utils.core_helpers import json_contains from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -29,6 +30,19 @@ from ..common_utils import DashScopeError, resolve_dashscope_openai_api_base DEFAULT_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-mode/v1" +_DECODED_JSON: Final = TypeAdapter(object) + + +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, key: str, /) -> object: ... + + +def _subscript(value: object, key: str) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + @with_config(ConfigDict(extra="allow", strict=True)) class _DashScopeUsage(TypedDict, total=False): @@ -149,7 +163,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, ) -> EmbeddingResponse: try: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise DashScopeError( status_code=raw_response.status_code, @@ -163,8 +177,8 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): original_response=response_json, ) - if "error" in response_json: - error: Final[object] = response_json["error"] + if json_contains(response_json, "error"): + error: Final = _subscript(response_json, "error") message: Final = error.get("message", str(error)) if isinstance(error, dict) else str(error) raise DashScopeError( status_code=raw_response.status_code, diff --git a/litellm/llms/dataforseo/search/transformation.py b/litellm/llms/dataforseo/search/transformation.py index 5cf7735325d..b88f7df0f91 100644 --- a/litellm/llms/dataforseo/search/transformation.py +++ b/litellm/llms/dataforseo/search/transformation.py @@ -4,13 +4,14 @@ Calls DataForSEO SERP API to search the web. DataForSEO API Reference: https://docs.dataforseo.com/v3/serp/google/organic/live/advanced/?bash """ -from collections.abc import Iterable, Mapping -from typing import Final, Literal +from collections.abc import Iterable, Mapping, Sized +from typing import Final, Literal, Protocol, runtime_checkable import httpx from pydantic import ConfigDict, TypeAdapter from litellm.constants import DEFAULT_DATAFORSEO_LOCATION_CODE +from litellm.litellm_core_utils.core_helpers import decoded_json_object, json_contains from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, @@ -21,6 +22,24 @@ from litellm.secret_managers.main import get_secret_str _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) + + +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, key: str | int, /) -> object: ... + + +def _subscript(value: object, key: str | int) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + +def _json_len(value: object) -> int: + if isinstance(value, Sized): + return len(value) + raise TypeError(f"object of type '{type(value).__name__}' has no len()") class DataForSEOSearchConfig(BaseSearchConfig): @@ -193,20 +212,20 @@ class DataForSEOSearchConfig(BaseSearchConfig): Returns: SearchResponse with standardized format """ - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) # Transform results to SearchResult objects results: Final = [] # DataForSEO wraps results in tasks array - if "tasks" in response_json and len(response_json["tasks"]) > 0: - task: Final = response_json["tasks"][0] + if json_contains(response_json, "tasks") and _json_len(_subscript(response_json, "tasks")) > 0: + task: Final = decoded_json_object(_subscript(_subscript(response_json, "tasks"), 0)) # Check if task was successful if task.get("status_code") == 20000 and "result" in task: # Result is an array, take first element - if len(task["result"]) > 0: - result: Final = _JSON_OBJECT.validate_python(task["result"][0]) + if _json_len(task["result"]) > 0: + result: Final = _JSON_OBJECT.validate_python(_subscript(task["result"], 0)) # Items contain the actual search results for raw_item in _JSON_VALUES.validate_python(result.get("items", [])): diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 74cd6f22704..b95219348e4 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -4,10 +4,11 @@ import json import math from collections.abc import Iterable, Mapping, Sequence from types import MappingProxyType -from typing import Any, Final, cast +from typing import Final, cast import httpx from pydantic import JsonValue, TypeAdapter +from typing_extensions import TypedDict as ExtraItemsTypedDict import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -510,6 +511,13 @@ def get_api_key_from_env() -> str | None: return get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY") +_COUNT_TOKENS_REQUEST_FIELDS: Final = frozenset({"model", "contents", "system_instruction", "tools"}) + + +class _ForwardedLitellmParams(ExtraItemsTypedDict, total=False, extra_items=object): + pass + + class GoogleAIStudioTokenCounter(BaseTokenCounter): """Token counter implementation for Google AI Studio provider.""" @@ -526,7 +534,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): model_to_use: str, messages: list[dict[str, object]] | None, contents: list[dict[str, object]] | None, - deployment: dict[str, Any] | None = None, + deployment: dict[str, object] | None = None, request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, @@ -537,6 +545,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): ) -> TokenCountResponse | None: import copy + from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.anthropic.messages.context_management.inference_view import ( InferenceRequestViewError, count_context_edits, @@ -553,11 +562,12 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): vertex_gemini_count_context, ) + litellm_params: Final = decoded_json_object((deployment or {}).get("litellm_params", {})) request_view: Final = await inference_request_view( messages or [], system, messages_format, - count_context_edits(context_management, tools, (deployment or {}).get("litellm_params", {}), model_to_use), + count_context_edits(context_management, tools, litellm_params, model_to_use), ) if isinstance(request_view, InferenceRequestViewError): return count_context_error_response(request_view, request_model, model_to_use) @@ -570,8 +580,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): if blocks_error is not None: return count_context_error_response(blocks_error, request_model, model_to_use) - deployment = deployment or {} - count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) + count_tokens_params_request: Final = copy.deepcopy(litellm_params) resolved_contents: Final = ( contents if contents is not None @@ -581,16 +590,15 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): ) if isinstance(resolved_contents, VertexGeminiCountContextError): return count_context_error_response(resolved_contents, request_model, model_to_use) - count_tokens_params: Final = { - "model": model_to_use, - "contents": resolved_contents, - "system_instruction": count_context.system_instruction, - "tools": count_context.tools, - } - count_tokens_params_request.update(count_tokens_params) try: result: Final = await GoogleAIStudioTokenCounter().acount_tokens( - **count_tokens_params_request, + **_ForwardedLitellmParams( + **{k: v for k, v in count_tokens_params_request.items() if k not in _COUNT_TOKENS_REQUEST_FIELDS} + ), + model=model_to_use, + contents=resolved_contents, + system_instruction=count_context.system_instruction, + tools=count_context.tools, ) except litellm.APIError as error: return provider_error_response(error, request_model, model_to_use) diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py index 4ed47323d4f..2c78d04eea5 100644 --- a/litellm/llms/gemini/count_tokens/handler.py +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -1,11 +1,12 @@ -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable import httpx from pydantic import ConfigDict, RootModel, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm +from litellm.litellm_core_utils.core_helpers import has_attributes, iter_items, json_contains from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.types.utils import LlmProviders @@ -26,8 +27,44 @@ class _GeminiCountTokensResponseRoot(RootModel[_GeminiCountTokensResponse]): model_config = ConfigDict(frozen=True, hide_input_in_errors=True) +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, key: str, /) -> object: ... + + +@runtime_checkable +class _SupportsSetItem(Protocol): + def __setitem__(self, key: str, value: object, /) -> None: ... + + +class _KeywordArguments(Protocol): + def keys(self) -> Iterable[str]: ... + + def __getitem__(self, key: str, /) -> object: ... + + +def _subscript(value: object, key: str) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + +def _item_assignment_target(value: object) -> _SupportsSetItem: + if isinstance(value, _SupportsSetItem): + return value + raise TypeError(f"'{type(value).__name__}' object does not support item assignment") + + +def _function_response_kwargs(value: object) -> _KeywordArguments: + if has_attributes(value, _KeywordArguments, "keys"): + return value + raise TypeError( + f"google.genai.types.FunctionResponse() argument after ** must be a mapping, not {type(value).__name__}" + ) + + class GoogleAIStudioTokenCounter: - def _clean_contents_for_gemini_api(self, contents: Any) -> object: + def _clean_contents_for_gemini_api(self, contents: Sequence[Mapping[str, object]] | None) -> object: """ Clean up contents to remove unsupported fields for the Gemini API. @@ -52,12 +89,14 @@ class GoogleAIStudioTokenCounter: for content in cleaned_contents: parts = content["parts"] - for part in parts: - if "functionResponse" in part: - function_response_data = part["functionResponse"] - function_response_part = FunctionResponse(**function_response_data) + for part in iter_items(parts): + if json_contains(part, "functionResponse"): + function_response_data = _subscript(part, "functionResponse") + function_response_part = FunctionResponse(**_function_response_kwargs(function_response_data)) function_response_part.id = None - part["functionResponse"] = function_response_part.model_dump(exclude_none=True) + _item_assignment_target(part)["functionResponse"] = function_response_part.model_dump( + exclude_none=True + ) return cleaned_contents diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index 50da73cc8b9..a02afa3a13f 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -9,13 +9,14 @@ Per OpenAPI spec (https://ai.google.dev/static/api/interactions.openapi.json): Requests use Api-Revision 2026-05-20 (`steps` schema). """ +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias import httpx from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, process_response_headers +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, has_attributes, process_response_headers from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo @@ -83,6 +84,18 @@ class _CancelPayloadSource(Protocol): def json(self) -> _CancelPayload: ... +class _ImageConfigItems(Protocol): + def keys(self) -> Iterable[str]: ... + + def __getitem__(self, key: str, /) -> object: ... + + +def _image_config_items(image_config: object) -> _ImageConfigItems: + if has_attributes(image_config, _ImageConfigItems, "keys"): + return image_config + raise TypeError(f"'{type(image_config).__name__}' object is not a mapping") + + def _interaction_body(response: _InteractionPayloadSource) -> _InteractionPayload: """Decode the body of an Interactions API interaction response.""" return response.json() @@ -241,7 +254,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): request_body["response_format"] = response_format # image_config moves out of generation_config into response_format. - generation_config: dict[str, Any] | None = optional_params.get("generation_config") + generation_config: Mapping[str, object] | None = optional_params.get("generation_config") if generation_config is not None: image_config = None if isinstance(generation_config, dict): @@ -255,7 +268,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): if image_config is not None: # Move image_config to response_format with type=image. - image_rf: Final[_JsonObject] = {"type": "image", **image_config} + image_rf: Final[_JsonObject] = {"type": "image", **_image_config_items(image_config)} existing_rf: Final = request_body.get("response_format") if existing_rf is None: request_body["response_format"] = image_rf diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 5324aa6f94e..08b5de10d41 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -4,9 +4,12 @@ GigaChat Streaming Response Handler import json import uuid -from collections.abc import Mapping, Sequence -from typing import Any, Final +from collections.abc import Awaitable, Mapping, Sequence +from typing import Final, Protocol +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.llms.gigachat.utils import convert_usage from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, @@ -14,13 +17,27 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import ChatCompletionUsageBlock, GenericStreamingChunk +_DECODED_JSON: Final = TypeAdapter(object) + + +class _SyncChunkSource(Protocol): + def __next__(self) -> object: ... + + +class _AsyncChunkSource(Protocol): + def __anext__(self) -> Awaitable[object]: ... + + +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + class GigaChatModelResponseIterator: """Iterator for GigaChat streaming responses.""" def __init__( self, - streaming_response: Any, + streaming_response: object, sync_stream: bool, json_mode: bool | None = False, ): @@ -28,7 +45,7 @@ class GigaChatModelResponseIterator: self.response_iterator = self.streaming_response self.json_mode = json_mode - def chunk_parser(self, chunk: Mapping[str, object]) -> GenericStreamingChunk: + def chunk_parser(self, chunk: _SupportsGet) -> GenericStreamingChunk: """Parse a single streaming chunk from GigaChat.""" choices: Sequence = chunk.get("choices") or () if not choices: @@ -106,14 +123,14 @@ class GigaChatModelResponseIterator: def __next__(self) -> GenericStreamingChunk: try: - chunk = self.response_iterator.__next__() + chunk: Final = require_attributes(self.response_iterator, _SyncChunkSource, "__next__").__next__() if isinstance(chunk, str): # Parse SSE format: data: {...} - chunk = chunk.removeprefix("data: ") - if chunk.strip() == "[DONE]": + payload: Final = chunk.removeprefix("data: ") + if payload.strip() == "[DONE]": raise StopIteration try: - chunk = json.loads(chunk) + decoded: Final = _DECODED_JSON.validate_python(json.loads(payload)) except json.JSONDecodeError: return GenericStreamingChunk( text="", @@ -123,7 +140,8 @@ class GigaChatModelResponseIterator: usage=None, index=0, ) - return self.chunk_parser(chunk) + return self.chunk_parser(require_attributes(decoded, _SupportsGet, "get")) + return self.chunk_parser(require_attributes(chunk, _SupportsGet, "get")) except StopIteration: raise @@ -132,14 +150,14 @@ class GigaChatModelResponseIterator: async def __anext__(self) -> GenericStreamingChunk: try: - chunk = await self.response_iterator.__anext__() + chunk: Final = await require_attributes(self.response_iterator, _AsyncChunkSource, "__anext__").__anext__() if isinstance(chunk, str): # Parse SSE format - chunk = chunk.removeprefix("data: ") - if chunk.strip() == "[DONE]": + payload: Final = chunk.removeprefix("data: ") + if payload.strip() == "[DONE]": raise StopAsyncIteration try: - chunk = json.loads(chunk) + decoded: Final = _DECODED_JSON.validate_python(json.loads(payload)) except json.JSONDecodeError: return GenericStreamingChunk( text="", @@ -149,6 +167,7 @@ class GigaChatModelResponseIterator: usage=None, index=0, ) - return self.chunk_parser(chunk) + return self.chunk_parser(require_attributes(decoded, _SupportsGet, "get")) + return self.chunk_parser(require_attributes(chunk, _SupportsGet, "get")) except StopAsyncIteration: raise diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 5d997babc6e..9174b77a9f0 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -1,11 +1,13 @@ import json import os -from typing import Final, Literal, Protocol, get_args +from collections.abc import Iterator +from typing import Final, Literal, Protocol, get_args, runtime_checkable import httpx from pydantic import TypeAdapter import litellm +from litellm.litellm_core_utils.core_helpers import as_object_list, iter_items, json_contains from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -38,6 +40,28 @@ class _SupportsTokenEncode(Protocol): def encode(self, text: str) -> list[int]: ... +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, key: str | int, /) -> object: ... + + +def _subscript(value: object, key: str | int) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + +def _unpack_pair(value: object) -> tuple[object, object]: + if not hasattr(value, "__iter__") and not hasattr(value, "__getitem__"): + raise TypeError(f"cannot unpack non-iterable {type(value).__name__} object") + first, second = iter_items(value) + return first, second + + +def _pairs(value: object) -> Iterator[tuple[object, object]]: + return (_unpack_pair(item) for item in iter_items(value)) + + def get_hf_task_embedding_for_model( model: str, task_type: str | None, api_base: str, timeout: float | httpx.Timeout | None = None ) -> object: @@ -183,15 +207,15 @@ class HuggingFaceEmbedding(BaseLLM): def _process_embedding_response( self, - embeddings: dict, + embeddings: object, model_response: EmbeddingResponse, model: str, input: list, encoding: _SupportsTokenEncode, ) -> EmbeddingResponse: output_data: Final = [] - if "similarities" in embeddings: - for idx, embedding in embeddings["similarities"]: + if json_contains(embeddings, "similarities"): + for idx, embedding in _pairs(_subscript(embeddings, "similarities")): output_data.append( { "object": "embedding", @@ -200,8 +224,9 @@ class HuggingFaceEmbedding(BaseLLM): } ) else: - for idx, embedding in enumerate(embeddings): - if isinstance(embedding, float) or isinstance(embedding, list) and isinstance(embedding[0], float): + for idx, embedding in enumerate(iter_items(embeddings)): + embedding_list = as_object_list(embedding) + if isinstance(embedding, float) or embedding_list is not None and isinstance(embedding_list[0], float): output_data.append( { "object": "embedding", @@ -214,7 +239,7 @@ class HuggingFaceEmbedding(BaseLLM): { "object": "embedding", "index": idx, - "embedding": embedding[0][0], # flatten list returned from hf + "embedding": _subscript(_subscript(embedding, 0), 0), # flatten list returned from hf } ) model_response.object = "list" @@ -288,10 +313,10 @@ class HuggingFaceEmbedding(BaseLLM): original_response=response, ) - embeddings: Final = response.json() + embeddings: Final = _DECODED_JSON.validate_python(response.json()) - if "error" in embeddings: - raise HuggingFaceError(status_code=500, message=embeddings["error"]) + if json_contains(embeddings, "error"): + raise HuggingFaceError(status_code=500, message=_subscript(embeddings, "error")) ## PROCESS RESPONSE ## return self._process_embedding_response( @@ -394,10 +419,10 @@ class HuggingFaceEmbedding(BaseLLM): original_response=response, ) - embeddings: Final = response.json() + embeddings: Final = _DECODED_JSON.validate_python(response.json()) - if "error" in embeddings: - raise HuggingFaceError(status_code=500, message=embeddings["error"]) + if json_contains(embeddings, "error"): + raise HuggingFaceError(status_code=500, message=_subscript(embeddings, "error")) ## PROCESS RESPONSE ## return self._process_embedding_response( diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 6e6a117227e..bb836eb1b35 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -365,7 +365,10 @@ class HuggingFaceEmbeddingConfig(BaseConfig): ) content: str = "" for chunk in streamed_response: - content += chunk["choices"][0]["delta"]["content"] + delta_content = chunk.choices[0].delta.content + if delta_content is None: + raise TypeError('can only concatenate str (not "NoneType") to str') + content += delta_content completion_response: Final[list[dict[str, str]]] = [{"generated_text": content}] ## LOGGING logging_obj.post_call( diff --git a/litellm/llms/litellm_proxy/skills/code_execution.py b/litellm/llms/litellm_proxy/skills/code_execution.py index fbc287589b3..af26b4aa095 100644 --- a/litellm/llms/litellm_proxy/skills/code_execution.py +++ b/litellm/llms/litellm_proxy/skills/code_execution.py @@ -15,11 +15,12 @@ import base64 import json from collections.abc import Mapping, Sequence from enum import Enum -from typing import Any, Final, Protocol, TypedDict +from typing import Final, Protocol, TypedDict from typing_extensions import ReadOnly from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import require_attributes class _ToolParameterSchema(TypedDict, total=False): @@ -237,17 +238,21 @@ class CodeExecutionHandler: execution_results: Final[list[_ExecutionResult]] = [] executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) - response: Any = None # Initialize to avoid possibly unbound error + response: _CompletionResponse | None = None # Initialize to avoid possibly unbound error for iteration in range(self.max_iterations): verbose_logger.debug("CodeExecutionHandler: Iteration %s/%s", iteration + 1, self.max_iterations) # Make LLM call - response = await litellm.acompletion( - model=model, - messages=current_messages, - tools=tools, - **kwargs, + response = require_attributes( + await litellm.acompletion( + model=model, + messages=current_messages, + tools=tools, + **kwargs, + ), + _CompletionResponse, + "choices", ) choice: _ResponseChoice = response.choices[0] diff --git a/litellm/llms/litellm_proxy/skills/transformation.py b/litellm/llms/litellm_proxy/skills/transformation.py index 73e26606497..3b488784443 100644 --- a/litellm/llms/litellm_proxy/skills/transformation.py +++ b/litellm/llms/litellm_proxy/skills/transformation.py @@ -10,8 +10,6 @@ Pattern follows litellm/llms/litellm_proxy/responses/transformation.py from collections.abc import Coroutine, Sequence from typing import TYPE_CHECKING, Final, Optional -from pydantic import JsonValue - from litellm.types.llms.anthropic_skills import ( DeleteSkillResponse, ListSkillsResponse, @@ -47,7 +45,7 @@ class LiteLLMSkillsTransformationHandler: file_content: bytes | None = None, file_name: str | None = None, file_type: str | None = None, - metadata: dict[str, JsonValue] | None = None, + metadata: dict[str, object] | None = None, user_id: str | None = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, _is_async: bool = False, @@ -130,7 +128,7 @@ class LiteLLMSkillsTransformationHandler: file_content: bytes | None = None, file_name: str | None = None, file_type: str | None = None, - metadata: dict[str, JsonValue] | None = None, + metadata: dict[str, object] | None = None, user_id: str | None = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, ) -> Skill: diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 14470922e57..fb707a39362 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -252,10 +252,11 @@ class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): # Build attributes with all available metadata # Exclude system fields and already-processed fields - attributes = {} - for key, value in result.items(): - if key not in ["id", "content", "distance", text_field]: - attributes[key] = value + attributes = { + key: value + for key, value in decoded_json_object(_DECODED_JSON.validate_python(result)).items() + if key not in ["id", "content", "distance", text_field] + } result_obj = VectorStoreSearchResult( score=score, diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 61d237de2d8..d13e46658ae 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -21,6 +21,7 @@ from litellm._logging import verbose_logger from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, strip_litellm_internal_message_fields, @@ -82,6 +83,7 @@ _MISTRAL_REQUEST_ROUTING: Final = TypeAdapter(MistralRequestRouting) _MISTRAL_CACHE_KEY_PARAMS: Final = TypeAdapter(MistralCacheKeyParams) _MISTRAL_REASONING_EFFORT_PARAMS: Final = TypeAdapter(MistralReasoningEffortParams) _MISTRAL_FILE_PART: Final = TypeAdapter(MistralFilePart) +_DECODED_JSON: Final = TypeAdapter(object) _REASONING_PROMPT_REASON: Final = "litellm sends Magistral a reasoning system prompt instead of the requested level." _NON_LEVEL_EFFORT_REASON: Final = ( "the Mistral API takes it only as a level such as 'high', so Mistral applies its default effort." @@ -1005,9 +1007,13 @@ class MistralConfig(OpenAIGPTConfig): logging_obj.model_call_details["response_headers"] = raw_response.headers # Handle Mistral-specific response transformations - response_data = raw_response.json() - response_data = self._handle_empty_content_response(response_data) - response_data = self._handle_content_list_to_str_conversion(response_data) + decoded_response: Final = _DECODED_JSON.validate_python(raw_response.json()) + if not is_str_keyed_dict(decoded_response): + missing_get: Final = f"{type(decoded_response).__name__!r} object has no attribute 'get'" + raise AttributeError(missing_get, name="get", obj=decoded_response) + response_data: Final = self._handle_content_list_to_str_conversion( + self._handle_empty_content_response(decoded_response) + ) final_response_obj: Final = cast( ModelResponse, diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 9163ddfa5de..846f5f27e0d 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -286,7 +286,7 @@ def handle_cohere_response( def handle_cohere_stream_chunk( - dict_chunk: Mapping[str, JsonValue], + dict_chunk: Mapping[str, object], prior_tool_calls_emitted: bool = False, prior_text_emitted: bool = False, ) -> ModelResponseStream: diff --git a/litellm/llms/oci/chat/generic.py b/litellm/llms/oci/chat/generic.py index 13cea58495d..5cf4c7aa535 100644 --- a/litellm/llms/oci/chat/generic.py +++ b/litellm/llms/oci/chat/generic.py @@ -8,6 +8,7 @@ parsing, and streaming chunk parsing for models served with import datetime import hashlib +from collections.abc import Mapping from typing import Final import httpx @@ -363,7 +364,7 @@ def handle_generic_response( return model_response -def handle_generic_stream_chunk(dict_chunk: dict) -> ModelResponseStream: +def handle_generic_stream_chunk(dict_chunk: Mapping) -> ModelResponseStream: """Parse a single GENERIC SSE chunk into a LiteLLM ModelResponseStream.""" # OCI streams tool calls progressively — early chunks may omit required fields. if dict_chunk.get("message") and dict_chunk["message"].get("toolCalls"): diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index f97f8f4ca59..256d0a74623 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -19,6 +19,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger from litellm.constants import DEFAULT_OCI_CHAT_MAX_TOKENS +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.prompt_templates.common_utils import ( merge_consecutive_system_messages, @@ -891,13 +892,14 @@ class OCIStreamWrapper(CustomStreamWrapper): if payload == "[DONE]": return None try: - dict_chunk: Final = json.loads(payload) + decoded_chunk: Final = _DECODED_JSON.validate_python(json.loads(payload)) except json.JSONDecodeError as e: raise OCIError( status_code=500, message=f"Chunk cannot be parsed as JSON: {e}", ) + dict_chunk: Final = decoded_json_object(decoded_chunk) if dict_chunk.get("apiFormat") == "COHERE": result: Final = handle_cohere_stream_chunk( dict_chunk, diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index fcf6a396a13..d6b9ff46410 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -1,6 +1,7 @@ import base64 import io -from typing import Final, overload +from collections.abc import Iterable +from typing import Final, Protocol, overload import httpx from openai.types.chat import ChatCompletionChunk @@ -8,6 +9,7 @@ from pydantic import BaseModel, TypeAdapter from litellm import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -19,6 +21,14 @@ _THINK_OFF_REASON: Final = "Ollama takes think only as on or off, and litellm ma _DECODED_JSON: Final = TypeAdapter(object) +class _JsonLookup(Protocol): + def get(self, key: str, default: object, /) -> object: ... + + +class _JsonItems(Protocol): + def items(self) -> Iterable[tuple[str, object]]: ... + + class OllamaResponseIdNormalizer: def __init__(self) -> None: self.response_id: str | None = None @@ -206,15 +216,15 @@ class OllamaModelInfo(BaseLLMModelInfo): return any(name.lower() in model_cost_keys for name in potential_model_names) @staticmethod - def _supports_function_calling(ollama_model_info: dict) -> bool: - _template: Final[str] = str(ollama_model_info.get("template", "") or "") + def _supports_function_calling(ollama_model_info: object) -> bool: + _template: Final[str] = str(require_attributes(ollama_model_info, _JsonLookup, "get").get("template", "") or "") return "tools" in _template.lower() @staticmethod - def _get_max_tokens(ollama_model_info: dict) -> int | None: - _model_info: Final[dict] = ollama_model_info.get("model_info", {}) + def _get_max_tokens(ollama_model_info: object) -> object: + _model_info: Final = require_attributes(ollama_model_info, _JsonLookup, "get").get("model_info", {}) - for key, value in _model_info.items(): + for key, value in require_attributes(_model_info, _JsonItems, "items").items(): if "context_length" in key: return value return None @@ -253,7 +263,7 @@ class OllamaModelInfo(BaseLLMModelInfo): "max_output_tokens": None, } - model_info: Final = response.json() + model_info: Final = _DECODED_JSON.validate_python(response.json()) max_tokens: Final = self._get_max_tokens(model_info) return { diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 1e70141f939..e1fc9a8d262 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -6,12 +6,13 @@ Ollama /chat/completion calls handled in llm_http_handler.py from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Any, Final, Protocol, TypedDict +from typing import Final, Protocol, TypedDict import httpx from typing_extensions import NotRequired, ReadOnly import litellm +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.types.mcp import has_header from litellm.types.utils import EmbeddingResponse @@ -23,6 +24,14 @@ class TokenEncoder(Protocol): def encode(self, text: str, /) -> Sequence[int]: ... +class _DebugLogger(Protocol): + def debug(self, message: str, /) -> object: ... + + +class _WarningLogger(Protocol): + def warning(self, message: str, /) -> object: ... + + class OllamaEmbeddingResponse(TypedDict): """Body of an Ollama ``/api/embed`` response.""" @@ -51,7 +60,7 @@ def _process_ollama_embedding_response( prompts: list[str], model: str, model_response: EmbeddingResponse, - logging_obj: Any, + logging_obj: object, encoding: TokenEncoder | None, ) -> EmbeddingResponse: output_data: Final = [] @@ -66,11 +75,15 @@ def _process_ollama_embedding_response( if encoding is not None: input_tokens = len(encoding.encode("".join(prompts))) if logging_obj: - logging_obj.debug("Ollama response missing prompt_eval_count; estimated with encoding.") + require_attributes(logging_obj, _DebugLogger, "debug").debug( + "Ollama response missing prompt_eval_count; estimated with encoding." + ) else: input_tokens = 0 if logging_obj: - logging_obj.warning("Missing prompt_eval_count and no encoding provided; defaulted to 0.") + require_attributes(logging_obj, _WarningLogger, "warning").warning( + "Missing prompt_eval_count and no encoding provided; defaulted to 0." + ) model_response.object = "list" model_response.data = output_data diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 739966c95d1..2cf48a75524 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -9,6 +9,7 @@ from collections.abc import Mapping, Sequence from typing import Final import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger @@ -18,6 +19,8 @@ from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ) +_DECODED_JSON: Final = TypeAdapter(object) + class OpenAICountTokensHandler(OpenAICountTokensConfig): """ @@ -34,7 +37,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, text: Mapping[str, object] | None = None, - ) -> Mapping[str, object]: + ) -> object: """ Handle a token counting request to OpenAI's Responses API. @@ -67,7 +70,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): api_key: str, api_base: str | None = None, timeout: float | httpx.Timeout | None = None, - ) -> Mapping[str, object]: + ) -> object: try: endpoint_url: Final = self.get_openai_count_tokens_endpoint(api_base) @@ -96,7 +99,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): message=error_text, ) - openai_response: Final = response.json() + openai_response: Final = _DECODED_JSON.validate_python(response.json()) verbose_logger.debug("OpenAI response: %s", openai_response) return openai_response diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 87a75dae95e..c82a2982a54 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -272,7 +272,7 @@ class OpenAITokenCounter(BaseTokenCounter): @staticmethod async def _counted( resolved: _CountCredential, - send: Callable[[str, str | None], Awaitable[Mapping[str, object]]], + send: Callable[[str, str | None], Awaitable[object]], model_to_use: str, request_model: str, ) -> TokenCountResponse | None: diff --git a/litellm/llms/recraft/image_edit/transformation.py b/litellm/llms/recraft/image_edit/transformation.py index 74bfa52c5e0..45b02a84379 100644 --- a/litellm/llms/recraft/image_edit/transformation.py +++ b/litellm/llms/recraft/image_edit/transformation.py @@ -1,10 +1,12 @@ from io import BufferedReader -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, runtime_checkable import httpx from httpx._types import RequestFiles +from pydantic import TypeAdapter from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str @@ -21,6 +23,24 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + + +@runtime_checkable +class _SupportsStrKeyLookup(Protocol): + def __getitem__(self, key: str, /) -> object: ... + + +def _json_field(value: object, key: str) -> object: + if isinstance(value, _SupportsStrKeyLookup): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + +def _image_object(image_data: object) -> ImageObject: + fields: Final = decoded_json_object(image_data) + return ImageObject(url=fields.get("url", None), b64_json=fields.get("b64_json", None)) + class RecraftImageEditConfig(BaseImageEditConfig): DEFAULT_BASE_URL: str = "https://external.api.recraft.ai" @@ -152,7 +172,7 @@ class RecraftImageEditConfig(BaseImageEditConfig): ) -> ImageResponse: model_response: Final = ImageResponse() try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image edit response: {e}", @@ -162,12 +182,7 @@ class RecraftImageEditConfig(BaseImageEditConfig): if not model_response.data: model_response.data = [] - for image_data in response_data["data"]: - model_response.data.append( - ImageObject( - url=image_data.get("url", None), - b64_json=image_data.get("b64_json", None), - ) - ) + for image_data in iter_items(_json_field(response_data, "data")): + model_response.data.append(_image_object(image_data)) return model_response diff --git a/litellm/llms/sap/credentials.py b/litellm/llms/sap/credentials.py index 74d065956cc..38a7cfedd34 100644 --- a/litellm/llms/sap/credentials.py +++ b/litellm/llms/sap/credentials.py @@ -3,7 +3,7 @@ from __future__ import annotations import json import os import tempfile -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone from pathlib import Path @@ -16,6 +16,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object, iter_items from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_httpx_client @@ -74,12 +75,12 @@ def _get_nested(d: object, path: Sequence[str]) -> object: return cur -def _load_json_env(var_name: str) -> Mapping[str, Sequence[Mapping[str, object]]] | None: +def _load_json_env(var_name: str) -> object: raw: Final = os.environ.get(var_name) if not raw: return None try: - return json.loads(raw) + return _DECODED_JSON.validate_python(json.loads(raw)) except json.JSONDecodeError: return None @@ -91,15 +92,23 @@ def _str_or_none(value: object) -> str | None: return None -def _load_vcap() -> Mapping[str, Sequence[Mapping[str, object]]]: +def _load_vcap() -> object: return _load_json_env(VCAP_SERVICES_ENV_VAR) or {} +def _vcap_service_groups(vcap: object) -> Iterable[object]: + services_by_name: Final = as_str_mapping(vcap) + if services_by_name is not None: + return services_by_name.values() + raise AttributeError(f"'{type(vcap).__name__}' object has no attribute 'values'", name="values", obj=vcap) + + def _get_vcap_service(label: str) -> Mapping[str, object] | None: - for services in _load_vcap().values(): - for svc in services: - if svc.get("label") == label: - return svc + for services in _vcap_service_groups(_load_vcap()): + for svc in iter_items(services): + service = decoded_json_object(svc) + if service.get("label") == label: + return service return None diff --git a/litellm/llms/snowflake/embedding/transformation.py b/litellm/llms/snowflake/embedding/transformation.py index 45bec1283a3..c5e61ad5573 100644 --- a/litellm/llms/snowflake/embedding/transformation.py +++ b/litellm/llms/snowflake/embedding/transformation.py @@ -1,9 +1,10 @@ from collections.abc import Mapping -from typing import Final +from typing import Final, Protocol, runtime_checkable import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -14,6 +15,29 @@ from litellm.types.utils import EmbeddingResponse from ..utils import SnowflakeBaseConfig, SnowflakeException _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) + + +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, key: str | int, /) -> object: ... + + +@runtime_checkable +class _SupportsSetItem(Protocol): + def __setitem__(self, key: str, value: object, /) -> None: ... + + +def _subscript(value: object, key: str | int) -> object: + if isinstance(value, _SupportsGetItem): + return value[key] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + +def _item_assignment_target(value: object) -> _SupportsSetItem: + if isinstance(value, _SupportsSetItem): + return value + raise TypeError(f"'{type(value).__name__}' object does not support item assignment") class SnowflakeEmbeddingConfig(SnowflakeBaseConfig, BaseEmbeddingConfig): @@ -54,10 +78,10 @@ class SnowflakeEmbeddingConfig(SnowflakeBaseConfig, BaseEmbeddingConfig): optional_params: dict, litellm_params: dict, ) -> EmbeddingResponse: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) # convert embeddings to 1d array - for item in response_json["data"]: - item["embedding"] = item["embedding"][0] + for item in iter_items(_subscript(response_json, "data")): + _item_assignment_target(item)["embedding"] = _subscript(_subscript(item, "embedding"), 0) returned_response: Final = EmbeddingResponse.model_validate(_JSON_OBJECT.validate_python(response_json)) returned_response.model = "snowflake/" + (returned_response.model or "") diff --git a/litellm/llms/soniox/realtime/transformation.py b/litellm/llms/soniox/realtime/transformation.py index ed7590b6f39..1bd0176b468 100644 --- a/litellm/llms/soniox/realtime/transformation.py +++ b/litellm/llms/soniox/realtime/transformation.py @@ -164,16 +164,19 @@ def _ignored_settings(update: TranscriptionSessionUpdate) -> tuple[str, ...]: def build_start_request( model: str, update: TranscriptionSessionUpdate | None, options: SonioxRealtimeOptions, drop_params: bool = False ) -> SonioxStartRequest: - return SonioxStartRequest( - **options.model_dump(exclude={"language_hints"}), - language_hints=_language_hints( - None if update is None else update.language, options.language_hints, options.language_hints_strict - ), - model=model, - audio_format="pcm_s16le", - sample_rate=_sample_rate(update), - num_channels=1, - enable_endpoint_detection=update is None or _uses_endpoint_detection(update, drop_params), + configured_options: Final[Mapping[str, object]] = options.model_dump(exclude={"language_hints"}) + return SonioxStartRequest.model_validate( + { + **configured_options, + "language_hints": _language_hints( + None if update is None else update.language, options.language_hints, options.language_hints_strict + ), + "model": model, + "audio_format": "pcm_s16le", + "sample_rate": _sample_rate(update), + "num_channels": 1, + "enable_endpoint_detection": update is None or _uses_endpoint_detection(update, drop_params), + } ) diff --git a/litellm/llms/valkey/vector_stores/transformation.py b/litellm/llms/valkey/vector_stores/transformation.py index cd95d9c95b7..bb9b2da3dbf 100644 --- a/litellm/llms/valkey/vector_stores/transformation.py +++ b/litellm/llms/valkey/vector_stores/transformation.py @@ -15,6 +15,7 @@ import httpx from pydantic import ConfigDict import litellm +from litellm.litellm_core_utils.core_helpers import attribute_items from litellm.llms.base_llm.vector_store.transformation import ( BaseDirectVectorStoreConfig, VectorStoreEmbeddingExecutor, @@ -33,7 +34,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from redis import Redis from redis.asyncio import Redis as AsyncRedis - from redis.commands.search.document import Document from redis.commands.search.query import Query from redis.commands.search.result import Result @@ -185,7 +185,7 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): return query_cls(knn_expr).return_fields(text_field, DISTANCE_FIELD_NAME).paging(0, k).dialect(2) @staticmethod - def _to_result(doc: "Document", text_field: str) -> VectorStoreSearchResult: + def _to_result(doc: object, text_field: str) -> VectorStoreSearchResult: content: Final = [VectorStoreResultContent(text=str(getattr(doc, text_field, "")), type="text")] return VectorStoreSearchResult( score=1.0 - float(getattr(doc, DISTANCE_FIELD_NAME)), @@ -196,7 +196,7 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): @classmethod def _to_response(cls, search_result: "Result", query_text: str, text_field: str) -> VectorStoreSearchResponse: - docs: Final = getattr(search_result, "docs", None) or () + docs: Final = attribute_items(search_result, "docs") data: Final = sorted( (cls._to_result(doc, text_field) for doc in docs), key=lambda result: result.get("score") or 0.0, diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 52345e1ff40..d0de71e0305 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1444,7 +1444,7 @@ class VertexAITokenCounter(BaseTokenCounter): async def count_tokens( self, model_to_use: str, - messages: list[dict[str, Any]] | None, + messages: list[dict[str, object]] | None, contents: list[dict[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 158ed979bad..1f03de4625c 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -1299,7 +1299,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _map_openai_to_vertex_params( self, - openai_request_body: dict[str, Any], + openai_request_body: dict[str, object], deployment_drop_params: object = None, ) -> dict[str, object]: """ diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 9999c75a91c..6e801ed5756 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -8,7 +8,7 @@ import os import re from collections.abc import Callable, Mapping, Sequence from functools import partial -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast from urllib.parse import quote import httpx @@ -18,7 +18,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params, has_attributes from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.prompt_templates.factory import ( GeminiMediaSource, @@ -271,8 +271,21 @@ def _image_url_payload_may_need_sync_gcs_metadata_fetch( return _gs_uri_requires_content_type_metadata(url) +class _SupportsModelDump(Protocol): + def model_dump(self, *, exclude_none: bool) -> object: ... + + +def _message_fields(raw: object) -> Mapping[object, object] | None: + if isinstance(raw, dict): + return raw + if not has_attributes(raw, _SupportsModelDump, "model_dump"): + return None + dumped: Final = raw.model_dump(exclude_none=False) + return dumped if isinstance(dumped, dict) else None + + def _openai_messages_may_need_sync_media_fetch( - messages: list[AllMessageValues], + messages: Sequence[object], custom_llm_provider: str = "vertex_ai", model: str | None = None, ) -> bool: @@ -283,10 +296,7 @@ def _openai_messages_may_need_sync_media_fetch( model is None or VertexGeminiConfig.supports_multimodal_function_responses(model) ) for raw in messages: - msg: Any = raw - if not isinstance(msg, dict) and hasattr(msg, "model_dump"): - msg = msg.model_dump(exclude_none=False) - if not isinstance(msg, dict): + if (msg := _message_fields(raw)) is None: continue images_field = msg.get("images") if isinstance(images_field, list): diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index 87b0dfdbcbc..894c33b5fbc 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -4,7 +4,7 @@ import os from collections.abc import Mapping, Sequence from io import BufferedRandom, BufferedReader, BytesIO from pathlib import Path -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import httpx from httpx._types import RequestFiles @@ -13,6 +13,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size @@ -39,6 +40,18 @@ class _VertexImagenPredictResponse(TypedDict, total=False): predictions: ReadOnly[Sequence[object]] +class _SupportsRead(Protocol): + def read(self) -> bytes | str: ... + + +def _image_items(image: object) -> Sequence[object] | None: + return image if isinstance(image, (list, tuple)) else None + + +def _image_fields(image: object) -> Mapping[object, object] | None: + return image if isinstance(image, dict) else None + + class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): """ Vertex AI Imagen Image Edit Configuration @@ -293,30 +306,30 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): return reference_images - def _read_all_bytes(self, image: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> bytes: + def _read_all_bytes(self, image: object, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> bytes: if depth > max_depth: raise ValueError( f"Max recursion depth {max_depth} reached while reading image bytes for Vertex AI Imagen image edit." ) - if isinstance(image, (list, tuple)): - for item in image: + if (items := _image_items(image)) is not None: + for item in items: if item is not None: return self._read_all_bytes(item, depth=depth + 1, max_depth=max_depth) raise ValueError("Unsupported image type for Vertex AI Imagen image edit.") - if isinstance(image, dict): + if (fields := _image_fields(image)) is not None: for key in ("data", "bytes", "content"): - if key in image and image[key] is not None: - value = image[key] + if key in fields and fields[key] is not None: + value = fields[key] if isinstance(value, str): try: return base64.b64decode(value) except Exception: continue return self._read_all_bytes(value, depth=depth + 1, max_depth=max_depth) - if "path" in image: - return self._read_all_bytes(image["path"], depth=depth + 1, max_depth=max_depth) + if "path" in fields: + return self._read_all_bytes(fields["path"], depth=depth + 1, max_depth=max_depth) if isinstance(image, bytes): return image @@ -350,7 +363,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): "Unsupported image input: filesystem paths are not accepted for " "Vertex AI Imagen image edit. Provide image bytes or a file-like object." ) - if hasattr(image, "read"): + if has_attributes(image, _SupportsRead, "read"): data = image.read() if isinstance(data, str): data = data.encode("utf-8") diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 8d39ae76456..28d6321453f 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -12,7 +12,7 @@ from typing import Any, Final from pydantic import JsonValue, TypeAdapter from litellm.anthropic_beta_headers_manager import aload_beta_headers_config, filter_and_transform_beta_headers -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.core_helpers import decoded_json_object, drops_unsupported_params from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import ( OBJECT_LIST_ADAPTER, @@ -34,6 +34,7 @@ _COUNT_TOKENS_PUBLISHERS: Final = ( _REQUEST_MESSAGES: Final = TypeAdapter(Sequence[object]) _CONTEXT_MANAGEMENT: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_DECODED_JSON: Final = TypeAdapter(object) def count_tokens_request_body( @@ -226,7 +227,7 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): raise ValueError(f"Token counting request failed with status {response.status_code}: {error_text}") # Parse response - result: Final = response.json() + result: Final = decoded_json_object(_DECODED_JSON.validate_python(response.json())) # Return token count # Vertex AI Anthropic returns: {"input_tokens": 123} diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index 080899d08b0..0dd3a395b54 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -9,15 +9,16 @@ import base64 import time from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, ClassVar, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, SupportsFloat, SupportsIndex, TypedDict, cast import httpx from httpx._types import FileContent, RequestFiles -from typing_extensions import ReadOnly +from typing_extensions import Buffer, ReadOnly import litellm from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.vertex_ai.common_utils import ( @@ -80,16 +81,16 @@ def veo_video_count_from_parameters(parameters: Mapping[str, object]) -> int | N def _build_vertex_video_usage_from_request_data( - request_data: dict[str, Any] | None, + request_data: Mapping[str, object] | None, ) -> dict[str, float | str]: """Build usage metadata (duration, resolution, video count) for video cost calculation.""" usage_data: Final[dict[str, float | str]] = {} if not request_data: return usage_data - parameters: Final = request_data.get("parameters", {}) + parameters: Final = decoded_json_object(request_data.get("parameters", {})) duration: Final = parameters.get("durationSeconds") or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS - if duration is not None: + if isinstance(duration, (str, Buffer, SupportsFloat, SupportsIndex)): try: usage_data["duration_seconds"] = float(duration) except (ValueError, TypeError): diff --git a/litellm/main.py b/litellm/main.py index bdb35823757..1d18100d1d5 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -40,6 +40,7 @@ import openai from openai import AsyncStream, Stream from openai.types.audio import TranscriptionStreamEvent from pydantic import BaseModel, TypeAdapter, ValidationError +from typing_extensions import TypedDict as ExtraItemsTypedDict from typing_extensions import assert_never, overload import litellm @@ -407,6 +408,10 @@ def _resolve_control_options(kwargs: Mapping[str, object], model: str) -> Contro return assert_never(control) +class _ForwardedCompletionParams(ExtraItemsTypedDict, total=False, extra_items=object): + pass + + @dataclass(frozen=True, slots=True) class _SupplementalProviderParams: params: GenericLiteLLMParams @@ -8081,7 +8086,7 @@ def text_completion( text_completion_response = TextCompletionResponse() - optional_params: Final[dict[str, Any]] = {} + optional_params: Final[dict[str, object]] = {} # default values for all optional params are none, litellm only passes them to the llm when they are set to non None values if best_of is not None: optional_params["best_of"] = best_of @@ -8148,7 +8153,7 @@ def text_completion( def process_prompt(i, individual_prompt): decoded_prompt: Final = tokenizer.decode(individual_prompt) - all_params: Final = {**kwargs, **optional_params} + all_params: Final = _ForwardedCompletionParams(**{**kwargs, **optional_params}) response: Final[TextCompletionResponse] = text_completion( model=model, prompt=decoded_prompt, @@ -8222,7 +8227,7 @@ def text_completion( messages=messages, *args, **kwargs, - **optional_params, + **_ForwardedCompletionParams(**optional_params), ) if kwargs.get("acompletion", False) is True: return response @@ -8566,7 +8571,7 @@ def transcription( ) -> ( TranscriptionResponse | Stream[TranscriptionStreamEvent] - | Coroutine[Any, Any, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] + | Coroutine[object, object, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] ): """ Calls openai + azure whisper endpoints. @@ -8655,7 +8660,7 @@ def transcription( TranscriptionResponse | Stream[TranscriptionStreamEvent] | AsyncStream[TranscriptionStreamEvent] - | Coroutine[Any, Any, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] + | Coroutine[object, object, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] ) provider_config: Final = ProviderConfigManager.get_provider_audio_transcription_config( diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 01c22dc6725..680b0e6ede4 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -371,8 +371,8 @@ class LazyFeatureMiddleware: fastapi_app.state.lazy_locks = {} @property - def _loaded(self) -> set: - return self._fastapi_app.state.lazy_loaded + def _loaded(self) -> set[str]: + return _lazy_loaded(self._fastapi_app) async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: # Short-circuit once every feature has loaded. diff --git a/litellm/proxy/admin_mcp.py b/litellm/proxy/admin_mcp.py index 32576ea8113..55b0777a27c 100644 --- a/litellm/proxy/admin_mcp.py +++ b/litellm/proxy/admin_mcp.py @@ -71,8 +71,12 @@ def _admin_route_allowed(auth: UserAPIKeyAuth, request: Request, method: str, pa return True +def _scope_app(scope: Mapping[str, object]) -> object: + return scope["app"] + + def _admin_routes(request: Request) -> Iterator[tuple[tuple[str, str], APIRoute]]: - app: Final = request.app + app: Final = _scope_app(request.scope) if not isinstance(app, FastAPI): return for route in app.routes: diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 1e10d628c1e..512a9166e31 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -9,13 +9,14 @@ import os from collections.abc import Mapping, MutableMapping from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Final, Literal, Protocol, cast from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest +from litellm.litellm_core_utils.core_helpers import attribute_or_none, has_attributes, iter_items from litellm.litellm_core_utils.hidden_params import get_hidden_params, set_hidden_param from litellm.llms.openai.workload_identity import resolve_openai_workload_identity_config from litellm.proxy._types import * @@ -82,6 +83,18 @@ if TYPE_CHECKING: router: Final = APIRouter() +class _UserBatchLister(Protocol): + async def list_user_batches( + self, + user_api_key_dict: UserAPIKeyAuth, + limit: int | None = None, + after: str | None = None, + provider: str | None = None, + target_model_names: str | None = None, + llm_router: Router | None = None, + ) -> Mapping[str, object]: ... + + def _hidden_param_string(hidden_params: Mapping[str, object], key: str) -> str: value: Final = hidden_params.get(key) return value if isinstance(value, str) else "" @@ -917,9 +930,9 @@ async def list_batches( # Try to use managed objects table for listing batches (returns encoded IDs). managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") - if managed_files_obj is not None and hasattr(managed_files_obj, "list_user_batches"): + if managed_files_obj is not None and has_attributes(managed_files_obj, _UserBatchLister, "list_user_batches"): verbose_proxy_logger.debug("Using managed objects table for batch listing") - response = await cast(Any, managed_files_obj).list_user_batches( + response = await managed_files_obj.list_user_batches( user_api_key_dict=user_api_key_dict, limit=limit, after=after, @@ -949,9 +962,9 @@ async def list_batches( # Encode batch IDs in the list response so clients can use # them for retrieve/cancel/file downloads through the proxy. - response_data: Final = getattr(response, "data", None) + response_data: Final = attribute_or_none(response, "data") if response_data: - for batch in response_data: + for batch in iter_items(response_data): encode_batch_response_ids(batch, model=model_param) verbose_proxy_logger.debug("Listed batches using model: %s", model_param) diff --git a/litellm/proxy/caching_routes.py b/litellm/proxy/caching_routes.py index 775d00ed53e..dcb56aef927 100644 --- a/litellm/proxy/caching_routes.py +++ b/litellm/proxy/caching_routes.py @@ -2,10 +2,12 @@ from collections.abc import Mapping from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import RedisCache +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import ProxyErrorTypes, ProxyException @@ -13,6 +15,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.caching import CachePingResponse, HealthCheckCacheParams masker: Final = SensitiveDataMasker() +_DECODED_JSON: Final = TypeAdapter(object) router: Final = APIRouter( prefix="/cache", @@ -149,7 +152,7 @@ async def cache_delete(request: Request): if litellm.cache is None: raise HTTPException(status_code=503, detail="Cache not initialized. litellm.cache is None") - request_data: Final = await request.json() + request_data: Final = decoded_json_object(_DECODED_JSON.validate_python(await request.json())) keys: Final = request_data.get("keys", None) if litellm.cache.type == "redis": diff --git a/litellm/proxy/client/chat.py b/litellm/proxy/client/chat.py index ccdef44f9e8..bba339faf7e 100644 --- a/litellm/proxy/client/chat.py +++ b/litellm/proxy/client/chat.py @@ -7,6 +7,8 @@ from pydantic import ConfigDict, TypeAdapter from .exceptions import StreamError, UnauthorizedError +_DECODED_JSON: Final = TypeAdapter(object) + _SSE_LINE: Final[TypeAdapter[bytes | bytearray]] = TypeAdapter( bytes | bytearray, config=ConfigDict(arbitrary_types_allowed=True, strict=True, hide_input_in_errors=True) ) @@ -59,7 +61,7 @@ class ChatClient: frequency_penalty: float | None = None, user: str | None = None, return_request: bool = False, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Create a chat completion. @@ -114,7 +116,7 @@ class ChatClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) diff --git a/litellm/proxy/client/cli/commands/http.py b/litellm/proxy/client/cli/commands/http.py index 529564c685c..27104052613 100644 --- a/litellm/proxy/client/cli/commands/http.py +++ b/litellm/proxy/client/cli/commands/http.py @@ -4,10 +4,13 @@ from typing import Final, Protocol import click import requests import rich +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict from ...http_client import HTTPClient +_DECODED_JSON: Final = TypeAdapter(object) + class _CliContext(TypedDict): base_url: ReadOnly[str] @@ -85,7 +88,7 @@ def request( json_data = None if json: try: - json_data = json_lib.loads(json) + json_data = _DECODED_JSON.validate_python(json_lib.loads(json)) except ValueError as e: raise click.BadParameter(f"Invalid JSON format: {e}") diff --git a/litellm/proxy/client/credentials.py b/litellm/proxy/client/credentials.py index 649fe529936..920aa238ace 100644 --- a/litellm/proxy/client/credentials.py +++ b/litellm/proxy/client/credentials.py @@ -3,9 +3,12 @@ from typing import Any, Final from urllib.parse import quote import requests +from pydantic import TypeAdapter from .exceptions import UnauthorizedError +_DECODED_JSON: Final = TypeAdapter(object) + class CredentialsManagementClient: def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): @@ -75,7 +78,7 @@ class CredentialsManagementClient: credential_values: Mapping[str, object], return_request: bool = False, display_name: str | None = None, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Create a new credential. @@ -111,7 +114,7 @@ class CredentialsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -121,7 +124,7 @@ class CredentialsManagementClient: self, credential_name: str, return_request: bool = False, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Delete a credential by name. @@ -148,7 +151,7 @@ class CredentialsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -159,7 +162,7 @@ class CredentialsManagementClient: credential_name: str, display_name: str | None, return_request: bool = False, - ) -> Mapping[str, object] | requests.Request: + ) -> object | requests.Request: url: Final = f"{self._base_url}/credentials/{quote(credential_name, safe='')}" data: Final[Mapping[str, object]] = {"display_name": display_name, "credential_info": {}} @@ -172,7 +175,7 @@ class CredentialsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -182,7 +185,7 @@ class CredentialsManagementClient: self, credential_name: str, return_request: bool = False, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Get a credential by name. @@ -209,7 +212,7 @@ class CredentialsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) diff --git a/litellm/proxy/client/health.py b/litellm/proxy/client/health.py index 685efa61d15..d40c64b5f41 100644 --- a/litellm/proxy/client/health.py +++ b/litellm/proxy/client/health.py @@ -1,4 +1,6 @@ -from typing import Any, Final +from typing import Final + +from litellm.litellm_core_utils.core_helpers import decoded_json_object from .http_client import HTTPClient @@ -19,7 +21,7 @@ class HealthManagementClient: """ self._http = HTTPClient(base_url=base_url, api_key=api_key, timeout=timeout) - def get_readiness(self) -> dict[str, Any]: + def get_readiness(self) -> object: """ Check the readiness of the LiteLLM proxy server. @@ -32,7 +34,7 @@ class HealthManagementClient: """ return self._http.request("GET", "/health/readiness") - def get_server_version(self) -> str | None: + def get_server_version(self) -> object: """ Get the LiteLLM server version from the readiness endpoint. @@ -40,4 +42,4 @@ class HealthManagementClient: Optional[str]: The server version if available, otherwise None. """ readiness: Final = self.get_readiness() - return readiness.get("litellm_version") + return decoded_json_object(readiness).get("litellm_version") diff --git a/litellm/proxy/client/http_client.py b/litellm/proxy/client/http_client.py index 96a170c75cb..e00529b1915 100644 --- a/litellm/proxy/client/http_client.py +++ b/litellm/proxy/client/http_client.py @@ -1,11 +1,14 @@ """HTTP client for making requests to the LiteLLM proxy server.""" from collections.abc import Mapping -from typing import Any, Final +from typing import Final import requests +from pydantic import TypeAdapter from typing_extensions import TypedDict, Unpack +_DECODED_JSON: Final = TypeAdapter(object) + class _RequestOptions(TypedDict, total=False, extra_items=object): pass @@ -32,10 +35,10 @@ class HTTPClient: uri: str, *, data: Mapping[str, object] | list | bytes | None = None, - json: Mapping[str, object] | list | None = None, + json: object = None, headers: dict[str, str] | None = None, **kwargs: Unpack[_RequestOptions], - ) -> Any: + ) -> object: """Make an HTTP request to the LiteLLM proxy server. This method is used to make generic requests to the LiteLLM proxy @@ -99,4 +102,4 @@ class HTTPClient: response.raise_for_status() # Parse and return JSON response - return response.json() + return _DECODED_JSON.validate_python(response.json()) diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index 028b338f412..dcd49cf375d 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -3,11 +3,14 @@ from collections.abc import Mapping from typing import Any, Final import requests +from pydantic import TypeAdapter from litellm.litellm_core_utils.secret_redaction import redact_string from .exceptions import UnauthorizedError +_DECODED_JSON: Final = TypeAdapter(object) + class KeysManagementClient: def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): @@ -122,7 +125,7 @@ class KeysManagementClient: budget_id: str | None = None, config: Mapping[str, object] | None = None, return_request: bool = False, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Generate an API key based on the provided data. @@ -179,7 +182,7 @@ class KeysManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -190,7 +193,7 @@ class KeysManagementClient: keys: builtins.list[str] | None = None, key_aliases: builtins.list[str] | None = None, return_request: bool = False, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Delete existing keys @@ -223,7 +226,7 @@ class KeysManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -239,7 +242,7 @@ class KeysManagementClient: key_alias: str | None = None, team_id: str | None = None, user_id: str | None = None, - ) -> dict[str, object] | requests.Request: + ) -> object | requests.Request: """ Update an existing API key's parameters. @@ -285,11 +288,11 @@ class KeysManagementClient: response: Final = session.send(request.prepare(), timeout=self._timeout) response_text = response.text response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except Exception: raise Exception(f"Error updating key: {response_text}") - def info(self, key: str, return_request: bool = False) -> dict[str, object] | requests.Request: + def info(self, key: str, return_request: bool = False) -> object | requests.Request: """ Get information about API keys. @@ -314,7 +317,7 @@ class KeysManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) except requests.exceptions.HTTPError as e: redacted_message: Final = redact_string(str(e)) if e.response.status_code == 401: diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 859322f061d..1a67ae1da15 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -5,9 +5,12 @@ from typing import Any, Final import requests from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.core_helpers import decoded_json_object + from .exceptions import NotFoundError, UnauthorizedError _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) +_DECODED_JSON: Final = TypeAdapter(object) class ModelsManagementClient: @@ -195,9 +198,11 @@ class ModelsManagementClient: # Find the matching model for model in models: - if (model_id and model.get("model_info", {}).get("id") == model_id) or ( - model_name and model.get("model_name") == model_name - ): + if ( + model_id + and decoded_json_object(_DECODED_JSON.validate_python(model.get("model_info", {}))).get("id") + == model_id + ) or (model_name and model.get("model_name") == model_name): return model # If we get here, no model was found diff --git a/litellm/proxy/client/teams.py b/litellm/proxy/client/teams.py index dc6fa1e228a..ad23892e998 100644 --- a/litellm/proxy/client/teams.py +++ b/litellm/proxy/client/teams.py @@ -4,9 +4,12 @@ import builtins from typing import Any, Final import requests +from pydantic import TypeAdapter from .exceptions import UnauthorizedError +_DECODED_JSON: Final = TypeAdapter(object) + class TeamsManagementClient: """Client for managing teams in LiteLLM proxy.""" @@ -80,7 +83,7 @@ class TeamsManagementClient: page_size: int = 10, sort_by: str | None = None, sort_order: str = "asc", - ) -> dict[str, object]: + ) -> object: """ Get a paginated list of teams with filtering and sorting options. @@ -125,7 +128,7 @@ class TeamsManagementClient: raise UnauthorizedError("Authentication failed. Check your API key.") response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) def get_available(self) -> builtins.list[dict[str, Any]]: """ diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e812b03385d..4d7c67972f2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -63,11 +63,14 @@ from litellm.litellm_core_utils.bug_report import ( strip_bug_report_notice, ) from litellm.litellm_core_utils.core_helpers import ( + attribute_or_none, finite_cost, get_or_create_metadata_bucket, + has_attributes, independent_snapshot, is_expected_client_error, redact_nested_match_and_regex_keys, + require_attributes, ) from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer from litellm.litellm_core_utils.get_supported_openai_params import ( @@ -674,6 +677,10 @@ class _ReadsHeaderValues(Protocol): def get(self, key: str, default: str = "") -> str: ... +class _NativeLoggingRelease(Protocol): + def release(self, success: bool, /) -> object: ... + + def _as_header_reader(headers: _ReadsHeaderValues) -> _ReadsHeaderValues: return headers @@ -693,6 +700,10 @@ def _as_success_dispatcher(logging_obj: _DispatchesSuccessHandlers) -> _Dispatch return logging_obj +class _ClosesAsync(Protocol): + def aclose(self) -> Awaitable[object]: ... + + def serialize_http_exception_detail( detail: object, ) -> tuple[str, dict | None]: @@ -4169,10 +4180,10 @@ class ProxyBaseLLMRequestProcessing: gating logic directly rather than reimplementing the finally block. """ if getattr(logging_obj, "call_type", None) in ("ocr", "aocr"): - pending: Final = getattr(logging_obj, "_native_pending_logging", None) + pending: Final = attribute_or_none(logging_obj, "_native_pending_logging") if pending is not None: logging_obj._native_pending_logging = None # rebind-ok: consume the native OCR release signal once - pending.release(not exception_raised) + require_attributes(pending, _NativeLoggingRelease, "release").release(not exception_raised) _enqueue_fn: Final = getattr(logging_obj, "enqueue_deferred_logging", None) if _enqueue_fn is None: return @@ -4604,7 +4615,7 @@ class ProxyBaseLLMRequestProcessing: async def finalize_streaming_generator_cleanup( request: Request | None, request_data: dict, - response: Any, + response: object, stream_completed: bool = False, client_disconnected: bool = False, user_api_key_dict: UserAPIKeyAuth | None = None, @@ -4641,7 +4652,7 @@ class ProxyBaseLLMRequestProcessing: ): await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) - if hasattr(response, "aclose"): + if has_attributes(response, _ClosesAsync, "aclose"): try: await response.aclose() except BaseException as e: # noqa: BLE001 diff --git a/litellm/proxy/common_utils/cache_pydantic_utils.py b/litellm/proxy/common_utils/cache_pydantic_utils.py index 34d46be56cd..0beeacd4d87 100644 --- a/litellm/proxy/common_utils/cache_pydantic_utils.py +++ b/litellm/proxy/common_utils/cache_pydantic_utils.py @@ -17,7 +17,7 @@ are encoded; pass a Pydantic model or convert with e.g. ``dataclasses.asdict`` f from __future__ import annotations from collections.abc import Mapping, Sequence -from typing import Any, TypeAlias, TypeVar +from typing import TypeAlias, TypeVar from pydantic import BaseModel, ValidationError @@ -67,7 +67,7 @@ class CacheCodec: return value @staticmethod - def deserialize(cached: Any, model_type: type[T]) -> T | None: + def deserialize(cached: object, model_type: type[T]) -> T | None: """ Decode a cache entry to ``model_type``. diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index d1787f3ff8b..c195779aa0a 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -719,7 +719,7 @@ def strip_callback_config(metadata: dict[str, object] | None) -> dict[str, objec return {k: v for k, v in metadata.items() if k not in _CALLBACK_CONFIG_SLOTS} -def encrypt_callback_vars(metadata: object) -> Any: +def encrypt_callback_vars(metadata: object) -> object: """Return a deep copy of metadata with callback_vars values encrypted at rest. Idempotent: a value that already decrypts cleanly is left unchanged so diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index e210c337922..08461aa3a44 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -2,7 +2,7 @@ import json import re from collections.abc import Collection, Mapping from types import MappingProxyType, UnionType -from typing import Annotated, Any, Final, Literal, Union, get_args, get_origin +from typing import Annotated, Any, Final, Literal, Protocol, Union, get_args, get_origin import orjson from fastapi import HTTPException, Request, UploadFile, status @@ -21,6 +21,7 @@ from litellm.constants import ( MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB, ) from litellm.integrations.otel.runtime import phase_event +from litellm.litellm_core_utils.core_helpers import as_object_list, has_attributes, require_attributes from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.callback_utils import ( get_metadata_variable_name_from_kwargs, @@ -37,6 +38,18 @@ _DECODED_JSON: Final = TypeAdapter(object) _JSON_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object]) +class _ReadableUpload(Protocol): + async def read(self) -> bytes: ... + + +class _UploadedFileInfo(Protocol): + @property + def filename(self) -> str | None: ... + + @property + def content_type(self) -> str | None: ... + + def resolve_realtime_route_model(model: str | None, intent: str | None, is_translation: bool) -> str | None: if model is not None: return model @@ -208,7 +221,7 @@ def _parse_form_metadata(raw_metadata: str) -> object: 400 because no provider or database can encode it. ``json.loads`` accepts the escape, so the multipart ``metadata`` field gets the same check.""" try: - parsed: Final[object] = json.loads(raw_metadata) + parsed: Final = _DECODED_JSON.validate_python(json.loads(raw_metadata)) json.dumps(parsed, ensure_ascii=False).encode("utf-8") except (json.JSONDecodeError, UnicodeEncodeError) as e: raise ProxyException( @@ -602,7 +615,7 @@ async def get_form_data(request: Request) -> dict[str, Any]: async def convert_upload_files_to_file_data( - form_data: dict[str, Any], + form_data: Mapping[str, object], ) -> dict[str, Any]: """ Convert FastAPI UploadFile objects to file data tuples for litellm. @@ -625,21 +638,24 @@ async def convert_upload_files_to_file_data( """ data: Final = {} for key, value in form_data.items(): - if isinstance(value, list): + values = as_object_list(value) + if values is not None: # Check if it's a list of UploadFile objects - if value and hasattr(value[0], "read"): + if values and has_attributes(values[0], _ReadableUpload, "read"): files = [] - for f in value: - file_content = await f.read() + for f in values: + file_content = await require_attributes(f, _ReadableUpload, "read").read() + file_info = require_attributes(f, _UploadedFileInfo, "filename", "content_type") # Create tuple: (filename, content, content_type) - files.append((f.filename, file_content, f.content_type)) + files.append((file_info.filename, file_content, file_info.content_type)) data[key] = files else: data[key] = value - elif hasattr(value, "read"): + elif has_attributes(value, _ReadableUpload, "read"): # Single UploadFile object - read and convert to list for consistency file_content = await value.read() - data[key] = [(value.filename, file_content, value.content_type)] + upload_info = require_attributes(value, _UploadedFileInfo, "filename", "content_type") + data[key] = [(upload_info.filename, file_content, upload_info.content_type)] else: # Regular form field data[key] = value diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index c7c73cfc2a8..91529ee36fc 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -24,6 +24,7 @@ T = TypeVar("T", bound=BaseModel) from litellm.constants import AUTH_OBJECTS_TARGET as AUTH_OBJECTS_TARGET # noqa: E402, PLC0414 # public re-export _HASHED_TOKEN_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}") +_CACHED_VALUE: Final = TypeAdapter(object) _BUDGET_WINDOWS_ABSENT_PREFIX: Final = "budget_windows_absent:" AuthCacheFactory: TypeAlias = Callable[[RedisCache, AuthCacheSettings], AuthCacheStore] @@ -217,7 +218,7 @@ class UserApiKeyCache(DualCache): if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) shared: Final = self._shared - cached: Final = ( + cached: Final = _CACHED_VALUE.validate_python( self._get_shared_sync(shared, key, local_only) if shared is not None else self.key_object_cache.get_cache( @@ -260,7 +261,7 @@ class UserApiKeyCache(DualCache): if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) shared: Final = self._shared - cached: Final = ( + cached: Final = _CACHED_VALUE.validate_python( await self._get_shared(shared, key, local_only) if shared is not None else await self.key_object_cache.async_get_cache( diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index a2ffdd4d59d..a15ffd57333 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -15,6 +15,7 @@ from fastapi import APIRouter, Depends, Request, Response from fastapi.responses import ORJSONResponse from pydantic import TypeAdapter +from litellm.llms.custom_httpx.container_handler import EndpointsConfig from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -28,13 +29,14 @@ from litellm.proxy.container_endpoints.ownership import ( ) _UPLOADED_FILE_FIELD: Final = TypeAdapter(object) +_ENDPOINTS_CONFIG: Final = TypeAdapter(EndpointsConfig) -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) + return _ENDPOINTS_CONFIG.validate_python(json.load(f)) def get_all_route_types() -> list[str]: @@ -57,7 +59,7 @@ def _get_container_provider_config(custom_llm_provider: str): def _create_handler_for_path_params( - path_params: list[str], + path_params: Sequence[str], route_type: str, returns_binary: bool = False, is_multipart: bool = False, diff --git a/litellm/proxy/db/health_check_latest.py b/litellm/proxy/db/health_check_latest.py index 024be7603d6..bc9120e3b4e 100644 --- a/litellm/proxy/db/health_check_latest.py +++ b/litellm/proxy/db/health_check_latest.py @@ -13,10 +13,9 @@ import json from collections.abc import Iterable, Mapping, Sequence from datetime import datetime, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final from pydantic import ConfigDict, JsonValue, TypeAdapter, field_validator -from typing_extensions import LiteralString from litellm._logging import verbose_proxy_logger from litellm.proxy.db.db_span import db_span @@ -78,17 +77,11 @@ class LatestHealthCheckRow(LiteLLMBaseModel): _ROWS_ADAPTER: Final = TypeAdapter(tuple[LatestHealthCheckRow, ...]) -class _HealthCheckRowsQuery(Protocol): - async def __call__(self, query: LiteralString, *args: object) -> Sequence[Mapping[str, object]]: ... - - -def _health_check_rows_query(prisma_client: PrismaClient) -> _HealthCheckRowsQuery: - return prisma_client.db.query_raw - - async def query_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]: + from litellm.proxy.utils import query_raw_rows + async with db_span("latest_health_checks", "LiteLLM_HealthCheckTable"): - rows: Final = await _health_check_rows_query(prisma_client)(LATEST_HEALTH_CHECKS_SQL) + rows: Final = await query_raw_rows(prisma_client.db, LATEST_HEALTH_CHECKS_SQL) return _ROWS_ADAPTER.validate_python(rows) @@ -106,10 +99,10 @@ async def fetch_latest_health_checks_for_models( if not model_names: return () try: + from litellm.proxy.utils import query_raw_rows + async with db_span("latest_health_checks", "LiteLLM_HealthCheckTable"): - rows: Final = await _health_check_rows_query(prisma_client)( - LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, list(model_names) - ) + rows: Final = await query_raw_rows(prisma_client.db, LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, list(model_names)) return _ROWS_ADAPTER.validate_python(rows) except Exception as query_err: # noqa: BLE001 # a paged model list must not fail on its health decoration verbose_proxy_logger.error("Error getting latest health checks for models: %s", query_err) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index be5de040fdf..6c4e2bbfdcf 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -12,7 +12,7 @@ import subprocess import time import urllib import urllib.parse -from collections.abc import Awaitable, Callable, Coroutine +from collections.abc import Awaitable, Callable, Coroutine, Mapping from datetime import datetime, timedelta from typing import TYPE_CHECKING, Final, Protocol, TypeVar @@ -36,8 +36,11 @@ from litellm.secret_managers.main import str_to_bool if TYPE_CHECKING: from prisma import Prisma + from prisma import models as prisma_models from prisma.types import HttpConfig + from litellm.repositories.prisma_protocols import TableActions + _ENVIRONMENT_VARIABLES_REVISION_MIGRATION: Final = "20261015120000_add_environment_variables_revision_trigger" __all__ = ( @@ -904,6 +907,21 @@ class PrismaWrapper: async def disconnect(self, timeout: float | timedelta | None = None) -> None: ... + async def execute_raw(self, query: str, *args: object) -> int: ... + + async def query_first(self, query: str, *args: object) -> Mapping[str, object] | None: ... + + @property + def litellm_deprecatedverificationtoken( + self, + ) -> TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]: ... + + @property + def litellm_mcpservertable(self) -> TableActions[prisma_models.LiteLLM_MCPServerTable]: ... + + @property + def litellm_shadowevaljob(self) -> TableActions[prisma_models.LiteLLM_ShadowEvalJob]: ... + def __getattr__(self, name: str): """ Proxy attribute access to the underlying Prisma client. diff --git a/litellm/proxy/db/proxy_worker_heartbeat.py b/litellm/proxy/db/proxy_worker_heartbeat.py index 3214c0974b4..d0f1b11ada7 100644 --- a/litellm/proxy/db/proxy_worker_heartbeat.py +++ b/litellm/proxy/db/proxy_worker_heartbeat.py @@ -97,11 +97,13 @@ async def count_live_proxy_workers(prisma_client: PrismaClient) -> int | None: counts on the primary: a lagging read replica must never undercount. """ try: + from litellm.proxy.utils import query_raw_rows + db: Final = prisma_client.db primary_db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) else db async with db_span("count_live_proxy_workers", "LiteLLM_ProxyWorkerHeartbeat"): rows: Final = _RAW_QUERY_RESULT.validate_python( - await primary_db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) + await query_raw_rows(primary_db, COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) ) return _COUNT_ROWS_ADAPTER.validate_python(rows)[0]["live_workers"] except Exception as count_err: # noqa: BLE001 # an unknown count must degrade to "warn", never to a 503 diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 50e6aa54da4..c71e2206f45 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -6,7 +6,7 @@ otherwise PrismaClient uses the writer-only PrismaWrapper directly. import asyncio import os -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import timedelta from typing import TYPE_CHECKING, Any, Final @@ -16,8 +16,11 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy.db.prisma_client import PrismaWrapper if TYPE_CHECKING: + from prisma import models as prisma_models from prisma.types import HttpConfig + from litellm.repositories.prisma_protocols import TableActions + # 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( @@ -302,6 +305,23 @@ class RoutingPrismaWrapper: raise RuntimeError("DATABASE_URL_READ_REPLICA not set; cannot recreate read replica client") await self._reader.recreate_prisma_client(reader_url, http_client=http_client) + if TYPE_CHECKING: + + async def execute_raw(self, query: str, *args: object) -> int: ... + + async def query_first(self, query: str, *args: object) -> Mapping[str, object] | None: ... + + @property + def litellm_deprecatedverificationtoken( + self, + ) -> TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]: ... + + @property + def litellm_mcpservertable(self) -> TableActions[prisma_models.LiteLLM_MCPServerTable]: ... + + @property + def litellm_shadowevaljob(self) -> TableActions[prisma_models.LiteLLM_ShadowEvalJob]: ... + def __getattr__(self, name: str) -> Any: if name in _TOP_LEVEL_READ_METHODS: return _DELEGATED_ATTRIBUTE.validate_python(getattr(self.read_target, name)) diff --git a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py index c4615a3cebf..cefcd4993eb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py @@ -14,12 +14,14 @@ import sys from typing import TYPE_CHECKING, Final, Literal from fastapi import HTTPException +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.logging_utils import ( convert_litellm_response_object_to_str, ) @@ -31,6 +33,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks GUARDRAIL_NAME: Final = "aporia" +_DECODED_JSON: Final = TypeAdapter(object) if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -128,8 +131,8 @@ class AporiaGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Aporia AI response: %s", response.text) if response.status_code == 200: # check if the response was flagged - _json_response: Final = response.json() - action: str = _json_response.get("action") # possible values are modify, passthrough, block, rephrase + _json_response: Final = decoded_json_object(_DECODED_JSON.validate_python(response.json())) + action: Final = _json_response.get("action") # possible values are modify, passthrough, block, rephrase if action == "block": raise HTTPException( status_code=400, diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py index 5f57969c310..dc1b9bb84a3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py @@ -3,6 +3,7 @@ from collections.abc import Mapping from typing import Any, Final, cast import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -29,6 +30,7 @@ JAVELIN_API_VERSION_STORED_BY_OLDER_RELEASES: Final = "v1" RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses}) _RESPONSES_API_CALL_TYPES: Final = RESPONSES_API_CALL_TYPES +_DECODED_JSON: Final = TypeAdapter(object) def resolve_content_safety_api_version(configured: str | None) -> str: @@ -62,9 +64,7 @@ class AzureGuardrailBase: self.api_base = api_base self.api_version: str | None = kwargs.get("api_version") - async def _post_to_content_safety( - self, endpoint_path: str, request_body: dict[str, object] - ) -> Mapping[str, object]: + async def _post_to_content_safety(self, endpoint_path: str, request_body: dict[str, object]) -> object: """POST to an Azure Content Safety endpoint with standard auth headers. Args: @@ -89,7 +89,7 @@ class AzureGuardrailBase: json=request_body, timeout=self.timeout, ) - response_json: Final[dict[str, object]] = response.json() + response_json: Final = _DECODED_JSON.validate_python(response.json()) verbose_proxy_logger.debug("Azure Content Safety response [%s]: %s", endpoint_path, response_json) return response_json diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 6a773f41f01..c3a979ad1bc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -16,7 +16,7 @@ import json import re import sys import time -from collections.abc import AsyncGenerator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence from datetime import datetime, timezone from itertools import accumulate, groupby from types import MappingProxyType @@ -33,7 +33,12 @@ from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + decoded_json_object, + iter_items, + redact_nested_match_and_regex_keys, +) from litellm.litellm_core_utils.litellm_logging import ( get_masked_values, ) @@ -1314,12 +1319,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) if httpx_response.status_code == 200: - _json_response: Final = httpx_response.json() + raw_json_response: Final = httpx_response.json() # check if the response was flagged verbose_proxy_logger.debug( "Bedrock AI response : %s", - redact_nested_match_and_regex_keys(_json_response), + redact_nested_match_and_regex_keys(raw_json_response), ) + _json_response: Final[BedrockGuardrailResponse] = raw_json_response bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response) if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response): self._log_apply_guardrail_attempt( @@ -1350,7 +1356,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): def _log_apply_guardrail_attempt( self, httpx_response: httpx.Response, - json_response: dict, # mutable-ok: raw AWS JSON payload + json_response: BedrockGuardrailResponse, request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper event_type: GuardrailEventHooks, start_time: "datetime", @@ -1365,7 +1371,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): billed_usages: Final[tuple[BedrockGuardrailUsage, ...]] = tuple(completed_chunk_usages) + ( (blocking_usage,) if isinstance(blocking_usage, dict) else () ) - logged_json_response: Final = ( + logged_json_response: Final[BedrockGuardrailResponse] = ( { **json_response, "usage": self._sum_usage_counters(billed_usages), @@ -2106,12 +2112,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): response: BedrockGuardrailChecksResponse, ) -> dict[str, object]: """Strip PII location offsets from a checks response before it is logged.""" - sanitized: Final[dict[str, Any]] = copy.deepcopy(dict(response)) - sensitive: Final = (sanitized.get("results") or {}).get("sensitiveInformation") or {} - for entry in sensitive.get("results") or []: - if isinstance(entry, dict): + sanitized: Final = copy.deepcopy(dict(response)) + sensitive: Final = decoded_json_object( + decoded_json_object(sanitized.get("results") or {}).get("sensitiveInformation") or {} + ) + for entry in iter_items(sensitive.get("results") or []): + if isinstance(pii_entry := as_str_mapping(entry), dict): for key in _BEDROCK_CHECKS_PII_LOCATION_KEYS: - entry.pop(key, None) + pii_entry.pop(key, None) return sanitized @staticmethod @@ -2205,7 +2213,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check if the guardrail would block content try: - _json_response: Final = response.json() + _json_response: Final[BedrockGuardrailResponse] = response.json() bedrock_guardrail_response: Final = BedrockGuardrailResponse(**_json_response) if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response): return "guardrail_intervened" @@ -2918,7 +2926,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: AsyncIterable[object], request_data: dict, ) -> AsyncGenerator[ModelResponseStream, None]: """ @@ -2970,7 +2978,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Collect all chunks to process them together started_at: Final = time.monotonic() - all_chunks: Final[list[ModelResponseStream]] = [] + all_chunks: Final[list[object]] = [] async for chunk in response: all_chunks.append(chunk) @@ -3079,18 +3087,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Return the reconstructed stream async for chunk in mock_response: yield chunk - elif raw_sse: - # Forwarding an unscannable stream would silently disable the guardrail, so fail closed. - # A raise cannot reach the client once a keepalive ping has flushed the headers, so the - # refusal travels as a frame, matching how a block is delivered above - for error_frame in anthropic_sse_error_frames( - f"{self.guardrail_name}: streamed response could not be assembled for scanning, blocking it" - ): - yield error_frame - return else: - for chunk in all_chunks: - yield chunk + for unassembled_chunk in ( + # Forwarding an unscannable stream would silently disable the guardrail, so fail closed. + # A raise cannot reach the client once a keepalive ping has flushed the headers, so the + # refusal travels as a frame, matching how a block is delivered above + anthropic_sse_error_frames( + f"{self.guardrail_name}: streamed response could not be assembled for scanning, blocking it" + ) + if raw_sse + else all_chunks + ): + yield unassembled_chunk def _extract_masked_texts_from_response(self, bedrock_guardrail_response: BedrockGuardrailResponse) -> list[str]: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 9b9e0949a1a..60905346fa1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -37,6 +37,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.prompt_templates.factory import ( get_attribute_or_key, get_tool_calls_from_response, @@ -73,7 +74,7 @@ if TYPE_CHECKING: ) BYPASS_HEADER: Final = "x-compresr-bypass" -_AGENTIC_LOOP_MODEL: Final = TypeAdapter(object) +_DECODED_JSON: Final = TypeAdapter(object) COMPRESR_RETRIEVE_TOOL_NAME: Final = "compresr_retrieve" DEFAULT_API_BASE: Final = "https://api.compresr.ai" DEFAULT_COMPRESSION_MODEL: Final = "latte_v2" @@ -150,6 +151,10 @@ class _JsonResponse(Protocol): def json(self) -> object: ... +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + def _json_body(response: _JsonResponse) -> object: return response.json() @@ -1121,8 +1126,13 @@ class CompresrGuardrail(CustomGuardrail): full_model_name = model if logging_obj is not None: - agentic_params: Final = getattr(logging_obj, "model_call_details", {}).get("agentic_loop_params", {}) - candidate: Final = _AGENTIC_LOOP_MODEL.validate_python(agentic_params.get("model", model)) + model_call_details: Final = require_attributes( + _DECODED_JSON.validate_python(getattr(logging_obj, "model_call_details", {})), _SupportsGet, "get" + ) + agentic_params: Final = require_attributes( + model_call_details.get("agentic_loop_params", {}), _SupportsGet, "get" + ) + candidate: Final = agentic_params.get("model", model) if isinstance(candidate, str) and candidate: full_model_name = candidate diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py index d507df24faa..0420bb2bb2a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -133,7 +133,7 @@ class DynamoAIGuardrails(CustomGuardrail): timeout=self.timeout, ) response.raise_for_status() - response_json: Final = response.json() + response_json: Final[DynamoAIResponse] = response.json() end_time = datetime.now() duration = (end_time - start_time).total_seconds() diff --git a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py index 98db3822092..1b9bfc725db 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py @@ -126,7 +126,7 @@ class EnkryptAIGuardrails(CustomGuardrail): timeout=self.timeout, ) response.raise_for_status() - response_json: Final = response.json() + response_json: Final[EnkryptAIResponse] = response.json() end_time = datetime.now() duration = (end_time - start_time).total_seconds() diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index e57d6e78378..6d4fb30dede 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -15,6 +15,7 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -51,6 +52,10 @@ class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): pass +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + class _GraySwanMonitorHTTPResponse(Protocol): def raise_for_status(self) -> object: ... @@ -614,9 +619,12 @@ class GraySwanGuardrail(CustomGuardrail): if not headers: headers = (request_data.get("metadata") or {}).get("headers") if not headers and logging_obj and getattr(logging_obj, "model_call_details", None): - headers = ( - (logging_obj.model_call_details or {}).get("litellm_params", {}).get("metadata", {}).get("headers") + litellm_params: Final = require_attributes( + _DECODED_JSON.validate_python((logging_obj.model_call_details or {}).get("litellm_params", {})), + _SupportsGet, + "get", ) + headers = require_attributes(litellm_params.get("metadata", {}), _SupportsGet, "get").get("headers") if not isinstance(headers, dict): return None diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py index cf6592a3e58..c2455123521 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py @@ -83,7 +83,8 @@ class GuardrailsAI(CustomGuardrail): timeout=self.timeout, ) verbose_proxy_logger.debug("guardrails_ai response: %s", response) - _json_response: Final = GuardrailsAIResponse(**response.json()) + decoded_response: Final[GuardrailsAIResponse] = response.json() + _json_response: Final = GuardrailsAIResponse(**decoded_response) if _json_response.get("validationPassed") is False: raise HTTPException( status_code=400, @@ -130,7 +131,8 @@ class GuardrailsAI(CustomGuardrail): }, ) - _json_response: Final = GuardrailsAIResponsePreCall(**response.json()) + decoded_response: Final[GuardrailsAIResponsePreCall] = response.json() + _json_response: Final = GuardrailsAIResponsePreCall(**decoded_response) response = _json_response.get("outputs", [])[0].get("data", [])[0] return response diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 0b762b2d5fe..f08d9112c33 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -22,6 +22,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.prompt_templates.factory import ( get_attribute_or_key, get_tool_calls_from_response, @@ -76,12 +77,17 @@ _HASH_CACHE_TTL_SECONDS: Final = 15 * 60 # untranslated messages can be read with concrete types (values pass through by # reference, so this is a shallow top-level reconstruction). _REQUEST_DATA_ADAPTER: Final = TypeAdapter(dict[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) class _JsonResponse(Protocol): def json(self) -> object: ... +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + def _json_body(response: _JsonResponse) -> object: return response.json() @@ -985,7 +991,12 @@ class HeadroomGuardrail(CustomGuardrail): full_model_name = model if logging_obj is not None: - agentic_params: Final = getattr(logging_obj, "model_call_details", {}).get("agentic_loop_params", {}) + model_call_details: Final = require_attributes( + _DECODED_JSON.validate_python(getattr(logging_obj, "model_call_details", {})), _SupportsGet, "get" + ) + agentic_params: Final = require_attributes( + model_call_details.get("agentic_loop_params", {}), _SupportsGet, "get" + ) candidate: Final = agentic_params.get("model", model) if isinstance(candidate, str) and candidate: full_model_name = candidate diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 3e5c54b4501..d59c4c316b3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -414,7 +414,8 @@ class LakeraAIGuardrail(CustomGuardrail): verbose_proxy_logger.debug( "Lakera AI v2 guard response: %s", _DECODED_JSON.validate_python(response.json()) ) - lakera_response = LakeraAIResponse(**response.json()) + decoded_response: Final[LakeraAIResponse] = response.json() + lakera_response = LakeraAIResponse(**decoded_response) return lakera_response, masked_entity_count except Exception as e: status = "guardrail_failed_to_respond" diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 71b8976bebf..b1e3d51bbc1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -30,6 +30,7 @@ from litellm.constants import ( DEFAULT_MAX_RECURSE_DEPTH, ) from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.path_utils import is_within, try_safe_join from litellm.proxy.guardrails.content_filter_data import ( @@ -147,10 +148,10 @@ class _CategoryConfigView(TypedDict): def _streamed_choice_content(choice: StreamingChoices | TextChoices) -> str | None: if isinstance(choice, TextChoices): return choice.text if isinstance(choice.text, str) else "" - delta: Final = getattr(choice, "delta", None) + delta: Final = attribute_or_none(choice, "delta") if not delta: return None - content: Final = getattr(delta, "content", None) + content: Final = attribute_or_none(delta, "content") return content if isinstance(content, str) else "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py index da937bd0325..0faddb47cc3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py @@ -11,7 +11,7 @@ import re from collections.abc import Iterator, Sequence from enum import Enum from re import Pattern -from typing import Any, Final +from typing import Final from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import ReadOnly, TypedDict @@ -33,7 +33,21 @@ class _PatternsFile(TypedDict): patterns: ReadOnly[Sequence[_PatternEntry]] +class _CategoryFileSummary(TypedDict, total=False): + category_name: ReadOnly[str] + display_name: ReadOnly[str] + description: ReadOnly[str] + default_action: ReadOnly[str] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _PatternExtraConfig(TypedDict, total=False): + keyword_pattern: ReadOnly[str] + allow_word_numbers: ReadOnly[bool] + + _PATTERNS_FILE: Final = TypeAdapter(_PatternsFile) +_PATTERN_EXTRAS: Final = TypeAdapter(_PatternExtraConfig) def _load_patterns_from_json() -> _PatternsFile: @@ -72,9 +86,11 @@ KNOWN_PATTERN_KEYS: Final = { "description", } -PATTERN_EXTRA_CONFIG: Final[dict[str, dict[str, Any]]] = {} +PATTERN_EXTRA_CONFIG: Final[dict[str, _PatternExtraConfig]] = {} for pattern_data in _PATTERNS_DATA["patterns"]: - extra_config = {key: value for key, value in pattern_data.items() if key not in KNOWN_PATTERN_KEYS} + extra_config = _PATTERN_EXTRAS.validate_python( + {key: value for key, value in pattern_data.items() if key not in KNOWN_PATTERN_KEYS} + ) PATTERN_EXTRA_CONFIG[pattern_data["name"]] = extra_config @@ -154,7 +170,7 @@ def _category_entry(categories_dir: str, filename: str) -> dict[str, str] | None if filename.endswith((".yaml", ".yml")): try: with open(category_file_path, "r") as f: - category_data = yaml.safe_load(f) + category_data: Final[_CategoryFileSummary] = yaml.safe_load(f) except Exception as e: from litellm._logging import verbose_proxy_logger diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py index 7c2d0dbc2fc..209a82f756f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py @@ -24,7 +24,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" "mode='pre_mcp_call'. JWT injection only fires for MCP tool calls." ) - optional_params: Final = getattr(litellm_params, "optional_params", None) + optional_params: Final[object] = getattr(litellm_params, "optional_params", None) def _get(key): if optional_params is not None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 2e3f7120726..9cb929dec92 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -81,6 +81,7 @@ import jwt from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey +from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger @@ -89,6 +90,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypesLiteral @@ -96,6 +98,8 @@ from litellm.types.utils import CallTypesLiteral if TYPE_CHECKING: from jwt.types import Options +_DECODED_JSON: Final = TypeAdapter(object) + class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): pass @@ -496,7 +500,7 @@ class MCPJWTSigner(CustomGuardrail): payload: Final[dict[str, object]] = jwt.decode(raw_token, signing_jwk.key, **decode_kwargs) return payload - async def _introspect_opaque_token(self, token: str) -> dict[str, object]: + async def _introspect_opaque_token(self, token: str) -> Mapping[str, object]: """ Perform RFC 7662 token introspection for opaque (non-JWT) tokens. @@ -522,7 +526,7 @@ class MCPJWTSigner(CustomGuardrail): timeout=self.timeout, ) resp.raise_for_status() - result: Final[dict[str, object]] = resp.json() + result: Final = decoded_json_object(_DECODED_JSON.validate_python(resp.json())) if not result.get("active", False): raise jwt.exceptions.ExpiredSignatureError( "MCPJWTSigner: incoming token is inactive (introspection returned active=false)" @@ -821,7 +825,7 @@ class MCPJWTSigner(CustomGuardrail): # ------------------------------------------------------------------ # FR-5: Verify incoming token before re-signing # ------------------------------------------------------------------ - jwt_claims: dict[str, object] | None = None + jwt_claims: Mapping[str, object] | None = None raw_token: Final[str | None] = hook_data.get("incoming_bearer_token") if self.access_token_discovery_uri and raw_token: diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py index 1956170ba67..1790f335548 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py @@ -3,12 +3,13 @@ import time import uuid from collections import OrderedDict from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import httpx from typing_extensions import NotRequired, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items, json_contains from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) @@ -36,6 +37,17 @@ class GraphTokenResponse(TypedDict): expires_in: NotRequired[int] +def _policy_action_blocks(action: object) -> bool: + policy_action: Final = decoded_json_object(action) + odata_type: Final = policy_action.get("@odata.type", "") + action_field: Final = policy_action.get("action", "") + + if json_contains(odata_type, "restrictAccessAction") or action_field == "restrictAccess": + restriction: Final = policy_action.get("restrictionAction", "") + return restriction == "block" + return False + + class PurviewGuardrailBase: """ Base class for Microsoft Purview guardrails. @@ -364,16 +376,11 @@ class PurviewGuardrailBase: # ------------------------------------------------------------------ @staticmethod - def _should_block(response: dict[str, Any]) -> bool: + def _should_block(response: Mapping[str, object]) -> bool: """Return True if any policyAction requires blocking.""" - for action in response.get("policyActions", []): - odata_type = action.get("@odata.type", "") - action_field = action.get("action", "") - - if "restrictAccessAction" in odata_type or action_field == "restrictAccess": - restriction = action.get("restrictionAction", "") - if restriction == "block": - return True + for action in iter_items(response.get("policyActions", [])): + if _policy_action_blocks(action): + return True return False # ------------------------------------------------------------------ diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index b9871b1e276..721064d1169 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,7 +1,8 @@ +import contextlib import time -from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Generator, Mapping, Sequence from enum import Enum, auto -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Protocol, runtime_checkable +from typing import TYPE_CHECKING, ClassVar, Final, Literal, Protocol, runtime_checkable import httpx from fastapi import HTTPException @@ -302,7 +303,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): def _log_request_debug( self, url: str, - body: dict, + body: Mapping[str, object], file_bytes: "bytes | None", file_type: "str | None", ) -> None: @@ -1162,11 +1163,17 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, - request_data: dict, + response: AsyncIterable[object], + request_data: dict[str, object], ) -> AsyncGenerator[ModelResponseStream, None]: """Process streaming response chunks.""" + async with contextlib.aclosing(self._moderated_stream(response, request_data)) as moderated: + async for moderated_item in moderated: + yield moderated_item + async def _moderated_stream( + self, response: AsyncIterable[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, diff --git a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py index 7c26c4c4c56..021af0f7586 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py +++ b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py @@ -6,17 +6,20 @@ # +-------------------------------------------------------------+ import os import uuid +from collections.abc import Sequence from typing import TYPE_CHECKING, Final, Literal, Optional import httpx from fastapi import HTTPException from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -31,6 +34,11 @@ if TYPE_CHECKING: _DECODED_JSON: Final = TypeAdapter(object) +class _OnyxGuardVerdict(TypedDict, total=False): + allowed: ReadOnly[bool] + violated_rules: ReadOnly[Sequence[str]] + + class OnyxGuardrail(CustomGuardrail): @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: @@ -69,7 +77,7 @@ class OnyxGuardrail(CustomGuardrail): payload: object, input_type: Literal["request", "response"], conversation_id: str, - ) -> dict: + ) -> _OnyxGuardVerdict: """ Call external Onyx Guard server for validation """ @@ -85,8 +93,8 @@ class OnyxGuardrail(CustomGuardrail): }, ) response.raise_for_status() - result: Final = response.json() - if not result.get("allowed", True): + result: Final[_OnyxGuardVerdict] = response.json() + if not decoded_json_object(_DECODED_JSON.validate_python(result)).get("allowed", True): detection_message = "Unknown violation" if "violated_rules" in result: detection_message = ", ".join(result["violated_rules"]) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py index ddc6485b7d5..50a81506bb1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py @@ -21,6 +21,7 @@ _NO_TOOL_INPUT: Final[Mapping[str, object]] = MappingProxyType({}) _DATA_URL_RE: Final = re.compile(r"^data:(?P[^;,]+)?(?P(?:;[^;,]+)*?)(?P;base64)?,", re.IGNORECASE) _URLSAFE_TO_STANDARD_B64: Final = str.maketrans("-_", "+/") _OBJECT_LIST: Final = TypeAdapter(list[object]) +_DECODED_JSON: Final = TypeAdapter(object) def str_or_none(value: object) -> str | None: @@ -240,7 +241,7 @@ def _parsed_tool_input(raw_arguments: str) -> Mapping[str, object]: if not raw_arguments: return _NO_TOOL_INPUT try: - parsed: Final = json.loads(raw_arguments) + parsed: Final = _DECODED_JSON.validate_python(json.loads(raw_arguments)) except (ValueError, TypeError): return _NO_TOOL_INPUT return as_str_mapping(parsed) or _NO_TOOL_INPUT @@ -325,13 +326,17 @@ def _declarations_in(message_index: int, message: object) -> Iterator[_Declarati yield _Declaration(call_id, message_index, name) +def _declaration_call_id(declaration: _Declaration) -> str: + return declaration.call_id + + def _declarations_by_call_id(messages: Sequence[object]) -> Mapping[str, tuple[_Declaration, ...]]: """Every named declaration of each tool call id, in message order, built in one pass.""" declarations: Final = chain.from_iterable( _declarations_in(index, message) for index, message in enumerate(messages) ) ordered: Final = sorted(declarations, key=attrgetter("call_id", "message_index")) - return MappingProxyType({call_id: tuple(group) for call_id, group in groupby(ordered, key=attrgetter("call_id"))}) + return MappingProxyType({call_id: tuple(group) for call_id, group in groupby(ordered, key=_declaration_call_id)}) def _resolve_tool_name( diff --git a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py index 4ec3ed21359..152515c50ed 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py @@ -4,6 +4,7 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Final from fastapi import HTTPException +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache @@ -11,6 +12,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.anthropic.common_utils import ( message_dict_with_server_tool_calls, message_dict_with_server_tool_calls_split_out, @@ -35,6 +37,8 @@ from litellm.types.utils import ( if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_DECODED_JSON: Final = TypeAdapter(object) + def _choice_with_message(choice: object, rewrite: Callable[[object], object]) -> object: mapping: Final = validated_object_mapping(choice) @@ -150,9 +154,13 @@ class PangeaHandler(CustomGuardrail): response.raise_for_status() result: Final = response.json() + decoded_result: Final = decoded_json_object(_DECODED_JSON.validate_python(result)) + verdict: Final = decoded_json_object(decoded_result.get("result", {})) - if result.get("result", {}).get("blocked"): - verbose_proxy_logger.warning("Pangea Guardrail (%s): Request blocked. Response: %s", hook_name, result) + if verdict.get("blocked"): + verbose_proxy_logger.warning( + "Pangea Guardrail (%s): Request blocked. Response: %s", hook_name, decoded_result + ) raise HTTPException( status_code=400, # Bad Request, indicating violation detail={ @@ -161,7 +169,7 @@ class PangeaHandler(CustomGuardrail): }, ) verbose_proxy_logger.debug( - "Pangea Guardrail (%s): Request passed. Response: %s", hook_name, result.get("result", {}).get("detectors") + "Pangea Guardrail (%s): Request passed. Response: %s", hook_name, verdict.get("detectors") ) return result diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 379c2c4b075..73ccd82d37f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -12,6 +12,7 @@ import os import re from collections.abc import AsyncIterable, Mapping, Sequence from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias from urllib.parse import urlparse @@ -27,6 +28,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.llms.anthropic.common_utils import ( anthropic_server_tool_call_params, tool_calls_with_server_tool_calls, @@ -143,6 +145,7 @@ class _ResponsesInputItem(LiteLLMBaseModel): _ResponsesInput: TypeAlias = str | tuple[_ResponsesInputItem, ...] | None _RESPONSES_INPUT: Final[TypeAdapter[_ResponsesInput]] = TypeAdapter(_ResponsesInput) _DECODED_ARGUMENTS: Final = TypeAdapter(object) +_NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({}) class _AirsToolEvent(TypedDict): @@ -901,7 +904,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): }, ) - def _prepare_metadata_from_request(self, data: dict[str, Any]) -> dict[str, object]: + def _prepare_metadata_from_request(self, data: Mapping[str, object]) -> dict[str, object]: """ Extract and prepare metadata from request data for PANW API call. @@ -915,8 +918,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): the profile linked to the API key (configured in Strata Cloud Manager). If both are provided, PANW API uses profile_id (profile_id takes precedence). """ - user_metadata: Final = data.get("metadata", {}) or {} - requester_meta: Final = user_metadata.get("requester_metadata", {}) or {} + user_metadata: Final = decoded_json_object(data.get("metadata") or _NO_METADATA) + requester_meta: Final = user_metadata.get("requester_metadata") or _NO_METADATA metadata: Final[dict[str, object]] = { "user": data.get("user") or "litellm_user", "model": data.get("model") or "unknown", @@ -924,7 +927,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): # Pass through PANW API fields (check requester_metadata fallback for /v1/messages routes) for key in ("profile_name", "profile_id", "user_ip", "app_name", "app_user"): - val = user_metadata.get(key) or requester_meta.get(key) + val = user_metadata.get(key) or decoded_json_object(requester_meta).get(key) if val: metadata[key] = val @@ -939,7 +942,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): data.get("litellm_trace_id") or user_metadata.get("litellm_trace_id") or user_metadata.get("trace_id") - or requester_meta.get("litellm_trace_id") + or decoded_json_object(requester_meta).get("litellm_trace_id") ) if trace_id: metadata["litellm_trace_id"] = trace_id diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index e136b9baaa8..91d196f184f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -8,6 +8,7 @@ # Standard library imports import json import os +from collections.abc import Sized from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import quote @@ -73,6 +74,12 @@ class _PillarProtectHTTPClient(Protocol): ) -> _PillarProtectHTTPResponse: ... +def _json_len(value: object) -> int: + if isinstance(value, Sized): + return len(value) + raise TypeError(f"object of type '{type(value).__name__}' has no len()") + + def _encode_json_for_header(data: object) -> str: """ JSON-serialize and URL-encode data for safe header transmission. @@ -725,7 +732,7 @@ class PillarGuardrail(CustomGuardrail): ) return payload - async def _call_pillar_api(self, headers: dict[str, str], payload: dict[str, Any]) -> _PillarProtectResponse: + async def _call_pillar_api(self, headers: dict[str, str], payload: dict[str, object]) -> _PillarProtectResponse: """ Call the Pillar API and return the response. @@ -737,7 +744,7 @@ class PillarGuardrail(CustomGuardrail): Pillar API response as dictionary """ verbose_proxy_logger.debug( - "Pillar Guardrail: Scanning %s messages for security threats", len(payload.get("messages", [])) + "Pillar Guardrail: Scanning %s messages for security threats", _json_len(payload.get("messages", [])) ) response: Final = await self.async_handler.post( url=f"{self.api_base}/api/v1/protect", diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index c76e8e9c084..7bd72b7323d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -7,6 +7,7 @@ import json import os +from collections.abc import Mapping from typing import Any, Final, Literal from fastapi import HTTPException @@ -17,6 +18,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params from litellm.llms.custom_httpx.http_handler import ( @@ -247,7 +249,7 @@ class QualifireGuardrail(CustomGuardrail): return api_tools if api_tools else None - def _check_if_flagged(self, result: dict[str, Any]) -> bool: + def _check_if_flagged(self, result: Mapping[str, object]) -> bool: """ Check if the Qualifire evaluation result indicates flagged content. @@ -255,12 +257,12 @@ class QualifireGuardrail(CustomGuardrail): A high score (close to 100) indicates GOOD content, low score indicates problems. """ # Check evaluation results for any flagged items - evaluation_results: Final = result.get("evaluationResults", []) or [] + evaluation_results: Final = result.get("evaluationResults") or () - for eval_result in evaluation_results: - results = eval_result.get("results", []) or [] - for r in results: - if r.get("flagged"): + for eval_result in iter_items(evaluation_results): + results = decoded_json_object(eval_result).get("results") or () + for r in iter_items(results): + if decoded_json_object(r).get("flagged"): return True return False @@ -384,7 +386,7 @@ class QualifireGuardrail(CustomGuardrail): timeout=self.timeout, ) response.raise_for_status() - result: Final = response.json() + result: Final = decoded_json_object(_DECODED_JSON.validate_python(response.json())) # Extract response info for logging qualifire_response: Final[dict[str, object]] = { diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index fa5251ba558..2e712f027cb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -110,7 +110,7 @@ class SingulrGuardrail(CustomGuardrail): return SingulrGuardrailConfigModel @staticmethod - def _metadata_containers(request_data: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]: + def _metadata_containers(request_data: Mapping[str, Any]) -> tuple[Mapping[str, object], ...]: litellm_params: Final = request_data.get("litellm_params") or _EMPTY_MAPPING return tuple( container @@ -124,7 +124,7 @@ class SingulrGuardrail(CustomGuardrail): ) @classmethod - def _resolve_metadata_value(cls, request_data: Mapping[str, object], key: str) -> str | None: + def _resolve_metadata_value(cls, request_data: Mapping[str, object], key: str) -> object: for container in cls._metadata_containers(request_data=request_data): value = container.get(key) if value: @@ -140,7 +140,7 @@ class SingulrGuardrail(CustomGuardrail): return None @classmethod - def _build_metadata(cls, request_data: Mapping[str, object]) -> Mapping[str, str] | None: + def _build_metadata(cls, request_data: Mapping[str, object]) -> Mapping[str, object] | None: fields: Final = ( "user_api_key_alias", "user_api_key_user_id", @@ -225,7 +225,7 @@ class SingulrGuardrail(CustomGuardrail): inputs: GenericGuardrailAPIInputs, texts: Sequence[str], structured_messages: Sequence[AllMessageValues], - request_data: Mapping[str, Any], + request_data: Mapping[str, object], ) -> GenericGuardrailAPIInputs: messages: Final = ( tuple(structured_messages) @@ -242,14 +242,16 @@ class SingulrGuardrail(CustomGuardrail): metadata: Final = self._build_metadata(request_data=request_data) - singulr_req_obj = SingulrGuardrailPayload( - correlation_id=request_data.get("litellm_call_id"), - model_name=inputs.get("model"), - guardrail_scope="request", - messages=messages, - images=images, - tools=tools, - metadata=metadata, + singulr_req_obj = SingulrGuardrailPayload.model_validate( + { + "correlation_id": request_data.get("litellm_call_id"), + "model_name": inputs.get("model"), + "guardrail_scope": "request", + "messages": messages, + "images": images, + "tools": tools, + "metadata": metadata, + } ) payload = singulr_req_obj.model_dump(mode="json") guardrail_resp = await self._call_api(payload) @@ -267,7 +269,7 @@ class SingulrGuardrail(CustomGuardrail): return inputs @staticmethod - def _mcp_tool_name(request_data: Mapping[str, Any]) -> str | None: + def _mcp_tool_name(request_data: Mapping[str, object]) -> object: return request_data.get("mcp_tool_name") or request_data.get("name") @staticmethod @@ -283,15 +285,17 @@ class SingulrGuardrail(CustomGuardrail): model: Final = request_data.get("model") return "mcp_tool_name" in request_data or (isinstance(model, str) and model.startswith(_MCP_MODEL_PREFIX)) - async def _apply_guardrail_on_mcp_request(self, request_data: Mapping[str, Any]) -> None: + async def _apply_guardrail_on_mcp_request(self, request_data: Mapping[str, object]) -> None: metadata: Final = self._build_metadata(request_data=request_data) - singulr_mcp_obj = SingulrMcpGuardrailPayload( - guardrail_scope="mcp_request", - tool_name=self._mcp_tool_name(request_data), - tool_arguments=self._mcp_arguments(request_data), - mcp_server_name=request_data.get("mcp_server_name"), - metadata=metadata, + singulr_mcp_obj = SingulrMcpGuardrailPayload.model_validate( + { + "guardrail_scope": "mcp_request", + "tool_name": self._mcp_tool_name(request_data), + "tool_arguments": self._mcp_arguments(request_data), + "mcp_server_name": request_data.get("mcp_server_name"), + "metadata": metadata, + } ) payload = singulr_mcp_obj.model_dump(mode="json") guardrail_resp = await self._call_api(payload) @@ -308,18 +312,20 @@ class SingulrGuardrail(CustomGuardrail): ) async def _apply_guardrail_on_mcp_response( - self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, Any] + self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, object] ) -> GenericGuardrailAPIInputs: if not texts: return inputs metadata: Final = self._build_metadata(request_data=request_data) - singulr_mcp_obj = SingulrMcpGuardrailPayload( - model_name=request_data.get("model"), - guardrail_scope="mcp_response", - tool_result=texts, - metadata=metadata, + singulr_mcp_obj = SingulrMcpGuardrailPayload.model_validate( + { + "model_name": request_data.get("model"), + "guardrail_scope": "mcp_response", + "tool_result": texts, + "metadata": metadata, + } ) payload = singulr_mcp_obj.model_dump(mode="json") guardrail_resp = await self._call_api(payload) @@ -358,7 +364,7 @@ class SingulrGuardrail(CustomGuardrail): ) async def _apply_guardrail_on_response( - self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, Any] + self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, object] ) -> GenericGuardrailAPIInputs: combined_texts: Final = "\n".join(texts) if texts else None @@ -377,14 +383,16 @@ class SingulrGuardrail(CustomGuardrail): metadata: Final = self._build_metadata(request_data=request_data) - singulr_resp_obj = SingulrGuardrailPayload( - correlation_id=request_data.get("litellm_call_id"), - guardrail_scope="response", - model_name=request_data.get("model"), - messages=request_data.get("messages"), - images=inputs.get("images"), - response=assistant_message, - metadata=metadata, + singulr_resp_obj = SingulrGuardrailPayload.model_validate( + { + "correlation_id": request_data.get("litellm_call_id"), + "guardrail_scope": "response", + "model_name": request_data.get("model"), + "messages": request_data.get("messages"), + "images": inputs.get("images"), + "response": assistant_message, + "metadata": metadata, + } ) payload = singulr_resp_obj.model_dump(mode="json") diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index c4fe6d6506b..b4d343b3932 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -39,6 +39,7 @@ from litellm.batches.batch_utils import ( from litellm.constants import BATCH_TPD_DESCRIPTOR_SUFFIX, BATCH_TPD_WINDOW_SECONDS from litellm.exceptions import RateLimitErrorCategory from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.proxy._types import ( ProxyErrorTypes, ProxyException, @@ -100,6 +101,7 @@ else: _BATCH_BODY_ADAPTER: Final = TypeAdapter(dict[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) _WINDOW_START_ADAPTER: Final[TypeAdapter[int | float | str | None]] = TypeAdapter(int | float | str | None) @@ -985,7 +987,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): for raw_line in iter_batch_input_lines(file_content_bytes): request_count += 1 try: - entry = json.loads(raw_line) + entry = _DECODED_JSON.validate_python(json.loads(raw_line)) except Exception: entry_total_tokens = estimate_batch_entry_tokens(raw_line) entry_output_tokens = self.parallel_request_limiter.no_max_tokens_output_floor( @@ -1008,7 +1010,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): output_tokens += entry_output_tokens try: - entry_total_tokens = count_entry_tokens(entry) + entry_total_tokens = count_entry_tokens(decoded_json_object(entry)) except Exception: entry_total_tokens = estimate_batch_entry_tokens(raw_line) total_tokens += entry_total_tokens diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 8f9aebecf30..8784c02ac41 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -2,13 +2,14 @@ import asyncio import json from collections.abc import Sequence from datetime import datetime, timezone -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import attribute_or_none, has_attributes from litellm.proxy._types import ( GenerateKeyRequest, GenerateKeyResponse, @@ -34,6 +35,10 @@ if TYPE_CHECKING: LITELLM_PREFIX_STORED_VIRTUAL_KEYS: Final = "litellm/" +class _ModelDumpable(Protocol): + def model_dump(self) -> object: ... + + class KeyManagementEventHooks: @staticmethod async def async_key_generated_hook( @@ -476,12 +481,13 @@ class KeyManagementEventHooks: verbose_proxy_logger.debug("Unable to load team metadata for team_id=%s: %s", team_id, exc) return None - metadata = getattr(team_obj, "metadata", None) - if metadata is None: + raw_metadata: Final = attribute_or_none(team_obj, "metadata") + if raw_metadata is None: return None - if hasattr(metadata, "model_dump"): - metadata = metadata.model_dump() + metadata: Final = ( + raw_metadata.model_dump() if has_attributes(raw_metadata, _ModelDumpable, "model_dump") else raw_metadata + ) if not isinstance(metadata, dict): return None diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index 6bd8417eb5e..576e3997db0 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -36,6 +36,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import has_attributes, iter_items, require_attributes from litellm.llms.litellm_proxy.skills.constants import LITELLM_SKILL_ID_PREFIX from litellm.llms.litellm_proxy.skills.prompt_injection import ( SkillPromptInjectionHandler, @@ -84,6 +85,11 @@ class _ChatCompletion(Protocol): def choices(self) -> Sequence[_ChatChoice]: ... +class _HasContent(Protocol): + @property + def content(self) -> object: ... + + @with_config(ConfigDict(extra="allow", strict=True)) class _CodeExecutionToolArgs(TypedDict, total=False): code: ReadOnly[str] @@ -105,6 +111,10 @@ def _first_choice(response: _ChatCompletion) -> _ChatChoice: return response.choices[0] +def _dict_or_none(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, dict) else None + + class SkillsInjectionHook(CustomLogger): """ Pre/Post-call hook that processes skills from container.skills parameter. @@ -456,26 +466,28 @@ class SkillsInjectionHook(CustomLogger): skill_files=all_skill_files, ) - def _extract_tool_calls(self, response: Any) -> list[dict[str, Any]]: + def _extract_tool_calls(self, response: object) -> list[dict[str, Any]]: """Extract tool calls from response, handling both formats.""" tool_calls: Final = [] # Get content - handle both dict and object responses - content = None - if isinstance(response, dict): - content = response.get("content", []) - elif hasattr(response, "content"): - content = response.content + response_mapping: Final = _dict_or_none(response) + content: Final = ( + response_mapping.get("content", []) + if response_mapping is not None + else (response.content if has_attributes(response, _HasContent, "content") else None) + ) # Anthropic/messages API format: response has "content" list with tool_use blocks if content: - for block in content: - if isinstance(block, dict) and block.get("type") == "tool_use": + for block in iter_items(content): + block_mapping = _dict_or_none(block) + if block_mapping is not None and block_mapping.get("type") == "tool_use": tool_calls.append( { - "id": block.get("id"), - "name": block.get("name"), - "input": block.get("input", {}), + "id": block_mapping.get("id"), + "name": block_mapping.get("name"), + "input": block_mapping.get("input", {}), } ) elif hasattr(block, "type") and getattr(block, "type", None) == "tool_use": @@ -488,7 +500,7 @@ class SkillsInjectionHook(CustomLogger): ) # OpenAI format: response has choices[0].message.tool_calls - if not tool_calls and hasattr(response, "choices") and response.choices: + if not tool_calls and has_attributes(response, _ChatCompletion, "choices") and response.choices: msg: Final = response.choices[0].message if hasattr(msg, "tool_calls") and msg.tool_calls: for tc in msg.tool_calls: @@ -770,11 +782,11 @@ print('No executable skill module found') executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) generated_files: Final[list[dict[str, object]]] = [] - current_response: Any = response + current_response: object = response for iteration in range(self.max_iterations): # OpenAI format response has choices[0].message - choice: _ChatChoice = _first_choice(current_response) + choice: _ChatChoice = _first_choice(require_attributes(current_response, _ChatCompletion, "choices")) assistant_message: _ChatMessage = choice.message stop_reason: str | None = choice.finish_reason diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index 5b7d7d04e24..f12756bc1e6 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -191,7 +191,7 @@ class _PROXY_MaxIterationsHandler(CustomLogger): """ if self.increment_script is not None: try: - result: Final = await self.increment_script( + result: Final[int] = await self.increment_script( keys=[cache_key], args=[self.ttl], ) @@ -207,7 +207,7 @@ class _PROXY_MaxIterationsHandler(CustomLogger): async def _in_memory_increment(self, cache_key: str) -> int: """Increment counter in in-memory cache with TTL.""" - current: Final = await self.internal_usage_cache.async_get_cache( + current: Final[int | None] = await self.internal_usage_cache.async_get_cache( key=cache_key, litellm_parent_otel_span=None, local_only=True, diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index c6e26999766..1830961dd06 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -6,7 +6,7 @@ Reduces context window size and improves tool selection accuracy. """ from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence -from typing import TYPE_CHECKING, Final, Optional +from typing import TYPE_CHECKING, Final, Optional, Protocol from fastapi import HTTPException from typing_extensions import ReadOnly, TypedDict @@ -18,6 +18,7 @@ from litellm.constants import ( DEFAULT_MCP_SEMANTIC_FILTER_TOP_K, ) from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import as_str_mapping, has_attributes from litellm.proxy.mcp_server.semantic_tool_filter import ( SemanticToolFilterContextWindowError, ) @@ -50,6 +51,35 @@ def _truncate_csv_at_tool_name_boundary(tool_names_csv: str, max_length: int) -> return head.rsplit(",", 1)[0] +class _PydanticTool(Protocol): + def model_dump(self, *, exclude_none: bool) -> Mapping[str, object]: ... + + +class _PydanticV1Tool(Protocol): + def dict(self, *, exclude_none: bool) -> Mapping[str, object]: ... + + +def _tool_as_dict(tool: object) -> object: + if has_attributes(tool, _PydanticTool, "model_dump"): + tool_dict: Final = tool.model_dump(exclude_none=True) + verbose_proxy_logger.debug( + "Converted Pydantic tool to dict: %s -> dict with keys: %s", + type(tool).__name__, + list(tool_dict.keys()), + ) + return tool_dict + if has_attributes(tool, _PydanticV1Tool, "dict"): + legacy_tool_dict: Final = tool.dict(exclude_none=True) + verbose_proxy_logger.debug("Converted Pydantic tool (v1) to dict: %s -> dict", type(tool).__name__) + return legacy_tool_dict + tool_mapping: Final = as_str_mapping(tool) + if tool_mapping is not None and isinstance(tool, dict): + verbose_proxy_logger.debug("Tool is already a dict with keys: %s", list(tool_mapping.keys())) + return tool_mapping + verbose_proxy_logger.warning("Tool is unknown type: %s, passing as-is", type(tool)) + return tool + + class SemanticToolFilterHook(CustomLogger): """ Pre-call hook that filters MCP tools semantically. @@ -92,7 +122,7 @@ class SemanticToolFilterHook(CustomLogger): self, tools: Iterable[Mapping[str, object]], user_api_key_dict: "UserAPIKeyAuth", - ) -> list[dict[str, object]]: + ) -> list[object]: """ Expand MCP references to actual tool definitions. @@ -119,26 +149,7 @@ class SemanticToolFilterHook(CustomLogger): ) # Convert Pydantic models to dicts for compatibility - openai_tools_as_dicts: Final[list[dict[str, object]]] = [] - for tool in openai_tools: - if hasattr(tool, "model_dump"): - tool_dict = tool.model_dump(exclude_none=True) - verbose_proxy_logger.debug( - "Converted Pydantic tool to dict: %s -> dict with keys: %s", - type(tool).__name__, - list(tool_dict.keys()), - ) - openai_tools_as_dicts.append(tool_dict) - elif hasattr(tool, "dict"): - tool_dict = tool.dict(exclude_none=True) - verbose_proxy_logger.debug("Converted Pydantic tool (v1) to dict: %s -> dict", type(tool).__name__) - openai_tools_as_dicts.append(tool_dict) - elif isinstance(tool, dict): - verbose_proxy_logger.debug("Tool is already a dict with keys: %s", list(tool.keys())) - openai_tools_as_dicts.append(tool) - else: - verbose_proxy_logger.warning("Tool is unknown type: %s, passing as-is", type(tool)) - openai_tools_as_dicts.append(tool) + openai_tools_as_dicts: Final = list(map(_tool_as_dict, openai_tools)) verbose_proxy_logger.debug( "Expanded %s MCP reference(s) to %s tools (all as dicts)", len(mcp_tools), len(openai_tools_as_dicts) @@ -149,8 +160,8 @@ class SemanticToolFilterHook(CustomLogger): async def _filter_expanded_tools( self, data: dict, - expanded_tools: list[dict[str, object]], - ) -> list[dict[str, object]]: + expanded_tools: list[object], + ) -> list[object]: """ Apply the semantic filter to expanded MCP tool definitions. diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 33a14bfbeda..6b36a6ac8bc 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -4743,7 +4743,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_tokens=reserved_tokens, ) ) - ptu_skips_socket_usage: Final = kwargs.get("call_type") == CallTypes.aresponses_websocket.value + call_type: Final[object] = kwargs.get("call_type") + ptu_skips_socket_usage: Final = call_type == CallTypes.aresponses_websocket.value pipeline_operations.extend( self._build_team_ptu_tpm_ops( standard_logging_metadata, # pyright: ignore[reportUnknownArgumentType] # untyped logging metadata diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 0269a784445..0ab3d1ca7af 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -5,9 +5,9 @@ This hook uses the DBSpendUpdateWriter to batch-write response IDs to the databa instead of writing immediately on each request. """ -from collections.abc import AsyncGenerator, Callable, Mapping +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypeVar +from typing import TYPE_CHECKING, Final, TypeVar from fastapi import HTTPException from pydantic import TypeAdapter, ValidationError @@ -460,8 +460,8 @@ class ResponsesIDSecurity(CustomLogger): return response async def async_post_call_streaming_iterator_hook( - self, user_api_key_dict: "UserAPIKeyAuth", response: Any, request_data: dict[str, object] - ) -> AsyncGenerator[BaseLiteLLMOpenAIResponseObject, None]: + self, user_api_key_dict: "UserAPIKeyAuth", response: AsyncIterable[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: advertise: Final = self._streamed_response_id_advertiser(user_api_key_dict, request_data) async for chunk in response: yield ( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..546521f9048 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -14,6 +14,7 @@ from typing import ( Any, Final, Literal, + Protocol, TypeAlias, cast, # noqa: TID251 # narrows untyped deployment and credential payloads ) @@ -49,7 +50,13 @@ from litellm.constants import ( X_LITELLM_DISABLE_CALLBACKS, ) from litellm.integrations.custom_guardrail import merged_guardrail_entries, without_server_streaming_classification -from litellm.litellm_core_utils.core_helpers import is_codex_user_agent, normalize_drop_params +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + attribute_or_none, + is_codex_user_agent, + normalize_drop_params, + require_attributes, +) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( TRUSTED_CALLBACK_VARS_FIELD, @@ -2664,11 +2671,11 @@ async def add_litellm_data_to_request( _raw_headers: Final[dict[str, str]] = RedactedDict(safe_get_request_headers(request)) - forward_llm_auth = False - if general_settings: - forward_llm_auth = general_settings.get("forward_llm_provider_auth_headers", False) - if not forward_llm_auth: - forward_llm_auth = getattr(litellm, "forward_llm_provider_auth_headers", False) + general_settings_view: Final = as_str_mapping(general_settings) + forward_llm_auth: Final = bool( + (general_settings_view.get("forward_llm_provider_auth_headers", False) if general_settings_view else False) + or attribute_or_none(litellm, "forward_llm_provider_auth_headers") + ) _headers: Final[dict[str, str]] = clean_headers( request.headers, litellm_key_header_name=( @@ -3562,6 +3569,10 @@ def _warn_stale_team_alias_once(warning_key: str, message: str, *args: str) -> N verbose_proxy_logger.warning(message, *args) +class _SupportsStartswith(Protocol): + def startswith(self, prefix: str, /) -> bool: ... + + def _update_model_if_team_alias_exists( data: dict, user_api_key_dict: UserAPIKeyAuth, @@ -3589,14 +3600,15 @@ def _update_model_if_team_alias_exists( exist (e.g. a gateway-level model group shared with the team). """ _model: Final = data.get("model") - if not _model or not user_api_key_dict.team_model_aliases or _model not in user_api_key_dict.team_model_aliases: + team_model_aliases: Final = as_str_mapping(user_api_key_dict.team_model_aliases) + if not _model or not team_model_aliases or _model not in team_model_aliases: return from litellm.proxy.proxy_server import llm_router # Skip alias rewrite if this model resolves to team-specific deployments # (team models use team_public_model_name, not model_aliases) - aliased_target: Final = user_api_key_dict.team_model_aliases[_model] + aliased_target: Final = team_model_aliases[_model] # Optional bypass for stale aliases from pre-PR deployments: # only enabled via feature flag to preserve backwards compatibility. @@ -3607,7 +3619,9 @@ def _update_model_if_team_alias_exists( enable_stale_alias_bypass: Final = _ENABLE_TEAM_STALE_ALIAS_BYPASS # Check if the alias points to a team-scoped UUID name # (format: "model_name_{team_id}_{uuid}") - is_stale_team_alias: Final = aliased_target.startswith(f"model_name_{user_api_key_dict.team_id}_") + is_stale_team_alias: Final = require_attributes(aliased_target, _SupportsStartswith, "startswith").startswith( + f"model_name_{user_api_key_dict.team_id}_" + ) if is_stale_team_alias and llm_router: if aliased_target not in llm_router.registry().index("model_name"): _warn_stale_team_alias_once( diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 3c733d08a2e..1e43a8375a6 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -236,8 +236,10 @@ _RAW_ROWS: Final = TypeAdapter(Sequence[dict[str, object]], config=ConfigDict(st async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) -> Sequence[Mapping[str, object]]: + from litellm.proxy.utils import query_raw_rows + async with db_span("auto_router_report_query", sql_relation(query)): - return _RAW_ROWS.validate_python(await prisma_client.db.query_raw(query, *args)) + return _RAW_ROWS.validate_python(await query_raw_rows(prisma_client.db, query, *args)) async def _execute_raw(prisma_client: "PrismaClient", query: str, *args: object) -> int: @@ -939,7 +941,7 @@ def _idle_router_groups( zero: Final = _summed_agg_row(()) idle: Final = frozenset( key - for key in (_strategy_router_key(deployment) for deployment in llm_router.registry().all() or ()) + for key in map(_strategy_router_key, llm_router.registry().all() or ()) if key is not None and key not in covered ) return tuple( diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index aaf5d30d2e1..a09045829b8 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -12,10 +12,11 @@ import asyncio import json from collections.abc import Mapping from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol from fastapi import APIRouter, Depends, Header, HTTPException from pydantic import Field, TypeAdapter +from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm._redis import _redis_kwargs_from_environment @@ -84,7 +85,11 @@ _REDACTED_VALUE: Final = "***REDACTED***" _URL_OVERRIDDEN_CONNECTION_FIELDS: Final[frozenset[str]] = frozenset({"host", "port", "db", "password", "username"}) -def _resolve_cache_url_precedence(settings: Mapping[str, object]) -> dict[str, Any]: +class _CacheSettings(TypedDict, total=False, extra_items=object): + pass + + +def _resolve_cache_url_precedence(settings: Mapping[str, object]) -> _CacheSettings: """Return cache settings with the url-vs-discrete-fields ambiguity resolved. When a full ``url`` is supplied it wins: the discrete @@ -97,8 +102,8 @@ def _resolve_cache_url_precedence(settings: Mapping[str, object]) -> dict[str, A url: Final = settings.get("url") has_url: Final = isinstance(url, str) and url.strip() != "" if not has_url or settings.get("redis_startup_nodes"): - return dict(settings) - return {k: v for k, v in settings.items() if k not in _URL_OVERRIDDEN_CONNECTION_FIELDS} + return _CacheSettings(**settings) + return _CacheSettings(**{k: v for k, v in settings.items() if k not in _URL_OVERRIDDEN_CONNECTION_FIELDS}) def _parse_stored_settings(cache_settings_value: object) -> dict[str, object]: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index a76dad8d02a..6c6118b2505 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -85,7 +85,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: ) from litellm.proxy.management_helpers.team_roster_sync import user_tx_db from litellm.proxy.management_helpers.utils import management_endpoint_wrapper -from litellm.proxy.utils import handle_exception_on_proxy, hash_password, run_off_loop_if_fips +from litellm.proxy.utils import handle_exception_on_proxy, hash_password, query_raw_rows, run_off_loop_if_fips from litellm.redis.stores.spend_counters import ( user_counter, ) @@ -1240,7 +1240,7 @@ async def user_info_v2( async def _fetch_admin_teams_and_keys_rows( prisma_client: "PrismaClient", sql_query: str ) -> Sequence[_AdminTeamAndKeyRows]: - return _ADMIN_TEAM_AND_KEY_ROWS_ADAPTER.validate_python(await prisma_client.db.query_raw(sql_query)) + return _ADMIN_TEAM_AND_KEY_ROWS_ADAPTER.validate_python(await query_raw_rows(prisma_client.db, sql_query)) async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f990f043dea..d4a0a35d9eb 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -40,6 +40,7 @@ from litellm.constants import ( MINIMUM_CUSTOM_KEY_LENGTH, UI_SESSION_TOKEN_TEAM_ID, ) +from litellm.litellm_core_utils.core_helpers import as_object_list from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.models.credentials import CredentialItem @@ -2058,7 +2059,7 @@ async def _check_project_key_limits( # Validate key models are a subset of project models if data.models and len(project_obj.models) > 0: - for m in data.models: + for m in as_object_list(data.models) or (): if m not in project_obj.models and m not in _INHERITED_MODEL_SENTINELS: raise HTTPException( status_code=400, diff --git a/litellm/proxy/management_endpoints/management_v1/budgets.py b/litellm/proxy/management_endpoints/management_v1/budgets.py index 19a5408496a..4c8d8fdf302 100644 --- a/litellm/proxy/management_endpoints/management_v1/budgets.py +++ b/litellm/proxy/management_endpoints/management_v1/budgets.py @@ -35,7 +35,7 @@ from litellm.proxy.list_api.list_framework import ( where_sql, ) from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import PrismaClient, query_raw_rows from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.proxy.management_endpoints.management_v1 import ( ListResponse, @@ -113,7 +113,7 @@ class PrismaBudgetListExecutor: async def count(self, where: tuple[Predicate, ...]) -> int: clauses, params = where_sql(where) sql: Final = f"SELECT COUNT(*) AS count FROM {BUDGET_TABLE}" + (f" WHERE {clauses}" if clauses else "") - counted: Final = _ROW_COUNTS.validate_python(await self.prisma_client.db.query_raw(sql, *params)) + counted: Final = _ROW_COUNTS.validate_python(await query_raw_rows(self.prisma_client.db, sql, *params)) return counted[0].count if counted else 0 async def find_many(self, plan: QueryPlan) -> Sequence[BudgetListItem]: @@ -130,7 +130,9 @@ class PrismaBudgetListExecutor: + f" ORDER BY {ordering}" + f" LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}" ) - return _BUDGET_ROWS.validate_python(await self.prisma_client.db.query_raw(sql, *params, plan.take, plan.skip)) + return _BUDGET_ROWS.validate_python( + await query_raw_rows(self.prisma_client.db, sql, *params, plan.take, plan.skip) + ) def _serialize(row: BudgetListItem) -> BudgetListItem: diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b8b926a3057..b3c1f48bf99 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -30,6 +30,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME, TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS +from litellm.litellm_core_utils.core_helpers import as_object_list, decoded_json_object from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.ptu_pricing import ( @@ -210,6 +211,7 @@ class KubernetesPodsResponse(BaseModel): CLEARABLE_LITELLM_PARAMS: Final = frozenset({"cache_control_injection_points", "litellm_credential_name"}) NULL_CLEARABLE_LITELLM_PARAMS: Final = frozenset((*SPECIAL_MODEL_INFO_PARAMS, *CLEARABLE_LITELLM_PARAMS)) +_STR_OBJECT_DICT: Final = TypeAdapter(dict[str, object]) async def update_team(*args, **kwargs): @@ -246,12 +248,12 @@ def _model_write_response( ) -> _ProxyModelRow | Mapping[str, object]: if member_write is None: return row - payload: Final = TypeAdapter(dict[str, object]).validate_json(row.model_dump_json()) + payload: Final = _STR_OBJECT_DICT.validate_json(row.model_dump_json()) stored_params: Final = payload.get("litellm_params") params: Final = ( - TypeAdapter(dict[str, object]).validate_json(stored_params) + _STR_OBJECT_DICT.validate_json(stored_params) if isinstance(stored_params, str) - else TypeAdapter(dict[str, object]).validate_python(stored_params) + else _STR_OBJECT_DICT.validate_python(stored_params) ) redacted: Final = redact_credentials_in_payload(params) return { @@ -1371,7 +1373,9 @@ def legacy_update_model_info_changes( value that differs from the stored override resets that override (``None``). """ stored: Final = db_model.model_info.model_dump(exclude_none=True) - requested: Final = incoming_model_info.model_dump(exclude_unset=True, exclude_none=True) + requested: Final = _STR_OBJECT_DICT.validate_python( + incoming_model_info.model_dump(exclude_unset=True, exclude_none=True) + ) sent: Final = {k: v for k, v in requested.items() if k not in _LEGACY_UPDATE_IGNORED_MODEL_INFO_FIELDS and v != ""} echoed: Final = frozenset( echoed_cost_map_fields( @@ -1381,7 +1385,7 @@ def legacy_update_model_info_changes( discovered_model_info, ) ) - kept: Final = without_server_derived_pricing(sent) + kept: Final[Mapping[str, object]] = without_server_derived_pricing(sent) return MappingProxyType( { **{k: getattr(incoming_model_info, k) for k in kept if k not in echoed}, @@ -1421,20 +1425,16 @@ async def _locked_legacy_update_model_info_json( return None stored_model_info: Final[object] = payload.get("model_info") stored: Final = ( - TypeAdapter(dict[str, object]).validate_json(stored_model_info) + _STR_OBJECT_DICT.validate_json(stored_model_info) if isinstance(stored_model_info, str) - else TypeAdapter(dict[str, object]).validate_python(stored_model_info or {}) - ) - return ( - TypeAdapter(dict[str, object]) - .dump_json( - { - **{k: v for k, v in stored.items() if k not in writes}, - **{k: v for k, v in writes.items() if v is not None}, - } - ) - .decode() + else _STR_OBJECT_DICT.validate_python(stored_model_info or {}) ) + return _STR_OBJECT_DICT.dump_json( + { + **{k: v for k, v in stored.items() if k not in writes}, + **{k: v for k, v in writes.items() if v is not None}, + } + ).decode() def _param_keys_to_remove(patch: updateDeployment) -> tuple[str, ...]: @@ -1543,7 +1543,7 @@ def update_db_model( prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel( model_name=merged_model_name, litellm_params=json.dumps(merged_litellm_params), - model_info=TypeAdapter(dict[str, object]).dump_json(merged_model_info).decode(), + model_info=_STR_OBJECT_DICT.dump_json(merged_model_info).decode(), ) if updated_patch.blocked is not None: @@ -2080,7 +2080,7 @@ async def add_model_to_db( slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None, ) -> "_ProxyModelRow | LiteLLM_ProxyModelTable": # encrypt litellm params # - _litellm_params_dict: Final = TypeAdapter(dict[str, object]).validate_python( + _litellm_params_dict: Final = _STR_OBJECT_DICT.validate_python( model_params.litellm_params.model_dump(exclude_none=True) ) if "complexity_router_config" in _litellm_params_dict: @@ -4000,19 +4000,20 @@ async def clear_cache() -> ReconcileOutcome: db_router_names: Final = set() - for model in current_models: - model_info = model.get("model_info", {}) + for model in as_object_list(current_models) or (): + deployment = decoded_json_object(model) + model_info = decoded_json_object(deployment.get("model_info", {})) if model_info.get("db_model", False): db_model_ids.append(model_info.get("id")) # Auto-router deployments (and only those) are wiped here, in the # same pass, so the reload rebuilds them -- see the comment below. - model_name = model.get("model_name") - if model_name is not None and str(model.get("litellm_params", {}).get("model", "")).startswith( - "auto_router/" - ): - db_router_names.add(model_name) + deployment_name = deployment.get("model_name") + if isinstance(deployment_name, str) and str( + decoded_json_object(deployment.get("litellm_params", {})).get("model", "") + ).startswith("auto_router/"): + db_router_names.add(deployment_name) router_model_id = model_info.get("id") - if router_model_id is not None: + if isinstance(router_model_id, str): llm_router.apply(ChangeSet(delete=[router_model_id])) else: # This is a config model, preserved by the reconcile below diff --git a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py index 1096954536a..d3854aa3ef3 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py @@ -6,11 +6,16 @@ based on user-provided attack examples and descriptions. import json from typing import Final +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items from litellm.proxy._types import ProxyErrorTypes, ProxyException +_DECODED_JSON: Final = TypeAdapter(object) + SUGGEST_TOOL: Final = { "type": "function", "function": { @@ -97,14 +102,19 @@ class AiPolicySuggester: "explanation": "No templates could be matched to your requirements.", } - result: Final = json.loads(tool_calls[0].function.arguments) + decoded: Final = _DECODED_JSON.validate_python(json.loads(tool_calls[0].function.arguments)) valid_ids: Final = {t["id"] for t in templates} - result["selected_templates"] = [ - s for s in result.get("selected_templates", []) if s.get("template_id") in valid_ids - ] - - return result + result: Final = decoded_json_object(decoded) + suggestion: Final = { + **result, + "selected_templates": [ + s + for s in iter_items(result.get("selected_templates", [])) + if decoded_json_object(s).get("template_id") in valid_ids + ], + } + return suggestion except Exception as e: verbose_proxy_logger.error("AI policy suggestion failed: %s", e) raise diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index aaa261eb19b..80df6a95cc2 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -65,7 +65,7 @@ from litellm.constants import ( MICROSOFT_USER_ID_ATTRIBUTE, MICROSOFT_USER_LAST_NAME_ATTRIBUTE, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -4856,7 +4856,7 @@ class MicrosoftSSOHandler: # Decode the JWT without signature verification # (signature is already verified by fastapi_sso) - decoded_token: Final = jwt.decode(id_token, options={"verify_signature": False}) + decoded_token: Final = decoded_json_object(jwt.decode(id_token, options={"verify_signature": False})) # Extract app_roles claim from the token ## check for both 'roles' and 'app_roles' claims diff --git a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py index 348f5815be6..c5016b49487 100644 --- a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py +++ b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py @@ -158,7 +158,9 @@ def _typed_table(repo: DailyTagSpendRepository | VerificationTokenRepository | U async def _query_raw(prisma_client: "PrismaClient", sql_query: str, *params: object) -> object: - return await prisma_client.db.query_raw(sql_query, *params) + from litellm.proxy.utils import query_raw_rows + + return await query_raw_rows(prisma_client.db, sql_query, *params) @router.get( diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index a719a5c67d4..27d7cff0572 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -139,6 +139,14 @@ def _object_permission_mapping(raw: object) -> Mapping[str, object]: return mapping if mapping is not None and isinstance(decoded, dict) else {} +def _decoded_tool_permissions(raw: str) -> Mapping[str, object]: + decoded: Final = _DECODED_JSON.validate_python(json.loads(raw)) + mapping: Final = as_str_mapping(decoded) + if mapping is not None: + return mapping + raise AttributeError(f"{type(decoded).__name__!r} object has no attribute 'keys'", name="keys", obj=decoded) + + async def handle_update_object_permission_common( data_json: dict, existing_object_permission_id: str | None, @@ -460,10 +468,11 @@ async def _resolve_team_allowed_mcp_servers( access_group_servers: Final[list[str]] = await MCPRequestHandler.get_mcp_servers_from_access_groups( team_object_permission.mcp_access_groups or [] ) - raw_tool_perms = team_object_permission.mcp_tool_permissions or {} - if isinstance(raw_tool_perms, str): - raw_tool_perms = json.loads(raw_tool_perms) - tool_perm_servers: Final[list[str]] = list(raw_tool_perms.keys()) + stored_tool_perms: Final = team_object_permission.mcp_tool_permissions or {} + tool_perms: Final = ( + _decoded_tool_permissions(stored_tool_perms) if isinstance(stored_tool_perms, str) else stored_tool_perms + ) + tool_perm_servers: Final[list[str]] = list(tool_perms.keys()) raw_servers: Final = set(direct_servers + access_group_servers + tool_perm_servers) resolved_servers: Final = await _resolve_mcp_server_identifiers_to_ids( identifiers=raw_servers, @@ -555,7 +564,7 @@ async def _get_grandfathered_key_mcp_server_ids( return frozenset() raw_tool_perms: Final = existing_object_permission.mcp_tool_permissions or {} tool_perm_keys: Final[frozenset[str]] = frozenset( - json.loads(raw_tool_perms).keys() if isinstance(raw_tool_perms, str) else raw_tool_perms.keys() + _decoded_tool_permissions(raw_tool_perms).keys() if isinstance(raw_tool_perms, str) else raw_tool_perms.keys() ) identifiers: Final = (frozenset(existing_object_permission.mcp_servers or []) | tool_perm_keys) - { SpecialMCPServerNames.no_mcp_servers.value, diff --git a/litellm/proxy/management_helpers/team_list_sorting.py b/litellm/proxy/management_helpers/team_list_sorting.py index caddd8afc08..a3a35857328 100644 --- a/litellm/proxy/management_helpers/team_list_sorting.py +++ b/litellm/proxy/management_helpers/team_list_sorting.py @@ -6,7 +6,7 @@ from typing import Final from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import PrismaClient, query_raw_rows COMPUTED_TEAM_SORTS: Final = MappingProxyType( { @@ -105,5 +105,5 @@ async def sorted_team_ids( f"ORDER BY {expression} {direction} NULLS LAST, t.team_id ASC " f"LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}" ) - rows: Final = _TEAM_IDS.validate_python(await prisma_client.db.query_raw(sql, *params, take, skip)) + rows: Final = _TEAM_IDS.validate_python(await query_raw_rows(prisma_client.db, sql, *params, take, skip)) return tuple(row.team_id for row in rows) diff --git a/litellm/proxy/management_helpers/user_list_sorting.py b/litellm/proxy/management_helpers/user_list_sorting.py index 407a188ff1a..aa77417311f 100644 --- a/litellm/proxy/management_helpers/user_list_sorting.py +++ b/litellm/proxy/management_helpers/user_list_sorting.py @@ -4,7 +4,7 @@ from typing import Final, Literal from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm.constants import UI_SESSION_TOKEN_TEAM_ID -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import PrismaClient, query_raw_rows class _UserId(BaseModel): @@ -63,7 +63,8 @@ async def sorted_user_ids( LIMIT $9 OFFSET $10 """ rows: Final = _USER_IDS.validate_python( - await prisma_client.db.query_raw( + await query_raw_rows( + prisma_client.db, sql, role or None, ids if user_ids else None, diff --git a/litellm/proxy/mcp_server/guardrail_translation/handler.py b/litellm/proxy/mcp_server/guardrail_translation/handler.py index 28c51148da5..b268ba893f5 100644 --- a/litellm/proxy/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/mcp_server/guardrail_translation/handler.py @@ -19,7 +19,7 @@ payload (name + arguments) so we just build the tool definition. """ from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from fastapi import HTTPException from mcp.types import Tool as MCPTool @@ -139,7 +139,7 @@ class MCPGuardrailTranslationHandler(BaseTranslation): async def process_input_messages( self, - data: dict[str, Any], + data: dict[str, object], guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> dict[str, object]: @@ -154,10 +154,12 @@ class MCPGuardrailTranslationHandler(BaseTranslation): # Convert MCP input via transform_mcp_tool_to_openai_tool, then map to litellm # ChatCompletionToolParam (openai SDK type has incompatible strict/cache_control). - mcp_tool: Final = MCPTool( - name=mcp_tool_name, - description=mcp_tool_description or "", - input_schema=dict(mcp_input_schema) if isinstance(mcp_input_schema, Mapping) else {}, + mcp_tool: Final = MCPTool.model_validate( + { + "name": mcp_tool_name, + "description": mcp_tool_description or "", + "input_schema": dict(mcp_input_schema) if isinstance(mcp_input_schema, Mapping) else {}, + } ) openai_tool: Final = transform_mcp_tool_to_openai_tool(mcp_tool) fn: Final = openai_tool["function"] diff --git a/litellm/proxy/mcp_server/oauth2_flow_backfill.py b/litellm/proxy/mcp_server/oauth2_flow_backfill.py index 1ef67df498f..7b72ba2a26f 100644 --- a/litellm/proxy/mcp_server/oauth2_flow_backfill.py +++ b/litellm/proxy/mcp_server/oauth2_flow_backfill.py @@ -39,15 +39,17 @@ from collections import Counter from collections.abc import Mapping, Sequence from typing import Final, Literal, Protocol -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy.mcp_server.db import ( # noqa: F401 # legacy module exports _decode_oauth_payload, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export decode_oauth_payload, decrypt_credentials, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import MCPUserCredentialsRepository from litellm.types.mcp import MCPCredentials OAuth2Flow = Literal["client_credentials", "authorization_code"] @@ -60,6 +62,7 @@ BackfillRule = Literal[ ] _BACKFILL_AUDIT_ACTOR: Final = "oauth2_flow_backfill" +_DECODED_JSON: Final = TypeAdapter(object) class _MCPServerRow(Protocol): @@ -108,21 +111,21 @@ def _mcp_server_table(prisma_client: PrismaClient) -> _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 + return MCPUserCredentialsRepository(prisma_client).table def _decrypted_credentials(raw_credentials: str | Mapping[str, JsonValue] | None) -> MCPCredentials | None: if raw_credentials is None: return None - parsed: JsonValue | Mapping[str, JsonValue] + parsed: Mapping[str, object] | None if isinstance(raw_credentials, str): try: - parsed = json.loads(raw_credentials) + parsed = as_str_mapping(_DECODED_JSON.validate_python(json.loads(raw_credentials))) except (ValueError, TypeError): return None else: - parsed = raw_credentials - if not isinstance(parsed, dict): + parsed = raw_credentials if isinstance(raw_credentials, dict) else None + if parsed is None: return None return decrypt_credentials(credentials=dict(parsed)) diff --git a/litellm/proxy/mcp_server/oauth2_token_cache.py b/litellm/proxy/mcp_server/oauth2_token_cache.py index 1e6e83c0f4b..4bb61f6a5bb 100644 --- a/litellm/proxy/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/mcp_server/oauth2_token_cache.py @@ -43,6 +43,7 @@ if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer _DECODED_JSON: Final = TypeAdapter(object) +_CACHED_TOKEN: Final = TypeAdapter(object) class MCPOAuth2TokenCache(InMemoryCache): @@ -259,8 +260,8 @@ class MCPPerUserTokenCache: from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 key: Final = self._cache_key(user_id, server_id) - encrypted: Final = await user_api_key_cache.async_get_cache(key) - if encrypted is None: + encrypted: Final = _CACHED_TOKEN.validate_python(await user_api_key_cache.async_get_cache(key)) + if not isinstance(encrypted, str): return None return self._codec().decode(encrypted) except Exception as exc: diff --git a/litellm/proxy/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/mcp_server/openapi_to_mcp_generator.py index 0fa904cb582..7282e59d746 100644 --- a/litellm/proxy/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/mcp_server/openapi_to_mcp_generator.py @@ -45,6 +45,7 @@ OPENAPI_TOOL_NAME_MAX_LEN: Final = 128 _OPENAPI_TOOL_NAME_MAX_LEN: Final = OPENAPI_TOOL_NAME_MAX_LEN _OPENAPI_DOCUMENT: Final = TypeAdapter(dict[str, object]) +_DECODED_SPEC: Final = TypeAdapter(object) _MAX_YAML_INT_LENGTH: Final = 1024 # YAML 1.2 core scalars; isolate these from PyYAML's global YAML 1.1 resolver. _YAML_SCALARS: Final = { @@ -225,6 +226,22 @@ def load_openapi_spec(filepath: str) -> _OpenAPISpec: return asyncio.run(load_openapi_spec_async(filepath)) +def _has_nonspecific_tag(event: ScalarEvent) -> bool: + match event: + case ScalarEvent(tag="!"): + return True + case _: + return False + + +def _scalar_text(node: ScalarNode) -> str: + match node: + case ScalarNode(value=str() as text): + return text + case _: + raise ValueError("Invalid YAML OpenAPI spec") + + class _OpenAPISafeLoader(yaml.SafeLoader): yaml_implicit_resolvers: ClassVar[Mapping[str | None, Sequence[tuple[str, re.Pattern[str]]]]] = { # pyright: ignore[reportIncompatibleVariableOverride] # PyYAML resolver tables are class state None: list(_YAML_SCALARS.items()) @@ -236,7 +253,7 @@ class _OpenAPISafeLoader(yaml.SafeLoader): raise ValueError("YAML aliases are not supported in OpenAPI specs") # noqa: TRY004 # callers handle malformed documents as ValueError node: Final = super().compose_node(parent, index) # PyYAML otherwise implicitly resolves even an explicit nonspecific ! tag. - if isinstance(event, ScalarEvent) and event.tag == "!" and isinstance(node, ScalarNode): + if isinstance(event, ScalarEvent) and _has_nonspecific_tag(event) and isinstance(node, ScalarNode): node.tag = "tag:yaml.org,2002:str" return node @@ -245,7 +262,7 @@ class _OpenAPISafeLoader(yaml.SafeLoader): if pattern is not None: if not isinstance(node, ScalarNode): raise ValueError("Invalid YAML OpenAPI spec") - scalar: Final[str] = node.value + scalar: Final = _scalar_text(node) if node.tag == "tag:yaml.org,2002:int" and len(scalar) > _MAX_YAML_INT_LENGTH: raise ValueError("YAML integer is too long in OpenAPI spec") if pattern.fullmatch(scalar) is None: @@ -281,10 +298,10 @@ def _validate_openapi_spec(value: object) -> _OpenAPISpec: def _parse_openapi_spec(content: str | bytes) -> _OpenAPISpec: try: - return _validate_openapi_spec(json.loads(content)) + return _validate_openapi_spec(_DECODED_SPEC.validate_python(json.loads(content))) except json.JSONDecodeError: try: - return _validate_openapi_spec(yaml.load(content, Loader=_OpenAPISafeLoader)) + return _validate_openapi_spec(_DECODED_SPEC.validate_python(yaml.load(content, Loader=_OpenAPISafeLoader))) except (yaml.YAMLError, RecursionError) as exc: raise ValueError("Invalid YAML OpenAPI spec") from exc diff --git a/litellm/proxy/mcp_server/operations.py b/litellm/proxy/mcp_server/operations.py index 66e46c73f98..dfc45417814 100644 --- a/litellm/proxy/mcp_server/operations.py +++ b/litellm/proxy/mcp_server/operations.py @@ -49,7 +49,7 @@ from mcp.types import ( ) from mcp.types import Tool as MCPTool from pydantic import AnyUrl, ConfigDict, Field, JsonValue, SecretStr, TypeAdapter -from typing_extensions import ReadOnly, TypedDict, assert_never +from typing_extensions import ReadOnly, TypedDict, Unpack, assert_never from litellm._logging import verbose_logger from litellm.constants import ( @@ -562,6 +562,7 @@ async def _dispatch_virtual_mcp_tool( is_error=True, ) + assert user_api_key_auth is not None # guaranteed by the flag check above args: Final = arguments or {} if name == MCP_TOOL_SEARCH_TOOL_NAME: return await handle_mcp_tool_search( @@ -576,7 +577,6 @@ async def _dispatch_virtual_mcp_tool( raw_headers=raw_headers, ) - assert user_api_key_auth is not None # guaranteed by the flag check above if name == AGENT_SEARCH_TOOL_NAME: return await handle_agent_search( query=str(args.get("query", "")), @@ -2926,6 +2926,10 @@ async def resolve_scoped_mcp_server( return _single_scoped_server(mcp_servers, allowed_mcp_servers) +class _CallMcpToolOptions(TypedDict, total=False, extra_items=object): + litellm_logging_obj: ReadOnly[LiteLLMLoggingObj | None] + + @client @catalog_operation(lambda: global_mcp_server_manager) async def call_mcp_tool( @@ -2943,7 +2947,7 @@ async def call_mcp_tool( request_state: str | None = None, incoming_bearer_token: SecretStr | None | EllipsisType = ..., guardrail_auth: UserAPIKeyAuth | None = None, - **kwargs: Any, + **kwargs: Unpack[_CallMcpToolOptions], ) -> CallToolResult | InputRequiredResult: """ Call a specific tool with the provided arguments (handles prefixed tool names). diff --git a/litellm/proxy/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/mcp_server/outbound_credentials/sso_assertion_store.py index fb4d838aff9..65ab6f6c609 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/mcp_server/outbound_credentials/sso_assertion_store.py @@ -30,6 +30,7 @@ from pydantic import ConfigDict, SecretStr, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_SSO_ASSERTION_CACHE_TTL_SECONDS +from litellm.repositories.table_repositories import MCPServerRepository, PrismaTableRepository from litellm.types.llms.base import LiteLLMBaseModel if TYPE_CHECKING: @@ -68,14 +69,18 @@ class _MCPServerTable(Protocol): async def find_first(self, *, where: Mapping[str, str]) -> object | None: ... +class _SSOAssertionRepository(PrismaTableRepository["LiteLLM_SSOIdentityAssertion"]): + table_name = "litellm_ssoidentityassertion" + + def _assertion_table(prisma_client: PrismaClient) -> _SSOAssertionTable: """The SSO assertion table, typed so the untyped prisma client surface stops here.""" - return prisma_client.db.litellm_ssoidentityassertion + return _SSOAssertionRepository(prisma_client).table 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 + return MCPServerRepository(prisma_client).table class SSOIdentityAssertion(LiteLLMBaseModel): diff --git a/litellm/proxy/mcp_server/utils.py b/litellm/proxy/mcp_server/utils.py index 3f72804ff35..92bea6053b0 100644 --- a/litellm/proxy/mcp_server/utils.py +++ b/litellm/proxy/mcp_server/utils.py @@ -16,6 +16,7 @@ from urllib.parse import quote from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.types.mcp_server.mcp_server_manager import MCPServer if typing.TYPE_CHECKING: @@ -702,8 +703,12 @@ class MCPMissingUserEnvVarsError(Exception): _ENV_VAR_PATTERN: Final = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}") +class _ModelDumpable(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + def parse_admin_env_vars( - env_vars: Iterable[Any] | None, + env_vars: Iterable[object] | None, ) -> tuple[dict[str, str], list[dict[str, Any]]]: """Split admin-configured env var entries into globals and per-user specs. @@ -723,7 +728,7 @@ def parse_admin_env_vars( for raw in env_vars: if raw is None: continue - if hasattr(raw, "model_dump"): + if has_attributes(raw, _ModelDumpable, "model_dump"): entry: Mapping[str, object] = raw.model_dump() elif isinstance(raw, dict): entry = raw diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index a6eee889fea..1ae0d475d8c 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -23,6 +23,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.constants import MAX_FILE_LIST_LIMIT from litellm.litellm_core_utils.batch_utils import batch_cost_is_final +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import ProxyException from litellm.repositories.managed_file_repository import ManagedFileRepository from litellm.repositories.table_repositories import ( @@ -597,7 +598,7 @@ def get_team_provider_credentials( return isinstance(public_model_name, str) and _key_allowlist_covers(public_model_name, key_model_allowlist) def _provider_credentials(model_id: str) -> dict[str, object] | None: - credentials: Final = registry.credentials(model_id, form="with_provider", team_id=team_id) + credentials: Final = as_str_mapping(registry.credentials(model_id, form="with_provider", team_id=team_id)) if credentials is not None and credentials.get("custom_llm_provider") == custom_llm_provider: return {key: value for key, value in credentials.items() if key != "model"} return None @@ -1545,6 +1546,10 @@ async def get_batch_from_database( return None, None +class _BatchCostJobScheduler(Protocol): + def get_job(self, job_id: str, /) -> object: ... + + def batch_cost_poller_is_active() -> bool: """ Whether the CheckBatchCost poller will account for a managed batch's cost itself. @@ -1565,10 +1570,10 @@ def batch_cost_poller_is_active() -> bool: try: import litellm.proxy.proxy_server as proxy_server_module - scheduler = getattr(proxy_server_module, "scheduler", None) + scheduler: Final[_BatchCostJobScheduler | None] = getattr(proxy_server_module, "scheduler", None) if scheduler is None: return False - job = scheduler.get_job("check_batch_cost_job") + job: Final = scheduler.get_job("check_batch_cost_job") if job is None: return False poller = getattr(getattr(job, "func", None), "__self__", None) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index aabb83b3a5e..e4e3fec855d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -131,6 +131,7 @@ def _responses_sub_route(path: str) -> tuple[str, ...] | None: _STR_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) class _AgentConfig(BaseModel): @@ -550,8 +551,11 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): ) = None handler_instance: Final = OpenAIPassthroughLoggingHandler() + logged_custom_llm_provider: Final = _DECODED_JSON.validate_python( + logging_obj.model_call_details.get("custom_llm_provider", "openai") + ) custom_llm_provider: Final = TypeAdapter(str).validate_python( - kwargs.get("custom_llm_provider", logging_obj.model_call_details.get("custom_llm_provider", "openai")) + kwargs.get("custom_llm_provider", logged_custom_llm_provider) ) if is_chat_completions: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 7f5983d5799..8111e46c0b7 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -68,6 +68,7 @@ from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, + require_attributes, ) from litellm.litellm_core_utils.initialize_dynamic_callback_params import validate_no_callback_env_reference from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY @@ -194,6 +195,13 @@ def _is_mergeable_client_metadata(key: str, value: object) -> bool: return not key.startswith("user_api_key_") and (key != "model_info" or isinstance(value, Mapping)) +def _request_body_fields(parsed: object) -> Mapping[str, object]: + fields: Final = as_str_mapping(parsed) + if fields is None: + raise TypeError(f"unsupported operand type(s) for |: '{type(parsed).__name__}' and 'dict'") + return fields + + class _RegisteredPassThroughRoute(TypedDict): endpoint_id: ReadOnly[str] path: ReadOnly[str | None] @@ -287,9 +295,9 @@ async def chat_completion_pass_through_endpoint( body: Final = await request.body() body_str: Final = body.decode() try: - data = ast.literal_eval(body_str) | data + data = {**_request_body_fields(_DECODED_JSON.validate_python(ast.literal_eval(body_str))), **data} except Exception: - data = json.loads(body_str) | data + data = {**_request_body_fields(_DECODED_JSON.validate_python(json.loads(body_str))), **data} data["adapter_id"] = adapter_id @@ -2237,15 +2245,13 @@ def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> di return restored_metadata -class _PassThroughRequestEnvelope(TypedDict, total=False): - query_params: Mapping[str, object] | None - custom_body: Mapping[str, object] | None - stream: bool | None +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... async def _parse_request_data_by_content_type( request: Request, -) -> tuple[object, object, None, bool | None]: +) -> tuple[object, object, None, object]: """ Parse request data based on content type. @@ -2264,10 +2270,12 @@ async def _parse_request_data_by_content_type( if "application/json" in content_type: # ✅ Handle JSON try: - body: _PassThroughRequestEnvelope = await request.json() - query_params_data = body.get("query_params") - custom_body_data = body.get("custom_body") - stream = body.get("stream") + json_body: Final = require_attributes( + _DECODED_JSON.validate_python(await request.json()), _SupportsGet, "get" + ) + query_params_data = json_body.get("query_params") + custom_body_data = json_body.get("custom_body") + stream = json_body.get("stream") except json.JSONDecodeError: # Handle requests with no body (e.g., DELETE requests) pass @@ -2275,14 +2283,16 @@ async def _parse_request_data_by_content_type( # ✅ Try to parse as JSON first (handles misconfigured clients sending JSON with multipart content-type) # If that fails, skip parsing - pass_through_request will handle actual multipart try: - body = await request.json() + multipart_json_body: Final = require_attributes( + _DECODED_JSON.validate_python(await request.json()), _SupportsGet, "get" + ) # Successfully parsed as JSON - treat as JSON body - query_params_data = body.get("query_params") - custom_body_data = body.get("custom_body") - stream = body.get("stream") + query_params_data = multipart_json_body.get("query_params") + custom_body_data = multipart_json_body.get("custom_body") + stream = multipart_json_body.get("stream") # If custom_body is not set, use the entire body - if custom_body_data is None and body: - custom_body_data = body + if custom_body_data is None and multipart_json_body: + custom_body_data = multipart_json_body except (json.JSONDecodeError, Exception): # Not JSON - this is actual multipart data # Skip parsing here to avoid consuming the request body stream diff --git a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py index bd8390d704e..07d4f7dae93 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py @@ -7,7 +7,7 @@ Handles guardrail execution for passthrough endpoints with: - Automatic inheritance from org/team/key levels when enabled """ -from typing import Any, Final +from typing import Final from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import merged_guardrail_entries @@ -243,14 +243,15 @@ class PassthroughGuardrailHandler: verbose_proxy_logger.debug("Passthrough guardrails config is empty, skipping") return None - temp_data: Final[dict[str, Any]] = {"metadata": {}} + inherited_metadata: Final = dict[str, list[object]]() + temp_data: Final = {"metadata": inherited_metadata} add_guardrails_from_key_or_team_metadata( key_metadata=user_api_key_dict.metadata, team_metadata=user_api_key_dict.team_metadata, data=temp_data, metadata_variable_name="metadata", ) - inherited_guardrails: Final[list[object]] = temp_data["metadata"].get("guardrails", []) + inherited_guardrails: Final = inherited_metadata.get("guardrails", []) guardrails_to_run: Final = merged_guardrail_entries(normalized_config, inherited_guardrails) verbose_proxy_logger.debug("Collected guardrails for passthrough endpoint: %s", guardrails_to_run) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 0c89627e74d..41aedf4b5cd 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -264,6 +264,7 @@ WORKER_APP: Final = "litellm_cli.main:proxy_worker_app" WORKER_LOG_CONFIG_ENV: Final = "LITELLM_WORKER_LOG_CONFIG" _LOG_CONFIG: Final = TypeAdapter[str | Mapping[str, object]](str | Mapping[str, object]) _LOG_CONFIG_FILE: Final = TypeAdapter[Mapping[str, object]](Mapping[str, object]) +_DECODED_JSON: Final = TypeAdapter[object](object) class ProxyInitializationHelpers: @@ -1332,8 +1333,7 @@ def run_server( ProxyInitializationHelpers.run_test_chat_completion(host, port, model, test) return else: - if headers: - headers = json.loads(headers) + worker_headers: Final = _DECODED_JSON.validate_python(json.loads(headers) if headers else headers) save_worker_config( model=model, alias=alias, @@ -1347,7 +1347,7 @@ def run_server( max_budget=max_budget, drop_params=drop_params, add_function_to_prompt=add_function_to_prompt, - headers=headers, + headers=worker_headers, save=save, config=config, use_queue=use_queue, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d21f5952ba2..e015740821d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -162,6 +162,7 @@ from litellm.proxy._types import ( ProxyException, ProxyLifespanState, SpecialModelNames, + SSOUserDefinedValues, SupportedDBObjectType, TeamDefaultSettings, TokenCountRequest, @@ -1136,7 +1137,6 @@ from litellm.types.utils import ( SelectTokenizerResponse, StandardLoggingPayload, ) -from litellm.types.utils import ModelInfo as ModelMapInfo from litellm.utils import add_custom_logger_callback_to_specific_event try: @@ -3200,16 +3200,23 @@ def _apply_background_response_settings(value: SettingsJsonValue) -> None: ) -user_custom_auth = None -user_custom_key_generate = None +_CustomKeyHook: TypeAlias = Callable[..., Awaitable[Mapping[str, object]]] + + +class _CustomSsoHandler(Protocol): + async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ... + + +user_custom_auth: Callable[..., Awaitable[object]] | None = None +user_custom_key_generate: _CustomKeyHook | None = None # Sentinel: prevents PKCE-no-Redis advisory from re-logging on config hot-reload. # Tests that need to reset it can patch 'litellm.proxy.proxy_server._pkce_no_redis_warning_emitted'. _pkce_no_redis_warning_emitted: bool = False _cp_no_redis_warning_emitted: bool = False -user_custom_key_update = None -user_custom_key_policy = None -user_custom_sso = None -user_custom_ui_sso_sign_in_handler = None +user_custom_key_update: _CustomKeyHook | None = None +user_custom_key_policy: _CustomKeyHook | None = None +user_custom_sso: _CustomSsoHandler | None = None +user_custom_ui_sso_sign_in_handler: object | None = None use_background_health_checks = None use_shared_health_check = None use_queue = False @@ -4546,9 +4553,13 @@ async def _add_spend_to_shared_global_proxy_spend(cache: UserApiKeyCache, respon async def _add_spend_to_global_proxy_spend(cache: UserApiKeyCache, response_cost: float) -> None: - local_before: Final = cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) + local_before: Final = _UNTYPED_VALUE.validate_python( + cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) + ) shared_total: Final = await _add_spend_to_shared_global_proxy_spend(cache, response_cost) - local_total: Final = cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) + local_total: Final = _UNTYPED_VALUE.validate_python( + cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) + ) if not isinstance(local_total, (int, float)): return if shared_total is not None and local_total != local_before and shared_total < local_total: @@ -5879,8 +5890,16 @@ class _DbCallbackRegistration: displaced: tuple[tuple[str, object], ...] = () +class _CallbackList(Protocol): + def __contains__(self, value: object, /) -> bool: ... + + def append(self, value: object, /) -> None: ... + + def _restore_displaced_callback(list_name: str, entry: object) -> None: - callback_list: Final = getattr(litellm, list_name) + callback_list: Final = require_attributes( + _UNTYPED_VALUE.validate_python(getattr(litellm, list_name)), _CallbackList, "append" + ) if entry not in callback_list: callback_list.append(entry) @@ -7375,9 +7394,11 @@ class ProxyConfig: custom_ui_sso_sign_in_handler: Final = general_settings.get("custom_ui_sso_sign_in_handler", None) if custom_ui_sso_sign_in_handler is not None: - user_custom_ui_sso_sign_in_handler = get_instance_fn( - value=custom_ui_sso_sign_in_handler, - config_file_path=config_file_path, + user_custom_ui_sso_sign_in_handler = _UNTYPED_VALUE.validate_python( + get_instance_fn( + value=custom_ui_sso_sign_in_handler, + config_file_path=config_file_path, + ) ) if enterprise_proxy_config is not None: @@ -8679,7 +8700,7 @@ class ProxyConfig: environment_read: Final = next(self._environment_read_numbers) sections: Final = tuple(self._settings_stores) responses: Final = await asyncio.gather(*(get_config_param(prisma_client, section) for section in sections)) - general_settings_response: Final = responses[sections.index("general_settings")] + general_settings_response: Final[_ConfigParamValueRow | None] = responses[sections.index("general_settings")] general_settings_value: Final[Mapping[str, SettingsJsonValue]] = ( self._prepared_db_settings_values("general_settings", general_settings_response.param_value) if general_settings_response is not None and general_settings_response.param_value is not None @@ -11036,7 +11057,7 @@ async def async_data_generator( await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( request=request, request_data=request_data, - response=response, + response=_UNTYPED_VALUE.validate_python(response), stream_completed=stream_completed, client_disconnected=client_disconnected, user_api_key_dict=user_api_key_dict, @@ -16020,10 +16041,9 @@ async def _count_request_tokens( deployment: dict[str, Any] | None = None if routed is None else dict(routed) litellm_model_name = None - model_info: ModelMapInfo | None = None + model_info: Final = None if routed is None else routed.get("model_info", {}) if deployment is not None: litellm_model_name = deployment.get("litellm_params", {}).get("model") - model_info = deployment.get("model_info", {}) deployment = { # rebind-ok: merge the named credential into a copy **deployment, "litellm_params": with_named_credential( @@ -16201,7 +16221,7 @@ async def _count_request_tokens( if model_info is not None: custom_tokenizer = cast( CustomHuggingfaceTokenizer | None, - model_info.get("custom_tokenizer", None), + _with_get(model_info).get("custom_tokenizer", None), ) system_message: Final = _system_message(local_system) prompt_turn: Final = prompt_only and not (system_message is None and not tools) @@ -20722,7 +20742,7 @@ async def update_config_general_settings( detail={"error": CommonProxyErrors.not_allowed_access.value}, ) - requested_value: Final[object] = data.field_value + requested_value: Final = _UNTYPED_VALUE.validate_python(data.field_value) if data.field_name in _GENERAL_SETTINGS_UI_LITELLM_FIELDS: return await _persist_general_settings_ui_litellm_field(data.field_name, requested_value, user_api_key_dict) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 8fd4fbfcb8b..df156ccedd6 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -21,7 +21,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy.db.db_span import db_span, db_spanned -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import PrismaClient, query_raw_rows from litellm.repositories.chunked_in import find_many_in from litellm.repositories.user_repository import UserRepository from litellm.types.llms.base import LiteLLMBaseModel @@ -207,7 +207,9 @@ async def _reverse_hash_key_metadata( warning: str, ) -> Mapping[str, KeyMetadataDict]: rows: Final = await _db_or_empty( - lambda: db_spanned("recover_key_metadata", table, lambda: prisma_client.db.query_raw(sql, sorted(wanted))), + lambda: db_spanned( + "recover_key_metadata", table, lambda: query_raw_rows(prisma_client.db, sql, sorted(wanted)) + ), warning, len(wanted), ) diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index 64b67ec55dd..e4a98619bfd 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -33,6 +33,7 @@ from litellm.constants import ( PTU_ROLLUP_MAX_BACKFILL_DAYS, PTU_SENTINEL_API_KEY, ) +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.ptu_pricing import ptu_terms from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled @@ -79,6 +80,10 @@ class _SentinelBatchManager(Protocol): async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + def _sentinel_batch(prisma_client: "PrismaClient") -> _SentinelBatchManager: batch: Final[_SentinelBatchManager] = prisma_client.db.batch_() return batch @@ -171,7 +176,7 @@ class _PTUDeployment: model_info: Mapping[str, object] -def _router_deployment(deployment: Mapping[str, object]) -> _PTUDeployment | None: +def _router_deployment(deployment: _SupportsGet) -> _PTUDeployment | None: """A router ``model_list`` entry in the shape the parser reads, else None. An id is required rather than defaulted because it keys the sentinel row: every @@ -421,8 +426,8 @@ def _config_deployments(router: "Router | None", *, owned_by_db: frozenset[str]) credential clone carries ``original_model_id`` and reuses its source's PTU config under a fresh id, so pricing it would bill one reservation once per distinct client key. """ - entries: Final = tuple((router.registry().all() if router is not None else None) or ()) - records: Final = tuple(_router_deployment(entry) for entry in entries) + entries: Final[tuple[object, ...]] = tuple((router.registry().all() if router is not None else None) or ()) + records: Final = tuple(_router_deployment(require_attributes(entry, _SupportsGet, "get")) for entry in entries) return tuple( record for record in records diff --git a/litellm/proxy/spend_tracking/spend_capture_rate.py b/litellm/proxy/spend_tracking/spend_capture_rate.py index 443010b65da..0e6bed9017e 100644 --- a/litellm/proxy/spend_tracking/spend_capture_rate.py +++ b/litellm/proxy/spend_tracking/spend_capture_rate.py @@ -89,8 +89,14 @@ async def captured_spend_by_day( end_date: date, ) -> Mapping[str, float]: """LiteLLM's tracked spend per UTC day (ISO date) for the given ``custom_llm_provider`` values.""" - rows: Final = await prisma_client.db.query_raw( - _CAPTURED_SPEND_BY_DAY_SQL, start_date.isoformat(), end_date.isoformat(), tuple(litellm_providers) + from litellm.proxy.utils import query_raw_rows + + rows: Final = await query_raw_rows( + prisma_client.db, + _CAPTURED_SPEND_BY_DAY_SQL, + start_date.isoformat(), + end_date.isoformat(), + tuple(litellm_providers), ) return MappingProxyType({row.date: row.spend for row in _CAPTURED_SPEND_ROWS.validate_python(rows)}) diff --git a/litellm/proxy/types_utils/utils.py b/litellm/proxy/types_utils/utils.py index 62b87dbbf36..4f9b7787682 100644 --- a/litellm/proxy/types_utils/utils.py +++ b/litellm/proxy/types_utils/utils.py @@ -206,8 +206,8 @@ async def _download_gcs_file_wrapper(bucket_name: str, object_key: str, local_fi def validate_custom_validate_return_type( - fn: Callable[..., Any] | None, -) -> Callable[..., Literal[True]] | None: + fn: Callable[..., object] | None, +) -> Callable[..., object] | None: if fn is None: return None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8e977c69557..d7c7045cbe1 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -162,10 +162,12 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import add_langfuse_trace_id_to_alert from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.litellm_core_utils.core_helpers import ( + attribute_or_none, coerce_token_limit, get_or_create_metadata_bucket, independent_snapshot, is_expected_client_error, + require_attributes, ) from litellm.litellm_core_utils.deployment_pricing import pricing_models from litellm.litellm_core_utils.duration_parser import get_budget_window_start @@ -315,7 +317,6 @@ if TYPE_CHECKING: from mcp.types import CallToolResult from opentelemetry.trace import Span as _Span from prisma import models as prisma_models - from prisma.actions import LiteLLM_DeprecatedVerificationTokenActions from prisma.client import TransactionManager from prisma.errors import UniqueViolationError from prisma.models import LiteLLM_DeprecatedVerificationToken @@ -364,6 +365,7 @@ _VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0 _VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"' _VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...]) _DELEGATED_ATTRIBUTE: Final = TypeAdapter(object) +_DECODED_JSON: Final = TypeAdapter(object) _ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"] _ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception @@ -405,7 +407,7 @@ def _update_many_batch(db: PrismaWrapper | RoutingPrismaWrapper) -> _UpdateManyB return db.batch_() -def _query_raw_rows( +def query_raw_rows( db: PrismaWrapper | RoutingPrismaWrapper, query: str, *args: object ) -> Awaitable[Sequence[Mapping[str, object]]]: return db.query_raw(query, *args) @@ -645,6 +647,24 @@ class _StreamIteratorHook(Protocol[_T]): def __call__(self, *, response: AsyncIterator[_T]) -> AsyncIterator[_T]: ... +class _BudgetAlertEmailer(Protocol): + async def budget_alerts( + self, + type: Literal[ + "token_budget", + "user_budget", + "soft_budget", + "max_budget_alert", + "team_budget", + "organization_budget", + "proxy_budget", + "projected_limit_exceeded", + "project_budget", + ], + user_info: CallInfo, + ) -> object: ... + + def _is_client_error_exception(exc: Exception) -> bool: if isinstance(exc, HTTPException): return exc.status_code < 500 @@ -1941,7 +1961,7 @@ class ProxyLogging: alerting=self.alerting, internal_usage_cache=self.internal_usage_cache.dual_cache, ) - self.email_logging_instance: Any | None = None + self.email_logging_instance: _BudgetAlertEmailer | None = None if BaseEmailLogger is not None: email_logger_class: Final = _get_email_logger_class() if email_logger_class is not None: @@ -5362,9 +5382,9 @@ async def _lookup_deprecated_key( _deprecated_key_cache.pop(hashed_token, None) try: - deprecated_keys_table: Final[ - LiteLLM_DeprecatedVerificationTokenActions[LiteLLM_DeprecatedVerificationToken] - ] = db.litellm_deprecatedverificationtoken + deprecated_keys_table: Final[TableActions[LiteLLM_DeprecatedVerificationToken]] = ( + db.litellm_deprecatedverificationtoken + ) deprecated_row: Final = await deprecated_keys_table.find_first( where={ "token": hashed_token, @@ -5415,11 +5435,20 @@ def _config_cache_key(param_name: str) -> str: return f"litellm_config:param:{param_name}" -def _pack_config_row(row: Any) -> dict[str, object]: +class _PackableConfigRow(Protocol): + @property + def param_name(self) -> object: ... + + @property + def param_value(self) -> object: ... + + +def _pack_config_row(row: object) -> dict[str, object]: + config_row: Final = require_attributes(row, _PackableConfigRow, "param_name", "param_value") return { - "param_name": row.param_name, - "param_value": row.param_value, - "reload_revision": getattr(row, "reload_revision", None), + "param_name": config_row.param_name, + "param_value": config_row.param_value, + "reload_revision": attribute_or_none(row, "reload_revision"), } @@ -6247,7 +6276,7 @@ class PrismaClient: LIMIT $1 OFFSET $2 """ - response = await _query_raw_rows(self.db, sql_query, limit, offset) + response = await query_raw_rows(self.db, sql_query, limit, offset) return response elif table_name == "spend": verbose_proxy_logger.debug("PrismaClient: get_data: table_name == 'spend'") @@ -7557,7 +7586,7 @@ class PrismaClient: # caller's own retried query, a stronger check than SELECT 1, # and a reader that fails to come back sets `_reader_unavailable` # so reads fall through to the writer just recreated here. - await _query_raw_rows(self.writer_db, "SELECT 1") + await query_raw_rows(self.writer_db, "SELECT 1") # A recreate can decline: the optimistic-lock guard no-ops when # the writer generation moved since cycle entry, and the routing # wrapper then leaves the reader untouched as well. Callers that @@ -7907,7 +7936,7 @@ class PrismaClient: async def _view_setup_gate_table_present(self) -> bool: rows: Final = _VIEW_SETUP_GATE_PROBE_ROWS.validate_python( - await _query_raw_rows(self.db, "SELECT to_regclass($1) IS NOT NULL AS present", _VIEW_SETUP_GATE_TABLE) + await query_raw_rows(self.db, "SELECT to_regclass($1) IS NOT NULL AS present", _VIEW_SETUP_GATE_TABLE) ) return rows[0]["present"] @@ -7917,7 +7946,7 @@ class PrismaClient: await asyncio.sleep(self._db_health_watchdog_interval_seconds) async with db_span("db_health_watchdog", None): await asyncio.wait_for( - _query_raw_rows(self.db, "SELECT 1"), + query_raw_rows(self.db, "SELECT 1"), timeout=self._db_health_watchdog_probe_timeout_seconds, ) if isinstance(self.db, RoutingPrismaWrapper) and self.db.writer_unavailable: @@ -7977,7 +8006,7 @@ class PrismaClient: async def _writer_is_read_only(self, writer: PrismaWrapper) -> bool: """True iff the pooled writer session answers reads but rejects writes (SQLSTATE 25006).""" rows: Final = _WRITER_WRITABILITY_PROBE_ROWS.validate_python( - await _query_raw_rows(writer, _WRITER_WRITABILITY_PROBE_SQL) + await query_raw_rows(writer, _WRITER_WRITABILITY_PROBE_SQL) ) return any(row.get("transaction_read_only") == "on" for row in rows) @@ -8004,7 +8033,7 @@ class PrismaClient: """ sql_query: Final = "SELECT 1" async with db_span("health_check", None): - response: Final[object] = await _query_raw_rows(wrapper, sql_query) + response: Final[object] = await query_raw_rows(wrapper, sql_query) return response async def _probe_answers_now(self, wrapper: PrismaWrapper) -> bool: @@ -8172,12 +8201,12 @@ class PrismaClient: verbose_proxy_logger.warning("Invalid response_time_ms value: %s", response_time_ms) return None - def _clean_details(self, details: dict | None) -> dict | None: + def _clean_details(self, details: dict | None) -> object: """Clean and validate details JSON""" if not isinstance(details, dict): return None try: - return safe_json_loads(safe_dumps(details)) + return _DECODED_JSON.validate_python(safe_json_loads(safe_dumps(details))) except Exception as e: verbose_proxy_logger.warning("Failed to clean details JSON: %s", e) return None diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index b6a0ebe3c58..c0d7ab8d6e1 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -1,18 +1,21 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import ( + TYPE_CHECKING, Annotated, - Any, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict Final, + Protocol, cast, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict ) from fastapi import APIRouter, Depends, HTTPException, Request, Response +from pydantic import TypeAdapter from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) +from litellm.litellm_core_utils.core_helpers import as_object_list, decoded_json_object, require_attributes from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -30,8 +33,25 @@ from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse from litellm.vector_stores.utils import VectorStoreWriteBody, callers_metadata_snapshot from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry +if TYPE_CHECKING: + from litellm import Router + router: Final = APIRouter() +_DECODED_JSON: Final = TypeAdapter(object) + + +class _CreatesManagedVectorStores(Protocol): + async def acreate_vector_store( + self, + create_request: Mapping[str, object], + llm_router: "Router", + target_model_names_list: Sequence[object], + litellm_parent_otel_span: object, + user_api_key_dict: UserAPIKeyAuth, + ) -> object: ... + + BLOCKED_QUERY_EMBEDDING_SELECTION_PARAMS: Final = frozenset( { "embedding_model", @@ -247,22 +267,23 @@ async def vector_store_create( ) # Check for target_model_names parameter - target_model_names: Final = data.pop("target_model_names", None) + target_model_names: Final = _DECODED_JSON.validate_python(data.pop("target_model_names", None)) if target_model_names: # Use managed vector stores for multi-model support - if isinstance(target_model_names, str): - target_model_names_list = [m.strip() for m in target_model_names.split(",")] - elif isinstance(target_model_names, list): - target_model_names_list = target_model_names - else: + target_model_names_list: Final = ( + [m.strip() for m in target_model_names.split(",")] + if isinstance(target_model_names, str) + else as_object_list(target_model_names) + ) + if target_model_names_list is None: raise HTTPException( status_code=400, detail="target_model_names must be a comma-separated string or list of model names", ) # Get managed vector stores hook - managed_vector_stores: Final[Any] = proxy_logging_obj.get_proxy_hook("managed_vector_stores") + managed_vector_stores: Final = proxy_logging_obj.get_proxy_hook("managed_vector_stores") if managed_vector_stores is None: raise HTTPException( status_code=500, @@ -276,8 +297,10 @@ async def vector_store_create( ) # Create vector store across multiple models - response: Final[object] = await managed_vector_stores.acreate_vector_store( - create_request=data, + response: Final = await require_attributes( + managed_vector_stores, _CreatesManagedVectorStores, "acreate_vector_store" + ).acreate_vector_store( + create_request=decoded_json_object(_DECODED_JSON.validate_python(data)), llm_router=llm_router, target_model_names_list=target_model_names_list, litellm_parent_otel_span=user_api_key_dict.parent_otel_span, diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 6bfccd2426b..4b9547300db 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -308,9 +308,7 @@ async def new_vector_store( # Extract and validate metadata metadata: Final = vector_store.get("vector_store_metadata") - validated_metadata: dict | None = None - if metadata is not None and isinstance(metadata, dict): - validated_metadata = metadata + validated_metadata: Final = metadata if isinstance(metadata, dict) else None new_vector_store: Final = await create_vector_store_in_db( vector_store_id=vector_store_id, @@ -382,7 +380,6 @@ async def list_vector_stores( from litellm.proxy.proxy_server import prisma_client vector_store_map: Final[dict[str, LiteLLM_ManagedVectorStore]] = {} - db_vector_store_ids: Final[set] = set() try: # Get vector stores from database first (source of truth) @@ -393,7 +390,7 @@ async def list_vector_stores( vector_store_id = vector_store.get("vector_store_id", None) if vector_store_id: vector_store_map[vector_store_id] = vector_store - db_vector_store_ids.add(vector_store_id) + db_vector_store_ids: Final = frozenset(vector_store_map) # Process in-memory vector stores if litellm.vector_store_registry is not None: @@ -647,8 +644,9 @@ async def update_vector_store( ) # Handle metadata serialization - if update_data.get("vector_store_metadata") is not None: - update_data["vector_store_metadata"] = safe_dumps(update_data["vector_store_metadata"]) + submitted_metadata: Final[object] = update_data.get("vector_store_metadata") + if submitted_metadata is not None: + update_data["vector_store_metadata"] = safe_dumps(submitted_metadata) # Handle litellm_params if provided. As with the create path, the # embedding-config auto-resolve previously persisted cleartext diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index 9d8318d5a83..696ad4ab170 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -1,7 +1,7 @@ from typing import TYPE_CHECKING, Any, Final import orjson -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm.types.videos.main import DecodedVideoId, VideoObject from litellm.types.videos.utils import ( @@ -14,6 +14,8 @@ from litellm.types.videos.utils import ( if TYPE_CHECKING: from litellm.router import Router +_DECODED_JSON: Final = TypeAdapter(object) + def extract_model_from_target_model_names(target_model_names: object) -> str | None: if isinstance(target_model_names, str): @@ -30,7 +32,7 @@ def video_reference_to_id(video_ref: object) -> str: if not isinstance(video_ref, str): return "" try: - parsed_ref: Final = orjson.loads(video_ref) + parsed_ref: Final = _DECODED_JSON.validate_python(orjson.loads(video_ref)) except orjson.JSONDecodeError: return video_ref return parsed_ref.get("id", "") if isinstance(parsed_ref, dict) else video_ref @@ -44,7 +46,7 @@ def get_custom_provider_from_data(data: dict[str, Any]) -> str | None: extra_body = data.get("extra_body") if isinstance(extra_body, str): try: - parsed_extra_body: Final = orjson.loads(extra_body) + parsed_extra_body: Final = _DECODED_JSON.validate_python(orjson.loads(extra_body)) if isinstance(parsed_extra_body, dict): extra_body = parsed_extra_body except Exception: diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 77bf6868bf1..2d86f4344db 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -36,6 +36,7 @@ from litellm.types.rag import ( RAGIngestResponse, ) from litellm.types.utils import ModelResponse +from litellm.types.vector_stores import VectorStoreSearchResponse from litellm.utils import client if TYPE_CHECKING: @@ -257,7 +258,7 @@ async def _execute_query_pipeline( {**provider_search_params, **kwargs, **filter_search_params, **store_search_params} ) with internal_sub_call(parent_request_id): - search_response: Final = await litellm.vector_stores.asearch( + search_response: Final[VectorStoreSearchResponse] = await litellm.vector_stores.asearch( vector_store_id=retrieval_config["vector_store_id"], query=query_text, max_num_results=retrieval_config.get("top_k", 10), diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index 9f608dbc99a..33994133ac0 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -1,5 +1,5 @@ from collections.abc import Sequence -from typing import Any, Final +from typing import Any, Final, TypeVar from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -7,6 +7,8 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessag from litellm.types.utils import ModelResponse from litellm.types.vector_stores import VectorStoreSearchResponse +_ItemT: Final = TypeVar("_ItemT") + class _ResultContentView(TypedDict): """Content entry carried by a vector store search result.""" @@ -22,12 +24,13 @@ class _SearchResultView(TypedDict): text: ReadOnly[NotRequired[str]] -class _SearchDataView(TypedDict): - results: ReadOnly[Sequence[_SearchResultView]] +class _StoreContentView(TypedDict): + type: ReadOnly[NotRequired[str | None]] + text: ReadOnly[NotRequired[str | None]] -class _ContextChunksView(TypedDict): - chunks: ReadOnly[Sequence[_SearchResultView | str | None]] +class _StoreResultView(TypedDict): + content: ReadOnly[NotRequired[Sequence[_StoreContentView] | None]] class _RerankResultView(TypedDict): @@ -42,6 +45,18 @@ class _MessageView(TypedDict): message: ReadOnly[object] +def _iterable_items(items: Sequence[_ItemT] | None) -> Sequence[_ItemT]: + if items is None: + raise TypeError("'NoneType' object is not iterable") + return items + + +def _sized_items(items: Sequence[_ItemT] | None) -> Sequence[_ItemT]: + if items is None: + raise TypeError("object of type 'NoneType' has no len()") + return items + + class RAGQuery: CONTENT_PREFIX_STRING = "Context:\n\n" @@ -70,16 +85,17 @@ class RAGQuery: return None @staticmethod - def build_context_message(context_chunks: Sequence[_SearchResultView | str | None]) -> ChatCompletionUserMessage: + def build_context_message( + context_chunks: Sequence[_StoreResultView | _SearchResultView | str | None] | None, + ) -> ChatCompletionUserMessage: """ Process search results and build a context message. """ context_content = RAGQuery.CONTENT_PREFIX_STRING - chunks: Final[_ContextChunksView] = {"chunks": context_chunks} - for chunk in chunks["chunks"]: + for chunk in _iterable_items(context_chunks): if isinstance(chunk, dict): - result_content: Sequence[_ResultContentView] | None = chunk.get("content") + result_content: Sequence[_StoreContentView | _ResultContentView] | None = chunk.get("content") if result_content: for content_item in result_content: content_text: str | None = content_item.get("text") @@ -123,27 +139,27 @@ class RAGQuery: @staticmethod def extract_documents_from_search( - search_response: Any, + search_response: VectorStoreSearchResponse, ) -> list[str | dict[str, object]]: """Extract text documents from vector store search response.""" documents: Final[list[str | dict[str, object]]] = [] - 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"): - documents.append(content["text"]) + results: Final[Sequence[_StoreResultView]] = _iterable_items(search_response.get("data", [])) + for result in results: + for content in _iterable_items(result.get("content", [])): + if content.get("type") == "text" and (content_text := content.get("text")): + documents.append(content_text) return documents @staticmethod - def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[_SearchResultView]: + def get_top_chunks_from_rerank( + search_response: VectorStoreSearchResponse, rerank_response: Any + ) -> list[_StoreResultView]: """Get the original search results corresponding to the top reranked results.""" - top_chunks: Final[list[_SearchResultView]] = [] - search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} - original_results: Final = search_data["results"] + top_chunks: Final[list[_StoreResultView]] = [] + original_results: Final = search_response.get("data", []) 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]) + if index is not None and index < len(sized_results := _sized_items(original_results)): + top_chunks.append(sized_results[index]) return top_chunks diff --git a/litellm/repositories/autorouter_session_repository.py b/litellm/repositories/autorouter_session_repository.py index 640fd26ffe3..e39087167e4 100644 --- a/litellm/repositories/autorouter_session_repository.py +++ b/litellm/repositories/autorouter_session_repository.py @@ -2,31 +2,25 @@ Repository for the auto-router per-session rollup (LiteLLM_AutoRouterSession). """ -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final from litellm.models.autorouter_session import LiteLLM_AutoRouterSession from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models -class _AutoRouterSessionDb(Protocol): - @property - def litellm_autoroutersession(self) -> TableActions["prisma_models.LiteLLM_AutoRouterSession"]: ... - - -class _PrismaClientView(Protocol): - @property - def db(self) -> _AutoRouterSessionDb: ... +class _AutoRouterSessionTable(PrismaTableRepository["prisma_models.LiteLLM_AutoRouterSession"]): + table_name = "litellm_autoroutersession" class AutoRouterSessionRepository(BaseRepository[LiteLLM_AutoRouterSession]): @property def table(self) -> TableActions["prisma_models.LiteLLM_AutoRouterSession"]: - client: Final[_PrismaClientView] = self.prisma_client - return client.db.litellm_autoroutersession + return _AutoRouterSessionTable(self._prisma_client).table @property def model_class(self) -> type[LiteLLM_AutoRouterSession]: diff --git a/litellm/repositories/base_repository.py b/litellm/repositories/base_repository.py index 92f8ba984f4..a0662a5fc33 100644 --- a/litellm/repositories/base_repository.py +++ b/litellm/repositories/base_repository.py @@ -45,12 +45,15 @@ class BaseRepository(ABC, Generic[T]): def __init__(self, prisma_client: object): self._prisma_client = prisma_client - @property - def prisma_client(self) -> Any: # any-ok: PrismaClient is an untyped runtime wrapper + def _connected_client(self) -> object: if self._prisma_client is None: raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") return self._prisma_client + @property + def prisma_client(self) -> Any: # any-ok: PrismaClient is an untyped runtime wrapper + return self._connected_client() + @property @abstractmethod def table(self) -> TableActions[DbRecord]: diff --git a/litellm/repositories/budget_repository.py b/litellm/repositories/budget_repository.py index 205646c8393..af6c8cfe6f5 100644 --- a/litellm/repositories/budget_repository.py +++ b/litellm/repositories/budget_repository.py @@ -3,28 +3,19 @@ Budget repository for database operations on LiteLLM_BudgetTable. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final from litellm.models.budget import LiteLLM_BudgetTable from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models -class _BudgetDb(Protocol): - """The single Prisma table this repository reaches for on ``prisma_client.db``.""" - - @property - def litellm_budgettable(self) -> TableActions["prisma_models.LiteLLM_BudgetTable"]: ... - - -class _PrismaClientView(Protocol): - """The one attribute this repository reads off the untyped Prisma client wrapper.""" - - @property - def db(self) -> _BudgetDb: ... +class _BudgetTable(PrismaTableRepository["prisma_models.LiteLLM_BudgetTable"]): + table_name = "litellm_budgettable" class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]): @@ -32,8 +23,7 @@ class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_BudgetTable"]: - client: Final[_PrismaClientView] = self.prisma_client - return client.db.litellm_budgettable + return _BudgetTable(self._prisma_client).table @property def model_class(self) -> type[LiteLLM_BudgetTable]: diff --git a/litellm/repositories/credentials_repository.py b/litellm/repositories/credentials_repository.py index cb1bc1d101b..d618b71ce52 100644 --- a/litellm/repositories/credentials_repository.py +++ b/litellm/repositories/credentials_repository.py @@ -7,11 +7,12 @@ so reads return the stored values verbatim. """ from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias from litellm.models.credentials import CredentialItem from litellm.repositories.base_repository import DbRecord, record_to_dict from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models @@ -29,27 +30,24 @@ class _PrismaClientView(Protocol): def db(self) -> _PrismaCredentialsDb: ... +class _CredentialsTableRepository(PrismaTableRepository["prisma_models.LiteLLM_CredentialsTable"]): + table_name = "litellm_credentialstable" + + class CredentialsRepository: """Repository for credentials database operations, keyed by credential name.""" - 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 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") - client: Final[_PrismaClientView] = self._prisma_client + client: Final[_PrismaClientView] = _CredentialsTableRepository(self._prisma_client).prisma_client return client @property def table(self) -> "_CredentialsTable": - from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync - - return wrap_table_actions_for_config_sync( - actions=self.prisma_client.db.litellm_credentialstable, - table_name="litellm_credentialstable", - ) + return _CredentialsTableRepository(self._prisma_client).table @staticmethod def _to_model(record: DbRecord | None) -> CredentialItem | None: diff --git a/litellm/repositories/object_permission_repository.py b/litellm/repositories/object_permission_repository.py index 99d6aeed88b..e90da64afcd 100644 --- a/litellm/repositories/object_permission_repository.py +++ b/litellm/repositories/object_permission_repository.py @@ -2,27 +2,19 @@ ObjectPermission repository for database operations on LiteLLM_ObjectPermissionTable. """ -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models -class _ObjectPermissionDb(Protocol): - @property - def litellm_objectpermissiontable(self) -> TableActions["prisma_models.LiteLLM_ObjectPermissionTable"]: ... - - -class _PrismaClientView(Protocol): - @property - def db(self) -> _ObjectPermissionDb: ... - - @property - def writer_db(self) -> _ObjectPermissionDb: ... +class _ObjectPermissionTable(PrismaTableRepository["prisma_models.LiteLLM_ObjectPermissionTable"]): + table_name = "litellm_objectpermissiontable" class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): @@ -34,9 +26,7 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_ObjectPermissionTable"]: - client: Final[_PrismaClientView] = self.prisma_client - database: Final = client.writer_db if self._use_writer else client.db - return database.litellm_objectpermissiontable + return _ObjectPermissionTable(self._prisma_client, use_writer=self._use_writer).table @property def model_class(self) -> type[LiteLLM_ObjectPermissionTable]: diff --git a/litellm/repositories/organization_repository.py b/litellm/repositories/organization_repository.py index 47eb8f4a609..1a773015414 100644 --- a/litellm/repositories/organization_repository.py +++ b/litellm/repositories/organization_repository.py @@ -3,28 +3,19 @@ Organization repository for database operations on LiteLLM_OrganizationTable. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final from litellm.models.organization import LiteLLM_OrganizationTable from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models -class _OrganizationDb(Protocol): - """The single Prisma table this repository reaches for on ``prisma_client.db``.""" - - @property - def litellm_organizationtable(self) -> TableActions["prisma_models.LiteLLM_OrganizationTable"]: ... - - -class _PrismaClientView(Protocol): - """The one attribute this repository reads off the untyped Prisma client wrapper.""" - - @property - def db(self) -> _OrganizationDb: ... +class _OrganizationTable(PrismaTableRepository["prisma_models.LiteLLM_OrganizationTable"]): + table_name = "litellm_organizationtable" class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]): @@ -32,8 +23,7 @@ class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_OrganizationTable"]: - client: Final[_PrismaClientView] = self.prisma_client - return client.db.litellm_organizationtable + return _OrganizationTable(self._prisma_client).table @property def model_class(self) -> type[LiteLLM_OrganizationTable]: diff --git a/litellm/repositories/project_repository.py b/litellm/repositories/project_repository.py index 0ea5009a5b6..ab19027a3d0 100644 --- a/litellm/repositories/project_repository.py +++ b/litellm/repositories/project_repository.py @@ -3,24 +3,19 @@ Project repository for database operations on LiteLLM_ProjectTable. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final from litellm.models.project import LiteLLM_ProjectTable from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models -class _ProjectDb(Protocol): - @property - def litellm_projecttable(self) -> TableActions["prisma_models.LiteLLM_ProjectTable"]: ... - - -class _PrismaClientView(Protocol): - @property - def db(self) -> _ProjectDb: ... +class _ProjectTable(PrismaTableRepository["prisma_models.LiteLLM_ProjectTable"]): + table_name = "litellm_projecttable" class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]): @@ -28,8 +23,7 @@ class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_ProjectTable"]: - client: Final[_PrismaClientView] = self.prisma_client - return client.db.litellm_projecttable + return _ProjectTable(self._prisma_client).table @property def model_class(self) -> type[LiteLLM_ProjectTable]: diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index d010455faa3..aa33da54572 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -10,10 +10,11 @@ methods; richer repositories live in their own modules. from collections.abc import Mapping, Sequence from datetime import timedelta from itertools import groupby -from typing import TYPE_CHECKING, Any, Final, Generic +from typing import TYPE_CHECKING, Any, Final, Generic, Protocol from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.repositories.accounting_delivery_repository import LedgerDatabase, LedgerTransaction from litellm.repositories.prisma_protocols import RowT_co, TableActions @@ -22,6 +23,16 @@ if TYPE_CHECKING: from prisma import models as prisma_models # noqa: F401 # used by quoted base-class subscripts +class _PrismaReaderClientView(Protocol): + @property + def db(self) -> object: ... + + +class _PrismaWriterClientView(Protocol): + @property + def writer_db(self) -> object: ... + + class PrismaTableRepository(Generic[RowT_co]): """Base for repositories that expose a single Prisma table.""" @@ -31,19 +42,26 @@ class PrismaTableRepository(Generic[RowT_co]): self._prisma_client = prisma_client self._use_writer = use_writer - @property - def prisma_client(self) -> Any: + def _connected_client(self) -> object: if self._prisma_client is None: raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") return self._prisma_client + @property + def prisma_client(self) -> Any: + return self._connected_client() + + def _selected_database(self) -> object: + client: Final = self._connected_client() + if self._use_writer: + return require_attributes(client, _PrismaWriterClientView, "writer_db").writer_db + return require_attributes(client, _PrismaReaderClientView, "db").db + @property def table(self) -> TableActions[RowT_co]: from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync - actions: Final[TableActions[RowT_co]] = getattr( - self.prisma_client.writer_db if self._use_writer else self.prisma_client.db, self.table_name - ) + actions: Final[TableActions[RowT_co]] = getattr(self._selected_database(), self.table_name) return wrap_table_actions_for_config_sync(actions=actions, table_name=self.table_name) diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index b89d87aa354..d076d148783 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.models.team import LiteLLM_TeamTable, Member from litellm.repositories.base_repository import ( BaseRepository, @@ -66,10 +67,12 @@ class _PrismaTeamDb(_TeamTables, Protocol): def tx(self) -> _TeamTransactionManager: ... -class _PrismaClientView(Protocol): +class _PrismaReaderClientView(Protocol): @property def db(self) -> _PrismaTeamDb: ... + +class _PrismaWriterClientView(Protocol): @property def writer_db(self) -> _PrismaTeamDb: ... @@ -95,8 +98,10 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): @property def _db(self) -> _PrismaTeamDb: - client: Final[_PrismaClientView] = self.prisma_client - return client.writer_db if self._use_writer else client.db + client: Final = self._connected_client() + if self._use_writer: + return require_attributes(client, _PrismaWriterClientView, "writer_db").writer_db + return require_attributes(client, _PrismaReaderClientView, "db").db @property def table(self) -> TableActions["prisma_models.LiteLLM_TeamTable"]: diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index 28fba2de564..f39d28a3039 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -12,10 +12,12 @@ from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.models.user import LiteLLM_UserTable, SCIMPlaceholder from litellm.repositories.base_repository import BaseRepository, DbRecord, record_to_dict from litellm.repositories.chunked_in import IN_LIST_CHUNK_SIZE from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import Prisma @@ -62,15 +64,13 @@ class _UserTransaction(_UserDb, Protocol): class _PrismaClientView(Protocol): - @property - def db(self) -> _UserDb: ... - - @property - def writer_db(self) -> _UserDb: ... - def tx(self) -> AbstractAsyncContextManager[_UserTransaction]: ... +class _UserTable(PrismaTableRepository["prisma_models.LiteLLM_UserTable"]): + table_name = "litellm_usertable" + + async def _count_billable_users(table: TableActions["prisma_models.LiteLLM_UserTable"]) -> int: from prisma import Json @@ -88,9 +88,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_UserTable"]: - client: Final[_PrismaClientView] = self.prisma_client - database: Final = client.writer_db if self._use_writer else client.db - return database.litellm_usertable + return _UserTable(self._prisma_client, use_writer=self._use_writer).table @property def model_class(self) -> type[LiteLLM_UserTable]: @@ -184,7 +182,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): where={"user_id": data["user_id"]}, data={"create": dict(data), "update": {}} ) - client: Final[_PrismaClientView] = self.prisma_client + client: Final = require_attributes(self._connected_client(), _PrismaClientView, "tx") async with client.tx() as transaction: await transaction.execute_raw("SET TRANSACTION ISOLATION LEVEL READ COMMITTED") await transaction.query_raw("SELECT 1 AS locked FROM pg_advisory_xact_lock(468530001)") diff --git a/litellm/repositories/verification_token_repository.py b/litellm/repositories/verification_token_repository.py index fb1f1c89e69..4b7ceca3212 100644 --- a/litellm/repositories/verification_token_repository.py +++ b/litellm/repositories/verification_token_repository.py @@ -9,6 +9,7 @@ from types import TracebackType from typing import TYPE_CHECKING, Final, Protocol from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.models.verification_token import ( LiteLLM_VerificationToken, ) @@ -76,8 +77,7 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): @property def _db(self) -> _PrismaVerificationTokenDb: - client: Final[_PrismaClientView] = self.prisma_client - return client.db + return require_attributes(self._connected_client(), _PrismaClientView, "db").db @property def table(self) -> TableActions["PrismaVerificationToken"]: diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 7e6ee84e2b2..9da89d96e73 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -4,13 +4,13 @@ from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, require_attributes from litellm.responses.litellm_completion_transformation.client_tool_search import ( bridge_client_tool_search, client_tool_search_loaded_tools, @@ -46,12 +46,24 @@ COLD_STORAGE_HANDLER: Final = ColdStorageHandler() ######################################################## +class _CarriesArguments(Protocol): + arguments: object + + +def _normalize_redacted_function_arguments(tool_call: object) -> None: + function: Final = attribute_or_none(tool_call, "function") + if function is None: + return + carrier: Final = require_attributes(function, _CarriesArguments, "arguments") + if carrier.arguments == REDACTED_BY_LITELLM: + carrier.arguments = REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER + + def _normalize_redacted_tool_call_arguments(message: Message) -> None: """Redaction stores the bare sentinel (invalid JSON) in tool-call arguments; normalize replayed history to "{}" so provider converters can parse it.""" for tool_call in message.tool_calls or []: - if (function := getattr(tool_call, "function", None)) is not None and function.arguments == REDACTED_BY_LITELLM: - function.arguments = REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER + _normalize_redacted_function_arguments(tool_call) function_call: Final = message.function_call if function_call is not None and function_call.arguments == REDACTED_BY_LITELLM: function_call.arguments = REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 0f8ba6ef98a..8b6c53a272d 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -3,10 +3,10 @@ import uuid from collections.abc import Mapping, Sequence from itertools import filterfalse from types import MappingProxyType -from typing import Any, Final, cast +from typing import Any, Final, Protocol, cast import litellm -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items, require_attributes from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params, served_from_cache from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend from litellm.llms.anthropic.common_utils import anthropic_server_call_fields, anthropic_server_tool_calls @@ -84,7 +84,11 @@ def _index_of_output_item_type(items: Sequence[object], item_type: str) -> int | ) -def _output_items_with_id(items: tuple[Any, ...], item_type: str, item_id: str | None) -> tuple[object, ...]: +class _CopiesWithUpdate(Protocol): + def model_copy(self, *, update: Mapping[str, object]) -> object: ... + + +def _output_items_with_id(items: tuple[object, ...], item_type: str, item_id: str | None) -> tuple[object, ...]: if item_id is None: return items @@ -93,7 +97,10 @@ def _output_items_with_id(items: tuple[Any, ...], item_type: str, item_id: str | return items return tuple( - item.model_copy(update={"id": item_id}) if index == target_index else item for index, item in enumerate(items) + require_attributes(item, _CopiesWithUpdate, "model_copy").model_copy(update={"id": item_id}) + if index == target_index + else item + for index, item in enumerate(items) ) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b49bbd23917..05379ffde3e 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -12,7 +12,6 @@ from itertools import accumulate, chain, groupby from types import MappingProxyType from typing import ( TYPE_CHECKING, - Any, Final, Literal, Protocol, @@ -22,12 +21,6 @@ from typing import ( runtime_checkable, ) -from openai.types.chat.chat_completion_named_tool_choice_param import ( - ChatCompletionNamedToolChoiceParam, -) -from openai.types.chat.chat_completion_named_tool_choice_param import ( - Function as NamedToolChoiceFunction, -) from openai.types.responses import ResponseFunctionToolCall, ResponseFunctionWebSearch, ResponseOutputRefusal from openai.types.responses.response_compaction_item import ResponseCompactionItem from openai.types.responses.response_create_params import ResponseInputParam @@ -56,7 +49,12 @@ from litellm.integrations.anthropic_cache_control_hook import ( whole_messages_within_cache_control_limit, with_cache_control, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params, normalize_drop_params +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + drops_unsupported_params, + normalize_drop_params, + require_attributes, +) from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, ) @@ -66,6 +64,7 @@ from litellm.litellm_core_utils.hidden_params import ( set_hidden_params, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict from litellm.litellm_core_utils.model_param_helper import embeds_thought_signature_in_tool_call_id from litellm.litellm_core_utils.prompt_templates.common_utils import ( invalid_type_message, @@ -386,6 +385,10 @@ class _HasId(Protocol): id: object +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + class _ResponsesToolCallItem(Protocol): name: object arguments: object @@ -984,8 +987,8 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_tool_choice( - tool_choice: Any, - ) -> str | dict[str, object] | None: + tool_choice: object, + ) -> object: """ Transform tool_choice from various formats to OpenAI Chat Completion format. @@ -1009,11 +1012,13 @@ class LiteLLMCompletionResponsesConfig: if isinstance(tool_choice, str): return tool_choice - if isinstance(tool_choice, dict): + if is_str_keyed_dict(tool_choice): tool_choice_type: Final = tool_choice.get("type") # If it has a function with name, it's standard OpenAI format - pass through - if tool_choice.get("function") and tool_choice.get("function", {}).get("name"): + if tool_choice.get("function") and require_attributes( + tool_choice.get("function", {}), _SupportsGet, "get" + ).get("name"): return tool_choice # Handle Cursor IDE dict formats without function name @@ -1028,17 +1033,13 @@ class LiteLLMCompletionResponsesConfig: elif tool_choice_type == "function": function_name: Final = tool_choice.get("name") if function_name: - return ChatCompletionNamedToolChoiceParam( - type="function", function=NamedToolChoiceFunction(name=function_name) - ) + return {"type": "function", "function": {"name": function_name}} return "required" elif tool_choice_type == "custom": custom: Final = tool_choice.get("custom") - custom_name = tool_choice.get("name") or (custom.get("name") if isinstance(custom, dict) else None) + custom_name = tool_choice.get("name") or (custom.get("name") if is_str_keyed_dict(custom) else None) if custom_name: - return ChatCompletionNamedToolChoiceParam( - type="function", function=NamedToolChoiceFunction(name=custom_name) - ) + return {"type": "function", "function": {"name": custom_name}} return "required" # Return as-is for unknown formats @@ -1062,7 +1063,7 @@ class LiteLLMCompletionResponsesConfig: return ToolChoiceCustomParam(type="custom", name=custom_name) case _, {"type": "function", "function": {"name": str(function_name)}}: return ToolChoiceFunctionParam(type="function", name=function_name) - case _, "none" | "auto" | "required" as normalized: + case _, normalized if normalized in ("none", "auto", "required"): return normalized case _, _: return "auto" diff --git a/litellm/responses/main.py b/litellm/responses/main.py index ae7c3fccffa..e085c4d7a2e 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -31,7 +31,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( ) from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params, require_attributes from litellm.litellm_core_utils.get_llm_provider_logic import ( declared_authenticating_provider, get_provider_affinity_header, @@ -1147,8 +1147,12 @@ def _apply_prompt_management_to_responses_call( ) -> tuple[str | ResponseInputParam, str, str | None, object]: """Returns the prompt-managed input, model and provider, plus the deployment metadata that still describes the upstream (``None`` once the prompt manager moved the request to another provider).""" - async_outcome: Final[_AsyncPromptManagementOutcome | None] = kwargs.pop("_async_prompt_merged_params", None) - if async_outcome is not None: + carried_outcome: Final[object] = kwargs.get("_async_prompt_merged_params") + kwargs.pop("_async_prompt_merged_params", None) + if carried_outcome is not None: + async_outcome: Final = require_attributes( + carried_outcome, _AsyncPromptManagementOutcome, "merged_optional_params" + ) for key, value in async_outcome.merged_optional_params.items(): local_vars[key] = value return input, model, custom_llm_provider, async_outcome.deployment_model_info diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index de5bfc964af..d045689cb56 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -455,7 +455,7 @@ class LiteLLM_Proxy_MCP_Handler: mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], litellm_trace_id: str | None = None, request_tags: list[str] | None = None, - ) -> tuple[list[Any], dict[str, str]]: + ) -> tuple[list[FunctionToolParam], dict[str, str]]: """ Centralized method to process MCP tools through the complete pipeline. diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index f4e75e01c65..ea6064b554d 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -1,6 +1,9 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast +from typing_extensions import ReadOnly +from typing_extensions import TypedDict as ExtraItemsTypedDict + from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.agentic_loop_settings import validated_max_agentic_loops @@ -39,6 +42,11 @@ else: MAX_MCP_TOOL_CALL_ROUNDS: Final = 5 +class _FollowUpRequest(ExtraItemsTypedDict, total=False, extra_items=object): + input: ReadOnly[str | ResponseInputParam] + model: ReadOnly[str] + + def _output_items(response: ResponsesAPIResponse) -> Sequence[object]: """Read a response's output items as plain objects; the field is a wide union of item models.""" return tuple(cast("Sequence[object]", response.output)) # cast-ok: items are only carried, never inspected @@ -1043,7 +1051,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.max_tool_call_rounds, ) - follow_up_response: Final = await aresponses(**follow_up_params) + follow_up_response: Final = await aresponses(**_FollowUpRequest(**follow_up_params)) # Route the follow-up through the same base_iterator machinery as # the initial call so its completion is checked for further tool diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 7fc8d5aed32..991cf6213c6 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -13,7 +13,7 @@ from typing_extensions import TypeIs # noqa: TID251 # narrows untyped wire pay import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, iter_items +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, decoded_json_object, iter_items from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value, is_nested_path from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.llms.base_llm.responses.transformation import ( @@ -91,8 +91,9 @@ def _is_object_dict( def _prompt_tokens_details_wrapper(details: object) -> PromptTokensDetailsWrapper: if _is_object_dict(details): input_details: Final = InputTokensDetails.model_validate(details) + input_fields: Final = decoded_json_object(input_details.model_dump(exclude={"cache_write_tokens_details"})) return PromptTokensDetailsWrapper( - **input_details.model_dump(exclude={"cache_write_tokens_details"}), + **input_fields, cache_creation_token_details=_cache_write_lifetimes(input_details.cache_write_tokens_details), ) return PromptTokensDetailsWrapper( diff --git a/litellm/router.py b/litellm/router.py index 4b51c16aa81..0d46326b904 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -15,6 +15,8 @@ from collections.abc import Callable, Mapping, Sequence from types import FunctionType, MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal +from pydantic import TypeAdapter + from litellm.router_config import RouterConfig, build_router_config from litellm.router_core import RouterCore, make_core from litellm.router_deprecated import DeprecatedRouterMembers @@ -203,6 +205,7 @@ def _router_init( _CORE_SIGNATURE: Final = inspect.signature(_router_init) +_PASSED_VALUE: Final = TypeAdapter(object) def router_arguments( @@ -214,7 +217,13 @@ def router_arguments( bound: Final = _CORE_SIGNATURE.bind_partial(None, *args, **kwargs).arguments except TypeError as error: raise TypeError(f"Router.__init__() {error}") from None - return MappingProxyType({name: value for name, value in bound.items() if name != "self"}) + return MappingProxyType( + { + name: value + for name, value in zip(bound, map(_PASSED_VALUE.validate_python, bound.values())) + if name != "self" + } + ) def _class_descriptor(cls: type, name: str) -> bool: diff --git a/litellm/router_config.py b/litellm/router_config.py index 1d38fe84d96..874fb314320 100644 --- a/litellm/router_config.py +++ b/litellm/router_config.py @@ -13,6 +13,8 @@ from dataclasses import dataclass, field from types import MappingProxyType from typing import TYPE_CHECKING, Final +from pydantic import TypeAdapter + import litellm if TYPE_CHECKING: @@ -120,6 +122,7 @@ def current_dynamic_config() -> RouterDynamicConfig: ) +_LIVE_VALUE: Final = TypeAdapter(object) _LIVE_SOURCES: Final = MappingProxyType( { "num_retries": "num_retries", @@ -159,7 +162,7 @@ class LiveDynamicConfig: source: Final = _LIVE_SOURCES.get(name) if source is None: raise AttributeError(name) - return getattr(litellm, source) + return _LIVE_VALUE.validate_python(getattr(litellm, source)) live_dynamic_config: Final = LiveDynamicConfig() diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index c2e1c197259..618ac4321d8 100644 --- a/litellm/router_deprecated.py +++ b/litellm/router_deprecated.py @@ -32,6 +32,8 @@ from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, S from types import MappingProxyType from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeVar, overload +from pydantic import TypeAdapter + from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet, RoutingContext from litellm.router_strategy.least_busy import picks_reserve @@ -67,6 +69,8 @@ DEPRECATED_ROUTER_MEMBER_USES: Final[Counter[str]] = Counter() # mutable-ok: th _INTERNAL_PREFIX: Final = "litellm.router" +_MEMBER_VALUE: Final = TypeAdapter(object) + # Router's own attribute hooks: a write or delete through an old name passes through them first. _ATTRIBUTE_HOOKS: Final = frozenset({"__getattribute__", "__setattr__", "__delattr__"}) @@ -935,7 +939,7 @@ def _name_members_after_router() -> None: # The old members were Router methods, so errors that name the method (a wrong-arity TypeError, # or "litellm.router.Router.() argument after ** must be a mapping") keep naming it the # way they did before the members moved here. - for member in vars(DeprecatedRouterMembers).values(): + for member in map(_MEMBER_VALUE.validate_python, vars(DeprecatedRouterMembers).values()): function = member.__func__ if isinstance(member, staticmethod) else member if inspect.isfunction(function): function.__qualname__ = f"Router.{function.__name__}" diff --git a/litellm/router_python_host.py b/litellm/router_python_host.py index 84e20270322..425d38d6233 100644 --- a/litellm/router_python_host.py +++ b/litellm/router_python_host.py @@ -19,6 +19,7 @@ from typing import ( cast, # noqa: TID251 # untyped litellm globals: the live config view and the callback lists ) +from pydantic import TypeAdapter from typing_extensions import overload import litellm @@ -53,6 +54,8 @@ CallbackList: TypeAlias = Literal[ Endpoint: TypeAlias = Callable[..., object] +_LITELLM_GLOBAL: Final = TypeAdapter(object) + class PythonHost: __slots__ = () @@ -126,7 +129,7 @@ class PythonHost: return litellm.logging_callback_manager def callback_list(self, name: CallbackList) -> list[object]: # mutable-ok: the process callback list itself - value: Final[object] = getattr(litellm, name) # pyright: ignore[reportAny] # an untyped litellm global + value: Final = _LITELLM_GLOBAL.validate_python(getattr(litellm, name)) return cast(list[object], value) # cast-ok: the process callback list is an untyped litellm global def set_callback_list(self, name: CallbackList, value: list[object]) -> None: # mutable-ok: assigned as given diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index 3427b81be68..ed402fd5262 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -14,10 +14,13 @@ from __future__ import annotations import hashlib import json from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable + +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, iter_items, require_attributes from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter from litellm.router_strategy.adaptive_router.classifier import classify_prompt from litellm.router_strategy.adaptive_router.config import ( @@ -41,7 +44,12 @@ _IDENTITY_FIELDS: Final = ( ) -def _resolve_session_key(kwargs: dict[str, Any]) -> str | None: +class _LoggedCallKwargs(TypedDict, total=False): + litellm_params: ReadOnly[Mapping[str, object]] + messages: ReadOnly[Sequence[Mapping[str, object]] | None] + + +def _resolve_session_key(kwargs: _LoggedCallKwargs) -> str | None: """Pick a stable per-conversation key for owner-cache attribution. Order: @@ -135,38 +143,70 @@ def _recent_tool_results( return results -def _assistant_content_and_tool_calls(response_obj: Any) -> tuple[object, Sequence[Mapping[str, object]]]: +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, index: int, /) -> object: ... + + +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + +class _SupportsModelDump(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +def _first_item(value: object) -> object: + if isinstance(value, _SupportsGetItem): + return value[0] + raise TypeError(f"'{type(value).__name__}' object is not subscriptable") + + +def _dict_fields(value: object) -> Mapping[str, object] | None: + fields: Final = as_str_mapping(value) + return fields if isinstance(value, dict) else None + + +def _attribute_or_dict_field(source: object, name: str) -> object: + attribute: Final = attribute_or_none(source, name) + fields: Final = _dict_fields(source) + return fields.get(name) if attribute is None and fields is not None else attribute + + +def _tool_call_fields(tool_call: object) -> Mapping[str, object]: + fields: Final = _dict_fields(tool_call) + if fields is not None: + return fields + try: + return require_attributes(tool_call, _SupportsModelDump, "model_dump").model_dump() + except Exception: + return {"name": getattr(tool_call, "name", ""), "arguments": ""} + + +def _assistant_content_and_tool_calls(response_obj: object) -> tuple[object, Sequence[Mapping[str, object]]]: """Return (assistant_text, tool_calls_list) extracted from a ModelResponse-ish object.""" if response_obj is None: return None, [] try: - choices: Final = getattr(response_obj, "choices", None) or response_obj.get("choices") + choices: Final = attribute_or_none(response_obj, "choices") or require_attributes( + response_obj, _SupportsGet, "get" + ).get("choices") except Exception: return None, [] if not choices: return None, [] - msg = choices[0] - msg = getattr(msg, "message", None) or (msg.get("message") if isinstance(msg, dict) else None) + first_choice: Final = _first_item(choices) + choice_fields: Final = _dict_fields(first_choice) + msg: Final = attribute_or_none(first_choice, "message") or ( + choice_fields.get("message") if choice_fields is not None else None + ) if msg is None: return None, [] - content = getattr(msg, "content", None) - if content is None and isinstance(msg, dict): - content = msg.get("content") - - raw_tool_calls = getattr(msg, "tool_calls", None) - if raw_tool_calls is None and isinstance(msg, dict): - raw_tool_calls = msg.get("tool_calls") - tool_calls: Final[list[dict[str, object]]] = [] - for tc in raw_tool_calls or []: - if isinstance(tc, dict): - tool_calls.append(tc) - else: - try: - tool_calls.append(tc.model_dump()) - except Exception: - tool_calls.append({"name": getattr(tc, "name", ""), "arguments": ""}) + content: Final = _attribute_or_dict_field(msg, "content") + raw_tool_calls: Final = _attribute_or_dict_field(msg, "tool_calls") + tool_calls: Final = [_tool_call_fields(tc) for tc in iter_items(raw_tool_calls or [])] return content, tool_calls @@ -214,7 +254,7 @@ class AdaptiveRouterPostCallHook(CustomLogger): async def _record( self, - kwargs: dict[str, Any], + kwargs: _LoggedCallKwargs, response_obj: object, response_status: int, ) -> None: diff --git a/litellm/router_strategy/adaptive_router/signals.py b/litellm/router_strategy/adaptive_router/signals.py index c7fffd05632..9159bc6c2c0 100644 --- a/litellm/router_strategy/adaptive_router/signals.py +++ b/litellm/router_strategy/adaptive_router/signals.py @@ -15,8 +15,9 @@ from __future__ import annotations import re from collections.abc import Mapping, Sequence from dataclasses import dataclass, field -from typing import Any, Final +from typing import Final, Protocol +from litellm.litellm_core_utils.core_helpers import as_str_mapping, require_attributes from litellm.router_strategy.adaptive_router.config import ( LOOP_REPEAT_THRESHOLD, MIN_TURNS_FOR_CLEAN_CREDIT, @@ -174,14 +175,23 @@ def _detect_failure(tool_results: Sequence[Mapping[str, object]]) -> bool: return False -def _signature(call: Mapping[str, Any]) -> str: +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +def _function_field(call: Mapping[str, object], key: str) -> object: + return require_attributes(call.get("function", {}), _SupportsGet, "get").get(key, "") + + +def _signature(call: Mapping[str, object]) -> str: """Stable signature for loop detection: name + sorted JSON-ish args.""" - name: Final = call.get("name") or call.get("function", {}).get("name", "") - call_args = call.get("arguments") - if call_args is None: - call_args = call.get("function", {}).get("arguments", "") - if isinstance(call_args, dict): - call_args = ",".join(f"{k}={call_args[k]}" for k in sorted(call_args.keys())) + name: Final = call.get("name") or _function_field(call, "name") + raw_args: Final = call.get("arguments") + call_args: Final = _function_field(call, "arguments") if raw_args is None else raw_args + args_mapping: Final = as_str_mapping(call_args) + if args_mapping is not None and isinstance(call_args, dict): + joined: Final = ",".join(f"{k}={args_mapping[k]}" for k in sorted(args_mapping.keys())) + return f"{name}({joined})" return f"{name}({call_args})" diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index 68e8701eb23..4e119c414a0 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -6,12 +6,13 @@ import asyncio from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, NamedTuple, Optional -from pydantic import ConfigDict +from pydantic import ConfigDict, TypeAdapter from litellm._logging import verbose_router_logger from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS from litellm.exceptions import Timeout as LiteLLMTimeout from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import decoded_json_object, iter_items from litellm.litellm_core_utils.internal_call_metadata import ( effective_turn_off_message_logging, forwarded_internal_call_metadata, @@ -41,6 +42,9 @@ else: LiteLLMRouterEncoder = Any +_DECODED_JSON: Final = TypeAdapter(object) + + class _CallerMetadata(LiteLLMBaseModel): model_config = ConfigDict(frozen=True, extra="ignore") @@ -123,14 +127,15 @@ class AutoRouter(CustomLogger): if self.auto_router_config is None: raise ValueError("No auto router config provided") auto_router_routes: Final[list[Route]] = [] - loaded_config: Final = json.loads(self.auto_router_config) - for route in loaded_config.get("routes", []): + loaded_config: Final = decoded_json_object(_DECODED_JSON.validate_python(json.loads(self.auto_router_config))) + for route in iter_items(loaded_config.get("routes", [])): + route_config = decoded_json_object(route) auto_router_routes.append( Route( - name=route.get("name"), - description=route.get("description"), - utterances=route.get("utterances", []), - score_threshold=route.get("score_threshold"), + name=route_config.get("name"), + description=route_config.get("description"), + utterances=route_config.get("utterances", []), + score_threshold=route_config.get("score_threshold"), ) ) return auto_router_routes diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 98dd39b40e3..c65820b9892 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -22,11 +22,11 @@ import asyncio import builtins import logging import time -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Mapping, MutableMapping, Sequence from datetime import datetime, timedelta, timezone from itertools import groupby from types import MappingProxyType -from typing import Any, Final +from typing import Any, Final, Protocol from pydantic import TypeAdapter @@ -44,6 +44,7 @@ from litellm.litellm_core_utils.batch_utils import ( ) from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, + require_attributes, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs @@ -51,7 +52,7 @@ from litellm.router_utils.cooldown_callbacks import ( get_prometheus_logger_from_callbacks, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.router import DeploymentTypedDict, LiteLLM_Params, RouterErrors +from litellm.types.router import LiteLLM_Params, RouterErrors from litellm.types.utils import BudgetConfig, GenericBudgetConfigType, StandardLoggingPayload from litellm.types.utils import BudgetConfig as GenericBudgetInfo @@ -60,6 +61,32 @@ _CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object]) _JOB_CHARGE_VALUES: Final = TypeAdapter(tuple[object, float]) +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +def _model_info_id(deployment: Mapping[str, object]) -> object: + return require_attributes(deployment.get("model_info", {}), _SupportsGet, "get").get("id") + + +def _deployment_budget(deployment: Mapping[str, object]) -> tuple[object, GenericBudgetInfo] | None: + litellm_params: Final = deployment.get("litellm_params", {}) + model_id: Final = require_attributes(deployment.get("model_info") or {}, _SupportsGet, "get").get("id") + params: Final = require_attributes(litellm_params, _SupportsGet, "get") + max_budget: Final = params.get("max_budget") + budget_duration: Final = params.get("budget_duration") + + verbose_router_logger.debug( + "Init Deployment Budget: max_budget: %s, budget_duration: %s, model_id: %s", + max_budget, + budget_duration, + model_id, + ) + if max_budget is None or budget_duration is None or model_id is None: + return None + return model_id, GenericBudgetInfo(time_period=budget_duration, budget_limit=max_budget) + + class _LiteLLMParamsDictView: """ Lightweight attribute view over `litellm_params` dict. @@ -124,7 +151,7 @@ class RouterBudgetLimiting(CustomLogger): self, dual_cache: DualCache, provider_budget_config: dict | None, - model_list: list[DeploymentTypedDict | dict[str, Any]] | None = None, + model_list: Sequence[Mapping[str, object]] | None = None, clock: Callable[[], float] = time.time, ): self.dual_cache = dual_cache @@ -137,7 +164,7 @@ class RouterBudgetLimiting(CustomLogger): self._retired = False self._flush_after_retire: asyncio.Task[None] | None = None self.provider_budget_config: GenericBudgetConfigType | None = provider_budget_config - self.deployment_budget_config: GenericBudgetConfigType | None = None + self.deployment_budget_config: Mapping[object, GenericBudgetInfo] | None = None self.tag_budget_config: GenericBudgetConfigType | None = None self._init_provider_budgets() self._init_deployment_budgets(model_list=model_list) @@ -251,9 +278,9 @@ class RouterBudgetLimiting(CustomLogger): def _filter_out_deployments_above_budget( self, potential_deployments: list[dict[str, object]], - healthy_deployments: list[dict[str, Any]], + healthy_deployments: Sequence[dict[str, object]], provider_configs: dict[str, GenericBudgetInfo], - deployment_configs: dict[str, GenericBudgetInfo], + deployment_configs: Mapping[object, GenericBudgetInfo], deployment_providers: list[str | None], spend_map: dict[str, float], request_tags: list[str], @@ -300,9 +327,9 @@ class RouterBudgetLimiting(CustomLogger): # Check deployment budget if self.deployment_budget_config and is_within_budget: _model_name = deployment.get("model_name") - _litellm_params = deployment.get("litellm_params") or {} + _litellm_params = require_attributes(deployment.get("litellm_params") or {}, _SupportsGet, "get") _litellm_model_name = _litellm_params.get("model") - model_id = deployment.get("model_info", {}).get("id") + model_id = _model_info_id(deployment) if model_id in deployment_configs: config = deployment_configs[model_id] current_spend = spend_map.get(f"deployment_spend:{model_id}:{config.budget_duration}", 0.0) @@ -334,12 +361,12 @@ class RouterBudgetLimiting(CustomLogger): async def _async_get_cache_keys_for_router_budget_limiting( self, - healthy_deployments: list[dict[str, Any]], + healthy_deployments: Sequence[Mapping[str, object]], request_kwargs: dict | None = None, ) -> tuple[ list[str], dict[str, GenericBudgetInfo], - dict[str, GenericBudgetInfo], + dict[object, GenericBudgetInfo], list[str | None], ]: """ @@ -354,7 +381,7 @@ class RouterBudgetLimiting(CustomLogger): """ cache_keys: Final[list[str]] = [] provider_configs: Final[dict[str, GenericBudgetInfo]] = {} - deployment_configs: Final[dict[str, GenericBudgetInfo]] = {} + deployment_configs: Final[dict[object, GenericBudgetInfo]] = {} deployment_providers: Final[list[str | None]] = [] # Resolve tags once before the loop (loop-invariant) @@ -378,7 +405,7 @@ class RouterBudgetLimiting(CustomLogger): # Check deployment budgets if self.deployment_budget_config: - model_id = deployment.get("model_info", {}).get("id") + model_id = _model_info_id(deployment) if model_id is not None: budget_config = self._get_budget_config_for_deployment(model_id) if budget_config is not None: @@ -768,7 +795,7 @@ class RouterBudgetLimiting(CustomLogger): def _get_budget_config_for_deployment( self, - model_id: str, + model_id: object, ) -> GenericBudgetInfo | None: if self.deployment_budget_config is None: return None @@ -784,7 +811,7 @@ class RouterBudgetLimiting(CustomLogger): return None return self.tag_budget_config.get(tag, None) - def _get_llm_provider_for_deployment(self, deployment: dict) -> str | None: + def _get_llm_provider_for_deployment(self, deployment: Mapping[str, object]) -> str | None: try: deployment_litellm_params: Final = deployment.get("litellm_params") or {} @@ -893,7 +920,7 @@ class RouterBudgetLimiting(CustomLogger): @staticmethod def should_init_router_budget_limiter( provider_budget_config: dict | None, - model_list: list[DeploymentTypedDict] | list[dict[str, Any]] | None = None, + model_list: Sequence[Mapping[str, object]] | None = None, ): """ Returns `True` if the router budget routing settings are set and RouterBudgetLimiting should be initialized @@ -913,7 +940,7 @@ class RouterBudgetLimiting(CustomLogger): return False for _model in model_list: - _litellm_params = _model.get("litellm_params", {}) + _litellm_params = require_attributes(_model.get("litellm_params", {}), _SupportsGet, "get") if _litellm_params.get("max_budget") or _litellm_params.get("budget_duration") is not None: return True return False @@ -943,37 +970,20 @@ class RouterBudgetLimiting(CustomLogger): def _init_deployment_budgets( self, - model_list: list[DeploymentTypedDict | dict[str, Any]] | None = None, + model_list: Sequence[Mapping[str, object]] | None = None, ): if model_list is None: return - for _model in model_list: - _litellm_params = _model.get("litellm_params", {}) - _model_info: dict = _model.get("model_info") or {} - _model_id = _model_info.get("id") - _max_budget = _litellm_params.get("max_budget") - _budget_duration = _litellm_params.get("budget_duration") - - verbose_router_logger.debug( - "Init Deployment Budget: max_budget: %s, budget_duration: %s, model_id: %s", - _max_budget, - _budget_duration, - _model_id, - ) - if _max_budget is not None and _budget_duration is not None and _model_id is not None: - _budget_config = GenericBudgetInfo( - time_period=_budget_duration, - budget_limit=_max_budget, - ) - if self.deployment_budget_config is None: - self.deployment_budget_config = {} - self.deployment_budget_config[_model_id] = _budget_config + for model_id, budget in filter(None, map(_deployment_budget, model_list)): + if not isinstance(self.deployment_budget_config, MutableMapping): + self.deployment_budget_config = dict(self.deployment_budget_config or {}) + self.deployment_budget_config[model_id] = budget verbose_router_logger.debug("Initialized Deployment Budget Config: %s", self.deployment_budget_config) def register_deployment_budget( self, - deployment: dict[str, Any] | DeploymentTypedDict, + deployment: Mapping[str, object], ) -> None: """ Register or refresh deployment-level budget config for a runtime-added deployment. @@ -981,7 +991,7 @@ class RouterBudgetLimiting(CustomLogger): self._init_deployment_budgets(model_list=[deployment]) def unregister_deployment_budget(self, model_id: str) -> None: - if self.deployment_budget_config is None: + if not isinstance(self.deployment_budget_config, MutableMapping): return self.deployment_budget_config.pop(model_id, None) if len(self.deployment_budget_config) == 0: diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index f004dfcc38d..600b9d7c455 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -28,7 +28,7 @@ from functools import partial from itertools import accumulate, chain, islice, takewhile from threading import Lock from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, cast from pydantic import BaseModel, Field, TypeAdapter, ValidationError, create_model from pydantic_core import ErrorDetails @@ -52,6 +52,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_parent_otel_span_from_kwargs, is_codex_user_agent, + require_attributes, ) from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata @@ -144,7 +145,7 @@ if TYPE_CHECKING: from litellm.router_api import RoutingContext from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter from litellm.router_strategy.savings_baseline import Baseline - from litellm.types.router import PreRoutingHookResponse + from litellm.types.router import DeploymentTypedDict, PreRoutingHookResponse else: RoutingContext = Any PreRoutingHookResponse = Any @@ -184,6 +185,15 @@ def _built_in_tier_or_none(tier_name: str) -> ComplexityTier | None: return ComplexityTier.__members__.get(tier_name) +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +def _declared_vision_support(deployment: Mapping[str, object]) -> object: + model_info: Final = deployment.get("model_info") + return require_attributes(model_info, _SupportsGet, "get").get("supports_vision") if model_info else None + + _CLASSIFICATION_TIER_CRITERIA: Final[Mapping[ComplexityTier, str]] = MappingProxyType( { ComplexityTier.NON_REASONING: ( @@ -3800,11 +3810,12 @@ class ComplexityRouter(CustomLogger): return True return False if is_vision_explicitly_disabled(model) else None - def deployment_verdict(deployment: Mapping[str, Any]) -> bool | None: - declared: Final = (deployment.get("model_info") or EMPTY_MAPPING).get("supports_vision") + def deployment_verdict(deployment: DeploymentTypedDict) -> bool | None: + declared: Final = _declared_vision_support(deployment) if declared is not None: return declared is True - return model_verdict((deployment.get("litellm_params") or EMPTY_MAPPING).get("model") or model_name) + litellm_params: Final = deployment.get("litellm_params") + return model_verdict((litellm_params.get("model") if litellm_params else None) or model_name) deployments: Final = self.ctx.registry().deployments(model_name=model_name) if not deployments: diff --git a/litellm/router_strategy/quality_router/quality_router.py b/litellm/router_strategy/quality_router/quality_router.py index fc8d5ab5fcb..c3b038b5f26 100644 --- a/litellm/router_strategy/quality_router/quality_router.py +++ b/litellm/router_strategy/quality_router/quality_router.py @@ -16,7 +16,7 @@ then cheapest `model_info.input_cost_per_token`). """ import math -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Optional from litellm._logging import verbose_router_logger @@ -162,7 +162,7 @@ class QualityRouter(CustomLogger): `_model_quality`, and `_model_cost`. Raises if any listed model is missing `litellm_routing_preferences`. """ - model_list: Final = self.ctx.registry().all() or [] + model_list: Final[Sequence[object]] = self.ctx.registry().all() or [] available: Final = set(self.config.available_models) # Track which available models we've matched so we can error on missing. diff --git a/litellm/router_utils/cooldown_callbacks.py b/litellm/router_utils/cooldown_callbacks.py index 25821703029..9e90b72ba1b 100644 --- a/litellm/router_utils/cooldown_callbacks.py +++ b/litellm/router_utils/cooldown_callbacks.py @@ -37,10 +37,10 @@ async def router_cooldown_event_callback( deployment_id, ) return - _litellm_params: Final = _deployment["litellm_params"] + _litellm_params: Final = _deployment.litellm_params temp_litellm_params = copy.deepcopy(_litellm_params) temp_litellm_params = dict(temp_litellm_params) - _model_name: Final = _deployment.get("model_name", None) or "" + _model_name: Final = _deployment.model_name or "" _api_base: Final = litellm.get_api_base(model=_model_name, optional_params=temp_litellm_params) or "" model_info: Final = _deployment["model_info"] model_id: Final = model_info.id diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 20f9e114747..068782f7a08 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, field from datetime import datetime from enum import Enum from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import httpx from pydantic import TypeAdapter @@ -798,7 +798,7 @@ async def run_async_fallback( fallback_depth: int, include_fallback_errors: bool = False, **kwargs, -) -> Any: +) -> object: """ Loops through all the fallback model groups and calls kwargs["original_function"] with the arguments and keyword arguments provided. diff --git a/litellm/router_utils/handle_error.py b/litellm/router_utils/handle_error.py index 97bf4710fca..a1e6d2b2cad 100644 --- a/litellm/router_utils/handle_error.py +++ b/litellm/router_utils/handle_error.py @@ -51,7 +51,7 @@ async def send_llm_exception_alert( # the proxy is already instrumented to send LLM API call failures return - litellm_debug_info: Final = getattr(original_exception, "litellm_debug_info", None) + litellm_debug_info: Final[str | None] = getattr(original_exception, "litellm_debug_info", None) exception_str = str(original_exception) if litellm_debug_info is not None: exception_str += litellm_debug_info diff --git a/litellm/router_utils/health_state_cache.py b/litellm/router_utils/health_state_cache.py index 1c9cb10a48b..c4cc9350cfa 100644 --- a/litellm/router_utils/health_state_cache.py +++ b/litellm/router_utils/health_state_cache.py @@ -9,6 +9,7 @@ import time from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final +from pydantic import TypeAdapter from typing_extensions import TypedDict from litellm import verbose_logger @@ -25,6 +26,9 @@ else: Span = Any +_CACHED_HEALTH_SNAPSHOT: Final = TypeAdapter(object) + + class DeploymentHealthStateValue(TypedDict): is_healthy: bool timestamp: float @@ -92,7 +96,7 @@ class DeploymentHealthCache: str(e), ) - def _extract_unhealthy_ids(self, raw: Any) -> set[str]: + def _extract_unhealthy_ids(self, raw: object) -> set[str]: """Given raw cache value, return set of non-stale unhealthy deployment IDs.""" if not raw or not isinstance(raw, dict): return set() @@ -109,7 +113,7 @@ class DeploymentHealthCache: async def async_get_unhealthy_deployment_ids(self, parent_otel_span: Span | None = None) -> set[str]: """Return set of deployment IDs currently marked unhealthy and not stale.""" try: - raw: Final = await self.cache.async_get_cache(key=self.CACHE_KEY) + raw: Final = _CACHED_HEALTH_SNAPSHOT.validate_python(await self.cache.async_get_cache(key=self.CACHE_KEY)) return self._extract_unhealthy_ids(raw) except Exception as e: verbose_logger.debug( @@ -122,7 +126,7 @@ class DeploymentHealthCache: def get_unhealthy_deployment_ids(self, parent_otel_span: Span | None = None) -> set[str]: """Sync version: return set of deployment IDs currently marked unhealthy and not stale.""" try: - raw: Final = self.cache.get_cache(key=self.CACHE_KEY) + raw: Final = _CACHED_HEALTH_SNAPSHOT.validate_python(self.cache.get_cache(key=self.CACHE_KEY)) return self._extract_unhealthy_ids(raw) except Exception as e: verbose_logger.debug( diff --git a/litellm/skills/main.py b/litellm/skills/main.py index 6f57bd1bc96..3262a05ca60 100644 --- a/litellm/skills/main.py +++ b/litellm/skills/main.py @@ -5,7 +5,7 @@ Provides create, list, get, and delete operations for skills import asyncio import contextvars -from collections.abc import Coroutine, Mapping +from collections.abc import Coroutine, Mapping, Sequence from functools import partial from typing import Any, Final @@ -14,6 +14,7 @@ import httpx import litellm from litellm.constants import request_timeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.llms.anthropic_skills import ( @@ -45,15 +46,15 @@ def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object]) -> Any | No def _get_skill_request_metadata( kwargs: Mapping[str, object], - extra_body: dict[str, Any] | None, -) -> dict[str, Any] | None: - if extra_body and isinstance(extra_body.get("metadata"), dict): - return extra_body["metadata"] + extra_body: Mapping[str, object] | None, +) -> dict[str, object] | None: + extra_body_metadata: Final = extra_body.get("metadata") if extra_body else None + if is_str_keyed_dict(extra_body_metadata): + return extra_body_metadata metadata: Final = kwargs.get("metadata") - if isinstance(metadata, dict) and isinstance(metadata.get("requester_metadata"), dict): - return metadata["requester_metadata"] - return None + requester_metadata: Final = metadata.get("requester_metadata") if is_str_keyed_dict(metadata) else None + return requester_metadata if is_str_keyed_dict(requester_metadata) else None def _get_litellm_skills_handler(): @@ -70,7 +71,7 @@ def _get_litellm_skills_handler(): @client async def acreate_skill( - files: list[Any] | None = None, + files: Sequence[object] | None = None, display_title: str | None = None, extra_headers: dict[str, str] | None = None, extra_query: Mapping[str, object] | None = None, @@ -133,7 +134,7 @@ async def acreate_skill( @client def create_skill( - files: list[Any] | None = None, + files: Sequence[object] | None = None, display_title: str | None = None, extra_headers: dict[str, str] | None = None, extra_query: Mapping[str, object] | None = None, diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py index 2b5b4d52274..daf65b77a85 100644 --- a/litellm/text_completions/streaming_iterator.py +++ b/litellm/text_completions/streaming_iterator.py @@ -1,6 +1,7 @@ from typing import Final from litellm.constants import VLLM_CHOICE_FIELDS +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact from litellm.types.utils import Choices, ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse @@ -43,9 +44,10 @@ class TextCompletionStreamWrapper: if chunk["choices"]: chat_choice: Final = chunk["choices"][0] - text_choices: Final = TextChoices( - **{field: chat_choice[field] for field in VLLM_CHOICE_FIELDS if field in chat_choice} + choice_fields: Final = decoded_json_object( + {field: chat_choice[field] for field in VLLM_CHOICE_FIELDS if field in chat_choice} ) + text_choices: Final = TextChoices(index=0, **choice_fields) delta: Final = chat_choice["delta"] text_choices["text"] = delta["content"] text_choices["reasoning_content"] = delta.get("reasoning_content") diff --git a/litellm/types/llms/anthropic_skills.py b/litellm/types/llms/anthropic_skills.py index d11fd32f88b..690c2226a73 100644 --- a/litellm/types/llms/anthropic_skills.py +++ b/litellm/types/llms/anthropic_skills.py @@ -2,6 +2,7 @@ Type definitions for Anthropic Skills API """ +from collections.abc import Sequence from typing import Any from typing_extensions import TypedDict @@ -16,7 +17,7 @@ class CreateSkillRequest(TypedDict, total=False): display_title: str | None """Display title for the skill (optional)""" - files: list[Any] | None + files: Sequence[object] | None """Files to upload for the skill. All files must be in the same top-level directory and must include a SKILL.md file at the root.""" diff --git a/litellm/utils.py b/litellm/utils.py index 13f22ddc36e..2c1a3892894 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -90,6 +90,8 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import ( ) from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, + has_attributes, + iter_items, normalize_drop_params, unbind_budget_reservation_from_callbacks, ) @@ -4506,8 +4508,9 @@ def get_optional_params_transcription( passed_params.pop("OPENAI_TRANSCRIPTION_PARAMS") passed_params.pop("model") - custom_llm_provider = passed_params.pop("custom_llm_provider") - drop_params = normalize_drop_params(passed_params.pop("drop_params")) + passed_params.pop("custom_llm_provider") + passed_params.pop("drop_params") + normalized_drop_params: Final = normalize_drop_params(drop_params) special_params: Final[Mapping[str, object]] = passed_params.pop("kwargs") from litellm.litellm_core_utils.audio_utils.transcription_streaming import ( transcription_stream_flag, @@ -4518,7 +4521,7 @@ def get_optional_params_transcription( passed_params[k] = v sends_through_openai_sdk: Final = custom_llm_provider in ("openai", "azure") - drops_unsupported_params: Final = drop_params is True or litellm.drop_params is True + drops_unsupported_params: Final = normalized_drop_params is True or litellm.drop_params is True transcription_params: Final = ( OPENAI_TRANSCRIPTION_PARAMS if sends_through_openai_sdk or drops_unsupported_params @@ -4570,7 +4573,7 @@ def get_optional_params_transcription( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=drop_params if drop_params is not None else False, + drop_params=normalized_drop_params if normalized_drop_params is not None else False, ) elif custom_llm_provider == "bedrock": from litellm.llms.bedrock.audio_transcription import ( @@ -4579,7 +4582,7 @@ def get_optional_params_transcription( ) match map_openai_transcription_params( - non_default_params, drop_params=drop_params is True or litellm.drop_params is True + non_default_params, drop_params=normalized_drop_params is True or litellm.drop_params is True ): case UnsupportedTranscriptionParams(message=message): raise UnsupportedParamsError(message=message, llm_provider="bedrock") @@ -4592,7 +4595,7 @@ def get_optional_params_transcription( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=drop_params if drop_params is not None else False, + drop_params=normalized_drop_params if normalized_drop_params is not None else False, ) optional_params = add_provider_specific_params_to_optional_params( @@ -6562,26 +6565,41 @@ def count_characters(text: str) -> int: _count_characters = count_characters +class _ResponsesOutputCarrier(Protocol): + @property + def output(self) -> object: ... + + +class _ResponsesContentCarrier(Protocol): + @property + def content(self) -> object: ... + + +class _ResponsesTextCarrier(Protocol): + @property + def text(self) -> str: ... + + def get_response_string(response_obj: ModelResponse | ModelResponseStream) -> str: # Handle Responses API streaming events if hasattr(response_obj, "type") and hasattr(response_obj, "response"): # This is a Responses API streaming event (e.g., ResponseCreatedEvent, ResponseCompletedEvent) # Extract text from the response object's output if available - responses_api_response: Final = getattr(response_obj, "response", None) - if responses_api_response and hasattr(responses_api_response, "output"): + responses_api_response: Final[object] = getattr(response_obj, "response", None) + if responses_api_response and has_attributes(responses_api_response, _ResponsesOutputCarrier, "output"): output_list: Final = responses_api_response.output # Use list accumulation to avoid O(n^2) string concatenation: # repeatedly doing `response_str += part` copies the full string each time # because Python strings are immutable, so total work grows with n^2. response_output_parts: Final[list[str]] = [] - for output_item in output_list: + for output_item in iter_items(output_list): # Handle output items with content array - if hasattr(output_item, "content"): - for content_part in output_item.content: - if hasattr(content_part, "text"): + if has_attributes(output_item, _ResponsesContentCarrier, "content"): + for content_part in iter_items(output_item.content): + if has_attributes(content_part, _ResponsesTextCarrier, "text"): response_output_parts.append(content_part.text) # Handle output items with direct text field - elif hasattr(output_item, "text"): + elif has_attributes(output_item, _ResponsesTextCarrier, "text"): response_output_parts.append(output_item.text) return "".join(response_output_parts) diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 598b56f4a79..b83c9ef48b1 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -437,7 +437,7 @@ class VectorStoreRegistry: # cast to VectorStoreConfig litellm_vector_store_config = LiteLLM_VectorStoreConfig(**vector_store_config) vector_store_name = litellm_vector_store_config.get("vector_store_name") - vector_store_litellm_params: dict[str, Any] = litellm_vector_store_config.get("litellm_params") or {} + vector_store_litellm_params = litellm_vector_store_config.get("litellm_params") or {} vector_store_id = vector_store_litellm_params.get("vector_store_id") if vector_store_id is None: diff --git a/litellm/videos/utils.py b/litellm/videos/utils.py index b0a78691a70..b7b663e1e00 100644 --- a/litellm/videos/utils.py +++ b/litellm/videos/utils.py @@ -1,4 +1,5 @@ -from typing import Any, Final, cast +from collections.abc import Mapping +from typing import Final, cast import litellm from litellm.llms.base_llm.videos.transformation import BaseVideoConfig @@ -47,7 +48,7 @@ class VideoGenerationRequestUtils: @staticmethod def get_requested_video_generation_optional_param( - params: dict[str, Any], + params: Mapping[str, object], ) -> VideoCreateOptionalRequestParams: """ Filter parameters to only include those defined in VideoCreateOptionalRequestParams. @@ -60,9 +61,8 @@ class VideoGenerationRequestUtils: """ params = dict(params or {}) - raw_kwargs = params.get("kwargs", {}) - if not isinstance(raw_kwargs, dict): - raw_kwargs = {} + nested_kwargs: Final = params.get("kwargs", {}) + raw_kwargs: Final[dict[str, object]] = nested_kwargs if isinstance(nested_kwargs, dict) else {} kwargs_extra_body: Final = raw_kwargs.pop("extra_body", None) top_level_extra_body: Final = params.get("extra_body") diff --git a/tests/unit/a2a_protocol/test_card_resolver.py b/tests/unit/a2a_protocol/test_card_resolver.py index 4cfa8dab45b..cdf651bd259 100644 --- a/tests/unit/a2a_protocol/test_card_resolver.py +++ b/tests/unit/a2a_protocol/test_card_resolver.py @@ -14,6 +14,7 @@ import pytest from litellm.a2a_protocol.card_resolver import ( LiteLLMA2ACardResolver, fix_agent_card_url, + get_agent_card_url, is_localhost_or_internal_url, normalize_agent_card_interfaces, resolve_a2a_protocol_version, @@ -117,6 +118,37 @@ def test_fix_agent_card_url_updates_interface_when_top_level_is_localhost(): assert result.supported_interfaces[0].url == "https://my-public-agent.example.com/" +@pytest.mark.parametrize( + ("card", "expected"), + [ + (SimpleNamespace(url="https://agent.example.com/"), "https://agent.example.com/"), + ( + SimpleNamespace(url="", supported_interfaces=[SimpleNamespace(url="https://interface.example.com/")]), + "https://interface.example.com/", + ), + ( + SimpleNamespace(supported_interfaces=[SimpleNamespace(url="https://interface.example.com/")]), + "https://interface.example.com/", + ), + (SimpleNamespace(url=""), None), + (SimpleNamespace(url=None, supported_interfaces=[]), None), + (SimpleNamespace(), None), + ], + ids=[ + "top_level_url", + "empty_url_falls_to_interface", + "interfaces_only", + "empty_url_without_interfaces", + "no_interfaces", + "bare", + ], +) +def test_get_agent_card_url_prefers_the_top_level_url_then_the_first_interface( + card: SimpleNamespace, expected: str | None +): + assert get_agent_card_url(card) == expected + + def test_normalize_agent_card_interfaces_downgrades_miscased_interfaces_to_the_0_3_dialect(): pb2 = pytest.importorskip("a2a.types.a2a_pb2") diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index e9a075af507..bd14b505323 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -259,6 +259,29 @@ def test_cache_rejects_unsupported_type() -> None: assert str(raised.value).split("Supported types: ")[1].split(", ") == [member.value for member in LiteLLMCacheType] +class _PingableInMemoryCache(InMemoryCache): + async def ping(self) -> bool: + return True + + async def delete_cache_keys(self, keys: list[str]) -> None: + for key in keys: + self.delete_cache(key) + + +@pytest.mark.asyncio +async def test_cache_ping_and_delete_cache_keys_reach_a_backend_that_supports_them() -> None: + backend: Final = _PingableInMemoryCache() + cache: Final = Cache(type="local", _backend=backend) + backend.set_cache("kept", "kept value") + backend.set_cache("dropped", "dropped value") + + assert await cache.ping() is True + await cache.delete_cache_keys(["dropped"]) + + assert backend.get_cache("dropped") is None + assert backend.get_cache("kept") == "kept value" + + def test_semantic_cache_scope_rejects_unknown_value(): with pytest.raises(ValueError, match="'team' is not a valid SemanticCacheScope"): _semantic_cache(semantic_cache_scope="team") diff --git a/tests/unit/caching/test_redis_semantic_cache.py b/tests/unit/caching/test_redis_semantic_cache.py index 1b2fbb4fd3d..9b6f6de0d7d 100644 --- a/tests/unit/caching/test_redis_semantic_cache.py +++ b/tests/unit/caching/test_redis_semantic_cache.py @@ -2,6 +2,7 @@ from collections.abc import Iterator from contextlib import contextmanager from importlib import import_module import sys +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -300,6 +301,28 @@ def test_redis_semantic_cache_matches_bytes_cache_key(): ) +_STORED_RESPONSE: Final = {"content": "Paris", "finish_reason": None, "cached": True} + + +@pytest.mark.parametrize( + ("cached_response", "expected"), + [ + (str(_STORED_RESPONSE), _STORED_RESPONSE), + (str(_STORED_RESPONSE).encode("utf-8"), _STORED_RESPONSE), + ("{'content': 'Paris'", None), + ], + ids=["python_repr", "python_repr_bytes", "unparseable"], +) +def test_redis_semantic_cache_reads_back_a_value_stored_as_its_python_repr( + cached_response: str | bytes, expected: dict[str, object] | None +): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + + assert redis_semantic_cache._get_cache_logic(cached_response) == expected + + def test_redis_semantic_cache_rejects_pre_isolation_unscoped_hit(): """Pre-isolation entries with no cache-key field cannot be safely reassigned to a caller's scope and are treated as misses.""" diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 0a9dc0ab3cf..2f236596960 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -4,7 +4,7 @@ import importlib import os import shutil import tempfile -from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator +from collections.abc import AsyncIterator, Callable, Container, Coroutine, Iterable, Iterator from dataclasses import dataclass, field from pathlib import Path from types import MappingProxyType @@ -596,3 +596,9 @@ def pytest_sessionfinish() -> None: _close_handler_if_needed(getattr(litellm, name, None)) _run_coroutine_if_needed(close_litellm_async_clients()) enable_socket() + + +def in_operator_error(value: Container[object]) -> str: + with pytest.raises(TypeError) as native: + _ = "key" in value + return str(native.value) diff --git a/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py b/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py index 13fef7a0945..ebe022072dc 100644 --- a/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py +++ b/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py @@ -23,7 +23,7 @@ from litellm import Router from litellm.caching.caching import DualCache from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback -from litellm.types.router import ModelInfo +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.types.utils import BudgetConfig, GenericBudgetConfigType @@ -243,11 +243,11 @@ async def test_router_cooldown_event_callback(): """ # Mock routing context mock_ctx = MagicMock() - mock_deployment = { - "litellm_params": {"model": "gpt-5-mini"}, - "model_name": "gpt-5-mini", - "model_info": ModelInfo(id="test-model-id"), - } + mock_deployment = Deployment( + model_name="gpt-5-mini", + litellm_params=LiteLLM_Params(model="gpt-5-mini"), + model_info=ModelInfo(id="test-model-id"), + ) mock_ctx.registry.return_value.deployment.return_value = mock_deployment # Create a real PrometheusLogger instance @@ -303,11 +303,11 @@ async def test_router_cooldown_event_callback_no_prometheus(): """ # Mock routing context mock_ctx = MagicMock() - mock_deployment = { - "litellm_params": {"model": "gpt-5-mini"}, - "model_name": "gpt-5-mini", - "model_info": ModelInfo(id="test-model-id"), - } + mock_deployment = Deployment( + model_name="gpt-5-mini", + litellm_params=LiteLLM_Params(model="gpt-5-mini"), + model_info=ModelInfo(id="test-model-id"), + ) mock_ctx.registry.return_value.deployment.return_value = mock_deployment await router_cooldown_event_callback( diff --git a/tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py b/tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py index a1a88653ee6..b90f563496e 100644 --- a/tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py +++ b/tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py @@ -2,6 +2,7 @@ import json import re from unittest.mock import MagicMock, patch +import httpx import pytest @@ -10,6 +11,7 @@ from litellm.integrations.bitbucket.bitbucket_prompt_manager import ( BitBucketPromptManager, BitBucketPromptTemplate, ) +from litellm.llms.custom_httpx.http_handler import HTTPHandler def test_bitbucket_prompt_template_creation(): @@ -502,3 +504,19 @@ rate: 0.5""" assert parsed["disabled"] is False assert parsed["count"] == 42 assert parsed["rate"] == 0.5 + + +def test_bitbucket_client_reads_repository_info_and_branches_from_the_api() -> None: + def bitbucket_api(request: httpx.Request) -> httpx.Response: + if request.url.path == "/2.0/repositories/test-workspace/test-repo": + return httpx.Response(200, json={"full_name": "test-workspace/test-repo", "is_private": True}) + if request.url.path == "/2.0/repositories/test-workspace/test-repo/refs/branches": + return httpx.Response(200, json={"values": [{"name": "main"}, {"name": "dev"}]}) + return httpx.Response(404) + + client = BitBucketClient({"workspace": "test-workspace", "repository": "test-repo", "access_token": "test-token"}) + client.http_handler = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(bitbucket_api))) + + assert client.get_repository_info() == {"full_name": "test-workspace/test-repo", "is_private": True} + assert [branch["name"] for branch in client.get_branches()] == ["main", "dev"] + assert client.test_connection() is True diff --git a/tests/unit/integrations/focus/test_focus_database.py b/tests/unit/integrations/focus/test_focus_database.py index 06240eac387..e5d046efb78 100644 --- a/tests/unit/integrations/focus/test_focus_database.py +++ b/tests/unit/integrations/focus/test_focus_database.py @@ -109,3 +109,19 @@ async def test_should_build_frame_from_rows_recovered_for_double_hashed_keys(mon assert result["api_key_alias"].to_list() == ["joined", "batch-worker"] assert result["team_id"].to_list() == ["team-0", "team-1"] + + +@pytest.mark.asyncio +async def test_table_info_returns_the_daily_user_spend_columns(monkeypatch: pytest.MonkeyPatch): + columns = [ + {"column_name": "id", "data_type": "text", "is_nullable": "NO"}, + {"column_name": "spend", "data_type": "double precision", "is_nullable": "YES"}, + ] + db, query_mock = _setup_db(monkeypatch, columns) + + info = await db.get_table_info() + + assert info == {"columns": columns, "table_name": "LiteLLM_DailyUserSpend"} + query_text, *params = query_mock.await_args.args + assert "WHERE table_name = 'LiteLLM_DailyUserSpend'" in query_text + assert params == [] diff --git a/tests/unit/integrations/test_email_alerting.py b/tests/unit/integrations/test_email_alerting.py new file mode 100644 index 00000000000..35de2b112c6 --- /dev/null +++ b/tests/unit/integrations/test_email_alerting.py @@ -0,0 +1,33 @@ +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock + +import pytest + +from litellm.integrations.email_alerting import get_all_team_member_emails +from litellm.proxy import proxy_server + + +@pytest.mark.asyncio +async def test_team_member_emails_are_the_set_emails_of_the_members_with_a_user_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + team: Final = SimpleNamespace( + members_with_roles=[{"user_id": "u-1", "role": "admin"}, {"role": "user"}, {"user_id": "u-2", "role": "user"}] + ) + query_raw: Final = AsyncMock( + return_value=[{"user_email": "a@example.com"}, {"user_email": None}, {}, {"user_email": "b@example.com"}] + ) + prisma_client: Final = SimpleNamespace( + db=SimpleNamespace( + litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=team)), + query_raw=query_raw, + ) + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + + emails: Final = await get_all_team_member_emails("team-1") + + assert emails == ["a@example.com", "b@example.com"] + assert query_raw.await_args is not None + assert query_raw.await_args.args[1] == ["u-1", "u-2"] diff --git a/tests/unit/integrations/test_galileo.py b/tests/unit/integrations/test_galileo.py index d067b31a651..80fc858331e 100644 --- a/tests/unit/integrations/test_galileo.py +++ b/tests/unit/integrations/test_galileo.py @@ -984,3 +984,32 @@ def test_galileo_output_str_of_a_chat_response_whose_message_is_not_a_litellm_me ) assert output == expected_output + + +GALILEO_MESSAGES: Final = [{"role": "user", "content": "weather in SF?"}] +GALILEO_TOOLS: Final = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] +GALILEO_FUNCTIONS: Final = [{"name": "get_time", "parameters": {}}] + + +@pytest.mark.parametrize( + ("optional_params", "expected_input"), + [ + ( + {"functions": GALILEO_FUNCTIONS, "tools": GALILEO_TOOLS, "temperature": 0.1}, + {"messages": GALILEO_MESSAGES, "functions": GALILEO_FUNCTIONS, "tools": GALILEO_TOOLS}, + ), + (None, {"messages": GALILEO_MESSAGES}), + ], +) +def test_galileo_error_span_input_carries_the_request_functions_and_tools( + optional_params: dict[str, object] | None, expected_input: dict[str, object] +) -> None: + _, output_text, logged_input = GalileoObserve()._get_galileo_input_output_content( + {"messages": GALILEO_MESSAGES, "optional_params": optional_params}, + None, + level="ERROR", + status_message="upstream 500", + ) + + assert output_text == "upstream 500" + assert logged_input == expected_input diff --git a/tests/unit/integrations/test_humanloop.py b/tests/unit/integrations/test_humanloop.py index 469dcef9fc9..daa87ee6874 100644 --- a/tests/unit/integrations/test_humanloop.py +++ b/tests/unit/integrations/test_humanloop.py @@ -1,15 +1,18 @@ import asyncio import importlib import os -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from typing import Final +import httpx import pytest import pytest_asyncio import litellm from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE +from litellm.integrations import humanloop from litellm.integrations.humanloop import HumanLoopPromptManager +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -30,6 +33,58 @@ def test_compile_prompt(): assert compiled_prompt[0]["content"] == "You are John. Answer questions as this person. Do not break character." +def _serve_humanloop_prompt(monkeypatch: pytest.MonkeyPatch, prompt_body: Mapping[str, object]) -> None: + def humanloop_api(request: httpx.Request) -> httpx.Response: + if request.url.path == "/v5/prompts/pr_1" and request.headers.get("X-Api-Key") == "hl-key": + return httpx.Response(200, json=prompt_body) + return httpx.Response(404, json={"detail": "not found"}) + + monkeypatch.setattr( + humanloop, + "get_httpx_client", + lambda: HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(humanloop_api))), + ) + + +@pytest.mark.parametrize( + ("template", "expected_template"), + [ + ({"role": "system", "content": "You are {{person}}."}, [{"role": "system", "content": "You are {{person}}."}]), + ( + [{"role": "system", "content": "Be brief."}, {"role": "user", "content": "{{question}}"}], + [{"role": "system", "content": "Be brief."}, {"role": "user", "content": "{{question}}"}], + ), + ], +) +def test_humanloop_prompt_from_api_lists_its_template_and_keeps_only_openai_params( + monkeypatch: pytest.MonkeyPatch, template: object, expected_template: list[dict[str, str]] +) -> None: + _serve_humanloop_prompt( + monkeypatch, + { + "id": "pr_1", + "name": "support-bot", + "template": template, + "model": "gpt-4o", + "temperature": 0.3, + "max_tokens": 64, + }, + ) + + prompt: Final = HumanLoopPromptManager()._get_prompt_from_id_api("pr_1", "hl-key") + + assert prompt["prompt_template"] == expected_template + assert prompt["model"] == "gpt-4o" + assert prompt["optional_params"] == {"temperature": 0.3, "max_tokens": 64} + + +def test_humanloop_prompt_from_api_rejects_a_template_that_is_not_messages(monkeypatch: pytest.MonkeyPatch) -> None: + _serve_humanloop_prompt(monkeypatch, {"id": "pr_1", "template": "plain text", "model": "gpt-4o"}) + + with pytest.raises(ValueError, match="Invalid template message type: "): + HumanLoopPromptManager()._get_prompt_from_id_api("pr_1", "hl-key") + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/integrations/test_lago.py b/tests/unit/integrations/test_lago.py new file mode 100644 index 00000000000..21e866fa747 --- /dev/null +++ b/tests/unit/integrations/test_lago.py @@ -0,0 +1,55 @@ +import json +from datetime import datetime, timezone +from typing import Final + +import httpx +import pytest + +from litellm.integrations.lago import LagoLogger +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.types.utils import ModelResponse, Usage + +CALL_TIME: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + +@pytest.fixture +def lago_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LAGO_API_KEY", "lago-key") + monkeypatch.setenv("LAGO_API_BASE", "https://lago.example") + monkeypatch.setenv("LAGO_API_EVENT_CODE", "llm-usage") + monkeypatch.delenv("LAGO_API_CHARGE_BY", raising=False) + + +def test_lago_event_bills_the_response_token_usage_to_the_end_user(lago_env: None) -> None: + sent: Final[list[httpx.Request]] = [] + + def lago(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={}) + + logger: Final = LagoLogger() + logger.sync_http_handler = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(lago))) + response: Final = ModelResponse(model="gpt-4o", usage=Usage(prompt_tokens=12, completion_tokens=5, total_tokens=17)) + + logger.log_success_event( + kwargs={ + "model": "gpt-4o", + "response_cost": 0.002, + "litellm_params": {"metadata": {}, "proxy_server_request": {"body": {"user": "end-user-1"}}}, + }, + response_obj=response, + start_time=CALL_TIME, + end_time=CALL_TIME, + ) + + assert [str(request.url) for request in sent] == ["https://lago.example/api/v1/events"] + event: Final = json.loads(sent[0].content)["event"] + assert event["external_subscription_id"] == "end-user-1" + assert event["code"] == "llm-usage" + assert event["properties"] == { + "model": "gpt-4o", + "response_cost": 0.002, + "prompt_tokens": 12, + "completion_tokens": 5, + "total_tokens": 17, + } diff --git a/tests/unit/integrations/test_langfuse_otel.py b/tests/unit/integrations/test_langfuse_otel.py index 71087b4de51..aa9a51a93ee 100644 --- a/tests/unit/integrations/test_langfuse_otel.py +++ b/tests/unit/integrations/test_langfuse_otel.py @@ -22,6 +22,7 @@ from opentelemetry.trace import Span import litellm from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, REDACTED_BY_LITELLM from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger +from litellm.integrations.langfuse.langfuse_otel_attributes import LangfuseLLMObsOTELAttributes from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME, OpenTelemetry, OpenTelemetryConfig from litellm.integrations.otel.model.destination import OtelDestination from litellm.integrations.otel.plumbing.context import set_request_destinations @@ -1636,3 +1637,31 @@ def test_a_request_without_a_langfuse_destination_still_exports_through_the_lega assert _spans_the_legacy_logger_exports((_LANGFUSE_DESTINATION.model_copy(update={"callback_name": "arize"}),)) == ( logged ) + + +OBSERVATION_MESSAGES: Final = [{"role": "user", "content": "weather in SF?"}] +OBSERVATION_TOOLS: Final = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] +OBSERVATION_FUNCTIONS: Final = [{"name": "get_time", "parameters": {}}] + + +@pytest.mark.parametrize( + ("request_kwargs", "expected_input"), + [ + ( + { + "messages": OBSERVATION_MESSAGES, + "optional_params": {"tools": OBSERVATION_TOOLS, "functions": OBSERVATION_FUNCTIONS, "top_p": 0.5}, + }, + {"messages": OBSERVATION_MESSAGES, "functions": OBSERVATION_FUNCTIONS, "tools": OBSERVATION_TOOLS}, + ), + ({"messages": OBSERVATION_MESSAGES}, {"messages": OBSERVATION_MESSAGES}), + ], +) +def test_observation_input_carries_the_request_functions_and_tools( + request_kwargs: Mapping[str, object], expected_input: Mapping[str, object] +) -> None: + span: Final = _TRACER.start_span("litellm_request") + + LangfuseLLMObsOTELAttributes.set_messages(span, request_kwargs) + + assert json.loads(str(span.attributes["langfuse.observation.input"])) == expected_input diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index 393905da6a2..c4fd7ed51c7 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -7569,6 +7569,57 @@ class TestOpentelemetryUnitTests: f"{SpanAttributes.LLM_COMPLETIONS.value}.1.function_call.name": "get_news", } + def test_chat_response_tool_calls_are_recorded_per_call_on_the_span(self): + from litellm.types.utils import ModelResponse + + response: Final = ModelResponse( + id="chatcmpl-tools", + model="gpt-4o", + choices=[ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + }, + { + "id": "call_time", + "type": "function", + "function": {"name": "get_time", "arguments": "{}"}, + }, + ], + }, + } + ], + ) + span: Final = TracerProvider().get_tracer("otel-tool-call-tests").start_span("litellm_request") + + OpenTelemetry().set_attributes( + span=span, + kwargs={ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "weather and time in SF?"}], + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "openai"}, + "standard_logging_object": {"id": "call-tools", "call_type": "completion", "metadata": {}}, + }, + response_obj=response, + ) + + completions: Final = SpanAttributes.LLM_COMPLETIONS.value + assert {key: value for key, value in span.attributes.items() if ".function_call." in key} == { + f"{completions}.0.function_call.name": "get_weather", + f"{completions}.0.function_call.arguments": '{"city": "SF"}', + f"{completions}.1.function_call.name": "get_time", + f"{completions}.1.function_call.arguments": "{}", + } + @pytest.mark.asyncio async def test_opentelemetry_integration(self): """ diff --git a/tests/unit/integrations/test_prompt_layer.py b/tests/unit/integrations/test_prompt_layer.py new file mode 100644 index 00000000000..4ebe2ce2fdb --- /dev/null +++ b/tests/unit/integrations/test_prompt_layer.py @@ -0,0 +1,75 @@ +import json +from datetime import datetime, timezone +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.integrations.prompt_layer import PromptLayerLogger +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +CALL_TIME: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) +REQUEST_KWARGS: Final = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "optional_params": {"temperature": 0.2}, + "litellm_params": {"metadata": {"pl_tags": ["prod"], "team": "search"}}, +} + + +def _log_to_promptlayer( + monkeypatch: pytest.MonkeyPatch, track_request_body: object +) -> tuple[list[tuple[str, object]], list[str]]: + posted: Final[list[tuple[str, object]]] = [] + logged: Final[list[str]] = [] + + def promptlayer(request: httpx.Request) -> httpx.Response: + posted.append((request.url.path, json.loads(request.content))) + if request.url.path == "/rest/track-request": + return httpx.Response(200, json=track_request_body) + return httpx.Response(200, json={"success": True}) + + monkeypatch.setenv("PROMPTLAYER_API_KEY", "pl-key") + monkeypatch.setattr( + litellm, "module_level_client", HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(promptlayer))) + ) + PromptLayerLogger().log_event( + kwargs=REQUEST_KWARGS, + response_obj={"id": "chatcmpl-1"}, + start_time=CALL_TIME, + end_time=CALL_TIME, + print_verbose=logged.append, + ) + return posted, logged + + +def test_prompt_layer_tracks_metadata_under_the_request_id_it_got_back(monkeypatch: pytest.MonkeyPatch) -> None: + posted, _ = _log_to_promptlayer(monkeypatch, {"success": True, "request_id": 42}) + + assert [path for path, _ in posted] == ["/rest/track-request", "/rest/track-metadata"] + track_request: Final = posted[0][1] + assert isinstance(track_request, dict) + assert track_request["tags"] == ["prod"] + assert track_request["kwargs"] == { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "temperature": 0.2, + } + assert posted[1][1] == {"request_id": 42, "api_key": "pl-key", "metadata": {"team": "search"}} + + +@pytest.mark.parametrize( + ("track_request_body", "logged_error"), + [ + ({"success": False, "request_id": 42}, "Promptlayer did not successfully log the response!"), + (["success", "request_id"], "'list' object has no attribute 'get'"), + ], +) +def test_prompt_layer_skips_metadata_and_logs_the_error_when_tracking_did_not_succeed( + monkeypatch: pytest.MonkeyPatch, track_request_body: object, logged_error: str +) -> None: + posted, logged = _log_to_promptlayer(monkeypatch, track_request_body) + + assert [path for path, _ in posted] == ["/rest/track-request"] + assert logged_error in logged[-1] diff --git a/tests/unit/integrations/test_wandb.py b/tests/unit/integrations/test_wandb.py index 9341c002378..dfddf683ad1 100644 --- a/tests/unit/integrations/test_wandb.py +++ b/tests/unit/integrations/test_wandb.py @@ -1,3 +1,8 @@ +import importlib.util +import sys +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import ModuleType, SimpleNamespace from typing import Final from unittest.mock import MagicMock @@ -21,3 +26,65 @@ def test_wandb_logging(): pass except Exception as e: print(e) + + +@dataclass(frozen=True) +class _TraceResult: + inputs: Mapping[str, object] + outputs: Mapping[str, object] + + +@dataclass(frozen=True) +class _TraceSpan: + name: str + attributes: Mapping[str, object] + start_time_ms: int + end_time_ms: int + span_kind: str + results: Sequence[_TraceResult] + + +@dataclass(frozen=True) +class _TraceTree: + root_span: _TraceSpan + model_dict: Mapping[str, object] + + +def _weights_biases_with_stub_trace_tree(monkeypatch: pytest.MonkeyPatch) -> ModuleType: + data_types: Final = ModuleType("wandb.sdk.data_types") + data_types.trace_tree = SimpleNamespace( + Result=_TraceResult, Span=_TraceSpan, WBTraceTree=_TraceTree, SpanKind=SimpleNamespace(LLM="LLM") + ) + for name in ("wandb", "wandb.sdk"): + monkeypatch.setitem(sys.modules, name, ModuleType(name)) + monkeypatch.setitem(sys.modules, "wandb.sdk.data_types", data_types) + spec: Final = importlib.util.find_spec("litellm.integrations.weights_biases") + assert spec is not None and spec.loader is not None + module: Final = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +def test_wandb_chat_completion_trace_lists_each_request_message_with_its_role(monkeypatch: pytest.MonkeyPatch) -> None: + weights_biases: Final = _weights_biases_with_stub_trace_tree(monkeypatch) + request: Final = { + "model": "gpt-4o", + "messages": [{"role": "system", "content": "Be brief."}, {"role": "user", "content": "Hi"}], + } + response: Final = { + "object": "chat.completion", + "model": "gpt-4o", + "created": 1767225600, + "choices": [{"message": {"role": "assistant", "content": "Hello!"}}], + } + + trace: Final = weights_biases.OpenAIRequestResponseResolver()(request, response, 1.5) + + assert trace.root_span.results == [ + _TraceResult( + inputs={"request": "\n\n**system**: Be brief.\n\n\n**user**: Hi\n"}, + outputs={"response": "\n\n**assistant**: Hello!\n"}, + ) + ] + assert trace.root_span.end_time_ms - trace.root_span.start_time_ms == 1500 diff --git a/tests/unit/integrations/test_weave_otel.py b/tests/unit/integrations/test_weave_otel.py index e8f06c00e41..a3a022d514b 100644 --- a/tests/unit/integrations/test_weave_otel.py +++ b/tests/unit/integrations/test_weave_otel.py @@ -1,13 +1,20 @@ +import json import os +from collections.abc import Mapping +from typing import Final from unittest.mock import MagicMock, patch import pytest +from opentelemetry.sdk.trace import TracerProvider +from litellm.integrations._types.open_inference import SpanAttributes as OpenInferenceSpanAttributes from litellm.integrations.weave.weave_otel import ( + WeaveLLMObsOTELAttributes, _set_weave_specific_attributes, get_weave_otel_config, ) from litellm.types.integrations.weave_otel import WeaveOtelConfig, WeaveSpanAttributes +from litellm.types.utils import ModelResponse def test_get_weave_otel_config(): @@ -154,3 +161,47 @@ def test_set_weave_specific_attributes_thread_id_and_is_turn(): mock_safe_set.assert_any_call( mock_span, WeaveSpanAttributes.IS_TURN.value, True ) + + +_TRACER: Final = TracerProvider().get_tracer("weave-otel-tests") +WEAVE_MESSAGES: Final = [{"role": "user", "content": "weather in SF?"}] +WEAVE_TOOLS: Final = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] +WEAVE_FUNCTIONS: Final = [{"name": "get_time", "parameters": {}}] + + +@pytest.mark.parametrize( + ("request_kwargs", "expected_input"), + [ + ( + { + "messages": WEAVE_MESSAGES, + "optional_params": {"tools": WEAVE_TOOLS, "functions": WEAVE_FUNCTIONS, "top_p": 0.5}, + }, + {"messages": WEAVE_MESSAGES, "functions": WEAVE_FUNCTIONS, "tools": WEAVE_TOOLS}, + ), + ({"messages": WEAVE_MESSAGES, "optional_params": None}, {"messages": WEAVE_MESSAGES}), + ], +) +def test_weave_input_value_carries_the_request_functions_and_tools( + request_kwargs: Mapping[str, object], expected_input: Mapping[str, object] +) -> None: + span: Final = _TRACER.start_span("litellm_request") + + WeaveLLMObsOTELAttributes.set_messages(span, request_kwargs) + + assert json.loads(str(span.attributes[OpenInferenceSpanAttributes.INPUT_VALUE])) == expected_input + + +def test_weave_output_value_is_the_dumped_model_response() -> None: + span: Final = _TRACER.start_span("litellm_request") + response: Final = ModelResponse( + id="chatcmpl-1", + model="gpt-4o", + choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Sunny."}}], + ) + + _set_weave_specific_attributes(span, {"model": "gpt-4o"}, response) + + output: Final = json.loads(str(span.attributes[OpenInferenceSpanAttributes.OUTPUT_VALUE])) + assert output["id"] == "chatcmpl-1" + assert output["choices"][0]["message"]["content"] == "Sunny." diff --git a/tests/unit/interactions/test_gemini_interactions_transformation.py b/tests/unit/interactions/test_gemini_interactions_transformation.py index ebc99e8e88b..2f167254b0f 100644 --- a/tests/unit/interactions/test_gemini_interactions_transformation.py +++ b/tests/unit/interactions/test_gemini_interactions_transformation.py @@ -598,3 +598,14 @@ class TestTransformRequestSchemaCoalescing: ) assert len(optional_params["response_format"]) == 1 assert body_retry["response_format"] == body["response_format"] + + def test_image_config_that_is_not_a_mapping_is_rejected(self, config): + with pytest.raises(TypeError, match=r"^'str' object is not a mapping$"): + config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="draw a sunset", + optional_params={"generation_config": {"image_config": "1:1"}}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) diff --git a/tests/unit/interactions/test_litellm_responses_bridge.py b/tests/unit/interactions/test_litellm_responses_bridge.py index 9330ae44ad8..a40b8ca0365 100644 --- a/tests/unit/interactions/test_litellm_responses_bridge.py +++ b/tests/unit/interactions/test_litellm_responses_bridge.py @@ -6,7 +6,7 @@ the litellm_responses bridge provider, which calls litellm.responses() internall """ from collections.abc import Mapping -from datetime import datetime +from datetime import datetime, timezone from typing import Final import pytest @@ -144,6 +144,26 @@ class TestBridgeResponseTransformation: assert datetime.fromisoformat(interaction.created).timestamp() == 1700000000 assert interaction.updated == interaction.created + @pytest.mark.parametrize( + ("created_at", "expected_created"), + [ + (datetime(2026, 1, 2, 3, 4, 5, tzinfo=timezone.utc), "2026-01-02T03:04:05+00:00"), + (None, None), + ], + ids=["datetime", "missing"], + ) + def test_responses_response_created_at_that_is_not_a_timestamp( + self, created_at: datetime | None, expected_created: str | None + ): + reply: Final = _responses_api_response("completed", None).model_copy(update={"created_at": created_at}) + + interaction: Final = LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_response( + reply, None + ) + + assert interaction.created == expected_created + assert interaction.updated == expected_created + def test_responses_response_with_a_status_that_is_not_text_is_refused(self): reply: Final = _responses_api_response("completed", None).model_copy(update={"status": 7}) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3964f472bd9..fd7ec0d6bca 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( calculate_cache_writing_cost, generic_cost_per_token, get_billed_token_rates, + get_regional_uplift_multiplier, get_token_type_cost_breakdown, ) from litellm.llms.gemini.image_generation.cost_calculator import ( @@ -2781,6 +2782,45 @@ def test_fedramp_uplift_defaults_to_one(model_info: ModelInfo, data_residency: s assert get_fedramp_uplift_multiplier(model_info, data_residency) == 1.0 +@pytest.mark.parametrize( + "model_info, data_residency", + [({f"regional_processing_uplift_multiplier_{r.value}": "1.25"}, r.value.upper()) for r in DataResidency], +) +def test_regional_uplift_reads_the_multiplier_named_by_each_residency(model_info: ModelInfo, data_residency: str): + assert get_regional_uplift_multiplier(model_info, data_residency) == 1.25 + + +@pytest.mark.parametrize( + "model_info", + [{"regional_processing_uplift_multiplier_us": value} for value in (1.25, "1.25", b"1.25", bytearray(b"1.25"))], +) +def test_regional_uplift_accepts_every_value_float_converts(model_info: ModelInfo): + assert get_regional_uplift_multiplier(model_info, "us") == 1.25 + + +class _IndexOnly: + def __index__(self) -> int: + return 2 + + +@pytest.mark.parametrize("model_info", [{"regional_processing_uplift_multiplier_us": _IndexOnly()}]) +def test_regional_uplift_accepts_an_index_only_value(model_info: ModelInfo): + assert get_regional_uplift_multiplier(model_info, "us") == 2.0 + + +@pytest.mark.parametrize( + "model_info, data_residency", + [ + ({"regional_processing_uplift_multiplier_eu": "not-a-number"}, "eu"), + ({"regional_processing_uplift_multiplier_eu": [1.25]}, "eu"), + ({"regional_processing_uplift_multiplier_apac": 1.25}, "apac"), + ({"regional_processing_uplift_multiplier_eu": 1.25}, None), + ], +) +def test_regional_uplift_defaults_to_one(model_info: ModelInfo, data_residency: str | None): + assert get_regional_uplift_multiplier(model_info, data_residency) == 1.0 + + def test_service_tier_suffixes_constant_in_sync_with_enum(): from litellm.litellm_core_utils.llm_cost_calc.utils import _SERVICE_TIER_SUFFIXES from litellm.types.utils import ServiceTier diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 5d6d14a4ca8..02476545cae 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -4,7 +4,7 @@ import datetime import json import logging import threading -from collections.abc import Iterator, Mapping +from collections.abc import Container, Iterator, Mapping from typing import Final import httpx @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_or_create_metadata_bucket, get_parent_otel_span_from_kwargs, get_provider_response_headers_from_hidden_params, + json_contains, map_finish_reason, normalize_drop_params, preserve_upstream_non_openai_attributes, @@ -33,6 +34,7 @@ from litellm.litellm_core_utils.core_helpers import ( redact_nested_match_and_regex_keys, response_cost_in_hidden_params, set_provider_response_headers_in_hidden_params, + subscript, unbind_budget_reservation_from_callbacks, without, ) @@ -855,6 +857,40 @@ def test_billed_usage_cost_keeps_the_upstream_charge_after_litellm_wrote_its_bil assert billed_usage_cost(unreported, "openai") is None +def test_subscript_reads_keys_and_indexes_and_rejects_values_without_item_access() -> None: + assert subscript({"role": "user"}, "role") == "user" + assert subscript(["system", "user"], 1) == "user" + with pytest.raises(TypeError, match="'int' object is not subscriptable"): + subscript(5, "role") + + +class _IterOnly: + def __iter__(self) -> Iterator[object]: + return iter(("error",)) + + +class _GetItemOnly: + def __getitem__(self, index: int) -> object: + return ("error",)[index] + + +@pytest.mark.parametrize( + "value", + [{"error": 1}, {"data": 1}, ["error"], "an error", "", _IterOnly(), _GetItemOnly()], +) +def test_json_contains_answers_what_the_in_operator_answers(value: Container[object]) -> None: + assert json_contains(value, "error") is ("error" in value) + + +@pytest.mark.parametrize("value", [5, 1.5, None, True, object()]) +def test_json_contains_raises_the_in_operators_own_error(value: Container[object]) -> None: + with pytest.raises(TypeError) as native: + _ = "error" in value + with pytest.raises(TypeError) as ours: + json_contains(value, "error") + assert str(ours.value) == str(native.value) + + def test_without_drops_the_named_keys_and_keeps_the_rest_in_order() -> None: source: Final = {"a": 1, "b": 2, "c": 3, "d": 4} diff --git a/tests/unit/litellm_core_utils/test_dd_tracing.py b/tests/unit/litellm_core_utils/test_dd_tracing.py index 251f1842e6e..e626d327989 100644 --- a/tests/unit/litellm_core_utils/test_dd_tracing.py +++ b/tests/unit/litellm_core_utils/test_dd_tracing.py @@ -1,10 +1,14 @@ import json +from dataclasses import dataclass +from typing import Final from unittest.mock import MagicMock, patch import pytest +from litellm.litellm_core_utils import dd_tracing from litellm.litellm_core_utils.dd_tracing import ( + set_active_span_tag, should_use_dd_profiler, _should_use_dd_tracer, ) @@ -97,3 +101,33 @@ def test_should_use_dd_profiler(): mock_get_secret.return_value = False assert should_use_dd_profiler() is False mock_get_secret.assert_called_once_with("USE_DDPROFILER", False) + + +@dataclass(frozen=True) +class _CurrentSpanTracer: + span: object + + def current_span(self) -> object: + return self.span + + +class _SetTagOnlySpan: + def __init__(self) -> None: + self.tags: dict[str, str] = {} + + def set_tag(self, key: str, value: str) -> None: + self.tags[key] = value + + +def test_set_active_span_tag_uses_set_tag_on_a_span_without_set_tag_str(monkeypatch: pytest.MonkeyPatch) -> None: + span: Final = _SetTagOnlySpan() + monkeypatch.setattr(dd_tracing, "tracer", _CurrentSpanTracer(span)) + + assert set_active_span_tag("litellm.call_id", "call-1") is True + assert span.tags == {"litellm.call_id": "call-1"} + + +def test_set_active_span_tag_reports_false_for_a_span_that_takes_no_tags(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(dd_tracing, "tracer", _CurrentSpanTracer(object())) + + assert set_active_span_tag("litellm.call_id", "call-1") is False diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index e8ea0c04f6f..58fd03e0cbc 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -2701,6 +2701,55 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests( dummy_logger.log_stream_event.assert_not_called() +def test_sync_success_callbacks_run_when_litellm_params_is_a_params_model( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.router import GenericLiteLLMParams + + class DummyLogger(CustomLogger): + pass + + logging_obj.stream = False + logging_obj.call_type = "completion" + logging_obj.model_call_details["litellm_params"] = GenericLiteLLMParams() + dummy_logger = DummyLogger() + dummy_logger.log_success_event = MagicMock() + logging_obj.dynamic_success_callbacks = [dummy_logger] + logging_obj.dynamic_async_success_callbacks = [] + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + + logging_obj.success_handler( + result=ModelResponse( + id="resp-123", + model="gpt-4o-mini", + choices=[{"message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop", "index": 0}], + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + ) + ) + + dummy_logger.log_success_event.assert_called_once() + + +@pytest.mark.parametrize(("extra_params", "expected"), [({}, True), ({"no-log": True}, False)]) +def test_should_run_callback_reads_no_log_from_a_params_model( + logging_obj: LitellmLogging, + monkeypatch: pytest.MonkeyPatch, + extra_params: dict[str, object], + expected: bool, +) -> None: + from litellm.types.router import GenericLiteLLMParams + + monkeypatch.setattr(litellm, "global_disable_no_log_param", False) + + should_run: Final = logging_obj.should_run_callback( + callback="langfuse", litellm_params=GenericLiteLLMParams(**extra_params), event_hook="success_handler" + ) + + assert should_run is expected + + @pytest.mark.asyncio @pytest.mark.parametrize("different_headers", [False, True]) async def test_completed_success_preserves_generic_logger_configuration( diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index cf8de9c40f5..94fd915844a 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -3523,6 +3523,34 @@ def test_is_chunk_non_empty_with_valid_tool_calls( ) +@pytest.mark.parametrize( + ("tool_calls", "non_empty"), + [([{"id": "call_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}], True), ([], False)], +) +def test_is_chunk_non_empty_reads_tool_calls_from_the_completion_obj( + initialized_custom_stream_wrapper: CustomStreamWrapper, tool_calls: list[dict[str, object]], non_empty: bool +): + assert ( + initialized_custom_stream_wrapper.is_chunk_non_empty( + completion_obj={"tool_calls": tool_calls}, + model_response=ModelResponseStream(), + response_obj={}, + ) + is non_empty + ) + + +def test_is_chunk_non_empty_rejects_completion_obj_tool_calls_without_a_length( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + with pytest.raises(TypeError, match="object of type 'generator' has no len"): + initialized_custom_stream_wrapper.is_chunk_non_empty( + completion_obj={"tool_calls": (tool_call for tool_call in ())}, + model_response=ModelResponseStream(), + response_obj={}, + ) + + def _make_chunk(content: Optional[str]) -> ModelResponseStream: return ModelResponseStream( id="test", @@ -4655,6 +4683,25 @@ def test_dispatch_triton_stream( assert initialized_custom_stream_wrapper.received_finish_reason == "stop" +@pytest.mark.parametrize( + "chunk", + [ + 'data: {"text_output": "triton text", "is_finished": true, "stop_reason": "stop", "generated_token_count": 2}', + b'data: {"text_output": "triton text", "is_finished": true, "stop_reason": "stop", "generated_token_count": 2}', + ], +) +def test_handle_triton_stream_parses_sse_text_output_lines( + initialized_custom_stream_wrapper: CustomStreamWrapper, chunk: str | bytes +): + assert initialized_custom_stream_wrapper.handle_triton_stream(chunk) == { + "text": "triton text", + "is_finished": True, + "finish_reason": "stop", + "prompt_tokens": 0, + "completion_tokens": 2, + } + + def test_dispatch_text_completion_openai_with_usage( initialized_custom_stream_wrapper: CustomStreamWrapper, ): diff --git a/tests/unit/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py b/tests/unit/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py index 01eff7f6bbf..1a10acd51ad 100644 --- a/tests/unit/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py +++ b/tests/unit/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py @@ -1,9 +1,16 @@ """Tests for litellm/llms/a2a/chat/guardrail_translation/handler.py.""" import json +from typing import Final, Literal +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey +from litellm.types.agents import LiteLLMSendMessageResponse +from litellm.types.utils import GenericGuardrailAPIInputs def _text_event(text: str) -> str: @@ -37,3 +44,53 @@ class TestA2AGuardrailHandlerStreamingScanKey: def test_json_items_that_are_not_objects_are_ignored(self): key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hi"), "[1, 2]", "5"]) assert key.texts == ("hi",) + + +SSN_REPLY: Final = {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": "SSN 123-45-6789"}]} + + +class _MaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Logging | None = None, + ) -> GenericGuardrailAPIInputs: + return {"texts": [text.replace("123-45-6789", "[SSN]") for text in inputs.get("texts", [])]} + + +class _BlockingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Logging | None = None, + ) -> GenericGuardrailAPIInputs: + raise ValueError(f"blocked {inputs.get('texts')}") + + +class TestA2AGuardrailHandlerOutputResponse: + async def test_dict_response_gets_the_guardrailed_text_in_place(self): + response: Final = {"jsonrpc": "2.0", "id": "req-1", "result": json.loads(json.dumps(SSN_REPLY))} + + processed: Final = await A2AGuardrailHandler().process_output_response( + response, _MaskingGuardrail(guardrail_name="mask-ssn") + ) + + assert processed is response + assert response["result"]["parts"][0]["text"] == "SSN [SSN]" + + async def test_pydantic_response_text_reaches_the_guardrail(self): + response: Final = LiteLLMSendMessageResponse(id="req-1", result=json.loads(json.dumps(SSN_REPLY))) + + with pytest.raises(ValueError, match=r"blocked \['SSN 123-45-6789'\]"): + await A2AGuardrailHandler().process_output_response(response, _BlockingGuardrail(guardrail_name="block")) + + async def test_response_of_an_unknown_type_is_returned_without_running_the_guardrail(self): + processed: Final = await A2AGuardrailHandler().process_output_response( + "SSN 123-45-6789", _BlockingGuardrail(guardrail_name="block") + ) + + assert processed == "SSN 123-45-6789" diff --git a/tests/unit/llms/anthropic/messages/test_stream_assembly.py b/tests/unit/llms/anthropic/messages/test_stream_assembly.py index f8cd00ac402..81b72bfbf06 100644 --- a/tests/unit/llms/anthropic/messages/test_stream_assembly.py +++ b/tests/unit/llms/anthropic/messages/test_stream_assembly.py @@ -7,6 +7,7 @@ from litellm.llms.anthropic.messages.stream_assembly import ( anthropic_message_delta_fields, anthropic_sse_error_events, assemble_anthropic_sse_stream, + collapse_pure_text_chunks, has_unreadable_json_event, is_anthropic_sse_stream, is_raw_sse_stream, @@ -379,3 +380,24 @@ def test_assembly_through_an_upstream_error_reads_past_the_error_frame(error_fra chunks: Final = (*ANTHROPIC_STREAM[:4], error_frame) assert assemble_anthropic_sse_stream(chunks, through_upstream_error=True) is not None + + +def test_collapsing_a_text_run_rejects_a_delta_whose_text_is_not_a_string(): + chunks: Final = ( + _frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}, + ), + _frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": 7}}, + ), + _frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + ) + + with pytest.raises(TypeError, match="sequence item 1: expected str instance, int found"): + collapse_pure_text_chunks(chunks) diff --git a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index 51ba2c34cd7..b14500bea1e 100644 --- a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -1,6 +1,7 @@ import base64 import json from collections.abc import Mapping +from io import BytesIO from typing import Final import httpx @@ -144,6 +145,27 @@ def test_flux2_image_edit_rejects_too_many_references(model: str, reference_imag ) +def _flux2_reference_request(image: object) -> dict[str, object]: + request, _ = AzureFoundryFlux2ImageEditConfig().transform_image_edit_request( + model="FLUX.2-pro", + prompt="Add a red hat", + image=image, + image_edit_optional_request_params={}, + litellm_params={}, + headers={}, + ) + return request + + +def test_flux2_image_edit_reads_a_file_like_reference_image(): + assert _flux2_reference_request(BytesIO(b"png-bytes"))["input_image"] == base64.b64encode(b"png-bytes").decode() + + +def test_flux2_image_edit_rejects_a_reference_image_it_cannot_read(): + with pytest.raises(ValueError, match="Unsupported image type: "): + _flux2_reference_request("cat.png") + + @pytest.mark.parametrize("dimensions", ({"size": "2048x1024"}, {"width": 2048, "height": 1024}, {"width": "2048", "height": "1024"})) @pytest.mark.usefixtures("local_model_cost_map") def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[str, int | str]): diff --git a/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py b/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py index 6ca70324162..e3aa01a609d 100644 --- a/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py +++ b/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py @@ -1,4 +1,17 @@ +import re +from datetime import datetime +from typing import Final + +import httpx +import pytest + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.image_edit.stability_transformation import BedrockStabilityImageEditConfig +from litellm.types.utils import ImageResponse +from tests.unit.conftest import in_operator_error + +INPAINT_MODEL: Final = "stability.stable-image-inpaint-v1:0" def test_transform_image_edit_request_returns_the_json_body_and_no_files(): @@ -12,3 +25,56 @@ def test_transform_image_edit_request_returns_the_json_body_and_no_files(): ) assert request == ({"output_format": "png", "prompt": "add a red hat", "image": "iVBORy1ieXRlcw=="}, {}) + + +def _edit_response(body: object) -> ImageResponse: + return BedrockStabilityImageEditConfig().transform_image_edit_response( + model=INPAINT_MODEL, + raw_response=httpx.Response(200, json=body), + logging_obj=Logging( + model=INPAINT_MODEL, + messages=[], + stream=False, + call_type="image_edit", + start_time=datetime(2026, 1, 1), + litellm_call_id="stability-edit-call", + function_id="stability-edit-function", + ), + ) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_transform_image_edit_response_returns_every_non_empty_image_in_order(): + response = _edit_response({"images": ["aW1hZ2Ux", "", "aW1hZ2Uy"], "finish_reasons": [None], "seeds": [7]}) + + assert [(image.b64_json, image.url) for image in response.data] == [("aW1hZ2Ux", None), ("aW1hZ2Uy", None)] + + +@pytest.mark.parametrize( + ("body", "status_code", "message"), + [ + ({"errors": ["prompt was flagged"]}, 200, "Bedrock Stability error: ['prompt was flagged']"), + ( + {"images": [], "finish_reasons": ["Filter reason: prompt"]}, + 400, + "Bedrock Stability error: Filter reason: prompt", + ), + ], +) +def test_transform_image_edit_response_raises_the_error_bedrock_reports(body: object, status_code: int, message: str): + with pytest.raises(BedrockError) as error: + _edit_response(body) + + assert (error.value.status_code, error.value.message) == (status_code, message) + + +@pytest.mark.parametrize( + ("body", "message"), + [ + (5, in_operator_error(5)), + ({"finish_reasons": 5}, "'int' object is not subscriptable"), + ], +) +def test_transform_image_edit_response_rejects_a_scalar_where_bedrock_sends_a_container(body: object, message: str): + with pytest.raises(TypeError, match=re.escape(message)): + _edit_response(body) diff --git a/tests/unit/llms/cohere/test_cohere_common_utils.py b/tests/unit/llms/cohere/test_cohere_common_utils.py new file mode 100644 index 00000000000..1d17824a181 --- /dev/null +++ b/tests/unit/llms/cohere/test_cohere_common_utils.py @@ -0,0 +1,18 @@ +from typing import Final + +import pytest + +from litellm.llms.cohere.common_utils import ModelResponseIterator + + +def test_stream_parses_object_chunks_and_rejects_a_chunk_that_is_not_an_object() -> None: + stream: Final = ModelResponseIterator( + streaming_response=iter(['{"text": "Hi", "index": 1}', "[1, 2]"]), + sync_stream=True, + ) + + first: Final = next(stream) + + assert (first["text"], first["index"], first["is_finished"]) == ("Hi", 1, False) + with pytest.raises(AttributeError, match=r"^'list' object has no attribute 'get'$"): + next(stream) diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index a758f4d4250..ba8fb3bf107 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -4253,6 +4253,7 @@ def test_video_generation_with_input_reference_keeps_file_multipart(): AZURE_AI_BASE = "https://myfoundry.services.ai.azure.com" AZURE_AI_CHAT_COMPLETIONS_URL = f"{AZURE_AI_BASE}/models/chat/completions" + def _a_tool_with_an_unsupported_field() -> dict: return { "type": "function", @@ -4260,6 +4261,7 @@ def _a_tool_with_an_unsupported_field() -> dict: "strict": True, } + A_COMPLETION = { "id": "chatcmpl-1", "object": "chat.completion", @@ -4459,6 +4461,45 @@ async def test_create_file_async_awaits_the_provider_credential_hook_instead_of_ assert client.post.call_args.kwargs["url"] == "https://files.example/v1/files" +def test_create_file_sync_sends_a_presigned_upload_with_its_own_method_url_headers_and_body(): + provider_config = Mock(spec=BaseFilesConfig) + provider_config.validate_environment.return_value = {"x-api-key": "unused-for-presigned"} + provider_config.get_complete_file_url.return_value = "https://files.example/v1/files" + provider_config.transform_create_file_request.return_value = { + "method": "PUT", + "url": "https://bucket.s3.example/batch.jsonl", + "headers": {"Authorization": "AWS4-HMAC-SHA256 signed"}, + "data": b'{"recordId": "req-1"}', + } + file_object = object() + provider_config.transform_create_file_response.return_value = file_object + upload_response = Mock(spec=httpx.Response) + client = Mock(spec=HTTPHandler) + client.put.return_value = upload_response + + result = BaseLLMHTTPHandler().create_file( + create_file_data={"file": b"{}", "purpose": "batch"}, + litellm_params={}, + provider_config=provider_config, + headers={}, + api_base=None, + api_key=None, + logging_obj=Mock(), + client=client, + timeout=30, + ) + + assert result is file_object + client.post.assert_not_called() + client.put.assert_called_once_with( + url="https://bucket.s3.example/batch.jsonl", + headers={"Authorization": "AWS4-HMAC-SHA256 signed"}, + data=b'{"recordId": "req-1"}', + timeout=30, + ) + assert provider_config.transform_create_file_response.call_args.kwargs["raw_response"] is upload_response + + @pytest.mark.asyncio async def test_create_batch_async_validates_credentials_off_the_event_loop(): provider_config = Mock(spec=BaseBatchesConfig) diff --git a/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py b/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py index 27afbee3173..5dd8b165a0f 100644 --- a/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py @@ -3,19 +3,20 @@ Unit tests for DashScope embedding transformation. """ import json +import re from typing import Final from unittest.mock import MagicMock import httpx import pytest - from litellm.llms.dashscope.common_utils import DashScopeError from litellm.llms.dashscope.embed.transformation import ( DEFAULT_API_BASE, DashScopeEmbeddingConfig, ) from litellm.types.utils import EmbeddingResponse +from tests.unit.conftest import in_operator_error WORKSPACE: Final = "https://llm-abc123.cn-beijing.maas.aliyuncs.com" @@ -171,3 +172,22 @@ def test_transform_embedding_response_error(): ) assert exc.value.status_code == 401 assert "Incorrect API key" in exc.value.message + + +def test_transform_embedding_response_rejects_a_scalar_body(): + raw = httpx.Response( + status_code=200, + content=b"5", + request=httpx.Request("POST", "https://example.com"), + ) + with pytest.raises(TypeError, match=f"^{re.escape(in_operator_error(5))}$"): + DashScopeEmbeddingConfig().transform_embedding_response( + model="text-embedding-v4", + raw_response=raw, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + api_key="sk-test", + request_data={"input": ["a"]}, + optional_params={}, + litellm_params={}, + ) diff --git a/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py b/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py index d94f6a53c89..c8071687b4f 100644 --- a/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py +++ b/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py @@ -4,6 +4,7 @@ Unit tests for DataForSEO Search functionality. import datetime import os +import re from unittest.mock import AsyncMock, patch import httpx @@ -15,6 +16,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +from tests.unit.conftest import in_operator_error @pytest.mark.asyncio @@ -172,6 +174,19 @@ def test_dataforseo_malformed_item_is_rejected_without_naming_its_position(): assert "429" not in str(exc_info.value) +@pytest.mark.parametrize( + ("payload", "message"), + [ + (5, in_operator_error(5)), + ({"tasks": 5}, "object of type 'int' has no len()"), + ({"tasks": [{"status_code": 20000, "result": 5}]}, "object of type 'int' has no len()"), + ], +) +def test_dataforseo_scalar_where_the_response_needs_a_container_raises_type_error(payload: object, message: str): + with pytest.raises(TypeError, match=re.escape(message)): + _transform(payload) + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 5f2af34e513..508bd264f19 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -1,4 +1,5 @@ import contextlib +import re from typing import Final from unittest.mock import AsyncMock, patch @@ -16,6 +17,7 @@ from litellm.llms.gemini.common_utils import ( ) from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter as GeminiCountTokensHandler +from tests.unit.conftest import in_operator_error _JSON_BODY: Final = TypeAdapter(dict[str, JsonValue]) @@ -498,6 +500,18 @@ class TestGoogleAIStudioTokenCounter: # Verify the contents are unchanged assert cleaned_contents == contents_without_function_response + @pytest.mark.parametrize( + ("part", "message"), + [ + (7, in_operator_error(7)), + (frozenset({"functionResponse"}), "'frozenset' object is not subscriptable"), + ({"functionResponse": "done"}, "FunctionResponse() argument after ** must be a mapping, not str"), + ], + ) + def test_clean_contents_for_gemini_api_rejects_a_malformed_part(self, part: object, message: str) -> None: + with pytest.raises(TypeError, match=re.escape(message)): + GeminiCountTokensHandler()._clean_contents_for_gemini_api([{"role": "user", "parts": [part]}]) + def _map_image_config(raw_image_config: str) -> dict[str, object]: return map_openai_image_params_to_gemini( diff --git a/tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py b/tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py index 35ca93319f5..913130e5d5c 100644 --- a/tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py +++ b/tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py @@ -2,6 +2,10 @@ Tests for litellm.llms.gigachat.chat.streaming """ +import json +from collections.abc import AsyncIterator +from typing import Final + from litellm.llms.gigachat.chat.streaming import GigaChatModelResponseIterator @@ -85,3 +89,32 @@ class TestChunkParserUsage: assert parsed["usage"]["prompt_tokens"] == 45 assert parsed["usage"]["total_tokens"] == 52 assert parsed["usage"]["prompt_tokens_details"]["cached_tokens"] == 20 + + +def _content_chunk(text: str) -> dict[str, object]: + return {"choices": [{"delta": {"content": text}, "index": 0, "finish_reason": None}]} + + +STREAM_ITEMS: Final = ( + f"data: {json.dumps(_content_chunk('Hel'))}", + _content_chunk("lo"), + "data: {not json", + "data: [DONE]", + f"data: {json.dumps(_content_chunk('after done'))}", +) + + +class TestStreamIteration: + def test_sync_stream_reads_sse_lines_and_parsed_chunks_until_done(self): + iterator = GigaChatModelResponseIterator(streaming_response=iter(STREAM_ITEMS), sync_stream=True) + + assert [chunk["text"] for chunk in iterator] == ["Hel", "lo", ""] + + async def test_async_stream_reads_sse_lines_and_parsed_chunks_until_done(self): + async def stream() -> AsyncIterator[object]: + for item in STREAM_ITEMS: + yield item + + iterator = GigaChatModelResponseIterator(streaming_response=stream(), sync_stream=False) + + assert [chunk["text"] async for chunk in iterator] == ["Hel", "lo", ""] diff --git a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py index 7040fbe5fcc..ef5b16606f4 100644 --- a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py @@ -1,4 +1,5 @@ import json +import re from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -7,6 +8,9 @@ import pytest import respx import litellm +from litellm.llms.huggingface.embedding.handler import HuggingFaceEmbedding +from litellm.types.utils import EmbeddingResponse +from tests.unit.conftest import in_operator_error MOCK_EMBEDDING_RESPONSE = [[0.1, 0.2, 0.3, 0.4, 0.5]] @@ -227,3 +231,46 @@ async def test_huggingface_embedding_sends_the_caller_timeout( assert lookup.called assert route.calls.last.request.extensions["timeout"]["read"] == 45 assert response.data[0]["embedding"] == MOCK_EMBEDDING_RESPONSE[0] + + +async def _embedding(sync_mode: bool, **kwargs: object) -> object: + return litellm.embedding(**kwargs) if sync_mode else await litellm.aembedding(**kwargs) + + +@pytest.mark.respx(assert_all_called=False) +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_huggingface_embedding_raises_the_error_the_endpoint_returns( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond(json={"max_position_embeddings": 512}) + respx_mock.get(HF_MODEL_LOOKUP_URL).respond(json={"pipeline_tag": "feature-extraction"}) + route: Final = respx_mock.post(HF_EMBED_URL).respond(json={"error": "Model BAAI/bge-m3 is currently loading"}) + request: Final = {"model": "huggingface/BAAI/bge-m3", "input": ["hello"], "api_key": "hf_test"} + + with pytest.raises( + litellm.APIError, match="HuggingfaceException - Model BAAI/bge-m3 is currently loading" + ) as exc_info: + await _embedding(sync_mode, **request) + + assert route.called + assert exc_info.value.status_code == 500 + + +@pytest.mark.parametrize( + ("embeddings", "message"), + [ + (5, in_operator_error(5)), + ([None], "'NoneType' object is not subscriptable"), + ({"similarities": [1]}, "cannot unpack non-iterable int object"), + ], +) +def test_process_embedding_response_rejects_a_malformed_body(embeddings: object, message: str) -> None: + with pytest.raises(TypeError, match=f"^{re.escape(message)}$"): + HuggingFaceEmbedding()._process_embedding_response( + embeddings=embeddings, + model_response=EmbeddingResponse(), + model="BAAI/bge-m3", + input=["hello"], + encoding=MagicMock(), + ) diff --git a/tests/unit/llms/litellm_proxy/skills/test_code_execution.py b/tests/unit/llms/litellm_proxy/skills/test_code_execution.py index d9c4470759f..a1cd6a4d9b2 100644 --- a/tests/unit/llms/litellm_proxy/skills/test_code_execution.py +++ b/tests/unit/llms/litellm_proxy/skills/test_code_execution.py @@ -1,3 +1,5 @@ +from typing import Final + import pytest from litellm.llms.litellm_proxy.skills.code_execution import ( @@ -104,3 +106,23 @@ class TestHandlerDefaults: handler = CodeExecutionHandler(max_iterations=falsy, sandbox_timeout=falsy) assert handler.max_iterations == DEFAULT_MAX_ITERATIONS assert handler.sandbox_timeout == DEFAULT_SANDBOX_TIMEOUT + + +class TestExecuteWithCodeExecution: + async def test_a_reply_without_tool_calls_ends_the_loop_after_one_call(self): + messages: Final = [{"role": "user", "content": "Say hi"}] + + outcome: Final = await CodeExecutionHandler().execute_with_code_execution( + model="gpt-5.6", + messages=messages, + tools=[LITELLM_CODE_EXECUTION_TOOL], + skill_files={}, + mock_response="Hi there", + ) + + response: Final = outcome["response"] + assert response is not None + assert response.choices[0].message.content == "Hi there" + assert outcome["messages"] == [*messages, {"role": "assistant", "content": "Hi there"}] + assert messages == [{"role": "user", "content": "Say hi"}] + assert (outcome["files"], outcome["execution_results"]) == ([], []) diff --git a/tests/unit/llms/mistral/test_mistral_chat_transformation.py b/tests/unit/llms/mistral/test_mistral_chat_transformation.py index 42673e17acf..696c53c0b5a 100644 --- a/tests/unit/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/unit/llms/mistral/test_mistral_chat_transformation.py @@ -740,6 +740,25 @@ class TestMistralEmptyContentHandling: assert MistralConfig._is_empty_assistant_message(message) is False +@pytest.mark.parametrize("body", [[], "ok"]) +def test_transform_response_rejects_a_body_that_is_not_a_json_object(body: object) -> None: + logging_obj: Final = MagicMock() + logging_obj.model_call_details = {} + + with pytest.raises(AttributeError, match=f"^'{type(body).__name__}' object has no attribute 'get'$"): + MistralConfig().transform_response( + model="mistral/mistral-large-latest", + raw_response=httpx.Response(200, json=body), + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "Hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + class TestMistralFileHandling: """Test suite for Mistral file handling functionality.""" diff --git a/tests/unit/llms/ollama/test_ollama_embedding.py b/tests/unit/llms/ollama/test_ollama_embedding.py index 97f6f8468f4..b32bb9c4415 100644 --- a/tests/unit/llms/ollama/test_ollama_embedding.py +++ b/tests/unit/llms/ollama/test_ollama_embedding.py @@ -1,4 +1,5 @@ import json +import logging from collections.abc import Mapping from types import MappingProxyType from typing import Final @@ -9,7 +10,11 @@ import pytest import respx import litellm -from litellm.llms.ollama.completion.handler import ollama_aembeddings, ollama_embeddings +from litellm.llms.ollama.completion.handler import ( + _process_ollama_embedding_response, + ollama_aembeddings, + ollama_embeddings, +) from litellm.types.utils import EmbeddingResponse @@ -121,6 +126,39 @@ def test_prompt_eval_fallback_when_missing(mock_embedding_response, mock_encodin assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] +@pytest.mark.parametrize( + ("use_encoding", "prompt_tokens", "level", "message"), + [ + (True, 5, logging.DEBUG, "Ollama response missing prompt_eval_count; estimated with encoding."), + (False, 0, logging.WARNING, "Missing prompt_eval_count and no encoding provided; defaulted to 0."), + ], +) +def test_missing_prompt_eval_count_is_reported_on_the_given_logger( + mock_encoding: MagicMock, + caplog: pytest.LogCaptureFixture, + use_encoding: bool, + prompt_tokens: int, + level: int, + message: str, +) -> None: + logger: Final = logging.getLogger("ollama-embedding-prompt-eval-test") + + with caplog.at_level(logging.DEBUG, logger=logger.name): + response: Final = _process_ollama_embedding_response( + response_json={"embeddings": [[0.1, 0.2]]}, + prompts=["hello"], + model="nomic-embed-text", + model_response=EmbeddingResponse(), + logging_obj=logger, + encoding=mock_encoding if use_encoding else None, + ) + + assert response.usage.prompt_tokens == prompt_tokens + assert [(record.levelno, record.getMessage()) for record in caplog.records if record.name == logger.name] == [ + (level, message) + ] + + @pytest.mark.parametrize( "api_base", [ diff --git a/tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py b/tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py index 97d65935a1b..bf119b01573 100644 --- a/tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py +++ b/tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py @@ -170,3 +170,11 @@ class TestRecraftImageEditTransformation: ) assert "Error transforming image edit response" in str(exc_info.value) + + def test_transform_image_edit_response_rejects_a_body_that_is_not_a_json_object(self): + with pytest.raises(TypeError, match="'int' object is not subscriptable"): + self.config.transform_image_edit_response( + model=self.model, + raw_response=httpx.Response(200, json=5), + logging_obj=self.logging_obj, + ) diff --git a/tests/unit/llms/sap/test_sap_fetch_creds.py b/tests/unit/llms/sap/test_sap_fetch_creds.py index cc52c8991f3..34c10be026b 100644 --- a/tests/unit/llms/sap/test_sap_fetch_creds.py +++ b/tests/unit/llms/sap/test_sap_fetch_creds.py @@ -160,3 +160,15 @@ def test_credentials_without_authentication_mode(monkeypatch): # validate_credentials should raise because no authentication mode is provided with pytest.raises(ValueError, match="SAP AI Core credentials are incomplete"): sap_credentials.validate_credentials(**creds) + + +@pytest.mark.parametrize("vcap_services", [["aicore"], "aicore"]) +def test_vcap_services_that_is_not_an_object_is_rejected(monkeypatch, vcap_services): + _prep_env(monkeypatch) + monkeypatch.setenv("VCAP_SERVICES", json.dumps(vcap_services)) + + with pytest.raises( + AttributeError, + match=f"^'{type(vcap_services).__name__}' object has no attribute 'values'$", + ): + sap_credentials.fetch_credentials() diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index aec68f1fa45..0539a939bb5 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -2962,6 +2962,30 @@ def test_convert_tool_response_with_url_image(content_type: str, nested: bool) - ) +def _assistant_with_image(url: str) -> Message: + return Message( + role="assistant", + content="Here is the picture", + images=[{"type": "image_url", "image_url": {"url": url}, "index": 0}], + ) + + +@pytest.mark.parametrize( + ("message", "needs_fetch"), + [ + (_assistant_with_image("gs://litellm-test-bucket/cat"), True), + (_assistant_with_image("gs://litellm-test-bucket/cat.png"), False), + ("gs://litellm-test-bucket/cat", False), + ], +) +def test_sync_media_fetch_check_reads_pydantic_messages_and_skips_other_objects( + message: object, needs_fetch: bool +) -> None: + from litellm.llms.vertex_ai.gemini.transformation import _openai_messages_may_need_sync_media_fetch + + assert _openai_messages_may_need_sync_media_fetch([message]) is needs_fetch + + CLAUDE_THINKING_SIGNATURE: Final = "EqQBCkYIBxgCKkCfQ2x0b3VkZS1zaWduYXR1cmUtbm90LW1pbnRlZC1ieS1nZW1pbmkSDJ3lXf5sD+QVqpFQmRoM" diff --git a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py index 6ae0bbd0d72..60d7629276f 100644 --- a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py +++ b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py @@ -2,7 +2,8 @@ import base64 import json import os from datetime import datetime -from io import BytesIO +from io import BytesIO, StringIO +from pathlib import Path from typing import Dict, Final from unittest.mock import MagicMock, patch @@ -512,6 +513,30 @@ class TestVertexAIImagenImageEditTransformation: == b"test_bytearray" ) + @pytest.mark.parametrize( + ("image", "expected"), + [ + ([None, b"first image"], b"first image"), + ({"data": base64.b64encode(b"from data").decode()}, b"from data"), + ({"bytes": None, "content": b"from content"}, b"from content"), + ({"path": BytesIO(b"from path")}, b"from path"), + (StringIO("from text reader"), b"from text reader"), + ], + ) + def test_read_all_bytes_unwraps_lists_dicts_and_readers(self, image: object, expected: bytes) -> None: + assert self.config._read_all_bytes(image) == expected + + @pytest.mark.parametrize( + ("image", "refusal"), + [ + ({"path": "/etc/passwd"}, "plain string values are not accepted"), + ({"path": Path("/etc/passwd")}, "filesystem paths are not accepted"), + ], + ) + def test_read_all_bytes_refuses_to_read_a_path_given_inside_a_dict(self, image: object, refusal: str) -> None: + with pytest.raises(ValueError, match=refusal): + self.config._read_all_bytes(image) + _IMAGEN_EDIT_IMAGE_KEY = "bytesBase64Encoded" diff --git a/tests/unit/proxy/client/test_credentials.py b/tests/unit/proxy/client/test_credentials.py index 9de990da487..d49223be710 100644 --- a/tests/unit/proxy/client/test_credentials.py +++ b/tests/unit/proxy/client/test_credentials.py @@ -316,3 +316,25 @@ def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server): client.list() assert time.monotonic() - started < 10 + + +@responses.activate +def test_update_display_name_returns_the_updated_credential(client): + updated = {"credential_name": "azure1", "credential_info": {"display_name": "Azure EU"}} + responses.add( + responses.PATCH, + f"{client._base_url}/credentials/azure1", + json=updated, + status=200, + match=[responses.matchers.json_params_matcher({"display_name": "Azure EU", "credential_info": {}})], + ) + + assert client.update_display_name("azure1", "Azure EU") == updated + + +@responses.activate +def test_update_display_name_unauthorized_error(client): + responses.add(responses.PATCH, f"{client._base_url}/credentials/azure1", status=401, json={"error": "Unauthorized"}) + + with pytest.raises(UnauthorizedError): + client.update_display_name("azure1", "Azure EU") diff --git a/tests/unit/proxy/client/test_keys.py b/tests/unit/proxy/client/test_keys.py index 3bff944144c..14a7891f367 100644 --- a/tests/unit/proxy/client/test_keys.py +++ b/tests/unit/proxy/client/test_keys.py @@ -524,3 +524,27 @@ def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server): client.list() assert time.monotonic() - started < 10 + + +@responses.activate +def test_update_sends_only_the_given_fields_and_returns_the_updated_key(client): + updated = {"key": "key-abc", "key_alias": "renamed", "models": ["gpt-4o"]} + responses.add( + responses.POST, + f"{client._base_url}/key/update", + json=updated, + status=200, + match=[ + responses.matchers.json_params_matcher({"key": "key-abc", "key_alias": "renamed", "models": ["gpt-4o"]}) + ], + ) + + assert client.update(key="key-abc", key_alias="renamed", models=["gpt-4o"]) == updated + + +@responses.activate +def test_update_failure_carries_the_server_response_text(client): + responses.add(responses.POST, f"{client._base_url}/key/update", status=400, body="key not found") + + with pytest.raises(Exception, match="Error updating key: key not found"): + client.update(key="key-abc") diff --git a/tests/unit/proxy/client/test_teams.py b/tests/unit/proxy/client/test_teams.py index b61091ca44b..7aff5c62ff5 100644 --- a/tests/unit/proxy/client/test_teams.py +++ b/tests/unit/proxy/client/test_teams.py @@ -2,7 +2,9 @@ import time import pytest import requests +import responses +from litellm.proxy.client.exceptions import UnauthorizedError from litellm.proxy.client.teams import TeamsManagementClient @@ -18,3 +20,42 @@ def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server): client.list() assert time.monotonic() - started < 10 + + +@responses.activate +def test_list_v2_sends_filters_and_returns_the_page(): + client = TeamsManagementClient(base_url="http://localhost:8000", api_key="sk-test") + page = {"teams": [{"team_id": "team-1"}], "total": 1, "page": 2, "page_size": 5, "total_pages": 1} + responses.add( + responses.GET, + "http://localhost:8000/v2/team/list", + json=page, + status=200, + match=[ + responses.matchers.query_param_matcher( + { + "page": "2", + "page_size": "5", + "sort_order": "desc", + "user_id": "user-1", + "team_alias": "search", + "sort_by": "created_at", + } + ) + ], + ) + + result = client.list_v2( + user_id="user-1", team_alias="search", page=2, page_size=5, sort_by="created_at", sort_order="desc" + ) + + assert result == page + + +@responses.activate +def test_list_v2_unauthorized_error(): + client = TeamsManagementClient(base_url="http://localhost:8000", api_key="sk-bad") + responses.add(responses.GET, "http://localhost:8000/v2/team/list", status=401, json={"error": "Unauthorized"}) + + with pytest.raises(UnauthorizedError): + client.list_v2() diff --git a/tests/unit/proxy/common_utils/test_http_parsing_utils.py b/tests/unit/proxy/common_utils/test_http_parsing_utils.py index 0fc7bb15a32..e4e32eeeac9 100644 --- a/tests/unit/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/unit/proxy/common_utils/test_http_parsing_utils.py @@ -9,7 +9,7 @@ import orjson import pytest from fastapi import HTTPException from fastapi.testclient import TestClient -from starlette.datastructures import FormData +from starlette.datastructures import FormData, Headers, UploadFile from starlette.requests import Request @@ -24,6 +24,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( safe_get_request_query_params, safe_set_request_parsed_body, coerce_numeric_form_fields, + convert_upload_files_to_file_data, get_form_data, get_request_body, get_tags_from_request_body, @@ -1764,3 +1765,24 @@ async def test_read_request_body_closes_the_otel_server_span_when_the_client_dis with pytest.raises(HTTPException): await read_request_body(request) assert (span.end.called, request.state.parent_otel_span) == (True, None), "the auth-time span must not dangle" + + +def _upload(filename: str, content: bytes, content_type: str) -> UploadFile: + return UploadFile(file=io.BytesIO(content), filename=filename, headers=Headers({"content-type": content_type})) + + +@pytest.mark.asyncio +async def test_convert_upload_files_to_file_data_reads_every_upload_into_file_tuples(): + form_data: Final = { + "files": [_upload("a.txt", b"alpha", "text/plain"), _upload("b.json", b"{}", "application/json")], + "file": _upload("skill.zip", b"PK", "application/zip"), + "tags": ["red", "blue"], + "display_title": "My skill", + } + + assert await convert_upload_files_to_file_data(form_data) == { + "files": [("a.txt", b"alpha", "text/plain"), ("b.json", b"{}", "application/json")], + "file": [("skill.zip", b"PK", "application/zip")], + "tags": ["red", "blue"], + "display_title": "My skill", + } diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py index c6a7fbd7cab..4ac3e769030 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py @@ -1,3 +1,4 @@ +import json from unittest.mock import AsyncMock, patch import httpx @@ -227,3 +228,57 @@ async def test_guardrails_ai_process_input(): # Should keep original content when no output fields are present assert result["messages"][0]["content"] == "Original message" + + +class _RecordingGuardrailsClient: + def __init__(self, status_code: int, body: dict[str, object]) -> None: + self.status_code = status_code + self.body = body + self.posts: list[dict[str, object]] = [] + + async def post(self, **kwargs: object) -> httpx.Response: + self.posts.append(kwargs) + return httpx.Response(self.status_code, json=self.body) + + +def _guardrails_ai() -> GuardrailsAI: + return GuardrailsAI(guardrail_name="test_guard", api_base="http://guardrails.test", guard_name="gibberish_guard") + + +@pytest.mark.asyncio +async def test_guardrails_ai_api_request_returns_the_passing_verdict(monkeypatch: pytest.MonkeyPatch) -> None: + verdict = {"callId": "call-1", "rawLlmOutput": "hello", "validatedOutput": "hello", "validationPassed": True} + client = _RecordingGuardrailsClient(200, verdict) + monkeypatch.setattr(litellm, "module_level_aclient", client) + + result = await _guardrails_ai().make_guardrails_ai_api_request(llm_output="hello", request_data={}) + + assert result == verdict + assert client.posts[0]["url"] == "http://guardrails.test/guards/gibberish_guard/validate" + assert json.loads(client.posts[0]["data"]) == {"llmOutput": "hello"} + + +@pytest.mark.asyncio +async def test_guardrails_ai_api_request_blocks_a_failed_validation(monkeypatch: pytest.MonkeyPatch) -> None: + verdict = {"callId": "call-2", "rawLlmOutput": "asdfgh", "validationPassed": False} + monkeypatch.setattr(litellm, "module_level_aclient", _RecordingGuardrailsClient(200, verdict)) + + with pytest.raises(HTTPException) as exc_info: + await _guardrails_ai().make_guardrails_ai_api_request(llm_output="asdfgh", request_data={}) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "Violated guardrail policy", "guardrails_ai_response": verdict} + + +@pytest.mark.asyncio +async def test_guardrails_ai_pre_call_request_returns_the_first_output_text(monkeypatch: pytest.MonkeyPatch) -> None: + body = {"outputs": [{"name": "text", "shape": [1], "datatype": "BYTES", "data": ["cleaned text"]}]} + client = _RecordingGuardrailsClient(200, body) + monkeypatch.setattr(litellm, "module_level_aclient", client) + + result = await _guardrails_ai().make_guardrails_ai_api_request_pre_call_request( + text_input="raw text", request_data={} + ) + + assert result == "cleaned text" + assert json.loads(client.posts[0]["data"])["inputs"][0]["data"] == ["raw text"] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 23288a6136f..dfb72b42000 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -10,6 +10,7 @@ from collections.abc import Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException @@ -2461,3 +2462,38 @@ async def test_lakera_post_call_masks_pii_and_allows(): result_dict = result.model_dump() assert result_dict["choices"][0]["message"]["content"] != "Your email is test@example.com" assert "[MASKED" in result_dict["choices"][0]["message"]["content"] + + +class _RecordingLakeraClient: + def __init__(self, body: dict[str, object]) -> None: + self.body = body + self.posts: list[dict[str, object]] = [] + + async def post(self, **kwargs: object) -> httpx.Response: + self.posts.append(kwargs) + return httpx.Response(200, json=self.body) + + +@pytest.mark.asyncio +async def test_call_v2_guard_returns_the_lakera_verdict_and_logs_it_without_the_payload() -> None: + guardrail: Final = LakeraAIGuardrail(api_key="test_key", guardrail_name="lakera-guard") + verdict: Final = { + "payload": [{"detector_type": "pii/email", "start": 0, "end": 16, "message_id": 0}], + "flagged": True, + "breakdown": [{"detector_type": "pii/email", "detected": True, "message_id": 0}], + } + client: Final = _RecordingLakeraClient(verdict) + guardrail.async_handler = client + request_data: Final[dict[str, object]] = {"metadata": {}} + + response, masked_entity_count = await guardrail.call_v2_guard( + messages=[{"role": "user", "content": "test@example.com"}], + request_data=request_data, + event_type="pre_call", + ) + + assert (response, masked_entity_count) == (verdict, {}) + assert client.posts[0]["url"] == "https://api.lakera.ai/v2/guard" + logged: Final = request_data["metadata"]["standard_logging_guardrail_information"][0] + assert logged["guardrail_status"] == "success" + assert logged["guardrail_response"] == {"flagged": True, "breakdown": verdict["breakdown"]} diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py index 34221720205..ea6892379c2 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py @@ -2803,3 +2803,11 @@ class TestRegistration: guardrail_class_registry["microsoft_purview"] is MicrosoftPurviewDLPGuardrail ) + + +@pytest.mark.parametrize("odata_type", [None, 7], ids=["null", "number"]) +def test_should_block_rejects_a_policy_action_whose_odata_type_is_not_text(odata_type: object) -> None: + response: Final = {"policyActions": [{"@odata.type": odata_type, "restrictionAction": "block"}]} + + with pytest.raises(TypeError, match=f"argument of type '{type(odata_type).__name__}'"): + PurviewGuardrailBase._should_block(response) diff --git a/tests/unit/proxy/guardrails/test_pillar_guardrails.py b/tests/unit/proxy/guardrails/test_pillar_guardrails.py index 86c1b1de8a1..1f167d34822 100644 --- a/tests/unit/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/unit/proxy/guardrails/test_pillar_guardrails.py @@ -1571,3 +1571,35 @@ async def test_mcp_call_masking( if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +class _RecordingPillarClient: + def __init__(self) -> None: + self.posted: list[dict[str, object]] = [] + + async def post(self, **kwargs: object) -> Response: + self.posted.append(kwargs) + return Response(200, json={"flagged": False}, request=Request("POST", "https://api.pillar.security")) + + +@pytest.mark.asyncio +async def test_messages_that_are_not_a_list_take_the_api_error_fallback_without_calling_pillar(env_setup): + guardrail = PillarGuardrail( + guardrail_name="pillar-fallback-block", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + fallback_on_error="block", + ) + pillar_client = _RecordingPillarClient() + guardrail.async_handler = pillar_client + + with pytest.raises(HTTPException) as excinfo: + await guardrail.async_pre_call_hook( + data={"model": "gpt-4", "messages": 5}, + cache=DualCache(), + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + ) + + assert excinfo.value.status_code == 503 + assert pillar_client.posted == [] diff --git a/tests/unit/proxy/hooks/test_key_management_event_hooks.py b/tests/unit/proxy/hooks/test_key_management_event_hooks.py index 8952951351c..4ba6fb95a9d 100644 --- a/tests/unit/proxy/hooks/test_key_management_event_hooks.py +++ b/tests/unit/proxy/hooks/test_key_management_event_hooks.py @@ -779,3 +779,34 @@ async def test_deleting_many_keys_writes_every_audit_row_in_one_insert(): assert {(row["table_name"], row["action"], row["changed_by"]) for row in rows} == { (LitellmTableNames.KEY_TABLE_NAME, "deleted", "admin") } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_metadata", "expected"), + [ + ({"secret_manager_settings": {"prefix": "team-a/"}}, {"prefix": "team-a/"}), + ({"secret_manager_settings": {}}, None), + ({"owner": "platform"}, None), + (None, None), + ], + ids=["team_settings", "empty_settings", "no_settings_key", "no_metadata"], +) +async def test_secret_manager_optional_params_come_from_the_team_metadata( + monkeypatch: pytest.MonkeyPatch, + team_metadata: dict[str, object] | None, + expected: dict[str, object] | None, +) -> None: + from litellm.models.team import LiteLLM_TeamTable + from litellm.proxy import proxy_server + + async def fake_get_team_object( + team_id: str, prisma_client: object, user_api_key_cache: object + ) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id=team_id, metadata=team_metadata) + + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", fake_get_team_object) + + assert await KeyManagementEventHooks._get_secret_manager_optional_params("team-a") == expected diff --git a/tests/unit/proxy/hooks/test_max_iterations_limiter.py b/tests/unit/proxy/hooks/test_max_iterations_limiter.py index 5eb499042b5..0fc6d1ce543 100644 --- a/tests/unit/proxy/hooks/test_max_iterations_limiter.py +++ b/tests/unit/proxy/hooks/test_max_iterations_limiter.py @@ -150,3 +150,42 @@ async def test_max_iterations_no_agent_id_passes(): call_type="", ) assert result is None + + +class _SharedCounterRedis: + def __init__(self) -> None: + self.counts: dict[str, int] = {} + self.script_args: list[list[int]] = [] + + def async_register_script(self, script: str): + async def increment(keys: list[str], args: list[int]) -> int: + self.script_args.append(args) + self.counts[keys[0]] = self.counts.get(keys[0], 0) + 1 + return self.counts[keys[0]] + + return increment + + +@pytest.mark.asyncio +async def test_max_iterations_counts_one_session_across_proxy_instances_through_redis(): + shared_redis = _SharedCounterRedis() + first, second = ( + PROXY_MaxIterationsHandler(internal_usage_cache=InternalUsageCache(DualCache(redis_cache=shared_redis))) + for _ in range(2) + ) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test-key-1234", metadata={"max_iterations": 3}) + data = {"metadata": {"session_id": "session-shared"}} + + for handler in (first, second, first): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type="" + ) + + with pytest.raises(HTTPException) as exc_info: + await second.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type="" + ) + + assert exc_info.value.status_code == 429 + assert "Current count: 4" in str(exc_info.value.detail) + assert shared_redis.script_args == [[first.ttl]] * 4 diff --git a/tests/unit/proxy/management_endpoints/test_user_agent_analytics_endpoints.py b/tests/unit/proxy/management_endpoints/test_user_agent_analytics_endpoints.py new file mode 100644 index 00000000000..a56c3095d03 --- /dev/null +++ b/tests/unit/proxy/management_endpoints/test_user_agent_analytics_endpoints.py @@ -0,0 +1,54 @@ +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock + +import pytest + +from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import ( + ActiveUsersAnalyticsResponse, + DistinctTagResponse, + DistinctTagsResponse, + TagActiveUsersResponse, + get_daily_active_users, + get_distinct_user_agent_tags, +) + + +def _prisma_returning(rows: list[dict[str, object]]) -> SimpleNamespace: + return SimpleNamespace(db=SimpleNamespace(query_raw=AsyncMock(return_value=rows))) + + +@pytest.mark.asyncio +async def test_distinct_user_agent_tags_are_read_from_the_raw_query_rows(monkeypatch: pytest.MonkeyPatch) -> None: + prisma: Final = _prisma_returning([{"tag": "User-Agent:curl", "usage_count": 7}, {"tag": "prod", "usage_count": 2}]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + response: Final = await get_distinct_user_agent_tags(user_api_key_dict=UserAPIKeyAuth()) + + assert response == DistinctTagsResponse( + results=[DistinctTagResponse(tag="User-Agent:curl"), DistinctTagResponse(tag="prod")] + ) + assert len(prisma.db.query_raw.await_args.args) == 1 + + +@pytest.mark.asyncio +async def test_daily_active_users_forwards_each_tag_filter_as_a_query_parameter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + prisma: Final = _prisma_returning([{"tag": "User-Agent:curl", "active_users": 3, "date": "2026-08-01"}]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + response: Final = await get_daily_active_users( + tag_filter=None, + tag_filters=["User-Agent:curl", "prod"], + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert response == ActiveUsersAnalyticsResponse( + results=[TagActiveUsersResponse(tag="User-Agent:curl", active_users=3, date="2026-08-01")] + ) + sql_query, *params = prisma.db.query_raw.await_args.args + assert params[2:] == ["User-Agent:curl", "prod"] + assert "dts.tag = $3 OR dts.tag = $4" in sql_query diff --git a/tests/unit/proxy/management_helpers/test_object_permission_utils.py b/tests/unit/proxy/management_helpers/test_object_permission_utils.py index 5814d8d2a97..12fb390e995 100644 --- a/tests/unit/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/unit/proxy/management_helpers/test_object_permission_utils.py @@ -1054,6 +1054,26 @@ async def test_resolve_team_allowed_mcp_servers_dict_tool_permissions( assert result == {"server-a"} +@pytest.mark.asyncio +@pytest.mark.parametrize("stored_tool_permissions", ["[]", "null", '"server-2"'], ids=["array", "null", "string"]) +@patch( + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_resolve_team_allowed_mcp_servers_rejects_stored_tool_permissions_that_are_not_an_object( + mock_access_groups, + stored_tool_permissions: str, +): + mock_perm = MagicMock(spec=LiteLLM_ObjectPermissionTable) + mock_perm.mcp_servers = ["server-1"] + mock_perm.mcp_access_groups = [] + mock_perm.mcp_tool_permissions = stored_tool_permissions + + with pytest.raises(AttributeError, match="object has no attribute 'keys'"): + await _resolve_team_allowed_mcp_servers(mock_perm) + + # ---- Tests for the all-proxy-mcpservers sentinel (team scoped to every server) ---- diff --git a/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py b/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py index c04323d7f7c..65e59fab0da 100644 --- a/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py +++ b/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py @@ -2238,3 +2238,62 @@ async def test_semantic_filter_hook_narrows_only_references_the_gateway_serves() ) assert narrowed == [{**gateway_reference, "allowed_tools": ["srv-tool_1"]}, external_tool] + + +@pytest.mark.asyncio +async def test_expand_mcp_tools_returns_every_expanded_tool_as_a_dict(): + from typing import Optional + + from pydantic import BaseModel, ConfigDict + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.proxy.mcp_server.semantic_tool_filter import SemanticMCPToolFilter + from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler + + class _FunctionTool(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str + name: str + description: Optional[str] = None + + class _LegacyFunctionTool: + def __init__(self, name: str) -> None: + self.name = name + + def dict(self, *, exclude_none: bool) -> dict[str, object]: + fields = {"type": "function", "name": self.name, "description": None} + return {key: value for key, value in fields.items() if value is not None or not exclude_none} + + plain_tool = {"type": "function", "name": "plain_tool", "parameters": {"type": "object"}} + unknown_tool = "opaque-tool" + pipeline_output = [ + _FunctionTool(type="function", name="v2_tool"), + _LegacyFunctionTool("legacy_tool"), + plain_tool, + unknown_tool, + ] + auth = UserAPIKeyAuth(api_key="hashed-key") + gateway_reference = {"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"} + + hook = SemanticToolFilterHook( + SemanticMCPToolFilter(embedding_model="text-embedding-3-small", litellm_router_instance=Mock()) + ) + with patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_to_openai_format", + new=AsyncMock(return_value=(pipeline_output, {})), + ) as process_tools: + expanded = await hook._expand_mcp_tools([gateway_reference], auth) + + assert expanded == [ + {"type": "function", "name": "v2_tool"}, + {"type": "function", "name": "legacy_tool"}, + plain_tool, + unknown_tool, + ] + assert process_tools.await_args.kwargs == { + "user_api_key_auth": auth, + "mcp_tools_with_litellm_proxy": [gateway_reference], + } diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index cd6c1ff11b5..45abb880f92 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -61,6 +61,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( WebsocketFrameModelRefused, WebsocketFrameModelVerdict, _build_passthrough_failure_request_payload, + _parse_request_data_by_content_type, _registered_pass_through_routes, _truncate_upstream_error_body, _update_metadata_with_tags_in_header, @@ -13242,3 +13243,106 @@ async def test_passthrough_header_and_callback_share_finalized_cost( finally: cleanup() await client.aclose() + + +def _stub_adapter_proxy(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + proxy_logging: Final = MagicMock() + proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + proxy_logging.post_call_failure_hook = AsyncMock() + + async def fake_add_litellm_data_to_request(**kwargs: object) -> object: + return kwargs["data"] + + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging) + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + return proxy_logging + + +@pytest.mark.asyncio +async def test_chat_completion_pass_through_endpoint_reads_a_body_only_json_can_parse( + monkeypatch: pytest.MonkeyPatch, +): + proxy_logging: Final = _stub_adapter_proxy(monkeypatch) + request: Final = MagicMock(spec=Request) + request.body = AsyncMock( + return_value=b'{"model": "unknown-model", "stream": false, "user": null, "messages": [{"role": "user", "content": "hi"}]}' + ) + + with pytest.raises(ProxyException): + await chat_completion_pass_through_endpoint( + fastapi_response=Response(), + request=request, + adapter_id="anthropic", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + forwarded: Final = proxy_logging.pre_call_hook.call_args.kwargs["data"] + assert (forwarded["model"], forwarded["stream"], forwarded["user"], forwarded["adapter_id"]) == ( + "unknown-model", + False, + None, + "anthropic", + ) + assert forwarded["messages"] == [{"role": "user", "content": "hi"}] + + +@pytest.mark.asyncio +async def test_chat_completion_pass_through_endpoint_rejects_a_body_that_is_not_an_object( + monkeypatch: pytest.MonkeyPatch, +): + proxy_logging: Final = _stub_adapter_proxy(monkeypatch) + request: Final = MagicMock(spec=Request) + request.body = AsyncMock(return_value=b'["not", "an", "object"]') + + with pytest.raises(ProxyException) as raised: + await chat_completion_pass_through_endpoint( + fastapi_response=Response(), + request=request, + adapter_id="anthropic", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert raised.value.message + proxy_logging.pre_call_hook.assert_not_awaited() + + +def _request_with_body(content_type: str, body: bytes) -> Request: + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + return Request( + { + "type": "http", + "method": "POST", + "path": "/pass-through", + "query_string": b"", + "headers": [(b"content-type", content_type.encode())], + }, + receive, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body", "expected"), + [ + (b'{"model": "m", "stream": true}', (None, {"model": "m", "stream": True}, None, True)), + ( + b'{"query_params": {"a": "1"}, "custom_body": {"b": 2}, "stream": false}', + ({"a": "1"}, {"b": 2}, None, False), + ), + (b"{}", (None, None, None, None)), + (b"[1, 2]", (None, None, None, None)), + (b'--x\r\nContent-Disposition: form-data; name="f"\r\n\r\nv\r\n--x--\r\n', (None, None, None, None)), + ], + ids=["json_body_is_the_custom_body", "explicit_envelope", "empty_object", "json_array", "real_multipart"], +) +async def test_parse_request_data_by_content_type_reads_json_sent_as_multipart( + body: bytes, expected: tuple[object, object, None, object] +): + request: Final = _request_with_body("multipart/form-data; boundary=x", body) + + assert await _parse_request_data_by_content_type(request) == expected diff --git a/tests/unit/proxy/test_caching_routes.py b/tests/unit/proxy/test_caching_routes.py index 8a59793495a..e139593976e 100644 --- a/tests/unit/proxy/test_caching_routes.py +++ b/tests/unit/proxy/test_caching_routes.py @@ -397,3 +397,20 @@ def test_extract_cache_params_is_empty_when_a_backend_attribute_name_is_not_a_st _use_redis_backed_cache(monkeypatch, {attribute_name: "redis.internal"}) assert _extract_cache_params() == {}, attribute_name + + +@pytest.fixture +def mock_redis_delete(mocker): + mock_cache = mocker.MagicMock() + mock_cache.type = "redis" + mock_cache.delete_cache_keys = mocker.AsyncMock() + mocker.patch.object(litellm, "cache", mock_cache) + return mock_cache + + +def test_cache_delete_removes_the_requested_keys(mock_redis_delete): + response = client.post("/cache/delete", headers={"Authorization": "Bearer sk-9876"}, json={"keys": ["k1", "k2"]}) + + assert response.status_code == 200 + assert response.json() == {"status": "success"} + mock_redis_delete.delete_cache_keys.assert_awaited_once_with(keys=["k1", "k2"]) diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index 05c4f9d8a67..3f3795dd33c 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -659,3 +659,14 @@ async def test_authoritative_combined_key_view_uses_writer_through_rotation( assert response.token == "current-token" reader.query_first.assert_not_awaited() assert writer.query_first.await_count == (2 if rotated else 1) + + +@pytest.mark.asyncio +async def test_get_data_all_users_pages_the_raw_user_query(prisma_client: PrismaClient) -> None: + rows: Final = [{"user_id": "u-1", "spend": 3.0, "key_aliases": ["alias-a"]}] + prisma_client.db.query_raw = AsyncMock(return_value=rows) + + result: Final = await prisma_client.get_data(table_name="user", query_type="find_all", limit=10, offset=20) + + assert result == rows + assert prisma_client.db.query_raw.await_args.args[1:] == (10, 20) diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 73d3414abb0..fa4a1656bc3 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -3450,3 +3450,131 @@ def test_proxy_vector_store_writes_send_only_the_callers_metadata_when_the_raw_b assert response.status_code == 200, response.text assert len(sent_bodies) == 1 assert sent_bodies[0].get("metadata") == CALLERS_METADATA + + +class _RecordingManagedVectorStores: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def acreate_vector_store(self, **kwargs: object) -> dict[str, object]: + self.calls.append(kwargs) + return {"id": "vs_managed", "object": "vector_store"} + + +def _vector_store_create_request(body: dict[str, object]) -> Request: + encoded = json.dumps(body).encode() + + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": encoded, "more_body": False} + + return Request( + { + "type": "http", + "method": "POST", + "path": "/v1/vector_stores", + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + }, + receive, + ) + + +def _install_managed_vector_stores( + monkeypatch: pytest.MonkeyPatch, managed: _RecordingManagedVectorStores +) -> MagicMock: + from litellm.proxy import proxy_server + + llm_router = MagicMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.get_proxy_hook.side_effect = lambda name: managed if name == "managed_vector_stores" else None + monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging_obj) + monkeypatch.setattr(proxy_server, "llm_router", llm_router) + return llm_router + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("target_model_names", "expected_targets"), + [("embed-a, embed-b", ["embed-a", "embed-b"]), (["embed-a"], ["embed-a"])], + ids=["comma_separated_string", "list"], +) +async def test_vector_store_create_with_target_model_names_creates_a_managed_store( + monkeypatch: pytest.MonkeyPatch, target_model_names: object, expected_targets: list[str] +) -> None: + from fastapi import Response + + from litellm.proxy.vector_store_endpoints.endpoints import vector_store_create + + managed = _RecordingManagedVectorStores() + llm_router = _install_managed_vector_stores(monkeypatch, managed) + + response = await vector_store_create( + request=_vector_store_create_request({"name": "docs", "target_model_names": target_model_names}), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert response == {"id": "vs_managed", "object": "vector_store"} + (call,) = managed.calls + assert call["target_model_names_list"] == expected_targets + assert call["llm_router"] is llm_router + assert call["create_request"]["name"] == "docs" + assert "target_model_names" not in call["create_request"] + + +@pytest.mark.asyncio +async def test_vector_store_create_rejects_target_model_names_that_are_not_a_string_or_list( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import Response + + from litellm.proxy.vector_store_endpoints.endpoints import vector_store_create + + managed = _RecordingManagedVectorStores() + _install_managed_vector_stores(monkeypatch, managed) + + with pytest.raises(HTTPException) as raised: + await vector_store_create( + request=_vector_store_create_request({"name": "docs", "target_model_names": 5}), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert raised.value.status_code == 400 + assert managed.calls == [] + + +@pytest.mark.asyncio +async def test_update_vector_store_persists_submitted_metadata_as_json() -> None: + from litellm.proxy.vector_store_endpoints.management_endpoints import update_vector_store + from litellm.types.vector_stores import VectorStoreUpdateRequest + + stored_row = MagicMock() + stored_row.model_dump = MagicMock(return_value={"vector_store_id": "vs_owned", "team_id": "team-A"}) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=stored_row) + mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock(return_value=stored_row) + + with ( + patch( + "litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.vector_store_endpoints.management_endpoints._check_vector_store_access", + new_callable=AsyncMock, + return_value=True, + ), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.vector_store_registry", None), + ): + await update_vector_store( + data=VectorStoreUpdateRequest( + vector_store_id="vs_owned", + vector_store_metadata={"owner": "search-team", "tier": 2}, + ), + user_api_key_dict=UserAPIKeyAuth(user_id="owner", team_id="team-A"), + ) + + written = mock_prisma_client.db.litellm_managedvectorstorestable.update.await_args.kwargs["data"] + assert json.loads(written["vector_store_metadata"]) == {"owner": "search-team", "tier": 2} diff --git a/tests/unit/responses/test_responses_utils.py b/tests/unit/responses/test_responses_utils.py index 0863abd1f1e..49ff7918849 100644 --- a/tests/unit/responses/test_responses_utils.py +++ b/tests/unit/responses/test_responses_utils.py @@ -415,8 +415,8 @@ class TestResponseAPILoggingUtils: assert result.prompt_tokens_details.cache_creation_tokens == 10059 assert result.prompt_tokens_details.cached_tokens == 0 - @pytest.mark.parametrize("as_model", [False, True]) - def test_transform_response_api_usage_prices_12h_cache_writes(self, as_model: bool) -> None: + @pytest.mark.parametrize("usage_shape", ["dict", "validated_model", "unvalidated_model"]) + def test_transform_response_api_usage_prices_12h_cache_writes(self, usage_shape: str) -> None: usage_data: Final = { "input_tokens": 5_000, "output_tokens": 0, @@ -426,7 +426,11 @@ class TestResponseAPILoggingUtils: "cache_write_tokens_details": {"ttl_30m_tokens": 3_000, "ttl_12h_tokens": 2_000}, }, } - usage_input: Final = ResponseAPIUsage(**usage_data) if as_model else usage_data + usage_input: Final = { + "dict": usage_data, + "validated_model": ResponseAPIUsage(**usage_data), + "unvalidated_model": ResponseAPIUsage.model_construct(**usage_data), + }[usage_shape] usage: Final = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_input) model_info: Final[ModelInfo] = { "input_cost_per_token": 1e-6, diff --git a/tests/unit/router_strategy/adaptive_router/test_hooks.py b/tests/unit/router_strategy/adaptive_router/test_hooks.py index 0ed22658144..d199f2e1922 100644 --- a/tests/unit/router_strategy/adaptive_router/test_hooks.py +++ b/tests/unit/router_strategy/adaptive_router/test_hooks.py @@ -1,5 +1,6 @@ """Unit tests for the AdaptiveRouterPostCallHook.""" +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest @@ -15,6 +16,7 @@ from litellm.router_strategy.adaptive_router.hooks import ( _resolve_session_key, ) from litellm.router_strategy.adaptive_router.signals import Turn +from litellm.types.utils import ModelResponse def _make_hook() -> AdaptiveRouterPostCallHook: @@ -236,6 +238,45 @@ async def test_hook_passes_tool_calls_through(): assert turn.tool_calls == [tc] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "expected_tool_calls"), + [ + ( + ModelResponse( + choices=[ + { + "message": { + "role": "assistant", + "content": "calling tool", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "search", "arguments": '{"q":"x"}'}, + } + ], + } + } + ] + ), + [{"id": "call_1", "type": "function", "function": {"name": "search", "arguments": '{"q":"x"}'}}], + ), + ( + _resp_with_content("calling tool", tool_calls=[SimpleNamespace(name="search")]), + [{"name": "search", "arguments": ""}], + ), + ], + ids=["model_response_tool_call", "tool_call_without_model_dump"], +) +async def test_hook_records_tool_call_objects_as_dicts(response: object, expected_tool_calls: list[dict[str, object]]): + hook = _make_hook() + await hook.async_log_success_event(_kwargs(chosen="fast"), response, 0.0, 1.0) + turn: Turn = hook.adaptive_router.record_turn.await_args.kwargs["turn"] + assert turn.assistant_content == "calling tool" + assert turn.tool_calls == expected_tool_calls + + # ---- _recent_tool_results ------------------------------------------------ diff --git a/tests/unit/router_strategy/adaptive_router/test_signals.py b/tests/unit/router_strategy/adaptive_router/test_signals.py index bd58b4c668d..f4b1298cd46 100644 --- a/tests/unit/router_strategy/adaptive_router/test_signals.py +++ b/tests/unit/router_strategy/adaptive_router/test_signals.py @@ -4,7 +4,7 @@ from typing import List, Tuple import pytest -from litellm.router_strategy.adaptive_router.config import TOOL_CALL_HISTORY_MAX +from litellm.router_strategy.adaptive_router.config import LOOP_REPEAT_THRESHOLD, TOOL_CALL_HISTORY_MAX from litellm.router_strategy.adaptive_router.signals import ( SessionState, SignalDelta, @@ -184,3 +184,33 @@ def test_apply_turn_is_o1_does_not_grow_history_unbounded(): Turn(tool_calls=[{"name": f"tool_{i}", "arguments": {}}]), ) assert len(state.tool_call_history) <= TOOL_CALL_HISTORY_MAX + + +def _function_tool_call(name: str, arguments: str) -> dict[str, object]: + return {"id": "call", "type": "function", "function": {"name": name, "arguments": arguments}} + + +@pytest.mark.parametrize( + ("names", "arguments", "expected_loops"), + [ + (("read_file",) * LOOP_REPEAT_THRESHOLD, ('{"path": "/x"}',) * LOOP_REPEAT_THRESHOLD, 1), + ( + tuple(f"tool_{i}" for i in range(LOOP_REPEAT_THRESHOLD)), + ('{"path": "/x"}',) * LOOP_REPEAT_THRESHOLD, + 0, + ), + ( + ("read_file",) * LOOP_REPEAT_THRESHOLD, + tuple(f'{{"path": "/{i}"}}' for i in range(LOOP_REPEAT_THRESHOLD)), + 0, + ), + ], + ids=["same_call", "different_names", "different_arguments"], +) +def test_loop_detection_reads_name_and_arguments_of_function_tool_calls( + names: Tuple[str, ...], arguments: Tuple[str, ...], expected_loops: int +): + state, _ = _replay( + [{"tool_calls": [_function_tool_call(name, args)]} for name, args in zip(names, arguments, strict=True)] + ) + assert state.loop_count == expected_loops diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 034b76507cf..2a5c3c81f18 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -4218,6 +4218,41 @@ def test_anthropic_text_disable_url_suffix_env_var( assert str(disabled_route.calls.last.request.url) == "https://api.example.com/custom/complete" +def test_huggingface_text_completion_sends_each_token_prompt_with_the_callers_params( + respx_mock: respx.MockRouter, +) -> None: + from litellm.rust_bridge.tokenizer import get_encoding + + encoding: Final = get_encoding("p50k_base") + respx_mock.get(url__startswith="https://huggingface.co/").respond(404) + endpoint: Final = respx_mock.post("https://hf.example.com/v1/chat/completions").respond( + json={ + "id": "chatcmpl-hf", + "object": "chat.completion", + "created": 1, + "model": "tgi", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}], + } + ) + + response: Final = litellm.text_completion( + model="huggingface/tgi", + prompt=[list(encoding.encode("Hello")), list(encoding.encode("world"))], + api_key="hf-test-key", + api_base="https://hf.example.com", + max_tokens=7, + ) + + assert [choice.text for choice in response.choices] == ["OK", "OK"] + sent: Final = sorted( + (json.loads(call.request.content) for call in endpoint.calls), key=lambda body: body["messages"][0]["content"] + ) + assert [(body["messages"], body["max_tokens"]) for body in sent] == [ + ([{"role": "user", "content": "Hello"}], 7), + ([{"role": "user", "content": "world"}], 7), + ] + + def test_image_edit_merges_headers_and_extra_headers(): from litellm.images.main import base_llm_http_handler diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index b0c55d3f0f2..5eaabdd2c59 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -11612,6 +11612,35 @@ async def test_success_enqueue_is_tracked_without_starting_callbacks_synchronous assert logging_obj.model_call_details["has_logged_async_success"] is True +def test_get_response_string_joins_the_text_parts_of_a_completed_responses_event() -> None: + from litellm.types.llms.openai import ResponseCompletedEvent + from litellm.utils import get_response_string + + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1700000000, + model="gpt-5.6", + object="response", + status="completed", + output=[ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [ + {"type": "output_text", "text": "Hello ", "annotations": []}, + {"type": "refusal", "refusal": "not this part"}, + {"type": "output_text", "text": "world", "annotations": []}, + ], + }, + {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "lookup", "arguments": "{}"}, + ], + ) + + assert get_response_string(ResponseCompletedEvent(type="response.completed", response=response)) == "Hello world" + + class _RecordingPreCallDeploymentHook(CustomLogger): def __init__(self) -> None: super().__init__()