diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 5c12649db0e..5b7e87407e0 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -1035,9 +1035,9 @@ jobs: run: | diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils else - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))" @@ -1294,7 +1294,7 @@ jobs: - name: Install dependencies if: steps.changes.outputs.decision != 'skip' run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard - name: Cache Prisma binaries if: steps.changes.outputs.decision != 'skip' @@ -1497,9 +1497,9 @@ jobs: RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }} run: | if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils else - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so fi diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 7db2de1f874..64fe03769b8 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.81" +version = "0.1.83" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.81" +version = "0.1.83" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index a12e1e27cfc..1c0c708ae7a 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.112" +version = "0.4.114" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -30,7 +30,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.112" +version = "0.4.114" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs index 6ccac4df659..0ded98f27c5 100644 --- a/litellm-rust/crates/python-bridge/src/router/host.rs +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -20,10 +20,10 @@ use std::{ use async_trait::async_trait; use litellm_router_core::{ - Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, Deployment, Error, Host, Output, - Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, RouterEvent, RouterRequest, - RouterResponse, RoutingContext, SecretRef, SelectStage, ThreadRandomness, - build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo, + Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Deployment, + Error, Host, Output, Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, + RouterEvent, RouterRequest, RouterResponse, RoutingContext, SecretRef, SelectStage, + ThreadRandomness, build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo, precall::PreCallFacts, }; use pyo3::{PyResult, Python, types::PyAnyMethods}; @@ -298,8 +298,9 @@ pub(crate) enum HostCallResponse { PreCall(PreCallFacts), /// The candidate models a `routing_plugins` select step reported. Models(Option>), - /// The ids a `callback_filters` select step left, or the error a filter raised. - Filtered(Result, AttemptError>), + /// The ids a `callback_filters` select step left with the one a filter prefers, or the error + /// a filter raised. + Filtered(Result<(Vec, Option), AttemptError>), /// What a pick's select or pre-route step, or a fallback check, raised on the host instead of /// answering: the walk takes it as the try's failure. Raised(AttemptError), @@ -352,8 +353,10 @@ impl HostCallResponse { got: self.kind(), }); } - if let (Self::Select(ids) | Self::Filtered(Ok(ids)), HostCall::Select { candidates, .. }) = - (self, call) + if let ( + Self::Select(ids) | Self::Filtered(Ok((ids, _))), + HostCall::Select { candidates, .. }, + ) = (self, call) && let Some(id) = ids.iter().find(|id| !candidates.contains(id)) { return Err(DriveError::UnknownCandidate { id: id.clone() }); @@ -890,7 +893,7 @@ impl Host for StepHost { &self, ctx: &RoutingContext<'_>, candidates: Vec>, - ) -> Result>, Error> { + ) -> Result { let (scope, group) = ctx.scope.wire(); let joined = |own: &[String], walked: &[String]| [own, walked].concat(); let call = HostCall::Select { @@ -909,8 +912,11 @@ impl Host for StepHost { let HostCallResponse::Filtered(result) = call_host(call).await else { unreachable!("{CHECKED}"); }; - let ids = result.map_err(Error::Attempt)?; - Ok(answered_by_id(&candidates, &ids)) + let (ids, preferred) = result.map_err(Error::Attempt)?; + Ok(CallbackFiltered { + kept: answered_by_id(&candidates, &ids), + preferred, + }) } async fn allow_fallback( diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs index 8157d0f5a32..6b58d63a714 100644 --- a/litellm-rust/crates/python-bridge/src/router/mod.rs +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -1181,7 +1181,8 @@ fn host_call_response( stage: Stage::CallbackFilters { .. }, .. } => { - // The remaining ids, or a failure mapping as an attempt reports one. + // The remaining ids, alone or paired with the id a filter prefers, or a failure + // mapping as an attempt reports one. let outcome = outcome.ok_or_else(missing)?; if let Ok(failure) = outcome.cast::() { return match attempt_outcome(failure)? { @@ -1191,7 +1192,10 @@ fn host_call_response( )), }; } - Ok(HostCallResponse::Filtered(Ok(outcome.extract()?))) + if let Ok(answer) = outcome.cast::() { + return Ok(HostCallResponse::Filtered(Ok(answer.extract()?))); + } + Ok(HostCallResponse::Filtered(Ok((outcome.extract()?, None)))) } HostCall::Select { stage: Stage::Strategy { records, .. }, diff --git a/litellm-rust/crates/router-core/src/host.rs b/litellm-rust/crates/router-core/src/host.rs index 0ef37873e6e..ae93288b0fa 100644 --- a/litellm-rust/crates/router-core/src/host.rs +++ b/litellm-rust/crates/router-core/src/host.rs @@ -405,6 +405,12 @@ pub enum SelectStage { OverrideCheck(RoutingStrategy), } +#[derive(Debug)] +pub struct CallbackFiltered { + pub kept: Vec>, + pub preferred: Option, +} + /// Routing events, delivered through [`Host::observe`]. They serialize as JSON objects tagged by /// `event` (the variant name), durations in seconds. #[derive(Clone, Debug, PartialEq, serde::Serialize)] @@ -714,14 +720,19 @@ pub trait Host: Send + Sync { } /// `async_callback_filter_deployments`: the host's routing callbacks narrow `candidates` in - /// turn. A filter that raises fails the try with [`Error::Attempt`], which the call retries and - /// falls back from like any attempt failure. A host without filters keeps every candidate. + /// turn, and may name one deployment they prefer, which the pick takes only when the filters + /// after the callbacks keep it. A filter that raises fails the try with [`Error::Attempt`], + /// which the call retries and falls back from like any attempt failure. A host without + /// filters keeps every candidate and prefers none. async fn callback_filters( &self, _ctx: &RoutingContext<'_>, candidates: Vec>, - ) -> Result>, Error> { - Ok(candidates) + ) -> Result { + Ok(CallbackFiltered { + kept: candidates, + preferred: None, + }) } /// The values of `refs`, one per ref, in order: what the deployment build substitutes for /// each `os.environ/` reference it found in entry `index` ([`crate::build`]). Asked once per diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs index 80c8f597157..f9b71c689c9 100644 --- a/litellm-rust/crates/router-core/src/lib.rs +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -124,8 +124,8 @@ pub use crate::{ error::{AttemptError, ConfigError, Error, FailureKind, Payload, Unsupported}, health::{Clock, Health, Outcome, SystemClock}, host::{ - Attempt, Blocked, BuildAsk, BuildPath, Built, Host, Output, Provided, RetryCause, - RouterEvent, RouterRequest, RoutingContext, SecretAt, SecretRef, SelectStage, + Attempt, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Host, Output, Provided, + RetryCause, RouterEvent, RouterRequest, RoutingContext, SecretAt, SecretRef, SelectStage, }, outcome::{OutcomeCallbacks, OutcomeWrite}, prerouting::{PreRoute, PreRouted}, diff --git a/litellm-rust/crates/router-core/src/select.rs b/litellm-rust/crates/router-core/src/select.rs index c7f3b1f9c31..75198c0a707 100644 --- a/litellm-rust/crates/router-core/src/select.rs +++ b/litellm-rust/crates/router-core/src/select.rs @@ -456,7 +456,7 @@ impl Core { let request_order = request.target_order.filter(|_| !inputs.fallback_hop); // `async_callback_filter_deployments` runs after the cooldown and blocked filters on every // async candidate list but a deployment named by id and a listed pass-through pick. - if settings.callback_filters + let preferred = if settings.callback_filters && !inputs.count_only && !resolution.kind.is_single() && request.pass_through != Some(PassThrough::Listed) @@ -472,8 +472,12 @@ impl Core { ctx.excluded_ids = &inputs.excluded_ids; ctx.skipped_ids = &inputs.skipped_ids; ctx.target_order = inputs.target_order.or(request_order); - healthy = self.host.callback_filters(&ctx, healthy).await?; - } + let filtered = self.host.callback_filters(&ctx, healthy).await?; + healthy = filtered.kept; + filtered.preferred + } else { + None + }; if inputs.count_only { return Ok(Candidates { name: name.to_owned(), @@ -648,8 +652,13 @@ impl Core { healthy.retain(|d| !d.cannot_run_batches()); } - if let Some(pinned) = request.pinned_deployment_id.as_deref() - && healthy.iter().any(|d| d.id == pinned) + if let Some(pinned) = [ + request.pinned_deployment_id.as_deref(), + preferred.as_deref(), + ] + .into_iter() + .flatten() + .find(|id| healthy.iter().any(|d| d.id == *id)) { healthy.retain(|d| d.id == pinned); } 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..2446b10d9db 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,12 +349,15 @@ _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 def _tool_search_replay_item(item: object) -> Mapping[str, object] | None: - raw: Final[object] = item.model_dump(exclude_unset=True) if isinstance(item, BaseModel) else item + raw: Final[object] = item.model_dump(exclude_unset=True, by_alias=True) if isinstance(item, BaseModel) else item if not isinstance(raw, dict): return None fields: Final = cast(dict[str, object], raw) # cast-ok: provider JSON item, keys forwarded verbatim @@ -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": @@ -2050,7 +2053,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): from pydantic import BaseModel if isinstance(str_line, BaseModel): - return self.chunk_parser(str_line.model_dump()) + return self.chunk_parser(str_line.model_dump(exclude_unset=True)) if not str_line or str_line.startswith("event:"): # ignore. 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 56ac9703112..1f2c1d4a7dd 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -13,10 +13,11 @@ import datetime import enum import sys from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol, SupportsIndex, SupportsInt, runtime_checkable import httpx import openai +from typing_extensions import Buffer import litellm from litellm.types.utils import LiteLLMCommonStrings @@ -1277,6 +1278,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, @@ -1292,8 +1306,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 @@ -1336,7 +1350,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 68876812c3e..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) @@ -291,6 +296,7 @@ def _guardrails_listed_in(metadata: object) -> list[str | dict[str, DynamicGuard class CustomGuardrail(CustomLogger): + mcp_block_behavior: Literal["tool_result", "raise"] | None = None # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path. use_native_during_call_hook: ClassVar[bool] = False 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/prompt_templates/tpm_media_utils.py b/litellm/litellm_core_utils/prompt_templates/tpm_media_utils.py index a4cb52d689c..a3e2d2e15d9 100644 --- a/litellm/litellm_core_utils/prompt_templates/tpm_media_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/tpm_media_utils.py @@ -202,7 +202,7 @@ def _tpm_embedding_leaf(value: object) -> _TPMMediaProjection: part: Final = as_str_mapping(value) if part is not None and part.get("type") == "file": return _tpm_embedding_file_block(part) - if not isinstance(value, str) or not value.startswith("data:"): + if not isinstance(value, str) or is_plain_embedding_text(value): return _tpm_media_leaf(value) if not _is_base64_data_uri(value): return _TPMMediaProjection(value) @@ -210,6 +210,10 @@ def _tpm_embedding_leaf(value: object) -> _TPMMediaProjection: return _TPMMediaProjection("" if media else value, media) +def is_plain_embedding_text(value: object) -> bool: + return isinstance(value, str) and not value.startswith("data:") + + def _is_json_array(value: object) -> bool: return isinstance(value, (list, tuple)) diff --git a/litellm/litellm_core_utils/realtime_guardrails.py b/litellm/litellm_core_utils/realtime_guardrails.py new file mode 100644 index 00000000000..382eecc5a8c --- /dev/null +++ b/litellm/litellm_core_utils/realtime_guardrails.py @@ -0,0 +1,316 @@ +import json +from collections.abc import AsyncIterator, Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_logger +from litellm.exceptions import ModifyResponseException + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + +_DEFAULT_VIOLATION_MESSAGE: Final = "I'm sorry, that request was blocked by the content filter." +BLOCKED_TOOL_OUTPUT: Final = json.dumps({"error": "Tool output blocked by content policy"}) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_JSON_ARRAY: Final = TypeAdapter(tuple[object, ...]) +_TEXTS: Final = TypeAdapter(tuple[str, ...]) +_TYPED_TEXT_PART_TYPES: Final = frozenset({"input_text", "output_text", "text"}) + + +@dataclass(frozen=True, slots=True) +class RealtimeGuardrailBlock: + message: str + ends_session_on_violation: bool + end_session_after_n_fails: int | None + + def ends_session(self, violation_count: int) -> bool: + return self.ends_session_on_violation or ( + self.end_session_after_n_fails is not None and violation_count >= self.end_session_after_n_fails + ) + + +@dataclass(frozen=True, slots=True) +class RealtimeGuardrailRewrite: + texts: tuple[str, ...] + + +RealtimeGuardrailVerdict: TypeAlias = RealtimeGuardrailBlock | RealtimeGuardrailRewrite | None + + +@dataclass(frozen=True, slots=True) +class RealtimeMessageText: + role: str + texts: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class RealtimeToolOutput: + texts: tuple[str, ...] + sanitized_event: str + + +RealtimeGuardedItem: TypeAlias = RealtimeMessageText | RealtimeToolOutput + + +def _json_object(value: object) -> Mapping[str, object] | None: + try: + return _JSON_OBJECT.validate_python(value) + except ValidationError: + return None + + +def _turn_copy(value: object) -> object: + mapping: Final = _json_object(value) + return value if mapping is None else dict(mapping) + + +def _typed_text(part: object) -> str | None: + part_object: Final = _json_object(part) + if part_object is None or part_object.get("type") not in _TYPED_TEXT_PART_TYPES: + return None + text: Final[object] = part_object.get("text", "") + return text if isinstance(text, str) else None + + +def _content_parts(content: object) -> Sequence[object]: + return _JSON_ARRAY.validate_python(content) if isinstance(content, list) else () + + +def _typed_texts(content: object) -> tuple[str, ...]: + return tuple(text for text in map(_typed_text, _content_parts(content)) if text is not None) + + +def _with_text(part: object, text: str | None) -> object: + part_object: Final = _json_object(part) + return part if text is None or part_object is None else {**part_object, "text": text} + + +def _tool_output(client_event: Mapping[str, object], item: Mapping[str, object]) -> RealtimeToolOutput | None: + output: Final[object] = item.get("output", "") + text: Final = output if isinstance(output, str) else json.dumps(output) + if not text: + return None + sanitized_event: Final = json.dumps({**client_event, "item": {**item, "output": BLOCKED_TOOL_OUTPUT}}) + return RealtimeToolOutput(texts=(text,), sanitized_event=sanitized_event) + + +def realtime_guarded_item(client_event: Mapping[str, object]) -> RealtimeGuardedItem | None: + if client_event.get("type") != "conversation.item.create": + return None + item: Final = _json_object(client_event.get("item")) + if item is None: + return None + if item.get("type") == "function_call_output": + return _tool_output(client_event, item) + texts: Final = _typed_texts(item.get("content")) + role: Final[object] = item.get("role") + return RealtimeMessageText(role=role if isinstance(role, str) else "user", texts=texts) if any(texts) else None + + +def rewritten_item(client_event: Mapping[str, object], texts: tuple[str, ...]) -> Mapping[str, object]: + item: Final = _json_object(client_event.get("item")) or {} + if item.get("type") == "function_call_output": + return {**item, "output": " ".join(texts)} + parts: Final = _content_parts(item.get("content")) + text_positions: Final = tuple(index for index, part in enumerate(parts) if _typed_text(part) is not None) + replacements: Final = dict(zip(text_positions, texts)) + merged_away: Final = frozenset(text_positions[len(texts) :]) + return { + **item, + "content": [ + _with_text(part, replacements.get(index)) for index, part in enumerate(parts) if index not in merged_away + ], + } + + +def _violation_text(error: Exception) -> str: + detail: Final[object] = getattr(error, "detail", None) + detail_object: Final = _json_object(detail) + if detail_object is not None: + reason: Final[object] = detail_object.get("error") + return reason if isinstance(reason, str) and reason else str(error) + if detail is not None: + return str(detail) + return str(error) or _DEFAULT_VIOLATION_MESSAGE + + +def _is_guardrail_block(error: Exception) -> bool: + return hasattr(error, "status_code") or isinstance(error, (ValueError, ModifyResponseException)) + + +def _screened_texts(screened: object, texts: tuple[str, ...]) -> tuple[str, ...]: + screened_object: Final = _json_object(screened) + try: + returned: Final = _TEXTS.validate_python(screened_object.get("texts") if screened_object else None) + except ValidationError: + return texts + return returned if len(returned) == len(texts) else texts + + +async def _guardrail_outcome( + guardrail: "CustomGuardrail", texts: tuple[str, ...], request_data: Mapping[str, object] +) -> RealtimeGuardrailBlock | tuple[str, ...]: + try: + screened: Final = await guardrail.apply_guardrail( + inputs={"texts": list(texts), "images": []}, + request_data=dict(request_data), + input_type="request", + ) + except Exception as error: + if not _is_guardrail_block(error): + verbose_logger.exception("[realtime guardrail] unexpected error in apply_guardrail: %s", error) + raise + return RealtimeGuardrailBlock( + message=guardrail.realtime_violation_message or _violation_text(error), + ends_session_on_violation=guardrail.on_violation == "end_session", + end_session_after_n_fails=guardrail.end_session_after_n_fails, + ) + return _screened_texts(screened, texts) + + +async def _screen( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] +) -> RealtimeGuardrailBlock | tuple[str, ...]: + if not guardrails: + return texts + outcome: Final = await _guardrail_outcome(guardrails[0], texts, request_data) + if isinstance(outcome, RealtimeGuardrailBlock): + return outcome + return await _screen(guardrails[1:], outcome, request_data) + + +async def _screened_until_block( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] +) -> AsyncIterator[RealtimeGuardrailBlock | str]: + for text in texts: + if isinstance(outcome := await _screen(guardrails, (text,), request_data), RealtimeGuardrailBlock): + yield outcome + return + for screened in outcome: + yield screened + + +async def _screen_each( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] +) -> RealtimeGuardrailBlock | tuple[str, ...]: + outcomes: Final = tuple([outcome async for outcome in _screened_until_block(guardrails, texts, request_data)]) + block: Final = next((outcome for outcome in outcomes if isinstance(outcome, RealtimeGuardrailBlock)), None) + return block if block is not None else tuple(outcome for outcome in outcomes if isinstance(outcome, str)) + + +def _joins(texts: tuple[str, ...]) -> tuple[str, ...]: + unspaced: Final = "".join(texts) + boundaries: Final = zip(texts, texts[1:]) + word_break_unknown: Final = any(not left[-1:].isspace() and not right[:1].isspace() for left, right in boundaries) + return (unspaced, " ".join(texts)) if word_break_unknown else (unspaced,) + + +async def _rewritten_parts( + guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object] +) -> RealtimeGuardrailBlock | tuple[str, ...]: + parts: Final = await _screen_each(guardrails, texts, request_data) + if isinstance(parts, RealtimeGuardrailBlock): + return parts + rejoined: Final = _joins(parts) + screened_rejoined: Final = await _screen_each(guardrails, rejoined, request_data) + if isinstance(screened_rejoined, RealtimeGuardrailBlock): + return screened_rejoined + masked_join: Final = next((text for text, join in zip(screened_rejoined, rejoined) if text != join), None) + return parts if masked_join is None else (masked_join,) + + +def _in_role_scope(guardrail: "CustomGuardrail", role: str) -> bool: + from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, + effective_skip_system_message_for_guardrail, + effective_skip_tool_message_for_guardrail, + role_out_of_guardrail_scope, + ) + + return not role_out_of_guardrail_scope( + role, + skip_system_message=effective_skip_system_message_for_guardrail(guardrail), + skip_tool_message=effective_skip_tool_message_for_guardrail(guardrail), + scan_only_tool_results=effective_scan_only_tool_results_for_guardrail(guardrail), + ) + + +def realtime_runner_guardrail( + callback: object, + check_data: Mapping[str, object], + event_hooks: Sequence["GuardrailEventHooks"], +) -> "CustomGuardrail | None": + from litellm.integrations.custom_guardrail import CustomGuardrail + + if not isinstance(callback, CustomGuardrail) or callback.use_native_lifecycle_hooks: + return None + runs: Final = any( + callback.should_run_guardrail(data=check_data, event_type=event_hook) for event_hook in event_hooks + ) + return callback if runs else None + + +def realtime_screening_guardrails( + role: str, + event_hooks: Sequence["GuardrailEventHooks"], + check_data: Mapping[str, object], + callbacks: Iterable[object], +) -> tuple["CustomGuardrail", ...]: + runner_guardrails: Final = (realtime_runner_guardrail(callback, check_data, event_hooks) for callback in callbacks) + guardrails: Final = { + id(guardrail): guardrail + for guardrail in runner_guardrails + if guardrail is not None and _in_role_scope(guardrail, role) + } + return tuple(guardrails.values()) + + +async def realtime_guardrail_verdict( + texts: tuple[str, ...], + role: str, + event_hooks: Sequence["GuardrailEventHooks"], + request_data: Mapping[str, object], + user_api_key_dict: object | None, + callbacks: Iterable[object], +) -> RealtimeGuardrailVerdict: + check_data: Final = {**request_data, "transcript": " ".join(texts), "stream": True} + in_scope: Final = realtime_screening_guardrails(role, event_hooks, check_data, callbacks) + guardrail_request_data: Final = { + **{key: _turn_copy(value) for key, value in request_data.items()}, + "user_api_key_dict": user_api_key_dict, + } + joined: Final = _joins(texts) + screened: Final = await _screen_each(in_scope, joined, guardrail_request_data) + if isinstance(screened, RealtimeGuardrailBlock): + return screened + if screened == joined: + return None + rewritten: Final = screened if len(texts) < 2 else await _rewritten_parts(in_scope, texts, guardrail_request_data) + if isinstance(rewritten, RealtimeGuardrailBlock): + return rewritten + return None if rewritten == texts else RealtimeGuardrailRewrite(texts=rewritten) + + +def guardrail_violation_error_event(message: str) -> str: + return json.dumps( + { + "type": "error", + "error": {"type": "guardrail_violation", "message": message, "code": "content_policy_violation"}, + } + ) + + +def guardrail_voice_prompt_events(message: str) -> tuple[str, str]: + prompt: Final = f"Say exactly the following message to the user, word for word, do not add anything else: {message}" + return ( + json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": prompt}]}, + } + ), + json.dumps({"type": "response.create"}), + ) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index b6af181980e..01aa7296b0f 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -47,13 +47,27 @@ from litellm.types.realtime import ALL_DELTA_TYPES from .litellm_logging import Logging as LiteLLMLogging from .realtime_errors import client_close_code, realtime_error_event, websocket_close_reason +from .realtime_guardrails import ( + RealtimeGuardedItem, + RealtimeGuardrailBlock, + RealtimeGuardrailRewrite, + RealtimeGuardrailVerdict, + RealtimeMessageText, + RealtimeToolOutput, + guardrail_violation_error_event, + guardrail_voice_prompt_events, + realtime_guarded_item, + realtime_guardrail_verdict, + realtime_runner_guardrail, + realtime_screening_guardrails, + rewritten_item, +) if TYPE_CHECKING: from openai.types.realtime.realtime_audio_formats import RealtimeAudioFormats from websockets.asyncio.client import ClientConnection from websockets.exceptions import ConnectionClosed - from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.types.guardrails import GuardrailEventHooks CLIENT_CONNECTION_CLASS = ClientConnection @@ -437,21 +451,6 @@ def _opted_out_of_always_on_guardrail(callback: object, request_data: Mapping[st ) -def _realtime_runner_guardrail( - callback: object, - check_data: Mapping[str, object], - event_hooks: Sequence["GuardrailEventHooks"], -) -> "CustomGuardrail | None": - from litellm.integrations.custom_guardrail import CustomGuardrail - - if not isinstance(callback, CustomGuardrail) or callback.use_native_lifecycle_hooks: - return None - runs: Final = any( - callback.should_run_guardrail(data=check_data, event_type=event_hook) for event_hook in event_hooks - ) - return callback if runs else None - - def translation_guardrails_apply(request_data: Mapping[str, object]) -> bool: from litellm.types.guardrails import GuardrailEventHooks @@ -479,7 +478,7 @@ def realtime_guardrails_apply(request_data: Mapping[str, object], callbacks: Seq check_data: Final = {**request_data, "stream": True} runner_event_hooks: Final = (GuardrailEventHooks.realtime_input_transcription, GuardrailEventHooks.pre_call) return any( - _realtime_runner_guardrail(callback, check_data, runner_event_hooks) is not None for callback in callbacks + realtime_runner_guardrail(callback, check_data, runner_event_hooks) is not None for callback in callbacks ) @@ -1663,16 +1662,60 @@ class RealTimeStreaming: async def run_realtime_guardrails( self, - transcript: str, + transcript: object, item_id: str | None = None, pre_block_backend_message: str | None = None, event_hooks: Sequence["GuardrailEventHooks"] | None = None, ) -> bool: """ - Run registered guardrails on realtime text (transcript, user message, tool output). - Returns True if blocked (synthetic warning already sent to client). Returns False if clean (caller should send response.create to the backend). + """ + from litellm.types.guardrails import GuardrailEventHooks + + if isinstance(transcript, str): + verdict: Final = await self.apply_realtime_guardrails( + (transcript,), "user", item_id, pre_block_backend_message, event_hooks + ) + return isinstance(verdict, RealtimeGuardrailBlock) + if realtime_screening_guardrails( + "user", + event_hooks or (GuardrailEventHooks.realtime_input_transcription,), + {**self.request_data, "stream": True}, + litellm.callbacks, + ): + raise TypeError(f"Realtime guardrails cannot screen a {type(transcript).__name__} transcript") + return False + + async def _screen_client_item(self, guarded_item: RealtimeGuardedItem | None) -> RealtimeGuardrailVerdict: + from litellm.types.guardrails import GuardrailEventHooks + + match guarded_item: + case RealtimeToolOutput(texts=texts, sanitized_event=sanitized_event): + return await self.apply_realtime_guardrails( + texts, + "tool", + pre_block_backend_message=sanitized_event, + event_hooks=[GuardrailEventHooks.pre_call], + ) + case RealtimeMessageText(role=role, texts=texts): + return await self.apply_realtime_guardrails(texts, role, event_hooks=[GuardrailEventHooks.pre_call]) + case None: + return None + + async def apply_realtime_guardrails( + self, + texts: tuple[str, ...], + role: str, + item_id: str | None = None, + pre_block_backend_message: str | None = None, + event_hooks: Sequence["GuardrailEventHooks"] | None = None, + ) -> RealtimeGuardrailVerdict: + """ + Run registered guardrails on realtime text (transcript, user message, tool output). + + A block has already been reported to the client and the backend when this returns. + A rewrite carries the masked texts the caller forwards in place of the originals. ``pre_block_backend_message`` (if provided) is sent to the backend BEFORE any of the guardrail's own backend messages when a block is @@ -1687,102 +1730,33 @@ class RealTimeStreaming: """ from litellm.types.guardrails import GuardrailEventHooks - if event_hooks is None: - event_hooks = [GuardrailEventHooks.realtime_input_transcription] - _realtime_event_types: Final = event_hooks - _check_data: Final = {**self.request_data, "transcript": transcript, "stream": True} - _already_run: Final[set] = set() - - runner_guardrails: Final = ( - _realtime_runner_guardrail(callback, _check_data, _realtime_event_types) for callback in litellm.callbacks + block: Final = await realtime_guardrail_verdict( + texts, + role, + event_hooks if event_hooks is not None else (GuardrailEventHooks.realtime_input_transcription,), + self.request_data, + self.user_api_key_dict, + litellm.callbacks, ) - for callback in runner_guardrails: - if callback is None or id(callback) in _already_run: - continue - _already_run.add(id(callback)) - try: - await callback.apply_guardrail( - inputs={"texts": [transcript], "images": []}, - request_data={"user_api_key_dict": self.user_api_key_dict}, - input_type="request", - ) - except Exception as e: - # Re-raise unexpected errors (no status_code/detail = programming bug, not a block). - # HTTPException and guardrail-raised exceptions have a status_code or detail attr. - is_guardrail_block = hasattr(e, "status_code") or isinstance(e, ValueError) - if not is_guardrail_block: - verbose_logger.exception( - "[realtime guardrail] unexpected error in apply_guardrail: %s", - e, - ) - raise - # Extract the human-readable error from the detail dict (HTTPException) - # or fall back to str(e) for plain ValueError. - detail: object | None = getattr(e, "detail", None) - if isinstance(detail, dict): - safe_msg = detail.get("error") or str(e) - elif detail is not None: - safe_msg = str(detail) - else: - safe_msg = str(e) or "I'm sorry, that request was blocked by the content filter." + if not isinstance(block, RealtimeGuardrailBlock): + return block + await self._send_blocked_turn_to_backend(pre_block_backend_message, item_id) + await self.websocket.send_text(guardrail_violation_error_event(block.message)) + if not self._is_transcription_session: + for voice_prompt_event in guardrail_voice_prompt_events(block.message): + await self._send_to_backend(voice_prompt_event) - # Use realtime_violation_message if configured; fall back to guardrail error text. - error_msg = getattr(callback, "realtime_violation_message", None) or safe_msg + self._violation_count += 1 + if block.ends_session(self._violation_count): + verbose_logger.warning("[realtime guardrail] ending session after violation %d", self._violation_count) + await self.backend_ws.close() - await self._send_blocked_turn_to_backend(pre_block_backend_message, item_id) - # Send the policy violation hint (shows as small gray status text in UI). - await self.websocket.send_text( - json.dumps( - { - "type": "error", - "error": { - "type": "guardrail_violation", - "message": error_msg, - "code": "content_policy_violation", - }, - } - ) - ) - if not self._is_transcription_session: - # Ask the LLM to voice the exact guardrail message so the - # user hears it as audio in voice sessions (not just text). - guardrail_prompt = ( - f"Say exactly the following message to the user, word for word, " - f"do not add anything else: {error_msg}" - ) - await self._send_to_backend( - json.dumps( - { - "type": "conversation.item.create", - "item": { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": guardrail_prompt}], - }, - } - ) - ) - await self._send_to_backend(json.dumps({"type": "response.create"})) - - self._violation_count += 1 - end_session_after: int | None = getattr(callback, "end_session_after_n_fails", None) - should_end = getattr(callback, "on_violation", None) == "end_session" or ( - end_session_after is not None and self._violation_count >= end_session_after - ) - if should_end: - verbose_logger.warning( - "[realtime guardrail] ending session after violation %d", - self._violation_count, - ) - await self.backend_ws.close() - - verbose_logger.warning( - "[realtime guardrail] BLOCKED transcript (violation %d): %r", - self._violation_count, - transcript[:80], - ) - return True - return False + verbose_logger.warning( + "[realtime guardrail] BLOCKED transcript (violation %d): %r", + self._violation_count, + " ".join(texts)[:80], + ) + return block async def _handle_provider_config_message(self, raw_response: str) -> None: """Process a backend message when a provider_config is set (transformed path).""" @@ -1866,7 +1840,7 @@ class RealTimeStreaming: self.store_message(event_str) self._capture_transcription_usage(event) await self._send_event_to_client(event, event_str) - blocked = await self.run_realtime_guardrails(cast(str, transcript)) + blocked = await self.run_realtime_guardrails(transcript) if not blocked and self._should_disable_vad_auto_response(): await self._send_to_backend(json.dumps({"type": "response.create"})) continue @@ -2267,92 +2241,22 @@ class RealTimeStreaming: guardrail_turn_detection_injected = False msg_type: str | None = None try: - from litellm.types.guardrails import GuardrailEventHooks - msg_obj = json.loads(message) client_event = msg_obj msg_type = client_event.get("type") - if msg_type == "conversation.item.create": - # Check user text messages for prompt injection - item = client_event.get("item", {}) - # Check function_call_output first so a client cannot - # bypass the tool-result guardrail by also setting - # role="user" on a function_call_output item. - if item.get("type") == "function_call_output": - # Tool results are client-controlled and fed to the - # model; check them with the same guardrail used for - # user text so an attacker cannot smuggle blocked - # content into a function_call_output. - output: object = item.get("output", "") - output_text = output if isinstance(output, str) else json.dumps(output) - if output_text: - # Build the sanitized function_call_output up - # front so we can hand it to the guardrail - # runner as the pre-block message. Providers - # that pair every toolCall with a toolResponse - # (e.g. Gemini/Vertex Live) require the - # toolResponse to arrive BEFORE any other - # client message — otherwise the guardrail's - # own clientContent would violate the - # pending-tool-call protocol contract and the - # backend could close the connection before - # the sanitized response ever lands. Dropping - # the blocked item outright would similarly - # leave such providers waiting indefinitely. - # The sanitized payload carries no blocked - # content — only a generic policy marker. - sanitized_msg = json.dumps( - { - **msg_obj, - "item": { - **item, - "output": json.dumps( - { - "error": "Tool output blocked by content policy", - } - ), - }, - } - ) - blocked = await self.run_realtime_guardrails( - output_text, - pre_block_backend_message=sanitized_msg, - event_hooks=[GuardrailEventHooks.pre_call], - ) - if blocked: - # ``_pending_guardrail_message`` is - # intentionally NOT set here. That flag - # exists to swallow the reflexive - # ``response.create`` an OpenAI client - # sends immediately after a user text - # message. In a tool-calling flow the - # client may not send a ``response.create`` - # at all (e.g. Gemini SDKs auto-respond), - # so leaving the flag set would - # incorrectly drop an unrelated - # ``response.create`` from a later - # interaction turn. - continue - elif item.get("role") == "user": - content_list: Sequence[object] = item.get("content", []) - texts = [ - c.get("text", "") - for c in content_list - if isinstance(c, dict) and c.get("type") == "input_text" - ] - combined_text = " ".join(texts) - if combined_text: - blocked = await self.run_realtime_guardrails( - combined_text, - event_hooks=[GuardrailEventHooks.pre_call], - ) - if blocked: - # Store the guardrail reason so the next response.create - # (sent automatically by the client) is rewritten to - # include it as response instructions. - self._pending_guardrail_message = combined_text - continue # don't forward the original blocked message + self._pending_guardrail_message = None + + match await self._screen_client_item(guarded_item := realtime_guarded_item(client_event)): + case RealtimeGuardrailBlock(): + if isinstance(guarded_item, RealtimeMessageText): + self._pending_guardrail_message = " ".join(guarded_item.texts) + continue + case RealtimeGuardrailRewrite(texts=masked_texts): + msg_obj["item"] = rewritten_item(client_event, masked_texts) + message = json.dumps(msg_obj) + case None: + pass if msg_type == "response.create" and self._pending_guardrail_message: # The guardrail already sent the synthetic AI bubble — drop this 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/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 8c402fe9b44..f6bad7f611b 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -117,7 +117,9 @@ class SensitiveDataMasker: elif key_is_sensitive and isinstance(item, str): masked_items.append(self.mask_value(item)) else: - masked_items.append(item if isinstance(item, (int, float, bool, str, list)) else str(item)) + masked_items.append( + item if item is None or isinstance(item, (int, float, bool, str, list)) else str(item) + ) return masked_items def mask_dict( @@ -144,7 +146,7 @@ class SensitiveDataMasker: str_value = str(v) if v is not None else "" masked_data[k] = self.mask_value(str_value) else: - masked_data[k] = v if isinstance(v, (int, float, bool, str, list)) else str(v) + masked_data[k] = v if v is None or isinstance(v, (int, float, bool, str, list)) else str(v) except Exception: masked_data[k] = "" diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 2ba95db16ef..0145d48ce9e 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 @@ -18,7 +18,7 @@ from openai.types import Completion 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 @@ -28,7 +28,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, ) @@ -153,6 +155,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: @@ -364,6 +373,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 @@ -450,9 +462,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 @@ -485,7 +496,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] = { @@ -973,7 +984,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": "", @@ -1107,7 +1118,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: @@ -1117,7 +1128,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 ( @@ -1876,7 +1887,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() @@ -1914,7 +1927,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 @@ -2867,8 +2880,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/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 8f1a08d2964..ca9529ba246 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -795,18 +795,21 @@ class AnthropicMessagesHandler(BaseTranslation): data: dict, # mutable-ok: the request body the adapter reads *, keep_cache_control: bool = False, + keep_tool_errors: bool = False, ) -> ChatCompletionRequest: """Translate Anthropic request to OpenAI chat completion format.""" ( chat_completion_compatible_request, _tool_name_mapping, - ) = LiteLLMAnthropicMessagesAdapter(keep_cache_control=keep_cache_control).translate_anthropic_to_openai( + ) = LiteLLMAnthropicMessagesAdapter( + keep_cache_control=keep_cache_control, keep_tool_errors=keep_tool_errors + ).translate_anthropic_to_openai( anthropic_message_request=cast(AnthropicMessagesRequest, scannable_tool_result_request(data)), preserve_midturn_system=True, ) return chat_completion_compatible_request - def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: + def get_structured_messages(self, data: dict, *, keep_tool_errors: bool = False) -> list[AllMessageValues] | None: """ Convert Anthropic messages request data to OpenAI-spec structured messages. @@ -815,7 +818,7 @@ class AnthropicMessagesHandler(BaseTranslation): messages: Final = data.get("messages") if messages is None: return None - chat_completion_compatible_request: Final = self._translate_to_openai(data) + chat_completion_compatible_request: Final = self._translate_to_openai(data, keep_tool_errors=keep_tool_errors) result: Final = cast( list[AllMessageValues], chat_completion_compatible_request.get("messages", []), 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/adapters/transformation.py b/litellm/llms/anthropic/messages/adapters/transformation.py index 43eaa37aba7..5e4fba6d706 100644 --- a/litellm/llms/anthropic/messages/adapters/transformation.py +++ b/litellm/llms/anthropic/messages/adapters/transformation.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast from pydantic import JsonValue, TypeAdapter import litellm -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.llms.anthropic.messages.convertible_blocks import tool_result_text_part from litellm.llms.anthropic.messages.utils import ( @@ -434,8 +434,9 @@ _CLAUDE_OVER_REASONING_EFFORT_PROVIDERS: Final = frozenset({LlmProviders.KYMA.va class LiteLLMAnthropicMessagesAdapter: - def __init__(self, *, keep_cache_control: bool = False) -> None: + def __init__(self, *, keep_cache_control: bool = False, keep_tool_errors: bool = False) -> None: self.keep_cache_control: Final = keep_cache_control + self.keep_tool_errors: Final = keep_tool_errors ### FOR [BETA] `/v1/messages` endpoint support @@ -615,11 +616,7 @@ class LiteLLMAnthropicMessagesAdapter: self._add_cache_control_if_applicable(content, doc_obj, model) new_user_content_list.append(doc_obj) elif content.get("type") == "tool_result": - tool_result = ChatCompletionToolMessage( - role="tool", - tool_call_id=content.get("tool_use_id", ""), - content=self._tool_result_content(content.get("content")), - ) + tool_result = self._tool_message(content) self._add_cache_control_if_applicable(content, tool_result, model) tool_message_list.append(tool_result) @@ -1432,6 +1429,15 @@ class LiteLLMAnthropicMessagesAdapter: return None return anthropic_image_source_to_openai_url(image_source) + def _tool_message(self, block: object) -> ChatCompletionToolMessage: + fields: Final = as_str_mapping(block) or {} + is_error: Final = fields.get("is_error") + tool_call_id: Final = tool_use_id if isinstance(tool_use_id := fields.get("tool_use_id"), str) else "" + content: Final = self._tool_result_content(fields.get("content")) + if self.keep_tool_errors and isinstance(is_error, bool): + return ChatCompletionToolMessage(role="tool", tool_call_id=tool_call_id, content=content, is_error=is_error) + return ChatCompletionToolMessage(role="tool", tool_call_id=tool_call_id, content=content) + def _tool_result_content(self, raw_content: object) -> ToolResultContent: if isinstance(raw_content, str): return raw_content 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 f32b3fc07ae..33687d08af5 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, Mapping -from typing import Final +from typing import Final, Protocol import httpx from openai import ( @@ -17,6 +17,7 @@ from pydantic import BaseModel, ConfigDict 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 ( @@ -61,6 +62,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 @@ -864,11 +870,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 6399ded9555..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 {} @@ -257,11 +267,12 @@ class BaseTranslation(ABC): """ return None - def get_structured_messages(self, data: dict) -> list["AllMessageValues"] | None: + def get_structured_messages(self, data: dict, *, keep_tool_errors: bool = False) -> list["AllMessageValues"] | None: """ Convert request data to OpenAI-spec structured messages. - Override in subclasses for format-specific conversion. + Override in subclasses for format-specific conversion. ``keep_tool_errors`` copies a provider's + tool-failure marker onto tool messages as ``is_error``; only for copies that are never sent to a model. Returns None if no convertible content is found. """ 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/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 49939399ec8..9f4fb8cfecc 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -6,12 +6,14 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic. import asyncio import contextlib +import functools import importlib.metadata import json -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence from dataclasses import dataclass, field +from enum import Enum, auto from types import MappingProxyType -from typing import Final, NoReturn, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Final, NoReturn, Protocol, runtime_checkable from pydantic import JsonValue, TypeAdapter @@ -24,11 +26,25 @@ from litellm.constants import ( BEDROCK_REALTIME_SDK_DISTRIBUTION, BEDROCK_REALTIME_SDK_SUPPORTED_RANGE, BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY, + REALTIME_SESSION_ENDED_ABNORMALLY_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.realtime_errors import realtime_error_event +from litellm.litellm_core_utils.realtime_guardrails import ( + RealtimeGuardedItem, + RealtimeGuardrailBlock, + RealtimeGuardrailRewrite, + RealtimeGuardrailVerdict, + RealtimeToolOutput, + guardrail_violation_error_event, + guardrail_voice_prompt_events, + realtime_guarded_item, + realtime_guardrail_verdict, + rewritten_item, +) from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeEventTypes from litellm.llms.base_llm.realtime.transformation import RealtimeSessionUpdateRejection from litellm.types.llms.bedrock import AwsAuthParams @@ -40,9 +56,13 @@ from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig, requested_modalities +if TYPE_CHECKING: + from litellm.types.guardrails import GuardrailEventHooks + _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _REQUEST_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) _SPEND_LOG_METADATA_KEYS: Final = ("litellm_metadata", "metadata") +GUARDRAIL_UNAVAILABLE_MESSAGE: Final = "A guardrail could not check this request, so the session was ended" _EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) @@ -259,11 +279,11 @@ def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseExc raise error from failure -def _parse_client_message(message: str) -> Mapping[str, JsonValue]: +def _parse_client_message(message: str) -> Mapping[str, JsonValue] | None: try: return _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message)) except ValueError: - return _EMPTY_JSON_OBJECT + return None async def _send_opening_message( @@ -282,6 +302,27 @@ async def _send_opening_message( client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt +async def _forward_client_message( + client_ws: RealtimeClientWebSocket, + bedrock_stream: BedrockBidirectionalStream, + transformation_config: BedrockRealtimeConfig, + model: str, + session_state: RealtimeResponseTransformInput, + send_to_bedrock: Callable[[str], Awaitable[None]], + client_message: str, +) -> None: + transformed_messages: Final = transformation_config.transform_realtime_request( + message=client_message, + model=model, + session_configuration_request=session_state.get("session_configuration_request"), + ) + if transformed_messages and not _is_session_committed(client_ws.scope): + await _send_opening_message(client_ws, bedrock_stream, client_message, transformed_messages, send_to_bedrock) + return + for bedrock_message in transformed_messages: + await send_to_bedrock(bedrock_message) + + async def _ack_session_update( client_ws: RealtimeClientWebSocket, transformation_config: BedrockRealtimeConfig, @@ -317,6 +358,161 @@ async def _reject_session_update( return True +async def _send_guardrail_violation(client_ws: RealtimeClientWebSocket, block: RealtimeGuardrailBlock) -> bool: + try: + await client_ws.send_text(guardrail_violation_error_event(block.message)) + except Exception as e: # noqa: BLE001 # any send failure means the client is gone + verbose_proxy_logger.debug("Bedrock Realtime: guardrail error not delivered: %s", e, exc_info=True) + return False + return True + + +async def _end_guarded_session(client_ws: RealtimeClientWebSocket) -> None: + verbose_proxy_logger.warning("Bedrock Realtime: ending session after a guardrail violation") + with contextlib.suppress(Exception): + await client_ws.close() + + +async def _end_unscreened_session(client_ws: RealtimeClientWebSocket) -> None: + with contextlib.suppress(Exception): + await client_ws.send_text(realtime_error_event(GUARDRAIL_UNAVAILABLE_MESSAGE, error_type="guardrail_error")) + with contextlib.suppress(Exception): + await client_ws.close() + + +@dataclass(frozen=True, slots=True) +class _GuardrailUnavailable: + pass + + +def _registered_callbacks() -> Iterable[object]: + return litellm.callbacks + + +class _ClientEventRoute(Enum): + FORWARD = auto() + DROP = auto() + END_SESSION = auto() + + +class BedrockRealtimeGuardrails: + """One Bedrock realtime session's guardrail checks, shared by both bridge directions. + + Nova Sonic has no way to cancel a response or remove a turn, so a blocked user transcript ends + the session before that turn's answer reaches the client. + """ + + def __init__( + self, + request_data: Mapping[str, object] = _EMPTY_JSON_OBJECT, + user_api_key_dict: object | None = None, + callbacks: Callable[[], Iterable[object]] = _registered_callbacks, + ) -> None: + self._request_data: Final = request_data + self._user_api_key_dict: Final = user_api_key_dict + self._callbacks: Final = callbacks + self._violation_count = 0 + self._drop_next_response_create = False + self._failed = False + + @property + def failed(self) -> bool: + return self._failed + + async def _verdict( + self, texts: tuple[str, ...], role: str, hook: "GuardrailEventHooks" + ) -> RealtimeGuardrailVerdict | _GuardrailUnavailable: + try: + return await realtime_guardrail_verdict( + texts, role, (hook,), self._request_data, self._user_api_key_dict, self._callbacks() + ) + except Exception as e: # noqa: BLE001 # a guardrail that cannot decide fails closed, never as a Bedrock failure + verbose_proxy_logger.warning("Bedrock Realtime: a guardrail failed, ending the session: %s", e) + self._failed = True + return _GuardrailUnavailable() + + def _ends_session(self, block: RealtimeGuardrailBlock) -> bool: + self._violation_count += 1 + return block.ends_session(self._violation_count) + + async def screen_client_event( + self, + client_event: Mapping[str, JsonValue], + client_ws: RealtimeClientWebSocket, + forward_to_bedrock: Callable[[str], Awaitable[None]], + ) -> _ClientEventRoute: + from litellm.types.guardrails import GuardrailEventHooks + + event_type: Final = client_event.get("type") + if event_type == "response.create" and self._drop_next_response_create: + self._drop_next_response_create = False + return _ClientEventRoute.DROP + if event_type == "conversation.item.create": + self._drop_next_response_create = False + guarded_item: Final = realtime_guarded_item(client_event) + if guarded_item is None: + return _ClientEventRoute.FORWARD + held: Final = json.dumps(client_event) + client_ws.scope[BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY] = held # rebind-ok: replayed on fallback + role: Final = "tool" if isinstance(guarded_item, RealtimeToolOutput) else "user" + match await self._verdict(guarded_item.texts, role, GuardrailEventHooks.pre_call): + case None: + return _ClientEventRoute.FORWARD + case RealtimeGuardrailRewrite(texts=texts): + await forward_to_bedrock(json.dumps({**client_event, "item": rewritten_item(client_event, texts)})) + return _ClientEventRoute.DROP + case RealtimeGuardrailBlock() as block: + return await self._route_blocked_item(block, guarded_item, client_ws, forward_to_bedrock) + case _GuardrailUnavailable(): + await _end_unscreened_session(client_ws) + return _ClientEventRoute.END_SESSION + + async def _route_blocked_item( + self, + block: RealtimeGuardrailBlock, + guarded_item: RealtimeGuardedItem, + client_ws: RealtimeClientWebSocket, + forward_to_bedrock: Callable[[str], Awaitable[None]], + ) -> _ClientEventRoute: + if isinstance(guarded_item, RealtimeToolOutput): + await forward_to_bedrock(guarded_item.sanitized_event) + if not await _send_guardrail_violation(client_ws, block): + return _ClientEventRoute.END_SESSION + if self._ends_session(block): + await _end_guarded_session(client_ws) + return _ClientEventRoute.END_SESSION + if isinstance(guarded_item, RealtimeToolOutput): + return _ClientEventRoute.DROP + for voice_prompt_event in guardrail_voice_prompt_events(f'"{block.message}"'): + await forward_to_bedrock(voice_prompt_event) + self._drop_next_response_create = True + return _ClientEventRoute.DROP + + async def continues_after( + self, provider_event: Mapping[str, object], client_ws: RealtimeClientWebSocket, client_audio_streamed: bool + ) -> bool: + from litellm.types.guardrails import GuardrailEventHooks + + transcript: Final = provider_event.get("transcript") + if ( + not client_audio_streamed + or provider_event.get("type") != "conversation.item.input_audio_transcription.completed" + or not isinstance(transcript, str) + or not transcript + ): + return True + match await self._verdict((transcript,), "user", GuardrailEventHooks.realtime_input_transcription): + case RealtimeGuardrailBlock() as block: + if await _send_guardrail_violation(client_ws, block): + await _end_guarded_session(client_ws) + return False + case _GuardrailUnavailable(): + await _end_unscreened_session(client_ws) + return False + case None | RealtimeGuardrailRewrite(): + return True + + async def _close_bedrock_input( bedrock_stream: BedrockBidirectionalStream, close_messages: Sequence[str], @@ -381,6 +577,8 @@ class BedrockRealtime(BaseAWSLLM): aws_external_id: str | None = None, aws_session_tags: object = None, drop_params: bool = False, + user_api_key_dict: object | None = None, + litellm_metadata: Mapping[str, object] | None = None, **kwargs: object, ): """ @@ -461,7 +659,15 @@ class BedrockRealtime(BaseAWSLLM): try: await self._run_session( - websocket, open_bidirectional_stream, model, logging_obj, pending_opening_message, drop_params + websocket, + open_bidirectional_stream, + model, + logging_obj, + pending_opening_message, + drop_params, + BedrockRealtimeGuardrails( + request_data={"litellm_metadata": litellm_metadata or {}}, user_api_key_dict=user_api_key_dict + ), ) finally: await _close_bedrock_client(bedrock_client) @@ -474,6 +680,7 @@ class BedrockRealtime(BaseAWSLLM): logging_obj: LiteLLMLogging, pending_opening_message: str | None, drop_params: bool = False, + guardrails: BedrockRealtimeGuardrails | None = None, ) -> None: transformation_config: Final = BedrockRealtimeConfig(drop_params=drop_params) @@ -496,6 +703,7 @@ class BedrockRealtime(BaseAWSLLM): "session_configuration_request": None, } + session_guardrails: Final = guardrails if guardrails is not None else BedrockRealtimeGuardrails() outcome: Final = await self._bridge( websocket, bedrock_stream, @@ -504,7 +712,10 @@ class BedrockRealtime(BaseAWSLLM): session_state, logging_obj, initial_message=pending_opening_message, + guardrails=session_guardrails, ) + if session_guardrails.failed: + logging_obj.model_call_details[REALTIME_SESSION_ENDED_ABNORMALLY_KEY] = True logged_events: Final = ( *outcome.logged_events, @@ -544,10 +755,12 @@ class BedrockRealtime(BaseAWSLLM): session_state: RealtimeResponseTransformInput, logging_obj: LiteLLMLogging, initial_message: str | None, + guardrails: BedrockRealtimeGuardrails | None = None, ) -> _BridgeOutcome: """Run both forwarding directions until the client leaves or either side fails.""" logged: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: events forwarded before a failure are still spend signals: Final = SessionEndSignals() + session_guardrails: Final = guardrails if guardrails is not None else BedrockRealtimeGuardrails() async def collect_logged_events() -> None: try: @@ -559,6 +772,7 @@ class BedrockRealtime(BaseAWSLLM): logging_obj, session_state, signals=signals, + guardrails=session_guardrails, ): logged.append(event) except Exception: @@ -575,6 +789,7 @@ class BedrockRealtime(BaseAWSLLM): logging_obj, initial_message, signals=signals, + guardrails=session_guardrails, ) ) bedrock_task: Final = asyncio.create_task(collect_logged_events()) @@ -610,6 +825,7 @@ class BedrockRealtime(BaseAWSLLM): initial_message: str | None = None, *, signals: SessionEndSignals, + guardrails: BedrockRealtimeGuardrails | None = None, ) -> None: """Forward messages from client WebSocket to Bedrock stream. @@ -631,10 +847,23 @@ class BedrockRealtime(BaseAWSLLM): raise _stream_ended_before_session_closed() from e verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200]) + forward_to_bedrock: Final = functools.partial( + _forward_client_message, + client_ws, + bedrock_stream, + transformation_config, + model, + session_state, + send_to_bedrock, + ) + session_guardrails: Final = guardrails if guardrails is not None else BedrockRealtimeGuardrails() try: async for message in _client_messages(client_ws, initial_message): verbose_proxy_logger.debug("Bedrock Realtime: Received from client: %s", message[:200]) parsed_client_message = _parse_client_message(message) + if parsed_client_message is None: + verbose_proxy_logger.warning("Bedrock Realtime: dropped a client message that is not valid JSON") + continue is_session_update = _json_str(parsed_client_message.get("type")) == "session.update" if is_session_update: rejection = transformation_config.session_update_rejection(parsed_client_message) @@ -643,18 +872,15 @@ class BedrockRealtime(BaseAWSLLM): break continue - transformed_messages = transformation_config.transform_realtime_request( - message=message, - model=model, - session_configuration_request=session_state.get("session_configuration_request"), - ) - if transformed_messages and not _is_session_committed(client_ws.scope): - await _send_opening_message( - client_ws, bedrock_stream, message, transformed_messages, send_to_bedrock - ) - else: - for bedrock_message in transformed_messages: - await send_to_bedrock(bedrock_message) + match await session_guardrails.screen_client_event( + parsed_client_message, client_ws, forward_to_bedrock + ): + case _ClientEventRoute.END_SESSION: + break + case _ClientEventRoute.DROP: + continue + case _ClientEventRoute.FORWARD: + await forward_to_bedrock(message) if is_session_update and not await _ack_session_update( client_ws, transformation_config, model, logging_obj, parsed_client_message @@ -679,6 +905,7 @@ class BedrockRealtime(BaseAWSLLM): session_state: RealtimeResponseTransformInput, *, signals: SessionEndSignals, + guardrails: BedrockRealtimeGuardrails | None = None, ) -> AsyncIterator[OpenAIRealtimeEvents]: """Forward messages from Bedrock to the client, yielding the ones to record for spend logging. @@ -694,6 +921,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200]) return True + session_guardrails: Final = guardrails if guardrails is not None else BedrockRealtimeGuardrails() output: Final = await bedrock_stream.await_output() while True: result = await output[1].receive() @@ -750,3 +978,7 @@ class BedrockRealtime(BaseAWSLLM): return if _should_log_event(openai_message): yield openai_message + if not await session_guardrails.continues_after( + openai_message, client_ws, transformation_config.client_audio_streamed + ): + return 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/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index c5488bc8a98..26246cfd7b6 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -103,7 +103,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): delivers_ended_stream_rewrites = True assembles_streamed_response = True - def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: + def get_structured_messages(self, data: dict, *, keep_tool_errors: bool = False) -> list[AllMessageValues] | None: """ Convert chat completions request data to OpenAI-spec structured messages. 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/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index aa4dea196ff..20b7506ae0a 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -777,7 +777,7 @@ class OpenAIResponsesHandler(BaseTranslation): delivers_ended_stream_rewrites = True assembles_streamed_response = True - def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: + def get_structured_messages(self, data: dict, *, keep_tool_errors: bool = False) -> list[AllMessageValues] | None: """ Convert Responses API request data to OpenAI-spec structured messages. 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 771f4cb3db4..77362a03757 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 @@ -8087,7 +8092,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 @@ -8154,7 +8159,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, @@ -8228,7 +8233,7 @@ def text_completion( messages=messages, *args, **kwargs, - **optional_params, + **_ForwardedCompletionParams(**optional_params), ) if kwargs.get("acompletion", False) is True: return response @@ -8572,7 +8577,7 @@ def transcription( ) -> ( TranscriptionResponse | Stream[TranscriptionStreamEvent] - | Coroutine[Any, Any, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] + | Coroutine[object, object, TranscriptionResponse | AsyncStream[TranscriptionStreamEvent]] ): """ Calls openai + azure whisper endpoints. @@ -8661,7 +8666,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/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4796ed3ac9c..9a3027475b7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -13865,12 +13865,12 @@ "azure_ai/jais-30b-chat": { "deprecation_date": "2026-01-30", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", - "input_cost_per_token": 0.0032, + "input_cost_per_token": 3.2e-06, "litellm_provider": "azure_ai", "max_input_tokens": 8192, "max_output_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.00971, + "output_cost_per_token": 9.71e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", "supports_pdf_input": false }, @@ -55160,8 +55160,8 @@ "watsonx/core42/jais-13b-chat": { "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.0005, - "output_cost_per_token": 0.002, + "input_cost_per_token": 1.8e-06, + "output_cost_per_token": 1.8e-06, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -64213,6 +64213,39 @@ "supports_tool_choice": true, "supports_vision": false }, + "tencent/hy4": { + "cache_read_input_token_cost": 4.2e-08, + "comment": "TokenHub bills Hy4 at 50% of list until 2026-12-31 23:59:59 Beijing time; these are the list rates.", + "input_cost_per_token": 8.34e-07, + "input_cost_per_token_cache_hit": 4.2e-08, + "litellm_provider": "tencent", + "max_input_tokens": 960000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.501e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true + }, + "tencent/hy4-preview": { + "cache_read_input_token_cost": 4.2e-08, + "comment": "TokenHub bills Hy4 preview at 50% of list until 2026-12-31 23:59:59 Beijing time; these are the list rates.", + "input_cost_per_token": 8.34e-07, + "input_cost_per_token_cache_hit": 4.2e-08, + "litellm_provider": "tencent", + "max_input_tokens": 960000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.501e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true + }, "tencent/minimax-m3": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 6e-08, @@ -70882,6 +70915,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token_above_272k_tokens": 1e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, "source": "https://docs.perplexity.ai/docs/agent-api/models" }, "perplexity/openai/gpt-5.4": { @@ -83107,13 +83141,16 @@ "mode": "chat", "max_input_tokens": 1000000, "max_output_tokens": 393216, - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 8e-07, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-1-flash-1017" + "source": "https://empiriolabs.ai/models/deepseek-v4-1-flash-1017", + "cache_read_input_token_cost": 2e-08, + "supports_prompt_caching": true, + "supports_function_calling": true }, "empiriolabs/deepseek-v4-flash-0731": { "litellm_provider": "empiriolabs", @@ -83125,7 +83162,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731" + "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-flash-0731:variant1": { "litellm_provider": "empiriolabs", @@ -83137,7 +83175,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731" + "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro": { "litellm_provider": "empiriolabs", @@ -83149,7 +83188,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro-0813": { "litellm_provider": "empiriolabs", @@ -83161,7 +83201,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro-0813" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro-0813", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro:variant1": { "litellm_provider": "empiriolabs", @@ -83179,7 +83220,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro:variant2": { "litellm_provider": "empiriolabs", @@ -83199,7 +83241,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro:variant3": { "litellm_provider": "empiriolabs", @@ -83211,7 +83254,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/fugu-max": { "litellm_provider": "empiriolabs", @@ -83409,7 +83453,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/glm-5-2" + "source": "https://empiriolabs.ai/models/glm-5-2", + "supports_function_calling": true }, "empiriolabs/glm-5-2:variant2": { "litellm_provider": "empiriolabs", @@ -83427,7 +83472,8 @@ "search_context_size_high": 0.013 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/glm-5-2" + "source": "https://empiriolabs.ai/models/glm-5-2", + "supports_function_calling": true }, "empiriolabs/glm-5-3": { "litellm_provider": "empiriolabs", @@ -83481,7 +83527,8 @@ "supports_reasoning": true, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-6" + "source": "https://empiriolabs.ai/models/kimi-k2-6", + "supports_function_calling": true }, "empiriolabs/kimi-k2-7-code": { "litellm_provider": "empiriolabs", @@ -83500,7 +83547,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-7-code" + "source": "https://empiriolabs.ai/models/kimi-k2-7-code", + "supports_function_calling": true }, "empiriolabs/kimi-k2-7-code-highspeed": { "litellm_provider": "empiriolabs", @@ -83519,7 +83567,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-7-code-highspeed" + "source": "https://empiriolabs.ai/models/kimi-k2-7-code-highspeed", + "supports_function_calling": true }, "empiriolabs/kimi-k2-7-code:variant1": { "litellm_provider": "empiriolabs", @@ -83540,7 +83589,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-7-code" + "source": "https://empiriolabs.ai/models/kimi-k2-7-code", + "supports_function_calling": true }, "empiriolabs/kimi-k3": { "litellm_provider": "empiriolabs", @@ -83559,7 +83609,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k3" + "source": "https://empiriolabs.ai/models/kimi-k3", + "supports_function_calling": true }, "empiriolabs/kimi-k3-1-pro": { "litellm_provider": "empiriolabs", @@ -83601,7 +83652,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/mimo-v2-6-flash" + "source": "https://empiriolabs.ai/models/mimo-v2-6-flash", + "supports_function_calling": true }, "empiriolabs/mimo-v2-6-pro": { "litellm_provider": "empiriolabs", @@ -83621,7 +83673,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/mimo-v2-6-pro" + "source": "https://empiriolabs.ai/models/mimo-v2-6-pro", + "supports_function_calling": true }, "empiriolabs/mimo-v2-6-pro-ultraspeed": { "litellm_provider": "empiriolabs", @@ -83635,7 +83688,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/mimo-v2-6-pro-ultraspeed" + "source": "https://empiriolabs.ai/models/mimo-v2-6-pro-ultraspeed", + "supports_function_calling": true }, "empiriolabs/minimax-m2-7": { "litellm_provider": "empiriolabs", @@ -83750,7 +83804,8 @@ "search_context_size_high": 0.14 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/mistral-small-4" + "source": "https://empiriolabs.ai/models/mistral-small-4", + "supports_function_calling": true }, "empiriolabs/muse-glimmer-30b": { "litellm_provider": "empiriolabs", @@ -83765,7 +83820,8 @@ "supports_reasoning": true, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/muse-glimmer-30b" + "source": "https://empiriolabs.ai/models/muse-glimmer-30b", + "supports_function_calling": true }, "empiriolabs/muse-spark-1-1": { "litellm_provider": "empiriolabs", @@ -83831,7 +83887,8 @@ "supports_system_messages": true, "supports_vision": true, "supports_prompt_caching": true, - "source": "https://empiriolabs.ai/models/nova-lite-1-0" + "source": "https://empiriolabs.ai/models/nova-lite-1-0", + "supports_function_calling": true }, "empiriolabs/nova-lite-2": { "litellm_provider": "empiriolabs", @@ -83845,7 +83902,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_prompt_caching": true, - "source": "https://empiriolabs.ai/models/nova-lite-2" + "source": "https://empiriolabs.ai/models/nova-lite-2", + "supports_function_calling": true }, "empiriolabs/nova-micro-1-0": { "litellm_provider": "empiriolabs", @@ -83857,7 +83915,8 @@ "cache_read_input_token_cost": 2.24e-08, "supports_system_messages": true, "supports_prompt_caching": true, - "source": "https://empiriolabs.ai/models/nova-micro-1-0" + "source": "https://empiriolabs.ai/models/nova-micro-1-0", + "supports_function_calling": true }, "empiriolabs/nova-pro-1-0": { "litellm_provider": "empiriolabs", @@ -83868,7 +83927,8 @@ "output_cost_per_token": 9.6e-06, "supports_system_messages": true, "supports_vision": true, - "source": "https://empiriolabs.ai/models/nova-pro-1-0" + "source": "https://empiriolabs.ai/models/nova-pro-1-0", + "supports_function_calling": true }, "empiriolabs/perplexity-advanced-deep-research": { "litellm_provider": "empiriolabs", @@ -83968,7 +84028,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b" + "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-27b": { "litellm_provider": "empiriolabs", @@ -83989,7 +84050,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-27b" + "source": "https://empiriolabs.ai/models/qwen3-5-27b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-35b-a3b": { "litellm_provider": "empiriolabs", @@ -84010,7 +84072,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b" + "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-397b-a17b": { "litellm_provider": "empiriolabs", @@ -84031,7 +84094,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b" + "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-4b": { "litellm_provider": "empiriolabs", @@ -84082,7 +84146,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-flash" + "source": "https://empiriolabs.ai/models/qwen3-5-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-5-flash:variant1": { "litellm_provider": "empiriolabs", @@ -84105,7 +84170,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-flash" + "source": "https://empiriolabs.ai/models/qwen3-5-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-5-omni-flash": { "litellm_provider": "empiriolabs", @@ -84123,7 +84189,8 @@ "search_context_size_high": 0.015 }, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash" + "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-5-omni-plus": { "litellm_provider": "empiriolabs", @@ -84142,7 +84209,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus" + "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-5-plus": { "litellm_provider": "empiriolabs", @@ -84151,8 +84219,8 @@ "max_output_tokens": 65536, "input_cost_per_token": 3.6e-07, "output_cost_per_token": 2.21e-06, - "input_cost_per_token_above_256k_tokens": 1.08e-06, - "output_cost_per_token_above_256k_tokens": 6.62e-06, + "input_cost_per_token_above_256k_tokens": 4.5e-07, + "output_cost_per_token_above_256k_tokens": 2.76e-06, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, @@ -84163,7 +84231,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-plus" + "source": "https://empiriolabs.ai/models/qwen3-5-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-5-plus:variant1": { "litellm_provider": "empiriolabs", @@ -84186,7 +84255,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-plus" + "source": "https://empiriolabs.ai/models/qwen3-5-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-6-27b": { "litellm_provider": "empiriolabs", @@ -84205,7 +84275,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-27b" + "source": "https://empiriolabs.ai/models/qwen3-6-27b", + "supports_function_calling": true }, "empiriolabs/qwen3-6-35b-a3b": { "litellm_provider": "empiriolabs", @@ -84219,7 +84290,8 @@ "supports_reasoning": true, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-35b-a3b" + "source": "https://empiriolabs.ai/models/qwen3-6-35b-a3b", + "supports_function_calling": true }, "empiriolabs/qwen3-6-flash": { "litellm_provider": "empiriolabs", @@ -84240,7 +84312,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-flash" + "source": "https://empiriolabs.ai/models/qwen3-6-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-6-flash:variant1": { "litellm_provider": "empiriolabs", @@ -84261,7 +84334,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-flash" + "source": "https://empiriolabs.ai/models/qwen3-6-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-6-plus": { "litellm_provider": "empiriolabs", @@ -84282,7 +84356,8 @@ "search_context_size_high": 0.026 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-plus" + "source": "https://empiriolabs.ai/models/qwen3-6-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-6-plus:variant1": { "litellm_provider": "empiriolabs", @@ -84303,7 +84378,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-plus" + "source": "https://empiriolabs.ai/models/qwen3-6-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-7-flash": { "litellm_provider": "empiriolabs", @@ -84330,7 +84406,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-flash" + "source": "https://empiriolabs.ai/models/qwen3-7-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-7-flash:variant1": { "litellm_provider": "empiriolabs", @@ -84357,7 +84434,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-flash" + "source": "https://empiriolabs.ai/models/qwen3-7-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-7-max": { "litellm_provider": "empiriolabs", @@ -84375,7 +84453,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-max" + "source": "https://empiriolabs.ai/models/qwen3-7-max", + "supports_function_calling": true }, "empiriolabs/qwen3-7-max:variant1": { "litellm_provider": "empiriolabs", @@ -84393,7 +84472,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-max" + "source": "https://empiriolabs.ai/models/qwen3-7-max", + "supports_function_calling": true }, "empiriolabs/qwen3-7-plus": { "litellm_provider": "empiriolabs", @@ -84414,7 +84494,8 @@ "search_context_size_high": 0.03 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-plus" + "source": "https://empiriolabs.ai/models/qwen3-7-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-7-plus:variant1": { "litellm_provider": "empiriolabs", @@ -84438,7 +84519,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-plus" + "source": "https://empiriolabs.ai/models/qwen3-7-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-8-27b": { "litellm_provider": "empiriolabs", @@ -84473,7 +84555,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-flash" + "source": "https://empiriolabs.ai/models/qwen3-8-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-8-max": { "litellm_provider": "empiriolabs", @@ -84492,7 +84575,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-max" + "source": "https://empiriolabs.ai/models/qwen3-8-max", + "supports_function_calling": true }, "empiriolabs/qwen3-8-max-0902": { "litellm_provider": "empiriolabs", @@ -84511,7 +84595,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-max-0902" + "source": "https://empiriolabs.ai/models/qwen3-8-max-0902", + "supports_function_calling": true }, "empiriolabs/qwen3-8-max:variant1": { "litellm_provider": "empiriolabs", @@ -84530,7 +84615,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-max" + "source": "https://empiriolabs.ai/models/qwen3-8-max", + "supports_function_calling": true }, "empiriolabs/qwen3-8-omni-flash": { "litellm_provider": "empiriolabs", @@ -84550,7 +84636,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/qwen3-8-omni-flash" + "source": "https://empiriolabs.ai/models/qwen3-8-omni-flash", + "supports_function_calling": true }, "empiriolabs/seed-2-0-code": { "litellm_provider": "empiriolabs", @@ -84564,7 +84651,8 @@ "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, - "source": "https://empiriolabs.ai/models/seed-2-0-code" + "source": "https://empiriolabs.ai/models/seed-2-0-code", + "supports_function_calling": true }, "empiriolabs/seed-2-0-lite": { "litellm_provider": "empiriolabs", @@ -84579,7 +84667,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-0-lite" + "source": "https://empiriolabs.ai/models/seed-2-0-lite", + "supports_function_calling": true }, "empiriolabs/seed-2-0-mini": { "litellm_provider": "empiriolabs", @@ -84594,7 +84683,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-0-mini" + "source": "https://empiriolabs.ai/models/seed-2-0-mini", + "supports_function_calling": true }, "empiriolabs/seed-2-0-pro": { "litellm_provider": "empiriolabs", @@ -84609,7 +84699,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-0-pro" + "source": "https://empiriolabs.ai/models/seed-2-0-pro", + "supports_function_calling": true }, "empiriolabs/seed-2-1-turbo": { "litellm_provider": "empiriolabs", @@ -84622,7 +84713,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-1-turbo" + "source": "https://empiriolabs.ai/models/seed-2-1-turbo", + "supports_function_calling": true }, "empiriolabs/skylark-embedding-vision": { "litellm_provider": "empiriolabs", 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/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index b8119cbe40c..8157c33d247 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -16540,6 +16540,22 @@ "description": "Character budget for each text sent to the guardrail's model. Implemented by guardrail='decision_model'.", "title": "Max Input Chars" }, + "mcp_block_behavior": { + "anyOf": [ + { + "enum": [ + "tool_result", + "raise" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "description": "How to handle guardrail blocks when executing MCP tools for an LLM request. 'raise' propagates the guardrail error to the client. 'tool_result' or None returns the block as a tool result and lets the LLM continue.", + "title": "Mcp Block Behavior" + }, "model": { "anyOf": [ { @@ -16597,7 +16613,7 @@ "type": "null" } ], - "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", + "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. Bedrock Nova Sonic cannot cancel a spoken turn's answer, so a blocked spoken turn always ends its session. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", "title": "On Violation" }, "only_scan_new_messages": { @@ -16673,7 +16689,7 @@ "type": "null" } ], - "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set.", + "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to the guardrail's block reason if not set.", "title": "Realtime Violation Message" }, "run_in_parallel": { @@ -19116,6 +19132,22 @@ "description": "Character budget for each text sent to the guardrail's model. Implemented by guardrail='decision_model'.", "title": "Max Input Chars" }, + "mcp_block_behavior": { + "anyOf": [ + { + "enum": [ + "tool_result", + "raise" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "description": "How to handle guardrail blocks when executing MCP tools for an LLM request. 'raise' propagates the guardrail error to the client. 'tool_result' or None returns the block as a tool result and lets the LLM continue.", + "title": "Mcp Block Behavior" + }, "metadata": { "anyOf": [ { @@ -19271,7 +19303,7 @@ "type": "null" } ], - "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", + "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. Bedrock Nova Sonic cannot cancel a spoken turn's answer, so a blocked spoken turn always ends its session. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", "title": "On Violation" }, "only_scan_new_messages": { @@ -19692,7 +19724,7 @@ "type": "null" } ], - "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set.", + "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to the guardrail's block reason if not set.", "title": "Realtime Violation Message" }, "rules": { @@ -51017,6 +51049,10 @@ ], "title": "Content" }, + "is_error": { + "title": "Is Error", + "type": "boolean" + }, "role": { "const": "tool", "title": "Role", 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/codex_catalog_source.py b/litellm/proxy/common_utils/codex_catalog_source.py new file mode 100644 index 00000000000..994ffe8e72e --- /dev/null +++ b/litellm/proxy/common_utils/codex_catalog_source.py @@ -0,0 +1,181 @@ +"""Codex's own catalog for the Codex release asking for `/v1/models`, fetched from that release's tag on GitHub. + +Codex fetches its catalog from chatgpt.com, which needs a ChatGPT login the proxy never sees, so the proxy reads +the same two files at the release's public tag instead, and falls back to the copy bundled with LiteLLM when the +fetch fails, the release has no tag (dev builds report 0.0.0) or ``LITELLM_OFFLINE_MODE`` is on. +""" + +from __future__ import annotations + +import asyncio +import re +import time +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from functools import cache +from types import MappingProxyType +from typing import Final, Protocol, TypeVar + +import httpx + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.offline_mode import is_offline_mode +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared factory has untyped optional parameters +) +from litellm.proxy.common_utils.codex_model_catalog import ( + CODEX_BUNDLED_MODELS_PATH, + CodexStockCatalog, + CodexStockModel, + bundled_codex_models, + codex_base_instructions, +) +from litellm.types.llms.custom_http import httpxSpecialProvider + +CODEX_RELEASE_FILE_URL: Final = ( + "https://raw.githubusercontent.com/openai/codex/rust-v{version}/codex-rs/models-manager/{file}" +) +CODEX_FETCH_TIMEOUT_SECONDS: Final = 2.0 +CODEX_FETCH_FAILURE_TTL_SECONDS: Final = 3600.0 +CODEX_CACHED_VERSIONS: Final = 8 +CODEX_REMEMBERED_FAILURES: Final = 256 +CODEX_FETCHES_PER_MINUTE: Final = 6 + +_RELEASE_VERSION: Final = re.compile(r"[0-9]+\.[0-9]+\.[0-9]+(?:-[0-9A-Za-z.]+)?") + + +@dataclass(frozen=True, slots=True) +class CodexStock: + models: Mapping[str, CodexStockModel] + instructions: str + + +_Value: Final = TypeVar("_Value") + + +class CodexReleaseClient(Protocol): + def get(self, url: str, *, timeout: float | None = None) -> Awaitable[httpx.Response]: ... + + +@cache +def bundled_codex_stock() -> CodexStock: + return CodexStock(models=bundled_codex_models(), instructions=codex_base_instructions()) + + +async def _release_file(client: CodexReleaseClient, version: str, file: str, timeout: float) -> str: + response: Final = await client.get(CODEX_RELEASE_FILE_URL.format(version=version, file=file), timeout=timeout) + return response.raise_for_status().content.decode("utf-8") + + +async def _release_stock(client: CodexReleaseClient, version: str, timeout: float) -> CodexStock: + models_json, instructions = await asyncio.wait_for( + asyncio.gather( + _release_file(client, version, "models.json", timeout), + _release_file(client, version, "prompt.md", timeout), + ), + timeout, + ) + catalog: Final = CodexStockCatalog.model_validate_json(models_json) + return CodexStock( + models=MappingProxyType({model.slug: model for model in catalog.models}), instructions=instructions + ) + + +class CodexCatalogSource: + """Codex's catalog per release: the one fetched from GitHub, else the bundled one. + + A fetched catalog is kept for good, a failed fetch is remembered for `failure_ttl` seconds so a blocked egress + costs one timeout per worker per version per hour. Only the `max_versions` most recently fetched catalogs and + the `max_failures` most recent failures are kept, apart, so versions that fail never push out ones that were + fetched. + Requests for a version already being fetched wait on that fetch, and a worker starts at most + `fetches_per_minute` fetches, serving the bundled catalog past that, so a caller cycling made-up versions + cannot turn every request into GitHub requests.""" + + def __init__( + self, + client: Callable[[], CodexReleaseClient], + *, + offline: Callable[[], bool] = is_offline_mode, + clock: Callable[[], float] = time.monotonic, + timeout: float = CODEX_FETCH_TIMEOUT_SECONDS, + failure_ttl: float = CODEX_FETCH_FAILURE_TTL_SECONDS, + max_versions: int = CODEX_CACHED_VERSIONS, + max_failures: int = CODEX_REMEMBERED_FAILURES, + fetches_per_minute: int = CODEX_FETCHES_PER_MINUTE, + ) -> None: + self._client: Final = client + self._offline: Final = offline + self._clock: Final = clock + self._timeout: Final = timeout + self._failure_ttl: Final = failure_ttl + self._max_versions: Final = max_versions + self._max_failures: Final = max_failures + self._fetches_per_minute: Final = fetches_per_minute + self._fetched: Mapping[str, CodexStock] = MappingProxyType({}) + self._failed_until: Mapping[str, float] = MappingProxyType({}) + self._in_flight: Mapping[str, asyncio.Future[CodexStock]] = MappingProxyType({}) + self._fetch_starts: tuple[float, ...] = () + + async def stock(self, client_version: str) -> CodexStock: + if self._offline() or _RELEASE_VERSION.fullmatch(client_version) is None: + return bundled_codex_stock() + known: Final = self._fetched.get(client_version) + if known is not None: + return known + if self._clock() < self._failed_until.get(client_version, float("-inf")): + return bundled_codex_stock() + pending: Final = self._in_flight.get(client_version) + if pending is not None: + return await asyncio.shield(pending) + if not self._take_fetch_slot(): + return bundled_codex_stock() + started: Final = asyncio.ensure_future(self._fetch_and_store(client_version)) + self._in_flight = MappingProxyType({**self._in_flight, client_version: started}) + return await asyncio.shield(started) + + def _take_fetch_slot(self) -> bool: + now: Final = self._clock() + recent: Final = tuple(start for start in self._fetch_starts if now - start < 60) + allowed: Final = len(recent) < self._fetches_per_minute + self._fetch_starts = (*recent, now) if allowed else recent + return allowed + + async def _fetch_and_store(self, version: str) -> CodexStock: + try: + fetched: Final = await self._fetch(version) + finally: + self._in_flight = MappingProxyType({key: task for key, task in self._in_flight.items() if key != version}) + if fetched is None: + self._failed_until = _latest( + self._failed_until, version, self._clock() + self._failure_ttl, self._max_failures + ) + return bundled_codex_stock() + self._fetched = _latest(self._fetched, version, fetched, self._max_versions) + return fetched + + async def _fetch(self, version: str) -> CodexStock | None: + try: + return await _release_stock(self._client(), version, self._timeout) + except Exception as e: # noqa: BLE001 # a failed best-effort fetch must never fail the model listing + verbose_proxy_logger.warning( + "Codex %s's model catalog could not be fetched from GitHub (%r), serving the catalog bundled " + "with LiteLLM (%s) for that version for the next %d seconds", + version, + e, + CODEX_BUNDLED_MODELS_PATH.name, + self._failure_ttl, + ) + return None + + +def _latest(stored: Mapping[str, _Value], version: str, value: _Value, limit: int) -> Mapping[str, _Value]: + others: Final = tuple((key, kept) for key, kept in stored.items() if key != version) + return MappingProxyType(dict((*others[max(0, len(others) - limit + 1) :], (version, value)))) + + +def _default_client() -> CodexReleaseClient: + return get_async_httpx_client(llm_provider=httpxSpecialProvider.CodexModelCatalog) + + +codex_catalog_source: Final = CodexCatalogSource(_default_client) diff --git a/litellm/proxy/common_utils/codex_model_catalog.py b/litellm/proxy/common_utils/codex_model_catalog.py index 01cce9a6ce6..6deee04e273 100644 --- a/litellm/proxy/common_utils/codex_model_catalog.py +++ b/litellm/proxy/common_utils/codex_model_catalog.py @@ -709,10 +709,16 @@ def codex_model_list_body( label_entries: Sequence[tuple[str, str]] | None = None, *, client_version: str | None = None, + stock: Mapping[str, CodexStockModel] | None = None, + instructions: str | None = None, ) -> str: - """The `/v1/models?client_version=...` body, logging the models Codex's byte limit left out.""" + """The `/v1/models?client_version=...` body, logging the models Codex's byte limit left out. `stock` and + `instructions` are the requesting Codex release's catalog and prompt, the bundled ones when None.""" body: Final = codex_models_response_json( - codex_catalog_rows(rows, entries, llm_router, team_id, label_entries), client_version=client_version + codex_catalog_rows(rows, entries, llm_router, team_id, label_entries), + stock=stock, + instructions=instructions, + client_version=client_version, ) if body.left_out: verbose_proxy_logger.warning( 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/exception_utils.py b/litellm/proxy/guardrails/exception_utils.py index 518c61d1fd3..68a29e5c2ef 100644 --- a/litellm/proxy/guardrails/exception_utils.py +++ b/litellm/proxy/guardrails/exception_utils.py @@ -13,16 +13,22 @@ def is_fastapi_http_exception(e: Exception, block_status_codes: Collection[int]) def enrich_http_exception_with_guardrail_context(exc: BaseException, callback: object) -> None: + if not hasattr(exc, "mcp_block_behavior"): + setattr( # noqa: B010 # exception types do not declare per-deployment MCP context + exc, "mcp_block_behavior", getattr(callback, "mcp_block_behavior", None) + ) try: from fastapi.exceptions import HTTPException except ImportError: return if not isinstance(exc, HTTPException): return + guardrail_name: Final[object] = getattr(callback, "guardrail_name", None) detail: Final = getattr(exc, "detail", None) if not isinstance(detail, dict): + if isinstance(guardrail_name, str): + setattr(exc, "guardrail_name", guardrail_name) # noqa: B010 # HTTPException does not declare guardrail_name return - guardrail_name: Final[object] = getattr(callback, "guardrail_name", None) if guardrail_name: detail.setdefault("guardrail_name", guardrail_name) event_hook: Final[object] = getattr(callback, "event_hook", None) 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/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 24b30e65e2b..2f2294db522 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -644,6 +644,9 @@ def _configure_callback_scoping( ) logging_only_scope: Final = None if logging_only_scope_error is not None else litellm_params.logging_only_scope custom_guardrail_callback.logging_only_scope = logging_only_scope + custom_guardrail_callback.mcp_block_behavior = ( # rebind-ok: initialize provider-independent callback configuration + litellm_params.mcp_block_behavior + ) custom_guardrail_callback.logging_only_continue_on_input_failure = ( logging_only_continue_error is None and logging_only_scope is None 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 42321fe1b1d..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( @@ -1112,6 +1114,7 @@ def _is_configured_pre_routing_strategy(llm_router: "Router", router_name: str) router_registry.plugins("complexity"), router_registry.plugins("adaptive"), router_registry.plugins("quality"), + router_registry.plugins("switchyard"), ) ) 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 2c7453ebb1c..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 ( @@ -142,6 +143,7 @@ from litellm.router_strategy.auto_router.auto_router_model_naming import ( count_capability_routers, gated_capabilities_of, is_complexity_router_model, + strategy_router_dependencies, validate_complexity_router_config_placement, validate_complexity_router_config_write, validate_strategy_router_model_write, @@ -209,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): @@ -245,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 { @@ -355,6 +358,22 @@ async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment return deployment_pydantic_obj +def _self_routing_violation( + model_name: str | None, + incoming_params: GenericLiteLLMParams | None, + existing_params: GenericLiteLLMParams | None, +) -> str | None: + if not model_name: + return None + effective_params: Final = { + **(_decrypted_litellm_params(existing_params) if existing_params is not None else {}), + **(_dump_without_none(incoming_params) if incoming_params is not None else {}), + } + if any(dependency.model_name == model_name for dependency in strategy_router_dependencies(effective_params)): + return f"A strategy router cannot route to itself: {model_name} is listed as one of its own models" + return None + + def _strategy_router_write_violation( incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None, @@ -374,6 +393,11 @@ def _strategy_router_write_violation( """ if incoming_params is None: return None + from litellm.router_strategy.switchyard_router.config import validate_switchyard_config_write + + switchyard_violation: Final = validate_switchyard_config_write(incoming_params.switchyard_router_config) + if switchyard_violation is not None: + return switchyard_violation config_violation: Final = validate_complexity_router_config_write( complexity_router_config=( _effective_complexity_router_config(incoming_params, existing_params) @@ -406,8 +430,11 @@ def _strategy_router_write_violation( def _raise_on_strategy_router_write_violation( incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None, + model_name: str | None = None, ) -> None: - violation = _strategy_router_write_violation(incoming_params=incoming_params, existing_params=existing_params) + violation: Final = _self_routing_violation(model_name, incoming_params, existing_params) or ( + _strategy_router_write_violation(incoming_params=incoming_params, existing_params=existing_params) + ) if violation is None: return raise ProxyException( @@ -1346,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( @@ -1356,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}, @@ -1396,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, ...]: @@ -1518,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: @@ -1751,6 +1776,7 @@ async def patch_model( _raise_on_strategy_router_write_violation( incoming_params=incoming_params, existing_params=remaining_params, + model_name=patch_data.model_name or db_model.model_name, ) written_team_id: Final = ( @@ -2054,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: @@ -3120,6 +3146,7 @@ async def add_new_model( _raise_on_strategy_router_write_violation( incoming_params=model_params.litellm_params, existing_params=None, + model_name=model_params.model_name, ) _raise_if_rate_limits_required_but_missing( @@ -3325,6 +3352,7 @@ async def update_model( _raise_on_strategy_router_write_violation( incoming_params=model_params.litellm_params, existing_params=deployment.litellm_params, + model_name=model_params.model_name or deployment.model_name, ) renamed_to: Final = ( model_params.model_name @@ -3972,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 @@ -4028,6 +4057,7 @@ async def clear_cache() -> ReconcileOutcome: llm_router.registry().plugins("complexity").pop(model_name, None) llm_router.registry().plugins("adaptive").pop(model_name, None) llm_router.registry().plugins("quality").pop(model_name, None) + llm_router.registry().plugins("switchyard").pop(model_name, None) # Reload only DB models. _add_deployment_locked, not add_deployment: this # coroutine already holds MODEL_RECONCILE_LOCK and asyncio.Lock is not 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..3297a298af4 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, @@ -521,6 +522,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( publish_auth_cache_invalidation, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy +from litellm.proxy.common_utils.codex_catalog_source import codex_catalog_source from litellm.proxy.common_utils.codex_model_catalog import CodexCatalogRow, codex_catalog_rows, codex_model_list_body from litellm.proxy.common_utils.config_includes import resolve_include_file_path, resolve_includes from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber @@ -1136,7 +1138,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 +3201,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 +4554,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 +5891,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 +7395,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 +8701,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 +11058,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, @@ -12899,6 +12921,7 @@ async def model_list( admin_listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above if client_version is not None: + admin_codex_stock: Final = await codex_catalog_source.stock(client_version) return Response( content=await asyncio.to_thread( codex_model_list_body, @@ -12907,6 +12930,8 @@ async def model_list( llm_router, team_id or user_api_key_dict.team_id, client_version=client_version, + stock=admin_codex_stock.models, + instructions=admin_codex_stock.instructions, ), media_type="application/json", ) @@ -12946,6 +12971,7 @@ async def model_list( entries: Final = key_listing.entries labels: Final = key_listing.labels if client_version is not None: + codex_stock: Final = await codex_catalog_source.stock(client_version) return Response( content=await asyncio.to_thread( codex_model_list_body, @@ -12955,6 +12981,8 @@ async def model_list( team_id or user_api_key_dict.team_id, labels, client_version=client_version, + stock=codex_stock.models, + instructions=codex_stock.instructions, ), media_type="application/json", ) @@ -16020,10 +16048,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 +16228,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 +20749,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 64bce0d21f2..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: @@ -2660,14 +2680,18 @@ class ProxyLogging: if callback is None: raise ValueError(f"No callback found for guardrail: {guardrail_name}") - return await self._execute_guardrail_hook( - callback=callback, - hook_type=hook_type, - data=data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - response=response, - ) + try: + return await self._execute_guardrail_hook( + callback=callback, + hook_type=hook_type, + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + response=response, + ) + except Exception as e: + enrich_http_exception_with_guardrail_context(e, callback) + raise async def _process_guardrail_callback( self, @@ -5358,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, @@ -5411,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"), } @@ -6243,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'") @@ -7553,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 @@ -7903,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"] @@ -7913,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: @@ -7973,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) @@ -8000,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: @@ -8168,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/realtime_api/main.py b/litellm/realtime_api/main.py index be1df4aa0f6..85d87554023 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -149,7 +149,7 @@ def _build_litellm_metadata(kwargs: Mapping[str, object]) -> dict[str, object]: _GUARDRAIL_OPT_OUT_METADATA_KEYS: Final = ("user_api_key_metadata", "user_api_key_team_metadata") -def _translation_litellm_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: +def _session_litellm_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: request_metadata: Final = _OBJECT_MAPPING.validate_python(kwargs.get("metadata") or {}) guardrail_opt_outs: Final = { key: request_metadata[key] for key in _GUARDRAIL_OPT_OUT_METADATA_KEYS if key in request_metadata @@ -218,7 +218,7 @@ def _translation_rejection_error( def _raise_if_http_translation_rejected(custom_llm_provider: str, model: str, kwargs: Mapping[str, object]) -> None: rejection: Final = _translation_rejection( - custom_llm_provider, _HTTP_TRANSLATION_PROVIDERS, _translation_litellm_metadata(kwargs) + custom_llm_provider, _HTTP_TRANSLATION_PROVIDERS, _session_litellm_metadata(kwargs) ) if rejection is None: return @@ -858,9 +858,7 @@ async def _arealtime( ) session_litellm_metadata: Final = ( - dict(_translation_litellm_metadata(kwargs)) - if realtime_mode == "translation" - else _build_litellm_metadata(kwargs) + _build_litellm_metadata(kwargs) if live_session_start is not None else dict(_session_litellm_metadata(kwargs)) ) translation_rejection: Final = ( _translation_rejection(_custom_llm_provider, _WEBSOCKET_TRANSLATION_PROVIDERS, session_litellm_metadata) @@ -1014,6 +1012,8 @@ async def _arealtime( aws_external_id=aws_external_id, aws_session_tags=aws_session_tags, drop_params=drop_params, + user_api_key_dict=kwargs.get("user_api_key_dict"), + litellm_metadata=session_litellm_metadata, ) elif _custom_llm_provider == "xai": api_base = ( 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 33ea85d7d53..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, @@ -199,6 +198,8 @@ ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool | AnthropicToolSearchToolRegex NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n" NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"}) +RESPONSES_WEB_SEARCH_TOOL_TYPES: Final = frozenset({"web_search", "web_search_preview"}) +RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM: Final = frozenset({"computer_use", "image_generation", "local_shell", "shell"}) _INCOMPLETE_REASON_BY_FINISH_REASON: Final[Mapping[str, Literal["max_output_tokens", "content_filter"]]] = ( MappingProxyType({"length": "max_output_tokens", "content_filter": "content_filter", "refusal": "content_filter"}) ) @@ -384,6 +385,10 @@ class _HasId(Protocol): id: object +class _SupportsGet(Protocol): + def get(self, key: str, /) -> object: ... + + class _ResponsesToolCallItem(Protocol): name: object arguments: object @@ -982,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. @@ -1007,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 @@ -1026,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 @@ -1060,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" @@ -2948,7 +2951,7 @@ class LiteLLMCompletionResponsesConfig: tool_type: Final = tool.get("type") if tool_type == "mcp": return ResponsesToolChatForm(chat_tools=(cast(OpenAIMcpServerTool, tool),), web_search_options=None) - if tool_type == "web_search_preview" or tool_type == "web_search": + if tool_type in RESPONSES_WEB_SEARCH_TOOL_TYPES: _search_context_size: Final[Literal["low", "medium", "high"]] = cast( Literal["low", "medium", "high"], tool.get("search_context_size") ) @@ -3001,7 +3004,7 @@ class LiteLLMCompletionResponsesConfig: return ResponsesToolChatForm(chat_tools=(client_tool_search_chat_tool(tool),), web_search_options=None) if tool_type == "tool_search": return ResponsesToolChatForm(chat_tools=(anthropic_tool_search_chat_tool(),), web_search_options=None) - if tool_type in ("computer_use", "image_generation", "local_shell", "shell"): + if tool_type in RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM: verbose_logger.warning( "Dropping Responses API tool of type '%s': it has no Chat Completions " "equivalent and the target provider would reject the request.", 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 1104aee353e..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. @@ -964,6 +964,8 @@ class LiteLLM_Proxy_MCP_Handler: error=e, ) verbose_logger.error("BlockedPiiEntityError in MCP tool call: %s", e) + if LiteLLM_Proxy_MCP_Handler.should_raise_mcp_guardrail_block(e): + raise error_message = f"Tool call blocked: PII entity '{getattr(e, 'entity_type', 'unknown')}' detected by guardrail '{getattr(e, 'guardrail_name', 'unknown')}'. {e}" tool_results.append( { @@ -980,6 +982,8 @@ class LiteLLM_Proxy_MCP_Handler: error=e, ) verbose_logger.error("GuardrailRaisedException in MCP tool call: %s", e) + if LiteLLM_Proxy_MCP_Handler.should_raise_mcp_guardrail_block(e): + raise error_message = ( f"Tool call blocked: Guardrail '{getattr(e, 'guardrail_name', 'unknown')}' violation. {e}" ) @@ -998,6 +1002,8 @@ class LiteLLM_Proxy_MCP_Handler: error=e, ) verbose_logger.error("HTTPException in MCP tool call: %s", e) + if LiteLLM_Proxy_MCP_Handler.should_raise_mcp_guardrail_block(e): + raise error_message = f"Tool call failed: {str(e.detail) if hasattr(e, 'detail') else str(e)}" tool_results.append( { @@ -1026,6 +1032,10 @@ class LiteLLM_Proxy_MCP_Handler: _execute_tool_calls = execute_tool_calls + @staticmethod + def should_raise_mcp_guardrail_block(error: Exception) -> bool: + return getattr(error, "mcp_block_behavior", None) == "raise" + @staticmethod def create_follow_up_messages_for_chat( original_messages: list[object], diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 189366feab4..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 @@ -550,6 +558,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Phase 4: Tool Execution Events if self.phase == "tool_execution": + if self._stream_error is not None: + self.phase = "finished" + return await self._anext_impl() # Emit any queued tool execution events if self.tool_execution_events: return self.tool_execution_events.pop(0) @@ -713,6 +724,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.phase = "tool_execution" await self._generate_tool_execution_events() + if self._stream_error is not None: + return None if not isinstance(response_obj, ResponsesAPIResponse): return chunk if self._follow_up_pending(): @@ -974,6 +987,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): except Exception as e: verbose_logger.error("Error in tool execution: %s", e) + if LiteLLM_Proxy_MCP_Handler.should_raise_mcp_guardrail_block(e): + self._stream_error = e + return import traceback traceback.print_exc() @@ -1035,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 e3c6ab73532..12c0aa398d3 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,11 +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"}), - "cache_write_tokens": input_details.cache_write_tokens, - }, + **{**input_fields, "cache_write_tokens": input_details.cache_write_tokens}, 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_api.py b/litellm/router_api.py index 1ade356f494..ca370488f5d 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -65,6 +65,7 @@ if TYPE_CHECKING: from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.quality_router.quality_router import QualityRouter + from litellm.router_strategy.switchyard_router.switchyard_router import SwitchyardRouter from litellm.router_utils.cooldown_cache import CooldownCache from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets as _AttemptedFallbackTargets from litellm.router_utils.health_state_cache import DeploymentHealthCache @@ -106,6 +107,7 @@ if TYPE_CHECKING: | dict[str, list[TaggedPreRoutingStrategy[ComplexityRouter]]] | dict[str, list[TaggedPreRoutingStrategy[QualityRouter]]] | dict[str, list[TaggedPreRoutingStrategy[AdaptiveRouter]]] + | dict[str, list[TaggedPreRoutingStrategy[SwitchyardRouter]]] | list[RoutingPlugin] | RouterBudgetLimiting | DualCache @@ -171,6 +173,7 @@ PluginKind: TypeAlias = Literal[ "complexity", "quality", "adaptive", + "switchyard", "routing", "budget_limiter", "prompt_caching_pins", @@ -1106,6 +1109,11 @@ class RouterRegistry: self, kind: Literal["adaptive"] ) -> "dict[str, list[TaggedPreRoutingStrategy[AdaptiveRouter]]]": ... # mutable-ok: the live dict + @overload + def plugins( + self, kind: Literal["switchyard"] + ) -> "dict[str, list[TaggedPreRoutingStrategy[SwitchyardRouter]]]": ... # mutable-ok: the live dict + @overload def plugins(self, kind: Literal["routing"]) -> "list[RoutingPlugin]": ... # mutable-ok: the live list @@ -1132,6 +1140,8 @@ class RouterRegistry: return r._core_quality_routers if kind == "adaptive": return r._core_adaptive_routers + if kind == "switchyard": + return r._core_switchyard_routers if kind == "routing": return r._core_routing_plugins if kind == "budget_limiter": diff --git a/litellm/router_config.py b/litellm/router_config.py index 36154af296c..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: @@ -45,6 +47,7 @@ class RouterDynamicConfig: context_window_fallbacks: Sequence[object] | None content_policy_fallbacks: Sequence[object] | None drop_params: bool | None + disable_token_counter: bool provider_list: Sequence[str] expose_router_debug_in_errors: bool redact_user_api_key_info: bool | None @@ -102,6 +105,7 @@ def current_dynamic_config() -> RouterDynamicConfig: context_window_fallbacks=litellm.context_window_fallbacks, content_policy_fallbacks=litellm.content_policy_fallbacks, drop_params=litellm.drop_params, + disable_token_counter=litellm.disable_token_counter, provider_list=litellm.provider_list, expose_router_debug_in_errors=litellm.expose_router_debug_in_errors, redact_user_api_key_info=litellm.redact_user_api_key_info, @@ -118,6 +122,7 @@ def current_dynamic_config() -> RouterDynamicConfig: ) +_LIVE_VALUE: Final = TypeAdapter(object) _LIVE_SOURCES: Final = MappingProxyType( { "num_retries": "num_retries", @@ -130,6 +135,7 @@ _LIVE_SOURCES: Final = MappingProxyType( "context_window_fallbacks": "context_window_fallbacks", "content_policy_fallbacks": "content_policy_fallbacks", "drop_params": "drop_params", + "disable_token_counter": "disable_token_counter", "provider_list": "provider_list", "expose_router_debug_in_errors": "expose_router_debug_in_errors", "redact_user_api_key_info": "redact_user_api_key_info", @@ -156,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_rust/core.py b/litellm/router_rust/core.py index 40ae7e0e554..3e222313f32 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -1177,6 +1177,7 @@ _HOSTED_STATE: Final = frozenset( "_core_auto_routers", "_core_complexity_routers", "_core_quality_routers", + "_core_switchyard_routers", "_core_adaptive_routers", "_core_routing_plugins", "_core_deployment_latency_map", @@ -1722,6 +1723,7 @@ class RustCore( run_plugins=partial(self._run_plugins, request_kwargs, messages) if self._hosted._core_routing_plugins # pyright: ignore[reportPrivateUsage] # the live plugin list else None, + call_type=None, ) def _core_log_retry( @@ -2715,7 +2717,7 @@ class RustCore( web_search: Final = not plan.sync and self._picks_async() and is_web_search_request(request) # The generic handler `_route_capability_pin` names: a call only some providers serve # drops the others, and a Decisions call carrying media skips the deployments refusing it. - handler: Final = plan.host_op.rsplit(".", 1)[-1] + handler: Final = plan.handler decisions_media: Final = handler in DECISIONS_MEDIA_CALL_TYPES and bool( request.get("audio") or request.get("videos") ) @@ -3113,6 +3115,7 @@ class RustCore( specific_deployment: object = _NO_INPUT, *, pass_through: bool = False, + call_type: str | None = None, ) -> object: """The host's pre-routing of one pick (router_rust.steps.prerouting): the model group to resolve instead of the step's ``model``, a custom strategy's deployment, or ``None``. ``pass_through`` @@ -3166,6 +3169,7 @@ class RustCore( fallback_depth=depth if isinstance(depth, int) else 0, router_strategy=normalize_strategy(r.routing_strategy), pass_through=pass_through, + call_type=call_type, enable_tag_filtering=r.enable_tag_filtering, match_any=r.tag_filtering_match_any, public_router=public if public is not None else self, @@ -5266,7 +5270,9 @@ class _Driver: async def pre_route(self, step: Mapping[str, object]) -> object: request: Final = self._request try: - return await self._core._pre_route(step, request, *self._plan.pick_input(request)) # pyright: ignore[reportPrivateUsage] # see above + return await self._core._pre_route( # pyright: ignore[reportPrivateUsage] # see above + step, request, *self._plan.pick_input(request), call_type=self._plan.handler + ) except Exception as error: # noqa: BLE001 # see select return self._pick_raised(error, logged=True) diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index e2c621ccb7d..7a5b52761e2 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -244,6 +244,11 @@ class Plan: # the request (`{**data, name: value, **kwargs}`). named: tuple[str, ...] = () + @property + def handler(self) -> str: + """The handler name the host op ends in, which route capabilities and pre-routing strategies key on.""" + return self.host_op.rsplit(".", 1)[-1] + @property def on_its_own_loop(self) -> bool: """``function_with_fallbacks`` runs these sync endpoints' whole walk on a private event loop.""" diff --git a/litellm/router_rust/invoke/walk.py b/litellm/router_rust/invoke/walk.py index 17e20119f43..ff8bc7eef5a 100644 --- a/litellm/router_rust/invoke/walk.py +++ b/litellm/router_rust/invoke/walk.py @@ -39,6 +39,7 @@ _WALK_KEYS: Final = ( _SERVED_KEYS: Final = ("model_info", "deployment") # The walk state `run_async_fallback` hands every hop as kwargs. _HOP_KWARGS: Final = ("fallback_depth", "max_fallbacks", "attempted_targets", "include_fallback_errors") +_RETRY_RECORD_KEYS: Final = ("previous_models", "request_retry_count") def _reentrant_hop_stamps(request: Mapping[str, object], field: str) -> Mapping[str, object]: @@ -179,11 +180,13 @@ class WalkMetadata: return sibling: Final = request.get("metadata" if self._field == "litellm_metadata" else "litellm_metadata") if isinstance(sibling, dict): - sibling.pop("attempted_fallbacks", None) # pyright: ignore[reportUnknownMemberType] # the caller's metadata dict - sibling.pop("original_model_group", None) # pyright: ignore[reportUnknownMemberType] # see above + for key in ("attempted_fallbacks", "original_model_group", *_RETRY_RECORD_KEYS): + sibling.pop(key, None) # pyright: ignore[reportUnknownMemberType] # the caller's metadata dict bucket: Final = request.get(self._field) if not isinstance(bucket, dict): return + for key in _RETRY_RECORD_KEYS: + bucket.pop(key, None) # pyright: ignore[reportUnknownMemberType] # see above bucket["attempted_fallbacks"] = 0 model: Final = request.get("model") if model is not None: diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index c96a9c1bbae..cd2d69d13a8 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -44,7 +44,10 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( ) from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck -from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck +from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( + PromptCachingDeploymentCheck, + new_responses_pins_in_flight, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging @@ -130,6 +133,7 @@ class OptionalChecks: redis_cache=cache.redis_cache, in_memory_cache=InMemoryCache(max_size_in_memory=PROMPT_CACHE_PIN_CACHE_MAX_SIZE), ) + self.prompt_caching_pins_in_flight: Final = new_responses_pins_in_flight() self._request_scoped_selectors: tuple[CustomLogger, ...] = () self._request_scoped_policies: tuple[CustomLogger, ...] = () # ``_slack_alerting_loggers``: the Slack alerting logger ``alerting_config`` builds, if any. @@ -167,6 +171,9 @@ class OptionalChecks: group: Final = hosted.get_routing_group(resolved) return group is not None and group.routing_strategy == "priority" + def _group_has_alternatives(self, model: str) -> bool: + return len(self._hosted._core_deployments_named(model) or ()) > 1 # pyright: ignore[reportPrivateUsage] # the host registry's name index + async def repin_session_after_success( self, litellm_params: Mapping[str, object], response: object, cooldowns: "CooldownCache" ) -> None: @@ -291,7 +298,10 @@ class OptionalChecks: callback: CustomLogger | None = None if name == "prompt_caching": callback = PromptCachingDeploymentCheck( - cache=self.prompt_caching_pin_cache, is_priority_group=self._is_priority_group + cache=self.prompt_caching_pin_cache, + is_priority_group=self._is_priority_group, + responses_pins_in_flight=self.prompt_caching_pins_in_flight, + group_has_alternatives=self._group_has_alternatives, ) elif name == "router_budget_limiting": if self.budget_limiter() is not None: @@ -448,16 +458,21 @@ class OptionalChecks: deployments: list[dict[str, object]], # mutable-ok: CustomLogger deployment-list contract messages: object, request: Mapping[str, object], - ) -> list[object]: # mutable-ok: as the filters return it + ) -> tuple[list[object], str | None]: # mutable-ok: as the filters return it """``async_callback_filter_deployments``: each routing callback narrows the deployments in - turn; a filter's error is raised for the Rust core to decide on.""" + turn, then the prompt caching check names the deployment it prefers among the rest, which + the Rust core takes only if its later filters keep it; a filter's error is raised for the + Rust core to decide on.""" remaining: list[object] = list(deployments) # mutable-ok: filter contract # rebind-ok: narrowed base: Final[object] = CustomLogger.async_filter_deployments # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts request_kwargs: Final = dict(request) - for callback in tuple(self.routing_callbacks()): + callbacks: Final = tuple(self.routing_callbacks()) + for callback in callbacks: # The base filter returns its input unchanged: only overriding filters are run. if not isinstance(callback, CustomLogger) or type(callback).async_filter_deployments is base: # pyright: ignore[reportUnknownMemberType] # see above continue + if isinstance(callback, PromptCachingDeploymentCheck): + continue remaining = await callback.async_filter_deployments( # pyright: ignore[reportUnknownMemberType, reportAssignmentType, reportUnknownVariableType] # see above model=model, healthy_deployments=remaining, @@ -465,7 +480,15 @@ class OptionalChecks: request_kwargs=request_kwargs, parent_otel_span=None, ) - return remaining + prompt_caching: Final = next((cb for cb in callbacks if isinstance(cb, PromptCachingDeploymentCheck)), None) + if prompt_caching is None: + return remaining, None + return remaining, await prompt_caching.async_preferred_deployment_id( + model, + remaining, # pyright: ignore[reportArgumentType] # RustCore's deployment dicts + messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + request_kwargs, + ) # -- release -------------------------------------------------------------------------------- @@ -670,15 +693,15 @@ class RustOptionalChecks: async def _answer_callback_filters( self, step: Mapping[str, object], request: Mapping[str, object], messages: object ) -> tuple[object, BaseException | None]: - """A ``callback_filters`` step: the remaining ids, or a filter's error as an attempt - failure with the error itself, which the caller keeps to raise.""" + """A ``callback_filters`` step: the remaining ids and the one a filter prefers, or a + filter's error as an attempt failure with the error itself, which the caller keeps to raise.""" from litellm.router_rust.core import attempt_failure candidates: Final = [str(id) for id in step.get("candidates") or ()] # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType, reportUnknownArgumentType] # the step's id list records: Final = self._records_of(candidates) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # RustCore's deployment dicts narrowed: Final = _request_with_failed_ids(request, step) try: - remaining: Final = await self._opt.filter_deployments( + remaining, preferred_id = await self._opt.filter_deployments( str(step.get("model_group")), records, # pyright: ignore[reportUnknownArgumentType] # see above messages, @@ -689,7 +712,8 @@ class RustOptionalChecks: # A filter's pin lands on the request kwargs the pick goes on to read. if narrowed is not request and narrowed.get(_ENCRYPTED_PIN) and isinstance(request, dict): request[_ENCRYPTED_PIN] = True # rebind-ok: the request dict the pick goes on to read - return filtered_ids(remaining, candidates), None # pyright: ignore[reportUnknownArgumentType] # see above + kept: Final = filtered_ids(remaining, candidates) # pyright: ignore[reportUnknownArgumentType] # see above + return (kept, preferred_id), None __all__ = [ diff --git a/litellm/router_rust/steps/precall.py b/litellm/router_rust/steps/precall.py index 227bb8de66f..c68c5fd5982 100644 --- a/litellm/router_rust/steps/precall.py +++ b/litellm/router_rust/steps/precall.py @@ -13,7 +13,9 @@ deployment records rather than a step, and those need not be in the Rust registr """ from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from functools import partial +from itertools import accumulate, groupby from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol @@ -21,6 +23,7 @@ from litellm._logging import verbose_router_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider +from litellm.litellm_core_utils.prompt_templates.tpm_media_utils import is_plain_embedding_text from litellm.litellm_core_utils.token_counter import ( offload_token_count, served_openai_models, @@ -38,6 +41,17 @@ if TYPE_CHECKING: # The chat params a deployment must take when the request sends them (`_PRE_CALL_CHECK_FILTERED_PARAMS`). _FILTERED_PARAMS: Final = frozenset({"response_format"}) +_CHARS_PER_OFFLOADED_COUNT: Final = 1 << 16 + + +@dataclass(frozen=True, slots=True) +class _TextBatch: + texts: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class _TokenIdBatch: + longest: int class PreCallRegistry(Protocol): @@ -59,35 +73,111 @@ def count_tokens( request_kwargs: Mapping[str, object] | None, tokenizer_model: str, deployments: Sequence[Mapping[str, object]] = (), + fits_within: int = 0, ) -> int: """``_count_pre_call_check_tokens``: messages, or Responses input normalized to messages, plus the tools and Anthropic ``system`` block the message list never carries. Embedding media an item of ``deployments`` reads as media is left out of the input: it never counts as text.""" from litellm.llms.anthropic.messages.native.utils import anthropic_system_to_openai_message - extras: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else MappingProxyType({}) - raw_instructions: Final = extras.get("instructions") - instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None - raw_tools: Final = extras.get("tools") - tools: Final = raw_tools if _is_list(raw_tools) and raw_tools else None + if messages is None and input is not None: + return _input_tokens(host, _text_input(input, deployments), request_kwargs, tokenizer_model, fits_within) + if messages is None: + raise ValueError("Either messages or input must be provided to count tokens") + extras: Final = _extras(request_kwargs) system_message: Final = anthropic_system_to_openai_message(extras.get("system")) counter: Final = host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # the host's overloads name litellm's untyped returns - if messages is not None: - counted_messages: Final = (system_message, *_items(messages)) if system_message is not None else messages - return counter(model=tokenizer_model, messages=counted_messages, tools=tools) - if input is not None: - from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, - ) + counted_messages: Final = (system_message, *_items(messages)) if system_message is not None else messages + return counter(model=tokenizer_model, messages=counted_messages, tools=_tools(extras)) - text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments)) - input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # its input type is untyped - input=text_input, # pyright: ignore[reportArgumentType] # str | list, as the request carries it - responses_api_request={"instructions": instructions} if instructions is not None else {}, - only_since_compaction=True, - ) - return counter(model=tokenizer_model, messages=input_messages, tools=tools) - raise ValueError("Either messages or input must be provided to count tokens") + +def _input_tokens( + host: "PythonHost", + text_input: object, + request_kwargs: Mapping[str, object] | None, + tokenizer_model: str, + fits_within: int, +) -> int: + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + counter: Final = host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # see count_tokens + match _embedding_batch(text_input): + case _TextBatch(texts): + return _longest_text_tokens(counter, tokenizer_model, texts, fits_within) + case _TokenIdBatch(longest): + return _token_id_count(host, longest) + case None: + pass + extras: Final = _extras(request_kwargs) + raw_instructions: Final = extras.get("instructions") + instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None + input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # its input type is untyped + input=text_input, # pyright: ignore[reportArgumentType] # str | list, as the request carries it + responses_api_request={"instructions": instructions} if instructions is not None else {}, + only_since_compaction=True, + ) + return counter(model=tokenizer_model, messages=input_messages, tools=_tools(extras)) + + +def _extras(request_kwargs: Mapping[str, object] | None) -> Mapping[str, object]: + return request_kwargs if request_kwargs is not None else MappingProxyType({}) + + +def _tools(extras: Mapping[str, object]) -> object: + raw_tools: Final = extras.get("tools") + return raw_tools if _is_list(raw_tools) and raw_tools else None + + +def _text_input(input: object, deployments: Sequence[Mapping[str, object]]) -> object: # the request field's name + text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments)) + return text_input + + +def _embedding_batch(text_input: object) -> _TextBatch | _TokenIdBatch | None: + items: Final = _items(text_input) if _is_list(text_input) else () + if not items: + return None + texts: Final = tuple(item for item in items if isinstance(item, str)) + if len(texts) == len(items): + return _TextBatch(texts) + if all(_is_token_id(item) for item in items): + return _TokenIdBatch(len(items)) + if all(_is_token_id_list(item) for item in items): + return _TokenIdBatch(max(len(_items(item)) for item in items)) + return None + + +def _longest_text_tokens( + counter: Callable[..., int], tokenizer_model: str, texts: Sequence[str], fits_within: int +) -> int: + return max(_text_tokens(counter, tokenizer_model, text, fits_within) for text in texts) + + +def _offload_slices(texts: Sequence[str]) -> tuple[tuple[str, ...], ...]: + ends: Final = accumulate(len(text) for text in texts) + slices: Final = groupby( + zip(ends, texts, strict=True), key=lambda end_and_text: (end_and_text[0] - 1) // _CHARS_PER_OFFLOADED_COUNT + ) + return tuple(tuple(text for _, text in texts_in_slice) for _, texts_in_slice in slices) + + +def _text_tokens(counter: Callable[..., int], tokenizer_model: str, text: str, fits_within: int) -> int: + size: Final = len(text.encode(errors="surrogatepass")) + return size if size <= fits_within else counter(model=tokenizer_model, text=text) + + +def _is_token_id(value: object) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +def _is_token_id_list(value: object) -> bool: + return _is_list(value) and all(_is_token_id(token) for token in _items(value)) + + +def _token_id_count(host: "PythonHost", count: int) -> int: + return 0 if host.dynamic().disable_token_counter else count def _max_input_tokens(registry: PreCallRegistry, model: str, deployment: Mapping[str, object]) -> int | None: @@ -246,16 +336,31 @@ class PreCallFacts: def _candidate_records(self) -> Sequence[Mapping[str, object]]: return tuple(self._record(id) for id in self._candidates) + def _fits_within(self) -> int: + return min((limit for limit in self._windows.values() if limit is not None), default=0) + + def _input_is_a_list(self) -> bool: + return self._messages is None and _is_list(self._input) + def _count(self) -> int | None: try: - return count_tokens( - self._host, - self._messages, - self._input, - self._request, - self._tokenizer_model(), - self._candidate_records(), - ) + tokenizer_model: Final = self._tokenizer_model() + if not self._input_is_a_list(): + return count_tokens( + self._host, + self._messages, + self._input, + self._request, + tokenizer_model, + self._candidate_records(), + self._fits_within(), + ) + text_input: Final = _text_input(self._input, self._candidate_records()) + match _embedding_batch(text_input): + case _TextBatch() if not tokenizer_model: + return 0 + case _: + return _input_tokens(self._host, text_input, self._request, tokenizer_model, self._fits_within()) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.error( "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e @@ -264,14 +369,39 @@ class PreCallFacts: async def _count_async(self) -> int | None: try: - return await offload_token_count(count_tokens)( - self._host, - self._messages, - self._input, - self._request, - self._tokenizer_model(), - self._candidate_records(), + tokenizer_model: Final = self._tokenizer_model() + if not self._input_is_a_list(): + return await offload_token_count(count_tokens)( + self._host, + self._messages, + self._input, + self._request, + tokenizer_model, + self._candidate_records(), + self._fits_within(), + ) + text_input: Final = ( + self._input + if all(map(is_plain_embedding_text, _items(self._input))) + else await offload_token_count(_text_input)(self._input, self._candidate_records()) ) + match _embedding_batch(text_input): + case _TextBatch() if not tokenizer_model: + return 0 + case _TextBatch(texts): + counter: Final = self._host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # see count_tokens + fits_within: Final = self._fits_within() + count_slice: Final = offload_token_count(_longest_text_tokens) + return max( + [ + await count_slice(counter, tokenizer_model, texts_in_slice, fits_within) + for texts_in_slice in _offload_slices(texts) + ] + ) + case _: + return await offload_token_count(_input_tokens)( + self._host, text_input, self._request, tokenizer_model, self._fits_within() + ) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.error( "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e diff --git a/litellm/router_rust/steps/prerouting.py b/litellm/router_rust/steps/prerouting.py index 63d15a1a80a..47bd1e69046 100644 --- a/litellm/router_rust/steps/prerouting.py +++ b/litellm/router_rust/steps/prerouting.py @@ -197,9 +197,13 @@ class PreRoutingStrategies: def quality(self) -> StrategyRegistry: return self._hosted._core_quality_routers # pyright: ignore[reportPrivateUsage] # see above + @property + def switchyard(self) -> StrategyRegistry: + return self._hosted._core_switchyard_routers # pyright: ignore[reportPrivateUsage] # see above + def all(self) -> tuple[StrategyRegistry, ...]: """In the order ``_select_pre_routing_strategy`` searches them.""" - return (self.auto, self.complexity, self.adaptive, self.quality) + return (self.switchyard, self.auto, self.complexity, self.adaptive, self.quality) def any_registered(self) -> bool: return any(self.all()) @@ -211,6 +215,7 @@ class PreRoutingStrategies: self._hosted._core_quality_routers = {} # pyright: ignore[reportPrivateUsage] # see above self._hosted._core_complexity_routers = {} # pyright: ignore[reportPrivateUsage] # see above self._hosted._core_auto_routers = {} # pyright: ignore[reportPrivateUsage] # see above + self._hosted._core_switchyard_routers = {} # pyright: ignore[reportPrivateUsage] # see above def register(self, deployment: Deployment) -> None: """The end of ``_add_deployment``: an auto, complexity or quality marker registers its router @@ -222,6 +227,8 @@ class PreRoutingStrategies: self._init_complexity_router(deployment) case "quality": self._init_quality_router(deployment) + case "switchyard": + self._init_switchyard_router(deployment) case _: pass @@ -231,7 +238,7 @@ class PreRoutingStrategies: return model_name: Final = deployment.model_name tags: Final = _deployment_tags(deployment) - for registry in (self.auto, self.complexity, self.quality): + for registry in (self.auto, self.complexity, self.quality, self.switchyard): _ = _unregister(registry, model_name, tags) if _unregister(self.adaptive, model_name, tags): self._sync_adaptive_router_hooks() @@ -463,6 +470,32 @@ class PreRoutingStrategies: ) self._register(self.quality, deployment, router, "Quality-router") + def _init_switchyard_router(self, deployment: Deployment) -> None: + from litellm.router_strategy.switchyard_router.config import SwitchyardRouterConfig, build_algorithm + from litellm.router_strategy.switchyard_router.switchyard_router import SwitchyardRouter + from litellm.types.decisions import OpenAIDecisionResponse + from litellm.types.utils import ModelResponse + + config: Final = SwitchyardRouterConfig.model_validate(deployment.litellm_params.switchyard_router_config) + if deployment.model_name in (config.models.capable, config.models.efficient, config.models.judge): + raise ValueError("Switchyard cannot route to itself") + ctx: Final = self._ctx + + async def complete_judge(request: Mapping[str, object]) -> ModelResponse: + result: Final = await ctx.acompletion(**request) # pyright: ignore[reportAttributeAccessIssue] # the core's RoutingContext + if not isinstance(result, ModelResponse): + raise TypeError("Switchyard judge must return a non-streaming completion") + return result + + async def decide(request: Mapping[str, object]) -> OpenAIDecisionResponse: + result: Final = await ctx.adecisions(**request) # pyright: ignore[reportAttributeAccessIssue] # see above + if not isinstance(result, OpenAIDecisionResponse): + raise TypeError("Switchyard decision judge must return an OpenAI-shaped decision response") + return result + + router: Final = SwitchyardRouter(deployment.model_name, config, build_algorithm(config), complete_judge, decide) + self._register(self.switchyard, deployment, router, "Switchyard") + def _init_adaptive_router(self, deployment: Deployment) -> None: from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook @@ -741,6 +774,7 @@ class PreRouter: public_router: object, run_plugins: Callable[[str], Awaitable[None]] | None = None, pass_through: bool = False, + call_type: str | None = None, ) -> str | None: """The start of ``_async_get_available_deployment``: run the hook, record its pick and pin the tier's litellm_params. The model group to resolve, or ``None`` to resolve ``model`` itself. @@ -759,6 +793,7 @@ class PreRouter: match_any=match_any, public_router=public_router, run_plugins=run_plugins, + call_type=call_type, ) routed: Final = response.model if response is not None else model if response is not None: @@ -782,6 +817,7 @@ class PreRouter: match_any: bool, public_router: object, run_plugins: Callable[[str], Awaitable[None]] | None, + call_type: str | None, ) -> "PreRoutingHookResponse | None": _restore_client_output_ceiling(request_kwargs) _record_routing_decision(request_kwargs, None) @@ -861,6 +897,7 @@ class PreRouter: messages=routing_messages if routing_messages is not None else messages, # pyright: ignore[reportArgumentType] # the request's messages, as given input=input, # pyright: ignore[reportArgumentType] # the request's input, as given specific_deployment=False, + call_type=call_type, ) hooked: Final[PreRoutingHookResponse | None] = routed traced: Final = ( diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py index 7c3abdcd45e..820db52d5a1 100644 --- a/litellm/router_rust/surface/registry.py +++ b/litellm/router_rust/surface/registry.py @@ -658,6 +658,7 @@ class HostRegistry(HostModelInfo): self._core_complexity_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same self._core_adaptive_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same self._core_quality_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same + self._core_switchyard_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same self.strategies: PreRoutingStrategies | None = None # The router's optional pre-call checks (deployment budgets, prompt caching pins); # RustCore installs its ``OptionalChecks`` (litellm/router_rust/steps/callbacks.py). diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index c3a65d6aa91..43869fdf832 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -184,6 +184,7 @@ class AdaptiveRouter: messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, + call_type: str | None = None, ) -> PreRoutingHookResponse | None: """ Plugin entry point invoked by `Router.async_pre_routing_hook` when the 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 b92c53fa93c..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 @@ -204,6 +209,7 @@ class AutoRouter(CustomLogger): messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, + call_type: str | None = None, ) -> Optional["PreRoutingHookResponse"]: """ This hook is called before the routing decision is made. diff --git a/litellm/router_strategy/auto_router/auto_router_model_naming.py b/litellm/router_strategy/auto_router/auto_router_model_naming.py index f5aeccd9305..47daa175c99 100644 --- a/litellm/router_strategy/auto_router/auto_router_model_naming.py +++ b/litellm/router_strategy/auto_router/auto_router_model_naming.py @@ -26,7 +26,7 @@ from litellm.router_strategy.complexity_router.config import ( AUTO_ROUTER_MODEL_PREFIX: Final = "auto_router/" -StrategyRouterKind = Literal["semantic", "complexity", "adaptive", "quality"] +StrategyRouterKind: TypeAlias = Literal["semantic", "complexity", "adaptive", "quality", "switchyard"] StrategyRouterDependencyRole: TypeAlias = Literal[ "tier", "default", "classifier", "embedding", "evaluation", "compactor" @@ -43,6 +43,7 @@ class StrategyRouterDependency: STRATEGY_ROUTER_PARAM_FIELDS: Final[frozenset[str]] = frozenset( { + "switchyard_router_config", "auto_router_config", "auto_router_config_path", "auto_router_default_model", @@ -57,6 +58,7 @@ STRATEGY_ROUTER_PARAM_FIELDS: Final[frozenset[str]] = frozenset( ) _REQUIRED_FIELD_GROUPS: Final[Mapping[StrategyRouterKind, tuple[tuple[str, ...], ...]]] = { + "switchyard": (("switchyard_router_config",),), "semantic": ( ("auto_router_config", "auto_router_config_path"), ("auto_router_default_model",), @@ -78,6 +80,8 @@ def classify_strategy_router_model(model: str) -> StrategyRouterKind | None: if not model.startswith(AUTO_ROUTER_MODEL_PREFIX): return None remainder: Final = model[len(AUTO_ROUTER_MODEL_PREFIX) :] + if remainder.startswith("switchyard"): + return "switchyard" if remainder.startswith("complexity_router"): return "complexity" if remainder.startswith("adaptive_router"): @@ -137,6 +141,13 @@ def strategy_router_dependencies( kind: Final = classify_strategy_router_model(str(litellm_params.get("model", ""))) if kind is None: return () + if kind == "switchyard": + models: Final = _mapping(_mapping(litellm_params.get("switchyard_router_config")).get("models")) + return ( + _named(models.get("capable"), "tier") + + _named(models.get("efficient"), "tier") + + _named(models.get("judge"), "classifier") + ) if kind == "semantic": return _named(litellm_params.get("auto_router_default_model"), "default") + _named( litellm_params.get("auto_router_embedding_model"), "embedding" 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 05dd4706b62..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: @@ -4490,6 +4501,7 @@ class ComplexityRouter(CustomLogger): messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, + call_type: str | None = None, ) -> PreRoutingHookResponse | None: """ Pre-routing hook called before the routing decision. diff --git a/litellm/router_strategy/quality_router/quality_router.py b/litellm/router_strategy/quality_router/quality_router.py index 2d9ad63e9de..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. @@ -330,6 +330,7 @@ class QualityRouter(CustomLogger): messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, + call_type: str | None = None, ) -> Optional["PreRoutingHookResponse"]: """Try keyword override first; fall back to complexity-tier routing.""" from litellm.types.router import PreRoutingHookResponse diff --git a/litellm/router_strategy/switchyard_router/__init__.py b/litellm/router_strategy/switchyard_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/router_strategy/switchyard_router/config.py b/litellm/router_strategy/switchyard_router/config.py new file mode 100644 index 00000000000..7e0d3ad3a94 --- /dev/null +++ b/litellm/router_strategy/switchyard_router/config.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, TypeAdapter, model_validator + +if TYPE_CHECKING: + from switchyard_rust.libsy import Algorithm, TaskClassifierConfig + + +class LlmClassifierConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + kind: Literal["llm"] = "llm" + base_threshold: float = Field(default=0.5, ge=0, le=1) + threshold_step: float = Field(default=0.0, ge=0, le=1) + recent_turn_window: int | None = Field(default=None, gt=0) + max_output_tokens: int = Field(default=4096, gt=0) + response_format_type: Literal["json_schema", "json_object"] = "json_schema" + prompt: str | None = None + + +class DecisionClassifierConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + kind: Literal["decision"] + cutoff: float = Field(default=0.6, ge=0, le=1) + instructions: str | Mapping[str, object] | None = None + evidence: Mapping[str, object] | None = None + + +_KIND: Final = TypeAdapter(Mapping[str, object]) + + +def _classifier_kind(value: object) -> str: + match value: + case LlmClassifierConfig() | DecisionClassifierConfig(): + return value.kind + case Mapping(): + return str(_KIND.validate_python(value).get("kind", "llm")) + case _: + return "llm" + + +ClassifierConfig: TypeAlias = Annotated[ + Annotated[LlmClassifierConfig, Tag("llm")] | Annotated[DecisionClassifierConfig, Tag("decision")], + Discriminator(_classifier_kind), +] + + +class StageOptions(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + picker: Literal["efficient_first", "capable_first"] = "efficient_first" + confidence_threshold: float | None = Field(default=None, ge=0, le=1) + recent_window: int | None = Field(default=None, gt=0) + tool_semantics: Mapping[str, tuple[str, ...]] | None = None + only_on_wrong_signal_escalation: bool = True + + +class StageModels(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", str_strip_whitespace=True) + + capable: str = Field(min_length=1) + efficient: str = Field(min_length=1) + judge: str | None = Field(default=None, min_length=1) + + @model_validator(mode="after") + def reject_recursive_targets(self) -> StageModels: + if any(name.startswith("auto_router/") for name in (self.capable, self.efficient, self.judge or "")): + raise ValueError("Switchyard targets must be concrete model groups, not auto-router markers") + if self.capable == self.efficient: + raise ValueError("The capable and efficient models must be different model groups") + return self + + +class SwitchyardRouterConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + algorithm: Literal["stage_router"] = "stage_router" + models: StageModels + options: StageOptions = Field(default_factory=StageOptions) + classifier: ClassifierConfig | None = None + + @model_validator(mode="after") + def validate_judge(self) -> SwitchyardRouterConfig: + if (self.classifier is None) != (self.models.judge is None): + raise ValueError("A judge model is required exactly when the classifier is enabled") + return self + + @property + def confidence_threshold(self) -> float: + if self.options.confidence_threshold is not None: + return self.options.confidence_threshold + return 0.7 if self.classifier is not None else 0.5 + + def runtime_models(self) -> Mapping[str, tuple[str, ...]]: + return { + "any": (self.models.capable, self.models.efficient), + "capable": (self.models.capable,), + "efficient": (self.models.efficient,), + **({"judge": (self.models.judge,)} if self.models.judge is not None else {}), + } + + +def _task_classifier(config: SwitchyardRouterConfig) -> TaskClassifierConfig | None: + from switchyard_rust.libsy import TaskClassifierConfig + + match config.classifier: + case None: + return None + case LlmClassifierConfig() as llm: + return TaskClassifierConfig( + base_threshold=llm.base_threshold, + threshold_step=llm.threshold_step, + recent_turn_window=llm.recent_turn_window, + max_output_tokens=llm.max_output_tokens, + response_format_type=llm.response_format_type, + prompt=llm.prompt, + ) + case DecisionClassifierConfig() as decision: + return TaskClassifierConfig.decision( + cutoff=decision.cutoff, + candidates={"capable": config.models.capable, "efficient": config.models.efficient}, + evidence=dict(decision.evidence or {}), + instructions=decision.instructions, + ) + + +def build_algorithm(config: SwitchyardRouterConfig) -> Algorithm: + from switchyard_rust.libsy import LlmFallback, stage_router + + classifier: Final = _task_classifier(config) + return stage_router( + picker=config.options.picker, + confidence_threshold=config.confidence_threshold, + recent_window=config.options.recent_window, + tool_semantics=dict(config.options.tool_semantics) if config.options.tool_semantics is not None else None, + only_on_wrong_signal_escalation=config.options.only_on_wrong_signal_escalation, + classifier=LlmFallback(config=classifier) if classifier is not None else None, + ) + + +def validate_switchyard_config_write(config: Mapping[str, object] | None) -> str | None: + if config is None: + return None + try: + build_algorithm(SwitchyardRouterConfig.model_validate(config)) + except ImportError: + return "Switchyard requires the optional dependency. Install litellm[switchyard] (nemo-switchyard==0.4.0a1)." + except (ValueError, TypeError) as error: + return f"Invalid switchyard_router_config: {error}" + return None diff --git a/litellm/router_strategy/switchyard_router/switchyard_router.py b/litellm/router_strategy/switchyard_router/switchyard_router.py new file mode 100644 index 00000000000..cb455bab337 --- /dev/null +++ b/litellm/router_strategy/switchyard_router/switchyard_router.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +import hashlib +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Final, Literal, TypeAlias, TypeVar + +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs +from litellm.litellm_core_utils.internal_call_metadata import ( + effective_turn_off_message_logging, + forwarded_internal_call_metadata, + parent_session_kwargs, +) +from litellm.router_strategy.switchyard_router.config import ( + DecisionClassifierConfig, + LlmClassifierConfig, + SwitchyardRouterConfig, +) +from litellm.router_strategy.switchyard_router.translation import ( + chat_messages, + decision_request, + decision_response, + judge_request, + judge_response, + mapping, + request_to_switchyard, +) +from litellm.types.decisions import OpenAIDecisionResponse +from litellm.types.router import PreRoutingHookResponse +from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN, ModelResponse, StandardLoggingRoutingDecision + +if TYPE_CHECKING: + from switchyard_rust.libsy import Algorithm, DecisionCall, ModelCall, RoutingOutcome + +JudgeCompletion: TypeAlias = Callable[[Mapping[str, object]], Awaitable[ModelResponse]] +DecisionCompletion: TypeAlias = Callable[[Mapping[str, object]], Awaitable[OpenAIDecisionResponse]] +FallbackReason: TypeAlias = Literal[ + "unsupported_request_type", + "unsupported_content", + "libsy_error", + "unexpected_judge_call", + "unexpected_decision_call", + "target_outside_config", + "no_decision", +] +_T = TypeVar("_T") +_SESSION_HEADER: Final = "x-switchyard-session-id" +_FORWARDED_HEADER_PREFIXES: Final = ("x-switchyard-", "x-codex-", "x-nemo-relay-", "x-dynamo-") +_FORWARDED_HEADERS: Final = frozenset( + {"x-session-id", "session-id", "thread-id", "x-task-id", "x-request-id", "x-client-request-id"} +) + + +async def _resolved(value: _T) -> _T: + return value + + +_SESSION_ALIASES: Final = frozenset( + {_SESSION_HEADER, "x-claude-code-session-id", "x-nemo-relay-session-id", "x-session-id", "session-id"} +) +_CODEX_TURN_METADATA_HEADER: Final = "x-codex-turn-metadata" + + +def _forwarded_header(name: str, session_scoped: bool) -> bool: + lowered: Final = name.lower() + if lowered in _SESSION_ALIASES or (lowered == _CODEX_TURN_METADATA_HEADER and not session_scoped): + return False + return lowered in _FORWARDED_HEADERS or lowered.startswith(_FORWARDED_HEADER_PREFIXES) + + +def _caller_key_hash(request_kwargs: Mapping[str, object]) -> str: + internal: Final = mapping(request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs)) or {}) + key: Final = internal.get("user_api_key_hash") + return str(key) if key is not None else "" + + +def _scoped_session(request_kwargs: Mapping[str, object]) -> str | None: + session: Final = parent_session_kwargs(request_kwargs).get("litellm_session_id") + if not session: + return None + return hashlib.sha256(f"{_caller_key_hash(request_kwargs)}\0{session}".encode()).hexdigest() + + +def switchyard_headers(request_kwargs: Mapping[str, object]) -> Mapping[str, str]: + proxy_request: Final = mapping(request_kwargs.get("proxy_server_request") or {}) + client: Final = mapping(proxy_request.get("headers") or {}) + session: Final = _scoped_session(request_kwargs) + forwarded: Final = { + name.lower(): value + for name, value in client.items() + if isinstance(value, str) and _forwarded_header(name, session is not None) + } + return {**forwarded, **({_SESSION_HEADER: session} if session else {})} + + +class SwitchyardRouter: + def __init__( + self, + model_name: str, + config: SwitchyardRouterConfig, + algorithm: Algorithm, + completion: JudgeCompletion, + decide: DecisionCompletion, + ) -> None: + self.model_name = model_name + self.config = config + self.algorithm = algorithm + self.completion = completion + self.decide = decide + + def _fallback( + self, reason: FallbackReason, messages: Sequence[Mapping[str, object]] | None + ) -> PreRoutingHookResponse: + return PreRoutingHookResponse( + model=self.config.models.capable, + messages=[dict(message) for message in messages] if messages is not None else None, + routing_decision={ + "router_type": "switchyard", + "router_model_name": self.model_name, + "routed_model": self.config.models.capable, + "switchyard_algorithm": self.config.algorithm, + "switchyard_fallback_reason": reason, + }, + ) + + def _internal_call( + self, + request: Mapping[str, object], + caller: Mapping[str, object], + bucket: Literal["metadata", "litellm_metadata"], + ) -> Mapping[str, object]: + caller_metadata: Final = mapping(caller.get(get_metadata_variable_name_from_kwargs(caller)) or {}) + return { + **request, + bucket: { + **forwarded_internal_call_metadata(caller_metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN), + INTERNAL_CALL_ORIGIN_METADATA_KEY: AUTOROUTER_CLASSIFIER_CALL_ORIGIN, + }, + "proxy_server_request": {"body": request}, + "num_retries": 0, + "turn_off_message_logging": effective_turn_off_message_logging(caller), + **parent_session_kwargs(caller), + } + + def _decision( + self, pick: str, evidence: object, messages: Sequence[Mapping[str, object]] | None + ) -> PreRoutingHookResponse: + decision: Final[StandardLoggingRoutingDecision] = { + "router_type": "switchyard", + "router_model_name": self.model_name, + "routed_model": pick, + "switchyard_algorithm": self.config.algorithm, + "switchyard_evidence": evidence, + } + return PreRoutingHookResponse( + model=pick, + messages=[dict(message) for message in messages] if messages is not None else None, + routing_decision=decision, + ) + + async def _answer_judge(self, call: ModelCall, caller: Mapping[str, object]) -> None: + try: + converted: Final = judge_response( + await self.completion( + self._internal_call(judge_request(call.request, call.models[0]), caller, "metadata") + ) + ) + except Exception as error: # noqa: BLE001 # judge/provider failures are inputs to Switchyard's fall-open decision + await call.fail(error) + else: + await call.respond(_resolved(converted)) + + async def _answer_decision(self, call: DecisionCall, caller: Mapping[str, object]) -> None: + try: + answer: Final = decision_response( + await self.decide( + self._internal_call(decision_request(call.request, call.model), caller, "litellm_metadata") + ), + call.request, + ) + except Exception as error: # noqa: BLE001 # decision/provider failures are inputs to Switchyard's fall-open decision + await call.fail(error) + else: + await call.respond(_resolved(answer)) + + def _outcome( + self, outcome: RoutingOutcome, messages: Sequence[Mapping[str, object]] | None + ) -> PreRoutingHookResponse: + pick: Final = outcome.selected_model_ids[0] + if pick not in (self.config.models.capable, self.config.models.efficient): + return self._fallback("target_outside_config", messages) + metadata: Final = outcome.metadata + evidence: Final[object] = metadata.evidence if metadata is not None else None + return self._decision(pick, evidence, messages) + + async def async_pre_routing_hook( + self, + model: str, + request_kwargs: Mapping[str, object], + messages: Sequence[Mapping[str, object]] | None = None, + input: object = None, + specific_deployment: bool | None = False, + call_type: str | None = None, + ) -> PreRoutingHookResponse | None: + from switchyard_rust.libsy import LibsyError, Step + + if model != self.model_name: + return None + try: + chat: Final = chat_messages(call_type, model, messages, input, request_kwargs) + request: Final = request_to_switchyard(chat) if chat is not None else None + except (ValueError, KeyError, TypeError): + return self._fallback("unsupported_content", messages) + if request is None: + return self._fallback("unsupported_request_type", messages) + judge: Final = self.config.models.judge + try: + async for step in self.algorithm.run_stream( + request, self.config.runtime_models(), headers=switchyard_headers(request_kwargs) + ): + match step: + case Step.CallModel(call) if isinstance( + self.config.classifier, LlmClassifierConfig + ) and call.models[:1] == [judge]: + await self._answer_judge(call, request_kwargs) + case Step.CallDecision(call) if ( + isinstance(self.config.classifier, DecisionClassifierConfig) and call.model == judge + ): + await self._answer_decision(call, request_kwargs) + case Step.CallModel(): + return self._fallback("unexpected_judge_call", messages) + case Step.CallDecision(): + return self._fallback("unexpected_decision_call", messages) + case Step.Done(outcome): + return self._outcome(outcome, messages) + except LibsyError: + return self._fallback("libsy_error", messages) + return self._fallback("no_decision", messages) diff --git a/litellm/router_strategy/switchyard_router/translation.py b/litellm/router_strategy/switchyard_router/translation.py new file mode 100644 index 00000000000..246f2ba3e41 --- /dev/null +++ b/litellm/router_strategy/switchyard_router/translation.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +import json +from collections.abc import Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Final, Protocol + +from pydantic import TypeAdapter + +from litellm.types.decisions import ( + OpenAIChoiceAnswer, + OpenAIDecisionResponse, + systemone_choice_key, +) +from litellm.types.utils import CallTypes, ModelResponse + +if TYPE_CHECKING: + from switchyard_rust.libsy import LlmResponse + +_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_SEQUENCE: Final = TypeAdapter(tuple[object, ...]) +_CALL_TYPES: Final = frozenset(call_type.value for call_type in CallTypes) + + +def mapping(value: object) -> Mapping[str, object]: + return _MAPPING.validate_python(value) + + +def sequence(value: object) -> tuple[object, ...]: + return _SEQUENCE.validate_python(value) + + +def text_content(value: object) -> tuple[Mapping[str, object], ...]: + if value is None: + return () + if isinstance(value, str): + return ({"type": "text", "text": value},) + blocks: Final = tuple(mapping(block) for block in sequence(value)) + if any(block.get("type") != "text" or not isinstance(block.get("text"), str) for block in blocks): + raise ValueError("Switchyard routing supports text content only; media and audio are unsupported") + return tuple({"type": "text", "text": block["text"]} for block in blocks) + + +def tool_call(value: object) -> Mapping[str, object]: + call: Final = mapping(value) + function: Final = mapping(call.get("function")) + arguments: Final = function.get("arguments") + if call.get("type") != "function" or not isinstance(arguments, str): + raise ValueError("Switchyard requires function calls with JSON arguments") + return { + "type": "tool_call", + "id": call["id"], + "name": function["name"], + "arguments": _MAPPING.validate_json(arguments), + } + + +def message_to_switchyard(value: object) -> Mapping[str, object]: + message: Final = mapping(value) + if message.get("audio") is not None: + raise ValueError("Switchyard routing does not support audio") + role: Final = message.get("role") + content: Final = text_content(message.get("content")) + if role == "tool": + return { + "role": "tool", + "content": ( + { + "type": "tool_result", + "tool_call_id": message["tool_call_id"], + "content": content, + "is_error": message.get("is_error"), + }, + ), + } + if role not in ("user", "system", "developer", "assistant"): + raise ValueError("Switchyard received an unsupported message role") + calls: Final = tuple(tool_call(call) for call in sequence(message.get("tool_calls") or ())) + return {"role": role, "content": (*content, *calls)} + + +class ChatConverter(Protocol): + def get_structured_messages( + self, + data: dict[str, object], # mutable-ok: matches the request converters' existing dict parameter + *, + keep_tool_errors: bool = False, + ) -> Sequence[object] | None: ... + + +def _instantiate(converter: Callable[[], ChatConverter]) -> ChatConverter: + return converter() + + +def chat_messages( + call_type: str | None, + model: str, + messages: Sequence[Mapping[str, object]] | None, + input: object, + request_kwargs: Mapping[str, object], +) -> tuple[Mapping[str, object], ...] | None: + from litellm.llms import load_guardrail_translation_mappings + + converter: Final = ( + load_guardrail_translation_mappings().get(CallTypes(call_type)) if call_type in _CALL_TYPES else None + ) + if converter is None: + return None + request: Final = { + **request_kwargs, + "model": model, + **({"messages": [dict(message) for message in messages]} if messages is not None else {}), + **({"input": input} if input is not None else {}), + } + return tuple(map(mapping, _instantiate(converter).get_structured_messages(request, keep_tool_errors=True) or ())) + + +def request_to_switchyard(messages: Sequence[Mapping[str, object]]) -> Mapping[str, object]: + if not messages: + raise ValueError("Switchyard requires at least one message") + return { + "model": "auto", + "messages": tuple(message_to_switchyard(message) for message in messages), + "stream": False, + } + + +def judge_message(value: object) -> Mapping[str, object]: + message: Final = mapping(value) + blocks: Final = text_content(message.get("content")) + return {"role": message["role"], "content": "".join(str(block["text"]) for block in blocks)} + + +def judge_request(request: Mapping[str, object], model: str) -> Mapping[str, object]: + output: Final = mapping(request.get("output") or {}) + sampling: Final = mapping(request.get("sampling") or {}) + options: Final = { + "max_tokens": output.get("max_output_tokens"), + "response_format": output.get("response_format"), + "temperature": sampling.get("temperature"), + "top_p": sampling.get("top_p"), + } + return { + "model": model, + "stream": False, + "messages": [ + judge_message(message) + for message in ( + *sequence(request.get("instructions") or ()), + *sequence(request.get("messages") or ()), + ) + ], + **{key: value for key, value in options.items() if value is not None}, + } + + +def judge_response(response: ModelResponse) -> LlmResponse.Agg: + from switchyard_rust.libsy import LlmResponse + + body: Final = _MAPPING.validate_json(response.model_dump_json(exclude_unset=True)) + choices: Final = tuple(mapping(choice) for choice in sequence(body["choices"])) + return LlmResponse.Agg( + { + "id": response.id, + "model": response.model, + "outputs": tuple(message_to_switchyard(choice["message"]) for choice in choices), + } + ) + + +def _decision_text(value: object) -> str: + return value if isinstance(value, str) else json.dumps(value) + + +def _decision_question(name: str, question: Mapping[str, object]) -> Mapping[str, object]: + kind: Final = mapping(question["kind"]) + if kind.get("type") != "choice": + raise ValueError("Switchyard decision routing asks choice questions only") + return { + "type": "choice", + "name": name, + "instructions": _decision_text(question.get("instructions")), + "choices": [ + { + "value": option["id"], + **({"description": _decision_text(option["description"])} if option.get("description") else {}), + } + for option in map(mapping, sequence(mapping(kind["data"])["options"])) + ], + } + + +def decision_request(request: Mapping[str, object], model: str) -> Mapping[str, object]: + questions: Final = mapping(request["questions"]) + return { + "model": model, + "input": _decision_text(request["context"]), + "questions": [_decision_question(name, mapping(question)) for name, question in questions.items()], + } + + +def _choice_value(answer: object) -> Mapping[str, object]: + if not isinstance(answer, OpenAIChoiceAnswer): + raise TypeError("The decision model refused or returned a non-choice answer") + return { + "type": "choice", + "data": { + "selected": systemone_choice_key(answer.choice), + "probabilities": {systemone_choice_key(p.value): p.probability for p in answer.probabilities}, + }, + } + + +def decision_response(response: OpenAIDecisionResponse, request: Mapping[str, object]) -> Mapping[str, object]: + names: Final = tuple(mapping(request["questions"])) + if len(response.answers) != len(names): + raise ValueError("The decision model did not answer every question") + return { + "id": None, + "model": response.model, + "answers": { + answer.name or name: { + "value": _choice_value(answer), + "provider_confidence": getattr(answer, "confidence", None), + } + for name, answer in zip(names, response.answers, strict=True) + }, + "usage": { + "input_tokens": response.usage.input_tokens, + "output_tokens": response.usage.output_tokens, + "total_tokens": response.usage.total_tokens, + }, + } 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/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index e5d944b3aef..2f0cbc0d474 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -4,23 +4,37 @@ Check if prompt caching is valid for a given deployment Route to previously cached model id, if valid """ -from collections.abc import Callable, Mapping, Sequence -from typing import Final, TypedDict, cast +import secrets +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final, NamedTuple, Protocol, TypedDict, cast +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly from litellm import verbose_logger from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT from litellm.integrations.anthropic_cache_control_hook import ( AllToolParamValues, AnthropicCacheControlHook, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.routing_common_utils import deployment_ids_failed_in_request from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam +from litellm.responses.litellm_completion_transformation.transformation import ( + RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM, + RESPONSES_WEB_SEARCH_TOOL_TYPES, + LiteLLMCompletionResponsesConfig, + ResponsesToolChatForm, + ResponseTools, +) +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam, ResponseInputParam from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import ( get_base_model_from_metadata, @@ -28,7 +42,7 @@ from litellm.utils import ( is_prompt_caching_valid_prompt, ) -from ..prompt_caching_cache import PromptCachingCache +from ..prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache class _PromptCachingSuccessEventKwargs(TypedDict, total=False): @@ -36,7 +50,165 @@ class _PromptCachingSuccessEventKwargs(TypedDict, total=False): litellm_params: ReadOnly[Mapping[str, object] | None] -def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int: +class _ResponsesPin(NamedTuple): + cache_key: str + marked_ttl_seconds: int + cacheable_deployment_ids: tuple[str, ...] + + +class ResponsesPinsInFlight(Protocol): + def get_cache(self, key: str) -> object: ... + + def set_cache(self, key: str, value: object) -> None: ... + + def delete_cache(self, key: str) -> None: ... + + +class _ResponsesPrompt(NamedTuple): + messages: list[AllMessageValues] # mutable-ok: the prefix helpers take chat messages as a list + tools: list[AllToolParamValues] | None # mutable-ok: the prefix helpers take chat tools as a list + key_tools: list[ChatCompletionToolParam] | None # mutable-ok: the prefix helpers take chat tools as a list + + +_RESPONSES_PIN_METADATA_KEY: Final = "responses_prompt_cache_pin" +_RESPONSES_PINS_IN_FLIGHT_LIMIT: Final = 10_000 +_RESPONSES_PIN_LIFETIME_SECONDS: Final = 3600 +_RETRY_RECORDS: Final = TypeAdapter(tuple[object, ...]) +_RESPONSES_CALL_TYPES: Final = frozenset((CallTypes.responses.value, CallTypes.aresponses.value)) +_RESPONSES_TOOL_TYPES_KEPT_AS_SENT: Final = RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM | RESPONSES_WEB_SEARCH_TOOL_TYPES + + +def new_responses_pins_in_flight() -> ResponsesPinsInFlight: + return InMemoryCache( + max_size_in_memory=_RESPONSES_PINS_IN_FLIGHT_LIMIT, default_ttl=_RESPONSES_PIN_LIFETIME_SECONDS + ) + + +def _key_chat_tool(source: Mapping[str, object], chat_tool: Mapping[str, object]) -> Mapping[str, object]: + function: Final = as_str_mapping(chat_tool.get("function")) + if source.get("type") != "function" or function is None: + return chat_tool + return {**chat_tool, "function": {key: source[key] for key in function if key in source}} + + +def _key_chat_tools( + sources: Sequence[Mapping[str, object]], forms: Sequence[ResponsesToolChatForm] +) -> Iterator[Mapping[str, object]]: + for source, form in zip(sources, forms, strict=True): + yield from (_key_chat_tool(source, chat_tool) for chat_tool in form.chat_tools) + + +def _responses_prompt(request_data: Mapping[str, object]) -> _ResponsesPrompt | None: + request_input: Final = request_data.get("input") + request_tools: Final = request_data.get("tools") + if request_data.get("previous_response_id") is not None or not isinstance(request_input, str | list): + return None + listed_tools: Final = ( + cast(ResponseTools, request_tools) # cast-ok: a bad tool raises inside the guard below and is caught + if isinstance(request_tools, list) + else None + ) + try: + tools_kept_as_sent: Final = tuple( + tool for tool in listed_tools or () if tool.get("type") in _RESPONSES_TOOL_TYPES_KEPT_AS_SENT + ) + convertible_tools: Final = ( + [tool for tool in listed_tools if tool.get("type") not in _RESPONSES_TOOL_TYPES_KEPT_AS_SENT] + if listed_tools is not None + else None + ) + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # its input type is untyped + input=cast(str | ResponseInputParam, request_input), # cast-ok: a bad item raises and is caught here + responses_api_request={"instructions": request_data.get("instructions")}, + item_cache_control=True, + only_since_compaction=True, + ) + forms: Final = LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(convertible_tools) + except Exception as error: # noqa: BLE001 # the provider judges the input, a routing preference never does + verbose_logger.debug("Prompt Caching: skipping an input with no chat form: %s", type(error).__name__) + return None + tools: Final = list(chain.from_iterable(form.chat_tools for form in forms)) + key_tools: Final = [*_key_chat_tools(convertible_tools or (), forms), *tools_kept_as_sent] + return _ResponsesPrompt( + messages=cast(list[AllMessageValues], messages), # cast-ok: the bridge's chat messages, read as plain JSON + tools=cast(list[AllToolParamValues], tools) or None, # cast-ok: the bridge's chat tools, read as plain JSON + key_tools=cast(list[ChatCompletionToolParam], key_tools) or None, # cast-ok: tools are hashed as plain JSON + ) + + +def _retry_records(metadata: Mapping[str, object]) -> tuple[object, ...]: + try: + return _RETRY_RECORDS.validate_python(metadata.get("previous_models", ())) + except ValidationError: + return () + + +def _record_deployment_ids(records: tuple[object, ...]) -> Iterator[str]: + deployment_ids: Final = ((as_str_mapping(record) or {}).get("deployment_id") for record in records) + return (deployment_id for deployment_id in deployment_ids if isinstance(deployment_id, str)) + + +def _retried_deployment_ids(metadata: Mapping[str, object]) -> frozenset[str] | None: + records: Final = _retry_records(metadata) + retry_count: Final = metadata.get("request_retry_count") + if type(retry_count) is int and retry_count > len(records): + return None + return frozenset(_record_deployment_ids(records)) + + +def _deployment_ids_tried_in_request(request_data: Mapping[str, object]) -> frozenset[str] | None: + retried: Final = tuple( + _retried_deployment_ids(as_str_mapping(request_data.get(bucket)) or {}) + for bucket in ("litellm_metadata", "metadata") + ) + known: Final = tuple(ids for ids in retried if ids is not None) + if len(known) < len(retried): + return None + return deployment_ids_failed_in_request(request_data).union(*known) + + +def _previous_response_model_id(request_data: Mapping[str, object]) -> str | None: + previous_response_id: Final = request_data.get("previous_response_id") + model_id: Final = ( + ResponsesAPIRequestUtils.get_model_id_from_response_id(previous_response_id) + if isinstance(previous_response_id, str) + else None + ) + tried: Final = _deployment_ids_tried_in_request(request_data) if model_id is not None else None + return model_id if model_id is not None and tried is not None and model_id not in tried else None + + +def _deployment_min_tokens(deployment: Mapping[str, Mapping[str, object]], fallback: int) -> int: + deployment_model: Final = deployment.get("litellm_params", {}).get("model") + return get_prompt_cache_min_tokens(model=deployment_model) if isinstance(deployment_model, str) else fallback + + +async def _cacheable_deployment_ids( + healthy_deployments: Sequence[Mapping[str, Mapping[str, object]]], + reached_minimum: int, + reaches: Callable[[int], Awaitable[bool]], +) -> tuple[str, ...]: + minimums: Final = tuple( + (str(deployment["model_info"]["id"]), _deployment_min_tokens(deployment, reached_minimum)) + for deployment in healthy_deployments + ) + reached: Final = MappingProxyType( + { + minimum: minimum <= reached_minimum or await reaches(minimum) + for minimum in frozenset(minimum for _, minimum in minimums) + } + ) + return tuple(sorted(deployment_id for deployment_id, minimum in minimums if reached[minimum])) + + +def _deployment_models(healthy_deployments: Sequence[Mapping[str, object]]) -> Iterator[str]: + models: Final = ( + (as_str_mapping(deployment.get("litellm_params")) or {}).get("model") for deployment in healthy_deployments + ) + return (model for model in models if isinstance(model, str) and model) + + +def _get_min_token_count_for_deployments(healthy_deployments: Sequence[Mapping[str, object]]) -> int: """ Returns the lowest minimum cacheable prefix across a model group. @@ -52,11 +224,7 @@ def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int losing a cache hit it had earned. The lowest can only cost a lookup that finds nothing. """ return min( - ( - get_prompt_cache_min_tokens(model=deployment["litellm_params"]["model"]) - for deployment in healthy_deployments - if deployment.get("litellm_params", {}).get("model") - ), + (get_prompt_cache_min_tokens(model=model) for model in _deployment_models(healthy_deployments)), default=DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, ) @@ -75,10 +243,20 @@ def _honors_one_hour_cache_ttl( class PromptCachingDeploymentCheck(CustomLogger): - def __init__(self, cache: DualCache, is_priority_group: Callable[[str], bool] | None = None): + def __init__( + self, + cache: DualCache, + is_priority_group: Callable[[str], bool] | None = None, + responses_pins_in_flight: ResponsesPinsInFlight | None = None, + group_has_alternatives: Callable[[str], bool] | None = None, + ): super().__init__() self.cache = cache self.is_priority_group = is_priority_group + self.group_has_alternatives: Final = group_has_alternatives + self.responses_pins_in_flight: Final = ( + responses_pins_in_flight if responses_pins_in_flight is not None else new_responses_pins_in_flight() + ) async def async_filter_deployments( self, @@ -88,43 +266,81 @@ class PromptCachingDeploymentCheck(CustomLogger): request_kwargs: dict | None = None, parent_otel_span: Span | None = None, ) -> list[dict]: - if self.is_priority_group is not None and self.is_priority_group(model): - return healthy_deployments - if request_kwargs is not None and request_kwargs.get("_target_order") is not None: - return healthy_deployments - if not healthy_deployments[1:]: - return healthy_deployments + preferred_id: Final = await self.async_preferred_deployment_id( + model, healthy_deployments, messages, request_kwargs + ) + preferred: Final = [ + deployment for deployment in healthy_deployments if deployment["model_info"]["id"] == preferred_id + ] + return preferred[:1] or healthy_deployments - request_data: Final = ( - cast( # cast-ok: request_kwargs is a bare dict; this only reads its keys - Mapping[str, object], request_kwargs - ) - if request_kwargs is not None + async def async_preferred_deployment_id( + self, + model: str, + healthy_deployments: Sequence[Mapping[str, Mapping[str, object]]], + messages: list[AllMessageValues] | None, # mutable-ok: the request's messages list, read as is + request_data: Mapping[str, object] | None, + ) -> str | None: + if self.is_priority_group is not None and self.is_priority_group(model): + return None + if request_data is not None and request_data.get("_target_order") is not None: + return None + pin_metadata: Final = ( + request_data.get("litellm_metadata") if messages is None and request_data is not None else None + ) + single_healthy: Final = not healthy_deployments[1:] + if single_healthy and not ( + isinstance(pin_metadata, dict) + and self.group_has_alternatives is not None + and self.group_has_alternatives(model) + ): + return None + previous_attempt_token: Final = (as_str_mapping(pin_metadata) or {}).get(_RESPONSES_PIN_METADATA_KEY) + previous_response_model_id: Final = ( + _previous_response_model_id(request_data) + if request_data is not None and isinstance(pin_metadata, dict) + else None + ) + if previous_response_model_id is not None: + return previous_response_model_id + responses_prompt: Final = ( + await offload_token_count(_responses_prompt)(request_data) + if request_data is not None and isinstance(pin_metadata, dict) else None ) request_tools: Final = ( - cast( # cast-ok: request_kwargs is untyped; the stand-down scan duck-types every tool it reads + responses_prompt.tools + if responses_prompt is not None + else cast( # cast-ok: request_kwargs is untyped; the stand-down scan duck-types every tool it reads list[AllToolParamValues] | None, request_data.get("tools") ) - if request_data is not None + if request_data is not None and messages is not None else None ) + request_messages: Final = responses_prompt.messages if responses_prompt is not None else messages prompt_messages: Final = ( PromptCachingCache.prepend_system_prompt( - messages, request_data.get("system") if request_data is not None else None + request_messages, request_data.get("system") if request_data is not None else None ) - if messages is not None + if request_messages is not None else None ) - if prompt_messages is not None and await offload_token_count(is_prompt_caching_valid_prompt)( - messages=prompt_messages, - model=model, - min_token_count=_get_min_token_count_for_deployments(healthy_deployments), - tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes the counter duck-types - list[ChatCompletionToolParam] | None, - request_tools, - ), - ): + + async def reaches(minimum: int) -> bool: + return await offload_token_count( + is_prompt_caching_valid_prompt + )( + messages=prompt_messages, + model=model, + min_token_count=minimum, + tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes the counter duck-types + list[ChatCompletionToolParam] | None, + request_tools, + ), + ) + + group_minimum: Final = _get_min_token_count_for_deployments(healthy_deployments) + if prompt_messages is not None and await reaches(group_minimum): prompt_cache: Final = PromptCachingCache( cache=self.cache, ) @@ -134,34 +350,49 @@ class PromptCachingDeploymentCheck(CustomLogger): ## be derived from the messages as they will be sent, not as they arrive here. affinity_messages: Final = AnthropicCacheControlHook.messages_with_default_injections( messages=prompt_messages, - models=( - deployment["litellm_params"]["model"] - for deployment in healthy_deployments - if isinstance(deployment.get("litellm_params"), dict) and deployment["litellm_params"].get("model") - ), + models=_deployment_models(healthy_deployments), tools=request_tools, enable_prompt_caching=( - request_kwargs.get("enable_prompt_caching") is True if request_kwargs is not None else None + request_data.get("enable_prompt_caching") is True if request_data is not None else None ), - request_kwargs=request_kwargs, + request_kwargs=request_data, ) - - model_id_dict: Final = await prompt_cache.async_get_model_id( - messages=affinity_messages, - tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes - Sequence[ChatCompletionToolParam] | None, - request_tools, - ), + prefix_tools: Final = ( + responses_prompt.key_tools + if responses_prompt is not None + else cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes + Sequence[ChatCompletionToolParam] | None, request_tools + ) ) - if model_id_dict is not None and model_id_dict["model_id"] not in deployment_ids_failed_in_request( - request_data or {} - ): - model_id: Final = model_id_dict["model_id"] - for deployment in healthy_deployments: - if deployment["model_info"]["id"] == model_id: - return [deployment] + positions: Final = await PromptCachingCache.async_prefix_positions(affinity_messages, prefix_tools) + if responses_prompt is not None and isinstance(pin_metadata, dict) and positions: + if isinstance(previous_attempt_token, str): + self.responses_pins_in_flight.delete_cache(previous_attempt_token) + pin_token: Final = secrets.token_hex(16) + self.responses_pins_in_flight.set_cache( + pin_token, + _ResponsesPin( + cache_key=positions[-1].cache_key, + marked_ttl_seconds=await offload_token_count(PromptCachingCache.pin_ttl_seconds)( + affinity_messages, prefix_tools + ), + cacheable_deployment_ids=await _cacheable_deployment_ids( + healthy_deployments, + group_minimum, + reaches, + ), + ), + ) + pin_metadata[_RESPONSES_PIN_METADATA_KEY] = pin_token + if single_healthy: + return None - return healthy_deployments + model_id_dict: Final = await prompt_cache.async_get_model_id_at(positions) + tried: Final = _deployment_ids_tried_in_request(request_data or {}) if model_id_dict is not None else None + if model_id_dict is not None and tried is not None and model_id_dict["model_id"] not in tried: + return model_id_dict["model_id"] + + return None def exempt_from_message_redaction(self) -> bool: return True @@ -173,6 +404,7 @@ class PromptCachingDeploymentCheck(CustomLogger): start_time: object, end_time: object, ) -> None: + recorded_pin: Final = self._take_responses_pin(kwargs) standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: @@ -180,6 +412,10 @@ class PromptCachingDeploymentCheck(CustomLogger): call_type: Final = standard_logging_object["call_type"] + if call_type in _RESPONSES_CALL_TYPES: + await self._pin_responses_request(recorded_pin, kwargs, standard_logging_object) + return + if ( call_type != CallTypes.completion.value and call_type != CallTypes.acompletion.value @@ -241,3 +477,33 @@ class PromptCachingDeploymentCheck(CustomLogger): ) return + + def _take_responses_pin(self, kwargs: _PromptCachingSuccessEventKwargs) -> object: + litellm_params: Final = kwargs.get("litellm_params") or {} + metadata: Final = litellm_params.get("litellm_metadata") + pin_token: Final = (as_str_mapping(metadata) or {}).get(_RESPONSES_PIN_METADATA_KEY) + if not isinstance(pin_token, str): + return None + if isinstance(metadata, dict): + metadata.pop(_RESPONSES_PIN_METADATA_KEY, None) # pyright: ignore[reportUnknownMemberType] # the caller's metadata dict + recorded_pin: Final = self.responses_pins_in_flight.get_cache(pin_token) + self.responses_pins_in_flight.delete_cache(pin_token) + return recorded_pin + + async def _pin_responses_request( + self, + recorded_pin: object, + kwargs: _PromptCachingSuccessEventKwargs, + standard_logging_object: StandardLoggingPayload, + ) -> None: + litellm_params: Final = kwargs.get("litellm_params") or {} + model_id: Final = standard_logging_object["model_id"] + if not isinstance(recorded_pin, _ResponsesPin) or model_id not in recorded_pin.cacheable_deployment_ids: + return + await PromptCachingCache(cache=self.cache).async_set_pin( + model_id, + recorded_pin.cache_key, + recorded_pin.marked_ttl_seconds + if _honors_one_hour_cache_ttl(standard_logging_object, litellm_params) + else PROMPT_CACHE_PIN_TTL_SECONDS, + ) diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 1d91255e28c..e01d744dfb6 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -432,18 +432,18 @@ class PromptCachingCache: if not positions: return - ttl: Final = ( + await self.async_set_pin( + model_id, + positions[-1].cache_key, await offload_token_count(PromptCachingCache.pin_ttl_seconds)(messages, tools) if honors_one_hour_ttl - else PROMPT_CACHE_PIN_TTL_SECONDS + else PROMPT_CACHE_PIN_TTL_SECONDS, ) - self.cache.in_memory_cache.delete_cache(positions[-1].cache_key) + + async def async_set_pin(self, model_id: str, cache_key: str, ttl: int) -> None: + self.cache.in_memory_cache.delete_cache(cache_key) with service_target(_PROMPT_CACHE_PINS_TARGET): - await self.cache.async_set_cache( - positions[-1].cache_key, - PromptCachingCacheValue(model_id=model_id), - ttl=ttl, - ) + await self.cache.async_set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=ttl) async def async_get_model_id( self, @@ -455,7 +455,10 @@ class PromptCachingCache: same way the provider cache does, so a breakpoint that moved forward since the last turn still lands on the deployment whose cache holds the earlier prefix. """ - cache_keys: Final = _lookback_keys(await PromptCachingCache.async_prefix_positions(messages, tools)) + return await self.async_get_model_id_at(await PromptCachingCache.async_prefix_positions(messages, tools)) + + async def async_get_model_id_at(self, positions: tuple[PrefixPosition, ...]) -> PromptCachingCacheValue | None: + cache_keys: Final = _lookback_keys(positions) if not cache_keys: return None 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/guardrails.py b/litellm/types/guardrails.py index 718a651bdee..e311f2ee5f9 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -1050,6 +1050,14 @@ GuardrailNumRetries: TypeAlias = Annotated[ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch update guardrails + mcp_block_behavior: Literal["tool_result", "raise"] | None = Field( + default=None, + description=( + "How to handle guardrail blocks when executing MCP tools for an LLM request. " + "'raise' propagates the guardrail error to the client. 'tool_result' or None " + "returns the block as a tool result and lets the LLM continue." + ), + ) api_key: str | None = Field(default=None, description="API key for the guardrail service") api_base: str | None = Field(default=None, description="Base URL for the guardrail service API") @@ -1162,12 +1170,13 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up description=( "For /v1/realtime sessions: 'warn' speaks the violation message and continues; " "'end_session' speaks the message and closes the connection. " + "Bedrock Nova Sonic cannot cancel a spoken turn's answer, so a blocked spoken turn always ends its session. " "For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning." ), ) realtime_violation_message: str | None = Field( default=None, - description="The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set.", + description="The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to the guardrail's block reason if not set.", ) # Model Armor params diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 9f63d27f5f3..c04430322c4 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -207,6 +207,7 @@ class DeploymentOptions: @dataclass(frozen=True, slots=True, kw_only=True) class SpecializedRouterOptions: + switchyard_router_config: Mapping[str, object] | None = None auto_router_config_path: str | None = None auto_router_config: str | None = None auto_router_default_model: str | None = None 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/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index 47f80c52845..4220a67bddb 100644 --- a/litellm/types/llms/custom_http.py +++ b/litellm/types/llms/custom_http.py @@ -34,6 +34,7 @@ class httpxSpecialProvider(str, Enum): ROICalculator = "roi_calculator" Sandbox = "sandbox" ModelCostMap = "model_cost_map" + CodexModelCatalog = "codex_model_catalog" PasswordBreachCheck = "password_breach_check" ASGI = "asgi" AgentHarness = "agent_harness" diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 06c28e371f4..36340a2e57b 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -65,6 +65,7 @@ from openai.types.responses.response_create_params import ( from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from openai.types.responses.response_function_web_search import ResponseFunctionWebSearch from pydantic import ( + BaseModel, ConfigDict, Discriminator, Field, @@ -957,6 +958,7 @@ class ChatCompletionToolMessage(TypedDict): role: Literal["tool"] content: str | Iterable[ToolMessageContentPart] # writable-ok: Pydantic warns on ReadOnly TypedDict fields tool_call_id: str + is_error: NotRequired[bool] # writable-ok: Pydantic warns on ReadOnly TypedDict fields class ChatCompletionFunctionMessage(TypedDict): @@ -1569,6 +1571,24 @@ def replayable_responses_item( } +def _with_wire_names(value: object, info: FieldSerializationInfo) -> object: + items: Final = _as_json_array(value) + if items is None: + return value + return [ + item.model_dump( + mode=info.mode, + by_alias=True, + exclude_unset=info.exclude_unset, + exclude_defaults=info.exclude_defaults, + exclude_none=info.exclude_none, + ) + if isinstance(item, BaseModel) + else item + for item in items + ] + + def replayable_responses_input(value: object, *, drops_set_created_by: bool = False) -> object: items: Final = _as_json_array(value) if items is None: @@ -1654,16 +1674,23 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): @field_serializer("output", mode="wrap") @classmethod def _serialize_output_as_replayable( - cls, value: builtins.object, handler: SerializerFunctionWrapHandler, _info: FieldSerializationInfo + cls, value: builtins.object, handler: SerializerFunctionWrapHandler, info: FieldSerializationInfo ): """ Drop the nulls the SDK types default in but providers never send: created_by on every output item (OpenAI rejects it when `output` is replayed as `input`), and status/content/encrypted_content on reasoning items (the OpenAI C# SDK crashes on status=null, issue #16824). """ - serialized: Final[object] = handler(value) + serialized: Final[object] = handler(_with_wire_names(value, info)) return replayable_responses_input(serialized) + @field_serializer("tools", mode="wrap") + @classmethod + def _serialize_tools_with_wire_names( # noqa: ANN206 # annotating it replaces the model's serialization schema + cls, value: builtins.object, handler: SerializerFunctionWrapHandler, info: FieldSerializationInfo + ): + return handler(_with_wire_names(value, info)) + @property def output_text(self) -> str: """ diff --git a/litellm/types/router.py b/litellm/types/router.py index 3dcfc694003..c0b182f0f20 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -659,6 +659,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): tag_regex: list[str] | None = None # auto-router params + switchyard_router_config: Mapping[str, object] | None = None auto_router_config_path: str | None = None auto_router_config: str | None = None auto_router_default_model: str | None = None @@ -1461,6 +1462,7 @@ class PreRoutingStrategy(Protocol): messages: list[dict[str, Any]] | None = None, input: "str | list[Any] | None" = None, specific_deployment: bool | None = False, + call_type: str | None = None, ) -> "PreRoutingHookResponse | None": ... diff --git a/litellm/types/utils.py b/litellm/types/utils.py index b082198d3f2..755a0a8cc7a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3382,7 +3382,10 @@ class StandardLoggingRoutingDecision(TypedDict, total=False): """Per-request provenance for a pre-routing strategy (auto-router) decision.""" router_model_name: str - router_type: ReadOnly[Literal["complexity", "adaptive", "quality", "semantic"]] + router_type: ReadOnly[Literal["complexity", "adaptive", "quality", "semantic", "switchyard"]] + switchyard_algorithm: ReadOnly[str] + switchyard_evidence: ReadOnly[object] + switchyard_fallback_reason: ReadOnly[str] router_config_id: ReadOnly[str] router_config_updated_at: ReadOnly[str] router_config_fingerprint: ReadOnly[str] @@ -3443,10 +3446,12 @@ class StandardLoggingKubernetesPodRouting(TypedDict): # so a redacted row stays explainable. `test_every_routing_decision_field_is_classified` # fails if a field is added to the record without being placed in one set or the other. PROMPT_QUOTING_ROUTING_DECISION_FIELDS: frozenset[str] = frozenset( - {"signals", "matched_keyword", "escalation_keyword", "classifier_crux"} + {"signals", "matched_keyword", "escalation_keyword", "classifier_crux", "switchyard_evidence"} ) DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset( { + "switchyard_algorithm", + "switchyard_fallback_reason", "router_model_name", "router_type", "router_config_id", 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/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4796ed3ac9c..9a3027475b7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -13865,12 +13865,12 @@ "azure_ai/jais-30b-chat": { "deprecation_date": "2026-01-30", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", - "input_cost_per_token": 0.0032, + "input_cost_per_token": 3.2e-06, "litellm_provider": "azure_ai", "max_input_tokens": 8192, "max_output_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.00971, + "output_cost_per_token": 9.71e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", "supports_pdf_input": false }, @@ -55160,8 +55160,8 @@ "watsonx/core42/jais-13b-chat": { "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.0005, - "output_cost_per_token": 0.002, + "input_cost_per_token": 1.8e-06, + "output_cost_per_token": 1.8e-06, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -64213,6 +64213,39 @@ "supports_tool_choice": true, "supports_vision": false }, + "tencent/hy4": { + "cache_read_input_token_cost": 4.2e-08, + "comment": "TokenHub bills Hy4 at 50% of list until 2026-12-31 23:59:59 Beijing time; these are the list rates.", + "input_cost_per_token": 8.34e-07, + "input_cost_per_token_cache_hit": 4.2e-08, + "litellm_provider": "tencent", + "max_input_tokens": 960000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.501e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true + }, + "tencent/hy4-preview": { + "cache_read_input_token_cost": 4.2e-08, + "comment": "TokenHub bills Hy4 preview at 50% of list until 2026-12-31 23:59:59 Beijing time; these are the list rates.", + "input_cost_per_token": 8.34e-07, + "input_cost_per_token_cache_hit": 4.2e-08, + "litellm_provider": "tencent", + "max_input_tokens": 960000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.501e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true + }, "tencent/minimax-m3": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 6e-08, @@ -70882,6 +70915,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token_above_272k_tokens": 1e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, "source": "https://docs.perplexity.ai/docs/agent-api/models" }, "perplexity/openai/gpt-5.4": { @@ -83107,13 +83141,16 @@ "mode": "chat", "max_input_tokens": 1000000, "max_output_tokens": 393216, - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 8e-07, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-1-flash-1017" + "source": "https://empiriolabs.ai/models/deepseek-v4-1-flash-1017", + "cache_read_input_token_cost": 2e-08, + "supports_prompt_caching": true, + "supports_function_calling": true }, "empiriolabs/deepseek-v4-flash-0731": { "litellm_provider": "empiriolabs", @@ -83125,7 +83162,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731" + "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-flash-0731:variant1": { "litellm_provider": "empiriolabs", @@ -83137,7 +83175,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731" + "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro": { "litellm_provider": "empiriolabs", @@ -83149,7 +83188,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro-0813": { "litellm_provider": "empiriolabs", @@ -83161,7 +83201,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro-0813" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro-0813", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro:variant1": { "litellm_provider": "empiriolabs", @@ -83179,7 +83220,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro:variant2": { "litellm_provider": "empiriolabs", @@ -83199,7 +83241,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/deepseek-v4-pro:variant3": { "litellm_provider": "empiriolabs", @@ -83211,7 +83254,8 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + "source": "https://empiriolabs.ai/models/deepseek-v4-pro", + "supports_function_calling": true }, "empiriolabs/fugu-max": { "litellm_provider": "empiriolabs", @@ -83409,7 +83453,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/glm-5-2" + "source": "https://empiriolabs.ai/models/glm-5-2", + "supports_function_calling": true }, "empiriolabs/glm-5-2:variant2": { "litellm_provider": "empiriolabs", @@ -83427,7 +83472,8 @@ "search_context_size_high": 0.013 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/glm-5-2" + "source": "https://empiriolabs.ai/models/glm-5-2", + "supports_function_calling": true }, "empiriolabs/glm-5-3": { "litellm_provider": "empiriolabs", @@ -83481,7 +83527,8 @@ "supports_reasoning": true, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-6" + "source": "https://empiriolabs.ai/models/kimi-k2-6", + "supports_function_calling": true }, "empiriolabs/kimi-k2-7-code": { "litellm_provider": "empiriolabs", @@ -83500,7 +83547,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-7-code" + "source": "https://empiriolabs.ai/models/kimi-k2-7-code", + "supports_function_calling": true }, "empiriolabs/kimi-k2-7-code-highspeed": { "litellm_provider": "empiriolabs", @@ -83519,7 +83567,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-7-code-highspeed" + "source": "https://empiriolabs.ai/models/kimi-k2-7-code-highspeed", + "supports_function_calling": true }, "empiriolabs/kimi-k2-7-code:variant1": { "litellm_provider": "empiriolabs", @@ -83540,7 +83589,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k2-7-code" + "source": "https://empiriolabs.ai/models/kimi-k2-7-code", + "supports_function_calling": true }, "empiriolabs/kimi-k3": { "litellm_provider": "empiriolabs", @@ -83559,7 +83609,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/kimi-k3" + "source": "https://empiriolabs.ai/models/kimi-k3", + "supports_function_calling": true }, "empiriolabs/kimi-k3-1-pro": { "litellm_provider": "empiriolabs", @@ -83601,7 +83652,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/mimo-v2-6-flash" + "source": "https://empiriolabs.ai/models/mimo-v2-6-flash", + "supports_function_calling": true }, "empiriolabs/mimo-v2-6-pro": { "litellm_provider": "empiriolabs", @@ -83621,7 +83673,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/mimo-v2-6-pro" + "source": "https://empiriolabs.ai/models/mimo-v2-6-pro", + "supports_function_calling": true }, "empiriolabs/mimo-v2-6-pro-ultraspeed": { "litellm_provider": "empiriolabs", @@ -83635,7 +83688,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/mimo-v2-6-pro-ultraspeed" + "source": "https://empiriolabs.ai/models/mimo-v2-6-pro-ultraspeed", + "supports_function_calling": true }, "empiriolabs/minimax-m2-7": { "litellm_provider": "empiriolabs", @@ -83750,7 +83804,8 @@ "search_context_size_high": 0.14 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/mistral-small-4" + "source": "https://empiriolabs.ai/models/mistral-small-4", + "supports_function_calling": true }, "empiriolabs/muse-glimmer-30b": { "litellm_provider": "empiriolabs", @@ -83765,7 +83820,8 @@ "supports_reasoning": true, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/muse-glimmer-30b" + "source": "https://empiriolabs.ai/models/muse-glimmer-30b", + "supports_function_calling": true }, "empiriolabs/muse-spark-1-1": { "litellm_provider": "empiriolabs", @@ -83831,7 +83887,8 @@ "supports_system_messages": true, "supports_vision": true, "supports_prompt_caching": true, - "source": "https://empiriolabs.ai/models/nova-lite-1-0" + "source": "https://empiriolabs.ai/models/nova-lite-1-0", + "supports_function_calling": true }, "empiriolabs/nova-lite-2": { "litellm_provider": "empiriolabs", @@ -83845,7 +83902,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_prompt_caching": true, - "source": "https://empiriolabs.ai/models/nova-lite-2" + "source": "https://empiriolabs.ai/models/nova-lite-2", + "supports_function_calling": true }, "empiriolabs/nova-micro-1-0": { "litellm_provider": "empiriolabs", @@ -83857,7 +83915,8 @@ "cache_read_input_token_cost": 2.24e-08, "supports_system_messages": true, "supports_prompt_caching": true, - "source": "https://empiriolabs.ai/models/nova-micro-1-0" + "source": "https://empiriolabs.ai/models/nova-micro-1-0", + "supports_function_calling": true }, "empiriolabs/nova-pro-1-0": { "litellm_provider": "empiriolabs", @@ -83868,7 +83927,8 @@ "output_cost_per_token": 9.6e-06, "supports_system_messages": true, "supports_vision": true, - "source": "https://empiriolabs.ai/models/nova-pro-1-0" + "source": "https://empiriolabs.ai/models/nova-pro-1-0", + "supports_function_calling": true }, "empiriolabs/perplexity-advanced-deep-research": { "litellm_provider": "empiriolabs", @@ -83968,7 +84028,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b" + "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-27b": { "litellm_provider": "empiriolabs", @@ -83989,7 +84050,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-27b" + "source": "https://empiriolabs.ai/models/qwen3-5-27b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-35b-a3b": { "litellm_provider": "empiriolabs", @@ -84010,7 +84072,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b" + "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-397b-a17b": { "litellm_provider": "empiriolabs", @@ -84031,7 +84094,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b" + "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b", + "supports_function_calling": true }, "empiriolabs/qwen3-5-4b": { "litellm_provider": "empiriolabs", @@ -84082,7 +84146,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-flash" + "source": "https://empiriolabs.ai/models/qwen3-5-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-5-flash:variant1": { "litellm_provider": "empiriolabs", @@ -84105,7 +84170,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-flash" + "source": "https://empiriolabs.ai/models/qwen3-5-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-5-omni-flash": { "litellm_provider": "empiriolabs", @@ -84123,7 +84189,8 @@ "search_context_size_high": 0.015 }, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash" + "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-5-omni-plus": { "litellm_provider": "empiriolabs", @@ -84142,7 +84209,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus" + "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-5-plus": { "litellm_provider": "empiriolabs", @@ -84151,8 +84219,8 @@ "max_output_tokens": 65536, "input_cost_per_token": 3.6e-07, "output_cost_per_token": 2.21e-06, - "input_cost_per_token_above_256k_tokens": 1.08e-06, - "output_cost_per_token_above_256k_tokens": 6.62e-06, + "input_cost_per_token_above_256k_tokens": 4.5e-07, + "output_cost_per_token_above_256k_tokens": 2.76e-06, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, @@ -84163,7 +84231,8 @@ "search_context_size_high": 0.015 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-plus" + "source": "https://empiriolabs.ai/models/qwen3-5-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-5-plus:variant1": { "litellm_provider": "empiriolabs", @@ -84186,7 +84255,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-5-plus" + "source": "https://empiriolabs.ai/models/qwen3-5-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-6-27b": { "litellm_provider": "empiriolabs", @@ -84205,7 +84275,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-27b" + "source": "https://empiriolabs.ai/models/qwen3-6-27b", + "supports_function_calling": true }, "empiriolabs/qwen3-6-35b-a3b": { "litellm_provider": "empiriolabs", @@ -84219,7 +84290,8 @@ "supports_reasoning": true, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-35b-a3b" + "source": "https://empiriolabs.ai/models/qwen3-6-35b-a3b", + "supports_function_calling": true }, "empiriolabs/qwen3-6-flash": { "litellm_provider": "empiriolabs", @@ -84240,7 +84312,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-flash" + "source": "https://empiriolabs.ai/models/qwen3-6-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-6-flash:variant1": { "litellm_provider": "empiriolabs", @@ -84261,7 +84334,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-flash" + "source": "https://empiriolabs.ai/models/qwen3-6-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-6-plus": { "litellm_provider": "empiriolabs", @@ -84282,7 +84356,8 @@ "search_context_size_high": 0.026 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-plus" + "source": "https://empiriolabs.ai/models/qwen3-6-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-6-plus:variant1": { "litellm_provider": "empiriolabs", @@ -84303,7 +84378,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-6-plus" + "source": "https://empiriolabs.ai/models/qwen3-6-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-7-flash": { "litellm_provider": "empiriolabs", @@ -84330,7 +84406,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-flash" + "source": "https://empiriolabs.ai/models/qwen3-7-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-7-flash:variant1": { "litellm_provider": "empiriolabs", @@ -84357,7 +84434,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-flash" + "source": "https://empiriolabs.ai/models/qwen3-7-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-7-max": { "litellm_provider": "empiriolabs", @@ -84375,7 +84453,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-max" + "source": "https://empiriolabs.ai/models/qwen3-7-max", + "supports_function_calling": true }, "empiriolabs/qwen3-7-max:variant1": { "litellm_provider": "empiriolabs", @@ -84393,7 +84472,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-max" + "source": "https://empiriolabs.ai/models/qwen3-7-max", + "supports_function_calling": true }, "empiriolabs/qwen3-7-plus": { "litellm_provider": "empiriolabs", @@ -84414,7 +84494,8 @@ "search_context_size_high": 0.03 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-plus" + "source": "https://empiriolabs.ai/models/qwen3-7-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-7-plus:variant1": { "litellm_provider": "empiriolabs", @@ -84438,7 +84519,8 @@ }, "supports_prompt_caching": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-7-plus" + "source": "https://empiriolabs.ai/models/qwen3-7-plus", + "supports_function_calling": true }, "empiriolabs/qwen3-8-27b": { "litellm_provider": "empiriolabs", @@ -84473,7 +84555,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-flash" + "source": "https://empiriolabs.ai/models/qwen3-8-flash", + "supports_function_calling": true }, "empiriolabs/qwen3-8-max": { "litellm_provider": "empiriolabs", @@ -84492,7 +84575,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-max" + "source": "https://empiriolabs.ai/models/qwen3-8-max", + "supports_function_calling": true }, "empiriolabs/qwen3-8-max-0902": { "litellm_provider": "empiriolabs", @@ -84511,7 +84595,8 @@ "search_context_size_high": 0.02 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-max-0902" + "source": "https://empiriolabs.ai/models/qwen3-8-max-0902", + "supports_function_calling": true }, "empiriolabs/qwen3-8-max:variant1": { "litellm_provider": "empiriolabs", @@ -84530,7 +84615,8 @@ "search_context_size_high": 0.01 }, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/qwen3-8-max" + "source": "https://empiriolabs.ai/models/qwen3-8-max", + "supports_function_calling": true }, "empiriolabs/qwen3-8-omni-flash": { "litellm_provider": "empiriolabs", @@ -84550,7 +84636,8 @@ }, "supports_response_schema": true, "supports_audio_input": true, - "source": "https://empiriolabs.ai/models/qwen3-8-omni-flash" + "source": "https://empiriolabs.ai/models/qwen3-8-omni-flash", + "supports_function_calling": true }, "empiriolabs/seed-2-0-code": { "litellm_provider": "empiriolabs", @@ -84564,7 +84651,8 @@ "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, - "source": "https://empiriolabs.ai/models/seed-2-0-code" + "source": "https://empiriolabs.ai/models/seed-2-0-code", + "supports_function_calling": true }, "empiriolabs/seed-2-0-lite": { "litellm_provider": "empiriolabs", @@ -84579,7 +84667,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-0-lite" + "source": "https://empiriolabs.ai/models/seed-2-0-lite", + "supports_function_calling": true }, "empiriolabs/seed-2-0-mini": { "litellm_provider": "empiriolabs", @@ -84594,7 +84683,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-0-mini" + "source": "https://empiriolabs.ai/models/seed-2-0-mini", + "supports_function_calling": true }, "empiriolabs/seed-2-0-pro": { "litellm_provider": "empiriolabs", @@ -84609,7 +84699,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-0-pro" + "source": "https://empiriolabs.ai/models/seed-2-0-pro", + "supports_function_calling": true }, "empiriolabs/seed-2-1-turbo": { "litellm_provider": "empiriolabs", @@ -84622,7 +84713,8 @@ "supports_vision": true, "supports_reasoning": true, "supports_response_schema": true, - "source": "https://empiriolabs.ai/models/seed-2-1-turbo" + "source": "https://empiriolabs.ai/models/seed-2-1-turbo", + "supports_function_calling": true }, "empiriolabs/skylark-embedding-vision": { "litellm_provider": "empiriolabs", diff --git a/pyproject.toml b/pyproject.toml index 97abe9d6376..c19e54ab230 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.109.0" +version = "1.110.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.16" @@ -54,6 +54,7 @@ Documentation = "https://docs.litellm.ai" # sub-packages and the opentelemetry trio move in lockstep, and grpcio is # supply-chain-pinned to a vetted, aged release. [project.optional-dependencies] +switchyard = ["nemo-switchyard==0.4.0a1"] proxy = [ "gunicorn>=23.0.0,<24.0", "uvicorn>=0.33.0,<1.0", @@ -84,8 +85,8 @@ proxy = [ "mcp>=2.3.0,<3", "httpx2>=2.10.0,<3", "pydantic>=2.12.0,<3", - "litellm-proxy-extras==0.4.112", - "litellm-enterprise==0.1.81", + "litellm-proxy-extras==0.4.114", + "litellm-enterprise==0.1.83", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", @@ -268,6 +269,7 @@ proxy-dev = [ "opentelemetry-instrumentation-fastapi==0.66b1", "azure-identity==1.25.2", "a2a-sdk==1.1.0", + "nemo-switchyard==0.4.0a1", "backports-zstd==1.8.0; python_version < '3.14'", "httptools>=0.7.1,<0.8", "semver>=3.0.4,<4", @@ -408,7 +410,7 @@ litellm-enterprise = { workspace = true } members = ["enterprise", "litellm-proxy-extras", "packaging/litellm-proxy"] [tool.commitizen] -version = "1.109.0" +version = "1.110.0" version_files = [ "pyproject.toml:^version", ] diff --git a/tests/integration/compatibility/test_codex_model_catalog.py b/tests/integration/compatibility/test_codex_model_catalog.py index 6aeeead0d90..a513088d082 100644 --- a/tests/integration/compatibility/test_codex_model_catalog.py +++ b/tests/integration/compatibility/test_codex_model_catalog.py @@ -36,7 +36,7 @@ pytestmark: Final = pytest.mark.shares_worker_state T = TypeVar("T") -CLIENT_VERSION: Final = "0.159.3" +CLIENT_VERSION: Final = "0.166.0" CONSISTENT_READS: Final = 6 CONVERGENCE_SECONDS: Final = 60 RELOAD_MARGIN_SECONDS: Final = 5.0 diff --git a/tests/integration/configuration/test_codex_catalog_boot.py b/tests/integration/configuration/test_codex_catalog_boot.py index 08c60debfcf..6564827e2c6 100644 --- a/tests/integration/configuration/test_codex_catalog_boot.py +++ b/tests/integration/configuration/test_codex_catalog_boot.py @@ -18,7 +18,7 @@ from tests.integration._support.process import owned_proxy, owned_proxy_process pytestmark: Final = pytest.mark.shares_worker_state -CLIENT_VERSION: Final = "0.159.3" +CLIENT_VERSION: Final = "0.166.0" FRESH_CONNECTION: Final = {"Connection": "close"} CONSISTENT_READS: Final = 8 CONVERGENCE_SECONDS: Final = 60 diff --git a/tests/integration/observability/_openinference_support.py b/tests/integration/observability/_openinference_support.py index 58785c89033..76ea8e48504 100644 --- a/tests/integration/observability/_openinference_support.py +++ b/tests/integration/observability/_openinference_support.py @@ -8,6 +8,7 @@ from contextlib import AbstractContextManager, ExitStack, contextmanager, nullco from dataclasses import dataclass, replace from itertools import chain from pathlib import Path +from types import MappingProxyType from typing import Final import httpx @@ -353,6 +354,9 @@ def _normalize_chat_caller_stream(events: Sequence[JsonValue]) -> tuple[dict[str return tuple({**_json_object_value(event), "created": 1} for event in events) +_FUNCTION_CALL_NULL_FIELDS: Final = MappingProxyType({"async": None, "caller": None, "namespace": None}) + + def _responses_caller_body(body: dict[str, JsonValue], model: str) -> dict[str, JsonValue]: usage: Final = _json_object_value(body["usage"]) output: Final = JSON_MESSAGES.validate_python(body["output"]) @@ -376,7 +380,9 @@ def _responses_caller_body(body: dict[str, JsonValue], model: str) -> dict[str, "truncation": None, "user": None, "store": None, - "output": [{**item, **({"namespace": None} if item.get("type") == "function_call" else {})} for item in output], + "output": [ + {**item, **(_FUNCTION_CALL_NULL_FIELDS if item.get("type") == "function_call" else {})} for item in output + ], "usage": { "input_tokens_details": None, "output_tokens_details": None, diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_spans.py b/tests/integration/observability/test_arize_otel_v2_openinference_spans.py index 6a7d6d99a84..351b77c89d6 100644 --- a/tests/integration/observability/test_arize_otel_v2_openinference_spans.py +++ b/tests/integration/observability/test_arize_otel_v2_openinference_spans.py @@ -263,9 +263,9 @@ def test_arize_otel_v2_a5_responses_sync_sdk(gateway: Gateway, tmp_path: Path) - extra_body={"metadata": {"trace_marker": marker}, "cache": {"no-cache": True}}, ) assert response.id.startswith("resp_"), response - assert _normalize_responses_caller_body(response.model_dump(mode="json", exclude_unset=True)) == ( - _responses_caller_response(_responses_response(marker), rig.model) - ), response + assert _normalize_responses_caller_body( + response.model_dump(mode="json", exclude_unset=True, by_alias=True) + ) == (_responses_caller_response(_responses_response(marker), rig.model)), response assert (response.status, response.model) == ("completed", rig.model), response assert ( response.output[0].type, @@ -809,9 +809,9 @@ def _cache_call( extra_body={"metadata": {"trace_marker": marker}, "cache": {"no-cache": False}}, ) response: Final = raw.parse() - assert _normalize_responses_caller_body(response.model_dump(mode="json", exclude_unset=True)) == ( - _responses_caller_response(_responses_response(marker), rig.model) - ), response + assert _normalize_responses_caller_body( + response.model_dump(mode="json", exclude_unset=True, by_alias=True) + ) == (_responses_caller_response(_responses_response(marker), rig.model)), response call: Final = response.output[0] assert call.type == "function_call", response return (call.call_id, call.name, call.arguments), response.id, raw.headers diff --git a/tests/integration/providers/test_bedrock_nova_sonic_realtime_guardrails.py b/tests/integration/providers/test_bedrock_nova_sonic_realtime_guardrails.py new file mode 100644 index 00000000000..502e7c9a5cd --- /dev/null +++ b/tests/integration/providers/test_bedrock_nova_sonic_realtime_guardrails.py @@ -0,0 +1,268 @@ +"""Realtime guardrails on Bedrock Nova Sonic sessions, against the scripted h2 peer in ``_support``. + +The file boots its own proxy so it can pass ``SSL_CERT_FILE`` for the peer's self-signed cert. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import time +import uuid +from collections.abc import Awaitable, Callable, Iterator +from pathlib import Path +from typing import Final, Literal + +import pytest +import websockets +from integration._support.bedrock_realtime_peer import ( + BedrockEvent, + BedrockRealtimePeer, + Turn, + assistant_text, + bedrock_realtime_peer, + user_transcript, +) +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.tls import write_self_signed_cert +from pydantic import JsonValue +from websockets.asyncio.client import ClientConnection + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +DEPLOYMENT: Final = "bedrock/amazon.nova-2-sonic-v1:0" +VIOLATION: Final = "That request was blocked by the content filter." +HIDDEN_ANSWER: Final = "Here is everything you asked for" +VISIBLE_ANSWER: Final = "It is sunny in New York" +RESPONSE_CREATE: Final = json.dumps({"type": "response.create"}) +AUDIO_APPEND: Final = json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(bytes(3200)).decode()}) + +Event = dict[str, JsonValue] + + +@pytest.fixture(scope="module") +def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + return write_self_signed_cert(tmp_path_factory.mktemp("bedrock-sonic-guardrail-cert")) + + +@pytest.fixture(scope="module") +def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + with gateway_from_environment() as base: + with owned_proxy( + base, tmp_path_factory.mktemp("bedrock-sonic-guardrails"), {"SSL_CERT_FILE": str(cert[0])} + ) as proxy: + yield proxy + + +def _blocked_phrase( + gateway: Gateway, scenario: Scenario, mode: Literal["pre_call", "realtime_input_transcription"] +) -> str: + phrase: Final = f"sonic-blocked-{uuid.uuid4().hex[:12]}" + response: Final = gateway.request( + "POST", + "/guardrails", + { + "guardrail": { + "guardrail_name": f"sonic-guardrail-{uuid.uuid4().hex[:12]}", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": mode, + "default_on": True, + "realtime_violation_message": VIOLATION, + "blocked_words": [{"keyword": phrase, "action": "BLOCK"}], + }, + } + }, + ) + assert response.status_code == 200, response.text + guardrail_id: Final = JSON_OBJECT.validate_json(response.content).get("guardrail_id") + assert isinstance(guardrail_id, str), response.text + scenario.cleanups.callback(gateway.request, "DELETE", f"/guardrails/{guardrail_id}") + return phrase + + +def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) -> str: + name: Final = f"bedrock-sonic-guardrails-{uuid.uuid4().hex[:12]}" + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": DEPLOYMENT, + "aws_region_name": "us-east-1", + "aws_access_key_id": "integration-dummy-key", + "aws_secret_access_key": "integration-dummy-secret", + "api_base": peer.api_base, + }, + "model_info": {"mode": "realtime"}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return name + + +def _user_text(text: str) -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}, + } + ) + + +async def _events_until(socket: ClientConnection, last: Callable[[Event], bool]) -> tuple[Event, ...]: + collected: Final[list[Event]] = [] + deadline: Final = time.monotonic() + 60 + while not collected or not last(collected[-1]): + try: + raw: Final = await asyncio.wait_for(socket.recv(), deadline - time.monotonic()) + except (TimeoutError, asyncio.TimeoutError): + raise AssertionError(f"timed out waiting; types={[event.get('type') for event in collected]}") + collected.append(JSON_OBJECT.validate_json(raw)) + return tuple(collected) + + +async def _events_until_closed(socket: ClientConnection) -> tuple[Event, ...]: + return tuple([JSON_OBJECT.validate_json(raw) async for raw in socket]) + + +def _is_done(event: Event) -> bool: + return event.get("type") == "response.done" + + +def _is_guardrail_error(event: Event) -> bool: + error: Final = event.get("error") + return event.get("type") == "error" and isinstance(error, dict) and error.get("type") == "guardrail_violation" + + +async def _session( + gateway: Gateway, model: str, key: str, exchange: Callable[[ClientConnection], Awaitable[tuple[Event, ...]]] +) -> tuple[Event, ...]: + proxy: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/realtime?model={model}", additional_headers={"Authorization": f"Bearer {key}"} + ) as socket: + created: Final = JSON_OBJECT.validate_json(await socket.recv()) + assert created.get("type") == "session.created", created + return await exchange(socket) + + +def _fields(event: BedrockEvent, name: str) -> dict[str, JsonValue]: + return object_value(object_value(event["event"]).get(name) or {}) + + +def _user_text_inputs(received: tuple[BedrockEvent, ...]) -> tuple[str, ...]: + content_starts: Final = tuple(_fields(event, "contentStart") for event in received) + user_blocks: Final = frozenset(start.get("contentName") for start in content_starts if start.get("role") == "USER") + text_inputs: Final = tuple(_fields(event, "textInput") for event in received) + return tuple( + string_value(text_input["content"]) + for text_input in text_inputs + if text_input.get("contentName") in user_blocks + ) + + +def _text_deltas(events: tuple[Event, ...]) -> str: + return "".join(string_value(event["delta"]) for event in events if event.get("type") == "response.text.delta") + + +def test_blocked_typed_text_never_reaches_nova_sonic_and_the_violation_is_voiced( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + turn: Final = Turn("user_input", user_transcript("ready") + assistant_text(VIOLATION)) + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + phrase: Final = _blocked_phrase(candidate, scenario, "pre_call") + model: Final = _deployment(candidate, scenario, peer) + + async def exchange(socket: ClientConnection) -> tuple[Event, ...]: + await socket.send(_user_text(f"please repeat {phrase}")) + await socket.send(RESPONSE_CREATE) + return await _events_until(socket, _is_done) + + events: Final = asyncio.run(_session(candidate, model, scenario.key(), exchange)) + + errors: Final = tuple(object_value(event["error"]) for event in events if _is_guardrail_error(event)) + assert [error["message"] for error in errors] == [VIOLATION], events + assert [error["code"] for error in errors] == ["content_policy_violation"], events + assert _text_deltas(events) == VIOLATION, events + received: Final = peer.received() + assert phrase not in json.dumps(received), received + assert [VIOLATION in text for text in _user_text_inputs(received)] == [True], received + + +def test_clean_typed_text_reaches_nova_sonic_unchanged(candidate: Gateway, cert: tuple[Path, Path]) -> None: + turn: Final = Turn("user_input", user_transcript("ready") + assistant_text(VISIBLE_ANSWER)) + text: Final = f"what is the weather in New York {uuid.uuid4().hex[:8]}" + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + _blocked_phrase(candidate, scenario, "pre_call") + model: Final = _deployment(candidate, scenario, peer) + + async def exchange(socket: ClientConnection) -> tuple[Event, ...]: + await socket.send(_user_text(text)) + await socket.send(RESPONSE_CREATE) + return await _events_until(socket, _is_done) + + events: Final = asyncio.run(_session(candidate, model, scenario.key(), exchange)) + + assert not any(_is_guardrail_error(event) for event in events), events + assert _text_deltas(events) == VISIBLE_ANSWER, events + assert _user_text_inputs(peer.received()) == (text,) + + +def test_blocked_transcript_ends_the_session_before_its_answer_reaches_the_client( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + with candidate.scenario() as scenario: + phrase: Final = _blocked_phrase(candidate, scenario, "realtime_input_transcription") + turn: Final = Turn("user_input", user_transcript(f"tell me {phrase}") + assistant_text(HIDDEN_ANSWER)) + with bedrock_realtime_peer((turn,), cert) as peer: + model: Final = _deployment(candidate, scenario, peer) + + async def exchange(socket: ClientConnection) -> tuple[Event, ...]: + await socket.send(AUDIO_APPEND) + await socket.send(_user_text("hello")) + events: Final = await asyncio.wait_for(_events_until_closed(socket), 60) + assert socket.close_code == 1000, (socket.close_code, events) + return events + + events: Final = asyncio.run(_session(candidate, model, scenario.key(), exchange)) + + assert HIDDEN_ANSWER not in json.dumps(events), events + assert not any(str(event.get("type")).startswith("response.") for event in events), events + transcripts: Final = [ + event["transcript"] + for event in events + if event.get("type") == "conversation.item.input_audio_transcription.completed" + ] + assert transcripts == [f"tell me {phrase}"], events + assert [object_value(event["error"])["message"] for event in events if _is_guardrail_error(event)] == [ + VIOLATION + ], events + + +def test_clean_transcript_under_a_transcript_guardrail_is_answered(candidate: Gateway, cert: tuple[Path, Path]) -> None: + spoken: Final = f"what is the weather in New York {uuid.uuid4().hex[:8]}" + with candidate.scenario() as scenario: + _blocked_phrase(candidate, scenario, "realtime_input_transcription") + turn: Final = Turn("user_input", user_transcript(spoken) + assistant_text(VISIBLE_ANSWER)) + with bedrock_realtime_peer((turn,), cert) as peer: + model: Final = _deployment(candidate, scenario, peer) + + async def exchange(socket: ClientConnection) -> tuple[Event, ...]: + await socket.send(AUDIO_APPEND) + await socket.send(_user_text("hello")) + return await _events_until(socket, _is_done) + + events: Final = asyncio.run(_session(candidate, model, scenario.key(), exchange)) + + assert not any(_is_guardrail_error(event) for event in events), events + assert _text_deltas(events) == VISIBLE_ANSWER, events diff --git a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py index 2712d6e30cd..89d6ff50d53 100644 --- a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py +++ b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py @@ -1591,13 +1591,6 @@ def _assert_voice_session_ungated(session: Session, observed: tuple[dict[str, Js assert _session_updates(observed) == (VOICE_UPDATE,), _session_updates(observed) -@pytest.mark.xfail( - strict=True, - reason=( - "the realtime route hands the transcript guardrail only the request's guardrails list, so the key and " - "team metadata that opts out of a default-on guardrail never reaches it" - ), -) def test_key_opted_out_of_the_transcript_guardrail_is_not_gated_by_the_prompt_guardrail( guardrail_proxy: OwnedProxy, ) -> None: @@ -1609,13 +1602,6 @@ def test_key_opted_out_of_the_transcript_guardrail_is_not_gated_by_the_prompt_gu _assert_voice_session_ungated(session, _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)) -@pytest.mark.xfail( - strict=True, - reason=( - "the realtime route hands the transcript guardrail only the request's guardrails list, so the key and " - "team metadata that opts out of a default-on guardrail never reaches it" - ), -) def test_team_opted_out_of_the_transcript_guardrail_is_not_gated(guardrail_proxy: OwnedProxy) -> None: with guardrail_proxy.gateway.scenario() as scenario: handle: Final = _scripted(scenario, _voice_scenario(BLOCKED_TRANSCRIPT)) diff --git a/tests/integration/routing/test_prompt_caching_affinity_tool_call_marks.py b/tests/integration/routing/test_prompt_caching_affinity_tool_call_marks.py index 9d4bf62adf5..38129013294 100644 --- a/tests/integration/routing/test_prompt_caching_affinity_tool_call_marks.py +++ b/tests/integration/routing/test_prompt_caching_affinity_tool_call_marks.py @@ -2,7 +2,7 @@ from pathlib import Path from typing import Final import pytest -from integration._support.client import Gateway, eventually +from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows, scratch_database from integration._support.process import owned_proxy_process from integration._support.wire import Wire, wire_server @@ -24,6 +24,7 @@ from pydantic import JsonValue _MODEL: Final = "affinity-claude" _FOLLOW_UPS: Final = 24 +_PINS_ROUTE: Final = "/memory-usage-in-mem-cache-items" _LONG_SYSTEM: Final = "Answer from the tool results below. " + "lorem ipsum " * 1500 @@ -91,3 +92,68 @@ def test_router_affinity_is_lost_when_auto_caching_stands_down_for_tool_call_mar assert [status for status, _, _ in follow_ups] == [200] * _FOLLOW_UPS, [text for _, _, text in follow_ups] assert {first_marker, *follow_up_markers} == served["left"] | served["right"] assert all(served[side] & set(follow_up_markers) for side in served), served + + +def _responses_first_turn(marker: str) -> list[JsonValue]: + return [ + {"role": "user", "content": "What is the weather in London?"}, + {"type": "function_call", "call_id": "call_london", "name": "lookup_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_london", "output": "sunny"}, + {"role": "user", "content": final_text(marker)}, + ] + + +def _responses_body(session: str, request_input: list[JsonValue]) -> dict[str, JsonValue]: + return { + "model": _MODEL, + "instructions": f"{_LONG_SYSTEM} session {session}", + "tools": [{"type": "function", "name": "lookup_weather", "parameters": {"type": "object", "properties": {}}}], + "input": request_input, + } + + +def _pins(gateway: Gateway) -> dict[str, JsonValue]: + return object_value(object_value(gateway.request("GET", _PINS_ROUTE).json())["prompt_caching_pin_cache"]) + + +@pytest.mark.timeout(240) +def test_responses_thread_holds_its_deployment_when_the_bridge_logs_chat_messages( + gateway: Gateway, tmp_path: Path +) -> None: + session: Final = new_marker() + first_marker: Final = new_marker() + follow_up_markers: Final = tuple(new_marker() for _ in range(_FOLLOW_UPS)) + first_turn: Final = _responses_first_turn(first_marker) + with scratch_database() as database_url, wire_server(anthropic_peer) as left, wire_server(anthropic_peer) as right: + config: Final = owned_config( + tmp_path, + [ + {**anthropic_deployment(_MODEL, left.url), "model_info": {"id": f"responses-left-{session}"}}, + {**anthropic_deployment(_MODEL, right.url), "model_info": {"id": f"responses-right-{session}"}}, + ], + router_settings={"optional_pre_call_checks": ["prompt_caching"]}, + ) + with ( + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": database_url}, + config=config, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=1, + ) as owned, + owned.gateway.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"enable_prompt_caching": True}) + first: Final = owned.gateway.request("POST", "/v1/responses", _responses_body(session, first_turn), key=key) + assert first.status_code == 200, first.text + eventually(lambda: _pins(owned.gateway), bool, seconds=30) + follow_ups: Final = tuple( + owned.gateway.request( + "POST", "/v1/responses", _responses_body(session, _follow_up(first_turn, marker)), key=key + ) + for marker in follow_up_markers + ) + served: Final = {"left": _served(left), "right": _served(right)} + assert [response.status_code for response in follow_ups] == [200] * _FOLLOW_UPS, follow_ups[0].text + assert {first_marker, *follow_up_markers} in tuple(served.values()), served 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/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 19aefb430c7..6f47699a3b4 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1075,6 +1075,48 @@ def test_streamed_tool_search_turn_replays_with_its_reasoning_item(): assert request["input"][5]["namespace"] == "get_weather" +def test_typed_tool_search_output_keeps_the_wire_name_of_async_for_the_replay() -> None: + tool_search_output: Final = { + "type": "tool_search_output", + "id": "tso_1", + "execution": "server", + "status": "completed", + "tools": [{"type": "function", "name": "get_weather", "parameters": {}, "strict": True, "async": True}], + } + answer: Final = ResponsesAPIResponse(id="resp_1", created_at=1, output=[tool_search_output]) + + message: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices(answer.output)[0].message + + assert message.provider_specific_fields == {"tool_search_items": [tool_search_output]} + + +def test_streamed_tool_search_output_replays_only_the_tool_fields_upstream_sent() -> None: + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIStreamEvents + + tool_search_output: Final = { + "type": "tool_search_output", + "id": "tso_1", + "execution": "server", + "status": "completed", + "tools": [ + {"type": "function", "name": "get_weather", "parameters": {}, "strict": None, "defer_loading": True}, + {"type": "namespace", "name": "crm", "description": "CRM", "tools": [{"type": "function", "name": "find"}]}, + ], + } + completed: Final = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse(id="resp_1", created_at=1, status="completed", output=[tool_search_output]), + ) + iterator: Final = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + + chunk: Final = iterator._handle_string_chunk(completed) + + assert chunk.choices[0].delta.provider_specific_fields == {"tool_search_items": [tool_search_output]} + + @pytest.mark.parametrize( "answer_item", [ diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 66646ed10f7..3107d1e1f43 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -5,7 +5,7 @@ import os import sys 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 functools import partial from importlib.metadata import version @@ -628,3 +628,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_realtime_guardrails.py b/tests/unit/litellm_core_utils/test_realtime_guardrails.py new file mode 100644 index 00000000000..26ab2887883 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_realtime_guardrails.py @@ -0,0 +1,481 @@ +import json +import sys +from typing import Final, Literal + +import pytest +from fastapi import HTTPException + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import CustomCodeGuardrail +from litellm.litellm_core_utils.realtime_guardrails import ( + RealtimeGuardrailBlock, + RealtimeGuardrailRewrite, + RealtimeMessageText, + realtime_guarded_item, + realtime_guardrail_verdict, + rewritten_item, +) +from litellm.types.guardrails import BlockedWord, ContentFilterAction, ContentFilterPattern, GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + + +class RaisingGuardrail(CustomGuardrail): + def __init__(self, error: Exception, realtime_violation_message: str | None = None) -> None: + super().__init__( + guardrail_name="raising", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + realtime_violation_message=realtime_violation_message, + ) + self.error = error + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + raise self.error + + +class MaskingGuardrail(CustomGuardrail): + def __init__(self, word: str) -> None: + super().__init__(guardrail_name=f"mask-{word}", event_hook=GuardrailEventHooks.pre_call, default_on=True) + self.word = word + self.screened: list[tuple[str, ...]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + texts: Final = inputs.get("texts", []) + self.screened.append(tuple(texts)) + return {**inputs, "texts": [text.replace(self.word, "[MASKED]") for text in texts]} + + +class LastTextRedactingGuardrail(MaskingGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + last: Final = inputs.get("texts", [""])[-1] + if self.word not in last: + return inputs + return {**inputs, "texts": [last.replace(self.word, "[MASKED]")]} + + +def _item_create(item: object) -> dict: + return {"type": "conversation.item.create", "item": item} + + +@pytest.mark.parametrize("role", ["user", "assistant", "system"]) +def test_only_the_typed_text_of_a_message_is_screened_whatever_its_role(role: str) -> None: + client_event: Final = _item_create( + { + "type": "message", + "role": role, + "content": [{"type": "input_audio", "audio": "AAAA"}, {"type": "input_text", "text": "hello"}], + } + ) + + assert realtime_guarded_item(client_event) == RealtimeMessageText(role=role, texts=("hello",)) + + +@pytest.mark.parametrize("part_type", ["output_text", "text"]) +def test_assistant_text_written_by_the_client_is_screened(part_type: str) -> None: + client_event: Final = _item_create( + {"type": "message", "role": "assistant", "content": [{"type": part_type, "text": "hello"}]} + ) + + assert realtime_guarded_item(client_event) == RealtimeMessageText(role="assistant", texts=("hello",)) + + +@pytest.mark.parametrize( + "client_event", + [ + pytest.param( + _item_create({"type": "message", "role": "user", "content": [{"type": "input_audio", "audio": "AAAA"}]}), + id="audio-only user message", + ), + pytest.param( + _item_create({"type": "function_call_output", "call_id": "call_1", "output": ""}), id="empty tool output" + ), + pytest.param(_item_create("hello"), id="item that is not an object"), + ], +) +def test_items_without_client_text_are_not_screened(client_event: dict) -> None: + assert realtime_guarded_item(client_event) is None + + +@pytest.mark.asyncio +async def test_a_plain_string_block_reason_becomes_the_violation_message() -> None: + guardrail: Final = RaisingGuardrail(HTTPException(status_code=400, detail="hello is not allowed")) + + block: Final = await realtime_guardrail_verdict( + ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + + assert isinstance(block, RealtimeGuardrailBlock) + assert block.message == "hello is not allowed" + + +@pytest.mark.asyncio +async def test_a_configured_violation_message_replaces_the_block_reason() -> None: + guardrail: Final = RaisingGuardrail(HTTPException(status_code=400, detail={"error": "keyword 'hello' detected"})) + configured: Final = RaisingGuardrail( + HTTPException(status_code=400, detail="hello is not allowed"), realtime_violation_message="Please rephrase" + ) + + block: Final = await realtime_guardrail_verdict( + ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + configured_block: Final = await realtime_guardrail_verdict( + ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (configured,) + ) + + assert isinstance(block, RealtimeGuardrailBlock) + assert block.message == "keyword 'hello' detected" + assert isinstance(configured_block, RealtimeGuardrailBlock) + assert configured_block.message == "Please rephrase" + + +@pytest.mark.asyncio +async def test_masked_texts_come_back_as_a_rewrite_and_reach_later_guardrails_masked() -> None: + first: Final = MaskingGuardrail("hunter2") + second: Final = MaskingGuardrail("swordfish") + + verdict: Final = await realtime_guardrail_verdict( + ("my hunter2", "and swordfish", "hi"), "user", (GuardrailEventHooks.pre_call,), {}, None, (first, second) + ) + + assert verdict == RealtimeGuardrailRewrite(texts=("my [MASKED]", "and [MASKED]", "hi")) + assert sorted(second.screened[:2]) == [("my [MASKED] and swordfish hi",), ("my [MASKED]and swordfishhi",)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("parts", "screened"), + [ + pytest.param( + ("Compare this photo ", "with that one"), + ("Compare this photo with that one",), + id="a space at the boundary", + ), + pytest.param( + ("Compare this photo", "with that one"), + ("Compare this photowith that one", "Compare this photo with that one"), + id="no space at the boundary", + ), + ], +) +async def test_a_clean_multi_part_item_sends_each_guardrail_only_the_joins_it_needs_one_text_per_call( + parts: tuple[str, ...], screened: tuple[str, ...] +) -> None: + guardrail: Final = MaskingGuardrail("hunter2") + + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + + assert verdict is None + assert sorted(guardrail.screened) == sorted((text,) for text in screened) + + +@pytest.mark.asyncio +async def test_texts_a_guardrail_leaves_unchanged_need_no_rewrite() -> None: + verdict: Final = await realtime_guardrail_verdict( + ("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail("hunter2"),) + ) + + assert verdict is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("parts", "phrase"), + [ + pytest.param(("tell me the secret", "plan"), "secret plan", id="split at a word break"), + pytest.param(("tell me about secret-", "project"), "secret-project", id="split mid-word"), + ], +) +async def test_a_phrase_split_across_typed_text_parts_is_still_blocked(parts: tuple[str, str], phrase: str) -> None: + client_event: Final = _item_create( + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": part} for part in parts]} + ) + guarded_item: Final = realtime_guarded_item(client_event) + assert isinstance(guarded_item, RealtimeMessageText) + + phrase_filter: Final = ContentFilterGuardrail( + guardrail_name="phrase-filter", + blocked_words=[BlockedWord(keyword=phrase, action=ContentFilterAction.BLOCK)], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + verdict: Final = await realtime_guardrail_verdict( + guarded_item.texts, guarded_item.role, (GuardrailEventHooks.pre_call,), {}, None, (phrase_filter,) + ) + + assert isinstance(verdict, RealtimeGuardrailBlock) + + +@pytest.mark.asyncio +async def test_a_mask_inside_one_part_keeps_every_part_in_place() -> None: + first_photo: Final = {"type": "input_image", "image_url": "data:image/png;base64,AAA"} + second_photo: Final = {"type": "input_image", "image_url": "data:image/png;base64,BBB"} + client_event: Final = _item_create( + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Compare this photo"}, + first_photo, + {"type": "input_text", "text": "with hunter2's photo"}, + second_photo, + ], + } + ) + guarded_item: Final = realtime_guarded_item(client_event) + assert isinstance(guarded_item, RealtimeMessageText) + + verdict: Final = await realtime_guardrail_verdict( + guarded_item.texts, "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail("hunter2"),) + ) + + assert isinstance(verdict, RealtimeGuardrailRewrite) + assert rewritten_item(client_event, verdict.texts)["content"] == [ + {"type": "input_text", "text": "Compare this photo"}, + first_photo, + {"type": "input_text", "text": "with [MASKED]'s photo"}, + second_photo, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("parts", "word", "masked"), + [ + pytest.param(("my secret", "plan is here"), "secret plan", "my [MASKED] is here", id="split at a word break"), + pytest.param(("my code is hun", "ter2 ok"), "hunter2", "my code is [MASKED] ok", id="split mid-word"), + ], +) +async def test_a_mask_across_parts_replaces_them_with_the_masked_text_the_model_would_have_read( + parts: tuple[str, ...], word: str, masked: str +) -> None: + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail(word),) + ) + + assert verdict == RealtimeGuardrailRewrite(texts=(masked,)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "parts", + [ + pytest.param(("abc", "123-45-6789 hunter2 rocks"), id="masks inside one part"), + pytest.param(("abc", "123-45-6789 hunter2", "rocks"), id="a mask inside a part and one across parts"), + ], +) +async def test_no_mask_one_text_part_gets_is_lost_when_its_joined_text_masks_differently( + parts: tuple[str, ...], +) -> None: + content_filter: Final = ContentFilterGuardrail( + guardrail_name="ssn-and-keyword-mask", + patterns=[ + ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.MASK) + ], + blocked_words=[BlockedWord(keyword="hunter2 rocks", action=ContentFilterAction.MASK)], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (content_filter,) + ) + + assert isinstance(verdict, RealtimeGuardrailRewrite) + forwarded: Final = " ".join(verdict.texts) + assert "123-45-6789" not in forwarded + assert "hunter2 rocks" not in forwarded + assert forwarded.startswith("abc") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("parts", "rewritten"), + [ + pytest.param(("my code is", "hunter2"), ("my code is", "[MASKED]"), id="mask inside one part"), + pytest.param(("my code is hun", "ter2"), ("my code is [MASKED]",), id="mask across parts"), + ], +) +async def test_a_guardrail_that_answers_with_one_text_still_masks_parts_joined_without_a_space( + parts: tuple[str, ...], rewritten: tuple[str, ...] +) -> None: + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (LastTextRedactingGuardrail("hunter2"),) + ) + + assert verdict == RealtimeGuardrailRewrite(texts=rewritten) + + +@pytest.mark.asyncio +async def test_a_mask_in_a_message_with_more_text_parts_than_the_recursion_limit_still_comes_back() -> None: + parts: Final = (*(f"part {index}." for index in range(sys.getrecursionlimit())), "my code is hunter2.") + + verdict: Final = await realtime_guardrail_verdict( + parts, "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail("hunter2"),) + ) + + assert verdict == RealtimeGuardrailRewrite(texts=(*parts[:-1], "my code is [MASKED].")) + + +class GuardrailLogRecordingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="log-recorder", event_hook=GuardrailEventHooks.pre_call, default_on=True) + self.guardrail_log_lengths: list[int] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + self.guardrail_log_lengths.append( + len(request_data["litellm_metadata"].get("standard_logging_guardrail_information", [])) + ) + return inputs + + +@pytest.mark.asyncio +async def test_a_guardrail_sees_only_the_current_turns_guardrail_log_however_long_the_session_runs() -> None: + session_request_data: Final = {"litellm_metadata": {"user_api_key_team_id": "team-a"}} + content_filter: Final = ContentFilterGuardrail( + guardrail_name="phrase-filter", + blocked_words=[BlockedWord(keyword="hunter2", action=ContentFilterAction.BLOCK)], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + recorder: Final = GuardrailLogRecordingGuardrail() + + for turn in ("hello", "how are you", "goodbye"): + await realtime_guardrail_verdict( + (turn,), "user", (GuardrailEventHooks.pre_call,), session_request_data, None, (content_filter, recorder) + ) + + assert recorder.guardrail_log_lengths == [1, 1, 1] + assert session_request_data == {"litellm_metadata": {"user_api_key_team_id": "team-a"}} + + +class ScopedGuardrail(RaisingGuardrail): + def __init__( + self, *, skip_system_message: bool | None, skip_tool_message: bool | None, scan_only_tool_results: bool | None + ) -> None: + super().__init__(HTTPException(status_code=400, detail="blocked")) + self.skip_system_message_in_guardrail = skip_system_message + self.skip_tool_message_in_guardrail = skip_tool_message + self.scan_only_tool_results = scan_only_tool_results + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "screened_roles"), + [ + pytest.param({}, ("user", "assistant", "system", "tool"), id="no scope"), + pytest.param({"skip_system_message": True}, ("user", "assistant", "tool"), id="skip system"), + pytest.param({"skip_tool_message": True}, ("user", "assistant", "system"), id="skip tool"), + pytest.param({"scan_only_tool_results": True}, ("tool",), id="only tool results"), + ], +) +async def test_a_guardrail_screens_only_the_roles_its_scope_allows( + scope: dict[str, bool], screened_roles: tuple[str, ...] +) -> None: + guardrail: Final = ScopedGuardrail( + skip_system_message=scope.get("skip_system_message"), + skip_tool_message=scope.get("skip_tool_message"), + scan_only_tool_results=scope.get("scan_only_tool_results"), + ) + + verdicts: Final = { + role: await realtime_guardrail_verdict( + ("hello",), role, (GuardrailEventHooks.pre_call,), {}, None, (guardrail,) + ) + for role in ("user", "assistant", "system", "tool") + } + + assert tuple(role for role, verdict in verdicts.items() if verdict is not None) == screened_roles + + +def test_a_rewrite_puts_the_masked_text_in_place_of_the_typed_text_parts() -> None: + audio: Final = {"type": "input_audio", "audio": "AAAA"} + client_event: Final = _item_create( + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "a"}, audio, {"type": "input_text", "text": "b"}], + } + ) + + assert rewritten_item(client_event, ("[MASKED] b",)) == { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "[MASKED] b"}, audio], + } + + +def test_a_rewrite_replaces_the_tool_output() -> None: + client_event: Final = _item_create({"type": "function_call_output", "call_id": "call_1", "output": "hunter2"}) + + assert rewritten_item(client_event, ("[MASKED]",)) == { + "type": "function_call_output", + "call_id": "call_1", + "output": "[MASKED]", + } + + +@pytest.mark.asyncio +async def test_a_guardrail_that_crashes_is_not_treated_as_a_block() -> None: + guardrail: Final = RaisingGuardrail(RuntimeError("guardrail backend down")) + + with pytest.raises(RuntimeError, match="guardrail backend down"): + await realtime_guardrail_verdict(("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,)) + + +_TEAM_POLICY_CODE: Final = ( + "def apply_guardrail(inputs, request_data, input_type):\n" + ' if request_data["team_id"] == "team-strict":\n' + ' return block("blocked for this team")\n' + " return allow()\n" +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_id", "blocked"), + [pytest.param("team-strict", True, id="policy team"), pytest.param("team-other", False, id="other team")], +) +async def test_a_team_scoped_guardrail_sees_the_session_team(team_id: str, blocked: bool) -> None: + team_policy = CustomCodeGuardrail( + custom_code=_TEAM_POLICY_CODE, guardrail_name="team-policy", event_hook="pre_call", default_on=True + ) + + verdict = await realtime_guardrail_verdict( + ("hello",), + "user", + (GuardrailEventHooks.pre_call,), + {"litellm_metadata": {"user_api_key_team_id": team_id}}, + None, + (team_policy,), + ) + + assert isinstance(verdict, RealtimeGuardrailBlock) is blocked diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index eec1dfcf459..7f205b0bd00 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -1056,6 +1056,117 @@ async def test_transcription_captured_in_backend_to_client(): assert logging_obj.model_call_details["messages"] == streaming.input_messages +@pytest.mark.asyncio +@pytest.mark.parametrize("transcript", [None, 123, ["a"], {"text": "a"}]) +async def test_a_non_string_transcript_does_not_end_the_session(transcript: object) -> None: + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + transcript_event: Final = json.dumps( + {"type": "conversation.item.input_audio_transcription.completed", "transcript": transcript, "item_id": "i"} + ) + later_event: Final = json.dumps({"type": "response.done", "response": {"id": "resp_1"}}) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[transcript_event.encode(), later_event.encode(), ConnectionClosed(None, None)] + ) + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + + await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() + + sent_types: Final = [json.loads(call.args[0])["type"] for call in client_ws.send_text.call_args_list] + assert sent_types == ["conversation.item.input_audio_transcription.completed", "response.done"] + backend_ws.send.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transcript", [None, 123, ["a"]]) +async def test_a_non_string_transcript_fails_the_session_when_a_transcript_guardrail_would_screen_it( + transcript: object, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + transcript_event: Final = json.dumps( + {"type": "conversation.item.input_audio_transcription.completed", "transcript": transcript, "item_id": "i"} + ) + later_event: Final = json.dumps({"type": "response.done", "response": {"id": "resp_1"}}) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[transcript_event.encode(), later_event.encode(), ConnectionClosed(None, None)] + ) + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + + close: Final = await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() + + assert close.code == 1011 + sent_types: Final = [json.loads(call.args[0])["type"] for call in client_ws.send_text.call_args_list] + assert sent_types == ["conversation.item.input_audio_transcription.completed"] + backend_ws.send.assert_not_called() + + +def _native_lifecycle_transcription_guardrail() -> CustomGuardrail: + class _NativeLifecycleGuardrail(CustomGuardrail): + use_native_lifecycle_hooks = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + return _NativeLifecycleGuardrail( + guardrail_name="test_native_guard", + event_hook=GuardrailEventHooks.realtime_input_transcription, + default_on=True, + ) + + +def _tool_results_only_transcription_guardrail() -> CustomGuardrail: + class _ToolResultsOnlyGuardrail(CustomGuardrail): + scan_only_tool_results = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + return _ToolResultsOnlyGuardrail( + guardrail_name="test_tool_results_guard", + event_hook=GuardrailEventHooks.realtime_input_transcription, + default_on=True, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "guardrail_factory", [_native_lifecycle_transcription_guardrail, _tool_results_only_transcription_guardrail] +) +async def test_a_non_string_transcript_keeps_the_session_when_no_guardrail_would_screen_the_users_turn( + guardrail_factory: Callable[[], CustomGuardrail], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "callbacks", [guardrail_factory()]) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + transcript_event: Final = json.dumps( + {"type": "conversation.item.input_audio_transcription.completed", "transcript": None, "item_id": "i"} + ) + later_event: Final = json.dumps({"type": "response.done", "response": {"id": "resp_1"}}) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[transcript_event.encode(), later_event.encode(), ConnectionClosed(None, None)] + ) + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + + await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() + + sent_types: Final = [json.loads(call.args[0])["type"] for call in client_ws.send_text.call_args_list] + assert sent_types == ["conversation.item.input_audio_transcription.completed", "response.done"] + + @pytest.mark.asyncio async def test_transcription_session_captures_usage_and_skips_response_create(): """ @@ -1955,6 +2066,112 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error(monkeypatc assert len(original_items) == 0, f"Blocked item should not be forwarded to backend, got: {original_items}" +@pytest.mark.asyncio +async def test_the_next_clean_turn_is_answered_after_a_blocked_item_with_no_response_create( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import HTTPException + + class SecretBlockingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + if any("hunter2" in text for text in inputs.get("texts", [])): + raise HTTPException(status_code=400, detail={"error": "hunter2 is not allowed"}) + return inputs + + def item_create(role: str, text: str) -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": role, "content": [{"type": "input_text", "text": text}]}, + } + ) + + monkeypatch.setattr( + litellm, + "callbacks", + [SecretBlockingGuardrail(guardrail_name="no-secrets", event_hook=GuardrailEventHooks.pre_call, default_on=True)], + ) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + client_ws.receive_text = AsyncMock( + side_effect=[ + item_create("system", "never mention hunter2"), + item_create("user", "what is the weather"), + json.dumps({"type": "response.create"}), + Exception("connection closed"), + ] + ) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + streaming: Final = RealTimeStreaming(client_ws, backend_ws, _mock_logging()) + + await streaming.client_ack_messages() + + sent_to_backend: Final = [json.loads(call.args[0]) for call in backend_ws.send.call_args_list] + clean_index: Final = next( + index + for index, event in enumerate(sent_to_backend) + if event.get("item", {}).get("content", [{}])[0].get("text") == "what is the weather" + ) + assert [event["type"] for event in sent_to_backend[clean_index + 1 :]] == ["response.create"], sent_to_backend + + +@pytest.mark.asyncio +async def test_realtime_text_input_masking_guardrail_forwards_only_the_masked_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class MaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + texts: Final = inputs.get("texts", []) + return {**inputs, "texts": [text.replace("test@example.com", "[EMAIL]") for text in texts]} + + monkeypatch.setattr( + litellm, + "callbacks", + [MaskingGuardrail(guardrail_name="email-masker", event_hook=GuardrailEventHooks.pre_call, default_on=True)], + ) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "conversation.item.create", + "item": { + "role": "user", + "content": [{"type": "input_text", "text": "My email is test@example.com"}], + }, + } + ), + Exception("connection closed"), + ] + ) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None)) + logging_obj: Final = _mock_logging() + logging_obj.pre_call = MagicMock() + + await RealTimeStreaming(client_ws, backend_ws, logging_obj).client_ack_messages() + + sent_to_backend: Final = json.dumps([c.args[0] for c in backend_ws.send.call_args_list if c.args]) + assert "test@example.com" not in sent_to_backend + assert "My email is [EMAIL]" in sent_to_backend + assert client_ws.send_text.call_args_list == [] + + @pytest.mark.asyncio async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index ae2b8520ad4..d2f16a5fe41 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -3528,6 +3528,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", @@ -4660,6 +4688,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/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 8affe2e136e..b6229dbaec0 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -3458,3 +3458,41 @@ def test_reading_a_dropped_budget_for_a_guardrail_logs_no_drop_warning(caplog: p assert messages == [{"role": "user", "content": "hi"}] assert not [record for record in caplog.records if "Dropping" in record.getMessage()] + + +FAILED_TOOL_RESULT_REQUEST: Final = { + "model": "claude-sonnet", + "max_tokens": 16, + "messages": [ + {"role": "user", "content": "Rename the helper"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "edit0", "name": "Edit", "input": {"file_path": "a.py"}}], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "edit0", "content": "String not found", "is_error": True} + ], + }, + ], +} + + +@pytest.mark.parametrize("keep_tool_errors", [True, False]) +def test_structured_messages_carry_tool_failures_only_when_asked(keep_tool_errors: bool) -> None: + messages: Final = AnthropicMessagesHandler().get_structured_messages( + FAILED_TOOL_RESULT_REQUEST, keep_tool_errors=keep_tool_errors + ) + assert messages is not None + (tool_message,) = [message for message in messages if message["role"] == "tool"] + assert tool_message.get("is_error") is (True if keep_tool_errors else None) + + +def test_request_sent_to_a_model_never_carries_tool_failure_marker() -> None: + from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter + + request, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + anthropic_message_request=FAILED_TOOL_RESULT_REQUEST # pyright: ignore[reportArgumentType] # a JSON request in the Anthropic shape + ) + assert "is_error" not in json.dumps(request["messages"], default=str) 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/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index cf2f25faaa9..1221667cb51 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -10,23 +10,38 @@ from typing import Final, Literal from unittest.mock import MagicMock import pytest +from fastapi import HTTPException from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict, Unpack import litellm from litellm.constants import ( BEDROCK_REALTIME_SDK_SUPPORTED_RANGE, + REALTIME_SESSION_ENDED_ABNORMALLY_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, WEBSOCKET_CLOSE_REASON_MAX_BYTES, ) +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.realtime_errors import websocket_close_reason +from litellm.litellm_core_utils.realtime_guardrails import BLOCKED_TOOL_OUTPUT from litellm.llms.bedrock.common_utils import BedrockError -from litellm.llms.bedrock.realtime.handler import BedrockRealtime, SessionEndSignals -from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.llms.bedrock.realtime.handler import ( + GUARDRAIL_UNAVAILABLE_MESSAGE, + BedrockRealtime, + BedrockRealtimeGuardrails, + SessionEndSignals, +) from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.realtime_api import main as realtime_main -from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import StandardLoggingPayload, StandardLoggingPayloadErrorInformation +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import ( + GenericGuardrailAPIInputs, + StandardLoggingPayload, + StandardLoggingPayloadErrorInformation, +) from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault @@ -223,10 +238,19 @@ class DisconnectingClientWS: return self._messages.pop(0) raise RuntimeError("client disconnected") - async def send_text(self, message): + async def send_text(self, message: str) -> None: self.sent_to_client.append(message) +class GuardrailErrorUndeliveredClientWS(DisconnectingClientWS): + """Client that disconnects just as the proxy tries to tell it about a guardrail violation""" + + async def send_text(self, message: str) -> None: + if "guardrail_violation" in message: + raise RuntimeError("client disconnected") + await super().send_text(message) + + class ClosableClientWS: def __init__(self): self.closed = False @@ -1052,7 +1076,16 @@ class TestBedrockRealtimeProviderFailurePropagation: assert other_provider.calls == 0 @pytest.mark.asyncio - async def test_readiness_failure_on_an_input_first_start_replays_the_input(self, stub_aws_sdk_client): + @pytest.mark.parametrize( + "guardrail_awaits_its_backend", [False, True], ids=["no guardrail", "a guardrail that awaits its backend"] + ) + async def test_readiness_failure_on_an_input_first_start_replays_the_input( + self, stub_aws_sdk_client, monkeypatch, guardrail_awaits_its_backend: bool + ): + guardrail: Final = AwaitingGuardrail( + "hunter2", guardrail_name="remote", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail] if guardrail_awaits_its_backend else []) handler = BedrockRealtime() websocket = RealtimeClientWS([self.USER_TEXT]) healthy_stream = SessionBoundBedrockStream() @@ -1687,5 +1720,614 @@ class TestBedrockRealtimeSdkImportErrors: assert BEDROCK_REALTIME_SDK_SUPPORTED_RANGE in close_reason +class SessionEndOptions(TypedDict, total=False): + on_violation: ReadOnly[str] + end_session_after_n_fails: ReadOnly[int] + + +class BlockedWordGuardrailOptions(SessionEndOptions, total=False): + guardrail_name: ReadOnly[str] + event_hook: ReadOnly[GuardrailEventHooks] + default_on: ReadOnly[bool] + realtime_violation_message: ReadOnly[str] + + +class BlockedWordGuardrail(CustomGuardrail): + def __init__(self, blocked_word: str, **kwargs: Unpack[BlockedWordGuardrailOptions]) -> None: + super().__init__(**kwargs) + self.blocked_word = blocked_word + self.checked: list[tuple[tuple[str, ...], object]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + self.checked.append((tuple(inputs["texts"]), request_data.get("user_api_key_dict"))) + if any(self.blocked_word in text for text in inputs["texts"]): + raise HTTPException(status_code=400, detail={"error": f"{self.blocked_word} is not allowed"}) + return inputs + + +class MaskingGuardrail(BlockedWordGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "texts": [text.replace(self.blocked_word, "[MASKED]") for text in inputs["texts"]]} + + +class AwaitingGuardrail(BlockedWordGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + await asyncio.sleep(0.05) + return await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + + +class SystemSkippingGuardrail(BlockedWordGuardrail): + skip_system_message_in_guardrail: Final = True + + +class CrashingGuardrail(BlockedWordGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLogging | None = None, + ) -> GenericGuardrailAPIInputs: + raise RuntimeError("guardrail backend down") + + +class TestBedrockRealtimeGuardrails: + """Realtime guardrails must screen Nova Sonic sessions the way they screen OpenAI realtime sessions""" + + MODEL: Final = "amazon.nova-sonic-v1:0" + VIOLATION: Final = "I can't help with that." + BLOCKED_TEXT: Final = "my secret is hunter2" + CLEAN_TEXT: Final = "what is the weather like" + RESPONSE_CREATE: Final = json.dumps({"type": "response.create"}) + + @staticmethod + def _guardrail( + event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call, + default_on: bool = True, + **kwargs: Unpack[SessionEndOptions], + ) -> BlockedWordGuardrail: + return BlockedWordGuardrail( + "hunter2", + guardrail_name="no_secrets", + event_hook=event_hook, + default_on=default_on, + realtime_violation_message=TestBedrockRealtimeGuardrails.VIOLATION, + **kwargs, + ) + + @staticmethod + def _user_text(text: str, role: str = "user") -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": role, "content": [{"type": "input_text", "text": text}]}, + } + ) + + @staticmethod + def _bedrock_events(stream: FakeBedrockStream) -> list[dict]: + return [json.loads(chunk.value.bytes_.decode("utf-8"))["event"] for chunk in stream.input_stream.sent] + + @classmethod + def _text_inputs(cls, stream: FakeBedrockStream) -> list[str]: + events: Final = cls._bedrock_events(stream) + user_contents: Final = { + event["contentStart"]["contentName"] + for event in events + if event.get("contentStart", {}).get("role") == "USER" + } + return [ + event["textInput"]["content"] + for event in events + if event.get("textInput", {}).get("contentName") in user_contents + ] + + @staticmethod + def _client_events(client_ws: DisconnectingClientWS | RealtimeClientWS) -> list[dict]: + return [json.loads(message) for message in client_ws.sent_to_client] + + @classmethod + def _guardrail_errors(cls, client_ws: DisconnectingClientWS | RealtimeClientWS) -> list[dict]: + return [ + event["error"] + for event in cls._client_events(client_ws) + if event["type"] == "error" and event["error"]["type"] == "guardrail_violation" + ] + + async def _forward_client( + self, + messages: list[str], + guardrail: BlockedWordGuardrail, + client_ws_type: type[DisconnectingClientWS] | type[ConnectedClientWS] = DisconnectingClientWS, + ) -> tuple[FakeBedrockStream, DisconnectingClientWS | ConnectedClientWS]: + stream: Final = FakeBedrockStream() + client_ws: Final = client_ws_type(messages) + await asyncio.wait_for( + BedrockRealtime()._forward_client_to_bedrock( + client_ws, + stream, + BedrockRealtimeConfig(), + self.MODEL, + {}, + signals=SessionEndSignals(), + guardrails=BedrockRealtimeGuardrails(callbacks=lambda: (guardrail,)), + ), + timeout=1, + ) + return stream, client_ws + + @staticmethod + def _user_transcript(text: str) -> tuple[str, ...]: + return ( + json.dumps( + { + "event": { + "contentStart": { + "role": "USER", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "FINAL"}), + } + } + } + ), + json.dumps({"event": {"textOutput": {"content": text}}}), + json.dumps({"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}), + ) + + @staticmethod + def _assistant_text(text: str) -> tuple[str, ...]: + return ( + json.dumps({"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}), + json.dumps({"event": {"textOutput": {"content": text}}}), + json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN"}}}), + ) + + async def _forward_bedrock( + self, + payloads: tuple[str, ...], + guardrail: BlockedWordGuardrail, + client_audio_streamed: bool = True, + client_ws_type: type[ConnectedClientWS] | type[GuardrailErrorUndeliveredClientWS] = ConnectedClientWS, + ) -> tuple[ConnectedClientWS | GuardrailErrorUndeliveredClientWS, list[object]]: + config: Final = BedrockRealtimeConfig() + config.client_audio_streamed = client_audio_streamed + client_ws: Final = client_ws_type([]) + logged: Final = [ + event + async for event in BedrockRealtime()._forward_bedrock_to_client( + ScriptedBedrockStream(payloads), + client_ws, + config, + self.MODEL, + FakeLogging(), + {}, + signals=_closing_input(), + guardrails=BedrockRealtimeGuardrails(callbacks=lambda: (guardrail,)), + ) + ] + return client_ws, logged + + @pytest.mark.asyncio + async def test_blocked_typed_message_never_reaches_bedrock_and_the_violation_is_voiced( + self, stub_aws_models: None + ) -> None: + guardrail: Final = self._guardrail() + + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT), self.RESPONSE_CREATE], guardrail + ) + clean_stream, _ = await self._forward_client( + [self._user_text(self.CLEAN_TEXT), self.RESPONSE_CREATE], guardrail + ) + + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert [f'"{self.VIOLATION}"' in text for text in self._text_inputs(stream)] == [True] + assert self._guardrail_errors(client_ws) == [ + {"type": "guardrail_violation", "message": self.VIOLATION, "code": "content_policy_violation"} + ] + assert guardrail.checked == [((self.BLOCKED_TEXT,), None), ((self.CLEAN_TEXT,), None)] + audio_inputs = [sum("audioInput" in event for event in self._bedrock_events(s)) for s in (stream, clean_stream)] + assert audio_inputs[0] == audio_inputs[1] > 0, "the client's own response.create must not start a 2nd turn" + + @pytest.mark.asyncio + @pytest.mark.parametrize("role", ["assistant", "system"]) + async def test_blocked_typed_message_is_screened_whatever_its_role(self, stub_aws_models: None, role: str) -> None: + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT, role=role)], self._guardrail() + ) + + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + + @pytest.mark.asyncio + async def test_a_system_message_nova_receives_as_a_user_turn_is_screened_by_a_system_skipping_guardrail( + self, stub_aws_models: None + ) -> None: + guardrail: Final = SystemSkippingGuardrail( + "hunter2", + guardrail_name="no_secrets", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + realtime_violation_message=self.VIOLATION, + ) + + stream, client_ws = await self._forward_client([self._user_text(self.BLOCKED_TEXT, role="system")], guardrail) + + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + + @pytest.mark.asyncio + async def test_the_next_clean_turn_is_answered_after_a_blocked_item_with_no_response_create( + self, stub_aws_models: None + ) -> None: + stream, _ = await self._forward_client( + [ + self._user_text(self.BLOCKED_TEXT, role="system"), + self._user_text(self.CLEAN_TEXT), + self.RESPONSE_CREATE, + ], + self._guardrail(), + ) + + events: Final = self._bedrock_events(stream) + clean_text_index: Final = next( + i for i, event in enumerate(events) if event.get("textInput", {}).get("content") == self.CLEAN_TEXT + ) + assert any("audioInput" in event for event in events[clean_text_index:]), "the client's response.create runs" + + @pytest.mark.asyncio + async def test_client_message_the_guardrail_cannot_parse_never_reaches_bedrock( + self, stub_aws_models: None + ) -> None: + lone_surrogate_text: Final = self._user_text(f"{self.BLOCKED_TEXT} \ud800") + + stream, _ = await self._forward_client( + [lone_surrogate_text, self._user_text(self.CLEAN_TEXT)], self._guardrail() + ) + + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert self._text_inputs(stream) == [self.CLEAN_TEXT] + + @pytest.mark.asyncio + async def test_clean_typed_message_reaches_bedrock_unchanged(self, stub_aws_models: None) -> None: + guardrail: Final = self._guardrail() + + stream, client_ws = await self._forward_client([self._user_text(self.CLEAN_TEXT)], guardrail) + + assert self._text_inputs(stream) == [self.CLEAN_TEXT] + assert client_ws.sent_to_client == [] + assert guardrail.checked == [((self.CLEAN_TEXT,), None)] + + @pytest.mark.asyncio + async def test_masked_typed_message_reaches_bedrock_masked_and_is_answered(self, stub_aws_models: None) -> None: + guardrail: Final = MaskingGuardrail( + "hunter2", guardrail_name="mask_secrets", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT), self.RESPONSE_CREATE], guardrail + ) + + events: Final = self._bedrock_events(stream) + assert "hunter2" not in json.dumps(events) + assert self._text_inputs(stream) == ["my secret is [MASKED]"] + assert client_ws.sent_to_client == [] + assert any("audioInput" in event for event in events), "the client's response.create runs" + + @pytest.mark.asyncio + async def test_masked_tool_output_reaches_bedrock_masked(self, stub_aws_models: None) -> None: + guardrail: Final = MaskingGuardrail( + "hunter2", guardrail_name="mask_secrets", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + tool_output: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "function_call_output", "call_id": "tool_1", "output": self.BLOCKED_TEXT}, + } + ) + + stream, client_ws = await self._forward_client([tool_output], guardrail) + + events: Final = self._bedrock_events(stream) + assert [event["toolResult"]["content"] for event in events if "toolResult" in event] == [ + "my secret is [MASKED]" + ] + assert client_ws.sent_to_client == [] + + @pytest.mark.asyncio + async def test_the_block_reason_is_spoken_when_no_message_is_configured(self, stub_aws_models: None) -> None: + guardrail: Final = BlockedWordGuardrail( + "hunter2", guardrail_name="no_secrets", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT), self.RESPONSE_CREATE], guardrail + ) + + assert self.BLOCKED_TEXT not in json.dumps(self._bedrock_events(stream)) + assert ['"hunter2 is not allowed"' in text for text in self._text_inputs(stream)] == [True] + assert [error["message"] for error in self._guardrail_errors(client_ws)] == ["hunter2 is not allowed"] + + @pytest.mark.asyncio + async def test_guardrail_that_crashes_ends_the_session_as_an_outage_without_forwarding_the_text( + self, stub_aws_models: None + ) -> None: + guardrail: Final = CrashingGuardrail( + "hunter2", guardrail_name="flaky", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + stream: Final = SessionBoundBedrockStream() + client_ws: Final = ConnectedClientWS([self._user_text(self.CLEAN_TEXT), self.RESPONSE_CREATE]) + logging_obj: Final = FakeLogging() + + async def open_stream() -> SessionBoundBedrockStream: + return stream + + await asyncio.wait_for( + BedrockRealtime()._run_session( + client_ws, + open_stream, + self.MODEL, + logging_obj, + None, + guardrails=BedrockRealtimeGuardrails(callbacks=lambda: (guardrail,)), + ), + timeout=2, + ) + + assert self.CLEAN_TEXT not in json.dumps(self._bedrock_events(stream)) + assert [event["error"] for event in self._client_events(client_ws) if event["type"] == "error"] == [ + {"type": "guardrail_error", "message": GUARDRAIL_UNAVAILABLE_MESSAGE} + ] + assert client_ws.closed + assert logging_obj.model_call_details.get(REALTIME_SESSION_ENDED_ABNORMALLY_KEY) is True + + @pytest.mark.asyncio + async def test_blocked_tool_output_is_replaced_and_answered_by_the_model_itself( + self, stub_aws_models: None + ) -> None: + tool_output: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "function_call_output", "call_id": "tool_1", "output": self.BLOCKED_TEXT}, + } + ) + + stream, client_ws = await self._forward_client([tool_output, self.RESPONSE_CREATE], self._guardrail()) + + events: Final = self._bedrock_events(stream) + assert "hunter2" not in json.dumps(events) + tool_result_index: Final = next(i for i, event in enumerate(events) if "toolResult" in event) + assert events[tool_result_index]["toolResult"]["content"] == BLOCKED_TOOL_OUTPUT + tool_starts: Final = [ + event["contentStart"] for event in events if event.get("contentStart", {}).get("type") == "TOOL" + ] + assert [start["toolResultInputConfiguration"]["toolUseId"] for start in tool_starts] == ["tool_1"] + assert self._text_inputs(stream) == [] + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + assert any("audioInput" in event for event in events[tool_result_index:]), "the client's response.create runs" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("event_hook", "default_on"), + [ + (GuardrailEventHooks.pre_call, False), + (GuardrailEventHooks.post_call, True), + (GuardrailEventHooks.realtime_input_transcription, True), + ], + ids=["not-default-on", "post-call-hook", "transcription-hook"], + ) + async def test_guardrail_that_does_not_apply_to_typed_text_changes_nothing( + self, stub_aws_models: None, event_hook: GuardrailEventHooks, default_on: bool + ) -> None: + guardrail: Final = self._guardrail(event_hook=event_hook, default_on=default_on) + + stream, client_ws = await self._forward_client([self._user_text(self.BLOCKED_TEXT)], guardrail) + + assert self._text_inputs(stream) == [self.BLOCKED_TEXT] + assert client_ws.sent_to_client == [] + assert guardrail.checked == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("guardrail_kwargs", "blocked_messages"), + [({"on_violation": "end_session"}, 1), ({"end_session_after_n_fails": 2}, 2)], + ids=["on-violation-end-session", "after-two-fails"], + ) + async def test_end_session_guardrail_closes_the_session( + self, stub_aws_models: None, guardrail_kwargs: SessionEndOptions, blocked_messages: int + ) -> None: + messages: Final = [self._user_text(self.BLOCKED_TEXT)] * blocked_messages + [self._user_text(self.CLEAN_TEXT)] + + stream, client_ws = await self._forward_client( + messages, self._guardrail(**guardrail_kwargs), client_ws_type=ConnectedClientWS + ) + + assert client_ws.closed + assert len(self._guardrail_errors(client_ws)) == blocked_messages + assert self.CLEAN_TEXT not in self._text_inputs(stream) + assert sum(self.VIOLATION in text for text in self._text_inputs(stream)) == blocked_messages - 1 + assert stream.input_stream.closed + + @pytest.mark.asyncio + @pytest.mark.parametrize("guardrail_kwargs", [{}, {"on_violation": "end_session"}], ids=["default", "end-session"]) + async def test_blocked_transcript_ends_the_session_before_its_answer_reaches_the_client( + self, guardrail_kwargs: SessionEndOptions + ) -> None: + payloads: Final = ( + *self._user_transcript(self.BLOCKED_TEXT), + *self._assistant_text("Your secret is safe with me"), + *self._user_transcript(self.CLEAN_TEXT), + *self._assistant_text("It is sunny"), + ) + guardrail: Final = self._guardrail( + event_hook=GuardrailEventHooks.realtime_input_transcription, **guardrail_kwargs + ) + + client_ws, _ = await self._forward_bedrock(payloads, guardrail) + + client_events: Final = self._client_events(client_ws) + assert client_ws.closed + assert not any(event["type"].startswith("response.") for event in client_events) + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + assert guardrail.checked == [((self.BLOCKED_TEXT,), None)] + + @pytest.mark.asyncio + async def test_blocked_transcript_answer_never_reaches_the_client_when_the_previous_answer_closes_late( + self, + ) -> None: + earlier_answer: Final = self._assistant_text("Earlier answer") + payloads: Final = ( + *earlier_answer[:2], + *self._user_transcript(self.BLOCKED_TEXT), + earlier_answer[2], + *self._assistant_text("Your secret is safe with me"), + ) + + client_ws, _ = await self._forward_bedrock( + payloads, self._guardrail(event_hook=GuardrailEventHooks.realtime_input_transcription) + ) + + client_events: Final = self._client_events(client_ws) + assert "Your secret is safe with me" not in json.dumps(client_events) + assert [event["delta"] for event in client_events if event["type"] == "response.text.delta"] == [ + "Earlier answer" + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("event_hook", "default_on", "client_audio_streamed"), + [ + (GuardrailEventHooks.realtime_input_transcription, True, False), + (GuardrailEventHooks.pre_call, True, True), + (GuardrailEventHooks.realtime_input_transcription, False, True), + ], + ids=["typed-text-session", "pre-call-hook", "not-default-on"], + ) + async def test_transcript_guardrail_that_does_not_apply_changes_nothing( + self, event_hook: GuardrailEventHooks, default_on: bool, client_audio_streamed: bool + ) -> None: + guardrail: Final = self._guardrail(event_hook=event_hook, default_on=default_on) + + client_ws, _ = await self._forward_bedrock( + (*self._user_transcript(self.BLOCKED_TEXT), *self._assistant_text("Your secret is safe with me")), + guardrail, + client_audio_streamed=client_audio_streamed, + ) + + client_events: Final = self._client_events(client_ws) + assert [event["delta"] for event in client_events if event["type"] == "response.text.delta"] == [ + "Your secret is safe with me" + ] + assert self._guardrail_errors(client_ws) == [] + assert guardrail.checked == [] + + @pytest.mark.asyncio + async def test_session_ends_when_the_client_leaves_before_hearing_about_a_blocked_message( + self, stub_aws_models: None + ) -> None: + stream, client_ws = await self._forward_client( + [self._user_text(self.BLOCKED_TEXT), self._user_text(self.CLEAN_TEXT)], + self._guardrail(), + client_ws_type=GuardrailErrorUndeliveredClientWS, + ) + + assert self._text_inputs(stream) == [] + assert self._client_events(client_ws) == [] + + @pytest.mark.asyncio + async def test_answers_stop_when_the_client_leaves_before_hearing_about_a_blocked_transcript(self) -> None: + guardrail: Final = self._guardrail(event_hook=GuardrailEventHooks.realtime_input_transcription) + + _, logged = await self._forward_bedrock( + (*self._user_transcript(self.BLOCKED_TEXT), *self._assistant_text("Your secret is safe with me")), + guardrail, + client_ws_type=GuardrailErrorUndeliveredClientWS, + ) + + assert not any(event["type"].startswith("response.") for event in logged) + + @pytest.mark.asyncio + async def test_proxy_request_guardrails_and_key_reach_the_bedrock_session( + self, stub_aws_sdk_client: dict[str, object], monkeypatch: pytest.MonkeyPatch + ) -> None: + class ProxyLogging(FakeLogging): + def update_from_kwargs(self, **kwargs: object) -> None: + pass + + guardrail: Final = self._guardrail(default_on=False) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + stream: Final = ScriptedBedrockStream([], receiver_type=DrainedThenOpenBedrockReceiver) + stub_aws_sdk_client["streams"] = [stream] + user_api_key_dict: Final = object() + client_ws: Final = DisconnectingClientWS([self._user_text(self.BLOCKED_TEXT)]) + + await asyncio.wait_for( + realtime_main._arealtime.__wrapped__( + model=f"bedrock/{self.MODEL}", + websocket=client_ws, + litellm_logging_obj=ProxyLogging(), + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + user_api_key_dict=user_api_key_dict, + metadata={"guardrails": ["no_secrets"]}, + ), + timeout=1, + ) + + assert guardrail.checked == [((self.BLOCKED_TEXT,), user_api_key_dict)] + assert "hunter2" not in json.dumps(self._bedrock_events(stream)) + assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION] + + @pytest.mark.asyncio + async def test_a_key_that_opted_out_of_global_guardrails_is_not_screened( + self, stub_aws_sdk_client: dict[str, object], monkeypatch: pytest.MonkeyPatch + ) -> None: + class ProxyLogging(FakeLogging): + def update_from_kwargs(self, **kwargs: object) -> None: + pass + + guardrail: Final = self._guardrail() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + stream: Final = ScriptedBedrockStream([], receiver_type=DrainedThenOpenBedrockReceiver) + stub_aws_sdk_client["streams"] = [stream] + client_ws: Final = DisconnectingClientWS([self._user_text(self.BLOCKED_TEXT)]) + + await asyncio.wait_for( + realtime_main._arealtime.__wrapped__( + model=f"bedrock/{self.MODEL}", + websocket=client_ws, + litellm_logging_obj=ProxyLogging(), + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + metadata={"user_api_key_metadata": {"disable_global_guardrails": True}}, + ), + timeout=1, + ) + + assert guardrail.checked == [] + assert self._guardrail_errors(client_ws) == [] + assert self.BLOCKED_TEXT in json.dumps(self._bedrock_events(stream)) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) 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_codex_catalog_source.py b/tests/unit/proxy/common_utils/test_codex_catalog_source.py new file mode 100644 index 00000000000..e3009e19dd7 --- /dev/null +++ b/tests/unit/proxy/common_utils/test_codex_catalog_source.py @@ -0,0 +1,257 @@ +import asyncio +import json +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from typing import Final + +import httpx +import pytest + +from litellm.proxy.common_utils.codex_catalog_source import ( + CODEX_RELEASE_FILE_URL, + CodexCatalogSource, + bundled_codex_stock, +) +from litellm.proxy.common_utils.codex_model_catalog import bundled_codex_models + +_TTL: Final = 3600.0 + + +def _release_files(slug: str, prompt: str) -> Mapping[str, bytes]: + entry: Final = {**bundled_codex_models()["gpt-5.5"].model_dump(mode="json", exclude_unset=True), "slug": slug} + return {"models.json": json.dumps({"models": [entry]}).encode(), "prompt.md": prompt.encode()} + + +@dataclass +class _GitHub: + """Serves each tagged release's files, 404 for any other tag, and records every URL asked for.""" + + releases: Mapping[str, Mapping[str, bytes]] + fail: Callable[[httpx.Request], httpx.Response] | None = None + requested: list[str] = field(default_factory=list) + + async def handle(self, request: httpx.Request) -> httpx.Response: + self.requested.append(str(request.url)) + if self.fail is not None: + return self.fail(request) + files: Final = { + CODEX_RELEASE_FILE_URL.format(version=version, file=name): body + for version, release in self.releases.items() + for name, body in release.items() # comprehension-ok: flattens release -> file in one lookup table + } + body: Final = files.get(str(request.url)) + return httpx.Response(404, text="404: Not Found") if body is None else httpx.Response(200, content=body) + + +@dataclass +class _Clock: + now: float = 0.0 + + def __call__(self) -> float: + return self.now + + +def _source( + github: _GitHub, + clock: _Clock | None = None, + *, + failure_ttl: float = _TTL, + max_versions: int = 8, + fetches_per_minute: int = 100, +) -> CodexCatalogSource: + client: Final = httpx.AsyncClient(transport=httpx.MockTransport(github.handle)) + return CodexCatalogSource( + lambda: client, + offline=lambda: False, + clock=clock or _Clock(), + failure_ttl=failure_ttl, + max_versions=max_versions, + fetches_per_minute=fetches_per_minute, + ) + + +async def test_release_catalog_and_prompt_are_served_for_the_asking_codex_version() -> None: + github: Final = _GitHub({"0.170.0": _release_files("gpt-from-0.170.0", "prompt of 0.170.0")}) + + stock: Final = await _source(github).stock("0.170.0") + + assert (list(stock.models), stock.instructions) == (["gpt-from-0.170.0"], "prompt of 0.170.0") + assert stock.models["gpt-from-0.170.0"].display_name == bundled_codex_models()["gpt-5.5"].display_name + + +async def test_each_codex_version_gets_its_own_release_and_a_missing_tag_does_not_spill_over() -> None: + github: Final = _GitHub( + { + "0.166.0": _release_files("gpt-from-0.166.0", "prompt of 0.166.0"), + "0.171.0-alpha.2": _release_files("gpt-from-alpha", "prompt of alpha"), + } + ) + source: Final = _source(github) + + served: Final = {version: await source.stock(version) for version in ("0.166.0", "0.171.0-alpha.2", "0.0.0")} + + assert {version: (list(stock.models), stock.instructions) for version, stock in served.items()} == { + "0.166.0": (["gpt-from-0.166.0"], "prompt of 0.166.0"), + "0.171.0-alpha.2": (["gpt-from-alpha"], "prompt of alpha"), + "0.0.0": (list(bundled_codex_models()), bundled_codex_stock().instructions), + } + + +def _unreachable(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("egress blocked", request=request) + + +def _server_error(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, request=request) + + +@pytest.mark.parametrize("fail", [_unreachable, _server_error], ids=["unreachable", "server-error"]) +async def test_a_failed_fetch_serves_the_bundled_catalog(fail: Callable[[httpx.Request], httpx.Response]) -> None: + github: Final = _GitHub({"0.170.0": _release_files("gpt-from-0.170.0", "p")}, fail=fail) + + assert await _source(github).stock("0.170.0") is bundled_codex_stock() + + +async def test_a_release_that_never_answers_serves_the_bundled_catalog_within_the_timeout() -> None: + async def hang(_: httpx.Request) -> httpx.Response: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + client: Final = httpx.AsyncClient(transport=httpx.MockTransport(hang)) + source: Final = CodexCatalogSource(lambda: client, offline=lambda: False, clock=_Clock(), timeout=0.01) + + assert await asyncio.wait_for(source.stock("0.170.0"), timeout=30) is bundled_codex_stock() + + +@pytest.mark.parametrize( + "files", + [ + {"models.json": b'{"models": [{"slug": "gpt-x"}]}', "prompt.md": b"p"}, + {"models.json": b"rate limited", "prompt.md": b"p"}, + {"models.json": _release_files("gpt-x", "p")["models.json"], "prompt.md": b"\xff\xfe"}, + {"models.json": _release_files("gpt-x", "p")["models.json"]}, + ], + ids=["entry-missing-fields", "not-json", "prompt-not-utf8", "prompt-missing"], +) +async def test_a_release_whose_files_do_not_validate_serves_the_bundled_catalog(files: Mapping[str, bytes]) -> None: + assert await _source(_GitHub({"0.170.0": files})).stock("0.170.0") is bundled_codex_stock() + + +async def test_a_failed_fetch_is_not_retried_until_the_failure_expires() -> None: + github: Final = _GitHub({}, fail=_unreachable) + clock: Final = _Clock() + source: Final = _source(github, clock, failure_ttl=_TTL) + + first: Final = await source.stock("0.170.0") + after_first: Final = len(github.requested) + clock.now = _TTL - 1 + within_ttl: Final = await source.stock("0.170.0") + after_within: Final = len(github.requested) + clock.now = _TTL + 1 + await source.stock("0.170.0") + + assert first is within_ttl is bundled_codex_stock() + assert (after_first > 0, after_within == after_first, len(github.requested) > after_within) == (True, True, True) + + +async def test_a_fetched_release_is_not_fetched_again() -> None: + github: Final = _GitHub({"0.170.0": _release_files("gpt-from-0.170.0", "p")}) + clock: Final = _Clock() + source: Final = _source(github, clock) + + first: Final = await source.stock("0.170.0") + fetches: Final = len(github.requested) + clock.now = _TTL * 100 + + assert await source.stock("0.170.0") is first + assert len(github.requested) == fetches + + +async def test_only_the_most_recent_versions_are_kept() -> None: + github: Final = _GitHub({version: _release_files(version, "p") for version in ("0.1.0", "0.2.0", "0.3.0")}) + source: Final = _source(github, max_versions=2) + for version in ("0.1.0", "0.2.0", "0.3.0"): + await source.stock(version) + fetches: Final = len(github.requested) + + await source.stock("0.3.0") + kept_was_cached: Final = len(github.requested) == fetches + await source.stock("0.1.0") + + assert (kept_was_cached, len(github.requested) > fetches) == (True, True) + + +async def test_versions_that_fail_never_push_out_a_fetched_one() -> None: + github: Final = _GitHub({"0.170.0": _release_files("gpt-from-0.170.0", "p")}) + source: Final = _source(github, max_versions=2) + fetched: Final = await source.stock("0.170.0") + for made_up in ("9.9.1", "9.9.2", "9.9.3"): + await source.stock(made_up) + fetches: Final = len(github.requested) + + assert await source.stock("0.170.0") is fetched + assert len(github.requested) == fetches + + +async def test_a_failed_version_keeps_its_cooldown_past_the_fetched_catalog_cap() -> None: + github: Final = _GitHub({}, fail=_unreachable) + source: Final = _source(github, max_versions=2) + for made_up in ("9.9.1", "9.9.2", "9.9.3", "9.9.4"): + await source.stock(made_up) + fetches: Final = len(github.requested) + + await source.stock("9.9.1") + + assert len(github.requested) == fetches, github.requested + + +async def test_a_worker_starts_at_most_its_fetch_budget_each_minute() -> None: + github: Final = _GitHub({}) + clock: Final = _Clock() + source: Final = _source(github, clock, fetches_per_minute=2) + for made_up in ("9.9.1", "9.9.2", "9.9.3"): + await source.stock(made_up) + within_minute: Final = len(github.requested) + clock.now = 61 + await source.stock("9.9.3") + + assert (within_minute, len(github.requested)) == (4, 6), github.requested + + +async def test_requests_for_a_version_being_fetched_share_that_fetch() -> None: + answer: Final = asyncio.Event() + files: Final = _release_files("gpt-from-0.170.0", "p") + requested: Final[list[str]] = [] + + async def slow_github(request: httpx.Request) -> httpx.Response: + requested.append(str(request.url)) + await answer.wait() + return httpx.Response(200, content=files[str(request.url).rsplit("/", 1)[1]]) + + client: Final = httpx.AsyncClient(transport=httpx.MockTransport(slow_github)) + source: Final = CodexCatalogSource(lambda: client, offline=lambda: False, clock=_Clock()) + waiting: Final = [asyncio.ensure_future(source.stock("0.170.0")) for _ in range(5)] + await asyncio.sleep(0) + answer.set() + served: Final = await asyncio.gather(*waiting) + + assert [list(stock.models) for stock in served] == [["gpt-from-0.170.0"]] * 5 + assert len(requested) == 2, requested + + +async def test_offline_mode_never_fetches(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_OFFLINE_MODE", "true") + github: Final = _GitHub({"0.170.0": _release_files("gpt-from-0.170.0", "p")}) + client: Final = httpx.AsyncClient(transport=httpx.MockTransport(github.handle)) + + stock: Final = await CodexCatalogSource(lambda: client, clock=_Clock()).stock("0.170.0") + + assert (stock, github.requested) == (bundled_codex_stock(), []) + + +@pytest.mark.parametrize("client_version", ["../../main", "0.170.0/../../x", "latest", "", "\u0660.\u0661.\u0662"]) +async def test_a_version_that_is_no_release_tag_is_never_fetched(client_version: str) -> None: + github: Final = _GitHub({}) + + assert await _source(github).stock(client_version) is bundled_codex_stock() + assert github.requested == [] diff --git a/tests/unit/proxy/common_utils/test_codex_model_catalog.py b/tests/unit/proxy/common_utils/test_codex_model_catalog.py index f0b3a8fabc7..d4cc2fcb006 100644 --- a/tests/unit/proxy/common_utils/test_codex_model_catalog.py +++ b/tests/unit/proxy/common_utils/test_codex_model_catalog.py @@ -1407,6 +1407,34 @@ def test_model_info_that_does_not_validate_declares_no_capability_but_keeps_its_ assert capabilities == CodexDeploymentCapabilities(priced_service_tier_keys=frozenset({priced})) +def test_model_list_body_takes_the_asking_releases_catalog_and_prompt() -> None: + router: Final = litellm.Router( + model_list=[ + {"model_name": "gpt-release-only", "litellm_params": {"model": "openai/gpt-release-only"}}, + {"model_name": "in-house", "litellm_params": {"model": "openai/in-house-coder"}}, + ] + ) + release_entry: Final = bundled_codex_models()["gpt-5.5"].model_copy( + update={"slug": "gpt-release-only", "display_name": "Release Only"} + ) + names: Final = ("gpt-release-only", "in-house") + + body: Final = codex_model_list_body( + _listing(*names), + tuple((name, name) for name in names), + router, + client_version="0.170.0", + stock={"gpt-release-only": release_entry}, + instructions="release prompt", + ) + + served: Final = {entry["slug"]: entry for entry in json.loads(body)["models"]} + assert (served["gpt-release-only"]["display_name"], served["in-house"]["base_instructions"]) == ( + "Release Only", + "release prompt", + ), served + + class _CollidingName(str): def __hash__(self) -> int: return 1 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/common_utils/test_openai_endpoint_utils.py b/tests/unit/proxy/common_utils/test_openai_endpoint_utils.py index 3e3e89f8117..3eb901144a3 100644 --- a/tests/unit/proxy/common_utils/test_openai_endpoint_utils.py +++ b/tests/unit/proxy/common_utils/test_openai_endpoint_utils.py @@ -77,9 +77,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( ), ], ) -def test_remove_sensitive_info_from_deployment( - model_config: dict, expected_config: dict -): +def test_remove_sensitive_info_from_deployment(model_config: dict, expected_config: dict): sanitized_config = remove_sensitive_info_from_deployment(model_config) assert sanitized_config == expected_config @@ -110,10 +108,7 @@ def test_remove_sensitive_info_from_deployment_with_excluded_keys(): sanitized_config = remove_sensitive_info_from_deployment( copy.deepcopy(base_config), excluded_keys={"litellm_credentials_name"} ) - assert ( - sanitized_config["litellm_params"]["litellm_credentials_name"] - == "my-credential-name" - ) + assert sanitized_config["litellm_params"]["litellm_credentials_name"] == "my-credential-name" # access_token should still be masked (not in excluded_keys) assert sanitized_config["litellm_params"]["access_token"] != "token-12345" @@ -121,3 +116,24 @@ def test_remove_sensitive_info_from_deployment_with_excluded_keys(): # api_key should still be removed (popped) regardless of excluded_keys assert "api_key" not in sanitized_config["litellm_params"] + + +def test_remove_sensitive_info_from_deployment_keeps_nested_nulls(): + router_config = { + "models": {"capable": "strong", "efficient": "fast", "judge": None}, + "options": {"confidence_threshold": None, "tool_semantics": None}, + "fallbacks": [None, "fast"], + } + deployment = { + "model_name": "route", + "litellm_params": { + "model": "auto_router/switchyard", + "api_key": "sk-sensitive-key-123", + "switchyard_router_config": copy.deepcopy(router_config), + }, + } + + sanitized = remove_sensitive_info_from_deployment(deployment) + + assert "api_key" not in sanitized["litellm_params"] + assert sanitized["litellm_params"]["switchyard_router_config"] == router_config 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_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 32c9535b4c8..2fdb77af8f9 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -1085,7 +1085,7 @@ class TestClearCache: Test that clear_cache successfully clears router model caches and reloads models. """ mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.registry.return_value.all.return_value = ["openai/gpt-4o", "openai/gpt-4o-mini"] @@ -1125,7 +1125,7 @@ class TestClearCache: # entries are auto_router/* deployments (so their router-map entries should be # cleared for reload); the config-defined router is preserved. mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.registry.return_value.all.return_value = [ { @@ -1209,7 +1209,7 @@ class TestClearCache: ) mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.registry.return_value.all.return_value = [ { @@ -1270,7 +1270,7 @@ class TestClearCachePreservesConfigRouters: ) mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.registry.return_value.all.return_value = [ { @@ -1315,7 +1315,7 @@ class TestClearCachePreservesConfigRouters: from litellm.proxy.management_endpoints.model_management_endpoints import clear_cache mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.registry.return_value.all.return_value = [ { @@ -1346,15 +1346,15 @@ class TestClearCachePreservesConfigRouters: assert "shared-name" in plugins["complexity"] @pytest.mark.asyncio - async def test_db_quality_and_adaptive_routers_are_evicted(self): - """The auto_router/ prefix also covers quality_router/ and adaptive_router/. Their + async def test_db_quality_adaptive_and_switchyard_routers_are_evicted(self): + """The auto_router/ prefix also covers quality_router/, adaptive_router/ and switchyard. Their registry entries must be popped too, or reload's init raises 'already exists' (quality) or leaves a stale entry (adaptive). """ from litellm.proxy.management_endpoints.model_management_endpoints import clear_cache mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.registry.return_value.all.return_value = [ { @@ -1367,11 +1367,17 @@ class TestClearCachePreservesConfigRouters: "model_info": {"id": "db-a", "db_model": True}, "litellm_params": {"model": "auto_router/adaptive_router/a1"}, }, + { + "model_name": "s1", + "model_info": {"id": "db-s", "db_model": True}, + "litellm_params": {"model": "auto_router/switchyard"}, + }, ] plugins["auto"] = {} plugins["complexity"] = {} plugins["quality"] = {"q1": MagicMock()} plugins["adaptive"] = {"a1": MagicMock()} + plugins["switchyard"] = {"s1": MagicMock()} mock_config = MagicMock() mock_config.add_deployment_locked = AsyncMock( @@ -1389,6 +1395,7 @@ class TestClearCachePreservesConfigRouters: assert "q1" not in plugins["quality"] assert "a1" not in plugins["adaptive"] + assert "s1" not in plugins["switchyard"] class TestDeleteModelClearsRouterRegistry: @@ -1503,7 +1510,7 @@ class TestDeleteModelClearsRouterRegistry: mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) mock_router = MagicMock() - plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}} + plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}} mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__ mock_router.apply.return_value = ApplyResult( version=1, @@ -11431,6 +11438,98 @@ async def test_add_new_model_reports_a_database_outage_during_the_insert_as_503( assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection") +@pytest.mark.parametrize("invalid", [ + {"algorithm": "random"}, + {"classifier": {}}, + {"options": {"capable_system_prompt": "ignored"}}, +]) +def test_switchyard_write_rejects_unloadable_configuration(invalid: dict[str, object]) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import _strategy_router_write_violation + from litellm.types.router import updateLiteLLMParams + + violation: Final = _strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + model="auto_router/switchyard", + switchyard_router_config={"models": {"capable": "strong", "efficient": "fast"}, **invalid}, + ), + existing_params=None, + ) + assert violation is not None + assert "Invalid switchyard_router_config" in violation + + +SWITCHYARD_PARAMS: Final = { + "model": "auto_router/switchyard", + "switchyard_router_config": {"models": {"capable": "strong", "efficient": "fast"}}, +} + + +@pytest.mark.parametrize( + ("model_name", "incoming", "existing"), + [ + ("strong", SWITCHYARD_PARAMS, None), + ("fast", {"model": "auto_router/switchyard"}, SWITCHYARD_PARAMS), + ("fast", None, SWITCHYARD_PARAMS), + ( + "smart", + { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "smart", "COMPLEX": "strong"}}, + }, + None, + ), + ], +) +def test_strategy_router_that_lists_itself_is_rejected_before_saving( + model_name: str, incoming: dict[str, object] | None, existing: dict[str, object] | None +) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _raise_on_strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + with pytest.raises(ProxyException) as exc_info: + _raise_on_strategy_router_write_violation( + incoming_params=updateLiteLLMParams(**incoming) if incoming is not None else None, + existing_params=updateLiteLLMParams(**existing) if existing is not None else None, + model_name=model_name, + ) + assert exc_info.value.code == "400" + assert model_name in exc_info.value.message + + +def test_strategy_router_with_other_targets_saves() -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import _self_routing_violation + from litellm.types.router import updateLiteLLMParams + + assert _self_routing_violation("route", updateLiteLLMParams(**SWITCHYARD_PARAMS), None) is None + + +def test_self_routing_is_caught_when_the_stored_model_is_encrypted(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _raise_on_strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-for-encryption") + stored: Final = updateLiteLLMParams( + **{**SWITCHYARD_PARAMS, "model": encrypt_value_helper("auto_router/switchyard")} + ) + assert stored.model != "auto_router/switchyard" + with pytest.raises(ProxyException) as exc_info: + _raise_on_strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + switchyard_router_config={"models": {"capable": "route", "efficient": "fast"}} + ), + existing_params=stored, + model_name="route", + ) + assert exc_info.value.code == "400" + + @pytest.mark.asyncio async def test_rename_rewrites_only_in_a_transaction_holding_every_locked_teams_advisory_lock(): """A team that starts listing the old name between the advisory snapshot and the row lock must not be 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/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 3b000ba741f..f74ff9b1fca 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -427,8 +427,7 @@ def test_enrich_http_exception_with_guardrail_context_dict_detail(): assert exc.detail["guardrail_mode"] == "post_call" -def test_enrich_http_exception_string_detail_noop(): - """L2: string-detail HTTPException is not mutated (can't add fields to a str).""" +def test_enrich_http_exception_string_detail_preserves_body_and_identifies_guardrail(): from litellm.proxy.guardrails.exception_utils import enrich_http_exception_with_guardrail_context class StubCallback: @@ -438,6 +437,7 @@ def test_enrich_http_exception_string_detail_noop(): exc = HTTPException(status_code=400, detail="Content blocked") enrich_http_exception_with_guardrail_context(exc, StubCallback()) assert exc.detail == "Content blocked" + assert getattr(exc, "guardrail_name") == "x" def test_enrich_http_exception_setdefault_does_not_overwrite(): 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/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py index a16661c0174..44ba9007d92 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -18,7 +18,9 @@ from openai.types.responses.tool_param import Mcp import litellm from litellm.caching.caching import DualCache +from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.mcp_server import operations as mcp_operations from litellm.proxy.common_utils.user_api_key_cache import object_permission_cache_key @@ -505,6 +507,139 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey assert post_call_failure_hook.await_args.kwargs.get("route") == "/responses/mcp/call_tool" +@pytest.mark.asyncio +@pytest.mark.parametrize("behavior", [None, "tool_result", "raise"]) +@pytest.mark.parametrize("error_kind", ["pii", "guardrail", "http", "string_http", "unattributed_http"]) +@pytest.mark.parametrize("mode", ["pre_mcp_call", "during_mcp_call"]) +async def test_mcp_block_behavior_preserves_failure_logging_and_propagates_only_opted_in_blocks( + monkeypatch: pytest.MonkeyPatch, + behavior: Literal["tool_result", "raise"] | None, + error_kind: str, + mode: str, +) -> None: + call_tool: Final = _setup_mcp_call_environment(monkeypatch) + failure_hook: Final = _setup_proxy_logging(monkeypatch) + callback: Final = CustomGuardrail(guardrail_name="blocked", event_hook=mode, default_on=True) + callback.mcp_block_behavior = behavior + callback.use_native_lifecycle_hooks = True + unrelated: Final = CustomGuardrail(guardrail_name="blocked", event_hook=mode, default_on=True) + unrelated.mcp_block_behavior = "tool_result" if behavior == "raise" else "raise" + monkeypatch.setattr(litellm, "callbacks", [unrelated, callback]) + error: Final = ( + BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="blocked") + if error_kind == "pii" + else GuardrailRaisedException(guardrail_name="blocked", message="denied") + if error_kind == "guardrail" + else HTTPException(status_code=400, detail="denied") + if error_kind == "string_http" + else HTTPException( + status_code=400, + detail={"error": "denied", **({"guardrail_name": "blocked"} if error_kind == "http" else {})}, + ) + ) + logger: Final = ProxyLogging(user_api_key_cache=DualCache()) + settings: Final = types.SimpleNamespace( + guardrail_list=[{"guardrail_name": "blocked", "callback": item} for item in (unrelated, callback)], + is_set=lambda name: name == "guardrail_list", + guardrail=lambda name: {"callback": callback}, + ) + monkeypatch.setattr( + sys.modules["litellm.proxy.proxy_server"], + "llm_router", + types.SimpleNamespace(settings=lambda: settings), + raising=False, + ) + auth: Final = UserAPIKeyAuth(api_key="test-key") + + async def block(**kwargs: object) -> None: + raise error + + callback.async_pre_call_hook = block + callback.async_moderation_hook = block + + async def call_with_guardrail(**kwargs: object) -> None: + if error_kind == "unattributed_http": + raise error + data: Final = {"mcp_arguments": {"query": "denied"}} + if mode == "pre_mcp_call": + await logger._process_guardrail_callback( + unrelated, data, auth, "call_mcp_tool", GuardrailEventHooks.pre_mcp_call + ) + else: + await logger._run_during_call_guardrail(callback, data, auth, auth, "call_mcp_tool") + + call_tool.side_effect = call_with_guardrail + tool_calls: Final = [{"id": "blocked-call", "function": {"name": "search", "arguments": "{}"}}] + + if behavior == "raise" and error_kind != "unattributed_http": + with pytest.raises(type(error)) as caught: + await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={"search": "local"}, tool_calls=tool_calls, user_api_key_auth=auth + ) + assert caught.value is error + else: + results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={"search": "local"}, tool_calls=tool_calls, user_api_key_auth=auth + ) + assert len(results) == 1 + assert results[0]["tool_call_id"] == "blocked-call" + expected_prefix: Final = "Tool call failed" if "http" in error_kind else "Tool call blocked" + assert results[0]["result"].startswith(expected_prefix) + failure_hook.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("behavior", [None, "tool_result", "raise"]) +async def test_initialized_content_filter_blocks_mcp_according_to_its_config( + monkeypatch: pytest.MonkeyPatch, behavior: Literal["tool_result", "raise"] | None +) -> None: + for name in ( + "callbacks", "input_callback", "success_callback", "failure_callback", + "_async_success_callback", "_async_failure_callback", + ): + monkeypatch.setattr(litellm, name, []) + handler: Final = InMemoryGuardrailHandler() + handler.initialize_guardrail( + guardrail={ + "guardrail_name": "configured-block", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_mcp_call", + "default_on": True, + "blocked_words": [{"keyword": "forbiddenmarker", "action": "BLOCK"}], + **({"mcp_block_behavior": behavior} if behavior is not None else {}), + }, + } + ) + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="local", name="local", transport=MCPTransport.http, url="http://upstream") + manager.registry = {server.server_id: server} + manager.tool_name_to_mcp_server_name_mapping = {"local-search": "local"} + client: Final = AsyncMock() + client.call_tool.return_value = CallToolResult(content=[TextContent(type="text", text="ok")]) + manager.create_mcp_client = AsyncMock(return_value=client) + logger: Final = ProxyLogging(user_api_key_cache=DualCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", logger) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", manager) + tool_calls: Final = [{"id": "blocked", "function": {"name": "search", "arguments": '{"query":"forbiddenmarker"}'}}] + kwargs: Final = { + "tool_server_map": {"search": "local"}, + "tool_calls": tool_calls, + "user_api_key_auth": UserAPIKeyAuth(api_key="test-key"), + "served_tools": [MCPTool(name="search", inputSchema={"type": "object"})], + } + if behavior == "raise": + with pytest.raises(HTTPException, match="forbiddenmarker"): + await LiteLLM_Proxy_MCP_Handler.execute_tool_calls(**kwargs) + else: + results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls(**kwargs) + assert len(results) == 1 + assert results[0]["tool_call_id"] == "blocked" + assert "Content blocked: keyword 'forbiddenmarker' detected" in results[0]["result"] + client.call_tool.assert_not_awaited() + + @pytest.mark.asyncio async def test_execute_tool_calls_logs_each_tool_call_under_its_own_id_on_the_request_trace( monkeypatch, diff --git a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py index f745b5bfadf..f113d00a471 100644 --- a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py @@ -6,9 +6,12 @@ from typing import Final, cast from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException from mcp.types import CallToolResult, TextContent import litellm # noqa: F401 - ensures litellm.responses.main is registered in sys.modules +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails.exception_utils import enrich_http_exception_with_guardrail_context from litellm.responses.mcp.mcp_streaming_iterator import ( MCPEnhancedStreamingIterator, ) @@ -293,6 +296,27 @@ async def test_continuation_id_is_final_round_not_interim_tool_call(monkeypatch) assert completed[-1].response.id != "resp-interim" +@pytest.mark.asyncio +async def test_opted_in_guardrail_block_emits_terminal_error_without_follow_up(monkeypatch): + call_tool: Final = _mock_mcp_environment(monkeypatch) + callback: Final = CustomGuardrail(guardrail_name="blocked") + callback.mcp_block_behavior = "raise" + error: Final = HTTPException(status_code=400, detail={"error": "guardrail denied", "guardrail_name": "blocked"}) + enrich_http_exception_with_guardrail_context(error, callback) + call_tool.side_effect = error + follow_up: Final = AsyncMock(return_value=_text_only_stream("unguarded follow-up")) + monkeypatch.setattr(responses_main_module, "aresponses", follow_up) + iterator: Final = _make_iterator([_completed_chunk([_function_call("call_1", "read_wiki_contents")])]) + + chunks: Final = [chunk async for chunk in iterator] + + assert chunks[-1].type == ResponsesAPIStreamEvents.ERROR + assert chunks[-1].error.code == "400" + assert "guardrail denied" in chunks[-1].error.message + assert all(chunk.type != ResponsesAPIStreamEvents.RESPONSE_COMPLETED for chunk in chunks) + follow_up.assert_not_awaited() + + @pytest.mark.asyncio async def test_follow_up_call_failure_emits_terminal_error_event(monkeypatch): """ diff --git a/tests/unit/responses/test_responses_utils.py b/tests/unit/responses/test_responses_utils.py index 037a6431fef..85e2a775bd9 100644 --- a/tests/unit/responses/test_responses_utils.py +++ b/tests/unit/responses/test_responses_utils.py @@ -429,8 +429,8 @@ class TestResponseAPILoggingUtils: assert usage.prompt_tokens_details.cache_creation_tokens == 50 assert usage.prompt_tokens_details.cache_write_tokens == 50 - @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, @@ -440,7 +440,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/router_strategy/switchyard_router/__init__.py b/tests/unit/router_strategy/switchyard_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_strategy/switchyard_router/test_config.py b/tests/unit/router_strategy/switchyard_router/test_config.py new file mode 100644 index 00000000000..c24023ddb79 --- /dev/null +++ b/tests/unit/router_strategy/switchyard_router/test_config.py @@ -0,0 +1,100 @@ +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.router_strategy.switchyard_router.config import SwitchyardRouterConfig, validate_switchyard_config_write + + +@pytest.mark.parametrize("classifier", [None, {}]) +def test_default_confidence_depends_on_classifier(classifier: dict[str, object] | None) -> None: + config: Final = SwitchyardRouterConfig.model_validate( + { + "models": { + "capable": "strong", + "efficient": "fast", + **({"judge": "judge"} if classifier is not None else {}), + }, + "classifier": classifier, + } + ) + assert config.confidence_threshold == (0.7 if classifier is not None else 0.5) + + +@pytest.mark.parametrize("classifier", [None, {}]) +def test_explicit_confidence_is_preserved(classifier: dict[str, object] | None) -> None: + config: Final = SwitchyardRouterConfig.model_validate( + { + "models": { + "capable": "strong", + "efficient": "fast", + **({"judge": "judge"} if classifier is not None else {}), + }, + "classifier": classifier, + "options": {"confidence_threshold": 0.63, "picker": "capable_first"}, + } + ) + assert config.confidence_threshold == 0.63 + + +@pytest.mark.parametrize( + "option", ["escalation_note", "deescalation_note", "capable_system_prompt", "efficient_system_prompt", "unknown"] +) +def test_request_rewrite_and_unknown_options_are_rejected(option: str) -> None: + with pytest.raises(ValidationError, match=option): + SwitchyardRouterConfig.model_validate( + {"models": {"capable": "strong", "efficient": "fast"}, "options": {option: "ignored"}} + ) + + +@pytest.mark.parametrize( + "bad", + [ + {"algorithm": "random"}, + {"classifier": {}}, + {"options": {"confidence_threshold": 2}}, + {"models": {"capable": "", "efficient": "fast"}}, + {"models": {"capable": "same", "efficient": "same"}}, + {"models": {"capable": "strong", "efficient": "fast", "judge": "j"}, "classifier": {"session_affinity": True}}, + {"models": {"capable": "strong", "efficient": "fast", "judge": "judge"}}, + {"models": {"capable": "strong", "efficient": "fast", "judge": "j"}, "classifier": {"kind": "other"}}, + { + "models": {"capable": "strong", "efficient": "fast", "judge": "j"}, + "classifier": {"kind": "decision", "cutoff": 1.5}, + }, + { + "models": {"capable": "strong", "efficient": "fast", "judge": "j"}, + "classifier": {"kind": "decision", "base_threshold": 0.5}, + }, + { + "models": {"capable": "strong", "efficient": "fast", "judge": "j"}, + "classifier": {"kind": "llm", "cutoff": 0.5}, + }, + ], +) +def test_invalid_config_cannot_be_saved(bad: dict[str, object]) -> None: + violation: Final = validate_switchyard_config_write({"models": {"capable": "strong", "efficient": "fast"}, **bad}) + assert violation is not None + assert "Invalid switchyard_router_config" in violation + + +def test_missing_package_rejects_save_with_install_instructions() -> None: + import sys + from unittest.mock import patch + + with patch.dict(sys.modules, {"switchyard_rust.libsy": None}): + violation: Final = validate_switchyard_config_write({"models": {"capable": "strong", "efficient": "fast"}}) + assert violation is not None + assert "litellm[switchyard]" in violation + + +@pytest.mark.parametrize( + "classifier", [{"kind": "decision"}, {"kind": "decision", "cutoff": 0.4, "evidence": {"a": 1}}] +) +def test_decision_classifier_config_can_be_saved(classifier: dict[str, object]) -> None: + assert ( + validate_switchyard_config_write( + {"models": {"capable": "strong", "efficient": "fast", "judge": "jev"}, "classifier": classifier} + ) + is None + ) diff --git a/tests/unit/router_strategy/switchyard_router/test_switchyard_router.py b/tests/unit/router_strategy/switchyard_router/test_switchyard_router.py new file mode 100644 index 00000000000..8a26b4610db --- /dev/null +++ b/tests/unit/router_strategy/switchyard_router/test_switchyard_router.py @@ -0,0 +1,438 @@ +import asyncio +import json +from itertools import chain +from collections.abc import AsyncIterator, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_api import ChangeSet +from litellm.router_strategy.switchyard_router.config import SwitchyardRouterConfig, build_algorithm +from litellm.router_strategy.switchyard_router.switchyard_router import SwitchyardRouter +from litellm.types.decisions import OpenAIDecisionResponse +from litellm.types.router import Deployment +from litellm.types.utils import ModelResponse + +libsy = pytest.importorskip("switchyard_rust.libsy") + +HELLO: Final = [{"role": "user", "content": "Hello"}] + + +async def unexpected_judge(request: Mapping[str, object]) -> ModelResponse: + raise AssertionError("This route must not call an LLM judge") + + +async def unexpected_decision(request: Mapping[str, object]) -> OpenAIDecisionResponse: + raise AssertionError("This route must not call a decision model") + + +def config( + *, classifier: Mapping[str, object] | None = None, picker: str = "efficient_first", threshold: float | None = None +) -> SwitchyardRouterConfig: + return SwitchyardRouterConfig.model_validate( + { + "models": { + "capable": "strong", + "efficient": "fast", + **({"judge": "judge"} if classifier is not None else {}), + }, + "options": {"picker": picker, **({"confidence_threshold": threshold} if threshold is not None else {})}, + **({"classifier": classifier} if classifier is not None else {}), + } + ) + + +def router( + route_config: SwitchyardRouterConfig, + *, + completion=unexpected_judge, + decide=unexpected_decision, + algorithm: object = None, +) -> SwitchyardRouter: + return SwitchyardRouter("route", route_config, algorithm or build_algorithm(route_config), completion, decide) + + +async def route( + route_config: SwitchyardRouterConfig, + messages: object, + *, + call_type: str | None = "acompletion", + request_kwargs: Mapping[str, object] | None = None, + **router_kwargs: object, +) -> tuple[str, Mapping[str, object]]: + result: Final = await router(route_config, **router_kwargs).async_pre_routing_hook( # pyright: ignore[reportArgumentType] # test doubles + "route", request_kwargs or {}, messages, call_type=call_type + ) + assert result is not None + assert result.routing_decision is not None + return result.model, result.routing_decision + + +def anthropic_turn(call_id: str, tool: str, result: str, marked: bool) -> tuple[dict[str, object], ...]: + error: Final = {"is_error": True} if marked else {} + return ( + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Next step", "signature": "sig"}, + {"type": "tool_use", "id": call_id, "name": tool, "input": {"file_path": "a.py"}}, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": call_id, "content": result, **error}]}, + ) + + +def responses_turn(call_id: str, tool: str, result: str, marked: bool) -> tuple[dict[str, object], ...]: + return ( + {"type": "function_call", "call_id": call_id, "name": tool, "arguments": '{"file_path":"a.py"}'}, + {"type": "function_call_output", "call_id": call_id, "output": result}, + ) + + +def chat_turn(call_id: str, tool: str, result: str, marked: bool) -> tuple[dict[str, object], ...]: + call: Final = {"id": call_id, "type": "function", "function": {"name": tool, "arguments": '{"file_path":"a.py"}'}} + return ( + {"role": "assistant", "content": None, "tool_calls": [call]}, + {"role": "tool", "tool_call_id": call_id, "content": result}, + ) + + +TURN_FORMATS: Final = {"anthropic_messages": anthropic_turn, "aresponses": responses_turn, "acompletion": chat_turn} + + +def transcript(endpoint: str, tool: str, results: tuple[str, ...], *, marked: bool = False) -> list[dict[str, object]]: + turn: Final = TURN_FORMATS[endpoint] + turns: Final = chain.from_iterable( + turn(f"call{index}", tool, result, marked) for index, result in enumerate(results) + ) + return [{"role": "user", "content": "Fix the helper"}, *turns] + + +def llm_router( + *, route_config: Mapping[str, object], judge_fails: bool = False, num_retries: int = 0 +) -> litellm.Router: + def deployment(name: str, reply: object) -> Mapping[str, object]: + return { + "model_name": name, + "litellm_params": {"model": f"openai/{name}", "api_key": "k", "mock_response": reply}, + } + + return litellm.Router( + num_retries=num_retries, + model_list=[ + deployment("strong", "strong"), + deployment("fast", "fast"), + *((deployment("judge", Exception("judge down")),) if judge_fails else ()), + switchyard_deployment(route_config), + ], + ) + + +def switchyard_deployment(route_config: Mapping[str, object]) -> dict[str, object]: + return { + "model_name": "route", + "litellm_params": {"model": "auto_router/switchyard", "switchyard_router_config": route_config}, + "model_info": {"id": "route-id"}, + } + + +@pytest.mark.asyncio +async def test_critical_tool_error_selects_capable_without_rewriting_messages() -> None: + messages: Final = transcript("acompletion", "bash", ("CUDA out of memory",)) + result: Final = await router(config()).async_pre_routing_hook("route", {}, messages, call_type="acompletion") + assert result is not None + assert result.model == "strong" + assert result.messages == messages + + +async def routed_reply(proxy: litellm.Router, endpoint: str, edits: tuple[str, ...]) -> object: + match endpoint: + case "anthropic_messages": + message: Final = await proxy.aanthropic_messages( + model="route", messages=transcript(endpoint, "Edit", edits), max_tokens=16 + ) + return message["content"][0]["text"] # pyright: ignore[reportIndexIssue] # non-streaming response + case "aresponses": + response: Final = await proxy.aresponses(model="route", input=transcript(endpoint, "Edit", edits)) + return response.output[0].content[0].text # pyright: ignore[reportAttributeAccessIssue] # non-streaming response + case _: + chat: Final = await proxy.acompletion(model="route", messages=transcript(endpoint, "Edit", edits)) + return chat.choices[0].message.content # pyright: ignore[reportAttributeAccessIssue] # non-streaming response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint", ["acompletion", "anthropic_messages", "aresponses"]) +async def test_router_routes_agent_traffic_from_every_endpoint(endpoint: str) -> None: + proxy: Final = llm_router(route_config={"models": {"capable": "strong", "efficient": "fast"}}) + assert await routed_reply(proxy, endpoint, ("Edit applied",) * 3) == "fast" + + +@pytest.mark.asyncio +async def test_router_follows_an_edited_switchyard_router_and_forgets_a_deleted_one() -> None: + proxy: Final = llm_router(route_config={"models": {"capable": "strong", "efficient": "fast"}}) + assert await routed_reply(proxy, "acompletion", ("Edit applied",) * 3) == "fast" + + swapped: Final = switchyard_deployment({"models": {"capable": "fast", "efficient": "strong"}}) + _ = proxy.apply(ChangeSet(upsert=[Deployment(**swapped)])) + assert await routed_reply(proxy, "acompletion", ("Edit applied",) * 3) == "strong" + + _ = proxy.apply(ChangeSet(delete=["route-id"])) + with pytest.raises(litellm.BadRequestError): + await proxy.acompletion(model="route", messages=HELLO) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("marked", "expected"), [(True, "strong"), (False, "fast")]) +async def test_anthropic_tool_failure_marker_reaches_switchyard(marked: bool, expected: str) -> None: + failures: Final = transcript("anthropic_messages", "Edit", ("String to replace not found.",) * 4, marked=marked) + model, decision = await route(config(), failures, call_type="anthropic_messages") + assert model == expected + assert "switchyard_fallback_reason" not in decision + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_kwargs", "expected"), + [ + ({"litellm_session_id": "s1"}, ["strong", "strong", "strong", "fast"]), + ( + {"proxy_server_request": {"headers": {"X-Session-Id": "s1", "session-id": "s1"}}}, + ["strong", "fast", "fast", "fast"], + ), + ({}, ["strong", "fast", "fast", "fast"]), + ], +) +async def test_capable_hold_follows_the_session(request_kwargs: Mapping[str, object], expected: list[str]) -> None: + switchyard: Final = router(config()) + turns: Final = (("CUDA out of memory",), *((("All tests passed",) * count) for count in (2, 3, 4))) + picks: Final = [await pick(switchyard, request_kwargs, results) for results in turns] + assert picks == expected + + +async def pick(switchyard: SwitchyardRouter, request_kwargs: Mapping[str, object], results: tuple[str, ...]) -> str: + result: Final = await switchyard.async_pre_routing_hook( + "route", request_kwargs, transcript("acompletion", "bash", results), call_type="acompletion" + ) + assert result is not None + return result.model + + +@pytest.mark.asyncio +async def test_a_session_hold_never_crosses_api_keys() -> None: + switchyard: Final = router(config()) + owner: Final = {"litellm_session_id": "shared", "metadata": {"user_api_key_hash": "owner"}} + other: Final = {"litellm_session_id": "shared", "metadata": {"user_api_key_hash": "other"}} + assert await pick(switchyard, owner, ("CUDA out of memory",)) == "strong" + assert await pick(switchyard, other, ("All tests passed",) * 2) == "fast" + assert await pick(switchyard, owner, ("All tests passed",) * 2) == "strong" + + +@pytest.mark.asyncio +async def test_a_client_metadata_key_hash_cannot_join_another_keys_session() -> None: + switchyard: Final = router(config()) + owner: Final = {"litellm_session_id": "shared", "litellm_metadata": {"user_api_key_hash": "owner"}} + spoofer: Final = { + "litellm_session_id": "shared", + "metadata": {"user_api_key_hash": "owner"}, + "litellm_metadata": {"user_api_key_hash": "spoofer"}, + } + assert await pick(switchyard, owner, ("CUDA out of memory",)) == "strong" + assert await pick(switchyard, spoofer, ("All tests passed",) * 2) == "fast" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("p_solve", "expected"), [(0.0, "strong"), (1.0, "fast")]) +async def test_llm_judge_verdict_selects_the_target(p_solve: float, expected: str) -> None: + async def judge(request: Mapping[str, object]) -> ModelResponse: + assert "autorouter_classifier" in str(request["metadata"]) + verdict: Final = {"crux": "c", "primary_rule": "none", "capability_boundary": "unmatched", "p_solve": p_solve} + return ModelResponse(choices=[{"message": {"role": "assistant", "content": json.dumps(verdict)}}]) + + model, _ = await route(config(classifier={}), HELLO, completion=judge) + assert model == expected + + +def decision_answer(advantage: float) -> OpenAIDecisionResponse: + return OpenAIDecisionResponse.model_validate( + { + "model": "jev", + "answers": [ + { + "type": "choice", + "name": "route", + "choice": "advantage", + "probabilities": [ + {"value": "advantage", "probability": advantage}, + {"value": "no_advantage", "probability": 1 - advantage}, + ], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 0, "total_tokens": 1}, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("parent_bucket", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize("classifier", [{}, {"kind": "decision"}]) +async def test_judge_calls_carry_the_callers_tags_and_origin_where_their_endpoint_reads_them( + parent_bucket: str, classifier: Mapping[str, object] +) -> None: + seen: Final[list[Mapping[str, object]]] = [] # mutable-ok: collects the judge request + + async def judge(request: Mapping[str, object]) -> ModelResponse: + seen.append(request) + verdict: Final = {"crux": "c", "primary_rule": "none", "capability_boundary": "unmatched", "p_solve": 1.0} + return ModelResponse(choices=[{"message": {"role": "assistant", "content": json.dumps(verdict)}}]) + + async def decide(request: Mapping[str, object]) -> OpenAIDecisionResponse: + seen.append(request) + return decision_answer(0.0) + + caller: Final = {parent_bucket: {"tags": ["eu-only"], "inherited_tags": ["team-eu"]}} + await route( + config(classifier=classifier, threshold=1.0), HELLO, completion=judge, decide=decide, request_kwargs=caller + ) + (request,) = seen + read_by_endpoint: Final = "litellm_metadata" if classifier else "metadata" + assert set(request) & {"metadata", "litellm_metadata"} == {read_by_endpoint} + metadata: Final = request[read_by_endpoint] + assert metadata["tags"] == ["eu-only"] # pyright: ignore[reportIndexIssue] # the judge request's metadata dict + assert metadata["inherited_tags"] == ["team-eu"] # pyright: ignore[reportIndexIssue] # same + assert metadata["internal_call_origin"] == "autorouter_classifier" # pyright: ignore[reportIndexIssue] # same + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("advantage", "expected"), [(0.61, "strong"), (0.6, "fast")]) +async def test_decision_model_advantage_must_exceed_the_cutoff(advantage: float, expected: str) -> None: + async def decide(request: Mapping[str, object]) -> OpenAIDecisionResponse: + (question,) = request["questions"] # pyright: ignore[reportGeneralTypeIssues] # the decisions request shape + assert {choice["value"] for choice in question["choices"]} == {"advantage", "no_advantage"} + return decision_answer(advantage) + + model, decision = await route( + config(classifier={"kind": "decision", "cutoff": 0.6}, threshold=1.0), HELLO, decide=decide + ) + assert model == expected + assert decision["switchyard_evidence"]["score"] == advantage + + +async def judge_down(request: Mapping[str, object]) -> ModelResponse: + raise ConnectionError("judge down") + + +async def decision_refused(request: Mapping[str, object]) -> OpenAIDecisionResponse: + return OpenAIDecisionResponse.model_validate( + { + "model": "jev", + "answers": [{"type": "refusal", "name": "route"}], + "usage": {"input_tokens": 1, "output_tokens": 0, "total_tokens": 1}, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("classifier", "router_kwargs", "picker", "expected"), + [ + ({}, {"completion": judge_down}, "efficient_first", "fast"), + ({}, {"completion": judge_down}, "capable_first", "strong"), + ({"kind": "decision"}, {"decide": decision_refused}, "efficient_first", "fast"), + ], +) +async def test_classifier_failure_falls_open_to_the_picker_default( + classifier: Mapping[str, object], router_kwargs: Mapping[str, object], picker: str, expected: str +) -> None: + model, decision = await route(config(classifier=classifier, picker=picker, threshold=1.0), HELLO, **router_kwargs) + assert model == expected + assert "switchyard_fallback_reason" not in decision + + +class EndsWithoutDecision: + async def _steps(self) -> AsyncIterator[object]: + for step in (): + yield step + + def run_stream(self, *args: object, **kwargs: object) -> AsyncIterator[object]: + return self._steps() + + +class RaisesLibsyError: + async def _steps(self) -> AsyncIterator[object]: + raise libsy.LibsyError("algorithm failed") + yield + + def run_stream(self, *args: object, **kwargs: object) -> AsyncIterator[object]: + return self._steps() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route_config", "messages", "call_type", "algorithm", "reason"), + [ + ( + config(), + [{"role": "user", "content": [{"type": "input_audio"}]}], + "acompletion", + None, + "unsupported_content", + ), + (config(), HELLO, None, None, "unsupported_request_type"), + (config(), HELLO, "acompletion", EndsWithoutDecision(), "no_decision"), + (config(), HELLO, "acompletion", RaisesLibsyError(), "libsy_error"), + ( + config(classifier={}, threshold=1.0), + HELLO, + "acompletion", + build_algorithm(config(classifier={"kind": "decision"}, threshold=1.0)), + "unexpected_decision_call", + ), + ( + config(classifier={"kind": "decision"}, threshold=1.0), + HELLO, + "acompletion", + build_algorithm(config(classifier={}, threshold=1.0)), + "unexpected_judge_call", + ), + ], +) +async def test_anything_switchyard_cannot_route_goes_to_capable( + route_config: SwitchyardRouterConfig, messages: object, call_type: str | None, algorithm: object, reason: str +) -> None: + model, decision = await route(route_config, messages, call_type=call_type, algorithm=algorithm) + assert model == "strong" + assert decision["switchyard_fallback_reason"] == reason + + +@pytest.mark.asyncio +async def test_failing_judge_is_called_once_even_when_the_router_retries(monkeypatch: pytest.MonkeyPatch) -> None: + attempts: Final[list[str]] = [] + + class CountAttempts(CustomLogger): + def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None: + attempts.append(model) + + monkeypatch.setattr(litellm, "callbacks", [CountAttempts()]) + proxy: Final = llm_router( + route_config={ + "models": {"capable": "strong", "efficient": "fast", "judge": "judge"}, + "classifier": {}, + "options": {"confidence_threshold": 1.0}, + }, + judge_fails=True, + num_retries=3, + ) + await proxy.acompletion(model="route", messages=HELLO) + assert attempts.count("judge") == 1 + + +@pytest.mark.asyncio +async def test_judge_cancellation_propagates() -> None: + async def cancelled(request: Mapping[str, object]) -> ModelResponse: + raise asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await route(config(classifier={}), HELLO, completion=cancelled) diff --git a/tests/unit/router_strategy/test_router_tag_routing.py b/tests/unit/router_strategy/test_router_tag_routing.py index 0d3876c64a3..e916e5b235d 100644 --- a/tests/unit/router_strategy/test_router_tag_routing.py +++ b/tests/unit/router_strategy/test_router_tag_routing.py @@ -3023,7 +3023,7 @@ class _RewriteToTierStrategy: self.rewrite_to = rewrite_to async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False + self, model, request_kwargs, messages=None, input=None, specific_deployment=False, call_type=None ): from litellm.types.router import PreRoutingHookResponse diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 5169642bc43..76e9a70614f 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -2,6 +2,9 @@ import asyncio import copy import os import functools +import json +import logging +import threading import uuid from typing import Final, cast @@ -10,7 +13,11 @@ import pytest import litellm from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, PROMPT_CACHE_LOOKBACK_POSITIONS +from litellm.constants import ( + DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, + PROMPT_CACHE_LOOKBACK_POSITIONS, + REDACTED_BY_LITELLM, +) from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup @@ -19,6 +26,7 @@ from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import _honors_one_hour_cache_ttl, _get_min_token_count_for_deployments, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import StandardLoggingPayload @@ -138,10 +146,13 @@ def test_bedrock_application_profile_uses_its_base_model_for_one_hour_ttl() -> N }, ) - assert _honors_one_hour_cache_ttl( - standard_logging_object, - {"model": None, "base_model": "bedrock/us.anthropic.claude-sonnet-5-5"}, - ) is True + assert ( + _honors_one_hour_cache_ttl( + standard_logging_object, + {"model": None, "base_model": "bedrock/us.anthropic.claude-sonnet-5-5"}, + ) + is True + ) def test_bedrock_application_profile_cache_hit_reads_base_model_from_metadata_model_info() -> None: @@ -244,6 +255,16 @@ async def test_async_filter_deployments_narrows_prompt_above_model_minimum(): assert filtered == [deployments[1]] +class _ReadTrackingMessages(list[AllMessageValues]): + def __init__(self, messages: list[AllMessageValues]) -> None: + super().__init__(messages) + self.read = False + + def __iter__(self): + self.read = True + return super().__iter__() + + class _PinLookupCounter(DualCache): def __init__(self) -> None: super().__init__() @@ -275,17 +296,19 @@ async def test_async_filter_deployments_skips_prefix_hash_for_a_single_deploymen prompt that hash alone is ~30 ms of GIL-holding work per request. """ cache = _PinLookupCounter() - check = PromptCachingDeploymentCheck(cache=cache) + check = PromptCachingDeploymentCheck(cache=cache, group_has_alternatives=lambda group: True) deployments = _deployments("anthropic/claude-opus-4-6") messages = _messages(word_count=5000) await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-1", messages=messages, tools=None) + tracked_messages = _ReadTrackingMessages(messages) filtered = await check.async_filter_deployments( - model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=messages + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=tracked_messages ) assert filtered == deployments assert cache.pin_lookups == 0 + assert tracked_messages.read is False two = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") assert await check.async_filter_deployments( @@ -313,10 +336,20 @@ async def test_async_filter_deployments_does_not_pin_when_target_order_is_set(): assert filtered == deployments +_RETRIES_ON_DEP_1: Final = tuple({"deployment_id": "dep-1"} for _ in range(4)) + + @pytest.mark.asyncio @pytest.mark.parametrize( "failed_in_request", - [{"_excluded_deployment_ids": ["dep-2"]}, {"_retry_skipped_deployment_ids": ("dep-2",)}], + [ + {"_excluded_deployment_ids": ["dep-2"]}, + {"_retry_skipped_deployment_ids": ("dep-2",)}, + {"metadata": {"previous_models": [{"deployment_id": "dep-2"}], "request_retry_count": 1}}, + {"litellm_metadata": {"previous_models": ({"deployment_id": "dep-2"},), "request_retry_count": 1}}, + {"metadata": {"previous_models": _RETRIES_ON_DEP_1, "request_retry_count": 5}}, + ], + ids=["excluded", "retry_skipped", "retried", "retried_litellm_metadata", "retries_rolled_over"], ) async def test_async_filter_deployments_does_not_pin_to_a_deployment_that_failed_this_request( failed_in_request: dict[str, object], @@ -334,6 +367,34 @@ async def test_async_filter_deployments_does_not_pin_to_a_deployment_that_failed ) == _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") +@pytest.mark.asyncio +@pytest.mark.parametrize( + "retry_metadata", + [ + {"previous_models": _RETRIES_ON_DEP_1, "request_retry_count": 4}, + {"previous_models": "dep-2"}, + {"previous_models": [None, "dep-2", {"deployment_id": None}], "request_retry_count": 3}, + ], + ids=["every_retry_on_record", "malformed_breadcrumbs", "malformed_records"], +) +async def test_async_filter_deployments_keeps_the_pin_while_every_retry_is_on_record( + retry_metadata: dict[str, object], +): + cache: Final = DualCache() + check: Final = PromptCachingDeploymentCheck(cache=cache) + messages: Final = _messages(word_count=5000) + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) + + filtered: Final = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=_deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6"), + messages=messages, + request_kwargs={"metadata": retry_metadata}, + ) + + assert [deployment["model_info"]["id"] for deployment in filtered] == ["dep-2"] + + @pytest.mark.asyncio async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is_lower(): """ @@ -1535,3 +1596,920 @@ async def test_pin_is_written_with_message_logging_off_while_loggers_still_see_r await _eventually(lambda: len(probe.logged_messages) == 5) redacted = [{"role": "user", "content": "redacted-by-litellm"}] assert probe.logged_messages == [(redacted, redacted)] * 5, probe.logged_messages + + +RESPONSES_BRIDGED_MODEL: Final = "anthropic/claude-sonnet-4-6" +RESPONSES_NATIVE_MODEL: Final = "openai/gpt-5.4-mini" +_RESPONSES_TOOL: Final = { + "type": "function", + "name": "get_time", + "parameters": {"type": "object", "properties": {}}, +} + + +def _marked_input_text(text: str) -> dict[str, object]: + return {"type": "input_text", "text": text, "cache_control": {"type": "ephemeral"}} + + +def _marked_responses_input() -> list[dict[str, object]]: + return [ + {"role": "system", "content": [_marked_input_text(LONG_PROMPT)]}, + {"role": "user", "content": "hello"}, + ] + + +def _responses_router(model: str) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": model, "api_key": "sk-fake"}, + "model_info": {"id": model_id}, + } + for model_id in ("dep-1", "dep-2") + ], + optional_pre_call_checks=["prompt_caching"], + ) + + +async def _serve_response(router: litellm.Router, request_input: object) -> str: + response = await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=copy.deepcopy(request_input), + instructions="Be terse.", + tools=[_RESPONSES_TOOL], + mock_response="ok", + ) + return response._hidden_params["model_id"] + + +async def _deployments_serving_repeats(router: litellm.Router, request_input: object) -> set[str]: + await _serve_response(router, request_input) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + return {await _serve_response(router, request_input) for _ in range(8)} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", [RESPONSES_BRIDGED_MODEL, RESPONSES_NATIVE_MODEL], ids=["bridged", "native"]) +async def test_repeated_responses_request_pins_to_one_deployment(model: str): + assert len(await _deployments_serving_repeats(_responses_router(model), _marked_responses_input())) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", [RESPONSES_BRIDGED_MODEL, RESPONSES_NATIVE_MODEL], ids=["bridged", "native"]) +async def test_responses_thread_stays_on_the_deployment_holding_its_prefix(model: str): + router = _responses_router(model) + first_turn = _marked_responses_input() + pinned_deployment = await _serve_response(router, first_turn) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + + later_turns = ( + [*first_turn, {"role": "assistant", "content": "ok"}, {"role": "user", "content": f"question {turn}"}] + for turn in range(8) + ) + + assert {await _serve_response(router, turn) for turn in later_turns} == {pinned_deployment} + + +@pytest.mark.asyncio +async def test_responses_request_pins_after_prompt_caching_is_turned_off_and_on_again(): + router = _responses_router(RESPONSES_NATIVE_MODEL) + await _serve_response(router, _marked_responses_input()) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + router.update_settings(optional_pre_call_checks=[]) + router.update_settings(optional_pre_call_checks=["prompt_caching"]) + router.prompt_caching_pin_cache.in_memory_cache.flush_cache() + + assert len(await _deployments_serving_repeats(router, _marked_responses_input())) == 1 + + +def _team_tagged_router(model: str) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": model, "api_key": "sk-fake", "tags": [team]}, + "model_info": {"id": model_id}, + } + for model_id, team in (("dep-1", "team-a"), ("dep-2", "team-b")) + ], + optional_pre_call_checks=["prompt_caching"], + enable_tag_filtering=True, + ) + + +async def _serve_team_response(router: litellm.Router, team: str) -> str: + response = await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=_marked_responses_input(), + instructions="Be terse.", + tools=[_RESPONSES_TOOL], + litellm_metadata={"tags": [team]}, + mock_response="ok", + ) + return response._hidden_params["model_id"] + + +async def _serve_team_completion(router: litellm.Router, team: str) -> str: + response = await router.acompletion( + model=MODEL_GROUP_ALIAS, + messages=_messages(word_count=5000), + metadata={"tags": [team]}, + mock_response="ok", + ) + return response._hidden_params["model_id"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model, serve", + [(RESPONSES_NATIVE_MODEL, _serve_team_response), (RESPONSES_BRIDGED_MODEL, _serve_team_completion)], + ids=["responses", "chat"], +) +async def test_prompt_cache_pin_never_overrides_the_request_tags(model: str, serve): + router = _team_tagged_router(model) + assert await serve(router, "team-a") == "dep-1" + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + + assert await serve(router, "team-b") == "dep-2" + + +@pytest.mark.asyncio +async def test_retry_moves_off_the_pinned_deployment_that_failed(): + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": { + "model": RESPONSES_BRIDGED_MODEL, + "api_key": "sk-fake", + "mock_response": mock_response, + "weight": weight, + }, + "model_info": {"id": model_id}, + } + for model_id, mock_response, weight in ( + ("dep-1", "litellm.InternalServerError", 1), + ("dep-2", "ok", 1000), + ) + ], + optional_pre_call_checks=["prompt_caching"], + num_retries=3, + retry_after=0, + disable_cooldowns=True, + ) + messages: Final = _messages(word_count=5000) + await PromptCachingCache(cache=router.prompt_caching_pin_cache).async_add_model_id( + model_id="dep-1", messages=messages, tools=None + ) + + response: Final = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages) + + assert response._hidden_params["model_id"] == "dep-2" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reused_as", ["metadata", "litellm_metadata"]) +async def test_pin_holds_for_a_metadata_dict_reused_after_a_request_failed_on_the_pinned_deployment(reused_as: str): + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": RESPONSES_BRIDGED_MODEL, "api_key": "sk-fake", "weight": weight}, + "model_info": {"id": model_id}, + } + for model_id, weight in (("dep-1", 1000), ("dep-2", 1)) + ], + optional_pre_call_checks=["prompt_caching"], + num_retries=6, + retry_after=0, + allowed_fails=100, + ) + router.cooldown_cache.add_deployment_to_cooldown( + model_id="dep-1", + original_exception=litellm.RateLimitError( + message="upstream 429", llm_provider="anthropic", model=RESPONSES_BRIDGED_MODEL + ), + exception_status=429, + cooldown_time=60, + ) + reused_metadata: Final[dict[str, object]] = {} + with pytest.raises(litellm.InternalServerError): + await router.acompletion( + model=MODEL_GROUP_ALIAS, + messages=[{"role": "user", "content": "hello"}], + metadata=reused_metadata, + mock_response="litellm.InternalServerError", + ) + router.cooldown_cache.cooldown_store.evict_local("dep-1") + messages: Final = _messages(word_count=5000) + await PromptCachingCache(cache=router.prompt_caching_pin_cache).async_add_model_id( + model_id="dep-2", messages=messages, tools=None + ) + + response: Final = await router.acompletion( + model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok", **{reused_as: reused_metadata} + ) + + assert response._hidden_params["model_id"] == "dep-2" + + +@pytest.mark.asyncio +async def test_responses_mark_on_the_input_item_pins_to_one_deployment(): + request_input = [ + {"role": "system", "content": LONG_PROMPT, "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "hello"}, + ] + + assert len(await _deployments_serving_repeats(_responses_router(RESPONSES_BRIDGED_MODEL), request_input)) == 1 + + +async def _read_responses_pin( + check: PromptCachingDeploymentCheck, + request: dict[str, object], + models: tuple[str, ...] = (AUTO_CACHING_MODEL, AUTO_CACHING_MODEL), +) -> tuple[str, ...]: + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=_deployments(*models), + messages=None, + request_kwargs=request, + ) + return tuple(str(deployment["model_info"]["id"]) for deployment in filtered) + + +async def _log_responses_success( + check: PromptCachingDeploymentCheck, + logged_messages: object, + metadata: dict[str, object], + model_id: str = "dep-1", + model_map_value: dict[str, float | None] | None = None, + call_type: str = "aresponses", +) -> None: + event = _success_kwargs( + call_type=call_type, + model=AUTO_CACHING_MODEL, + model_id=model_id, + request_messages=logged_messages, + request_kwargs={}, + ) + await check.async_log_success_event( + kwargs={ + **event, + "standard_logging_object": { + **event["standard_logging_object"], + "model_map_information": {"model_map_key": AUTO_CACHING_MODEL, "model_map_value": model_map_value}, + }, + "litellm_params": {"litellm_metadata": metadata}, + }, + response_obj=None, + start_time=None, + end_time=None, + ) + + +async def _pinned_deployment_ids( + first_request: dict[str, object], logged_messages: object, next_request: dict[str, object] +) -> tuple[str, ...]: + check = PromptCachingDeploymentCheck(cache=DualCache()) + metadata: dict[str, object] = {} + await _read_responses_pin(check, {**first_request, "litellm_metadata": metadata}) + await _log_responses_success(check, logged_messages, metadata) + return await _read_responses_pin(check, {**next_request, "litellm_metadata": {}}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "logged_messages", + [ + _marked_responses_input(), + [ + {"role": "system", "content": "Be terse."}, + {"role": "system", "content": LONG_PROMPT, "cache_control": {"type": "ephemeral"}}, + {"role": "assistant", "tool_calls": [{"id": "call_1", "type": "function"}]}, + {"role": "tool", "tool_call_id": "call_1", "content": "12:00"}, + ], + [{"role": "user", "content": REDACTED_BY_LITELLM}], + ], + ids=["client_input", "bridged_chat_messages", "redacted"], +) +async def test_responses_pin_does_not_depend_on_what_the_success_event_logged( + logged_messages: list[dict[str, object]], +): + request = {"input": _marked_responses_input(), "instructions": "Be terse.", "tools": [_RESPONSES_TOOL]} + + pinned = await _pinned_deployment_ids(request, logged_messages, request) + + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "forged_pin", + [["forged-redis-key", 3600], {"cache_key": "forged-redis-key", "marked_ttl_seconds": 3600}], + ids=["array", "object"], +) +async def test_responses_pin_sent_in_the_request_metadata_is_never_written(forged_pin: object): + cache = DualCache() + + await _log_responses_success( + PromptCachingDeploymentCheck(cache=cache), + _marked_responses_input(), + {"responses_prompt_cache_pin": forged_pin}, + ) + + assert not cache.in_memory_cache.cache_dict + + +@pytest.mark.asyncio +async def test_interleaved_responses_requests_each_pin_their_own_prompt(): + check = PromptCachingDeploymentCheck(cache=DualCache()) + terse = {"input": _marked_responses_input(), "instructions": "Be terse."} + verbose = {"input": _marked_responses_input(), "instructions": "Be verbose."} + terse_metadata: dict[str, object] = {} + verbose_metadata: dict[str, object] = {} + await _read_responses_pin(check, {**terse, "litellm_metadata": terse_metadata}) + await _read_responses_pin(check, {**verbose, "litellm_metadata": verbose_metadata}) + + await _log_responses_success(check, _marked_responses_input(), terse_metadata, model_id="dep-1") + await _log_responses_success(check, _marked_responses_input(), verbose_metadata, model_id="dep-2") + + assert await _read_responses_pin(check, {**terse, "litellm_metadata": {}}) == ("dep-1",) + assert await _read_responses_pin(check, {**verbose, "litellm_metadata": {}}) == ("dep-2",) + + +@pytest.mark.asyncio +async def test_retried_responses_request_leaves_no_pin_in_flight_once_it_succeeds(): + pins_in_flight = InMemoryCache() + check = PromptCachingDeploymentCheck(cache=DualCache(), responses_pins_in_flight=pins_in_flight) + metadata: dict[str, object] = {} + request = {"input": _marked_responses_input(), "litellm_metadata": metadata} + await _read_responses_pin(check, request) + await _read_responses_pin(check, request) + + await _log_responses_success(check, _marked_responses_input(), metadata) + + assert not pins_in_flight.cache_dict + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["acompact_responses", "acreate_interaction"]) +async def test_request_sharing_the_responses_input_shape_leaves_no_pin_behind(call_type: str): + cache = DualCache() + pins_in_flight = InMemoryCache() + check = PromptCachingDeploymentCheck(cache=cache, responses_pins_in_flight=pins_in_flight) + metadata: dict[str, object] = {} + await _read_responses_pin(check, {"input": _marked_responses_input(), "litellm_metadata": metadata}) + assert pins_in_flight.cache_dict + + await _log_responses_success(check, _marked_responses_input(), metadata, call_type=call_type) + + assert not pins_in_flight.cache_dict + assert not cache.in_memory_cache.cache_dict + + +@pytest.mark.asyncio +async def test_responses_request_served_by_the_only_healthy_deployment_pins_it(): + cache = _PinLookupCounter() + check = PromptCachingDeploymentCheck(cache=cache, group_has_alternatives=lambda group: True) + request = {"input": _marked_responses_input(), "instructions": "Be terse."} + only_healthy = [_deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)[1]] + metadata: dict[str, object] = {} + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=only_healthy, + messages=None, + request_kwargs={**request, "litellm_metadata": metadata}, + ) + assert filtered == only_healthy + assert cache.pin_lookups == 0 + + await _log_responses_success(check, _marked_responses_input(), metadata, model_id="dep-2") + + assert await _read_responses_pin(check, {**request, "litellm_metadata": {}}) == ("dep-2",) + + +@pytest.mark.asyncio +async def test_responses_request_to_a_single_deployment_group_skips_the_prefix_hash(): + cache = _PinLookupCounter() + pins_in_flight = InMemoryCache() + check = PromptCachingDeploymentCheck( + cache=cache, responses_pins_in_flight=pins_in_flight, group_has_alternatives=lambda group: False + ) + deployments = _deployments(AUTO_CACHING_MODEL) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=None, + request_kwargs={"input": _marked_responses_input(), "litellm_metadata": {}}, + ) + + assert filtered == deployments + assert cache.pin_lookups == 0 + assert not pins_in_flight.cache_dict + + +@pytest.mark.asyncio +async def test_responses_request_to_a_single_deployment_router_group_writes_no_pin(monkeypatch): + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": RESPONSES_NATIVE_MODEL, "api_key": "sk-fake"}, + "model_info": {"id": "dep-1"}, + } + ], + optional_pre_call_checks=["prompt_caching"], + ) + logged_after_the_check: list[str] = [] + + class _LoggedAfterTheCheck(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + logged_after_the_check.append(kwargs["litellm_call_id"]) + + monkeypatch.setattr(litellm, "callbacks", [*litellm.callbacks, _LoggedAfterTheCheck()]) + + assert await _serve_response(router, _marked_responses_input()) == "dep-1" + + assert await _eventually(lambda: logged_after_the_check) + assert not router.prompt_caching_pin_cache.in_memory_cache.cache_dict + + +def _distinct_marked_responses_input(index: int) -> list[dict[str, object]]: + return [ + {"role": "system", "content": [_marked_input_text(f"prompt {index} {LONG_PROMPT}")]}, + {"role": "user", "content": "hello"}, + ] + + +@pytest.mark.asyncio +async def test_responses_prompts_served_while_a_peer_cools_down_stay_pinned_once_it_recovers(): + router = _responses_router(RESPONSES_NATIVE_MODEL) + pins = router.prompt_caching_pin_cache.in_memory_cache.cache_dict + router.cooldown_cache.add_deployment_to_cooldown( + model_id="dep-1", + original_exception=litellm.RateLimitError( + message="upstream 429", llm_provider="openai", model=RESPONSES_NATIVE_MODEL + ), + exception_status=429, + cooldown_time=60, + ) + prompts = [_distinct_marked_responses_input(index) for index in range(8)] + for prompt in prompts: + pins_before = len(pins) + assert await _serve_response(router, prompt) == "dep-2" + assert await _eventually(lambda before=pins_before: len(pins) > before) + router.cooldown_cache.cooldown_store.evict_local("dep-1") + unpinned = { + ( + await router.aresponses(model=MODEL_GROUP_ALIAS, input=f"unpinned {index}", mock_response="ok") + )._hidden_params["model_id"] + for index in range(16) + } + assert unpinned == {"dep-1", "dep-2"} + + assert {await _serve_response(router, prompt) for prompt in prompts} == {"dep-2"} + + +@pytest.mark.asyncio +async def test_responses_with_other_instructions_are_a_different_prefix(): + pinned = await _pinned_deployment_ids( + {"input": _marked_responses_input(), "instructions": "Be terse."}, + _marked_responses_input(), + {"input": _marked_responses_input(), "instructions": "Be verbose."}, + ) + + assert len(pinned) == 2 + + +@pytest.mark.asyncio +async def test_responses_with_another_tool_set_are_a_different_prefix(): + other_tool = {**_RESPONSES_TOOL, "name": "get_weather"} + + pinned = await _pinned_deployment_ids( + {"input": _marked_responses_input(), "tools": [_RESPONSES_TOOL]}, + _marked_responses_input(), + {"input": _marked_responses_input(), "tools": [other_tool]}, + ) + + assert len(pinned) == 2 + + +@pytest.mark.asyncio +async def test_responses_custom_tools_with_another_grammar_are_a_different_prefix(): + custom_tool: Final = { + "type": "custom", + "name": "run_sql", + "description": "Run a query", + "format": {"type": "grammar", "syntax": "lark", "definition": "start: SELECT"}, + } + other_grammar: Final = {**custom_tool, "format": {"type": "grammar", "syntax": "lark", "definition": "start: UPDATE"}} + + pinned = await _pinned_deployment_ids( + {"input": _marked_responses_input(), "tools": [custom_tool]}, + _marked_responses_input(), + {"input": _marked_responses_input(), "tools": [other_grammar]}, + ) + + assert len(pinned) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "parameters", + [{"type": "object", "properties": {}}, {"properties": {"city": {"type": "string"}}}], + ids=["typed_parameters", "untyped_parameters"], +) +async def test_chat_and_responses_requests_sending_the_same_tool_share_one_prefix(parameters: dict[str, object]): + described_tool: Final = {**_RESPONSES_TOOL, "description": "Get the time", "parameters": parameters} + chat_tool: Final = { + "type": "function", + "function": {key: value for key, value in described_tool.items() if key != "type"}, + } + router: Final = _responses_router(RESPONSES_BRIDGED_MODEL) + await router.acompletion( + model=MODEL_GROUP_ALIAS, + messages=[ + {"role": "system", "content": [{**_marked_input_text(LONG_PROMPT), "type": "text"}]}, + {"role": "user", "content": "hello"}, + ], + tools=[chat_tool], + mock_response="ok", + ) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + chat_pin_keys: Final = set(router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + router.prompt_caching_pin_cache.in_memory_cache.flush_cache() + + await router.aresponses( + model=MODEL_GROUP_ALIAS, input=_marked_responses_input(), tools=[described_tool], mock_response="ok" + ) + + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + assert set(router.prompt_caching_pin_cache.in_memory_cache.cache_dict) == chat_pin_keys + + +@pytest.mark.asyncio +async def test_responses_marked_tool_set_long_enough_to_cache_pins_on_its_own(): + marked_tool = {**_RESPONSES_TOOL, "description": LONG_PROMPT, "cache_control": {"type": "ephemeral"}} + request = {"input": "hello", "tools": [marked_tool]} + + pinned = await _pinned_deployment_ids(request, [{"role": "user", "content": "hello"}], request) + + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +async def test_responses_continuation_turn_is_left_to_the_deployment_holding_its_previous_response(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + new_thread = {"input": "hello", "instructions": LONG_PROMPT} + + pinned = await _pinned_deployment_ids( + new_thread, [{"role": "user", "content": "hello"}], {**new_thread, "previous_response_id": "resp_1"} + ) + + assert len(pinned) == 2 + + +@pytest.mark.asyncio +async def test_responses_answer_from_a_fallback_group_keeps_the_primary_groups_pin(): + check = PromptCachingDeploymentCheck(cache=DualCache()) + request = {"input": _marked_responses_input()} + assert await _serve_turn(check, request["input"]) == 2 + fallback_hop_metadata: dict[str, object] = {} + await _read_responses_pin(check, {**request, "litellm_metadata": fallback_hop_metadata}) + + await _log_responses_success(check, request["input"], fallback_hop_metadata, model_id="fallback-dep") + + pinned = await _read_responses_pin(check, {**request, "litellm_metadata": {}}) + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "answering_deployment,expected_pin", + [("dep-1", ("dep-1", "dep-2")), ("dep-2", ("dep-2",))], + ids=["below_its_minimum", "above_its_minimum"], +) +async def test_responses_pin_goes_only_to_a_deployment_whose_model_caches_the_prefix( + answering_deployment: str, expected_pin: tuple[str, ...] +): + mixed_group = ("anthropic/claude-opus-4-5", "anthropic/claude-sonnet-4-5") + check = PromptCachingDeploymentCheck(cache=DualCache()) + request = {"input": [{"role": "user", "content": [_marked_input_text("word " * 1400)]}]} + metadata: dict[str, object] = {} + await _read_responses_pin(check, {**request, "litellm_metadata": metadata}, mixed_group) + + await _log_responses_success(check, request["input"], metadata, model_id=answering_deployment) + + pinned = await _read_responses_pin(check, {**request, "litellm_metadata": {}}, mixed_group) + assert pinned == expected_pin + + +@pytest.mark.asyncio +async def test_embedding_request_never_takes_a_responses_pin(): + check = PromptCachingDeploymentCheck(cache=DualCache()) + pinned_input = _marked_responses_input() + assert await _serve_turn(check, pinned_input) == 2 + assert len(await _read_responses_pin(check, {"input": pinned_input, "litellm_metadata": {}})) == 1 + + filtered = await _read_responses_pin(check, {"input": pinned_input, "metadata": {}}) + + assert len(filtered) == 2 + + +class _InputReadOnlyOnceTheLoopTurns(list[dict[str, object]]): + def __init__(self, items: list[dict[str, object]], loop_turned: threading.Event) -> None: + super().__init__(items) + self.loop_turned: Final = loop_turned + self.loop_turned_before_first_read: bool | None = None + + def __iter__(self): + if self.loop_turned_before_first_read is None: + self.loop_turned_before_first_read = self.loop_turned.wait(timeout=5) + return super().__iter__() + + +@pytest.mark.asyncio +async def test_responses_thread_is_converted_off_the_event_loop(): + loop_turned = threading.Event() + thread = _InputReadOnlyOnceTheLoopTurns(_marked_responses_input(), loop_turned) + asyncio.get_running_loop().call_soon(loop_turned.set) + + await _read_responses_pin( + PromptCachingDeploymentCheck(cache=DualCache()), {"input": thread, "litellm_metadata": {}} + ) + + assert thread.loop_turned_before_first_read is True + + +@pytest.mark.asyncio +async def test_responses_long_instructions_count_toward_the_caching_minimum(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + request = {"input": "hello", "instructions": LONG_PROMPT} + + pinned = await _pinned_deployment_ids(request, [{"role": "user", "content": "hello"}], request) + + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_map_value,pinned_past_five_minutes", + [ + pytest.param({"cache_creation_input_token_cost_above_1hr": 6.6e-06}, True, id="one-hour-cache-writes"), + pytest.param({"cache_creation_input_token_cost_above_1hr": None}, False, id="no-one-hour-cache-writes"), + ], +) +async def test_responses_one_hour_pin_follows_the_deployments_one_hour_cache_support( + model_map_value: dict[str, float | None], pinned_past_five_minutes: bool +): + clock = _ManualClock() + check = PromptCachingDeploymentCheck(cache=DualCache(in_memory_cache=InMemoryCache(clock=clock))) + marked = {**_marked_input_text(LONG_PROMPT), "cache_control": {"type": "ephemeral", "ttl": "1h"}} + request = {"input": [{"role": "user", "content": [marked]}]} + metadata: dict[str, object] = {} + await _read_responses_pin(check, {**request, "litellm_metadata": metadata}) + await _log_responses_success(check, request["input"], metadata, model_map_value=model_map_value) + + clock.advance(PROMPT_CACHE_PIN_TTL_SECONDS + 1) + + assert (len(await _read_responses_pin(check, {**request, "litellm_metadata": {}})) == 1) is pinned_past_five_minutes + + +@pytest.mark.asyncio +async def test_responses_turn_without_input_leaves_routing_to_its_previous_response(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + request = {"previous_response_id": "resp_1", "instructions": LONG_PROMPT} + + pinned = await _pinned_deployment_ids(request, None, request) + + assert len(pinned) == 2 + + +async def _serve_turn(check: PromptCachingDeploymentCheck, request_input: object) -> int: + metadata: dict[str, object] = {} + pinned = await _read_responses_pin(check, {"input": request_input, "litellm_metadata": metadata}) + await _log_responses_success(check, request_input, metadata) + return len(pinned) + + +@pytest.mark.asyncio +async def test_responses_thread_longer_than_the_lookback_window_keeps_its_deployment(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + check = PromptCachingDeploymentCheck(cache=DualCache()) + turns = [ + {"role": "system", "content": LONG_PROMPT}, + *( + {"role": "assistant" if index % 2 else "user", "content": f"turn {index}"} + for index in range(2 * PROMPT_CACHE_LOOKBACK_POSITIONS) + ), + ] + + served = [await _serve_turn(check, turns[:end]) for end in range(2, len(turns) + 1, 2)] + + assert served[1:] == [1] * (len(served) - 1) + + +@pytest.mark.asyncio +async def test_responses_failover_writes_the_pin_for_the_deployment_that_answered(): + check = PromptCachingDeploymentCheck(cache=DualCache()) + request = {"input": _marked_responses_input()} + metadata: dict[str, object] = {} + await _read_responses_pin(check, {**request, "litellm_metadata": metadata}) + await _read_responses_pin(check, {**request, "litellm_metadata": metadata}) + await _log_responses_success(check, _marked_responses_input(), metadata, model_id="dep-2") + + pinned = await _read_responses_pin(check, {**request, "litellm_metadata": {}}) + + assert pinned == ("dep-2",) + + +@pytest.mark.asyncio +async def test_responses_pin_covers_only_what_follows_the_newest_compaction(): + compaction = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "summary"}), + } + prefix = {"role": "user", "content": [_marked_input_text(LONG_PROMPT)]} + + first_input = [{"role": "user", "content": "dropped turn"}, compaction, prefix] + + pinned = await _pinned_deployment_ids( + {"input": first_input}, + first_input, + {"input": [{"role": "user", "content": "another dropped turn"}, compaction, prefix]}, + ) + + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +async def test_responses_image_in_the_prefix_pins_to_one_deployment(): + request_input = [ + { + "role": "user", + "content": [{"type": "input_image", "image_url": ONE_PIXEL_PNG}, _marked_input_text(LONG_PROMPT)], + }, + {"role": "user", "content": "hello"}, + ] + + assert len(await _deployments_serving_repeats(_responses_router(RESPONSES_BRIDGED_MODEL), request_input)) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_input", + [[{"role": "system", "content": LONG_PROMPT}, {"role": "user", "content": "hello"}], LONG_PROMPT], + ids=["item_list", "plain_string"], +) +async def test_responses_auto_cached_prefix_pins_to_one_deployment(monkeypatch, request_input: object): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + + assert len(await _deployments_serving_repeats(_responses_router(RESPONSES_BRIDGED_MODEL), request_input)) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_input", + [ + [LONG_PROMPT, LONG_PROMPT], + [{"role": "user", "content": 7}], + 7, + [{"type": "web_search_call"}], + [{"type": "web_search_call", "id": "ws_1", "status": "completed", "action": {"type": "search"}}], + [{"type": ["message"], "role": "user", "content": LONG_PROMPT}], + [{"role": "user", "content": [{"type": "input_image"}]}], + ], + ids=[ + "embedding_strings", + "non_text_item_content", + "integer", + "web_search_call_without_id", + "search_action_without_query", + "unhashable_item_type", + "image_without_a_source", + ], +) +async def test_input_that_is_not_a_responses_prompt_never_narrows_or_raises(request_input: object): + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + metadata: dict[str, object] = {} + + filtered = await _read_responses_pin(check, {"input": request_input, "litellm_metadata": metadata}) + await _log_responses_success(check, request_input, metadata) + + assert filtered == ("dep-1", "dep-2") + assert not cache.in_memory_cache.cache_dict + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "next_tools,expected_pin", + [([{"type": "image_generation"}], ("dep-1",)), ([{"type": "web_search"}], ("dep-1", "dep-2"))], + ids=["same_native_tool", "other_native_tool"], +) +async def test_responses_native_tools_are_part_of_the_prefix( + next_tools: list[dict[str, object]], expected_pin: tuple[str, ...] +): + request_input = _marked_responses_input() + + pinned = await _pinned_deployment_ids( + {"input": request_input, "tools": [{"type": "image_generation"}]}, + request_input, + {"input": request_input, "tools": next_tools}, + ) + + assert pinned == expected_pin + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native_tool_type", ["shell", "image_generation", "computer_use", "local_shell"]) +async def test_responses_native_tool_routed_to_a_native_model_logs_no_dropped_tool_warning( + caplog: pytest.LogCaptureFixture, native_tool_type: str +): + router = _responses_router(RESPONSES_NATIVE_MODEL) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=_marked_responses_input(), + tools=[{"type": native_tool_type}], + mock_response="ok", + ) + + assert not [record for record in caplog.records if "Dropping Responses API tool" in record.getMessage()] + + +class _MetadataSerializingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.serialized_metadata: list[str] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.serialized_metadata.append(json.dumps(kwargs["litellm_params"]["metadata"])) + + +@pytest.mark.asyncio +async def test_logger_serializing_the_request_metadata_never_sees_the_prompt_cache_key(monkeypatch): + serializing_logger = _MetadataSerializingLogger() + monkeypatch.setattr(litellm, "callbacks", [serializing_logger]) + router = _responses_router(RESPONSES_NATIVE_MODEL) + + assert len(await _deployments_serving_repeats(router, _marked_responses_input())) == 1 + + assert await _eventually(lambda: len(serializing_logger.serialized_metadata) == 9) + logged_metadata = "".join(serializing_logger.serialized_metadata) + pin_keys = tuple(router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + assert pin_keys + assert not [pin_key for pin_key in pin_keys if pin_key.split(":")[1] in logged_metadata] + + +@pytest.mark.asyncio +async def test_responses_call_hands_back_the_callers_metadata_without_the_pin_token(): + router: Final = _responses_router(RESPONSES_BRIDGED_MODEL) + caller_metadata: Final[dict[str, object]] = {"tags": ["team-a"]} + + await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=_marked_responses_input(), + litellm_metadata=caller_metadata, + mock_response="ok", + ) + + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + assert "responses_prompt_cache_pin" not in caller_metadata + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failed_in_request,expected_deployments", + [({}, ("dep-2",)), ({"_excluded_deployment_ids": ["dep-2"]}, ("dep-1", "dep-2"))], + ids=["previous_deployment_healthy", "previous_deployment_failed_this_request"], +) +async def test_responses_continuation_turn_goes_to_the_deployment_that_answered_its_previous_response( + failed_in_request: dict[str, object], expected_deployments: tuple[str, ...] +): + previous_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id("anthropic", "dep-2", "resp_x") + request: dict[str, object] = { + "input": "hello", + "previous_response_id": previous_response_id, + "litellm_metadata": {}, + } + + pinned = await _read_responses_pin( + PromptCachingDeploymentCheck(cache=DualCache()), {**request, **failed_in_request} + ) + + assert pinned == expected_deployments + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_tool", ["get_time", None, 42], ids=["string", "null", "number"]) +async def test_responses_request_with_a_malformed_tool_is_left_to_the_provider(bad_tool: object): + request: dict[str, object] = {"input": _marked_responses_input(), "tools": [bad_tool], "litellm_metadata": {}} + + pinned = await _read_responses_pin(PromptCachingDeploymentCheck(cache=DualCache()), request) + + assert pinned == ("dep-1", "dep-2") diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 31fef208a8c..d285d6fa619 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -81,6 +81,7 @@ _IN_MEMORY_ONLY_CALLERS: Final = frozenset( "litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py", "litellm/responses/litellm_completion_transformation/transformation.py", "litellm/router_utils/client_initalization_utils.py", + "litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py", "litellm/router_utils/router_callbacks/track_deployment_metrics.py", "litellm/secret_managers/cyberark_secret_manager.py", "litellm/secret_managers/google_secret_manager.py", diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index f803d9e0b26..3dfaf86b29c 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_router/core/scripted_host.py b/tests/unit/test_router/core/scripted_host.py index 33164bd4408..066faad4db8 100644 --- a/tests/unit/test_router/core/scripted_host.py +++ b/tests/unit/test_router/core/scripted_host.py @@ -166,6 +166,7 @@ def dynamic_config(**changes: object) -> RouterDynamicConfig: # kwargs-ok: Rout context_window_fallbacks=None, content_policy_fallbacks=None, drop_params=None, + disable_token_counter=False, provider_list=litellm.provider_list, expose_router_debug_in_errors=False, redact_user_api_key_info=None, diff --git a/tests/unit/test_router/core/test_construction_precall_characterization_precall.py b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py index 9cdb3f55d17..c9e04117f4d 100644 --- a/tests/unit/test_router/core/test_construction_precall_characterization_precall.py +++ b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py @@ -27,6 +27,7 @@ PICKS: Final = 60 # three-way shuffles miss a deployment in 20 picks about once SHORT: Final = [{"role": "user", "content": "hi"}] LONG_TEXT: Final = " ".join(["word"] * 100) LONG: Final = [{"role": "user", "content": LONG_TEXT}] +OPENAI_EMBEDDER: Final = "text-embedding-3-small" EMBEDDING_IMAGE: Final = f"data:image/png;base64,{base64.b64encode(bytes(range(256)) * 1024).decode()}" IMAGE: Final = [ { @@ -86,6 +87,14 @@ def _small_and_big( ) +def _small_and_big_openai_embedders(make_core: MakeCore, small: int) -> RouterCore: + counted_as: Final = {"model_info": {"base_model": OPENAI_EMBEDDER}} + return make_core( + model_list=[windowed("g", "small", small, **counted_as), windowed("g", "big", 10000, **counted_as)], + enable_pre_call_checks=True, + ) + + def _json_mode_group(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: constructor arguments return make_core( model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL), deployment("g", "new", model=JSON_MODE_MODEL)], @@ -279,6 +288,43 @@ async def test_an_embedding_image_stays_out_of_the_context_window_on_a_backend_t assert await _picked(core, times=2, input=EMBEDDING_IMAGE) == {"image"} +@pytest.mark.asyncio +async def test_an_embedding_input_over_the_smaller_window_in_bytes_but_not_in_tokens_fits_it( + make_core: MakeCore, +) -> None: + short_words: Final = "word " * 8 + assert litellm.token_counter(model=OPENAI_EMBEDDER, text=short_words) <= 10 < len(short_words.encode()) + assert await _picked(_small_and_big_openai_embedders(make_core, small=10), input=["hi", short_words]) == { + "small", + "big", + } + + +@pytest.mark.asyncio +async def test_an_embedding_input_that_fits_the_window_in_characters_but_not_in_tokens_is_counted( + make_core: MakeCore, +) -> None: + dense: Final = "\ua66e" * 40 + assert len(dense) <= 50 < litellm.token_counter(model=OPENAI_EMBEDDER, text=dense) + assert await _picked(_small_and_big_openai_embedders(make_core, small=50), times=10, input=["hi", dense]) == {"big"} + + +@pytest.mark.asyncio +async def test_an_embedding_batch_is_never_refused_on_the_default_tokenizers_count(make_core: MakeCore) -> None: + dense: Final = "\ua66e" * 40 + assert litellm.token_counter(text=dense) > 50 + assert await _picked(_small_and_big(make_core, small=50), input=["hi", dense]) == {"small", "big"} + + +@pytest.mark.asyncio +async def test_an_embedding_input_with_a_lone_surrogate_is_counted(make_core: MakeCore) -> None: + unpaired: Final = "\ud800" + "word " * 20 + assert litellm.token_counter(model=OPENAI_EMBEDDER, text=unpaired) > 10 + assert await _picked(_small_and_big_openai_embedders(make_core, small=10), times=10, input=["hi", unpaired]) == { + "big" + } + + @pytest.mark.asyncio async def test_an_embedding_image_counts_as_text_on_a_text_only_backend(make_core: MakeCore) -> None: core: Final = make_core( diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 8f8676103b2..ca61fcd0905 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -4396,6 +4396,48 @@ def test_context_window_pre_call_check_counts_an_embedding_image_as_text_on_a_te ) +@pytest.mark.parametrize( + ("backend", "fits"), + [ + pytest.param("vertex_ai/gemini-embedding-2", True, id="image-backend-without-an-openai-tokenizer"), + pytest.param("text-embedding-3-small", False, id="text-only-backend-counts-the-image"), + ], +) +def test_context_window_pre_call_check_on_a_caption_next_to_an_embedding_image(backend: str, fits: bool, caplog): + caption: Final = "word " * 40 + router: Final = _embedding_context_router(backend, window=20) + deployments: Final = router.registry().deployments("embedder") + + with caplog.at_level(logging.ERROR, logger="LiteLLM Router"): + if fits: + kept: Final = router._pre_call_checks( + model="embedder", healthy_deployments=deployments, input=[caption, _EMBEDDING_IMAGE_INPUT] + ) + assert len(kept) == 1 + else: + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="embedder", healthy_deployments=deployments, input=[caption, _EMBEDDING_IMAGE_INPUT] + ) + + assert "failed to count tokens" not in caplog.text + + +@pytest.mark.asyncio +async def test_an_async_pre_call_check_reads_embedding_media_off_the_event_loop(): + from tests.unit.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags + + router: Final = _embedding_context_router("vertex_ai/gemini-embedding-2", window=20) + caption_and_images: Final = ["a red square", *[_EMBEDDING_IMAGE_INPUT] * 400] + + deployment, took, lags = await timed_with_loop_lags( + lambda: router.select("embedder", {}, input=caption_and_images) + ) + + assert deployment["litellm_params"]["model"] == "vertex_ai/gemini-embedding-2" + assert_loop_stayed_free(took, lags) + + async def test_a_wildcard_route_checks_the_context_window_with_the_requested_openai_models_tokenizer(): limit: Final = ( litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT) @@ -4547,6 +4589,173 @@ async def test_aresponses_enforces_context_window_pre_call_check(): ) +def _embedding_router(max_input_tokens: int) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "sk-test", + "mock_response": [0.1, 0.2], + }, + "model_info": {"max_input_tokens": max_input_tokens}, + } + ], + enable_pre_call_checks=True, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "embedding_input", + [ + pytest.param(["one two three four five six"] * 3, id="texts"), + pytest.param([list(range(6))] * 3, id="token-id-lists"), + pytest.param(list(range(10)), id="token-ids-at-the-limit"), + ], +) +async def test_aembedding_accepts_a_batch_whose_inputs_each_fit_the_context_window(embedding_input, caplog): + router: Final = _embedding_router(max_input_tokens=10) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Router"): + response: Final = await router.aembedding(model="embed", input=embedding_input) + + assert response.data + assert "failed to count tokens" not in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "embedding_input", + [ + pytest.param(["one", " ".join(["word"] * 20)], id="texts"), + pytest.param(list(range(11)), id="token-ids"), + pytest.param([[1], list(range(11))], id="token-id-lists"), + ], +) +async def test_aembedding_rejects_a_batch_with_one_input_over_the_context_window(embedding_input): + router: Final = _embedding_router(max_input_tokens=10) + + with pytest.raises(litellm.ContextWindowExceededError): + await router.aembedding(model="embed", input=embedding_input) + + +def test_embedding_rejects_a_batch_with_one_input_over_the_context_window(): + router: Final = _embedding_router(max_input_tokens=10) + + with pytest.raises(litellm.ContextWindowExceededError): + router.embedding(model="embed", input=["one", " ".join(["word"] * 20)]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "embedding_input", + [ + pytest.param(["one", " ".join(["word"] * 20)], id="texts"), + pytest.param(list(range(11)), id="token-ids"), + pytest.param([[1], list(range(11))], id="token-id-lists"), + ], +) +async def test_aembedding_skips_the_context_window_check_when_token_counting_is_disabled(embedding_input, monkeypatch): + monkeypatch.setattr(litellm, "disable_token_counter", True) + router: Final = _embedding_router(max_input_tokens=10) + + response: Final = await router.aembedding(model="embed", input=embedding_input) + + assert response.data + + +async def _until_every_count_slot_is_taken() -> None: + from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + from litellm.litellm_core_utils.token_counter import _count_offload_limiter_for_this_loop + + while _count_offload_limiter_for_this_loop().borrowed_tokens < TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_a_chat_pre_call_count_does_not_wait_behind_long_embedding_batches(): + from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + + window: Final = {"max_input_tokens": 10000} + router: Final = litellm.Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "sk-test", + "mock_response": [0.1], + }, + "model_info": window, + }, + { + "model_name": "chat", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "mock_response": "hi"}, + "model_info": window, + }, + ], + enable_pre_call_checks=True, + ) + batch: Final = ["word " * 2100] * 76 + batches: Final = tuple( + asyncio.create_task(router.aembedding(model="embed", input=batch)) + for _ in range(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) + ) + await asyncio.wait_for(_until_every_count_slot_is_taken(), timeout=30) + chat: Final = asyncio.create_task(router.acompletion(model="chat", messages=[{"role": "user", "content": "hi"}])) + + first_done, _ = await asyncio.wait((chat, *batches), return_when=asyncio.FIRST_COMPLETED) + await asyncio.gather(chat, *batches) + + assert first_done == {chat} + + +def test_a_routing_plugin_counts_an_embedding_text_batch_by_its_longest_text(): + long_text: Final = "word " * 4000 + context: Final = RouterContext(_embedding_router(max_input_tokens=10)) + + assert context.count_pre_call_check_tokens(None, ["one", long_text]) == litellm.token_counter(text=long_text) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_model_cost_map", "isolate_router_model_cost_state") +@pytest.mark.parametrize( + "prompt_kwargs", + [ + pytest.param({"instructions": "you are a helpful assistant. " * 20}, id="instructions"), + pytest.param( + { + "tools": [ + { + "type": "function", + "name": "lookup", + "description": "look up a record by its identifier. " * 20, + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + } + ] + }, + id="tools", + ), + ], +) +async def test_aresponses_with_empty_input_still_counts_its_other_prompt_content(prompt_kwargs): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "small-ctx", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, + "model_info": {"max_input_tokens": 50}, + } + ], + enable_pre_call_checks=True, + ) + + with pytest.raises(litellm.ContextWindowExceededError): + await router.aresponses(model="small-ctx", input=[], **prompt_kwargs) + + def test_get_deployment_model_info_base_model_flow(): """Test that get_deployment_model_info correctly handles the base model flow""" from unittest.mock import patch @@ -8797,7 +9006,7 @@ class TestConsumedRequestTagsStamp: self.rewrite_to = rewrite_to async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False + self, model, request_kwargs, messages=None, input=None, specific_deployment=False, call_type=None ): from litellm.types.router import PreRoutingHookResponse @@ -8888,7 +9097,7 @@ class TestClaudeCodeSubagentSessionRouterBinding: self.routed_model = routed_model async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False + self, model, request_kwargs, messages=None, input=None, specific_deployment=False, call_type=None ): from litellm.types.router import PreRoutingHookResponse @@ -9942,7 +10151,7 @@ class TestModelGroupAliasReachesPreRoutingStrategies: class _RewriteStrategy: async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False + self, model, request_kwargs, messages=None, input=None, specific_deployment=False, call_type=None ): from litellm.types.router import PreRoutingHookResponse @@ -10059,7 +10268,7 @@ class TestTeamPublicNameReachesPreRoutingStrategies: self.rewrite_to = rewrite_to async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False + self, model, request_kwargs, messages=None, input=None, specific_deployment=False, call_type=None ): from litellm.types.router import PreRoutingHookResponse @@ -10296,7 +10505,7 @@ class TestAutoRouterCompressionDecoupling: self.received_messages: list[dict] | None = None async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False + self, model, request_kwargs, messages=None, input=None, specific_deployment=False, call_type=None ): from litellm.types.router import PreRoutingHookResponse 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__() diff --git a/tests/unit/types/llms/test_types_llms_openai.py b/tests/unit/types/llms/test_types_llms_openai.py index 03d6948950f..a3399c0654e 100644 --- a/tests/unit/types/llms/test_types_llms_openai.py +++ b/tests/unit/types/llms/test_types_llms_openai.py @@ -579,6 +579,35 @@ def test_shell_call_output_the_sdk_cannot_type_serializes_as_sent(): assert json.loads(response.model_dump_json())["output"] == [item] +@pytest.mark.parametrize( + "dump", + [ + lambda response: response.model_dump(exclude_none=True), + lambda response: json.loads(response.model_dump_json(exclude_none=True)), + ], + ids=["python", "json"], +) +def test_function_call_and_tool_dump_async_under_its_wire_name(dump) -> None: + call: Final = { + "id": "fc_1", + "type": "function_call", + "call_id": "call_1", + "name": "get_weather", + "arguments": "{}", + "status": "completed", + "async": True, + } + tool: Final = {"type": "function", "name": "get_weather", "parameters": {}, "strict": True, "async": False} + response: Final = ResponsesAPIResponse.model_validate( + {"id": "resp_1", "created_at": 1, "object": "response", "status": "completed", "output": [call], "tools": [tool]} + ) + + dumped: Final = dump(response) + + assert (dumped["output"], dumped["tools"]) == ([call], [tool]) + assert ResponsesAPIResponse.model_validate(dumped).model_dump(exclude_none=True)["output"] == [call] + + def test_normalize_fine_tuning_job_dict_maps_azure_pending(): from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index d71791b6f11..aae5dfe6472 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -168,6 +168,7 @@ OPTION_NAMES: Final = ( "tag_regex", "max_file_size_mb", "silent_model", + "switchyard_router_config", "auto_router_config_path", "auto_router_config", "auto_router_default_model", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index f9304bab0d2..d13bc3d3181 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -109,6 +109,7 @@ export interface AutoRouterDeployment extends AutoRouterCandidateDeployment { model?: string | null; base_model?: string | null; custom_llm_provider?: string | null; + switchyard_router_config?: unknown; complexity_router_config?: unknown; complexity_router_default_model?: string | null; auto_router_config?: unknown; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts index 2cdc13606bb..0fe320a318a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts @@ -87,6 +87,14 @@ const configManaged = (label: string, config: Record): Presenta /** How each strategy renders itself, given its own config object. */ const PRESENTERS: Record) => Presentation> = { + switchyard: (config) => ({ + typeLabel: "Switchyard", + targets: dedupe( + [asRecord(config.models).capable, asRecord(config.models).efficient].filter( + (model): model is string => typeof model === "string", + ), + ), + }), complexity: (config) => ({ typeLabel: complexityTypeLabel(config), targets: dedupe(Object.values(asRecord(config.tiers)).flatMap(normalizeTierModels)), @@ -131,7 +139,8 @@ export const toAutoRouterRow = ( canDelete: canDelete && mayActOnRow, editBlockedReason, createdAt: info.created_at ?? undefined, - defaultModel: (params[strategy.defaultModelKey] as string | null | undefined) ?? null, + defaultModel: + strategy.kind === "switchyard" ? null : (params[strategy.defaultModelKey] as string | null | undefined) ?? null, deployment, ...PRESENTERS[strategy.kind](asRecord(params[strategy.configKey])), }; diff --git a/ui/litellm-dashboard/src/components/add_model/AutoRouterClassifierTabs.tsx b/ui/litellm-dashboard/src/components/add_model/AutoRouterClassifierTabs.tsx index bd8b4eb9854..0dd80854312 100644 --- a/ui/litellm-dashboard/src/components/add_model/AutoRouterClassifierTabs.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AutoRouterClassifierTabs.tsx @@ -135,6 +135,7 @@ interface AutoRouterClassifierTabsProps { children: React.ReactNode; judgeModel?: React.ReactNode; decisionModel?: React.ReactNode; + switchyard?: { selected: boolean; onSelect: () => void; fields: React.ReactNode }; } const AutoRouterClassifierTabs: React.FC = ({ @@ -143,6 +144,7 @@ const AutoRouterClassifierTabs: React.FC = ({ children, judgeModel, decisionModel, + switchyard, }) => { const { t } = useTranslation(); const id = useId(); @@ -159,14 +161,16 @@ const AutoRouterClassifierTabs: React.FC = ({ jev: "jev", custom: "custom", }; - const family = familyByType[classifierType]; + const family = switchyard?.selected ? "switchyard" : familyByType[classifierType]; const hasCustomTiers = Boolean(value.custom_tier_set); const changeType = (next: ClassifierType) => { if (next !== classifierType) onChange(transitionClassifierType(value, next)); }; const changeFamily = (next: unknown) => { + if (next === "switchyard") return switchyard?.onSelect(); if (next === family) return; + if (next === familyByType[classifierType]) return onChange(value); if (next === "heuristics") changeType("heuristic"); if (next === "llm") changeType("llm"); if (next === "jev") changeType("jev"); @@ -207,7 +211,11 @@ const AutoRouterClassifierTabs: React.FC = ({ - + {[ { value: "heuristics", @@ -220,6 +228,9 @@ const AutoRouterClassifierTabs: React.FC = ({ label: t("addModel.autoRouterClassifierTabs.families.decisionModel"), description: t("addModel.autoRouterClassifierTabs.families.decisionModelDescription"), }, + ...(switchyard + ? [{ value: "switchyard", label: "Switchyard", description: "Route by the conversation’s tool history." }] + : []), ].map((option) => ( + {family === "switchyard" && switchyard?.fields} {family === "jev" && (decisionModel ?? (
@@ -355,7 +367,7 @@ const AutoRouterClassifierTabs: React.FC = ({

)} - {hasCustomTiers && ( + {family !== "switchyard" && hasCustomTiers && (

{t("addModel.autoRouterClassifierTabs.restoreStandardTiers")}

)} {availability.data?.error && !availability.isChecking && ( @@ -372,7 +384,7 @@ const AutoRouterClassifierTabs: React.FC = ({

)} - {children} + {family !== "switchyard" && children} ); }; diff --git a/ui/litellm-dashboard/src/components/add_model/BlockedReasonTooltip.tsx b/ui/litellm-dashboard/src/components/add_model/BlockedReasonTooltip.tsx new file mode 100644 index 00000000000..59588cf478f --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/BlockedReasonTooltip.tsx @@ -0,0 +1,17 @@ +import React from "react"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; + +const BlockedReasonTooltip: React.FC<{ reason: string | null; children: React.ReactElement }> = ({ + reason, + children, +}) => + reason === null ? ( + children + ) : ( + + + {reason} + + ); + +export default BlockedReasonTooltip; diff --git a/ui/litellm-dashboard/src/components/add_model/SwitchyardConfigFields.tsx b/ui/litellm-dashboard/src/components/add_model/SwitchyardConfigFields.tsx new file mode 100644 index 00000000000..2043e392928 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/SwitchyardConfigFields.tsx @@ -0,0 +1,342 @@ +"use client"; + +import { useId, useState } from "react"; +import { Input } from "@/components/ui/input"; +import { Textarea } from "@/components/ui/textarea"; +import { Slider } from "@/components/ui/slider"; +import { Switch } from "@/components/ui/switch"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import ModelChoiceCombobox, { type ModelChoice } from "./ModelChoiceCombobox"; +import { + stageConfidence, + toggleClassifier, + setClassifierKind, + parseToolSemantics, + type SwitchyardConfig, +} from "./switchyard_config"; + +interface Props { + value: SwitchyardConfig; + onChange: (value: SwitchyardConfig) => void; + models: ModelChoice[]; + decisionModels: ModelChoice[]; + onToolSemanticsError: (error: string | null) => void; +} + +const isStructured = (instructions: unknown): instructions is Record => + typeof instructions === "object" && instructions !== null; +const instructionsText = (instructions: string | Record | null | undefined): string => + isStructured(instructions) ? JSON.stringify(instructions, null, 2) : instructions ?? ""; + +export default function SwitchyardConfigFields({ + value, + onChange, + models, + decisionModels, + onToolSemanticsError, +}: Props) { + const id = useId(); + const [toolSemantics, setToolSemantics] = useState( + value.options.tool_semantics ? JSON.stringify(value.options.tool_semantics, null, 2) : "", + ); + const classifier = value.classifier; + const llm = classifier?.kind === "llm" ? classifier : undefined; + const decision = classifier?.kind === "decision" ? classifier : undefined; + const setOption = (key: "recent_window", text: string) => + onChange({ ...value, options: { ...value.options, [key]: text === "" ? undefined : Number(text) } }); + const setClassifierNumber = ( + key: "base_threshold" | "threshold_step" | "max_output_tokens" | "recent_turn_window", + text: string, + ) => { + if (llm) onChange({ ...value, classifier: { ...llm, [key]: text === "" ? undefined : Number(text) } }); + }; + const changeToolSemantics = (text: string) => { + setToolSemantics(text); + try { + const parsed = parseToolSemantics(text); + onChange({ ...value, options: { ...value.options, tool_semantics: parsed } }); + onToolSemanticsError(null); + } catch { + onToolSemanticsError("Tool semantics must be JSON mapping observe, mutate, plan, or new to arrays of tool names"); + } + }; + return ( +
+
+
+ + Heuristic + +

+ Stage Router +

+
+
+ + + {value.options.picker === "capable_first" && ( +

+ Switchyard has not benchmarked this bias, so it has no calibrated confidence thresholds. +

+ )} +
+ {(["capable", "efficient"] as const).map((tier) => ( +
+ + onChange({ ...value, models: { ...value.models, [tier]: model ?? "" } })} + choices={models} + placeholder="Select a model" + ariaInvalid={undefined} + ariaDescribedBy={undefined} + /> +
+ ))} +
+
+
+
+ +

+ Ask a judge model when the tool history does not give a confident routing decision. +

+
+ onChange(toggleClassifier(value, enabled))} + /> +
+ {classifier && ( +
+
+ + +
+
+ + onChange({ ...value, models: { ...value.models, judge: judge ?? "" } })} + choices={decision ? decisionModels : models} + placeholder={decision ? "Select a decision model" : "Select a judge model"} + ariaInvalid={undefined} + ariaDescribedBy={undefined} + /> +
+
+ )} +
+
+
+ + Confidence threshold + + {stageConfidence(value).toFixed(2)} +
+ + onChange({ + ...value, + options: { ...value.options, confidence_threshold: Array.isArray(next) ? next[0] : next }, + }) + } + /> +

+ Below this confidence, use {classifier ? "the judge model" : "the bias’s default model"}. Critical errors can + still select the capable model. +

+
+
+ Advanced settings +
+

Stage

+
+ + setOption("recent_window", event.target.value)} + /> +
+
+ +