diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_autorouter_savings_estimate_coverage/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_autorouter_savings_estimate_coverage/migration.sql new file mode 100644 index 00000000000..88e404b189d --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_autorouter_savings_estimate_coverage/migration.sql @@ -0,0 +1,5 @@ +ALTER TABLE "LiteLLM_AutoRouterSession" +ADD COLUMN IF NOT EXISTS "savings_estimated_turns" INTEGER NOT NULL DEFAULT 0, +ADD COLUMN IF NOT EXISTS "savings_estimated_actual_spend" DOUBLE PRECISION NOT NULL DEFAULT 0, +ADD COLUMN IF NOT EXISTS "savings_estimated_saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0, +ADD COLUMN IF NOT EXISTS "savings_estimated_baseline_models" JSONB NOT NULL DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915010000_add_autorouter_baseline_state/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915010000_add_autorouter_baseline_state/migration.sql new file mode 100644 index 00000000000..1720ee03843 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915010000_add_autorouter_baseline_state/migration.sql @@ -0,0 +1,36 @@ +CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterBaselineComparison" ( + "scope" TEXT PRIMARY KEY, + "api_key" TEXT NOT NULL, + "session_id" TEXT NOT NULL, + "router_name" TEXT NOT NULL, + "initial_equivalent" BOOLEAN NOT NULL, + "revision" BIGINT NOT NULL DEFAULT 0, + "published_revision" BIGINT NOT NULL DEFAULT 0, + "history" TEXT, + "attempted_at" TIMESTAMP(3), + "retired" BOOLEAN NOT NULL DEFAULT FALSE, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_scope" + ON "LiteLLM_AutoRouterBaselineComparison" ("api_key", "session_id", "router_name"); +CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_updated" + ON "LiteLLM_AutoRouterBaselineComparison" ("updated_at"); +CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_dirty" + ON "LiteLLM_AutoRouterBaselineComparison" ("attempted_at", "updated_at", "scope") + WHERE NOT "retired" AND "revision" <> "published_revision"; + +CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterBaselineObservation" ( + "request_id" TEXT PRIMARY KEY, + "scope" TEXT NOT NULL, + "started_at" DOUBLE PRECISION NOT NULL, + "revision" BIGINT NOT NULL, + "data" TEXT NOT NULL, + "publication" TEXT, + "conflicted" BOOLEAN NOT NULL DEFAULT FALSE +); + +CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_event_order" + ON "LiteLLM_AutoRouterBaselineObservation" ("scope", "started_at", "request_id"); +CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_event_revision" + ON "LiteLLM_AutoRouterBaselineObservation" ("scope", "revision", "started_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index c4606796ebf..d2032cec0d0 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1551,6 +1551,36 @@ model LiteLLM_AdaptiveRouterSession { @@index([last_activity_at], map: "idx_adaptive_router_session_activity") } +model LiteLLM_AutoRouterBaselineComparison { + scope String @id + api_key String + session_id String + router_name String + initial_equivalent Boolean + revision BigInt @default(0) + published_revision BigInt @default(0) + history String? + attempted_at DateTime? + retired Boolean @default(false) + updated_at DateTime @default(now()) + + @@index([api_key, session_id, router_name], map: "idx_autorouter_baseline_scope") + @@index([updated_at], map: "idx_autorouter_baseline_updated") +} + +model LiteLLM_AutoRouterBaselineObservation { + request_id String @id + scope String + started_at Float + revision BigInt + data String + publication String? + conflicted Boolean @default(false) + + @@index([scope, started_at, request_id], map: "idx_autorouter_baseline_event_order") + @@index([scope, revision, started_at], map: "idx_autorouter_baseline_event_revision") +} + model LiteLLM_AutoRouterSession { api_key String session_id String @@ -1577,6 +1607,10 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + savings_estimated_turns Int @default(0) + savings_estimated_actual_spend Float @default(0) + savings_estimated_saved_spend Float @default(0) + savings_estimated_baseline_models Json @default("{}") classifier_cost Float @default(0) classifier_cost_recorded_turns Int @default(0) tier_turns Json @default("{}") diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ebab2a118fc..2720cf01f2e 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2030,7 +2030,7 @@ dependencies = [ ] [[package]] -name = "litellm-callbacks-legacy" +name = "litellm-callbacks-legacy-python" version = "0.1.0" dependencies = [ "litellm-auth", @@ -2192,7 +2192,7 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-auth-gcp", - "litellm-callbacks-legacy", + "litellm-callbacks-legacy-python", "litellm-core", "litellm-core-utils", "litellm-host-python", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index fa2bdb4224c..a6185632871 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -11,7 +11,7 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-core = { path = "crates/core" } litellm-host = { path = "crates/host" } -litellm-callbacks-legacy = { path = "crates/callbacks-legacy" } +litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" } litellm-framing = { path = "crates/framer" } litellm-auth = { path = "crates/auth" } litellm-auth-aws = { path = "crates/auth-aws" } diff --git a/litellm-rust/crates/callbacks-legacy/AGENTS.md b/litellm-rust/crates/callbacks-legacy-python/AGENTS.md similarity index 100% rename from litellm-rust/crates/callbacks-legacy/AGENTS.md rename to litellm-rust/crates/callbacks-legacy-python/AGENTS.md diff --git a/litellm-rust/crates/callbacks-legacy/Cargo.toml b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml similarity index 90% rename from litellm-rust/crates/callbacks-legacy/Cargo.toml rename to litellm-rust/crates/callbacks-legacy-python/Cargo.toml index 023c13d912b..fe19578e04d 100644 --- a/litellm-rust/crates/callbacks-legacy/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "litellm-callbacks-legacy" +name = "litellm-callbacks-legacy-python" version = "0.1.0" edition.workspace = true license.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json similarity index 100% rename from litellm-rust/crates/callbacks-legacy/python_contract.json rename to litellm-rust/crates/callbacks-legacy-python/python_contract.json diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs similarity index 99% rename from litellm-rust/crates/callbacks-legacy/src/adapter.rs rename to litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 883a35f0df5..e1742190205 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -19,9 +19,9 @@ use serde_json::Value; use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, deferred::{PendingLogging, PendingSuccess}, - finalize, is_internal_call, - legacy_python::Streaming, - prepare, setup, + finalize, is_internal_call, prepare, + python::Streaming, + setup, }; /// What the legacy contract needs to know about the route it is logging. diff --git a/litellm-rust/crates/callbacks-legacy/src/call.rs b/litellm-rust/crates/callbacks-legacy-python/src/call.rs similarity index 100% rename from litellm-rust/crates/callbacks-legacy/src/call.rs rename to litellm-rust/crates/callbacks-legacy-python/src/call.rs diff --git a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs similarity index 99% rename from litellm-rust/crates/callbacks-legacy/src/callbacks.rs rename to litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs index 5f04224e6d7..7caa787dd9d 100644 --- a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs @@ -6,8 +6,8 @@ use litellm_host::event::{RequestContext, WireRequest}; use litellm_host_python::to_py; use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; -use crate::legacy_python::{Logging, Wrapper}; use crate::logger::PythonLogger; +use crate::python::{Logging, Wrapper}; pub trait LegacyCallbacks { /// `Logging.update_from_kwargs`: what the logger is told about the request it is diff --git a/litellm-rust/crates/callbacks-legacy/src/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs similarity index 100% rename from litellm-rust/crates/callbacks-legacy/src/deferred.rs rename to litellm-rust/crates/callbacks-legacy-python/src/deferred.rs diff --git a/litellm-rust/crates/callbacks-legacy/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs similarity index 98% rename from litellm-rust/crates/callbacks-legacy/src/lib.rs rename to litellm-rust/crates/callbacks-legacy-python/src/lib.rs index eaa1a8b714e..44393792d1f 100644 --- a/litellm-rust/crates/callbacks-legacy/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -13,9 +13,9 @@ mod adapter; mod call; mod callbacks; mod deferred; -mod legacy_python; mod logger; mod preparation; +mod python; #[cfg(test)] #[path = "../tests/support.rs"] mod test_support; diff --git a/litellm-rust/crates/callbacks-legacy/src/logger.rs b/litellm-rust/crates/callbacks-legacy-python/src/logger.rs similarity index 94% rename from litellm-rust/crates/callbacks-legacy/src/logger.rs rename to litellm-rust/crates/callbacks-legacy-python/src/logger.rs index 061941f05b9..38f3bf29828 100644 --- a/litellm-rust/crates/callbacks-legacy/src/logger.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/logger.rs @@ -5,7 +5,7 @@ use pyo3::{ types::{PyDict, PyTuple}, }; -use crate::legacy_python::{self, Wrapper}; +use crate::python::{self, Wrapper}; /// The `Logging` instance one call fans out through. pub struct PythonLogger { @@ -90,7 +90,7 @@ impl DeploymentHooks { kwargs: &Py, call_type: &str, ) -> PyResult> { - legacy_python::DeploymentHooks::BeforeDeploymentCall + python::DeploymentHooks::BeforeDeploymentCall .call(py, (kwargs, call_type)) .map(Bound::unbind) } @@ -101,7 +101,7 @@ impl DeploymentHooks { response: &Option>, call_type: &str, ) -> PyResult> { - legacy_python::DeploymentHooks::AfterDeploymentSuccess + python::DeploymentHooks::AfterDeploymentSuccess .call(py, (kwargs, response, call_type)) .map(Bound::unbind) } @@ -112,7 +112,7 @@ impl DeploymentHooks { error: &Py, call_type: &str, ) -> PyResult> { - legacy_python::DeploymentHooks::AfterDeploymentFailure + python::DeploymentHooks::AfterDeploymentFailure .call(py, (kwargs, error, call_type)) .map(Bound::unbind) } diff --git a/litellm-rust/crates/callbacks-legacy/src/preparation.rs b/litellm-rust/crates/callbacks-legacy-python/src/preparation.rs similarity index 99% rename from litellm-rust/crates/callbacks-legacy/src/preparation.rs rename to litellm-rust/crates/callbacks-legacy-python/src/preparation.rs index fa1ff9acd4d..aab654c9893 100644 --- a/litellm-rust/crates/callbacks-legacy/src/preparation.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/preparation.rs @@ -3,7 +3,7 @@ use pyo3::{ types::{PyDict, PyList}, }; -use crate::legacy_python::Wrapper; +use crate::python::Wrapper; struct CredentialEntry<'py>(Bound<'py, PyAny>); diff --git a/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs similarity index 97% rename from litellm-rust/crates/callbacks-legacy/src/legacy_python.rs rename to litellm-rust/crates/callbacks-legacy-python/src/python.rs index 7f5c77c1735..cb609d52878 100644 --- a/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -1,7 +1,7 @@ use pyo3::prelude::*; use strum::{IntoStaticStr, VariantArray}; -const MODULE: &str = "litellm.rust_bridge.legacy_callbacks"; +const MODULE: &str = "litellm.rust_bridge.callbacks_legacy_python"; /// Every litellm Python internal the native call still borrows, grouped by the subsystem it /// belongs to. Rust drives the call; these exist only so behaviour that Python owns today @@ -9,7 +9,7 @@ const MODULE: &str = "litellm.rust_bridge.legacy_callbacks"; /// A group is deleted once Rust owns that subsystem, so this enum only shrinks. Calling a /// user's own callback is not borrowing and does not belong here. /// -/// `litellm/rust_bridge/legacy_callbacks.py` is the only Python module behind it, and +/// `litellm/rust_bridge/callbacks_legacy_python.py` is the only Python module behind it, and /// `python_contract.json` pins each function's parameters on both sides. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum LegacyPython { diff --git a/litellm-rust/crates/callbacks-legacy/tests/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs similarity index 100% rename from litellm-rust/crates/callbacks-legacy/tests/deferred.rs rename to litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs diff --git a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs similarity index 100% rename from litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs rename to litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy-python/tests/payload.rs similarity index 100% rename from litellm-rust/crates/callbacks-legacy/tests/payload.rs rename to litellm-rust/crates/callbacks-legacy-python/tests/payload.rs diff --git a/litellm-rust/crates/callbacks-legacy/tests/support.rs b/litellm-rust/crates/callbacks-legacy-python/tests/support.rs similarity index 95% rename from litellm-rust/crates/callbacks-legacy/tests/support.rs rename to litellm-rust/crates/callbacks-legacy-python/tests/support.rs index d3cc32e301f..d0c02fa6da5 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/tests/support.rs @@ -5,12 +5,12 @@ use pyo3::types::{PyDict, PyTuple}; use crate::{LegacyLogging, LegacySurface, PublicCall}; -/// The parameters of every `legacy_callbacks` function, as the real module declares them. -/// `tests/test_litellm/rust_bridge/test_legacy_callbacks.py` pins this file to the Python +/// The parameters of every `callbacks_legacy_python` function, as the real module declares them. +/// `tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python /// signatures, and [`namespace`] binds every fake call against it. pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); -/// Stand-ins for `legacy_callbacks`, the only Python module the crate calls. Tests +/// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests /// share one interpreter and run concurrently, so each fake is installed idempotently and /// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). /// Every fake is bound against the contract first, so a call the real module would reject @@ -23,10 +23,10 @@ import sys import traceback import types -for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'): +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'): sys.modules.setdefault(name, types.ModuleType(name)) -legacy = sys.modules['litellm.rust_bridge.legacy_callbacks'] +legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python'] CONTRACT = json.loads(python_contract) diff --git a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs similarity index 100% rename from litellm-rust/crates/callbacks-legacy/tests/terminal.rs rename to litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index a19a709e60c..79e78d150e0 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,9 +1,9 @@ - Target invariants, not completion claims; these supersede the crate guidance below where they conflict - Keep this crate the product-specific PyO3 consumer of `litellm-host-python` - Own registration, input projection, the route host and the caller callables it answers operations with (file readers, token providers), public response/error construction and the per-call composition of machine, route host and callback contract - - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, re-aliasing unchanged body keys) lives in `litellm-callbacks-legacy` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy + - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, re-aliasing unchanged body keys) lives in `litellm-callbacks-legacy-python` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy - Value-oriented execution, sync waiting, nested-runtime checks, signal polling and panic containment live in `litellm-host-python`; native async work uses `pyo3-async-runtimes`, Serde output uses `Pythonized` - - Core owns typed native state, the route machine, provider preparation/I/O and normalization; the host driver owns terminal events; the legacy adapter in `litellm-callbacks-legacy` owns `Logging` dispatch policy + - Core owns typed native state, the route machine, provider preparation/I/O and normalization; the host driver owns terminal events; the legacy adapter in `litellm-callbacks-legacy-python` owns `Logging` dispatch policy - Python, Rust SDK and gateway use one lifecycle-bearing core route entrypoint; provider helpers stay private, never bridge-accessible transport drivers - Built-in provider/config/secret/auth/document preparation stays in Rust; caller-authored callbacks and focused Python-file reads run only at core-selected points - Target GIL-enabled CPython explicitly with `#[pymodule(gil_used = true)]`; detach Rust-only work diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 8d31855f2fa..3a4a579efa3 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -18,7 +18,7 @@ panic-test = [] [dependencies] bytes.workspace = true litellm-auth.workspace = true -litellm-callbacks-legacy.workspace = true +litellm-callbacks-legacy-python.workspace = true litellm-core.workspace = true litellm-core-utils.workspace = true litellm-auth-gcp.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index 8c42315ac59..b606293f79f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -1,7 +1,9 @@ mod host; use host::MessagesRouteHost; -use litellm_callbacks_legacy::{LegacySurface, PassThroughStream, PublicCall, run_legacy_call}; +use litellm_callbacks_legacy_python::{ + LegacySurface, PassThroughStream, PublicCall, run_legacy_call, +}; use litellm_core::messages::route::{messages_machine, supports}; use pyo3::{ prelude::*, diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index 9be3171f70b..e518f972bac 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -7,7 +7,7 @@ use std::sync::{Arc, LazyLock}; use host::OcrRouteHost; use litellm_auth_gcp::VertexAuth; -use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call}; +use litellm_callbacks_legacy_python::{LegacySurface, PublicCall, run_legacy_call}; use litellm_core::ocr::route::ocr_machine; use litellm_core_utils::settings::ProcessEnvironment; use litellm_llms::base_llm::ocr::{ diff --git a/litellm/experimental_mcp_client/Readme.md b/litellm/experimental_mcp_client/Readme.md index 0c7b0aa76b9..14decce0256 100644 --- a/litellm/experimental_mcp_client/Readme.md +++ b/litellm/experimental_mcp_client/Readme.md @@ -15,3 +15,9 @@ Upgrade SDK1-dependent libraries before installing them alongside `litellm[mcp]` The shared unit-test workflow runs the MCP integration suite once, with SDK2 in the gateway environment and an isolated SDK1 peer. Keep the SDK1 list/call compatibility test while SDK1 clients are supported; remove it when that support is explicitly retired and the client migration is documented See the official [SDK migration guide](https://py.sdk.modelcontextprotocol.io/migration/) for Python API changes + +## HTTP redirects + +For streamable HTTP POST requests, the MCP SDK follows method-preserving redirects such as HTTP 307/308 within the configured endpoint's origin. Redirects to another path on the same scheme, host and port work. The SDK also permits an HTTP-to-HTTPS upgrade on the same host using the default ports + +Redirects to a different origin are rejected before the destination receives a request or credentials. Configure the final MCP endpoint URL directly if the server redirects to a different host or port. Setting the HTTP client's `follow_redirects` option does not override the SDK's policy diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 009089e0a8e..ba8addbaaa0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -212,6 +212,7 @@ if TYPE_CHECKING: from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config from litellm.litellm_core_utils.llm_cost_calc.utils import BilledTokenRates from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector + from litellm.proxy.hooks.autorouter_baseline_cache import BaselineCacheContext, CapturedBaselineObservation try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -501,6 +502,8 @@ class Logging(LiteLLMLoggingBaseClass): litellm_request_debug: bool = False streamed_anthropic_message_id: str | None = None classifier_input: Mapping[str, JsonValue] | None = None + baseline_cache_context: "BaselineCacheContext | None" = None + baseline_observation: "CapturedBaselineObservation | None" = None def __init__( self, @@ -508,7 +511,7 @@ class Logging(LiteLLMLoggingBaseClass): messages, stream, call_type, - start_time, + start_time: datetime.datetime, litellm_call_id: str, function_id: str, litellm_trace_id: str | None = None, @@ -2181,6 +2184,7 @@ class Logging(LiteLLMLoggingBaseClass): logging_result, start_time, end_time, + build_logging_payload: bool = True, ): """Resolve hidden params, compute response cost, and emit the standard logging payload.""" hidden_params: Final = getattr(logging_result, "_hidden_params", {}) @@ -2205,6 +2209,9 @@ class Logging(LiteLLMLoggingBaseClass): else: self.model_call_details["response_cost"] = self._response_cost_calculator(result=logging_result) + if not build_logging_payload: + return + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( logging_result, start_time, end_time ) @@ -2215,6 +2222,19 @@ class Logging(LiteLLMLoggingBaseClass): if standard_logging_payload is not None: emit_standard_logging_payload(standard_logging_payload) + async def _prepare_baseline_cache_estimate(self, response_obj: object) -> None: + if self.baseline_cache_context is None: + return + from litellm.proxy.hooks.autorouter_baseline_cache import finalize_baseline_cache + + await finalize_baseline_cache(self, response_obj) + + async def invalidate_baseline_cache_estimate(self, reason: str, *, completed: bool = False) -> None: + """Invalidate uncertain attempts; retire the reservation at logical completion.""" + from litellm.proxy.hooks.autorouter_baseline_cache import invalidate_baseline_cache + + await invalidate_baseline_cache(self, reason, completed=completed) + def _build_standard_logging_payload( self, init_response_obj: object, start_time: Any, end_time: Any ) -> StandardLoggingPayload | None: @@ -2266,6 +2286,7 @@ class Logging(LiteLLMLoggingBaseClass): end_time=None, cache_hit=None, standard_logging_object: StandardLoggingPayload | None = None, + build_logging_payload: bool = True, ): try: if start_time is None: @@ -2303,6 +2324,7 @@ class Logging(LiteLLMLoggingBaseClass): logging_result=logging_result, start_time=start_time, end_time=end_time, + build_logging_payload=build_logging_payload, ) elif standard_logging_object is not None: self.model_call_details["standard_logging_object"] = standard_logging_object @@ -3051,8 +3073,17 @@ class Logging(LiteLLMLoggingBaseClass): result=result, cache_hit=cache_hit, standard_logging_object=kwargs.get("standard_logging_object", None), + build_logging_payload=self.baseline_cache_context is None, ) + if self.stream is not True and self.baseline_cache_context is not None: + await self._prepare_baseline_cache_estimate(result) + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( + result, start_time, end_time + ) + if (prepared_payload := self.model_call_details.get("standard_logging_object")) is not None: + emit_standard_logging_payload(prepared_payload) + ## BUILD COMPLETE STREAMED RESPONSE if "async_complete_streaming_response" in self.model_call_details: return # break out of this. @@ -3097,6 +3128,8 @@ class Logging(LiteLLMLoggingBaseClass): self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) + await self._prepare_baseline_cache_estimate(complete_streaming_response) + ## STANDARDIZED LOGGING PAYLOAD try: self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( @@ -3125,6 +3158,7 @@ class Logging(LiteLLMLoggingBaseClass): # Only build standard_logging_object if not already built by # _success_handler_helper_fn if self.model_call_details.get("standard_logging_object") is None: + await self._prepare_baseline_cache_estimate(result) ## STANDARDIZED LOGGING PAYLOAD self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( result, start_time, end_time @@ -3631,6 +3665,8 @@ class Logging(LiteLLMLoggingBaseClass): """ Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions. """ + if self.baseline_cache_context is not None: + await self.invalidate_baseline_cache_estimate("failed_request") await self.special_failure_handlers(exception=exception) if not self.should_run_logging(event_type="async_failure"): # prevent double logging return @@ -6153,6 +6189,8 @@ def _autorouter_savings_for_payload( model_id: str | None, usage_object: Mapping[str, object] | None, cost_breakdown: Mapping[str, object] | None, + baseline_usage: Usage | None = None, + baseline_provenance: Literal["observed_initial", "modeled"] | None = None, ) -> float | None: """The auto-router savings figure for the payload, or ``None`` when there is none. @@ -6171,6 +6209,8 @@ def _autorouter_savings_for_payload( model_id=model_id, usage_object=usage_object, cost_breakdown=cost_breakdown, + baseline_usage=baseline_usage, + baseline_provenance=baseline_provenance, ) except Exception as e: # noqa: BLE001 # a savings figure must never fail request logging verbose_logger.debug("autorouter savings skipped on logging payload: %s", e) @@ -6347,13 +6387,18 @@ def get_standard_logging_object_payload( model_name = response_model_name request_cost_breakdown: Final = cost_breakdown_with_guardrail(logging_obj.cost_breakdown, guardrail_cost) - autorouter_savings: Final = _autorouter_savings_for_payload( - request_metadata=metadata, - model=model_name, - custom_llm_provider=custom_llm_provider, - model_id=_model_id, - usage_object=usage_dict, - cost_breakdown=request_cost_breakdown, + captured_baseline: Final = logging_obj.baseline_observation + autorouter_savings: Final = ( + None + if status != "success" or cache_hit or logging_obj.baseline_cache_context is not None + else _autorouter_savings_for_payload( + request_metadata=metadata, + model=model_name, + custom_llm_provider=custom_llm_provider, + model_id=_model_id, + usage_object=usage_dict, + cost_breakdown=request_cost_breakdown, + ) ) payload: Final[StandardLoggingPayload] = StandardLoggingPayload( @@ -6400,6 +6445,26 @@ def get_standard_logging_object_payload( response_cost=response_cost, cost_breakdown=request_cost_breakdown, autorouter_savings=autorouter_savings, + autorouter_savings_estimate=( + { + "version": 3, + "status": "unknown", + "reason": "pending_projection", + } # mutable-ok: spend-log JSON serialization requires plain mappings + if captured_baseline is not None + else ( + { # mutable-ok: spend-log JSON serialization requires plain mappings + "version": 1, + "status": "estimated" if autorouter_savings is not None else "unknown", + "reason": "uncached_usage" if autorouter_savings is not None else "baseline_unavailable", + } + if metadata.get("routing_decision") + else None + ) + ), + autorouter_baseline_observation=( + captured_baseline.model_dump_json() if captured_baseline is not None else None + ), total_tokens=usage_dict.get("total_tokens", 0), prompt_tokens=usage_dict.get("prompt_tokens", 0), completion_tokens=usage_dict.get("completion_tokens", 0), diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index dc8bbc9edac..359b8bb08c9 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -4,7 +4,8 @@ Calling + translation logic for anthropic's `/v1/messages` endpoint import copy import json -from collections.abc import Callable +from collections.abc import Callable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Union, cast import httpx @@ -31,7 +32,6 @@ from litellm.types.llms.anthropic import ( ContentBlockStop, MessageBlockDelta, MessageStartBlock, - UsageDelta, ) from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, @@ -557,6 +557,7 @@ class ModelResponseIterator: self.tool_index = -1 self.json_mode = json_mode self.speed = speed + self._cumulative_usage: Mapping[str, object] = MappingProxyType({}) # rewritten-name -> caller's original. Built per-request from the # forward map in AnthropicConfig._build_request_tool_name_maps; only # contains entries we actually rewrote, so a tool legitimately named @@ -631,10 +632,12 @@ class ModelResponseIterator: return True return False - def _handle_usage(self, anthropic_usage_chunk: dict | UsageDelta) -> Usage: + def _handle_usage(self, anthropic_usage_chunk: Mapping[str, object]) -> Usage: + # message_delta usage is cumulative but may omit fields reported at message_start. + self._cumulative_usage = MappingProxyType({**self._cumulative_usage, **anthropic_usage_chunk}) reasoning_content: Final = "".join(self.reasoning_content_chunks) if self.reasoning_content_chunks else None usage: Final = AnthropicConfig().calculate_usage( - usage_object=cast(dict, anthropic_usage_chunk), + usage_object=self._cumulative_usage, reasoning_content=reasoning_content, speed=self.speed, ) diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index 38cd429d99a..dd2135f4918 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -4,9 +4,11 @@ Anthropic CountTokens API handler. Uses httpx for HTTP requests instead of the Anthropic SDK. """ -from typing import Any, Final +from collections.abc import Mapping +from typing import Final import httpx +from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import verbose_logger @@ -16,6 +18,8 @@ from litellm.llms.anthropic.count_tokens.transformation import ( ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +_COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue]) + class AnthropicCountTokensHandler(AnthropicCountTokensConfig): """ @@ -27,13 +31,14 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): async def handle_count_tokens_request( self, model: str, - messages: list[dict[str, Any]], + messages: list[dict[str, JsonValue]], api_key: str, api_base: str | None = None, timeout: float | httpx.Timeout | None = None, - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, - ) -> dict[str, Any]: + tools: list[dict[str, JsonValue]] | None = None, + system: JsonValue = None, + optional_params: Mapping[str, JsonValue] | None = None, + ) -> dict[str, JsonValue]: """ Handle a CountTokens request using httpx. @@ -52,7 +57,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): """ try: # Validate the request - self.validate_request(model, messages) + self.validate_request(model, messages, system=system, tools=tools) verbose_logger.debug("Processing Anthropic CountTokens request for model: %s", model) @@ -62,6 +67,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): messages=messages, tools=tools, system=system, + optional_params=optional_params, ) verbose_logger.debug("Transformed request: %s", request_body) @@ -97,7 +103,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): message=error_text, ) - anthropic_response: Final = response.json() + anthropic_response: Final = _COUNT_RESPONSE.validate_json(response.content) verbose_logger.debug("Anthropic response: %s", anthropic_response) diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 12581b9f658..fb12747cec0 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -4,10 +4,17 @@ Anthropic CountTokens API transformation logic. This module handles the transformation of requests to Anthropic's CountTokens API format. """ -from typing import Any, Final +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +from pydantic import JsonValue, TypeAdapter from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION +_COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue]) +COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config") + class AnthropicCountTokensConfig: """ @@ -31,27 +38,31 @@ class AnthropicCountTokensConfig: def transform_request_to_count_tokens( self, model: str, - messages: list[dict[str, Any]], - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, - ) -> dict[str, Any]: + messages: list[dict[str, JsonValue]], + tools: list[dict[str, JsonValue]] | None = None, + system: JsonValue = None, + optional_params: Mapping[str, JsonValue] | None = None, + ) -> dict[str, JsonValue]: # mutable-ok: provider transport requires JSON dictionaries """ Transform request to Anthropic CountTokens format. Includes optional system and tools fields for accurate token counting. """ - request: Final[dict[str, Any]] = { - "model": model, - "messages": messages, - } - - if system is not None: - request["system"] = system - - if tools is not None: - request["tools"] = tools - - return request + options: Final[Mapping[str, JsonValue]] = optional_params or MappingProxyType({}) + return _COUNT_REQUEST.validate_python( + MappingProxyType( + { + "model": model, + "messages": messages, + **MappingProxyType( + {key: value for key, value in (("system", system), ("tools", tools)) if value is not None} + ), + **MappingProxyType( + {key: value for key, value in options.items() if key in COUNT_TOKEN_OPTION_NAMES} + ), + } + ) + ) def get_required_headers(self, api_key: str) -> dict[str, str]: """ @@ -76,7 +87,14 @@ class AnthropicCountTokensConfig: headers, _ = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key) return headers - def validate_request(self, model: str, messages: list[dict[str, Any]]) -> None: + def validate_request( + self, + model: str, + messages: Sequence[Mapping[str, JsonValue]], + *, + system: JsonValue = None, + tools: list[dict[str, JsonValue]] | None = None, + ) -> None: """ Validate the incoming count tokens request. @@ -90,7 +108,7 @@ class AnthropicCountTokensConfig: if not model: raise ValueError("model parameter is required") - if not messages: + if not messages and not system and not tools: raise ValueError("messages parameter is required") if not isinstance(messages, list): diff --git a/litellm/llms/anthropic/prompt_cache_prediction.py b/litellm/llms/anthropic/prompt_cache_prediction.py index e69a02bd93a..447cefb1c45 100644 --- a/litellm/llms/anthropic/prompt_cache_prediction.py +++ b/litellm/llms/anthropic/prompt_cache_prediction.py @@ -1,10 +1,11 @@ from __future__ import annotations +import asyncio import hashlib import json from collections.abc import Mapping, Sequence -from dataclasses import dataclass -from itertools import accumulate +from dataclasses import dataclass, field +from itertools import accumulate, groupby from types import MappingProxyType from typing import Annotated, Final, Literal, Protocol, TypeAlias @@ -14,9 +15,14 @@ from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictInt, TypeAda import litellm from litellm.llms.anthropic.common_utils import AnthropicModelInfo, is_anthropic_oauth_key from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler -from litellm.llms.anthropic.experimental_pass_through.messages.transformation import DEFAULT_ANTHROPIC_API_VERSION +from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + DEFAULT_ANTHROPIC_API_VERSION, + AnthropicMessagesConfig, +) from litellm.types.router import LiteLLM_Params from litellm.types.utils import ModelResponse +from litellm.utils import supports_thinking_cache_preservation _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _HEADERS: Final = TypeAdapter(dict[str, str]) @@ -100,10 +106,7 @@ _Block: TypeAlias = Annotated[_Text | _ToolUse | _ToolResult, Field(discriminato class _Message(_StrictModel): role: Literal["user", "assistant"] - content: str | Annotated[tuple[_Block, ...], Field(strict=False)] - - def blocks(self) -> tuple[_Text | _ToolUse | _ToolResult, ...]: - return (_Text(type="text", text=self.content),) if isinstance(self.content, str) else tuple(self.content) + content: Annotated[str, Field(min_length=1, pattern=r"\S")] | Annotated[tuple[_Block, ...], Field(strict=False)] class _Tool(_StrictModel): @@ -113,10 +116,7 @@ class _Tool(_StrictModel): type: Literal["custom"] | None = None -class _Request(_StrictModel): - messages: tuple[_Message, ...] = Field(min_length=1, strict=False) - system: str | Annotated[tuple[_ResultText, ...], Field(strict=False)] | None = None - tools: Annotated[tuple[_Tool, ...], Field(strict=False)] | None = None +class _RequestOptions(_StrictModel): model: str | None = None max_tokens: int | None = None stream: bool | None = None @@ -127,6 +127,289 @@ class _Request(_StrictModel): metadata: Mapping[str, JsonValue] | None = None +class _Request(_RequestOptions): + messages: tuple[_Message, ...] = Field(min_length=1, strict=False) + system: str | Annotated[tuple[_ResultText, ...], Field(strict=False)] | None = None + tools: Annotated[tuple[_Tool, ...], Field(strict=False)] | None = None + + +class _Thinking(_StrictModel): + type: Literal["thinking"] + thinking: str + signature: str = Field(min_length=1) + + +_PlanBlock: TypeAlias = Annotated[_Text | _ToolUse | _ToolResult | _Thinking, Field(discriminator="type")] + + +class _PlanMessage(_StrictModel): + role: Literal["user", "assistant", "system"] + content: str | Annotated[tuple[_PlanBlock, ...], Field(strict=False)] + + +class _PlanTool(_Tool): + cache_control: _CacheControl | None = None + + +class _PlanRequest(_RequestOptions): + messages: tuple[_PlanMessage, ...] = Field(min_length=1, strict=False) + system: str | Annotated[tuple[_Text, ...], Field(strict=False)] | None = None + tools: Annotated[tuple[_PlanTool, ...], Field(strict=False)] | None = None + cache_control: _CacheControl | None = None + thinking: Mapping[str, JsonValue] | None = None + tool_choice: Mapping[str, JsonValue] | None = None + output_config: Mapping[str, JsonValue] | None = None + speed: Literal["fast", "standard"] | None = None + service_tier: Literal["auto", "standard_only"] | None = None + + +@dataclass(frozen=True, slots=True) +class CacheBoundary: + fingerprint: str + prefix_body: Mapping[str, JsonValue] = field(repr=False) + ttl_seconds: int + lookback_fingerprints: tuple[str, ...] + content_fingerprint: str = "" + lookback_content_fingerprints: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class PromptCachePlan: + full_body: Mapping[str, JsonValue] = field(repr=False) + breakpoints: tuple[CacheBoundary, ...] + + +@dataclass(frozen=True, slots=True) +class UnsupportedCachePlan: + reason: Literal[ + "unsupported_prompt_shape", + "conflicting_cache_ttl", + "too_many_cache_breakpoints", + "invalid_cache_ttl_order", + "unsupported_thinking_cache_semantics", + "token_count_unavailable", + "inconsistent_prefix_token_count", + ] + + +@dataclass(frozen=True, slots=True) +class CountedBreakpoint: + fingerprint: str + ttl_seconds: int + prefix_tokens: int + lookback_fingerprints: tuple[str, ...] + content_fingerprint: str = "" + lookback_content_fingerprints: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class CountedPromptCachePlan: + total_tokens: int + breakpoints: tuple[CountedBreakpoint, ...] + + +@dataclass(frozen=True, slots=True) +class _Position: + section: Literal["tools", "system", "messages"] + message_index: int + role: str + block: Mapping[str, JsonValue] + marker: _CacheControl | None + + +def _content_blocks(content: JsonValue) -> tuple[Mapping[str, JsonValue], ...]: + if isinstance(content, str): + return (MappingProxyType({"type": "text", "text": content}),) + return tuple(_JSON_OBJECT.validate_python(block) for block in content) if isinstance(content, list) else () + + +def _position( + section: Literal["tools", "system", "messages"], + message_index: int, + role: str, + block: Mapping[str, JsonValue], +) -> _Position: + control: Final = block.get("cache_control") + return _Position( + section, + message_index, + role, + MappingProxyType({key: value for key, value in block.items() if key != "cache_control"}), + _CacheControl.model_validate(control) if control is not None else None, + ) + + +def _positions(body: Mapping[str, JsonValue]) -> tuple[_Position, ...]: + tools: Final = body.get("tools") + messages: Final = body.get("messages") + return ( + *tuple( + _position("tools", -1, "", _JSON_OBJECT.validate_python(tool)) + for tool in (tools if isinstance(tools, list) else ()) + ), + *tuple(_position("system", -1, "", block) for block in _content_blocks(body.get("system"))), + *tuple( + _position("messages", message_index, str(message.get("role")), block) + for message_index, raw_message in enumerate(messages if isinstance(messages, list) else ()) + for message in (_JSON_OBJECT.validate_python(raw_message),) + for block in _content_blocks(message.get("content")) + ), + ) + + +def _prefix_body( + body: Mapping[str, JsonValue], + positions: tuple[_Position, ...], + last_index: int, +) -> Mapping[str, JsonValue]: + prefix: Final = positions[: last_index + 1] + sections: Final = MappingProxyType( + { + section: _count_objects(tuple(position.block for position in prefix if position.section == section)) + for section in ("tools", "system") + if any(position.section == section for position in prefix) + } + ) + messages: Final = tuple( + MappingProxyType( + _JSON_OBJECT.validate_python( + MappingProxyType( + {"role": group[0].role, "content": _count_objects(tuple(position.block for position in group))} + ) + ) + ) + for _, values in groupby( + (position for position in prefix if position.section == "messages"), + key=lambda position: position.message_index, + ) + for group in (tuple(values),) + ) + return MappingProxyType( + _JSON_OBJECT.validate_python( + MappingProxyType( + { + **MappingProxyType({key: body[key] for key in COUNT_TOKEN_OPTION_NAMES if key in body}), + **sections, + "messages": _count_objects(messages), + } + ) + ) + ) + + +def _position_group(position: _Position, index: int) -> tuple[str, int, str | int]: + block_type: Final = position.block.get("type") + return ( + position.section, + position.message_index, + block_type if isinstance(block_type, str) and block_type in ("tool_use", "tool_result") else index, + ) + + +def _chain_digest(previous: str, current: str) -> str: + return _digest((previous, current)) + + +def _cacheable_position(position: _Position) -> bool: + block_type: Final = position.block.get("type") + if block_type == "thinking": + return False + text: Final = position.block.get("text") + return block_type != "text" or (isinstance(text, str) and bool(text.strip())) + + +def _entry_fingerprint(fingerprint: str, ttl_seconds: int) -> str: + return _digest(("native-cache-prefix-v2", fingerprint, ttl_seconds)) + + +def parse_cache_plan(body: Mapping[str, JsonValue]) -> PromptCachePlan | UnsupportedCachePlan: + try: + request: Final = _PlanRequest.model_validate(body) + positions: Final = _positions(body) + except ValidationError: + return UnsupportedCachePlan("unsupported_prompt_shape") + explicit: Final = tuple( + (index, position.marker) for index, position in enumerate(positions) if position.marker is not None + ) + automatic_index: Final = next( + (index for index in reversed(range(len(positions))) if _cacheable_position(positions[index])), None + ) + automatic_existing: Final = next((marker for index, marker in explicit if index == automatic_index), None) + if ( + request.cache_control is not None + and automatic_existing is not None + and automatic_existing != request.cache_control + ): + return UnsupportedCachePlan("conflicting_cache_ttl") + automatic: Final = ( + ((automatic_index, request.cache_control),) + if (request.cache_control is not None and automatic_index is not None and automatic_existing is None) + else () + ) + markers: Final = tuple(sorted((*explicit, *automatic), key=lambda value: value[0])) + if len(markers) > 4: + return UnsupportedCachePlan("too_many_cache_breakpoints") + ttls: Final = tuple(3600 if marker.ttl == "1h" else 300 for _, marker in markers) + if any(first < second for first, second in zip(ttls, ttls[1:])): + return UnsupportedCachePlan("invalid_cache_ttl_order") + settings: Final = MappingProxyType( + { + key: body[key] + for key in ("thinking", "output_config", "speed") + if key in body and not (key == "speed" and body[key] == "standard") + } + ) + hashes: Final = tuple( + accumulate( + ( + _digest( + ( + position.section, + position.message_index, + position.role, + position.block, + body.get("tool_choice") if position.section == "messages" else None, + ) + ) + for position in positions + ), + _chain_digest, + initial=_digest(settings), + ) + )[1:] + groups: Final = tuple( + tuple(index for index, _ in values) + for _, values in groupby( + enumerate(positions), + key=lambda item: _position_group(item[1], item[0]), + ) + ) + return PromptCachePlan( + full_body=MappingProxyType(dict(body)), + breakpoints=tuple( + CacheBoundary( + fingerprint=_entry_fingerprint(hashes[index], ttl), + prefix_body=_prefix_body(body, positions, index), + ttl_seconds=ttl, + lookback_fingerprints=tuple( + _entry_fingerprint(hashes[earlier], ttl) + for group in reversed(tuple(group for group in groups if group[0] <= index)[-20:]) + for earlier in reversed(group) + if earlier <= index + ), + content_fingerprint=hashes[index], + lookback_content_fingerprints=tuple( + hashes[earlier] + for group in reversed(tuple(group for group in groups if group[0] <= index)[-20:]) + for earlier in reversed(group) + if earlier <= index + ), + ) + for (index, _), ttl in zip(markers, ttls) + ), + ) + + @dataclass(frozen=True, slots=True) class PromptPrefix: prefix_body: Mapping[str, JsonValue] @@ -137,68 +420,28 @@ class PromptPrefix: def _digest(value: object) -> str: return hashlib.sha256( - json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode() + json.dumps(value, default=_json_object, separators=(",", ":"), ensure_ascii=False).encode() ).hexdigest() -def _next_digest(previous: str, boundary: tuple[int, str, Mapping[str, JsonValue]]) -> str: - return _digest((previous, boundary)) +def _json_object(value: object) -> dict[str, JsonValue]: # mutable-ok: JSON serialization requires a dictionary + return _JSON_OBJECT.validate_python(value) def parse_prompt(body: Mapping[str, JsonValue]) -> PromptPrefix | None: try: - request: Final = _Request.model_validate(body) - blocks: Final = tuple(message.blocks() for message in request.messages) + _Request.model_validate(body) except ValidationError: return None - markers: Final = tuple( - (message_index, block_index, block.cache_control) - for message_index, message_blocks in enumerate(blocks) - for block_index, block in enumerate(message_blocks) - if block.cache_control is not None - ) - if len(markers) != 1: + plan: Final = parse_cache_plan(body) + if isinstance(plan, UnsupportedCachePlan) or len(plan.breakpoints) != 1: return None - message_end, block_end, marker = markers[0] - normalized: Final = _JSON_OBJECT.validate_python(request.model_dump(mode="json", exclude_none=True)) - context: Final = MappingProxyType({key: normalized[key] for key in ("system", "tools") if key in normalized}) - boundaries: Final = tuple( - ( - message_index, - request.messages[message_index].role, - _JSON_OBJECT.validate_python( - block.model_dump(mode="json", exclude=MappingProxyType({"cache_control": True}), exclude_none=True) - ), - ) - for message_index, message_blocks in enumerate(blocks[: message_end + 1]) - for block_index, block in enumerate(message_blocks) - if message_index < message_end or block_index <= block_end - ) - hashes: Final = tuple( - accumulate(boundaries, _next_digest, initial=_digest((_JSON_OBJECT.validate_python(context), marker.ttl))) - )[1:] - prefix_messages: Final = tuple( - _Message( - role=request.messages[message_index].role, - content=tuple( - block - for block_index, block in enumerate(message_blocks) - if message_index < message_end or block_index <= block_end - ), - ) - for message_index, message_blocks in enumerate(blocks[: message_end + 1]) - ) + prefix: Final = plan.breakpoints[0] return PromptPrefix( - prefix_body=MappingProxyType( - _JSON_OBJECT.validate_python( - _Request(messages=prefix_messages, system=request.system, tools=request.tools).model_dump( - mode="json", exclude_none=True - ) - ) - ), - fingerprint=hashes[-1], - fingerprints=tuple(reversed(hashes[-20:])), - ttl_seconds=3600 if marker.ttl == "1h" else 300, + prefix_body=prefix.prefix_body, + fingerprint=prefix.fingerprint, + fingerprints=prefix.lookback_fingerprints, + ttl_seconds=prefix.ttl_seconds, ) @@ -246,6 +489,9 @@ class _CountBody(BaseModel): messages: Sequence[Mapping[str, JsonValue]] tools: Sequence[Mapping[str, JsonValue]] | None = None system: str | Sequence[Mapping[str, JsonValue]] | None = None + thinking: Mapping[str, JsonValue] | None = None + tool_choice: Mapping[str, JsonValue] | None = None + output_config: Mapping[str, JsonValue] | None = None class _CountResult(BaseModel): @@ -262,16 +508,36 @@ def _count_objects( return [dict(value) for value in values] # mutable-ok: serialize read-only inputs at the provider API boundary -async def count_prompt_tokens(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: - native: Final = _CountBody.model_validate(body) +def _messages_url(model: str, api_key: str, api_base: str | None) -> str: + return AnthropicMessagesConfig().get_complete_url( # pyright: ignore[reportUnknownMemberType] # canonical native URL owner takes legacy JSON arguments + api_base=api_base, + api_key=api_key, + model=model, + optional_params=_JSON_OBJECT.validate_python(MappingProxyType({})), + litellm_params=_JSON_OBJECT.validate_python(MappingProxyType({})), + ) + + +async def count_prompt_tokens( + model: str, + api_key: str, + body: Mapping[str, JsonValue], + api_base: str | None = None, +) -> int | None: try: + native: Final = _CountBody.model_validate(body) + count_url: Final = _messages_url(model, api_key, api_base) + "/count_tokens" result: Final = _CountResult.model_validate( await _counter.handle_count_tokens_request( model=model, messages=_count_objects(native.messages), tools=_count_objects(native.tools) if native.tools is not None else None, - system=native.system, + system=_JSON_OBJECT.validate_python(MappingProxyType({"system": native.system}))["system"], api_key=api_key, + api_base=count_url, + optional_params=_JSON_OBJECT.validate_python( + MappingProxyType({key: body[key] for key in COUNT_TOKEN_OPTION_NAMES if key in body}) + ), timeout=15.0, ) ) @@ -280,10 +546,55 @@ async def count_prompt_tokens(model: str, api_key: str, body: Mapping[str, JsonV return result.input_tokens +async def count_cache_plan( + model: str, + api_key: str, + plan: PromptCachePlan, + token_counter: TokenCounter = count_prompt_tokens, +) -> CountedPromptCachePlan | UnsupportedCachePlan: + if any(position.block.get("type") == "thinking" for position in _positions(plan.full_body)): + if not supports_thinking_cache_preservation(model, "anthropic"): + return UnsupportedCachePlan("unsupported_thinking_cache_semantics") + total: Final = await token_counter(model, api_key, plan.full_body) + if total is None: + return UnsupportedCachePlan("token_count_unavailable") + counts: Final = tuple( + await asyncio.gather(*(token_counter(model, api_key, marker.prefix_body) for marker in plan.breakpoints)) + ) + if any(value is None for value in counts): + return UnsupportedCachePlan("token_count_unavailable") + known: Final = tuple(value for value in counts if value is not None) + if any(value < 0 for value in (total, *known)) or any( + first > second for first, second in zip(known, (*known[1:], total)) + ): + return UnsupportedCachePlan("inconsistent_prefix_token_count") + return CountedPromptCachePlan( + total, + tuple( + CountedBreakpoint( + marker.fingerprint, + marker.ttl_seconds, + count, + marker.lookback_fingerprints, + marker.content_fingerprint, + marker.lookback_content_fingerprints, + ) + for marker, count in zip(plan.breakpoints, known) + ), + ) + + @dataclass(frozen=True, slots=True) class NativePredictionTarget: model: str - api_key: str + api_key: str = field(repr=False) + api_base: str | None = None + + +def supported_baseline_recipient(target: NativePredictionTarget, wire: httpx.Request) -> bool: + return wire.headers.get("x-api-key") == target.api_key and wire.url == httpx.URL( + _messages_url(target.model, target.api_key, target.api_base) + ) @dataclass(frozen=True, slots=True) @@ -297,11 +608,26 @@ class UnsupportedPredictionTarget: def resolve_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget | UnsupportedPredictionTarget: + return _resolve_prediction_target(params, allow_configured_endpoint=False) + + +def resolve_baseline_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget | UnsupportedPredictionTarget: + return _resolve_prediction_target(params, allow_configured_endpoint=True) + + +def _resolve_prediction_target( + params: LiteLLM_Params, + *, + allow_configured_endpoint: bool, +) -> NativePredictionTarget | UnsupportedPredictionTarget: configured_options: Final = frozenset(params.model_dump(exclude_defaults=True, exclude_none=True)) if configured_options - _DEPLOYMENT_OPTIONS: return UnsupportedPredictionTarget("unsupported_deployment_configuration") api_base: Final = AnthropicModelInfo.get_api_base(params.api_base) - if api_base not in ("https://api.anthropic.com", "https://api.anthropic.com/v1/messages"): + if not allow_configured_endpoint and api_base not in ( + "https://api.anthropic.com", + "https://api.anthropic.com/v1/messages", + ): return UnsupportedPredictionTarget("unsupported_provider_endpoint") try: model, provider, _, _ = litellm.get_llm_provider( @@ -314,7 +640,7 @@ def resolve_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget api_key: Final = AnthropicModelInfo.get_api_key(params.api_key) if api_key is None or not _supported_provider_key(api_key): return UnsupportedPredictionTarget("unsupported_provider_credentials") - return NativePredictionTarget(model=model, api_key=api_key) + return NativePredictionTarget(model=model, api_key=api_key, api_base=api_base) def _supported_provider_key(api_key: str) -> bool: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8e0cdf547a2..49a332e62bb 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2163,6 +2163,8 @@ class BaseLLMHTTPHandler: e=e, litellm_params=litellm_params_dict ) if should_retry and not hit_max_attempt: + if logging_obj.baseline_cache_context is not None: + await logging_obj.invalidate_baseline_cache_estimate("retried_request") verbose_logger.debug( "Anthropic /v1/messages: invalid thinking signature; " "stripping thinking blocks and retrying (attempt %s/%s).", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7cf858ed9ff..4b0f5e8b49a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14510,6 +14510,7 @@ "supports_native_structured_output": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_sampling_params": false, @@ -14547,6 +14548,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14698,6 +14700,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14727,6 +14730,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14759,6 +14763,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14796,6 +14801,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14831,6 +14837,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14869,6 +14876,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14986,6 +14994,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -15027,6 +15036,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, diff --git a/litellm/models/autorouter_session.py b/litellm/models/autorouter_session.py index c7126236ec3..ddce2b5ef81 100644 --- a/litellm/models/autorouter_session.py +++ b/litellm/models/autorouter_session.py @@ -8,6 +8,8 @@ maintains per (api_key, session_id, router_name). from collections.abc import Mapping from datetime import datetime +from pydantic import Field + from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -22,18 +24,25 @@ class LiteLLM_AutoRouterSession(LiteLLMPydanticObjectBase): turns: int spend: float saved_spend: float + savings_estimated_turns: int = 0 + savings_estimated_actual_spend: float = 0.0 + savings_estimated_saved_spend: float = 0.0 + savings_estimated_baseline_models: Mapping[str, int] = Field(default_factory=dict) classifier_cost: float tier_turns: Mapping[str, int] baseline_models: Mapping[str, int] @property def baseline_model(self) -> str | None: - """The baseline most of this session's turns were priced against, or None when no turn recorded one. + """The baseline most covered turns were priced against, or None when none were estimated. A router reconfigured mid-session leaves turns priced against two baselines; the row keeps both counts, and the label is the one that priced the most money-carrying turns rather than whatever the router is configured with now. """ - if not self.baseline_models: + if not self.savings_estimated_baseline_models: return None - return max(self.baseline_models, key=lambda model: (self.baseline_models[model], model)) + return max( + self.savings_estimated_baseline_models, + key=lambda model: (self.savings_estimated_baseline_models[model], model), + ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2bf7b1ee803..56b3f590210 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -13,6 +13,7 @@ from pydantic import ( ConfigDict, Field, Json, + JsonValue, PositiveInt, field_validator, model_validator, @@ -3941,6 +3942,7 @@ class SpendLogsRouterMetadata(TypedDict): class SpendLogsMetadata(TypedDict): + autorouter_baseline_observation: ReadOnly[str | None] """ Specific metadata k,v pairs logged to spendlogs for easier cost tracking """ @@ -3981,7 +3983,8 @@ class SpendLogsMetadata(TypedDict): original_model_group: ReadOnly[str | None] # Model group requested before any fallbacks cost_breakdown: CostBreakdown | None # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.) compression_savings: CompressionSavingsMetadata | None - autorouter_savings: ReadOnly[float | None] # stamped by the logging payload; None = not auto-routed + autorouter_savings: ReadOnly[float | None] + autorouter_savings_estimate: ReadOnly[Mapping[str, JsonValue] | None] litellm_gateway_injected_cache: ReadOnly[str | None] router_metadata: ReadOnly[SpendLogsRouterMetadata | None] # None = deployment not flagged internal_router_model azure_spillover: ReadOnly[AzureSpillover | None] # None = Azure did not report spillover diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py index 47be3888a58..09dd062c888 100644 --- a/litellm/proxy/client/cli/commands/statusline_script.py +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -32,6 +32,7 @@ import unicodedata import urllib.error import urllib.request from collections.abc import Callable, Mapping +from math import isfinite from pathlib import Path from types import MappingProxyType from typing import IO, Final, NamedTuple, Protocol @@ -43,6 +44,7 @@ FETCH_TIMEOUT_SECONDS: Final = 3 BAR_WIDTH: Final = 24 BAR_FULL: Final = "\u2588" BAR_EMPTY: Final = "\u2591" +SEPARATOR: Final = " \u00b7 " TRANSCRIPT_SCAN_LIMIT_BYTES: Final = 4 * 1024 * 1024 CLAUDE_BASE_URL_ENV_KEYS: Final = ("ANTHROPIC_BASE_URL",) CLAUDE_API_KEY_ENV_KEYS: Final = ("ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY") @@ -63,8 +65,11 @@ class Session(NamedTuple): router_name: str last_model: str spend: float - baseline_spend: float + baseline_spend: float | None baseline_model: str | None + turns: int | None = None + savings_estimated_turns: int | None = None + savings_estimated_actual_spend: float | None = None class Credentials(NamedTuple): @@ -205,17 +210,38 @@ def _session_from_payload(payload: Mapping[str, object]) -> Session | None: router_name: Final = printable(payload.get("router_name")) last_model: Final = printable(payload.get("last_model")) spend: Final = payload.get("spend") - baseline_spend: Final = payload.get("baseline_spend") + baseline_spend: Final = payload.get("savings_estimated_baseline_spend", payload.get("baseline_spend")) + turns: Final = payload.get("turns") + estimated_turns: Final = payload.get("savings_estimated_turns") + estimated_actual: Final = payload.get("savings_estimated_actual_spend") if not router_name or not last_model: return None - if not isinstance(spend, (int, float)) or not isinstance(baseline_spend, (int, float)): + if not isinstance(spend, (int, float)) or isinstance(spend, bool) or not isfinite(spend): + return None + if baseline_spend is not None and ( + not isinstance(baseline_spend, (int, float)) or isinstance(baseline_spend, bool) or not isfinite(baseline_spend) + ): return None return Session( router_name=router_name, last_model=last_model, spend=float(spend), - baseline_spend=float(baseline_spend), + baseline_spend=float(baseline_spend) if baseline_spend is not None else None, baseline_model=printable(payload.get("baseline_model")) or None, + turns=turns if isinstance(turns, int) and not isinstance(turns, bool) and turns >= 0 else None, + savings_estimated_turns=( + estimated_turns + if isinstance(estimated_turns, int) and not isinstance(estimated_turns, bool) and estimated_turns >= 0 + else (0 if estimated_turns is not None else None) + ), + savings_estimated_actual_spend=( + float(estimated_actual) + if isinstance(estimated_actual, (int, float)) + and not isinstance(estimated_actual, bool) + and isfinite(estimated_actual) + and estimated_actual >= 0 + else None + ), ) @@ -314,15 +340,36 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo return f"{code}{text}{RESET}" if use_color else text routed: Final = paint(BOLD, f"Routed to: {model}") - if session is None or session.baseline_model is None or session.baseline_spend <= 0: + if session is None: return routed + if session.savings_estimated_turns == 0 or session.baseline_spend is None: + return f"{routed}{SEPARATOR}Savings unavailable" + if session.baseline_model is None or session.baseline_spend <= 0: + return routed + if session.savings_estimated_turns is not None and ( + session.savings_estimated_actual_spend is None + or session.turns is None + or session.savings_estimated_turns > session.turns + ): + return f"{routed}{SEPARATOR}Savings unavailable" + compared_spend: Final = ( + session.savings_estimated_actual_spend + if session.savings_estimated_turns is not None and session.savings_estimated_actual_spend is not None + else session.spend + ) + coverage: Final = ( + f"{SEPARATOR}{session.savings_estimated_turns} of {session.turns} turns estimated" + if session.savings_estimated_turns is not None + else "" + ) reference: Final = baseline_label(session.baseline_model, config_dir) - pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100 - delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}") - peak: Final = max(session.spend, session.baseline_spend) + pct: Final = round((session.baseline_spend - compared_spend) / session.baseline_spend * 100) + sign: Final = "-" if pct > 0 else "+" if pct < 0 else "" + delta: Final = paint(LITELLM_COLOR, f"{sign}{abs(pct)}% vs {reference}") + peak: Final = max(compared_spend, session.baseline_spend) label_width: Final = max(_display_width(session.router_name), _display_width(reference)) rows: Final = ( - (session.router_name, session.spend, LITELLM_COLOR), + (session.router_name, compared_spend, LITELLM_COLOR), (reference, session.baseline_spend, BASELINE_COLOR), ) lines: Final = ( @@ -331,7 +378,7 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo f"{paint(DIM, f'${amount:.2f}')}" for label, amount, color in rows ) - return "\n".join((f"{routed} {delta}", *lines)) + return "\n".join((f"{routed} {delta}{coverage}", *lines)) def color_enabled(env: Mapping[str, str]) -> bool: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6f769e6971a..9484fd7c723 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3749,6 +3749,14 @@ class ProxyBaseLLMRequestProcessing: "async_streaming_data_generator: error closing response stream: %s", e, ) + logging_obj: Final = request_data.get("litellm_logging_obj") + if ( + not stream_completed + and isinstance(logging_obj, LiteLLMLoggingObj) + and logging_obj.baseline_cache_context is not None + and logging_obj.model_call_details.get("prompt_cache_response_complete") is not True + ): + await logging_obj.invalidate_baseline_cache_estimate("incomplete_response", completed=True) @staticmethod async def async_streaming_data_generator( diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index 3a61da164d0..0d812ee812a 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -26,6 +26,7 @@ from typing import TYPE_CHECKING, Final, NamedTuple from litellm._logging import verbose_proxy_logger from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES +from litellm.proxy.db.create_views import SupportsExecuteRaw if TYPE_CHECKING: from litellm.proxy._types import SpendLogsPayload @@ -75,6 +76,9 @@ SELECT COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens, COALESCE(SUM(spend), 0)::float8 AS spend, COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend, + COALESCE(SUM(savings_estimated_turns), 0)::int AS savings_estimated_turns, + COALESCE(SUM(savings_estimated_actual_spend), 0)::float8 AS savings_estimated_actual_spend, + COALESCE(SUM(savings_estimated_saved_spend), 0)::float8 AS savings_estimated_saved_spend, COALESCE(SUM(classifier_cost), 0)::float8 AS classifier_cost, COALESCE(SUM(classifier_cost_recorded_turns), 0)::int AS classifier_cost_recorded_turns, COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds @@ -104,6 +108,9 @@ class AutoRouterTurnTransaction: cache_touched: bool tier: str | None = None baseline_model: str | None = None + savings_estimated_turns: int = 0 + savings_estimated_actual_spend: float = 0.0 + savings_estimated_saved_spend: float = 0.0 class TurnCacheFacts(NamedTuple): @@ -215,13 +222,18 @@ def build_autorouter_turn_transaction( turn_at: Final = _turn_time_utc(str(payload.get("startTime") or "")) if turn_at is None: return None - from litellm.proxy.spend_tracking.savings import classifier_cost_from_decision + from litellm.proxy.spend_tracking.savings import ( + classifier_cost_from_decision, + recorded_estimated_autorouter_savings, + ) usage_object_raw: Final = metadata.get("usage_object") cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None) tier_raw: Final = routing_decision.get("tier") baseline_raw: Final = routing_decision.get("savings_baseline_model") classifier_cost: Final = classifier_cost_from_decision(routing_decision) + actual_spend: Final = float(payload.get("spend") or 0.0) + (classifier_cost or 0.0) + estimated_savings: Final = recorded_estimated_autorouter_savings(metadata) return AutoRouterTurnTransaction( api_key=api_key, session_id=bounded_session_id(session_id), @@ -232,13 +244,16 @@ def build_autorouter_turn_transaction( model=model, turn_at=turn_at, total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0), - spend=float(payload.get("spend") or 0.0) + (classifier_cost or 0.0), + spend=actual_spend, saved_spend=saved_spend, classifier_cost=classifier_cost or 0.0, covered=cache.covered, cache_hit=cache.read_tokens > 0, cache_ttl_seconds=cache.write_ttl_seconds, cache_touched=cache.touched, + savings_estimated_turns=int(estimated_savings is not None), + savings_estimated_actual_spend=actual_spend if estimated_savings is not None else 0.0, + savings_estimated_saved_spend=estimated_savings if estimated_savings is not None else 0.0, ) @@ -263,6 +278,10 @@ _BASELINE: Final = f"{_p('baseline_model')}::text" _BASELINE_DELTA: Final = ( f"(CASE WHEN {_BASELINE} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_BASELINE}, 1) END)" ) +_ESTIMATED_BASELINE: Final = f"{_p('savings_estimated_turns')}::int = 1 AND {_BASELINE} IS NOT NULL" +_ESTIMATED_BASELINE_DELTA: Final = ( + f"(CASE WHEN {_ESTIMATED_BASELINE} THEN jsonb_build_object({_BASELINE}, 1) ELSE '{{}}'::jsonb END)" +) _IN_ORDER: Final = f"{_TURN_AT}::timestamp >= t.last_turn_at" _SAME: Final = f"{_IN_ORDER} AND t.last_model = {_MODEL}" @@ -281,7 +300,8 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t ( same_model_turns, same_model_hits, first_visit_turns, first_visit_hits, return_turns, return_hits, return_expired_misses, return_within_ttl_misses, ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns, - baseline_models + baseline_models, savings_estimated_turns, savings_estimated_actual_spend, savings_estimated_saved_spend, + savings_estimated_baseline_models ) VALUES ( {_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp, @@ -292,13 +312,18 @@ VALUES ( (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END), (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END), {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8, - {_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA} + {_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA}, + {_p("savings_estimated_turns")}::int, {_p("savings_estimated_actual_spend")}::float8, + {_p("savings_estimated_saved_spend")}::float8, {_ESTIMATED_BASELINE_DELTA} ) ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET turns = t.turns + 1, total_tokens = t.total_tokens + EXCLUDED.total_tokens, spend = t.spend + EXCLUDED.spend, saved_spend = t.saved_spend + EXCLUDED.saved_spend, + savings_estimated_turns = t.savings_estimated_turns + EXCLUDED.savings_estimated_turns, + savings_estimated_actual_spend = t.savings_estimated_actual_spend + EXCLUDED.savings_estimated_actual_spend, + savings_estimated_saved_spend = t.savings_estimated_saved_spend + EXCLUDED.savings_estimated_saved_spend, classifier_cost = t.classifier_cost + EXCLUDED.classifier_cost, classifier_cost_recorded_turns = t.classifier_cost_recorded_turns + 1, covered_turns = t.covered_turns + EXCLUDED.covered_turns, @@ -331,6 +356,10 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET baseline_models = (CASE WHEN {_BASELINE} IS NOT NULL THEN t.baseline_models || jsonb_build_object({_BASELINE}, COALESCE((t.baseline_models ->> {_BASELINE})::int, 0) + 1) ELSE t.baseline_models END), + savings_estimated_baseline_models = (CASE WHEN {_ESTIMATED_BASELINE} + THEN t.savings_estimated_baseline_models || jsonb_build_object( + {_BASELINE}, COALESCE((t.savings_estimated_baseline_models ->> {_BASELINE})::int, 0) + 1) + ELSE t.savings_estimated_baseline_models END), first_turn_at = LEAST(t.first_turn_at, EXCLUDED.first_turn_at), last_turn_at = GREATEST(t.last_turn_at, EXCLUDED.last_turn_at) """ @@ -348,6 +377,10 @@ def _upsert_params(transaction: AutoRouterTurnTransaction) -> tuple[str | float return tuple(_as_sql_param(getattr(transaction, name)) for name in _UPSERT_PARAM_FIELDS) +async def write_autorouter_turn(db: SupportsExecuteRaw, transaction: AutoRouterTurnTransaction) -> None: + await db.execute_raw(UPSERT_AUTOROUTER_SESSION_SQL, *_upsert_params(transaction)) + + async def _upsert_turn_with_retry( prisma_client: PrismaClient, transaction: AutoRouterTurnTransaction, @@ -355,7 +388,7 @@ async def _upsert_turn_with_retry( ) -> None: for attempt in range(n_retry_times + 1): try: - await prisma_client.db.execute_raw(UPSERT_AUTOROUTER_SESSION_SQL, *_upsert_params(transaction)) + await write_autorouter_turn(prisma_client.db, transaction) except DB_RETRY_SAFE_ERROR_TYPES: if attempt >= n_retry_times: raise diff --git a/litellm/proxy/db/baseline_accounting.py b/litellm/proxy/db/baseline_accounting.py new file mode 100644 index 00000000000..8622cb9e481 --- /dev/null +++ b/litellm/proxy/db/baseline_accounting.py @@ -0,0 +1,640 @@ +from __future__ import annotations + +import asyncio +import json +from collections.abc import AsyncIterator, Callable, Sequence +from datetime import datetime, timedelta +from functools import reduce +from itertools import groupby +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, Protocol, cast + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, model_validator +from typing_extensions import Self + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.db.autorouter_session_rollup import ( + AutoRouterTurnTransaction, + write_autorouter_turn, +) +from litellm.proxy.db.create_views import SupportsRawQueries +from litellm.proxy.db.daily_spend_bulk_upsert import ( + DAILY_SPEND_TABLES, + DailySpendEntity, + SpendRow, + build_bulk_upsert, + merge_by_conflict_key, +) +from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper +from litellm.proxy.spend_tracking.baseline_accounting import ( + BaselineEstimate, + BaselineHistory, + BaselineObservation, + advance_baseline_history, +) +from litellm.proxy.spend_tracking.savings import BaselineCosts, BaselineCostSnapshot, price_baseline_comparison + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +class DailyBaselineTarget(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + entity: DailySpendEntity + entity_id: str | None + + +class DailyBaselineAttribution(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + date: str + api_key: str + model: str | None = None + custom_llm_provider: str | None = None + model_group: str | None = None + endpoint: str | None = None + mcp_namespaced_tool_name: str | None = None + targets: tuple[DailyBaselineTarget, ...] = () + + def adjustment(self, target: DailyBaselineTarget, savings_delta: float, request_id: str) -> SpendRow: + table: Final = DAILY_SPEND_TABLES[target.entity] + return MappingProxyType( + { + "date": self.date, + "api_key": self.api_key, + "model": self.model, + "custom_llm_provider": self.custom_llm_provider, + "model_group": self.model_group, + "endpoint": self.endpoint, + "mcp_namespaced_tool_name": self.mcp_namespaced_tool_name, + table.entity_id_column: target.entity_id, + "request_id": request_id, + "autorouter_savings_spend": savings_delta, + } + ) + + +class BaselineAccountingRecord(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + scope: str = Field(pattern=r"^autorouter-baseline:v3:[a-f0-9]{64}$") + api_key: str = Field(min_length=1) + session_id: str = Field(min_length=1, max_length=256) + router_name: str = Field(min_length=1) + baseline_model: str = Field(min_length=1) + observation: BaselineObservation + pricing: BaselineCostSnapshot + turn: AutoRouterTurnTransaction | None + daily: DailyBaselineAttribution | None + + @model_validator(mode="after") + def consistent_turn(self) -> Self: + turn: Final = self.turn + if turn is not None and ( + (turn.api_key, turn.session_id, turn.router_name, turn.baseline_model) + != (self.api_key, self.session_id, self.router_name, self.baseline_model) + or turn.spend != self.pricing.actual_spend + self.pricing.classifier_cost + or any( + ( + turn.saved_spend, + turn.savings_estimated_turns, + turn.savings_estimated_actual_spend, + turn.savings_estimated_saved_spend, + ) + ) + ): + raise ValueError("Baseline observation must own an unestimated turn with matching scope and actual cost") + return self + + +class BaselinePublication(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + version: Literal[3] = 3 + comparison_id: str + comparison_started_at: float + status: Literal["estimated", "unknown"] + reason: str + provenance: Literal["observed_identical", "modeled"] | None = None + actual_spend: float | None = None + baseline_spend: float | None = None + input_tokens: int | None = None + cache_read_input_tokens: int | None = None + cache_creation_5m_input_tokens: int | None = None + cache_creation_1h_input_tokens: int | None = None + + @property + def costs(self) -> BaselineCosts | None: + if self.status != "estimated" or self.actual_spend is None or self.baseline_spend is None: + return None + return BaselineCosts(self.actual_spend, self.baseline_spend) + + +def baseline_publication( + record: BaselineAccountingRecord, estimate: BaselineEstimate, first_at: float +) -> BaselinePublication: + costs: Final = price_baseline_comparison(record.pricing, estimate.usage, estimate.provenance) + details: Final = estimate.usage.prompt_tokens_details if estimate.usage is not None else None + writes: Final = details.cache_creation_token_details if details is not None else None + return BaselinePublication( + comparison_id=record.scope, + comparison_started_at=first_at, + status="estimated" if costs is not None else "unknown", + reason=estimate.reason if costs is not None or estimate.usage is None else "pricing_unavailable", + provenance=estimate.provenance if costs is not None else None, + actual_spend=costs.actual if costs is not None else None, + baseline_spend=costs.baseline if costs is not None else None, + input_tokens=details.text_tokens if details is not None else None, + cache_read_input_tokens=details.cached_tokens if details is not None else None, + cache_creation_5m_input_tokens=writes.ephemeral_5m_input_tokens if writes is not None else None, + cache_creation_1h_input_tokens=writes.ephemeral_1h_input_tokens if writes is not None else None, + ) + + +class _Comparison(BaseModel): + revision: int + published_revision: int + initial_equivalent: bool + retired: bool + history: str | None + + +class _StoredRecord(BaseModel): + data: str + publication: str | None + conflicted: bool + started_at: float + + +class _Change(BaseModel): + request_id: str + publication: BaselinePublication + api_key: str + session_id: str + router_name: str + baseline_model: str + covered_delta: int + actual_delta: float + savings_delta: float + daily: DailyBaselineAttribution | None + + +class _TransactionManager(Protocol): + async def __aenter__(self) -> SupportsRawQueries: ... + + async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... + + +class _TransactionalDatabase(Protocol): + def tx(self, *, timeout: timedelta) -> _TransactionManager: ... + + +_COMPARISONS: Final = TypeAdapter(tuple[_Comparison, ...]) +_RECORDS: Final = TypeAdapter(tuple[_StoredRecord, ...]) +_HISTORY: Final = TypeAdapter(BaselineHistory) +_PAGE_TIMESTAMPS: Final = 128 +_TRANSACTION_TIMEOUT: Final = timedelta(seconds=10) + +_CREATE_COMPARISON: Final = """ +INSERT INTO "LiteLLM_AutoRouterBaselineComparison" + (scope, api_key, session_id, router_name, initial_equivalent) +VALUES ($1, $2, $3, $4, NOT EXISTS ( + SELECT 1 FROM "LiteLLM_AutoRouterSession" + WHERE api_key = $2 AND session_id = $3 AND router_name = $4 +)) ON CONFLICT (scope) DO NOTHING +""" +_LOCK_COMPARISON: Final = """ +SELECT revision, published_revision, initial_equivalent, retired, history +FROM "LiteLLM_AutoRouterBaselineComparison" WHERE scope = $1 FOR UPDATE +""" +_INSERT_RECORD: Final = """ +INSERT INTO "LiteLLM_AutoRouterBaselineObservation" + (request_id, scope, started_at, revision, data) +VALUES ($1, $2, $3::float8, $4::bigint, $5) +ON CONFLICT (request_id) DO NOTHING +""" +_MARK_CONFLICT: Final = """ +UPDATE "LiteLLM_AutoRouterBaselineObservation" +SET conflicted = TRUE, revision = $4::bigint +WHERE request_id = $1 AND scope = $2 AND data <> $3 AND NOT conflicted +""" +_READ_PAGE: Final = """ +WITH times AS ( + SELECT DISTINCT started_at FROM "LiteLLM_AutoRouterBaselineObservation" + WHERE scope = $1 AND revision > $2::bigint + AND ($3::float8 IS NULL OR started_at > $3::float8) + AND ($5::float8 IS NULL OR ( + started_at >= $5::float8 AND publication::jsonb->>'status' = 'estimated' + )) + ORDER BY started_at LIMIT $4::int +) +SELECT data, publication, conflicted, started_at +FROM "LiteLLM_AutoRouterBaselineObservation" +WHERE scope = $1 AND revision > $2::bigint + AND started_at IN (SELECT started_at FROM times) + AND ($5::float8 IS NULL OR publication::jsonb->>'status' = 'estimated') +ORDER BY started_at, request_id +""" +_UPDATE_LOGS: Final = """ +WITH changes AS ( + SELECT request_id, publication::jsonb AS publication + FROM jsonb_to_recordset($1::jsonb) AS x(request_id text, publication jsonb) +) +UPDATE "LiteLLM_SpendLogs" AS logs +SET metadata = (COALESCE(logs.metadata::jsonb, '{}'::jsonb) - 'autorouter_baseline_observation') || jsonb_build_object( + 'autorouter_savings_estimate', changes.publication, + 'autorouter_savings', CASE WHEN changes.publication->>'status' = 'estimated' THEN + (changes.publication->>'baseline_spend')::float8 - (changes.publication->>'actual_spend')::float8 + ELSE NULL END +) +FROM changes WHERE logs.request_id = changes.request_id +""" +_UPDATE_PUBLICATIONS: Final = """ +UPDATE "LiteLLM_AutoRouterBaselineObservation" AS observations +SET publication = x.publication::text +FROM jsonb_to_recordset($1::jsonb) AS x(request_id text, publication jsonb) +WHERE observations.request_id = x.request_id +""" +_UPDATE_SESSIONS: Final = """ +WITH changes AS ( + SELECT * FROM jsonb_to_recordset($1::jsonb) AS x( + api_key text, session_id text, router_name text, baseline_model text, + covered_delta int, actual_delta float8, savings_delta float8 + ) +), totals AS ( + SELECT api_key, session_id, router_name, SUM(covered_delta)::int AS covered_delta, + SUM(actual_delta) AS actual_delta, SUM(savings_delta) AS savings_delta + FROM changes GROUP BY api_key, session_id, router_name +), models AS ( + SELECT api_key, session_id, router_name, jsonb_object_agg(baseline_model, delta) AS deltas + FROM ( + SELECT api_key, session_id, router_name, baseline_model, SUM(covered_delta)::int AS delta + FROM changes GROUP BY api_key, session_id, router_name, baseline_model + ) grouped GROUP BY api_key, session_id, router_name +) +UPDATE "LiteLLM_AutoRouterSession" AS session +SET saved_spend = session.saved_spend + totals.savings_delta, + savings_estimated_turns = session.savings_estimated_turns + totals.covered_delta, + savings_estimated_actual_spend = session.savings_estimated_actual_spend + totals.actual_delta, + savings_estimated_saved_spend = session.savings_estimated_saved_spend + totals.savings_delta, + savings_estimated_baseline_models = ( + SELECT COALESCE(jsonb_object_agg(key, value), '{}'::jsonb) FROM ( + SELECT key, SUM(value::int)::int AS value FROM ( + SELECT * FROM jsonb_each_text(session.savings_estimated_baseline_models) + UNION ALL SELECT * FROM jsonb_each_text(models.deltas) + ) combined GROUP BY key HAVING SUM(value::int) > 0 + ) counts + ) +FROM totals JOIN models USING (api_key, session_id, router_name) +WHERE session.api_key = totals.api_key AND session.session_id = totals.session_id + AND session.router_name = totals.router_name +""" + + +def _primary_transaction(client: PrismaClient) -> _TransactionManager: + primary: Final = cast(_TransactionalDatabase, writer_wrapper(client.db)) + return primary.tx(timeout=_TRANSACTION_TIMEOUT) + + +def _serialized(model: BaseModel) -> str: + return json.dumps(model.model_dump(mode="json"), sort_keys=True, separators=(",", ":")) + + +def _change(record: BaselineAccountingRecord, old: BaselinePublication | None, new: BaselinePublication) -> _Change: + previous: Final = old.costs if old is not None else None + current: Final = new.costs + return _Change( + request_id=record.observation.request_id, + publication=new, + api_key=record.api_key, + session_id=record.session_id, + router_name=record.router_name, + baseline_model=record.baseline_model, + covered_delta=int(current is not None) - int(previous is not None), + actual_delta=(current.actual if current is not None else 0.0) + - (previous.actual if previous is not None else 0.0), + savings_delta=(current.savings if current is not None else 0.0) + - (previous.savings if previous is not None else 0.0), + daily=record.daily, + ) + + +def _project_group( + previous: tuple[BaselineHistory, tuple[_Change, ...]], stored: Sequence[_StoredRecord] +) -> tuple[BaselineHistory, tuple[_Change, ...]]: + history, prior_changes = previous + records: Final = tuple(BaselineAccountingRecord.model_validate_json(item.data) for item in stored) + observations: Final = tuple( + record.observation.model_copy( + update=MappingProxyType( + {"outcome": "uncertain", "baseline_equivalent": False, "reason": "conflicting_observation"} + ) + ) + if row.conflicted + else record.observation + for record, row in zip(records, stored) + ) + advanced, estimates = advance_baseline_history(history, observations) + publications: Final = tuple( + baseline_publication( + record, estimate, advanced.first_at if advanced.first_at is not None else observations[0].started_at + ) + for record, estimate in zip(records, estimates) + ) + changes: Final = tuple( + _change(record, old, publication) + for record, row, publication in zip(records, stored, publications) + for old in (BaselinePublication.model_validate_json(row.publication) if row.publication else None,) + if publication != old + ) + return advanced, (*prior_changes, *changes) + + +async def _publish(db: SupportsRawQueries, changes: Sequence[_Change]) -> None: + if not changes: + return + serialized: Final = json.dumps(tuple(change.model_dump(mode="json") for change in changes), separators=(",", ":")) + await db.execute_raw(_UPDATE_LOGS, serialized) + await db.execute_raw(_UPDATE_SESSIONS, serialized) + for entity, table in DAILY_SPEND_TABLES.items(): + if adjustments := tuple( + change.daily.adjustment(target, change.savings_delta, change.request_id) + for change in changes + if change.daily is not None and change.savings_delta != 0 + for target in change.daily.targets + if target.entity == entity + ): + statement, values = build_bulk_upsert(table, merge_by_conflict_key(table, adjustments)) + await db.execute_raw(statement, *values) + await db.execute_raw(_UPDATE_PUBLICATIONS, serialized) + + +class BaselineAccountingStore: + def __init__(self, transaction: Callable[[], _TransactionManager]) -> None: + self.transaction: Final = transaction + + @classmethod + def for_client(cls, client: PrismaClient) -> BaselineAccountingStore: + def transaction() -> _TransactionManager: + return _primary_transaction(client) + + return cls(transaction) + + async def append( + self, record: BaselineAccountingRecord + ) -> Literal["recorded", "retired", "conflict", "unavailable"]: + try: + async with self.transaction() as db: + await db.execute_raw("SET LOCAL statement_timeout = 5000") + await db.execute_raw("SET LOCAL lock_timeout = 1000") + await db.execute_raw( + _CREATE_COMPARISON, record.scope, record.api_key, record.session_id, record.router_name + ) + rows: Final = _COMPARISONS.validate_python(tuple(await db.query_raw(_LOCK_COMPARISON, record.scope))) + if not rows: + return "unavailable" + revision: Final = rows[0].revision + 1 + data: Final = _serialized(record) + inserted: Final = await db.execute_raw( + _INSERT_RECORD, + record.observation.request_id, + record.scope, + record.observation.started_at, + revision, + data, + ) + if inserted and record.turn is not None: + await write_autorouter_turn(db, record.turn) + conflicted: Final = ( + 0 + if inserted + else await db.execute_raw( + _MARK_CONFLICT, record.observation.request_id, record.scope, data, revision + ) + ) + canonical: Final = ( + _RECORDS.validate_python( + tuple( + await db.query_raw( + 'SELECT data, publication, conflicted, started_at FROM "LiteLLM_AutoRouterBaselineObservation" ' + "WHERE request_id=$1 AND scope=$2", + record.observation.request_id, + record.scope, + ) + ) + ) + if not inserted + else () + ) + if not inserted and not canonical: + return "conflict" + if rows[0].retired: + await _publish( + db, + ( + _change( + BaselineAccountingRecord.model_validate_json(canonical[0].data) + if canonical + else record, + BaselinePublication.model_validate_json(canonical[0].publication) + if canonical and canonical[0].publication is not None + else None, + BaselinePublication( + comparison_id=record.scope, + comparison_started_at=canonical[0].started_at + if canonical + else record.observation.started_at, + status="unknown", + reason="comparison_retired", + ), + ), + ), + ) + return "retired" + if inserted or conflicted: + await self._withdraw( + db, record.scope, canonical[0].started_at if canonical else record.observation.started_at + ) + await db.execute_raw( + 'UPDATE "LiteLLM_AutoRouterBaselineComparison" SET revision = $2::bigint, ' + "updated_at = CURRENT_TIMESTAMP, attempted_at = NULL WHERE scope = $1", + record.scope, + revision, + ) + return "recorded" + except Exception: # noqa: BLE001 # accounting failure must not change inference or actual billing + verbose_proxy_logger.warning("Auto-router baseline observation could not be persisted") + return "unavailable" + + async def _pages( + self, db: SupportsRawQueries, scope: str, after_revision: int, withdraw_from: float | None = None + ) -> AsyncIterator[tuple[_StoredRecord, ...]]: + cursor: float | None = None + while page := _RECORDS.validate_python( + tuple(await db.query_raw(_READ_PAGE, scope, after_revision, cursor, _PAGE_TIMESTAMPS, withdraw_from)) + ): + yield page + cursor = page[-1].started_at # rebind-ok: keyset pagination advances after each complete timestamp group + + async def _withdraw(self, db: SupportsRawQueries, scope: str, started_at: float) -> None: + async for page in self._pages(db, scope, 0, withdraw_from=started_at): + await _publish( + db, + tuple( + _change( + BaselineAccountingRecord.model_validate_json(row.data), + previous, + BaselinePublication( + comparison_id=scope, + comparison_started_at=min(previous.comparison_started_at, started_at), + status="unknown", + reason="pending_projection", + ), + ) + for row in page + if row.publication is not None + for previous in (BaselinePublication.model_validate_json(row.publication),) + ), + ) + + async def retire_before(self, cutoff: datetime, batch_size: int, timeout_ms: int) -> None: + async with self.transaction() as db: + await db.execute_raw(f"SET LOCAL statement_timeout = {max(1, timeout_ms)}") + await db.execute_raw(f"SET LOCAL lock_timeout = {max(1, timeout_ms)}") + await db.execute_raw( + 'WITH expired AS (SELECT scope FROM "LiteLLM_AutoRouterBaselineComparison" ' + "WHERE NOT retired AND updated_at < $1::timestamptz ORDER BY updated_at " + "LIMIT $2::int FOR UPDATE SKIP LOCKED) " + 'UPDATE "LiteLLM_AutoRouterBaselineComparison" AS comparison ' + "SET retired=TRUE, history=NULL FROM expired WHERE comparison.scope=expired.scope", + cutoff, + batch_size, + ) + await db.execute_raw( + 'DELETE FROM "LiteLLM_AutoRouterBaselineObservation" WHERE request_id IN (' + 'SELECT event.request_id FROM "LiteLLM_AutoRouterBaselineObservation" AS event ' + 'JOIN "LiteLLM_AutoRouterBaselineComparison" AS comparison USING (scope) ' + "WHERE comparison.retired AND comparison.updated_at < $1::timestamptz " + "LIMIT $2::int)", + cutoff, + batch_size, + ) + + async def project(self, scope: str) -> Literal["published", "unchanged", "unavailable"]: + try: + async with self.transaction() as db: + await db.execute_raw("SET LOCAL statement_timeout = 5000") + await db.execute_raw("SET LOCAL lock_timeout = 1000") + rows: Final = _COMPARISONS.validate_python(tuple(await db.query_raw(_LOCK_COMPARISON, scope))) + if not rows or rows[0].retired or rows[0].revision == rows[0].published_revision: + return "unchanged" + missing_log: Final = await db.query_raw( + 'SELECT 1 FROM "LiteLLM_AutoRouterBaselineObservation" AS observation ' + 'WHERE scope=$1 AND publication IS NULL AND NOT EXISTS (SELECT 1 FROM "LiteLLM_SpendLogs" AS log ' + "WHERE log.request_id=observation.request_id) LIMIT 1", + scope, + ) + if missing_log: + return "unavailable" + state: Final = rows[0] + checkpoint: Final = ( + _HISTORY.validate_json(state.history) + if state.history is not None + else BaselineHistory(equivalent=state.initial_equivalent) + ) + changed: Final = await db.query_raw( + 'SELECT 1 FROM "LiteLLM_AutoRouterBaselineObservation" ' + "WHERE scope = $1 AND revision > $2::bigint AND started_at <= $3::float8 LIMIT 1", + scope, + state.published_revision, + checkpoint.last_at, + ) + history = BaselineHistory(equivalent=state.initial_equivalent) if changed else checkpoint + async for page in self._pages(db, scope, 0 if changed else state.published_revision): + history, updates = reduce( + _project_group, + (tuple(group) for _, group in groupby(page, key=lambda item: item.started_at)), + (history, ()), + ) + await _publish(db, updates) + await db.execute_raw( + 'UPDATE "LiteLLM_AutoRouterBaselineComparison" ' + "SET published_revision = revision, history = $2 WHERE scope = $1", + scope, + _HISTORY.dump_json(history).decode(), + ) + return "published" + except Exception: # noqa: BLE001 # rollback leaves the durable revision dirty for a later flush + verbose_proxy_logger.warning("Auto-router baseline projection remains pending") + return "unavailable" + + +class _Scope(BaseModel): + scope: str + + +_SCOPES: Final = TypeAdapter(tuple[_Scope, ...]) +_CLAIM_DIRTY: Final = """ +WITH candidates AS ( + SELECT scope FROM "LiteLLM_AutoRouterBaselineComparison" + WHERE NOT retired AND revision <> published_revision + AND (attempted_at IS NULL OR attempted_at < CURRENT_TIMESTAMP - INTERVAL '30 seconds') + ORDER BY attempted_at NULLS FIRST, updated_at, scope LIMIT 32 FOR UPDATE SKIP LOCKED +) +UPDATE "LiteLLM_AutoRouterBaselineComparison" AS comparison +SET attempted_at = CURRENT_TIMESTAMP FROM candidates +WHERE comparison.scope = candidates.scope RETURNING comparison.scope +""" + + +async def _flush_records( + store: BaselineAccountingStore, records: Sequence[BaselineAccountingRecord] +) -> tuple[BaselineAccountingRecord, ...]: + slots: Final = asyncio.Semaphore(4) + + async def append(record: BaselineAccountingRecord) -> bool: + async with slots: + return await store.append(record) == "unavailable" + + failed: Final = await asyncio.gather(*(append(record) for record in records)) + return tuple(record for record, retry in zip(records, failed) if retry) + + +async def flush_baseline_accounting(client: PrismaClient) -> None: + from litellm.proxy.utils import request_spend_log_flush + + store: Final = BaselineAccountingStore.for_client(client) + async with client.baseline_accounting_lock: + batch: Final = tuple(client.baseline_accounting_transactions[:32]) + client.baseline_accounting_transactions = client.baseline_accounting_transactions[ + 32: + ] # rebind-ok: drain under lock + more_queued: Final = bool(client.baseline_accounting_transactions) + try: + remaining: Final = await asyncio.wait_for(_flush_records(store, batch), timeout=5) + except (Exception, asyncio.CancelledError) as error: # noqa: BLE001 # unknown acknowledgements can be replayed safely + async with client.baseline_accounting_lock: + client.baseline_accounting_transactions.extend(batch) + if isinstance(error, asyncio.CancelledError): + raise + return + async with client.baseline_accounting_lock: + client.baseline_accounting_transactions.extend(remaining) + if more_queued and len(remaining) < len(batch): + request_spend_log_flush(client) + try: + async with store.transaction() as db: + await db.execute_raw("SET LOCAL statement_timeout = 1000") + scopes: Final = _SCOPES.validate_python(tuple(await db.query_raw(_CLAIM_DIRTY))) + slots: Final = asyncio.Semaphore(4) + + async def project(item: _Scope) -> str: + async with slots: + return await store.project(item.scope) + + outcomes: Final = await asyncio.wait_for(asyncio.gather(*(project(item) for item in scopes)), timeout=5) + if len(scopes) == 32 and "published" in outcomes: + request_spend_log_flush(client) + except Exception: # noqa: BLE001 # durable dirty comparisons remain eligible after the retry interval + verbose_proxy_logger.warning("Auto-router baseline projection will retry on a later spend flush") diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index 108b0e884ba..eb130a5196f 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -14,6 +14,8 @@ from itertools import groupby from types import MappingProxyType from typing import Final, Literal +from pydantic import TypeAdapter + DailySpendEntity = Literal["user", "team", "org", "tag", "end_user", "agent"] SqlValue = str | int | float | None @@ -43,6 +45,36 @@ DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingP } ) +_ENTITY_INPUT_KEYS: Final[Mapping[DailySpendEntity, str]] = MappingProxyType( + { + "user": "user", + "team": "team_id", + "org": "organization_id", + "end_user": "end_user", + "agent": "agent_id", + "tag": "request_tags", + } +) +_TAGS: Final = TypeAdapter(tuple[str, ...]) + + +def daily_spend_entity_ids(payload: Mapping[str, object], entity: DailySpendEntity) -> tuple[str | None, ...]: + key: Final = _ENTITY_INPUT_KEYS[entity] + if key not in payload: + return () + value: Final = payload[key] + if entity == "tag": + if value is None: + return () + tags: Final = _TAGS.validate_json(value) if isinstance(value, str) else _TAGS.validate_python(value) + return tuple(dict.fromkeys(tags)) + if value is None: + return (None,) if entity == "user" else () + if not isinstance(value, str) or (entity == "end_user" and not value): + return () + return (value,) + + # The unique constraint's columns after the entity id, in constraint order. A NULL can # never match itself in a unique index, so every one of these is normalized to '': the # conflict target has to be NULL-free or the row is re-inserted on every single flush. diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 165486a4669..ba92c1e4f65 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -18,6 +18,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, overload from urllib.parse import quote, unquote +from pydantic import TypeAdapter from typing_extensions import LiteralString, ReadOnly, TypedDict import litellm @@ -51,6 +52,7 @@ from litellm.proxy.common_utils.user_api_key_cache import project_cache_key from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, build_bulk_upsert, + daily_spend_entity_ids, merge_by_conflict_key, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( @@ -82,6 +84,8 @@ from litellm.repositories.prisma_protocols import BatchTable from litellm.types.utils import CallTypes if TYPE_CHECKING: + from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction + from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution from litellm.proxy.utils import PrismaClient, ProxyLogging else: PrismaClient = Any @@ -89,6 +93,7 @@ else: RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) +_SPEND_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) def _org_member_transaction_key(org_id: str, user_id: str) -> str: @@ -579,25 +584,31 @@ class DBSpendUpdateWriter: metadata_raw: Final = payload.get("metadata") if not metadata_raw: return - metadata: Final = json.loads(metadata_raw) - if not isinstance(metadata, dict) or not metadata.get("routing_decision"): + metadata: Final = _SPEND_METADATA_ADAPTER.validate_json(metadata_raw) + routing_decision: Final = metadata.get("routing_decision") + if not isinstance(routing_decision, Mapping) or not routing_decision: return from litellm.proxy.db.autorouter_session_rollup import ( build_autorouter_turn_transaction, ) usage_object_raw: Final = metadata.get("usage_object") + cost_breakdown: Final = metadata.get("cost_breakdown") + savings_estimate: Final = metadata.get("autorouter_savings_estimate") savings_spend: Final = compute_savings_spend( model=payload.get("model"), custom_llm_provider=payload.get("custom_llm_provider"), compression_saved_tokens=0, gateway_injected_cache=marks_gateway_injection(metadata, payload.get("model_id")), - routing_decision=metadata.get("routing_decision"), + routing_decision=routing_decision, usage_object=usage_object_raw if isinstance(usage_object_raw, dict) else None, model_id=payload.get("model_id"), llm_router=get_llm_router, - cost_breakdown=metadata.get("cost_breakdown"), + cost_breakdown=cost_breakdown if isinstance(cost_breakdown, Mapping) else None, recorded_autorouter_savings=metadata.get("autorouter_savings"), + recorded_autorouter_savings_estimate=( + savings_estimate if isinstance(savings_estimate, Mapping) else None + ), billed_at=payload.get("endTime"), ) transaction: Final = build_autorouter_turn_transaction( @@ -605,6 +616,11 @@ class DBSpendUpdateWriter: metadata=metadata, saved_spend=savings_spend.autorouter, ) + try: + if await self._enqueue_baseline_accounting(payload, metadata, transaction, prisma_client): + return + except Exception: # noqa: BLE001 # optional baseline capture must preserve the original actual-spend rollup + verbose_proxy_logger.warning("Auto-router baseline observation was unavailable; actual turn retained") if transaction is None: return async with prisma_client._autorouter_turn_transactions_lock: @@ -612,6 +628,95 @@ class DBSpendUpdateWriter: except Exception as e: # noqa: BLE001 # a metrics enqueue must never fail the spend write verbose_proxy_logger.debug("_enqueue_autorouter_turn_transaction error (non-blocking): %s", e) + async def _enqueue_baseline_accounting( + self, + payload: SpendLogsPayload, + metadata: Mapping[str, object], + turn: "AutoRouterTurnTransaction | None", + prisma_client: "PrismaClient", + ) -> bool: + from litellm.proxy.db.baseline_accounting import ( + BaselineAccountingRecord, + ) + from litellm.proxy.hooks.autorouter_baseline_cache import CapturedBaselineObservation + from litellm.proxy.spend_tracking.savings import baseline_cost_snapshot + + serialized: Final = metadata.get("autorouter_baseline_observation") + if not isinstance(serialized, str): + return False + captured: Final = CapturedBaselineObservation.model_validate_json(serialized) + if captured.api_key != payload["api_key"] or captured.session_id != payload["session_id"]: + return False + decision: Final = _SPEND_METADATA_ADAPTER.validate_python( + metadata.get("routing_decision") or MappingProxyType({}) + ) + breakdown: Final = _SPEND_METADATA_ADAPTER.validate_python( + metadata.get("cost_breakdown") or MappingProxyType({}) + ) + daily: Final = await self._baseline_daily_attribution(payload, prisma_client) + record: Final = BaselineAccountingRecord( + scope=captured.scope, + api_key=captured.api_key, + session_id=captured.session_id, + router_name=captured.router_name, + baseline_model=captured.baseline_model, + observation=captured.observation.model_copy(update=MappingProxyType({"request_id": payload["request_id"]})), + pricing=baseline_cost_snapshot(captured.model, captured.prices, payload["spend"], breakdown, decision), + turn=turn, + daily=daily, + ) + async with prisma_client.baseline_accounting_lock: + if len(prisma_client.baseline_accounting_transactions) >= 10000: + verbose_proxy_logger.warning("Auto-router baseline observation queue is full") + return False + prisma_client.baseline_accounting_transactions.append(record) + from litellm.proxy.utils import request_spend_log_flush + + request_spend_log_flush(prisma_client) + return True + + async def _baseline_daily_attribution( + self, + payload: SpendLogsPayload, + prisma_client: "PrismaClient", + ) -> "DailyBaselineAttribution | None": + from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution, DailyBaselineTarget + + normalized: Final = cast(SpendLogsPayload, MappingProxyType({**payload, "end_user_id": payload["end_user"]})) + bases: Final = tuple( + zip( + DAILY_SPEND_TABLES, + await asyncio.gather( + *( + self._common_add_spend_log_transaction_to_daily_transaction( # pyright: ignore[reportUnknownMemberType] # legacy payload union; this caller supplies a validated spend payload + normalized, + prisma_client, + "request_tags" if entity == "tag" else entity, + ) + for entity in DAILY_SPEND_TABLES + ) + ), + ) + ) + base: Final = next((base for _, base in bases if base is not None), None) + if base is None: + return None + return DailyBaselineAttribution( + date=base["date"], + api_key=base["api_key"], + model=base.get("model"), + custom_llm_provider=base.get("custom_llm_provider"), + model_group=base.get("model_group"), + endpoint=base.get("endpoint"), + mcp_namespaced_tool_name=base.get("mcp_namespaced_tool_name"), + targets=tuple( + DailyBaselineTarget(entity=entity, entity_id=identity) + for entity, values in bases + if values is not None + for identity in daily_spend_entity_ids(payload, entity) + ), + ) + def _enqueue_tool_registry_upsert( self, kwargs: dict | None, @@ -2322,21 +2427,13 @@ class DBSpendUpdateWriter: prisma_client: PrismaClient, type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user", ) -> BaseDailySpendTransaction | None: - common_expected_keys: Final = ["startTime", "api_key"] - if type == "user": - expected_keys = ["user", *common_expected_keys] - elif type == "team": - expected_keys = ["team_id", *common_expected_keys] - elif type == "org": - expected_keys = ["organization_id", *common_expected_keys] - elif type == "request_tags": - expected_keys = ["request_tags", *common_expected_keys] - elif type == "end_user": - expected_keys = ["end_user_id", *common_expected_keys] - elif type == "agent": - expected_keys = ["agent_id", *common_expected_keys] - else: - raise ValueError(f"Invalid type: {type}") + entity: Final = "tag" if type == "request_tags" else type + identity_payload: Final = ( + MappingProxyType({**payload, "end_user": payload.get("end_user_id")}) if type == "end_user" else payload + ) + if not daily_spend_entity_ids(identity_payload, entity): + return None + expected_keys: Final = ("startTime", "api_key") if not all(key in payload for key in expected_keys): verbose_proxy_logger.debug( "Missing expected keys: %s, in payload, skipping from daily_user_spend_transactions", expected_keys @@ -2399,6 +2496,7 @@ class DBSpendUpdateWriter: usage_object=usage_obj, cost_breakdown=_metadata.get("cost_breakdown"), recorded_autorouter_savings=_metadata.get("autorouter_savings"), + recorded_autorouter_savings_estimate=_metadata.get("autorouter_savings_estimate"), billed_at=payload.get("endTime"), ) timed_duration_ms: Final = _timed_request_duration_ms(payload, request_status, is_internal_call) @@ -2597,14 +2695,10 @@ class DBSpendUpdateWriter: verbose_proxy_logger.debug("request_tags is None for request. Skipping incrementing tag spend.") return - request_tags: Sequence[str] = [] - if isinstance(payload["request_tags"], str): - request_tags = json.loads(payload["request_tags"]) - elif isinstance(payload["request_tags"], list): - request_tags = payload["request_tags"] - else: - raise ValueError(f"Invalid request_tags: {payload['request_tags']}") + request_tags: Final = daily_spend_entity_ids(payload, "tag") for tag in request_tags: + if tag is None: + continue endpoint_str = base_daily_transaction.get("endpoint") or "" daily_transaction_key = f"{tag}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" daily_transaction = DailyTagSpendTransaction( diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index e97e9f6e683..b28a653c9aa 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -549,6 +549,17 @@ class SpendLogCleanup: Prune auto-router session rollup rows, which carry their own retention horizon. """ session_cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=float(retention_seconds)) + from litellm.proxy.db.baseline_accounting import BaselineAccountingStore + + if remaining_ms := self._remaining_timeout_ms(deadline)(): + try: + await BaselineAccountingStore.for_client(prisma_client).retire_before( + session_cutoff, + self.batch_size, + remaining_ms, + ) + except Exception: # noqa: BLE001 # retained observations are retried by the next cleanup job + verbose_proxy_logger.warning("Auto-router baseline retention remains pending") sessions_result: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff, deadline) verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_result.rows_deleted) return (sessions_result,) diff --git a/litellm/proxy/hooks/__init__.py b/litellm/proxy/hooks/__init__.py index a504c2ba102..0e78a0843cd 100644 --- a/litellm/proxy/hooks/__init__.py +++ b/litellm/proxy/hooks/__init__.py @@ -2,6 +2,7 @@ import os from typing import Final, Literal from . import * +from .autorouter_baseline_cache import AutoRouterBaselineCache from .cache_control_check import _PROXY_CacheControlCheck from .litellm_skills import SkillsInjectionHook from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler @@ -25,6 +26,7 @@ PROXY_HOOKS: Final = { "max_budget_per_session_limiter": _PROXY_MaxBudgetPerSessionHandler, "sensitive_data_routing": _PROXY_SensitiveDataRoutingHandler, "prompt_cache_prediction": PromptCacheObserver, + "autorouter_baseline_cache": AutoRouterBaselineCache, } ## FEATURE FLAG HOOKS ## diff --git a/litellm/proxy/hooks/autorouter_baseline_cache.py b/litellm/proxy/hooks/autorouter_baseline_cache.py new file mode 100644 index 00000000000..8cea7d0e364 --- /dev/null +++ b/litellm/proxy/hooks/autorouter_baseline_cache.py @@ -0,0 +1,344 @@ +from __future__ import annotations + +import asyncio +import hashlib +import json +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass, replace +from datetime import datetime +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +import httpx +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter + +from litellm._logging import verbose_proxy_logger +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import ( + get_litellm_metadata_from_kwargs, # pyright: ignore[reportUnknownVariableType] # legacy metadata boundary validated below +) +from litellm.llms.anthropic.prompt_cache_prediction import ( + CountedPromptCachePlan, + NativePredictionTarget, + TokenCounter, + UnsupportedCachePlan, + UnsupportedPredictionTarget, + count_cache_plan, + count_prompt_tokens, + parse_cache_plan, + resolve_baseline_prediction_target, + supported_baseline_recipient, + supported_prediction_headers, +) +from litellm.proxy.spend_tracking.baseline_accounting import BaselineObservation +from litellm.proxy.spend_tracking.savings import ( + _effective_model_info, # pyright: ignore[reportPrivateUsage] # existing deployment-price owner + _proxy_llm_router, # pyright: ignore[reportPrivateUsage] # existing optional proxy-router owner +) +from litellm.types.router import BaselineRouteStamp +from litellm.types.utils import CallTypes, ModelInfo, Usage +from litellm.utils import get_prompt_cache_min_tokens + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.utils import PrismaClient + from litellm.router import Router + +_METADATA: Final = TypeAdapter(Mapping[str, object]) +_PRICES: Final[TypeAdapter[ModelInfo | None]] = TypeAdapter(ModelInfo | None) +_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue]) +_COUNT_TIMEOUT: Final = 3.0 +_MAX_COUNTS: Final = 4096 + + +class CapturedBaselineObservation(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + scope: str + api_key: str + session_id: str + router_name: str + baseline_model: str + model: str + prices: ModelInfo | None + observation: BaselineObservation + + +@dataclass(frozen=True, slots=True) +class BaselineCacheContext: + collector: AutoRouterBaselineCache + capture: CapturedBaselineObservation + target: NativePredictionTarget | UnsupportedPredictionTarget + baseline_deployment_id: str + invalidated: str | None = None + + +class _Metadata(BaseModel): + model_config = ConfigDict(strict=True, arbitrary_types_allowed=True) + route: BaselineRouteStamp = Field(alias="_autorouter_baseline_route") + user_api_key_hash: str = Field(min_length=1) + session_id: str | None = None + + +class _WireEvent(BaseModel): + model_config = ConfigDict(strict=True, arbitrary_types_allowed=True) + httpx_response: httpx.Response + api_call_start_time: datetime + completion_start_time: datetime + custom_llm_provider: str + stream: bool = False + prompt_cache_response_complete: bool = False + + +class _ResponseUsage(BaseModel): + model_config = ConfigDict(strict=True, from_attributes=True) + usage: Usage | None = None + + +def _digest(value: object) -> str: + return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + +class AutoRouterBaselineCache(CustomLogger): + def __init__( + self, + prisma_client: PrismaClient | None, + router: Callable[[], Router | None] = _proxy_llm_router, + token_counter: TokenCounter | None = None, + clock: Callable[[], float] = time.time, + ) -> None: + super().__init__() # pyright: ignore[reportUnknownMemberType] # legacy callback constructor + self.router: Final = router + self.token_counter: Final = token_counter + self.clock: Final = clock + self.count_slots: Final = asyncio.Semaphore(8) + self.counts: Mapping[str, tuple[int, float]] = MappingProxyType({}) + + async def async_pre_call_deployment_hook(self, kwargs: Mapping[str, object], call_type: CallTypes | None) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = kwargs.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging) or call_type != CallTypes.anthropic_messages: + return + try: + metadata: Final = _METADATA.validate_python( + get_litellm_metadata_from_kwargs( + {"litellm_params": kwargs} # mutable-ok: legacy metadata owner requires a dictionary + ) + ) + if metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY): + return + if logging_obj.baseline_cache_context is not None: + await invalidate_baseline_cache(logging_obj, "retried_request") + return + request: Final = _Metadata.model_validate(metadata) + session: Final = kwargs.get("litellm_session_id") or request.session_id or logging_obj.litellm_session_id + if not isinstance(session, str) or not session or len(session) > 256: + return + router: Final = self.router() + deployment: Final = router.get_deployment(request.route.baseline_deployment_id) if router else None + if deployment is None: + return + target: Final = resolve_baseline_prediction_target(deployment.litellm_params) + prices: Final = _PRICES.validate_python( + _effective_model_info(router, request.route.baseline_deployment_id, request.route.baseline_model) + ) + scope: Final = "autorouter-baseline:v3:" + _digest( + ( + request.user_api_key_hash, + session, + request.route.router_name, + request.route.baseline_deployment_id, + deployment.litellm_params.model_dump(mode="json"), + prices, + ) + ) + started: Final = logging_obj.start_time.timestamp() + capture: Final = CapturedBaselineObservation( + scope=scope, + api_key=request.user_api_key_hash, + session_id=session, + router_name=request.route.router_name, + baseline_model=request.route.baseline_model, + model=target.model if isinstance(target, NativePredictionTarget) else request.route.baseline_model, + prices=prices, + observation=BaselineObservation( + request_id=logging_obj.litellm_call_id, + started_at=started, + available_at=started, + outcome="uncertain", + baseline_equivalent=False, + reason="incomplete_response", + ), + ) + logging_obj.baseline_cache_context = BaselineCacheContext( + self, capture, target, request.route.baseline_deployment_id + ) + except Exception: # noqa: BLE001 # optional observation cannot fail inference + verbose_proxy_logger.warning("Auto-router baseline observation could not be initialized") + + async def _count(self, target: NativePredictionTarget, body: Mapping[str, JsonValue]) -> int | None: + key: Final = _digest((target.model, target.api_key, target.api_base, _JSON_BODY.validate_python(body))) + now: Final = self.clock() + cached: Final = self.counts.get(key) + if cached is not None and cached[1] > now: + return cached[0] + async with self.count_slots: + tokens: Final = ( + await self.token_counter(target.model, target.api_key, body) + if self.token_counter is not None + else await count_prompt_tokens(target.model, target.api_key, body, api_base=target.api_base) + ) + if tokens is None or tokens < 0: + return None + retained: Final = tuple((k, v) for k, v in self.counts.items() if v[1] > now and k != key)[-(_MAX_COUNTS - 1) :] + self.counts = MappingProxyType(dict((*retained, (key, (tokens, now + 3600))))) + return tokens + + async def plan( + self, target: NativePredictionTarget, wire: httpx.Request, body: Mapping[str, JsonValue], usage: Usage | None + ) -> tuple[CountedPromptCachePlan | None, str | None]: + if not supported_prediction_headers(wire.headers): + return None, "unsupported_request_headers" + plan: Final = parse_cache_plan(body) + if isinstance(plan, UnsupportedCachePlan): + return None, plan.reason + details: Final = usage.prompt_tokens_details if usage is not None else None + if ( + not plan.breakpoints + and details is not None + and ((details.cached_tokens or 0) + (details.cache_creation_tokens or 0)) + ): + return None, "implicit_cache_without_breakpoints" + + async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + return await self._count(target, body) + + try: + counted: Final = await asyncio.wait_for( + count_cache_plan(target.model, target.api_key, plan, token_counter=count), timeout=_COUNT_TIMEOUT + ) + return (None, counted.reason) if isinstance(counted, UnsupportedCachePlan) else (counted, None) + except TimeoutError: + return None, "token_count_timeout" + except Exception: # noqa: BLE001 # token counting cannot fail a completed request + return None, "token_count_unavailable" + + +async def invalidate_baseline_cache(logging_obj: Logging, reason: str, *, completed: bool = False) -> None: + context: Final = logging_obj.baseline_cache_context + if context is not None: + logging_obj.baseline_cache_context = replace( + context, invalidated=reason + ) # rebind-ok: request-owned retry marker + logging_obj.baseline_observation = context.capture.model_copy( + update=MappingProxyType( + { # rebind-ok: capture uncertainty for failure logging + "observation": context.capture.observation.model_copy( + update=MappingProxyType( + { + "available_at": max(context.capture.observation.started_at, context.collector.clock()), + "reason": reason, + } + ) + ), + } + ) + ) + + +async def finalize_baseline_cache(logging_obj: Logging, response_obj: object) -> None: + context: Final = logging_obj.baseline_cache_context + if context is None: + return + try: + capture: Final = await _capture(context, logging_obj, response_obj) + if logging_obj.baseline_cache_context is context: + logging_obj.baseline_observation = capture # rebind-ok: attach only to the captured request owner + except Exception: # noqa: BLE001 # observation failures must preserve inference and billing + await invalidate_baseline_cache(logging_obj, "observation_unavailable") + + +async def _capture( + context: BaselineCacheContext, logging_obj: Logging, response_obj: object +) -> CapturedBaselineObservation: + original: Final = context.capture.observation + details: Final = _METADATA.validate_python(logging_obj.model_call_details) + if details.get("cache_hit") is True: + return context.capture.model_copy( + update=MappingProxyType( + { + "observation": original.model_copy( + update=MappingProxyType({"outcome": "response_cache", "reason": "response_cache_hit"}) + ) + } + ) + ) + event: Final = _WireEvent.model_validate(details) + wire: Final = event.httpx_response.request + usage: Final = _ResponseUsage.model_validate(response_obj).usage + complete: Final = ( + event.custom_llm_provider == "anthropic" + and event.httpx_response.status_code == 200 + and (not event.stream or event.prompt_cache_response_complete) + ) + started: Final = original.started_at + available: Final = event.completion_start_time.timestamp() + if context.invalidated or not complete or not started <= available <= context.collector.clock(): + return context.capture.model_copy( + update=MappingProxyType( + { + "observation": original.model_copy( + update=MappingProxyType( + { + "available_at": max(started, context.collector.clock()), + "reason": context.invalidated or "incomplete_response", + } + ) + ) + } + ) + ) + target: Final = context.target + if isinstance(target, UnsupportedPredictionTarget) or not supported_baseline_recipient(target, wire): + return context.capture.model_copy( + update=MappingProxyType( + { + "observation": original.model_copy( + update=MappingProxyType( + { + "available_at": available, + "reason": target.reason + if isinstance(target, UnsupportedPredictionTarget) + else "unsupported_baseline_recipient", + } + ) + ) + } + ) + ) + body: Final = _JSON_BODY.validate_json(wire.content) + same: Final = ( + logging_obj.get_router_model_id() == context.baseline_deployment_id and body.get("model") == target.model + ) + plan, reason = await context.collector.plan(target, wire, body, usage) + minimum: Final = get_prompt_cache_min_tokens(target.model) + return context.capture.model_copy( + update=MappingProxyType( + { + "observation": BaselineObservation( + request_id=original.request_id, + started_at=started, + available_at=available, + outcome="complete", + baseline_equivalent=same, + usage=usage, + plan=plan, + minimum_cache_tokens=minimum, + reason=reason, + ) + } + ) + ) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 200ed6c3bf3..a6d5a17d73e 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -556,6 +556,9 @@ class _SessionAggRow(BaseModel): total_tokens: int spend: float saved_spend: float + savings_estimated_turns: int = 0 + savings_estimated_actual_spend: float = 0.0 + savings_estimated_saved_spend: float = 0.0 classifier_cost: float classifier_cost_recorded_turns: int session_seconds: float @@ -582,9 +585,19 @@ def _cache_bucket(turns: int, hits: int) -> AutoRouterCacheBucket: return AutoRouterCacheBucket(turns=turns, hits=hits, hit_rate_pct=_pct(hits, turns)) +def _savings_cohort( + turns: int, estimated_turns: int, actual_spend: float, saved_spend: float +) -> tuple[float | None, float | None]: + if turns > 0 and estimated_turns == 0: + return None, None + return saved_spend, actual_spend + saved_spend + + def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals: return_misses: Final = row.return_turns - row.return_hits - baseline_spend: Final = row.spend + row.saved_spend + saved_spend, baseline_spend = _savings_cohort( + row.turns, row.savings_estimated_turns, row.savings_estimated_actual_spend, row.savings_estimated_saved_spend + ) sessions: Final = row.sessions return AutoRouterBenchmarkTotals( sessions=sessions, @@ -593,11 +606,15 @@ def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals: avg_session_seconds=row.session_seconds / sessions if sessions else 0.0, avg_tokens_per_session=row.total_tokens / sessions if sessions else 0.0, spend=row.spend, - saved_spend=row.saved_spend, + savings_estimated_turns=row.savings_estimated_turns, + savings_estimated_actual_spend=row.savings_estimated_actual_spend, + saved_spend=saved_spend, classifier_cost=row.classifier_cost if row.classifier_cost_recorded_turns == row.turns else None, baseline_spend=baseline_spend, - saved_pct=_pct(row.saved_spend, baseline_spend), - saved_per_session=row.saved_spend / sessions if sessions else 0.0, + saved_pct=_pct(saved_spend, baseline_spend) if saved_spend is not None and baseline_spend is not None else None, + saved_per_session=(row.savings_estimated_saved_spend / sessions if sessions else 0.0) + if row.savings_estimated_turns == row.turns + else None, cache=AutoRouterCacheStats( coverage_pct=_pct(row.covered_turns, row.turns), hit_rate_pct=_pct(row.cache_hits, row.covered_turns), @@ -627,6 +644,8 @@ def _benchmark_group(row: _SessionAggRow) -> AutoRouterBenchmarkGroup: avg_tokens_per_session=totals.avg_tokens_per_session, spend=totals.spend, saved_spend=totals.saved_spend, + savings_estimated_turns=totals.savings_estimated_turns, + savings_estimated_actual_spend=totals.savings_estimated_actual_spend, classifier_cost=totals.classifier_cost, baseline_spend=totals.baseline_spend, saved_pct=totals.saved_pct, @@ -658,6 +677,9 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow: total_tokens=sum(row.total_tokens for row in rows), spend=sum(row.spend for row in rows), saved_spend=sum(row.saved_spend for row in rows), + savings_estimated_turns=sum(row.savings_estimated_turns for row in rows), + savings_estimated_actual_spend=sum(row.savings_estimated_actual_spend for row in rows), + savings_estimated_saved_spend=sum(row.savings_estimated_saved_spend for row in rows), classifier_cost=sum(row.classifier_cost for row in rows), classifier_cost_recorded_turns=sum(row.classifier_cost_recorded_turns for row in rows), session_seconds=sum(row.session_seconds for row in rows), @@ -807,6 +829,9 @@ async def get_auto_router_session( raise HTTPException( status_code=404, detail=f"No auto-routed turns recorded for session {session_id!r} under this key" ) + saved_spend, baseline_spend = _savings_cohort( + row.turns, row.savings_estimated_turns, row.savings_estimated_actual_spend, row.savings_estimated_saved_spend + ) return AutoRouterSessionResponse( session_id=session_id, router_name=row.router_name, @@ -814,10 +839,13 @@ async def get_auto_router_session( turns=row.turns, last_model=row.last_model, spend=row.spend, - saved_spend=row.saved_spend, - baseline_spend=row.spend + row.saved_spend, + savings_estimated_turns=row.savings_estimated_turns, + savings_estimated_actual_spend=row.savings_estimated_actual_spend, + saved_spend=saved_spend, + baseline_spend=baseline_spend if row.savings_estimated_turns == row.turns else None, + savings_estimated_baseline_spend=baseline_spend, baseline_model=row.baseline_model, - baseline_models=row.baseline_models, + baseline_models=row.savings_estimated_baseline_models, ) diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 94a59828451..e395f56194f 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.utils import get_custom_url from litellm.repositories.table_repositories import ClaudeCodePluginRepository +from litellm.router_strategy.complexity_router.fuse_presets import FusePresetCatalog, get_fuse_presets from litellm.types.agents import AgentCard from litellm.types.mcp import MCPPublicServer from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -424,6 +425,14 @@ async def get_complexity_scorer_defaults() -> ComplexityScorerDefaults: ) +@router.get( + "/public/complexity_router/fuse_presets", + response_model=FusePresetCatalog, +) +async def get_public_fuse_presets() -> FusePresetCatalog: + return get_fuse_presets() + + @router.get( "/public/litellm_model_cost_map", tags=["public", "model management"], diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index c4606796ebf..d2032cec0d0 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1551,6 +1551,36 @@ model LiteLLM_AdaptiveRouterSession { @@index([last_activity_at], map: "idx_adaptive_router_session_activity") } +model LiteLLM_AutoRouterBaselineComparison { + scope String @id + api_key String + session_id String + router_name String + initial_equivalent Boolean + revision BigInt @default(0) + published_revision BigInt @default(0) + history String? + attempted_at DateTime? + retired Boolean @default(false) + updated_at DateTime @default(now()) + + @@index([api_key, session_id, router_name], map: "idx_autorouter_baseline_scope") + @@index([updated_at], map: "idx_autorouter_baseline_updated") +} + +model LiteLLM_AutoRouterBaselineObservation { + request_id String @id + scope String + started_at Float + revision BigInt + data String + publication String? + conflicted Boolean @default(false) + + @@index([scope, started_at, request_id], map: "idx_autorouter_baseline_event_order") + @@index([scope, revision, started_at], map: "idx_autorouter_baseline_event_revision") +} + model LiteLLM_AutoRouterSession { api_key String session_id String @@ -1577,6 +1607,10 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + savings_estimated_turns Int @default(0) + savings_estimated_actual_spend Float @default(0) + savings_estimated_saved_spend Float @default(0) + savings_estimated_baseline_models Json @default("{}") classifier_cost Float @default(0) classifier_cost_recorded_turns Int @default(0) tier_turns Json @default("{}") diff --git a/litellm/proxy/spend_tracking/baseline_accounting.py b/litellm/proxy/spend_tracking/baseline_accounting.py new file mode 100644 index 00000000000..5980fb66211 --- /dev/null +++ b/litellm/proxy/spend_tracking/baseline_accounting.py @@ -0,0 +1,348 @@ +"""Pure, chronological cache accounting for the recorded baseline comparison. + +Observation collection, pricing and durable publication belong to their existing +owners. Replaying these values in event order is independent of callback order. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from itertools import groupby +from math import isfinite +from types import MappingProxyType +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, Field + +from litellm.llms.anthropic.prompt_cache_prediction import CountedBreakpoint, CountedPromptCachePlan +from litellm.types.utils import CacheCreationTokenDetails, PromptTokensDetailsWrapper, Usage + +MAX_CACHE_TTL: Final = 3600 +MAX_CACHE_ENTRIES: Final = 1024 + + +class BaselineObservation(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + version: Literal[3] = 3 + request_id: str = Field(min_length=1) + started_at: float = Field(allow_inf_nan=False, ge=0) + available_at: float = Field(allow_inf_nan=False, ge=0) + outcome: Literal["complete", "uncertain", "response_cache"] + baseline_equivalent: bool + usage: Usage | None = None + plan: CountedPromptCachePlan | None = None + minimum_cache_tokens: int = Field(default=0, ge=0) + reason: str | None = None + + +@dataclass(frozen=True, slots=True) +class BaselineEstimate: + request_id: str + reason: str + provenance: Literal["observed_identical", "modeled"] | None = None + usage: Usage | None = None + + +@dataclass(frozen=True, slots=True) +class CacheEntry: + fingerprint: str + content_fingerprint: str + tokens: int + ttl_seconds: int + available_at: float + expires_at: float + uncertain: bool = False + + +@dataclass(frozen=True, slots=True) +class BaselineHistory: + first_at: float | None = None + last_at: float | None = None + equivalent: bool = True + uncertain_before: float = 0.0 + entries: tuple[CacheEntry, ...] = () + blocked_until: float = 0.0 + + +def _complete_usage(usage: Usage | None) -> bool: + if usage is None or usage.prompt_tokens < 0 or usage.completion_tokens < 0: + return False + details: Final = usage.prompt_tokens_details + if details is None: + return False + values: Final = (details.text_tokens, details.cached_tokens, details.cache_creation_tokens) + if any(value is None or value < 0 for value in values): + return False + split: Final = details.cache_creation_token_details + writes: Final = details.cache_creation_tokens or 0 + return ( + usage.total_tokens == usage.prompt_tokens + usage.completion_tokens + and sum(value or 0 for value in values) == usage.prompt_tokens + and ( + writes == 0 + or ( + split is not None + and split.ephemeral_5m_input_tokens is not None + and split.ephemeral_1h_input_tokens is not None + and min(split.ephemeral_5m_input_tokens, split.ephemeral_1h_input_tokens) >= 0 + and split.ephemeral_5m_input_tokens + split.ephemeral_1h_input_tokens == writes + ) + ) + ) + + +def _valid_plan(plan: CountedPromptCachePlan | None) -> bool: + if plan is None or plan.total_tokens < 0 or len(plan.breakpoints) > 4: + return False + return all( + marker.fingerprint + and marker.content_fingerprint + and marker.fingerprint in marker.lookback_fingerprints + and marker.content_fingerprint in marker.lookback_content_fingerprints + and marker.ttl_seconds in (300, 3600) + and 0 <= marker.prefix_tokens <= plan.total_tokens + for marker in plan.breakpoints + ) and all( + left.prefix_tokens <= right.prefix_tokens and left.ttl_seconds >= right.ttl_seconds + for left, right in zip(plan.breakpoints, plan.breakpoints[1:]) + ) + + +def _markers(observation: BaselineObservation) -> tuple[CountedBreakpoint, ...]: + return ( + tuple( + marker + for marker in observation.plan.breakpoints + if marker.prefix_tokens >= observation.minimum_cache_tokens + ) + if observation.plan is not None + else () + ) + + +def _matches(entry: CacheEntry, markers: tuple[CountedBreakpoint, ...], started: float) -> bool: + return entry.available_at <= started < entry.expires_at and any( + entry.fingerprint in marker.lookback_fingerprints + and entry.tokens <= marker.prefix_tokens + and entry.ttl_seconds == marker.ttl_seconds + for marker in markers + ) + + +def _ambiguous(entry: CacheEntry, markers: tuple[CountedBreakpoint, ...], started: float) -> bool: + return entry.available_at <= started < entry.expires_at and any( + entry.content_fingerprint in marker.lookback_content_fingerprints + and (entry.uncertain or entry.ttl_seconds != marker.ttl_seconds) + for marker in markers + ) + + +def _usage_with_cache(usage: Usage, total: int, read: int, write_5m: int, write_1h: int) -> Usage: + writes: Final = write_5m + write_1h + original_details: Final = usage.prompt_tokens_details or PromptTokensDetailsWrapper() + details: Final = original_details.model_copy( + deep=True, + update=MappingProxyType( + { + "text_tokens": total - read - writes, + "cached_tokens": read, + "cache_creation_tokens": writes, + "cache_write_tokens": writes, + "cache_creation_token_details": CacheCreationTokenDetails( + ephemeral_5m_input_tokens=write_5m, + ephemeral_1h_input_tokens=write_1h, + ), + } + ), + ) + return Usage.model_validate( + { # mutable-ok: Usage only runs its normalizing constructor for a plain dictionary + **usage.model_dump(), + "prompt_tokens": total, + "total_tokens": total + usage.completion_tokens, + "prompt_tokens_details": details, + "cache_read_input_tokens": read, + "cache_creation_input_tokens": writes, + }, + ) + + +def _estimate(history: BaselineHistory, observation: BaselineObservation, equivalent: bool) -> BaselineEstimate: + if observation.outcome != "complete" or not _complete_usage(observation.usage): + return BaselineEstimate(observation.request_id, observation.reason or observation.outcome) + usage: Final = observation.usage + if usage is None: + return BaselineEstimate(observation.request_id, "missing_usage") + if equivalent and observation.baseline_equivalent: + return BaselineEstimate( + observation.request_id, "identical_baseline_path", "observed_identical", usage.model_copy(deep=True) + ) + if observation.started_at < history.blocked_until: + return BaselineEstimate(observation.request_id, "concurrent_uncertainty") + plan: Final = observation.plan + if not _valid_plan(plan) or plan is None: + return BaselineEstimate(observation.request_id, observation.reason or "unsupported_cache_plan") + markers: Final = _markers(observation) + if any(_ambiguous(entry, markers, observation.started_at) for entry in history.entries): + return BaselineEstimate(observation.request_id, "cache_ttl_changed") + read: Final = max( + ( + entry.tokens + for entry in history.entries + if not entry.uncertain and _matches(entry, markers, observation.started_at) + ), + default=0, + ) + end: Final = markers[-1].prefix_tokens if markers else 0 + if read < end and observation.started_at < history.uncertain_before + max(marker.ttl_seconds for marker in markers): + return BaselineEstimate(observation.request_id, "history_unavailable") + one_hour: Final = max( + (marker.prefix_tokens for marker in markers if marker.ttl_seconds == 3600 and marker.prefix_tokens > read), + default=read, + ) + expired: Final = any( + entry.expires_at <= observation.started_at + and any(entry.fingerprint in marker.lookback_fingerprints for marker in markers) + for entry in history.entries + ) + reason: Final = ( + "cache_prefix_available" + if read + else "cache_prefix_expired" + if expired + else "cache_prefix_cold" + if markers + else "below_cache_minimum" + if plan.breakpoints + else "no_cache_breakpoints" + ) + return BaselineEstimate( + observation.request_id, + reason, + "modeled", + _usage_with_cache(usage, plan.total_tokens, read, end - one_hour, one_hour - read), + ) + + +def _writes(history: BaselineHistory, observation: BaselineObservation) -> tuple[CacheEntry, ...]: + if ( + observation.outcome != "complete" + or observation.started_at < history.blocked_until + or not _complete_usage(observation.usage) + or not _valid_plan(observation.plan) + ): + return () + markers: Final = _markers(observation) + ambiguous: Final = tuple(entry for entry in history.entries if _ambiguous(entry, markers, observation.started_at)) + hit: Final = ( + max( + ( + entry + for entry in history.entries + if not entry.uncertain and _matches(entry, markers, observation.started_at) + ), + key=lambda entry: entry.tokens, + default=None, + ) + if not ambiguous + else None + ) + refresh: Final = ( + ( + CacheEntry( + hit.fingerprint, + hit.content_fingerprint, + hit.tokens, + hit.ttl_seconds, + observation.available_at, + observation.started_at + hit.ttl_seconds, + ), + ) + if hit is not None and all(marker.fingerprint != hit.fingerprint for marker in markers) + else () + ) + return ( + *refresh, + *( + CacheEntry( + marker.fingerprint, + marker.content_fingerprint, + marker.prefix_tokens, + marker.ttl_seconds, + observation.available_at, + observation.started_at + max((marker.ttl_seconds, *(entry.ttl_seconds for entry in ambiguous))), + uncertain=bool(ambiguous), + ) + for marker in markers + ), + ) + + +def _entry_key(entry: CacheEntry) -> tuple[str, str, int, int, bool]: + return entry.fingerprint, entry.content_fingerprint, entry.tokens, entry.ttl_seconds, entry.uncertain + + +def _compact_entries(entries: tuple[CacheEntry, ...], started: float) -> tuple[CacheEntry, ...]: + ordered: Final = sorted((entry for entry in entries if entry.expires_at >= started - MAX_CACHE_TTL), key=_entry_key) + return tuple( + retained + for _, values in groupby(ordered, key=_entry_key) + for group in (tuple(values),) + for retained in ( + max( + (entry for entry in group if entry.available_at <= started), + key=lambda entry: entry.expires_at, + default=None, + ), + *(entry for entry in group if entry.available_at > started), + ) + if retained is not None + ) + + +def advance_baseline_history( + history: BaselineHistory, + simultaneous: Sequence[BaselineObservation], +) -> tuple[BaselineHistory, tuple[BaselineEstimate, ...]]: + """Apply one request-start timestamp; ties cannot manufacture initial equality. + + The storage owner groups and orders observations before calling this function. + Equal timestamps are evaluated against the same preceding cache snapshot. + """ + if not simultaneous: + return history, () + started: Final = simultaneous[0].started_at + valid_order: Final = ( + isfinite(started) + and all(item.started_at == started and item.available_at >= started for item in simultaneous) + and (history.last_at is None or started > history.last_at) + ) + if not valid_order: + return history, tuple(BaselineEstimate(item.request_id, "invalid_observation_order") for item in simultaneous) + first: Final = started if history.first_at is None else history.first_at + uncertain: Final = max(history.uncertain_before, first) + relevant: Final = tuple(item for item in simultaneous if item.outcome != "response_cache") + equivalent: Final = history.equivalent and all(item.baseline_equivalent for item in relevant) + before: Final = BaselineHistory( + first, history.last_at, equivalent, uncertain, history.entries, history.blocked_until + ) + estimates: Final = tuple(_estimate(before, item, equivalent) for item in simultaneous) + invalidated: Final = any( + item.outcome != "complete" or not _complete_usage(item.usage) or not _valid_plan(item.plan) for item in relevant + ) + blocked: Final = max((history.blocked_until, *(item.available_at for item in relevant if invalidated))) + entries: Final = _compact_entries( + () if invalidated else (*history.entries, *(entry for item in relevant for entry in _writes(before, item))), + started, + ) + overflow: Final = len(entries) > MAX_CACHE_ENTRIES + return BaselineHistory( + first_at=first, + last_at=started, + equivalent=equivalent, + uncertain_before=max(started, blocked) if invalidated or overflow else uncertain, + entries=() if overflow else entries, + blocked_until=blocked, + ), estimates diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index 7d9b6514a34..b7a2ac62844 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -10,7 +10,11 @@ have been aggregated across models. from collections.abc import Callable, Mapping from datetime import datetime -from typing import TYPE_CHECKING, Final, NamedTuple +from math import isclose, isfinite +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, NamedTuple + +from pydantic import BaseModel, ConfigDict, Field import litellm from litellm._logging import verbose_proxy_logger @@ -65,7 +69,7 @@ def _resolve_model(model: str | None, custom_llm_provider: str | None) -> _Model return None try: resolved_model, provider, _, _ = litellm.get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) - except Exception as e: # noqa: BLE001 # get_llm_provider raises for unroutable names; degrade to zero savings + except Exception as e: # noqa: BLE001 # get_llm_provider raises for unroutable names; degrade to an unavailable estimate verbose_proxy_logger.debug( "savings: cannot resolve provider for model=%s custom_llm_provider=%s (%s)", model, custom_llm_provider, e ) @@ -118,6 +122,68 @@ class PricingBasis(NamedTuple): _STANDARD_RATES: Final = PricingBasis() +class BaselineCostSnapshot(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + model: str + provider: str + prices: ModelInfo | None + basis: PricingBasis = _STANDARD_RATES + actual_spend: float = Field(allow_inf_nan=False, ge=0) + actual_token_cost: float | None = Field(default=None, allow_inf_nan=False, ge=0) + classifier_cost: float = Field(default=0.0, allow_inf_nan=False, ge=0) + + +def baseline_cost_snapshot( + model: str, + prices: ModelInfo | None, + actual_spend: float, + cost_breakdown: Mapping[str, object] | None, + routing_decision: Mapping[str, object] | None, +) -> BaselineCostSnapshot: + return BaselineCostSnapshot( + model=model, + provider="anthropic", + prices=prices, + actual_spend=actual_spend, + basis=_pricing_basis(cost_breakdown), + actual_token_cost=_recorded_token_cost(cost_breakdown), + classifier_cost=classifier_cost_from_decision(routing_decision) or 0.0, + ) + + +class BaselineCosts(NamedTuple): + actual: float + baseline: float + + @property + def savings(self) -> float: + return self.baseline - self.actual + + +def price_baseline_comparison( + snapshot: BaselineCostSnapshot, + baseline_usage: Usage | None, + provenance: Literal["observed_identical", "modeled"] | None, +) -> BaselineCosts | None: + if baseline_usage is None or provenance is None: + return None + actual: Final = snapshot.actual_spend + snapshot.classifier_cost + if provenance == "observed_identical": + return BaselineCosts(actual=actual, baseline=snapshot.actual_spend) + if snapshot.prices is None or snapshot.actual_token_cost is None: + return None + token_cost: Final = _cost_of_usage( + _ModelIdentity(snapshot.model, snapshot.provider), baseline_usage, snapshot.prices, snapshot.basis + ) + if token_cost is None or not isfinite(token_cost) or token_cost < 0: + return None + baseline: Final = snapshot.actual_spend + token_cost - snapshot.actual_token_cost + if not isfinite(baseline) or baseline < 0: + return None + return BaselineCosts(actual=actual, baseline=baseline) + + def _pricing_basis(cost_breakdown: Mapping[str, object] | None) -> PricingBasis: """The basis recorded on a request, defaulting to standard rates when absent. @@ -225,56 +291,16 @@ def _baseline_cache_rate_keys(baseline_info: ModelInfo | None) -> tuple[bool, bo ) -def _baseline_usage(usage: Usage, conversation_continuing: bool, baseline_info: ModelInfo | None = None) -> Usage: - """The same request as a single-model baseline would have met it. - - The baseline is one model serving every turn, so whether it had this prompt cached - is simply whether the conversation was already underway. On a continuing - conversation it wrote the prompt on an earlier turn and would only read it now, so - the cache tokens move into the read bucket and whatever this request paid to write - counts against the saving; that write is what switching models costs. - - On a conversation's first turn nothing was cached anywhere, for any model. The - baseline would have written the same prompt, so the cache buckets stay where they are - and both arms carry the write at their own rates, unless the baseline has no rate for - a bucket, in which case those tokens are its plain input. Charging the write to this case - too, which is all a single rollup row can support, understates a first turn to a - few percent of its value and can render a profitable route as a loss. - - A continuing turn that mostly read from cache is the third case: the selected model - was already warm, so it is the one that has been serving this conversation and the - baseline's cache holds exactly what its does. The tokens written are the turn's own - growth, new to every model, and the baseline would have paid to write them too. - Moving them would forgive the baseline a write it really owes and shrink the - reported saving. "Mostly read" rather than "read anything" on purpose: a switch onto - a model holding a small prefix of this prompt still writes most of it, and must keep - counting that write against the saving. - - Only the cache buckets move. Every other field the request was priced on travels - through untouched, audio and image and video counts among them, because the baseline - is this same request served by a model that happened to be warm; naming the fields to - keep instead would price the baseline on a request that never ran, and would go stale - the next time a priced field is added. - """ +def _baseline_usage(usage: Usage, baseline_info: ModelInfo | None = None) -> Usage: cache_read, cache_creation = _cache_token_split(usage) details: Final = usage.prompt_tokens_details if details is None or (cache_read <= 0 and cache_creation <= 0): return usage - - # The tokens this request paid to write move into the cached count and the creation - # charge is dropped: on one model that cache was already warm, so the baseline would - # have read them rather than paying to create them. The 5m/1h breakdown goes with - # them; left behind it re-charges the write. - warm: Final = conversation_continuing and cache_creation > 0 and cache_read <= cache_creation - reads = cache_read + cache_creation if warm else cache_read - writes = 0 if warm else cache_creation - prices_reads, prices_writes = _baseline_cache_rate_keys(baseline_info) - reads = reads if prices_reads else 0 - writes = writes if prices_writes else 0 + reads: Final = cache_read if prices_reads else 0 + writes: Final = cache_creation if prices_writes else 0 if (reads, writes) == (cache_read, cache_creation): return usage - other_modalities: Final = sum( (getattr(details, field, 0) or 0) for field in ("audio_tokens", "image_tokens", "video_tokens") ) @@ -309,64 +335,47 @@ def compute_autorouter_savings( cost_breakdown: Mapping[str, object] | None = None, baseline_deployment_id: str | None = None, selected_deployment_id: str | None = None, -) -> float: - """Net dollars the router saved, or cost, by serving this request on ``selected_model``. - - Signed on purpose. Switching models leaves the new one with a cold cache, so the - request pays a cache-creation charge that staying on one model would not have - incurred; when that charge outweighs the cheaper rates, routing lost money and the - dashboard has to be able to say so. Zero when both sides resolve to the same - deployment, or when either cannot be resolved or priced. - - Only one side of this subtraction is a counterfactual. What the request cost on the - model that served it is a number the operator was actually billed, and the cost - calculator already wrote it down, so ``cost_breakdown`` is read rather than - re-derived. Recomputing it means restating every pricing dimension the biller - applied, and each one omitted is a silent disagreement with the ``spend`` column - beside it; a request billed at a priority tier recomputed at standard rates reads as - half its real cost. - - The baseline has no such record, since it never ran, so it is priced through the same - cost engine on the basis the biller used for this request. An operator running that - one model instead of the router would have sent this request to the same tier and the - same region, because both are properties of the request and the deployment's - contract, not of which model the router happened to pick. - - ``conversation_continuing`` says whether the baseline would already have had this - prompt cached. It defaults to True because that is the conservative reading: a - request whose shape the router could not determine is charged the write and - under-claims rather than inflating a savings figure. - """ - # No provider argument for the baseline on purpose: it arrives from the routing - # metadata as a single self-describing string, already qualified by the auto-router, - # so there is no second field that could disagree with it. + baseline_usage: Usage | None = None, + baseline_provenance: Literal["observed_initial", "modeled"] | None = None, +) -> float | None: + """Price established baseline usage; conversation shape cannot establish cache warmth.""" baseline: Final = _resolve_model(baseline_model, None) selected: Final = _resolve_model(selected_model, selected_provider) if baseline is None or selected is None: - return 0.0 - same_target: Final = ( - baseline_deployment_id == selected_deployment_id - if baseline_deployment_id and selected_deployment_id - else baseline == selected - ) - if same_target: - return 0.0 + return None + if baseline_usage is None and any(_cache_token_split(usage)): + return None basis: Final = _pricing_basis(cost_breakdown) effective_baseline_info: Final = baseline_info if baseline_info is not None else _model_info(baseline) + modeled_usage: Final = baseline_usage if baseline_usage is not None else usage baseline_cost: Final = _cost_of_usage( - baseline, - _baseline_usage(usage, conversation_continuing, effective_baseline_info), - effective_baseline_info, - basis, + baseline, _baseline_usage(modeled_usage, effective_baseline_info), effective_baseline_info, basis + ) + recorded_selected_cost: Final = _recorded_token_cost(cost_breakdown) + selected_cost: Final = ( + recorded_selected_cost + if recorded_selected_cost is not None + else _cost_of_usage(selected, usage, selected_info, basis) ) - # Falls back to pricing the request only when the biller recorded nothing, which is - # every row written before the breakdown carried its basis. - selected_cost = _recorded_token_cost(cost_breakdown) - if selected_cost is None: - selected_cost = _cost_of_usage(selected, usage, selected_info, basis) if baseline_cost is None or selected_cost is None: - return 0.0 - return baseline_cost - selected_cost + return None + if baseline_provenance == "observed_initial": + same_prices: Final = effective_baseline_info == ( + selected_info if selected_info is not None else _model_info(selected) + ) + equivalent: Final = ( + baseline_usage is not None + and baseline_usage == usage + and baseline == selected + and bool(baseline_deployment_id) + and baseline_deployment_id == selected_deployment_id + and same_prices + and recorded_selected_cost is not None + and isclose(baseline_cost, recorded_selected_cost, rel_tol=1e-9, abs_tol=1e-12) + ) + return 0.0 if equivalent else None + difference: Final = baseline_cost - selected_cost + return difference if isfinite(difference) else None def _usage_from_spend_log(usage_object: Mapping[str, object] | None) -> Usage | None: @@ -463,11 +472,23 @@ def _proxy_llm_router() -> "Router | None": def _numeric_savings(value: object) -> float | None: """``value`` as a recorded savings figure, or ``None`` when it is not one.""" - if isinstance(value, bool) or not isinstance(value, (int, float)): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not isfinite(value): return None return float(value) +def recorded_estimated_autorouter_savings(metadata: Mapping[str, object]) -> float | None: + estimate: Final = metadata.get("autorouter_savings_estimate") + if ( + not isinstance(estimate, Mapping) + or type(estimate.get("version")) is not int + or estimate.get("version") not in (1, 2, 3) + or estimate.get("status") != "estimated" + ): + return None + return _numeric_savings(metadata.get("autorouter_savings")) + + def classifier_cost_from_decision(routing_decision: Mapping[str, object] | None) -> float | None: """The LLM-classifier cost a routing decision recorded, or ``None`` when it holds none. @@ -490,22 +511,10 @@ def autorouter_savings_for_request( model_id: str | None = None, llm_router: "Callable[[], Router | None] | None" = None, cost_breakdown: Mapping[str, object] | None = None, + baseline_usage: Usage | None = None, + baseline_provenance: Literal["observed_initial", "modeled"] | None = None, ) -> float | None: - """Auto-router savings for one request, net of the classifier call that routed it, - or ``None`` when the driver is off. - - ``None`` and ``0.0`` are different facts: ``None`` means this request cannot carry a - figure at all (no routing decision, no baseline, unusable usage), while ``0.0`` is a - real figure for a routed request whose baseline resolved to the served deployment. - Never raises: pricing failures inside degrade to zero, and the driver-off cases - return ``None``, so this is safe on the logging path where a raise would fail the - request's logging. - - The classifier deduction lives here, at the figure's one computation owner, rather - than in any reader: the stamped ``autorouter_savings`` is then already net, so the - session rollup, the daily tables and every logging consumer agree without each - re-deriving the deduction, and the recorded-figure-wins path cannot deduct twice. - """ + """Return net savings for established usage, or None when the estimate is unavailable.""" usage: Final = _usage_from_spend_log(usage_object) if usage is None or not model: return None @@ -522,15 +531,16 @@ def autorouter_savings_for_request( selected_model=model, selected_provider=custom_llm_provider, usage=usage, - # Absent means the router never recorded a shape, which is the conservative - # reading: charge the cache write rather than claim a first turn's saving. - conversation_continuing=decision.get("conversation_continuing") is not False, selected_info=_effective_model_info(router_instance, model_id, model or ""), baseline_info=_effective_model_info(router_instance, baseline_id, baseline_model or ""), cost_breakdown=cost_breakdown, baseline_deployment_id=baseline_id, selected_deployment_id=model_id, + baseline_usage=baseline_usage, + baseline_provenance=baseline_provenance, ) + if gross is None: + return None classifier_cost: Final = classifier_cost_from_decision(decision) return gross if classifier_cost is None else gross - classifier_cost @@ -542,6 +552,8 @@ def autorouter_savings_for_logging_payload( model_id: str | None, usage_object: Mapping[str, object] | None, cost_breakdown: Mapping[str, object] | None, + baseline_usage: Usage | None = None, + baseline_provenance: Literal["observed_initial", "modeled"] | None = None, ) -> float | None: """The figure the logging payload records for a request, or ``None`` when none should be. @@ -561,6 +573,8 @@ def autorouter_savings_for_logging_payload( model_id=model_id, llm_router=_proxy_llm_router, cost_breakdown=cost_breakdown, + baseline_usage=baseline_usage, + baseline_provenance=baseline_provenance, ) @@ -575,6 +589,7 @@ def compute_savings_spend( llm_router: "Callable[[], Router | None] | None" = None, cost_breakdown: Mapping[str, object] | None = None, recorded_autorouter_savings: object = None, + recorded_autorouter_savings_estimate: Mapping[str, object] | None = None, billed_at: datetime | str | None = None, ) -> SavingsSpend: """ @@ -604,11 +619,9 @@ def compute_savings_spend( figure is normally the smaller of the two, being a subset of the same requests, but not always: a request that only writes cache and never reads it has negative net savings, and dropping such a request from the attributed figure can lift it above - the total. Auto-router savings compare the - served ``model`` against the counterfactual baseline the router recorded on - its ``routing_decision``, and are zero unless the two differ. That record - also says whether the conversation was already underway, which is what tells - a mid-conversation switch from a first turn. + the total. Auto-router savings compare established baseline usage against the + recorded selected-model cost. Versioned unknown estimates contribute no dollars + to this subtotal and are excluded from the separately reported coverage cohort. ``llm_router`` is passed as a provider rather than a router because every spend write calls this and only auto-routed ones need one, so looking it up eagerly at the call @@ -653,10 +666,21 @@ def compute_savings_spend( # The figure the logging path recorded wins, before the usage gate on purpose: a row # whose usage no longer parses still carries the number computed when it did. - recorded_savings: Final = _numeric_savings(recorded_autorouter_savings) + recorded_savings: Final = ( + recorded_estimated_autorouter_savings( + MappingProxyType( + { + "autorouter_savings": recorded_autorouter_savings, + "autorouter_savings_estimate": recorded_autorouter_savings_estimate, + } + ) + ) + if recorded_autorouter_savings_estimate is not None + else _numeric_savings(recorded_autorouter_savings) + ) autorouter: Final = ( recorded_savings - if recorded_savings is not None + if recorded_savings is not None or recorded_autorouter_savings_estimate is not None else autorouter_savings_for_request( model=model, custom_llm_provider=custom_llm_provider, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 9756844b587..8f85ecdd480 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -9,7 +9,7 @@ from functools import reduce from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Protocol, cast, runtime_checkable -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue import litellm from litellm._logging import verbose_proxy_logger @@ -137,7 +137,15 @@ def _get_router_metadata_for_spend_log( ) -_STAMPED_METADATA_KEYS: Final = frozenset(("router_metadata", "azure_spillover")) +_STAMPED_METADATA_KEYS: Final = frozenset( + ( + "router_metadata", + "azure_spillover", + "autorouter_savings", + "autorouter_savings_estimate", + "autorouter_baseline_observation", + ) +) def _get_spend_logs_metadata( @@ -156,6 +164,8 @@ def _get_spend_logs_metadata( cost_breakdown: CostBreakdown | None = None, litellm_call_id: str | None = None, autorouter_savings: float | None = None, + autorouter_savings_estimate: Mapping[str, JsonValue] | None = None, + autorouter_baseline_observation: str | None = None, router_metadata: SpendLogsRouterMetadata | None = None, azure_spillover: AzureSpillover | None = None, ) -> SpendLogsMetadata: @@ -196,6 +206,8 @@ def _get_spend_logs_metadata( cost_breakdown=None, compression_savings=None, autorouter_savings=autorouter_savings, + autorouter_savings_estimate=autorouter_savings_estimate, + autorouter_baseline_observation=autorouter_baseline_observation, litellm_gateway_injected_cache=None, litellm_call_id=litellm_call_id, router_metadata=router_metadata, @@ -207,7 +219,12 @@ def _get_spend_logs_metadata( # Filter the metadata dictionary to include only the specified keys clean_metadata: Final = SpendLogsMetadata( - **{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key not in _STAMPED_METADATA_KEYS}, + **MappingProxyType( + {key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key not in _STAMPED_METADATA_KEYS} + ), + autorouter_savings=autorouter_savings, + autorouter_savings_estimate=autorouter_savings_estimate, + autorouter_baseline_observation=autorouter_baseline_observation, router_metadata=router_metadata, azure_spillover=azure_spillover, ) @@ -231,7 +248,6 @@ def _get_spend_logs_metadata( clean_metadata["cold_storage_object_key"] = cold_storage_object_key clean_metadata["litellm_overhead_time_ms"] = litellm_overhead_time_ms clean_metadata["cost_breakdown"] = cost_breakdown - clean_metadata["autorouter_savings"] = autorouter_savings clean_metadata["litellm_call_id"] = litellm_call_id return clean_metadata @@ -660,6 +676,16 @@ def get_logging_payload( autorouter_savings=( standard_logging_payload.get("autorouter_savings", None) if standard_logging_payload is not None else None ), + autorouter_savings_estimate=( + standard_logging_payload.get("autorouter_savings_estimate") + if standard_logging_payload is not None + else None + ), + autorouter_baseline_observation=( + standard_logging_payload.get("autorouter_baseline_observation") + if standard_logging_payload is not None + else None + ), litellm_call_id=litellm_call_id, router_metadata=_get_router_metadata_for_spend_log( metadata=metadata, diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index a972f08b8bf..7bdadeadf86 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -506,6 +506,16 @@ class WebSearchInterceptionSettings(BaseModel): class WebSearchInterceptionSettingsResponse(SettingsResponse): """Response model for web search interception settings""" + active_on_this_pod: bool = Field( + default=False, + description=( + "Whether the process answering this request has the interception callback " + "registered. Read-only: it reports what is running here, while values.enabled " + "is the cluster-wide setting, and the two disagree while a pod is still " + "applying a change or failed to apply it." + ), + ) + def _with_websearch_enabled_resolved(config: Mapping[str, object]) -> dict[str, object]: """ @@ -1496,11 +1506,22 @@ async def get_websearch_interception_settings( config: Final = await proxy_config.get_config() - return await _get_settings_with_schema( + from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, + ) + + settings: Final = await _get_settings_with_schema( settings_key="websearch_interception_params", settings_class=WebSearchInterceptionSettings, config=_with_websearch_enabled_resolved(config), ) + return WebSearchInterceptionSettingsResponse( + values=settings["values"], + field_schema=settings["field_schema"], + active_on_this_pod=bool( + litellm.logging_callback_manager.get_custom_loggers_for_type(WebSearchInterceptionLogger) + ), + ) @router.patch( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 62710d570db..f6f437bea75 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -248,6 +248,7 @@ if TYPE_CHECKING: from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction + from litellm.proxy.db.baseline_accounting import BaselineAccountingRecord from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction from litellm.repositories.prisma_protocols import TableActions from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline @@ -3040,6 +3041,10 @@ class ProxyLogging: Otherwise, returns None and the original exception is used. """ + logging_obj: Final[object] = request_data.get("litellm_logging_obj") # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # legacy request data is narrowed to Logging below + if isinstance(logging_obj, Logging) and logging_obj.baseline_cache_context is not None: + await logging_obj.invalidate_baseline_cache_estimate("failed_request", completed=True) + ### ALERTING ### await self.update_request_status(litellm_call_id=request_data.get("litellm_call_id", ""), status="fail") if AlertType.llm_exceptions in self.alert_types and not _is_client_error_exception(original_exception): @@ -4192,6 +4197,10 @@ class PrismaClient: http_client: "HttpConfig | None" = None, ): ## init logging object + self.baseline_accounting_transactions: list[ + BaselineAccountingRecord + ] = [] # mutable-ok: locked background queue + self.baseline_accounting_lock: Final = asyncio.Lock() self.proxy_logging_obj = proxy_logging_obj self.token_auth: DatabaseTokenAuth | None = resolve_database_token_auth() verbose_proxy_logger.debug("Creating Prisma Client..") @@ -7160,7 +7169,15 @@ async def _total_queued_spend_transactions(prisma_client: PrismaClient) -> int: autorouter_queue_size: Final = len(prisma_client.autorouter_turn_transactions) from litellm.proxy.db.shadow_eval_funnel import pending_shadow_eval_funnel_events - return spend_queue_size + tool_queue_size + autorouter_queue_size + pending_shadow_eval_funnel_events() + async with prisma_client.baseline_accounting_lock: + baseline_queue_size: Final = len(prisma_client.baseline_accounting_transactions) + return ( + spend_queue_size + + tool_queue_size + + autorouter_queue_size + + baseline_queue_size + + pending_shadow_eval_funnel_events() + ) async def update_daily_tag_spend( @@ -7225,7 +7242,10 @@ async def update_spend_logs_job( # Atomically pop batch from queue. The tool usage queue counts toward the # emptiness check: a spend-log write failure aborts a run before the tool # drain below, and those entries must not strand once the spend queue drains. + from litellm.proxy.db.baseline_accounting import flush_baseline_accounting + if await _total_queued_spend_transactions(prisma_client) == 0: + await flush_baseline_accounting(prisma_client) return logs_to_process: Final = await dequeue_spend_logs(prisma_client, MAX_LOGS_PER_INTERVAL) @@ -7282,6 +7302,8 @@ async def update_spend_logs_job( tool_tracking_err, ) + await flush_baseline_accounting(prisma_client) + async with prisma_client._autorouter_turn_transactions_lock: autorouter_turns_to_process: Final = prisma_client.autorouter_turn_transactions[:MAX_LOGS_PER_INTERVAL] remaining_autorouter_turns: Final = prisma_client.autorouter_turn_transactions[ @@ -7404,7 +7426,9 @@ async def _monitor_spend_logs_queue( proxy_logging_obj=proxy_logging_obj, ) else: - # Exponential backoff when no logs to process + from litellm.proxy.db.baseline_accounting import flush_baseline_accounting + + await flush_baseline_accounting(prisma_client) current_interval = min(current_interval * backoff_multiplier, max_backoff) if await _wait_for_spend_log_flush_request(flush_requested, current_interval): diff --git a/litellm/router.py b/litellm/router.py index 3f3daaf2eae..300b069a464 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13723,6 +13723,20 @@ class Router: to the deployment that actually served the request. Every attempt therefore writes or clears, never just writes. """ + from litellm.types.router import BaselineRouteStamp + + baseline_model: Final = routing_decision.get("savings_baseline_model") if routing_decision else None + baseline_id: Final = routing_decision.get("savings_baseline_deployment_id") if routing_decision else None + router_name: Final = routing_decision.get("router_model_name") if routing_decision else None + Router._stamp_or_clear_metadata_key( + request_kwargs=request_kwargs, + key="_autorouter_baseline_route", + value=( + BaselineRouteStamp(router_name, baseline_model, baseline_id) + if router_name and baseline_model and baseline_id + else None + ), + ) Router._stamp_or_clear_metadata_key( request_kwargs=request_kwargs, key="routing_decision", diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index c29f3b3a542..a3d6ccbd437 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -75,6 +75,7 @@ from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, ModelResponse, RoutingDecisionCause, + StandardLoggingHeuristicV2Forecast, StandardLoggingRoutingDecision, StandardLoggingRoutingDecisionTierBoundaries, ) @@ -1043,6 +1044,7 @@ class ClassificationOutcome(NamedTuple): capability_forecast: CapabilityClassifierForecast | None = None llm_v2_forecast: LLMV2Decision | None = None jev_verdict: JevVerdict | None = None + heuristic_v2_forecast: StandardLoggingHeuristicV2Forecast | None = None def _with_signal(outcome: ClassificationOutcome, signal: str | None) -> ClassificationOutcome: @@ -1075,6 +1077,8 @@ def _with_classifier_forecast( decision: StandardLoggingRoutingDecision, outcome: ClassificationOutcome ) -> StandardLoggingRoutingDecision: """Attach validated forecasts and their applied policy to the routing decision.""" + if outcome.heuristic_v2_forecast is not None: + return {**decision, "heuristic_v2_forecast": outcome.heuristic_v2_forecast} if outcome.jev_verdict is not None: forecasted_decision: Final[StandardLoggingRoutingDecision] = { **decision, @@ -1772,6 +1776,7 @@ class ComplexityRouter(CustomLogger): conversation_continuing: bool = True, tier_litellm_params: Mapping[str, object] | None = None, context_escalation_original_tier: ComplexityTier | str | None = None, + heuristic_v2_forecast: StandardLoggingHeuristicV2Forecast | None = None, ) -> StandardLoggingRoutingDecision: """Assemble the per-request provenance record for this router's decision. @@ -1831,7 +1836,9 @@ class ComplexityRouter(CustomLogger): masked_tier_litellm_params: Final = mask_credentials_in_payload(tier_litellm_params) if isinstance(masked_tier_litellm_params, Mapping): decision["tier_litellm_params"] = masked_tier_litellm_params - return decision + return ( + decision if heuristic_v2_forecast is None else {**decision, "heuristic_v2_forecast": heuristic_v2_forecast} + ) async def aclassify( self, @@ -1888,6 +1895,15 @@ class ComplexityRouter(CustomLogger): score=None, signals=(f"request-type:{request_type.value}", *probability_signals), cause="heuristic_v2", + heuristic_v2_forecast=StandardLoggingHeuristicV2Forecast( + probabilities={ + candidate.value: prediction.probabilities[index] + for index, candidate in enumerate(TIER_SEVERITY_ORDER, start=1) + }, + threshold=predictor.routing_threshold, + predicted_tier=tier.value, + request_type=request_type.value, + ), ) async def _classify_heuristic_first( @@ -3553,6 +3569,7 @@ class ComplexityRouter(CustomLogger): context_escalation_original_tier=( decision.get("context_escalation_original_tier") if decision is not None else None ), + heuristic_v2_forecast=decision.get("heuristic_v2_forecast") if decision is not None else None, ) from litellm.types.router import PreRoutingHookResponse as HookResponse @@ -3732,6 +3749,7 @@ class ComplexityRouter(CustomLogger): conversation_continuing=bool(decision.get("conversation_continuing", True)), tier_litellm_params=self._litellm_params_for_model(candidate_tier, new_model), context_escalation_original_tier=decision.get("context_escalation_original_tier"), + heuristic_v2_forecast=decision.get("heuristic_v2_forecast"), ) return response.model_copy( update={ # mutable-ok: model_copy types update as a plain dict @@ -3776,6 +3794,7 @@ class ComplexityRouter(CustomLogger): conversation_continuing=bool(decision.get("conversation_continuing", True)), tier_litellm_params=self._litellm_params_for_model(None, default_model), context_escalation_original_tier=decision.get("context_escalation_original_tier"), + heuristic_v2_forecast=decision.get("heuristic_v2_forecast"), ) return response.model_copy( update={ # mutable-ok: model_copy types update as a plain dict diff --git a/litellm/router_strategy/complexity_router/fuse_presets.json b/litellm/router_strategy/complexity_router/fuse_presets.json new file mode 100644 index 00000000000..4006366dc25 --- /dev/null +++ b/litellm/router_strategy/complexity_router/fuse_presets.json @@ -0,0 +1,100 @@ +{ + "version": "2026-09-17-v1", + "models": [ + { + "id": "gpt-6-astra-v1", + "label": "GPT-6 Astra", + "model": "gpt-6-astra", + "text": "OpenAI model for demanding end-to-end work, including reasoning, coding, research, and document tasks", + "sources": ["https://developers.openai.com/api/docs/models/gpt-6-astra"] + }, + { + "id": "gpt-5.6-sol-v1", + "label": "GPT-5.6 Sol", + "model": "gpt-5.6-sol", + "text": "OpenAI model for complex professional work, supporting reasoning and tool calling", + "sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-sol"] + }, + { + "id": "gpt-5.6-luna-v1", + "label": "GPT-5.6 Luna", + "model": "gpt-5.6-luna", + "text": "OpenAI model for high-volume workloads, supporting reasoning and tool calling", + "sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-luna"] + }, + { + "id": "gpt-5.6-terra-v1", + "label": "GPT-5.6 Terra", + "model": "gpt-5.6-terra", + "text": "OpenAI general-purpose model supporting reasoning, text and image input, and tool calling", + "sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-terra"] + }, + { + "id": "claude-haiku-4-5-v1", + "label": "Claude Haiku 4.5", + "model": "claude-haiku-4-5", + "text": "Anthropic latency-focused model supporting text and image input, tool use, and extended thinking", + "sources": ["https://platform.claude.com/docs/en/models/haiku-4-5/overview"] + }, + { + "id": "claude-sonnet-5-v1", + "label": "Claude Sonnet 5", + "model": "claude-sonnet-5", + "text": "Anthropic model balancing speed and capability, with adaptive thinking and tool use", + "sources": ["https://platform.claude.com/docs/en/models/sonnet-5/overview"] + }, + { + "id": "claude-opus-5-v1", + "label": "Claude Opus 5", + "model": "claude-opus-5", + "text": "Anthropic model for complex agentic coding and enterprise work, with adaptive thinking", + "sources": ["https://platform.claude.com/docs/en/models/opus-5/overview"] + }, + { + "id": "claude-fable-5-v1", + "label": "Claude Fable 5", + "model": "claude-fable-5", + "text": "Anthropic model for demanding reasoning and long-running agent tasks, with always-on adaptive thinking", + "sources": ["https://platform.claude.com/docs/en/models/fable-5/introducing-claude-fable-5-and-claude-mythos-5"] + }, + { + "id": "claude-fable-5-1-v1", + "label": "Claude Fable 5.1", + "model": "claude-fable-5-1", + "text": "Anthropic model for demanding reasoning, long-running agentic coding, and multistep research, with always-on adaptive thinking", + "sources": ["https://platform.claude.com/docs/en/models/fable-5-1/overview"] + } + ], + "harnesses": [ + { + "id": "unspecified-v1", + "label": "Unspecified runtime", + "text": "Agent runtime is unspecified. Assess the task using the supplied context without assuming repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://code.claude.com/docs/en/how-claude-code-works", "https://mini-swe-agent.com/latest/faq/"] + }, + { + "id": "claude-code-v1", + "label": "Claude Code", + "text": "Claude Code supplies an agent loop with context management and configured tools. Available actions depend on the session's tools, permissions, and execution environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://code.claude.com/docs/en/how-claude-code-works"] + }, + { + "id": "codex-cli-v1", + "label": "Codex CLI", + "text": "Codex CLI supplies a terminal-based coding agent. File operations, command execution, and integrations depend on the session's tools, permissions, and sandbox. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://learn.chatgpt.com/docs/codex/cli", "https://learn.chatgpt.com/codex/permissions"] + }, + { + "id": "opencode-v1", + "label": "OpenCode", + "text": "OpenCode supplies a configurable agent runtime. Available actions depend on the selected agent, tools, permissions, and execution environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://opencode.ai/docs/agents/"] + }, + { + "id": "mini-swe-agent-v1", + "label": "mini-SWE-agent", + "text": "The standard mini-SWE-agent setup uses a bash-only action interface and separate command executions. Available commands and resources depend on its configured environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://mini-swe-agent.com/latest/faq/"] + } + ] +} diff --git a/litellm/router_strategy/complexity_router/fuse_presets.py b/litellm/router_strategy/complexity_router/fuse_presets.py new file mode 100644 index 00000000000..66a96ec5ad5 --- /dev/null +++ b/litellm/router_strategy/complexity_router/fuse_presets.py @@ -0,0 +1,52 @@ +from functools import lru_cache +from importlib.resources import files +from typing import Annotated, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, StringConstraints + +ProfileText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=4000)] + + +class FuseModelPreset(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + id: str + label: str + text: ProfileText + sources: tuple[str, ...] = Field(min_length=1) + model: str + + +class FuseHarnessPreset(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + id: str + label: str + text: ProfileText + sources: tuple[str, ...] = Field(min_length=1) + + +class FusePresetCatalog(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + version: str + models: tuple[FuseModelPreset, ...] + harnesses: tuple[FuseHarnessPreset, ...] + + +@lru_cache(maxsize=1) +def get_fuse_presets() -> FusePresetCatalog: + return FusePresetCatalog.model_validate_json( + files(__package__).joinpath("fuse_presets.json").read_text(encoding="utf-8") + ) + + +def resolve_fuse_profile(text: str | None, preset_id: str | None, kind: Literal["model", "harness"]) -> str | None: + if preset_id is None: + return text + catalog: Final = get_fuse_presets() + presets: Final = catalog.models if kind == "model" else catalog.harnesses + preset: Final = next((entry for entry in presets if entry.id == preset_id), None) + if preset is None: + return None + return text if text is not None else preset.text diff --git a/litellm/router_strategy/complexity_router/llm_v2.py b/litellm/router_strategy/complexity_router/llm_v2.py index 2f545a65aaa..18351237e65 100644 --- a/litellm/router_strategy/complexity_router/llm_v2.py +++ b/litellm/router_strategy/complexity_router/llm_v2.py @@ -7,15 +7,15 @@ from dataclasses import dataclass from sys import float_info from typing import Annotated, Final, Literal, TypeAlias -from pydantic import BaseModel, ConfigDict, Field, StrictFloat, StringConstraints, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field, StrictFloat, StringConstraints, TypeAdapter, model_validator from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.base_utils import ( type_to_response_format_param, # pyright: ignore[reportUnknownVariableType] # legacy output validated below ) +from litellm.router_strategy.complexity_router.fuse_presets import ProfileText, resolve_fuse_profile ShortText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=512)] -ProfileText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=4000)] class _SolverProfile(TypedDict): @@ -139,20 +139,41 @@ class LLMV2Config(BaseModel): efficient_tier: str = "SIMPLE" capable_tier: str = "REASONING" - efficient_profile: ProfileText - capable_profile: ProfileText - harness: ProfileText + efficient_profile: ProfileText | None = None + capable_profile: ProfileText | None = None + harness: ProfileText | None = None + efficient_profile_preset: str | None = None + capable_profile_preset: str | None = None + harness_preset: str | None = None max_quality_gap: float = Field(ge=0.0, le=1.0, description="Maximum estimated success loss allowed for efficient.") max_output_tokens: int = Field(default=1024, ge=1) response_format: Literal["json_schema", "json_object"] = "json_schema" calibration: LLMV2Calibration | None = None + @model_validator(mode="after") + def validate_profiles(self) -> LLMV2Config: + self._profile_texts() + return self + + def _profile_texts(self) -> tuple[str, str, str]: + efficient: Final = resolve_fuse_profile(self.efficient_profile, self.efficient_profile_preset, "model") + capable: Final = resolve_fuse_profile(self.capable_profile, self.capable_profile_preset, "model") + harness: Final = resolve_fuse_profile(self.harness, self.harness_preset, "harness") + if efficient is None: + raise ValueError("efficient_profile requires text or a known efficient_profile_preset") + if capable is None: + raise ValueError("capable_profile requires text or a known capable_profile_preset") + if harness is None: + raise ValueError("harness requires text or a known harness_preset") + return efficient, capable, harness + def system_prompt(self, efficient_model: str, capable_model: str) -> str: + efficient, capable, harness = self._profile_texts() profiles: Final[_SolverProfiles] = { "prompt_version": LLM_V2_PROMPT_VERSION, - "harness": self.harness, - "efficient": {"model": efficient_model, "profile": self.efficient_profile}, - "capable": {"model": capable_model, "profile": self.capable_profile}, + "harness": harness, + "efficient": {"model": efficient_model, "profile": efficient}, + "capable": {"model": capable_model, "profile": capable}, } schema: Final = ( "\n\nResponse JSON schema:\n" + json.dumps(LLMV2Verdict.model_json_schema()) diff --git a/litellm/rust_bridge/legacy_callbacks.py b/litellm/rust_bridge/callbacks_legacy_python.py similarity index 100% rename from litellm/rust_bridge/legacy_callbacks.py rename to litellm/rust_bridge/callbacks_legacy_python.py diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 9f29f27e41d..fd2202a1156 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -199,13 +199,20 @@ class AutoRouterBenchmarkTotals(BaseModel): description="Recorded LLM classifier cost already included in spend; null when any session turns predate " "subtotal recording, and zero for an empty window" ) - saved_spend: float = Field( - description="Signed dollars saved versus each router's savings baseline (derived from its hardest " - "tier, or the configured override), from the same per-request savings record the usage tab reads" + savings_estimated_turns: int = Field( + description="Turns covered by the current savings estimator; legacy estimates are excluded" + ) + savings_estimated_actual_spend: float = Field( + description="Actual spend, including classifier cost, for covered turns only" + ) + saved_spend: float | None = Field( + description="Signed savings for covered turns only; null when traffic has no current estimates" + ) + baseline_spend: float | None = Field(description="Estimated single-model cost for covered turns only") + saved_pct: float | None = Field(description="Covered savings over covered baseline spend, as a percentage") + saved_per_session: float | None = Field( + description="Average session savings; unavailable unless every turn is covered" ) - baseline_spend: float = Field(description="spend plus saved_spend: the estimated single-model cost") - saved_pct: float = Field(description="saved_spend over baseline_spend, as a percentage") - saved_per_session: float cache: AutoRouterCacheStats @@ -236,16 +243,27 @@ class AutoRouterSessionResponse(BaseModel): turns: int = Field(description="Auto-routed turns the rollup has recorded for this session so far") last_model: str = Field(description="The deployment model the most recent turn was routed to") spend: float = Field(description="What the session's routed traffic actually cost, classifier calls included") - saved_spend: float = Field(description="Estimated savings against the baseline, net of classifier cost") - baseline_spend: float = Field(description="spend plus saved_spend: the estimated single-model cost") + savings_estimated_turns: int = Field( + description="Turns covered by the current savings estimator; legacy estimates are excluded" + ) + savings_estimated_actual_spend: float = Field( + description="Actual spend, including classifier cost, for covered turns only" + ) + saved_spend: float | None = Field(description="Estimated savings for covered turns only, net of classifier cost") + baseline_spend: float | None = Field( + description="Estimated single-model cost; unavailable unless every turn is covered" + ) + savings_estimated_baseline_spend: float | None = Field( + description="Estimated single-model cost for covered turns only" + ) baseline_model: str | None = Field( - description="The savings baseline most of this session's turns were priced against, recorded turn by " + description="The savings baseline most covered turns were priced against, recorded turn by " "turn, so it still names the counterfactual after the router is reconfigured or removed. None when no " "turn recorded one: rows from before the baseline was recorded, and adaptive and quality routers, " "which derive no baseline and so report no savings" ) baseline_models: Mapping[str, int] = Field( - description="Turns priced against each baseline model; more than one entry means the router's " + description="Covered turns priced against each baseline model; more than one entry means the router's " "baseline changed mid-session and baseline_spend mixes both" ) diff --git a/litellm/types/router.py b/litellm/types/router.py index adadb053ab2..aef64c09417 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -1057,6 +1057,13 @@ class TaggedPreRoutingStrategy(Generic[_PreRoutingStrategyT_co]): strategy: _PreRoutingStrategyT_co +@dataclass(frozen=True, slots=True) +class BaselineRouteStamp: + router_name: str + baseline_model: str + baseline_deployment_id: str + + @dataclass(frozen=True, slots=True) class ConsumedRequestTagsStamp: """The model group a tagged router rewrote to, plus the request tags spent selecting it.""" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4c46bee13d4..b725acf6906 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -146,6 +146,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_assistant_prefill: bool | None supports_prompt_caching: bool | None supports_prompt_cache_breakpoint: ReadOnly[bool | None] + supports_thinking_cache_preservation: ReadOnly[bool | None] supports_computer_use: bool | None supports_audio_input: bool | None supports_embedding_image_input: bool | None @@ -2974,6 +2975,13 @@ LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN: Final[InternalCallOrigin] = "llm_as_a_judg BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN: Final[InternalCallOrigin] = "background_response_cost_poll" +class StandardLoggingHeuristicV2Forecast(TypedDict): + probabilities: ReadOnly[Mapping[str, float]] + threshold: ReadOnly[float] + predicted_tier: ReadOnly[str] + request_type: ReadOnly[str] + + class StandardLoggingRoutingDecision(TypedDict, total=False): """Per-request provenance for a pre-routing strategy (auto-router) decision.""" @@ -2992,6 +3000,7 @@ class StandardLoggingRoutingDecision(TypedDict, total=False): classifier_cost: float classifier_probabilities: ReadOnly[Mapping[str, float]] classifier_confidence: ReadOnly[float] + heuristic_v2_forecast: ReadOnly[StandardLoggingHeuristicV2Forecast] classifier_crux: str # writable-ok: added only when a capability verdict is available classifier_primary_rule: str # writable-ok: added only when a capability verdict is available classifier_capability_boundary: str # writable-ok: added only when a capability verdict is available @@ -3037,6 +3046,7 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset( "classifier_cost", "classifier_probabilities", "classifier_confidence", + "heuristic_v2_forecast", "classifier_primary_rule", "classifier_capability_boundary", "classifier_p_solve", @@ -3450,7 +3460,9 @@ class StandardLoggingPayload(ClassifierAudit): stream: bool | None response_cost: float cost_breakdown: CostBreakdown | None # Detailed cost breakdown - autorouter_savings: ReadOnly[float | None] # None = not an auto-routed caller request; 0.0 is a real figure + autorouter_savings: ReadOnly[float | None] + autorouter_savings_estimate: ReadOnly[Mapping[str, JsonValue] | None] + autorouter_baseline_observation: ReadOnly[str | None] response_cost_failure_debug_info: StandardLoggingModelCostFailureDebugInformation | None status: StandardLoggingPayloadStatus status_fields: StandardLoggingPayloadStatusFields diff --git a/litellm/utils.py b/litellm/utils.py index 48d13bc16af..b724313641f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1881,6 +1881,7 @@ def client(original_function): # Type assertion: logging_obj is guaranteed to be non-None after function_setup assert logging_obj is not None, "logging_obj should not be None after function_setup" + kwargs["litellm_logging_obj"] = logging_obj modified_kwargs: Final = await async_pre_call_deployment_hook(kwargs, call_type) if modified_kwargs is not None: kwargs = modified_kwargs @@ -2848,6 +2849,14 @@ def supports_prompt_cache_breakpoint(model: str, custom_llm_provider: str | None ) +def supports_thinking_cache_preservation(model: str, custom_llm_provider: str | None = None) -> bool: + return _supports_factory( + model=model, + custom_llm_provider=custom_llm_provider, + key="supports_thinking_cache_preservation", + ) + + def supports_computer_use(model: str, custom_llm_provider: str | None = None) -> bool: """ Check if the given model supports computer use and return a boolean value. @@ -5822,6 +5831,7 @@ def _get_model_info_helper( supports_assistant_prefill=None, supports_prompt_caching=None, supports_prompt_cache_breakpoint=None, + supports_thinking_cache_preservation=None, supports_computer_use=None, supports_pdf_input=None, ) @@ -6094,6 +6104,7 @@ def _get_model_info_helper( supports_assistant_prefill=_model_info.get("supports_assistant_prefill", None), supports_prompt_caching=_model_info.get("supports_prompt_caching", None), supports_prompt_cache_breakpoint=_model_info.get("supports_prompt_cache_breakpoint", None), + supports_thinking_cache_preservation=_model_info.get("supports_thinking_cache_preservation", None), supports_audio_input=_model_info.get("supports_audio_input", None), supports_audio_output=_model_info.get("supports_audio_output", None), supports_pdf_input=_model_info.get("supports_pdf_input", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7cf858ed9ff..4b0f5e8b49a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14510,6 +14510,7 @@ "supports_native_structured_output": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_sampling_params": false, @@ -14547,6 +14548,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14698,6 +14700,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14727,6 +14730,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14759,6 +14763,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14796,6 +14801,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14831,6 +14837,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14869,6 +14876,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -14986,6 +14994,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, @@ -15027,6 +15036,7 @@ "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, "supports_reasoning": true, "supports_response_schema": true, "supports_native_structured_output": true, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 44b2569defd..509f957b8d1 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -911,6 +911,9 @@ "supports_system_messages": { "type": "boolean" }, + "supports_thinking_cache_preservation": { + "type": "boolean" + }, "supports_tool_choice": { "type": "boolean" }, diff --git a/pyproject.toml b/pyproject.toml index 1295feabb43..821f885dbfc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -302,6 +302,7 @@ editable-profile = "dev" include = [ "litellm/proxy/_experimental/out/**", "litellm/router_strategy/complexity_router/artifacts/*.json", + "litellm/router_strategy/complexity_router/fuse_presets.json", "litellm/proxy/client/cli/commands/codex_base_instructions.md", ] exclude = [ diff --git a/schema.prisma b/schema.prisma index c4606796ebf..d2032cec0d0 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1551,6 +1551,36 @@ model LiteLLM_AdaptiveRouterSession { @@index([last_activity_at], map: "idx_adaptive_router_session_activity") } +model LiteLLM_AutoRouterBaselineComparison { + scope String @id + api_key String + session_id String + router_name String + initial_equivalent Boolean + revision BigInt @default(0) + published_revision BigInt @default(0) + history String? + attempted_at DateTime? + retired Boolean @default(false) + updated_at DateTime @default(now()) + + @@index([api_key, session_id, router_name], map: "idx_autorouter_baseline_scope") + @@index([updated_at], map: "idx_autorouter_baseline_updated") +} + +model LiteLLM_AutoRouterBaselineObservation { + request_id String @id + scope String + started_at Float + revision BigInt + data String + publication String? + conflicted Boolean @default(false) + + @@index([scope, started_at, request_id], map: "idx_autorouter_baseline_event_order") + @@index([scope, revision, started_at], map: "idx_autorouter_baseline_event_revision") +} + model LiteLLM_AutoRouterSession { api_key String session_id String @@ -1577,6 +1607,10 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + savings_estimated_turns Int @default(0) + savings_estimated_actual_spend Float @default(0) + savings_estimated_saved_spend Float @default(0) + savings_estimated_baseline_models Json @default("{}") classifier_cost Float @default(0) classifier_cost_recorded_turns Int @default(0) tier_turns Json @default("{}") diff --git a/tests/integration/_support/mcp.py b/tests/integration/_support/mcp.py index bdf60becbaa..d924ee6dad0 100644 --- a/tests/integration/_support/mcp.py +++ b/tests/integration/_support/mcp.py @@ -9,7 +9,7 @@ import httpx from integration._support.asgi import asgi_server from integration._support.client import Gateway, Scenario from integration._support.database import read_rows -from mcp.server.mcpserver import MCPServer +from mcp.server.fastmcp import FastMCP from mcp.server.transport_security import TransportSecuritySettings from mcp_tests.mcp_e2e_upstream_server import add, multiply from starlette.requests import Request @@ -27,7 +27,12 @@ class McpPeer: @contextmanager def mcp_peer() -> Iterator[McpPeer]: - service: Final = MCPServer("integration-math") + service: Final = FastMCP( + "integration-math", + stateless_http=True, + json_response=True, + transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), + ) service.add_tool(add) service.add_tool(multiply) @@ -35,11 +40,7 @@ def mcp_peer() -> Iterator[McpPeer]: def fail() -> str: raise ValueError("synthetic tool failure") - app: Final = service.streamable_http_app( - stateless_http=True, - json_response=True, - transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), - ) + app: Final = service.streamable_http_app() observed: Final[queue.Queue[dict[str, object]]] = queue.Queue() async def capture(scope: Scope, receive: Receive, send: Send) -> None: @@ -93,7 +94,9 @@ def tool_names(gateway: Gateway, key: str, identity: str) -> dict[str, str]: } -def call_tool(gateway: Gateway, key: str, identity: str, name: str, arguments: dict[str, object]) -> httpx.Response: +def call_tool( + gateway: Gateway, key: str, identity: str, name: str, arguments: dict[str, object] +) -> httpx.Response: return gateway.client.post( "/mcp-rest/tools/call", headers={"x-litellm-api-key": key}, diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 3f1ecab3489..fe7b6dfe7ac 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -1311,27 +1311,6 @@ ], "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[us.anthropic.claude-opus-5-v1:0-stream_full_usage]": [ "quota_management.spend_tracking.scripted_wire.logs_cost" - ], - "tests/integration/mcp/test_mcp_lifecycle.py::test_health_intersects_route_restricted_key_grants_in_both_management_modes": [ - "other.mcp.health.restricted_keys_intersect_grants_in_both_modes" - ], - "tests/integration/mcp/test_mcp_lifecycle.py::test_warm_credential_removal_rejects_without_upstream_traffic": [ - "other.mcp.credentials.warm_removal_fails_closed_without_upstream_traffic" - ], - "tests/integration/observability/test_guardrail_effects.py::test_request_selected_mcp_guardrail_blocks_direct_and_virtual_calls": [ - "other.mcp.guardrails.request_selection_blocks_resolved_tool_without_execution" - ], - "tests/integration/mcp/test_oauth_configuration.py::test_same_url_oauth_credentials_and_revocation_are_isolated_by_user_and_server[revoke]": [ - "other.mcp.oauth.same_url_credentials_are_isolated_by_user_and_server" - ], - "tests/integration/mcp/test_oauth_configuration.py::test_same_url_oauth_credentials_and_revocation_are_isolated_by_user_and_server[expire]": [ - "other.mcp.oauth.same_url_credentials_are_isolated_by_user_and_server" - ], - "tests/integration/mcp/test_mcp_lifecycle.py::test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution[anonymous]": [ - "other.mcp.permissions.same_url_servers_enforce_discovery_and_execution" - ], - "tests/integration/mcp/test_mcp_lifecycle.py::test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution[bearer]": [ - "other.mcp.permissions.same_url_servers_enforce_discovery_and_execution" ] }, "browser": { diff --git a/tests/integration/mcp/README.md b/tests/integration/mcp/README.md deleted file mode 100644 index 6260d128a2c..00000000000 --- a/tests/integration/mcp/README.md +++ /dev/null @@ -1,28 +0,0 @@ -# MCP security regression coverage - -[LIT-4506](https://linear.app/litellm-ai/issue/LIT-4506) tracks ten gateway guards and the later JWT/OAuth acceptance. This inventory distinguishes executable assertions from unresolved coverage. A listed test counts as verified only when its exact commit has an executed, passing result - -Run the controlled gateway cases through `python tests/integration/run.py extensions`. They use real HTTP, PostgreSQL, scoped non-master keys and an SDK upstream. The existing runner supplies test entitlement; these tests do not validate licenses or external-provider consent. Canonical nodes and contract IDs live in `../contracts.json` - -| Requested guard | Existing or added coverage | Remaining limitation and owner | -| --- | --- | --- | -| 1. Discovery scoped by org/team/user/key | `test_mcp_lifecycle.py` checks the exact key-granted catalog and health visibility in both management modes. [PR #38680](https://github.com/BerriAI/litellm/pull/38680) adds team/org/user toolset E2E assertions | Per-principal native MCP coverage is not established by REST results; reuse #38680 rather than duplicate it | -| 2. Users cannot attach unauthorized servers to their own keys | Existing live probes are recorded on LIT-4506; they are not durable endpoint regression tests | Own-key create/update escalation and its permission-validator boundary remain on existing management security tickets, including [LIT-4502](https://linear.app/litellm-ai/issue/LIT-4502). A generic route denial does not prove that validator ran | -| 3. UI/API permission parity | Existing dashboard tests cover admin operations | The same non-admin actor must be tested through browser and API; admin UI tests do not establish parity. Retained with [LIT-3644](https://linear.app/litellm-ai/issue/LIT-3644) | -| 4. Server ID determines identity | `test_mcp_lifecycle.py` grants one of two servers sharing a URL and denies calls to the other, using explicit server IDs for direct REST calls and server-qualified search results for virtual calls, with and without bearer credentials | Virtual calls identify the target by the searched tool name, not the REST `server_id` field. Bare names such as `add` are ambiguous across servers; duplicate aliases/names and unprefixed protocol routing remain with [LIT-4500](https://linear.app/litellm-ai/issue/LIT-4500) | -| 5. Same-URL servers do not share credentials | `test_oauth_configuration.py` crosses two gateway users with two server IDs and four distinct stored OAuth tokens. It checks actual upstream headers and successful results, then invalidates only one tuple | Controlled stored-token tests do not prove separate external-provider accounts or consent flows | -| 6. OAuth never falls back to anonymous | OAuth isolation variants remove a stored token or expire it without refresh, require separate list/call auth failures and no upstream requests, and preserve all other valid tuples. `test_mcp_lifecycle.py` also covers warm static-header removal and OBO without a caller JWT | External upstream revocation, refresh/reauthorization and aggregate challenges remain with [LIT-4501](https://linear.app/litellm-ai/issue/LIT-4501), [LIT-3433](https://linear.app/litellm-ai/issue/LIT-3433), [LIT-4422](https://linear.app/litellm-ai/issue/LIT-4422) and [LIT-4436](https://linear.app/litellm-ai/issue/LIT-4436) | -| 7. Stateful HTTP/session continuity | Legacy public-client tests exercise initialized sessions | No claim here proves upstream session state continuity; retained under [LIT-3143](https://linear.app/litellm-ai/issue/LIT-3143) | -| 8. Production guardrails/hooks run | `../observability/test_guardrail_effects.py` checks selected pre-call guards on direct and virtual execution, key/team/request selection, allowed results and zero denied executions. The two legacy test-owned dispatcher files are removed | This does not establish every post-call/output-scanning or concurrent hook contract | -| 9. Permissions enforced at discovery and execution | Exact key catalog and virtual search results plus forbidden direct/virtual calls in `test_mcp_lifecycle.py`; existing `../compatibility/test_persisted_toolsets.py` checks tool-level ceiling, denied sibling and allowed control | All principal/transport combinations are not established; link #38680's evidence for its additional principal cases | -| 10. Stateless/stateful matrix | These controlled peers use stateless HTTP upstreams | Stateful combinations depend on LIT-3143 and shared conformance runs. Modern-agent/legacy-upstream interaction remains deferred; legacy passes do not establish modern conformance | - -## Additional JWT/OAuth acceptance - -[LIT-3467 / PR #41909](https://github.com/BerriAI/litellm/pull/41909) owns one shared real login/consent, immediate list/call and cold-restart implementation, with aggregate SSO and explicitly configured per-server JWT variants. Reuse that implementation and its protected login secret; do not create another browser bootstrap here. Credit its exact-commit evidence separately from these controlled credential tests - -The two-user/two-server cases here create non-admin users and scoped API keys through management APIs. They store synthetic upstream OAuth credentials through the real credential endpoint and assert the actual bearer at the owned upstream. This deliberately isolates credential lookup, expiry and revocation from consent. No gateway API key may replace the expected upstream token - -Gateway JWT precedence, invalid/expired gateway JWTs, inactive-user denial, and their MCP-specific interaction with isolated credential lookup remain unverified by these API-key cases. General JWT unit/API tests are useful existing coverage but do not substitute for those MCP outcomes. Real-provider auth failures should extend LIT-3467's settled helpers; its explicit-header case must not be described as an uninterrupted Authorization-only OAuth flow - -[PR #41718 / LIT-7737](https://github.com/BerriAI/litellm/pull/41718) owns dependency and public-client compatibility checks. This suite consumes the merged SDK2 API and keeps the existing dependency constraints. Its result must be reported independently of an installation-matrix pass diff --git a/tests/integration/mcp/test_mcp_lifecycle.py b/tests/integration/mcp/test_mcp_lifecycle.py index b32cf97605f..7ded23794be 100644 --- a/tests/integration/mcp/test_mcp_lifecycle.py +++ b/tests/integration/mcp/test_mcp_lifecycle.py @@ -1,18 +1,14 @@ -import json import uuid from contextlib import ExitStack -from pathlib import Path from typing import Final import pytest -import yaml from hypothesis import strategies as st from hypothesis.stateful import RuleBasedStateMachine, invariant, rule, run_state_machine_as_test from integration._support.client import Gateway from integration._support.database import read_rows from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests -from integration._support.process import owned_proxy from integration._support.mcp import call_tool, mcp_peer, register_mcp, tool_names @@ -57,7 +53,7 @@ def test_tool_error_remains_error_and_healthy_sibling_returns_value(gateway: Gat failure: Final = call_tool(gateway, key, identity, names["fail"], {}) assert failure.status_code == 200, failure.text assert failure.json()["isError"] is True - assert failure.json()["content"][0]["text"] == "Error executing tool fail" + assert "synthetic tool failure" in failure.json()["content"][0]["text"] healthy: Final = call_tool(gateway, key, identity, names["multiply"], {"a": 3, "b": 5}) assert healthy.status_code == 200, healthy.text assert healthy.json()["isError"] is False @@ -125,182 +121,3 @@ def test_generated_mcp_edits_preserve_actual_headers_and_tool_results(gateway: G self.resources.close() run_state_machine_as_test(Servers, settings=LIFECYCLE_SETTINGS) - - -@pytest.mark.covers("other.mcp.health.restricted_keys_intersect_grants_in_both_modes") -def test_health_intersects_route_restricted_key_grants_in_both_management_modes( - gateway: Gateway, tmp_path: Path -) -> None: - for mode in ("restricted", "view_all"): - config = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) - config["general_settings"]["user_mcp_management_mode"] = mode - path = tmp_path / f"health-{mode}.yaml" - path.write_text(yaml.safe_dump(config)) - with ( - owned_proxy(gateway, tmp_path, {}, config=path) as candidate, - mcp_peer() as peer, - candidate.scenario() as scenario, - ): - first = register_mcp(scenario, peer, "health" + uuid.uuid4().hex) - second = register_mcp(scenario, peer, "health" + uuid.uuid4().hex) - control = scenario.key(object_permission={"mcp_servers": [first]}) - names = tool_names(candidate, control, first) - healthy = call_tool(candidate, control, first, names["add"], {"a": 3, "b": 5}) - assert healthy.status_code == 200 and healthy.json()["content"][0]["text"] == "8", healthy.text - for grants in ([first], [second], []): - key = scenario.key( - allowed_routes=["/v1/mcp/server", "/v1/mcp/server/health"], - object_permission={"mcp_servers": grants}, - ) - listed = candidate.request("GET", "/v1/mcp/server", key=key) - assert listed.status_code == 200, listed.text - assert {row["server_id"] for row in listed.json()} == set(grants), listed.text - for requested in (None, [second], [first, second]): - response = candidate.client.get( - "/v1/mcp/server/health", - headers={"Authorization": f"Bearer {key}"}, - params=[] if requested is None else [("server_ids", identity) for identity in requested], - ) - assert response.status_code == 200, response.text - expected = set(grants) if requested is None else set(grants).intersection(requested) - assert {row["server_id"] for row in response.json()} == expected, response.text - assert all(row["status"] == "healthy" for row in response.json()) - - -@pytest.mark.covers("other.mcp.credentials.warm_removal_fails_closed_without_upstream_traffic") -def test_warm_credential_removal_rejects_without_upstream_traffic(gateway: Gateway) -> None: - with mcp_peer() as peer, gateway.scenario() as scenario: - identity = register_mcp( - scenario, - peer, - "credentials" + uuid.uuid4().hex, - auth_type="bearer_token", - static_headers={"Authorization": "Bearer synthetic-upstream-credential"}, - ) - key = scenario.key(object_permission={"mcp_servers": [identity]}) - names = tool_names(gateway, key, identity) - warm = call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5}) - assert warm.status_code == 200 and warm.json()["content"][0]["text"] == "8", warm.text - calls = tuple(item for item in peer.drain() if item["body"].get("method") == "tools/call") - assert len(calls) == 1 - assert calls[0]["headers"][b"authorization"] == b"Bearer synthetic-upstream-credential" - removed = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "static_headers": {}}) - assert removed.status_code == 202, removed.text - stored = gateway.request("GET", f"/v1/mcp/server/{identity}") - assert stored.status_code == 200, stored.text - assert stored.json()["auth_type"] == "bearer_token" - assert not stored.json().get("static_headers"), stored.text - peer.drain() - for operation in ("list", "call"): - rejected = ( - gateway.client.get( - "/mcp-rest/tools/list", params={"server_id": identity}, headers={"x-litellm-api-key": key} - ) - if operation == "list" - else call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5}) - ) - assert rejected.status_code == 500, rejected.text - if operation == "list": - assert rejected.json()["detail"]["error"] == "internal", rejected.text - assert "Failed to list tools from server" in rejected.json()["detail"]["message"], rejected.text - else: - assert "requires a usable upstream credential" in rejected.text, rejected.text - assert peer.drain() == (), "missing static credential escaped to upstream" - changed = gateway.request( - "PUT", - "/v1/mcp/server", - { - "server_id": identity, - "auth_type": "oauth2_token_exchange", - "token_exchange_endpoint": peer.url + "/token", - "credentials": {"client_id": "synthetic-client"}, - }, - ) - assert changed.status_code == 202, changed.text - peer.drain() - rejected_subject = call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5}) - assert rejected_subject.status_code == 401, rejected_subject.text - assert peer.drain() == (), "virtual key cannot supply an OBO subject token" - control_id = register_mcp(scenario, peer, "control" + uuid.uuid4().hex, auth_type="none") - control_key = scenario.key(object_permission={"mcp_servers": [control_id]}) - control_names = tool_names(gateway, control_key, control_id) - control = call_tool(gateway, control_key, control_id, control_names["multiply"], {"a": 3, "b": 5}) - assert control.status_code == 200 and control.json()["content"][0]["text"] == "15", control.text - - -@pytest.mark.parametrize("authenticated", (False, True), ids=("anonymous", "bearer")) -@pytest.mark.covers("other.mcp.permissions.same_url_servers_enforce_discovery_and_execution") -def test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution( - gateway: Gateway, authenticated: bool -) -> None: - with mcp_peer() as peer, gateway.scenario() as scenario: - aliases: Final = tuple("scope" + uuid.uuid4().hex for _ in range(2)) - servers: Final = tuple( - register_mcp( - scenario, - peer, - alias, - auth_type="bearer_token" if authenticated else "none", - static_headers={ - "X-Integration-Server": alias, - **({"Authorization": f"Bearer synthetic-{alias}"} if authenticated else {}), - }, - ) - for alias in aliases - ) - for virtual in (False, True): - keys: Final = tuple( - scenario.key(object_permission={"mcp_servers": [server], "mcp_tool_search_enabled": virtual}) - for server in servers - ) - for server, alias, key in zip(servers, aliases, keys): - catalog: Final = gateway.request("GET", "/mcp-rest/tools/list", key=key) - assert catalog.status_code == 200, catalog.text - if virtual: - assert {tool["name"] for tool in catalog.json()["tools"]} == { - "mcp_tool_search", - "mcp_tool_call", - "agent_search", - "skill_search", - }, catalog.text - search: Final = gateway.request( - "POST", - "/mcp-rest/tools/call", - {"name": "mcp_tool_search", "arguments": {"query": "add", "top_k": 10}}, - key=key, - ) - assert search.status_code == 200 and search.json()["isError"] is False, search.text - assert [tool["name"] for tool in json.loads(search.json()["content"][0]["text"])] == [ - f"{alias}-add" - ], search.text - else: - assert {tool["mcp_info"]["server_id"] for tool in catalog.json()["tools"]} == {server} - assert {tool["name"] for tool in catalog.json()["tools"]} == {"add", "multiply", "fail"} - for server_index, caller_index in ((0, 0), (1, 0), (1, 1)): - peer.drain() - response: Final = gateway.request( - "POST", - "/mcp-rest/tools/call", - { - "name": "mcp_tool_call" if virtual else "add", - **({} if virtual else {"server_id": servers[server_index]}), - "arguments": ( - {"tool_name": f"{aliases[server_index]}-add", "arguments": {"a": 3, "b": 5}} - if virtual - else {"a": 3, "b": 5} - ), - }, - key=keys[caller_index], - ) - observed: Final = peer.drain() - if server_index != caller_index: - assert response.status_code == 403 and "not allowed" in response.text, response.text - assert observed == (), "forbidden server reached the upstream" - continue - assert response.status_code == 200 and response.json()["isError"] is False, response.text - assert response.json()["content"][0]["text"] == "8", response.text - calls: Final = tuple(item for item in observed if item["body"].get("method") == "tools/call") - assert len(calls) == 1 - assert calls[0]["headers"][b"x-integration-server"] == aliases[server_index].encode() - expected_auth: Final = f"Bearer synthetic-{aliases[server_index]}".encode() if authenticated else None - assert all(item["headers"].get(b"authorization") == expected_auth for item in observed) diff --git a/tests/integration/mcp/test_oauth_configuration.py b/tests/integration/mcp/test_oauth_configuration.py index 4c46c706054..45d407f2423 100644 --- a/tests/integration/mcp/test_oauth_configuration.py +++ b/tests/integration/mcp/test_oauth_configuration.py @@ -2,14 +2,14 @@ import json import queue import uuid from urllib.parse import parse_qs, urlsplit -from typing import Final, Literal +from typing import Final from pathlib import Path import pytest from integration._support.client import Gateway, eventually from integration._support.database import read_rows -from integration._support.mcp import McpPeer, call_tool, mcp_peer, register_mcp, tool_names +from integration._support.mcp import McpPeer, register_mcp from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server @@ -102,87 +102,3 @@ def test_partial_discovery_and_unrelated_edit_keep_actual_authorization_destinat "PUT", "/v1/mcp/server", {"server_id": identity, "server_name": alias + "renamed"} ) assert updated.status_code == 202, updated.text - - -@pytest.mark.covers("other.mcp.oauth.same_url_credentials_are_isolated_by_user_and_server") -@pytest.mark.parametrize("transition", ("revoke", "expire")) -def test_same_url_oauth_credentials_and_revocation_are_isolated_by_user_and_server( - gateway: Gateway, - transition: Literal["revoke", "expire"], -) -> None: - with mcp_peer() as peer, gateway.scenario() as scenario: - servers: Final = tuple( - register_mcp( - scenario, - peer, - "oauth" + uuid.uuid4().hex, - auth_type="oauth2", - oauth2_flow="authorization_code", - authorization_url=peer.url + "/authorize", - token_url=peer.url + "/token", - credentials={"client_id": "synthetic-oauth-client"}, - ) - for _ in range(2) - ) - users: Final = tuple(scenario.user(user_role="internal_user") for _ in range(2)) - keys: Final = tuple( - scenario.key(user_id=user, object_permission={"mcp_servers": list(servers)}) for user in users - ) - for user_index, key in enumerate(keys): - for server_index, server_id in enumerate(servers): - stored: Final = gateway.request( - "POST", - f"/v1/mcp/server/{server_id}/oauth-user-credential", - {"access_token": f"synthetic-user-{user_index}-server-{server_index}", "expires_in": 3600}, - key=key, - ) - assert stored.status_code == 200 and stored.json()["has_credential"] is True, stored.text - scenario.cleanups.callback( - gateway.request, - "DELETE", - f"/v1/mcp/server/{server_id}/oauth-user-credential", - key=key, - ) - names: Final = tuple(tool_names(gateway, keys[0], server) for server in servers) - for generation in range(2): - for user_index, key in enumerate(keys): - for server_index, server_id in enumerate(servers): - peer.drain() - discovery: Final = gateway.request( - "GET", - "/mcp-rest/tools/list", - key=key, - params={"server_id": server_id}, - ) - call: Final = call_tool(gateway, key, server_id, names[server_index]["add"], {"a": 3, "b": 5}) - observed: Final = peer.drain() - if generation == 1 and user_index == 0 and server_index == 0: - for rejected in (discovery, call): - assert rejected.status_code == 401, rejected.text - assert rejected.json() == {"detail": "Unauthorized"}, rejected.text - assert "resource_metadata=" in rejected.headers["www-authenticate"] - assert observed == (), "unusable credentials must not fall back to another user or server" - else: - assert discovery.status_code == 200, discovery.text - assert {tool["name"] for tool in discovery.json()["tools"]} == set(names[server_index].values()) - assert call.status_code == 200 and call.json()["isError"] is False, call.text - assert call.json()["content"][0]["text"] == "8", call.text - calls: Final = tuple(item for item in observed if item["body"].get("method") == "tools/call") - assert len(calls) == 1 - expected: Final = f"Bearer synthetic-user-{user_index}-server-{server_index}".encode() - assert calls[0]["headers"][b"authorization"] == expected - assert all(item["headers"].get(b"authorization") == expected for item in observed) - if generation == 0: - changed: Final = gateway.request( - "DELETE" if transition == "revoke" else "POST", - f"/v1/mcp/server/{servers[0]}/oauth-user-credential", - None - if transition == "revoke" - else { - "access_token": "synthetic-expired-user-0-server-0", - "expires_in": -60, - }, - key=keys[0], - ) - assert changed.status_code == 200, changed.text - assert changed.json()["has_credential"] is (transition == "expire"), changed.text diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 5a79b619906..645af77526f 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -8,7 +8,6 @@ import yaml from integration._support.client import Gateway, eventually from integration._support.database import read_rows -from integration._support.mcp import mcp_peer, register_mcp, tool_names from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server @@ -144,73 +143,3 @@ def test_guardrail_denial_prevents_provider_and_preserves_allowed_control(gatewa ) assert len(observed.get("/__observations").json()["requests"]) == 1 assert len(policy.drain()) == 2 - - -@pytest.mark.covers("other.mcp.guardrails.request_selection_blocks_resolved_tool_without_execution") -def test_request_selected_mcp_guardrail_blocks_direct_and_virtual_calls(gateway: Gateway, tmp_path: Path) -> None: - guardrail = "mcp-policy-" + uuid.uuid4().hex - config = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) - config["guardrails"] = [ - { - "guardrail_name": guardrail, - "litellm_params": { - "guardrail": "custom_code", - "mode": "pre_mcp_call", - "default_on": False, - "custom_code": ( - "def apply_guardrail(inputs, request_data, input_type):\n" - ' if inputs.get("tools", [{}])[0].get("function", {}).get("name") == "add":\n' - ' return block("integration resolved add denied")\n' - " return allow()\n" - ), - }, - } - ] - path = tmp_path / "mcp-guardrail.yaml" - path.write_text(yaml.safe_dump(config)) - with ( - owned_proxy(gateway, tmp_path, {}, config=path) as candidate, - mcp_peer() as peer, - candidate.scenario() as scenario, - ): - identity = register_mcp(scenario, peer, "guardrail" + uuid.uuid4().hex) - permission = {"mcp_servers": [identity], "mcp_tool_search_enabled": True} - key = scenario.key(object_permission=permission) - key_selected = scenario.key(object_permission=permission, guardrails=[guardrail]) - team = scenario.team(guardrails=[guardrail], object_permission={"mcp_servers": [identity]}) - team_selected = scenario.key(team_id=team, object_permission=permission) - catalog_key = scenario.key(object_permission={"mcp_servers": [identity]}) - names = tool_names(candidate, catalog_key, identity) - assert set(names) == {"add", "multiply", "fail"} - for virtual in (False, True): - for caller, selected, tool, expected in ( - (key, [], "add", 8), - (key, [guardrail], "add", None), - (key_selected, [], "add", None), - (team_selected, [], "add", None), - (key, [guardrail], "multiply", 15), - ): - arguments = {"a": 3, "b": 5} - peer.drain() - response = candidate.client.post( - "/mcp-rest/tools/call", - headers={"x-litellm-api-key": caller}, - json={ - "server_id": identity, - "name": "mcp_tool_call" if virtual else names[tool], - "arguments": {"tool_name": names[tool], "arguments": arguments} if virtual else arguments, - "guardrails": selected, - }, - ) - calls = tuple(item for item in peer.drain() if item["body"].get("method") == "tools/call") - if expected is None: - assert response.status_code == 400, response.text - assert "integration resolved add denied" in response.text, response.text - assert calls == (), "pre-call denial must prevent upstream execution" - else: - assert response.status_code == 200, response.text - assert response.json()["isError"] is False - assert response.json()["content"][0]["text"] == str(expected), response.text - assert len(calls) == 1 - assert calls[0]["body"]["params"]["name"] == tool - assert calls[0]["body"]["params"]["arguments"] == arguments diff --git a/tests/mcp_tests/mcp_e2e_upstream_server.py b/tests/mcp_tests/mcp_e2e_upstream_server.py index 3361163badf..28fb0846481 100644 --- a/tests/mcp_tests/mcp_e2e_upstream_server.py +++ b/tests/mcp_tests/mcp_e2e_upstream_server.py @@ -1,6 +1,6 @@ """Deterministic upstream MCP server for the mcp e2e suite. -A tiny MCP server exposing `add` and `multiply` over streamable-http so the +A tiny FastMCP server exposing `add` and `multiply` over streamable-http so the suite has a self-hosted, offline upstream to register and exercise. DNS-rebinding protection is turned off because the litellm container reaches this over the compose network by service name (`mcp-upstream:8090`), not localhost, and the @@ -9,10 +9,15 @@ stack is an isolated throwaway. Bind host/port come from MCP_HOST/MCP_PORT. import os -from mcp.server.mcpserver import MCPServer +from mcp.server.fastmcp import FastMCP from mcp.server.transport_security import TransportSecuritySettings -mcp: MCPServer = MCPServer("e2e-math") +mcp: FastMCP = FastMCP( + "e2e-math", + host=os.getenv("MCP_HOST", "0.0.0.0"), + port=int(os.getenv("MCP_PORT", "8090")), + transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), +) @mcp.tool() @@ -28,12 +33,7 @@ def multiply(a: int, b: int) -> int: def main() -> None: - mcp.run( - transport="streamable-http", - host=os.getenv("MCP_HOST", "0.0.0.0"), - port=int(os.getenv("MCP_PORT", "8090")), - transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), - ) + mcp.run(transport="streamable-http") if __name__ == "__main__": diff --git a/tests/mcp_tests/test_mcp_guardrails.py b/tests/mcp_tests/test_mcp_guardrails.py new file mode 100644 index 00000000000..04401992449 --- /dev/null +++ b/tests/mcp_tests/test_mcp_guardrails.py @@ -0,0 +1,770 @@ +""" +Test file for MCP Guardrails Feature + +This file tests the MCP guardrails functionality for both pre and during MCP call hooks, +including various guardrail types and proper exception handling. +""" + +import asyncio +import pytest +from datetime import datetime +from typing import Optional, Dict, Any +from unittest.mock import MagicMock, AsyncMock, patch + +# Add the project root to the path + +import litellm +from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.caching.caching import DualCache +from litellm.types.mcp import ( + MCPPreCallRequestObject, + MCPPreCallResponseObject, + MCPDuringCallRequestObject, + MCPDuringCallResponseObject, +) +from litellm.types.llms.base import HiddenParams +from litellm.types.guardrails import GuardrailEventHooks +from fastapi import HTTPException + + +class MockPiiGuardrail(CustomGuardrail): + """Mock PII guardrail that raises BlockedPiiEntityError""" + + def __init__(self, should_block: bool = True, entity_type: str = "EMAIL_ADDRESS"): + super().__init__() + self.should_block = should_block + self.entity_type = entity_type + self.guardrail_name = "mock-pii-guardrail" + self.call_count = 0 + + def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool: + """Always run for testing""" + return True + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: str, + ): + """Mock pre-call hook that raises BlockedPiiEntityError""" + self.call_count += 1 + + if self.should_block: + raise BlockedPiiEntityError( + entity_type=self.entity_type, + guardrail_name=self.guardrail_name, + ) + return None + + +class MockContentGuardrail(CustomGuardrail): + """Mock content guardrail that raises GuardrailRaisedException""" + + def __init__(self, should_block: bool = True): + super().__init__() + self.should_block = should_block + self.guardrail_name = "mock-content-guardrail" + self.call_count = 0 + + def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool: + """Always run for testing""" + return True + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: str, + ): + """Mock pre-call hook that raises GuardrailRaisedException""" + self.call_count += 1 + + if self.should_block: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, message="Content violates policy" + ) + return None + + +class MockHttpGuardrail(CustomGuardrail): + """Mock HTTP guardrail that raises HTTPException""" + + def __init__(self, should_block: bool = True): + super().__init__() + self.should_block = should_block + self.guardrail_name = "mock-http-guardrail" + self.call_count = 0 + + def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool: + """Always run for testing""" + return True + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: str, + ): + """Mock pre-call hook that raises HTTPException""" + self.call_count += 1 + + if self.should_block: + raise HTTPException( + status_code=400, detail={"error": "Violated guardrail policy"} + ) + return None + + +class MockDuringCallGuardrail(CustomGuardrail): + """Mock guardrail for during-call testing""" + + def __init__(self, should_block: bool = True): + super().__init__() + self.should_block = should_block + self.guardrail_name = "mock-during-guardrail" + self.call_count = 0 + + def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool: + """Always run for testing""" + return True + + async def async_moderation_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + call_type: str, + ): + """Mock during-call hook that raises exceptions""" + self.call_count += 1 + + if self.should_block: + raise BlockedPiiEntityError( + entity_type="PHONE_NUMBER", + guardrail_name=self.guardrail_name, + ) + return None + + +class MockProxyLogging: + """Mock proxy logging object for testing MCP guardrails""" + + def __init__(self, guardrails: Optional[list] = None): + self.guardrails = guardrails if guardrails is not None else [] + self.call_details = {"user_api_key_cache": DualCache()} + self.dynamic_success_callbacks = [] + self.call_count = 0 + + def get_combined_callback_list(self, dynamic_success_callbacks, global_callbacks): + """Return the guardrails for testing""" + return self.guardrails + + def _convert_mcp_to_llm_format(self, request_obj, kwargs: dict) -> dict: + """Convert MCP tool call to LLM message format""" + tool_call_content = ( + f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}" + ) + + return { + "messages": [{"role": "user", "content": tool_call_content}], + "model": kwargs.get("model", "mcp-tool-call"), + "user_api_key_user_id": kwargs.get("user_api_key_user_id"), + "user_api_key_team_id": kwargs.get("user_api_key_team_id"), + } + + def _convert_llm_result_to_mcp_response(self, llm_result, request_obj): + """Convert LLM result back to MCP response format""" + return None # For testing, we don't need to convert back + + def _parse_pre_mcp_call_hook_response(self, response, original_request): + """Parse pre MCP call hook response""" + return response + + async def async_pre_mcp_tool_call_hook( + self, + kwargs: dict, + request_obj: Any, + start_time: datetime, + end_time: datetime, + ) -> Optional[Any]: + """Mock pre MCP tool call hook""" + self.call_count += 1 + + # Simulate the actual hook logic + for guardrail in self.guardrails: + if isinstance(guardrail, CustomGuardrail): + try: + synthetic_data = self._convert_mcp_to_llm_format( + request_obj, kwargs + ) + + # Check if guardrail should run + if not guardrail.should_run_guardrail( + synthetic_data, GuardrailEventHooks.pre_mcp_call + ): + continue + + result = await guardrail.async_pre_call_hook( + user_api_key_dict=kwargs.get("user_api_key_auth"), + cache=self.call_details["user_api_key_cache"], + data=synthetic_data, + call_type="mcp_call", + ) + if result is not None: + return self._parse_pre_mcp_call_hook_response( + result, request_obj + ) + except ( + BlockedPiiEntityError, + GuardrailRaisedException, + HTTPException, + ) as e: + # Re-raise guardrail exceptions + raise e + except Exception as e: + # Log non-guardrail exceptions as non-blocking + print( + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {str(e)}" + ) + + return None + + async def async_during_mcp_tool_call_hook( + self, + kwargs: dict, + request_obj: Any, + start_time: datetime, + end_time: datetime, + ) -> Optional[Any]: + """Mock during MCP tool call hook""" + self.call_count += 1 + + # Simulate the actual hook logic + for guardrail in self.guardrails: + if isinstance(guardrail, CustomGuardrail): + try: + synthetic_data = self._convert_mcp_to_llm_format( + request_obj, kwargs + ) + result = await guardrail.async_moderation_hook( + data=synthetic_data, + user_api_key_dict=kwargs.get("user_api_key_auth"), + call_type="mcp_call", + ) + if result is not None: + return result + except ( + BlockedPiiEntityError, + GuardrailRaisedException, + HTTPException, + ) as e: + # Re-raise guardrail exceptions + raise e + except Exception as e: + # Log non-guardrail exceptions as non-blocking + print( + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {str(e)}" + ) + + return None + + +@pytest.fixture +def mock_user_api_key(): + """Mock user API key for testing""" + return UserAPIKeyAuth(api_key="test_key", user_id="test_user") + + +@pytest.fixture +def mock_cache(): + """Mock cache for testing""" + return DualCache() + + +@pytest.fixture +def mock_pii_guardrail(): + """Mock PII guardrail that blocks""" + return MockPiiGuardrail(should_block=True) + + +@pytest.fixture +def mock_pii_guardrail_allow(): + """Mock PII guardrail that allows""" + return MockPiiGuardrail(should_block=False) + + +@pytest.fixture +def mock_content_guardrail(): + """Mock content guardrail that blocks""" + return MockContentGuardrail(should_block=True) + + +@pytest.fixture +def mock_http_guardrail(): + """Mock HTTP guardrail that blocks""" + return MockHttpGuardrail(should_block=True) + + +@pytest.fixture +def mock_during_guardrail(): + """Mock during-call guardrail that blocks""" + return MockDuringCallGuardrail(should_block=True) + + +@pytest.fixture +def mock_proxy_logging(): + """Mock proxy logging object""" + return MockProxyLogging() + + +class TestMCPGuardrailsPreCall: + """Test MCP guardrails for pre-call hooks""" + + @pytest.mark.asyncio + async def test_pii_guardrail_blocks_pre_call( + self, mock_pii_guardrail, mock_user_api_key, mock_cache + ): + """Test that PII guardrail properly blocks pre-call""" + proxy_logging = MockProxyLogging([mock_pii_guardrail]) + + # Create MCP request + request_obj = MCPPreCallRequestObject( + tool_name="email_tool", + arguments={"email": "test@example.com"}, + server_name="email_server", + user_api_key_auth=mock_user_api_key.model_dump(), + hidden_params=HiddenParams(), + ) + + kwargs = { + "name": "email_tool", + "arguments": {"email": "test@example.com"}, + "server_name": "email_server", + "user_api_key_auth": mock_user_api_key, + } + + # Test that BlockedPiiEntityError is raised + with pytest.raises(BlockedPiiEntityError) as excinfo: + await proxy_logging.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + # Verify the error details + assert excinfo.value.entity_type == "EMAIL_ADDRESS" + assert excinfo.value.guardrail_name == "mock-pii-guardrail" + assert mock_pii_guardrail.call_count == 1 + + @pytest.mark.asyncio + async def test_pii_guardrail_allows_pre_call( + self, mock_pii_guardrail_allow, mock_user_api_key, mock_cache + ): + """Test that PII guardrail allows pre-call when configured to allow""" + proxy_logging = MockProxyLogging([mock_pii_guardrail_allow]) + + request_obj = MCPPreCallRequestObject( + tool_name="email_tool", + arguments={"email": "test@example.com"}, + server_name="email_server", + user_api_key_auth=mock_user_api_key.model_dump(), + hidden_params=HiddenParams(), + ) + + kwargs = { + "name": "email_tool", + "arguments": {"email": "test@example.com"}, + "server_name": "email_server", + "user_api_key_auth": mock_user_api_key, + } + + # Test that no exception is raised + result = await proxy_logging.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is None + assert mock_pii_guardrail_allow.call_count == 1 + + @pytest.mark.asyncio + async def test_content_guardrail_blocks_pre_call( + self, mock_content_guardrail, mock_user_api_key, mock_cache + ): + """Test that content guardrail properly blocks pre-call""" + proxy_logging = MockProxyLogging([mock_content_guardrail]) + + request_obj = MCPPreCallRequestObject( + tool_name="content_tool", + arguments={"content": "sensitive content"}, + server_name="content_server", + user_api_key_auth=mock_user_api_key.model_dump(), + hidden_params=HiddenParams(), + ) + + kwargs = { + "name": "content_tool", + "arguments": {"content": "sensitive content"}, + "server_name": "content_server", + "user_api_key_auth": mock_user_api_key, + } + + # Test that GuardrailRaisedException is raised + with pytest.raises(GuardrailRaisedException) as excinfo: + await proxy_logging.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + # Verify the error details + assert "Content violates policy" in str(excinfo.value) + assert excinfo.value.guardrail_name == "mock-content-guardrail" + assert mock_content_guardrail.call_count == 1 + + @pytest.mark.asyncio + async def test_http_guardrail_blocks_pre_call( + self, mock_http_guardrail, mock_user_api_key, mock_cache + ): + """Test that HTTP guardrail properly blocks pre-call""" + proxy_logging = MockProxyLogging([mock_http_guardrail]) + + request_obj = MCPPreCallRequestObject( + tool_name="http_tool", + arguments={"url": "http://example.com"}, + server_name="http_server", + user_api_key_auth=mock_user_api_key.model_dump(), + hidden_params=HiddenParams(), + ) + + kwargs = { + "name": "http_tool", + "arguments": {"url": "http://example.com"}, + "server_name": "http_server", + "user_api_key_auth": mock_user_api_key, + } + + # Test that HTTPException is raised + with pytest.raises(HTTPException) as excinfo: + await proxy_logging.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + # Verify the error details + assert excinfo.value.status_code == 400 + assert "Violated guardrail policy" in str(excinfo.value.detail) + assert mock_http_guardrail.call_count == 1 + + @pytest.mark.asyncio + async def test_multiple_guardrails_pre_call( + self, mock_pii_guardrail, mock_content_guardrail, mock_user_api_key, mock_cache + ): + """Test multiple guardrails - first one should block""" + proxy_logging = MockProxyLogging([mock_pii_guardrail, mock_content_guardrail]) + + request_obj = MCPPreCallRequestObject( + tool_name="test_tool", + arguments={"email": "test@example.com"}, + server_name="test_server", + user_api_key_auth=mock_user_api_key.model_dump(), + hidden_params=HiddenParams(), + ) + + kwargs = { + "name": "test_tool", + "arguments": {"email": "test@example.com"}, + "server_name": "test_server", + "user_api_key_auth": mock_user_api_key, + } + + # Test that first guardrail blocks + with pytest.raises(BlockedPiiEntityError): + await proxy_logging.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + # Verify only first guardrail was called + assert mock_pii_guardrail.call_count == 1 + assert mock_content_guardrail.call_count == 0 + + +class TestMCPGuardrailsDuringCall: + """Test MCP guardrails for during-call hooks""" + + @pytest.mark.asyncio + async def test_during_call_guardrail_blocks( + self, mock_during_guardrail, mock_user_api_key, mock_cache + ): + """Test that during-call guardrail properly blocks execution""" + proxy_logging = MockProxyLogging([mock_during_guardrail]) + + request_obj = MCPDuringCallRequestObject( + tool_name="phone_tool", + arguments={"phone": "555-123-4567"}, + server_name="phone_server", + start_time=datetime.now().timestamp(), + hidden_params=HiddenParams(), + ) + + kwargs = { + "name": "phone_tool", + "arguments": {"phone": "555-123-4567"}, + "server_name": "phone_server", + } + + # Test that BlockedPiiEntityError is raised + with pytest.raises(BlockedPiiEntityError) as excinfo: + await proxy_logging.async_during_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + # Verify the error details + assert excinfo.value.entity_type == "PHONE_NUMBER" + assert excinfo.value.guardrail_name == "mock-during-guardrail" + assert mock_during_guardrail.call_count == 1 + + +class TestMCPGuardrailsIntegration: + """Test MCP guardrails integration with MCP server manager""" + + @pytest.mark.asyncio + async def test_mcp_server_manager_with_guardrails(self): + """Test MCP server manager with guardrail integration""" + + mock_proxy_logging = MockProxyLogging([MockPiiGuardrail(should_block=True)]) + + # Test that guardrail exception is properly raised in the hook + with pytest.raises(BlockedPiiEntityError): + await mock_proxy_logging.async_pre_mcp_tool_call_hook( + kwargs={ + "name": "email_tool", + "arguments": {"email": "test@example.com"}, + }, + request_obj=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + @pytest.mark.asyncio + async def test_guardrail_exception_propagation(self): + """Test that guardrail exceptions properly propagate through the system""" + # Test BlockedPiiEntityError + with pytest.raises(BlockedPiiEntityError): + raise BlockedPiiEntityError( + entity_type="EMAIL_ADDRESS", guardrail_name="test-guardrail" + ) + + # Test GuardrailRaisedException + with pytest.raises(GuardrailRaisedException): + raise GuardrailRaisedException( + guardrail_name="test-guardrail", message="Test message" + ) + + # Test HTTPException + with pytest.raises(HTTPException): + raise HTTPException(status_code=400, detail={"error": "Test error"}) + + +class TestMCPGuardrailsErrorHandling: + """Test MCP guardrails error handling scenarios""" + + @pytest.mark.asyncio + async def test_non_guardrail_exception_logging(self, mock_user_api_key, mock_cache): + """Test that non-guardrail exceptions are logged as non-blocking""" + + class MockFailingGuardrail(CustomGuardrail): + def should_run_guardrail( + self, data: dict, event_type: GuardrailEventHooks + ) -> bool: + return True + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: str, + ): + raise Exception("Non-guardrail error") + + proxy_logging = MockProxyLogging([MockFailingGuardrail()]) + + request_obj = MCPPreCallRequestObject( + tool_name="test_tool", + arguments={"test": "data"}, + server_name="test_server", + user_api_key_auth=mock_user_api_key.model_dump(), + hidden_params=HiddenParams(), + ) + + kwargs = { + "name": "test_tool", + "arguments": {"test": "data"}, + "server_name": "test_server", + "user_api_key_auth": mock_user_api_key, + } + + # Test that non-guardrail exceptions are handled gracefully + result = await proxy_logging.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + # Should return None (not raise exception) + assert result is None + + @pytest.mark.asyncio + async def test_guardrail_should_not_run(self, mock_user_api_key, mock_cache): + """Test that guardrails don't run when should_run_guardrail returns False""" + + class MockConditionalGuardrail(CustomGuardrail): + def should_run_guardrail( + self, data: dict, event_type: GuardrailEventHooks + ) -> bool: + return False # Don't run + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: str, + ): + raise BlockedPiiEntityError("EMAIL_ADDRESS", "test-guardrail") + + proxy_logging = MockProxyLogging([MockConditionalGuardrail()]) + + request_obj = MCPPreCallRequestObject( + tool_name="test_tool", + arguments={"test": "data"}, + server_name="test_server", + user_api_key_auth=mock_user_api_key.model_dump(), + hidden_params=HiddenParams(), + ) + + kwargs = { + "name": "test_tool", + "arguments": {"test": "data"}, + "server_name": "test_server", + "user_api_key_auth": mock_user_api_key, + } + + # Test that guardrail doesn't run and no exception is raised + result = await proxy_logging.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + # Should return None (guardrail didn't run) + assert result is None + + +class TestMCPGuardrailsEdgeCases: + """Test MCP guardrails edge cases and error conditions""" + + @pytest.mark.asyncio + async def test_empty_guardrails_list(self, mock_user_api_key, mock_cache): + """Test behavior with empty guardrails list""" + proxy_logging = MockProxyLogging([]) # No guardrails + + request_obj = MCPPreCallRequestObject( + tool_name="test_tool", + arguments={"test": "data"}, + server_name="test_server", + user_api_key_auth=mock_user_api_key.model_dump(), + hidden_params=HiddenParams(), + ) + + kwargs = { + "name": "test_tool", + "arguments": {"test": "data"}, + "server_name": "test_server", + "user_api_key_auth": mock_user_api_key, + } + + # Should return None without any issues + result = await proxy_logging.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is None + + @pytest.mark.asyncio + async def test_guardrail_with_invalid_data(self, mock_user_api_key, mock_cache): + """Test guardrail behavior with invalid data""" + + class MockInvalidDataGuardrail(CustomGuardrail): + def should_run_guardrail( + self, data: dict, event_type: GuardrailEventHooks + ) -> bool: + return True + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: str, + ): + # Try to access invalid data + invalid_data = data.get("invalid_key", {}) + if invalid_data.get("should_fail"): + raise BlockedPiiEntityError("EMAIL_ADDRESS", "test-guardrail") + return None + + proxy_logging = MockProxyLogging([MockInvalidDataGuardrail()]) + + request_obj = MCPPreCallRequestObject( + tool_name="test_tool", + arguments={"test": "data"}, + server_name="test_server", + user_api_key_auth=mock_user_api_key.model_dump(), + hidden_params=HiddenParams(), + ) + + kwargs = { + "name": "test_tool", + "arguments": {"test": "data"}, + "server_name": "test_server", + "user_api_key_auth": mock_user_api_key, + } + + # Should handle invalid data gracefully + result = await proxy_logging.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is None + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/mcp_tests/test_mcp_hooks.py b/tests/mcp_tests/test_mcp_hooks.py new file mode 100644 index 00000000000..6dac7da6d07 --- /dev/null +++ b/tests/mcp_tests/test_mcp_hooks.py @@ -0,0 +1,475 @@ +""" +Test file for MCP Hook Architecture + +This file demonstrates the new MCP hook system with comprehensive examples +and validation tests. +""" + +import asyncio +import pytest +from datetime import datetime +from typing import Optional + +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.mcp import ( + MCPPreCallRequestObject, + MCPPreCallResponseObject, + MCPDuringCallRequestObject, + MCPDuringCallResponseObject, + MCPPostCallResponseObject, +) +from litellm.types.llms.base import HiddenParams + + +class TestMCPAccessControlHook(CustomLogger): + """Test hook for access control functionality""" + + def __init__(self): + self.allowed_tools = {"github/create_issue", "zapier/send_email"} + self.blocked_users = {"user123", "user456"} + self.call_count = 0 + + async def async_pre_mcp_tool_call_hook( + self, kwargs, request_obj: MCPPreCallRequestObject, start_time, end_time + ) -> Optional[MCPPreCallResponseObject]: + """Test access control validation""" + self.call_count += 1 + + tool_name = request_obj.tool_name + user_id = kwargs.get("user_api_key_auth", {}).get("user_id") + + # Check if user is blocked + if user_id in self.blocked_users: + return MCPPreCallResponseObject( + should_proceed=False, + error_message=f"User {user_id} is not authorized to use MCP tools", + ) + + # Check if tool is allowed + if tool_name not in self.allowed_tools: + return MCPPreCallResponseObject( + should_proceed=False, + error_message=f"Tool {tool_name} is not authorized", + ) + + return None # Allow execution to proceed + + +class TestMCPCostTrackingHook(CustomLogger): + """Test hook for cost tracking functionality""" + + def __init__(self): + self.cost_map = { + "github/create_issue": 0.10, + "zapier/send_email": 0.05, + "default": 0.01, + } + self.call_count = 0 + + async def async_post_mcp_tool_call_hook( + self, kwargs, response_obj: MCPPostCallResponseObject, start_time, end_time + ) -> Optional[MCPPostCallResponseObject]: + """Test cost calculation after tool execution""" + self.call_count += 1 + + tool_name = kwargs.get("name", "") + cost = self.cost_map.get(tool_name, self.cost_map["default"]) + + # Set the response cost + response_obj.hidden_params.response_cost = cost + + return response_obj + + +class TestMCPMonitoringHook(CustomLogger): + """Test hook for real-time monitoring functionality""" + + def __init__(self): + self.max_execution_time = 30.0 # seconds + self.call_count = 0 + + async def async_during_mcp_tool_call_hook( + self, kwargs, request_obj: MCPDuringCallRequestObject, start_time, end_time + ) -> Optional[MCPDuringCallResponseObject]: + """Test execution time monitoring""" + self.call_count += 1 + + tool_name = request_obj.tool_name + execution_time = (datetime.now() - start_time).total_seconds() + + # Check if execution is taking too long + if execution_time > self.max_execution_time: + return MCPDuringCallResponseObject( + should_continue=False, + error_message=f"Tool {tool_name} execution timeout after {execution_time}s", + ) + + return None # Allow execution to continue + + +class TestMCPArgumentValidationHook(CustomLogger): + """Test hook for argument validation functionality""" + + def __init__(self): + self.call_count = 0 + + async def async_pre_mcp_tool_call_hook( + self, kwargs, request_obj: MCPPreCallRequestObject, start_time, end_time + ) -> Optional[MCPPreCallResponseObject]: + """Test argument validation and sanitization""" + self.call_count += 1 + + tool_name = request_obj.tool_name + arguments = request_obj.arguments.copy() # Create a copy to modify + + # Example: Validate GitHub issue creation + if tool_name == "github/create_issue": + if not arguments.get("title"): + return MCPPreCallResponseObject( + should_proceed=False, error_message="GitHub issue title is required" + ) + + # Sanitize the title + title = arguments["title"] + if len(title) > 100: + title = title[:97] + "..." + arguments["title"] = title + + # Example: Validate email sending + elif tool_name == "zapier/send_email": + if not arguments.get("to"): + return MCPPreCallResponseObject( + should_proceed=False, error_message="Email recipient is required" + ) + + return MCPPreCallResponseObject( + should_proceed=True, modified_arguments=arguments + ) + + +# Test fixtures +@pytest.fixture +def access_control_hook(): + return TestMCPAccessControlHook() + + +@pytest.fixture +def cost_tracking_hook(): + return TestMCPCostTrackingHook() + + +@pytest.fixture +def monitoring_hook(): + return TestMCPMonitoringHook() + + +@pytest.fixture +def argument_validation_hook(): + return TestMCPArgumentValidationHook() + + +# Test cases +class TestMCPHooks: + """Test cases for MCP hook functionality""" + + @pytest.mark.asyncio + async def test_access_control_hook_allowed_tool(self, access_control_hook): + """Test that allowed tools pass validation""" + kwargs = { + "user_api_key_auth": {"user_id": "user789"}, + "name": "github/create_issue", + } + request_obj = MCPPreCallRequestObject( + tool_name="github/create_issue", + arguments={"title": "Test issue"}, + user_api_key_auth={"user_id": "user789"}, + ) + + result = await access_control_hook.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is None # Should allow execution + assert access_control_hook.call_count == 1 + + @pytest.mark.asyncio + async def test_access_control_hook_blocked_user(self, access_control_hook): + """Test that blocked users are rejected""" + kwargs = { + "user_api_key_auth": {"user_id": "user123"}, + "name": "github/create_issue", + } + request_obj = MCPPreCallRequestObject( + tool_name="github/create_issue", + arguments={"title": "Test issue"}, + user_api_key_auth={"user_id": "user123"}, + ) + + result = await access_control_hook.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is not None + assert result.should_proceed is False + assert "not authorized" in result.error_message + + @pytest.mark.asyncio + async def test_access_control_hook_unauthorized_tool(self, access_control_hook): + """Test that unauthorized tools are rejected""" + kwargs = { + "user_api_key_auth": {"user_id": "user789"}, + "name": "unauthorized_tool", + } + request_obj = MCPPreCallRequestObject( + tool_name="unauthorized_tool", + arguments={"param": "value"}, + user_api_key_auth={"user_id": "user789"}, + ) + + result = await access_control_hook.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is not None + assert result.should_proceed is False + assert "not authorized" in result.error_message + + @pytest.mark.asyncio + async def test_cost_tracking_hook(self, cost_tracking_hook): + """Test cost tracking functionality""" + kwargs = {"name": "github/create_issue"} + response_obj = MCPPostCallResponseObject( + mcp_tool_call_response=[], hidden_params=HiddenParams() + ) + + result = await cost_tracking_hook.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is not None + assert result.hidden_params.response_cost == 0.10 + assert cost_tracking_hook.call_count == 1 + + @pytest.mark.asyncio + async def test_cost_tracking_hook_default_cost(self, cost_tracking_hook): + """Test default cost assignment""" + kwargs = {"name": "unknown_tool"} + response_obj = MCPPostCallResponseObject( + mcp_tool_call_response=[], hidden_params=HiddenParams() + ) + + result = await cost_tracking_hook.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is not None + assert result.hidden_params.response_cost == 0.01 # Default cost + + @pytest.mark.asyncio + async def test_monitoring_hook_normal_execution(self, monitoring_hook): + """Test monitoring hook with normal execution time""" + kwargs = {"name": "test_tool"} + request_obj = MCPDuringCallRequestObject( + tool_name="test_tool", arguments={}, start_time=datetime.now().timestamp() + ) + + result = await monitoring_hook.async_during_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is None # Should allow execution to continue + assert monitoring_hook.call_count == 1 + + @pytest.mark.asyncio + async def test_argument_validation_hook_valid_github_issue( + self, argument_validation_hook + ): + """Test argument validation for valid GitHub issue""" + kwargs = {"name": "github/create_issue"} + request_obj = MCPPreCallRequestObject( + tool_name="github/create_issue", arguments={"title": "Valid issue title"} + ) + + result = await argument_validation_hook.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is not None + assert result.should_proceed is True + assert result.modified_arguments == {"title": "Valid issue title"} + assert argument_validation_hook.call_count == 1 + + @pytest.mark.asyncio + async def test_argument_validation_hook_missing_title( + self, argument_validation_hook + ): + """Test argument validation for missing GitHub issue title""" + kwargs = {"name": "github/create_issue"} + request_obj = MCPPreCallRequestObject( + tool_name="github/create_issue", arguments={} # Missing title + ) + + result = await argument_validation_hook.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is not None + assert result.should_proceed is False + assert "title is required" in result.error_message + + @pytest.mark.asyncio + async def test_argument_validation_hook_long_title_sanitization( + self, argument_validation_hook + ): + """Test argument validation with title sanitization""" + kwargs = {"name": "github/create_issue"} + long_title = "A" * 150 # Very long title + request_obj = MCPPreCallRequestObject( + tool_name="github/create_issue", arguments={"title": long_title} + ) + + result = await argument_validation_hook.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is not None + assert result.should_proceed is True + assert len(result.modified_arguments["title"]) == 100 # Truncated + assert result.modified_arguments["title"].endswith("...") + + @pytest.mark.asyncio + async def test_argument_validation_hook_email_validation( + self, argument_validation_hook + ): + """Test argument validation for email sending""" + kwargs = {"name": "zapier/send_email"} + request_obj = MCPPreCallRequestObject( + tool_name="zapier/send_email", + arguments={"to": "test@example.com", "subject": "Test"}, + ) + + result = await argument_validation_hook.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is not None + assert result.should_proceed is True + assert result.modified_arguments == { + "to": "test@example.com", + "subject": "Test", + } + + @pytest.mark.asyncio + async def test_argument_validation_hook_missing_email_recipient( + self, argument_validation_hook + ): + """Test argument validation for missing email recipient""" + kwargs = {"name": "zapier/send_email"} + request_obj = MCPPreCallRequestObject( + tool_name="zapier/send_email", + arguments={"subject": "Test"}, # Missing 'to' field + ) + + result = await argument_validation_hook.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert result is not None + assert result.should_proceed is False + assert "recipient is required" in result.error_message + + +# Integration test +class TestMCPHookIntegration: + """Integration tests for MCP hook system""" + + @pytest.mark.asyncio + async def test_hook_chain_execution(self): + """Test that multiple hooks can work together""" + access_hook = TestMCPAccessControlHook() + cost_hook = TestMCPCostTrackingHook() + validation_hook = TestMCPArgumentValidationHook() + + # Test data + kwargs = { + "user_api_key_auth": {"user_id": "user789"}, + "name": "github/create_issue", + } + request_obj = MCPPreCallRequestObject( + tool_name="github/create_issue", + arguments={"title": "Integration test issue"}, + user_api_key_auth={"user_id": "user789"}, + ) + + # Execute pre-hooks + access_result = await access_hook.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + validation_result = await validation_hook.async_pre_mcp_tool_call_hook( + kwargs=kwargs, + request_obj=request_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + # Both hooks should allow execution + assert access_result is None + assert validation_result is not None + assert validation_result.should_proceed is True + + # Simulate post-hook execution + response_obj = MCPPostCallResponseObject( + mcp_tool_call_response=[], hidden_params=HiddenParams() + ) + + cost_result = await cost_hook.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert cost_result is not None + assert cost_result.hidden_params.response_cost == 0.10 + + +if __name__ == "__main__": + # Run the tests + pytest.main([__file__, "-v"]) diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index e8caa241a53..f3c68b489a5 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -11,6 +11,7 @@ from datetime import datetime, timedelta, timezone from typing import Final import pytest +from prisma import Prisma from litellm.proxy.db.autorouter_session_rollup import ( AUTOROUTER_BENCHMARKS_SQL, @@ -43,6 +44,7 @@ async def _turn( classifier_cost: float = 0.0, tier: "str | None" = None, baseline: "str | None" = None, + estimated: bool = True, ) -> None: touched: Final = 1 if (hit or ttl is not None or not covered) else 0 await db.execute_raw( @@ -63,6 +65,9 @@ async def _turn( touched, tier, baseline, + int(estimated), + spend if estimated else 0.0, + saved if estimated else 0.0, ) @@ -208,6 +213,9 @@ async def test_subtotal_coverage_survives_legacy_and_rolling_writers(db, writers assert row["saved_spend"] == pytest.approx(0.02 * len(writers)) assert row["classifier_cost"] == pytest.approx(0.004 * sum(writers)) assert row["classifier_cost_recorded_turns"] == sum(writers) + assert row["savings_estimated_turns"] == sum(writers) + assert row["savings_estimated_actual_spend"] == pytest.approx(0.01 * sum(writers)) + assert row["savings_estimated_saved_spend"] == pytest.approx(0.02 * sum(writers)) groups: Final = await db.query_raw( AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key ) @@ -217,6 +225,32 @@ async def test_subtotal_coverage_survives_legacy_and_rolling_writers(db, writers assert groups[0]["turns"] == len(writers) assert groups[0]["spend"] == row["spend"] assert groups[0]["saved_spend"] == row["saved_spend"] + assert groups[0]["savings_estimated_turns"] == sum(writers) + assert groups[0]["savings_estimated_actual_spend"] == row["savings_estimated_actual_spend"] + assert groups[0]["savings_estimated_saved_spend"] == row["savings_estimated_saved_spend"] + + +async def test_unknown_and_legacy_turns_preserve_actual_spend_without_entering_the_estimated_cohort(db: Prisma) -> None: + key: Final = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, spend=0.25, saved=-0.05, baseline="opus") + await _turn( + db, key, "B", T0 + timedelta(seconds=1), spend=0.7, saved=0, baseline="sonnet", estimated=False + ) + await _legacy_turn(db, key, T0 + timedelta(seconds=2)) + + row: Final = await _row(db, key) + assert row["saved_spend"] == pytest.approx(-0.03) + assert row["savings_estimated_baseline_models"] == {"opus": 1} + groups: Final = await db.query_raw( + AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key + ) + assert len(groups) == 1 + for actual in (row, groups[0]): + assert actual["turns"] == 3 + assert actual["spend"] == pytest.approx(0.96) + assert actual["savings_estimated_turns"] == 1 + assert actual["savings_estimated_actual_spend"] == pytest.approx(0.25) + assert actual["savings_estimated_saved_spend"] == pytest.approx(-0.05) async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): diff --git a/tests/proxy_behavior/spend/test_baseline_accounting.py b/tests/proxy_behavior/spend/test_baseline_accounting.py new file mode 100644 index 00000000000..e187a44c29d --- /dev/null +++ b/tests/proxy_behavior/spend/test_baseline_accounting.py @@ -0,0 +1,263 @@ +import asyncio +import json +import uuid +from collections.abc import AsyncIterator, Callable +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from typing import Final + +import pytest +from prisma import Prisma + +import litellm +from litellm.llms.anthropic.prompt_cache_prediction import CountedBreakpoint, CountedPromptCachePlan +from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction +from litellm.proxy.db.baseline_accounting import ( + BaselineAccountingRecord, + BaselineAccountingStore, + DailyBaselineAttribution, + DailyBaselineTarget, +) +from litellm.proxy.db.create_views import SupportsRawQueries +from litellm.proxy.spend_tracking.baseline_accounting import BaselineObservation +from litellm.proxy.spend_tracking.savings import BaselineCostSnapshot +from litellm.types.utils import Usage + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +@asynccontextmanager +async def _transaction(db: Prisma, *, before_commit: bool = False, after_commit: bool = False) -> AsyncIterator[SupportsRawQueries]: + async with db.tx() as tx: + yield tx + if before_commit: + raise RuntimeError("injected pre-commit interruption") + if after_commit: + raise RuntimeError("injected lost commit acknowledgement") + + +def _store(db: Prisma, **faults: bool) -> BaselineAccountingStore: + def transaction(): + return _transaction(db, **faults) + + return BaselineAccountingStore(transaction) + + +@pytest.fixture +def record() -> Callable[..., BaselineAccountingRecord]: + run: Final = uuid.uuid4().hex + marker: Final = CountedBreakpoint("prefix", 3600, 6000, ("prefix",), "content", ("content",)) + usage: Final = Usage( + prompt_tokens=6200, completion_tokens=30, total_tokens=6230, + cache_creation_input_tokens=6000, cache_read_input_tokens=0, + prompt_tokens_details={ + "text_tokens": 200, "cached_tokens": 0, "cache_creation_tokens": 6000, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 6000}, + }, + ) + + def create(label: str = "first", started: float = 10000.0, identical: bool = True) -> BaselineAccountingRecord: + return BaselineAccountingRecord( + scope="autorouter-baseline:v3:" + run * 2, api_key=run, session_id=run, + router_name="test-router", baseline_model="anthropic/claude-opus-5", + observation=BaselineObservation( + request_id=run + label, started_at=started, available_at=started + 0.1, + outcome="complete", baseline_equivalent=identical, usage=usage, + plan=CountedPromptCachePlan(6200, (marker,)), minimum_cache_tokens=4096, + ), + pricing=BaselineCostSnapshot( + model="claude-opus-5", provider="anthropic", + prices=litellm.get_model_info("claude-opus-5", custom_llm_provider="anthropic"), + actual_spend=0.17, actual_token_cost=0.17, + ), + turn=AutoRouterTurnTransaction( + api_key=run, session_id=run, router_name="test-router", router_type="heuristic", + model="claude-opus-5", turn_at=datetime.fromtimestamp(started, timezone.utc), + total_tokens=6230, spend=0.17, saved_spend=0.0, classifier_cost=0.0, + covered=True, cache_hit=False, cache_ttl_seconds=3600, cache_touched=True, + baseline_model="anthropic/claude-opus-5", + ), + daily=DailyBaselineAttribution( + date="2026-09-15", api_key=run, model="claude-opus-5", custom_llm_provider="anthropic", + targets=tuple(DailyBaselineTarget(entity=entity, entity_id=run) for entity in ("user", "team", "org", "end_user", "agent", "tag")), + ), + ) + + return create + + +async def _log(db: Prisma, record: BaselineAccountingRecord) -> None: + await db.execute_raw( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id,call_type,api_key,spend,"startTime","endTime") ' + "VALUES ($1, 'anthropic_messages', $2, 0.17, to_timestamp($3::float8), to_timestamp($3::float8))", + record.observation.request_id, record.api_key, record.observation.started_at, + ) + + +async def _session(db: Prisma, record: BaselineAccountingRecord): + rows: Final = await db.query_raw('SELECT * FROM "LiteLLM_AutoRouterSession" WHERE api_key=$1', record.api_key) + return rows[0] + + +async def test_late_replay_updates_all_projections_without_rebilling(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: + store: Final = _store(db) + late: Final = record("late", 10001.0) + early: Final = record("early", identical=False) + await _log(db, late) + assert await store.append(late) == "recorded" + assert await store.project(late.scope) == "published" + before: Final = await _session(db, late) + assert before["savings_estimated_actual_spend"] == before["spend"] == 0.17 + assert before["saved_spend"] == 0.0 + await _log(db, early) + assert await store.append(early) == "recorded" + pending: Final = await _session(db, late) + assert pending["spend"] == 0.34 and pending["savings_estimated_turns"] == 0 + assert pending["saved_spend"] == pending["savings_estimated_actual_spend"] == 0.0 + waiting: Final = await db.query_raw('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=$1', late.observation.request_id) + assert waiting[0]["metadata"]["autorouter_savings"] is None + assert waiting[0]["metadata"]["autorouter_savings_estimate"]["reason"] == "pending_projection" + assert await store.project(early.scope) == "published" + after: Final = await _session(db, late) + assert after["spend"] == 0.34 and after["turns"] == 2 + assert after["savings_estimated_actual_spend"] == 0.17 and after["savings_estimated_turns"] == 1 + logs: Final = await db.query_raw('SELECT spend, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=$1', late.observation.request_id) + assert logs[0]["spend"] == 0.17 + assert logs[0]["metadata"]["autorouter_savings_estimate"]["provenance"] == "modeled" + assert after["saved_spend"] == pytest.approx(logs[0]["metadata"]["autorouter_savings"]) + for table in ("DailyUserSpend", "DailyTeamSpend", "DailyOrganizationSpend", "DailyEndUserSpend", "DailyAgentSpend", "DailyTagSpend"): + rows: Final = await db.query_raw(f'SELECT spend,api_requests,autorouter_savings_spend FROM "LiteLLM_{table}" WHERE api_key=$1', late.api_key) + assert rows[0]["spend"] == rows[0]["api_requests"] == 0 + assert rows[0]["autorouter_savings_spend"] == pytest.approx(after["saved_spend"]) + + +async def test_commit_ack_loss_and_concurrent_duplicate_delivery_are_idempotent(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: + event: Final = record() + await _log(db, event) + assert await _store(db, after_commit=True).append(event) == "unavailable" + store: Final = _store(db) + assert set(await asyncio.gather(*(store.append(event) for _ in range(4)))) == {"recorded"} + assert await store.project(event.scope) == "published" + assert await store.project(event.scope) == "unchanged" + session: Final = await _session(db, event) + assert session["turns"] == session["savings_estimated_turns"] == 1 + assert session["spend"] == session["savings_estimated_actual_spend"] == 0.17 + + +async def test_publication_rollback_keeps_dirty_revision_for_retry(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: + event: Final = record() + await _log(db, event) + store: Final = _store(db) + assert await store.append(event) == "recorded" + assert await _store(db, before_commit=True).project(event.scope) == "unavailable" + session: Final = await _session(db, event) + assert session["spend"] == 0.17 and session["savings_estimated_turns"] == 0 + revisions: Final = await db.query_raw('SELECT revision,published_revision FROM "LiteLLM_AutoRouterBaselineComparison" WHERE scope=$1', event.scope) + assert revisions[0]["revision"] > revisions[0]["published_revision"] + assert await store.project(event.scope) == "published" + assert (await _session(db, event))["savings_estimated_turns"] == 1 + + +async def test_conflicting_duplicate_cannot_restore_an_observed_estimate(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: + event: Final = record() + await _log(db, event) + store: Final = _store(db) + assert await store.append(event) == "recorded" + assert await store.project(event.scope) == "published" + conflict: Final = event.model_copy(update={"observation": event.observation.model_copy(update={"baseline_equivalent": False, "started_at": 20000.0, "available_at": 20001.0})}) + assert await store.append(conflict) == "recorded" + assert (await _session(db, event))["savings_estimated_turns"] == 0 + assert await store.append(event) == "recorded" + assert await store.project(event.scope) == "published" + session: Final = await _session(db, event) + assert session["turns"] == 1 and session["savings_estimated_turns"] == 0 + rows: Final = await db.query_raw('SELECT publication FROM "LiteLLM_AutoRouterBaselineObservation" WHERE request_id=$1', event.observation.request_id) + assert json.loads(rows[0]["publication"])["reason"] == "conflicting_observation" + + +async def test_retired_history_never_recreates_an_initial_zero(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: + original: Final = record() + await _log(db, original) + store: Final = _store(db) + assert await store.append(original) == "recorded" + assert await store.project(original.scope) == "published" + await db.execute_raw('UPDATE "LiteLLM_AutoRouterBaselineComparison" SET updated_at=to_timestamp(0) WHERE scope=$1', original.scope) + await store.retire_before(datetime(2000, 1, 1, tzinfo=timezone.utc), 1000, 1000) + next_turn: Final = record("after-retention", 20000.0) + await _log(db, next_turn) + assert await store.append(next_turn) == "retired" + assert await store.project(original.scope) == "unchanged" + after: Final = await _session(db, original) + assert after["turns"] == 2 and after["spend"] == 0.34 + assert after["savings_estimated_turns"] == 1 and after["savings_estimated_actual_spend"] == 0.17 + + +async def test_native_observation_enters_spend_pipeline_once_with_shared_daily_attribution( + db: Prisma, record: Callable[..., BaselineAccountingRecord], monkeypatch: pytest.MonkeyPatch, +) -> None: + import os + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + from litellm.proxy.hooks.autorouter_baseline_cache import CapturedBaselineObservation + from litellm.proxy.utils import PrismaClient, ProxyLogging + + event: Final = record("routed", identical=False) + capture: Final = CapturedBaselineObservation( + scope=event.scope, api_key=event.api_key, session_id=event.session_id, + router_name=event.router_name, baseline_model=event.baseline_model, + model=event.pricing.model, prices=event.pricing.prices, observation=event.observation, + ) + metadata: Final = { + "routing_decision": {"router_model_name": event.router_name, "savings_baseline_model": event.baseline_model}, + "usage_object": event.observation.usage.model_dump(), + "cost_breakdown": {"input_cost": 0.16, "output_cost": 0.01}, + "autorouter_savings": None, "autorouter_savings_estimate": {"version": 3, "status": "unknown", "reason": "pending_projection"}, + "autorouter_baseline_observation": capture.model_dump_json(), + } + payload: Final = { + "request_id": event.observation.request_id, "api_key": event.api_key, "session_id": event.session_id, + "startTime": datetime.fromtimestamp(event.observation.started_at, timezone.utc).isoformat(), + "endTime": datetime.fromtimestamp(event.observation.available_at, timezone.utc).isoformat(), + "spend": 0.17, "prompt_tokens": 6200, "completion_tokens": 30, "model": event.pricing.model, + "model_group": event.router_name, "model_id": "baseline", "custom_llm_provider": "anthropic", + "call_type": "anthropic_messages", "status": "success", "metadata": json.dumps(metadata), + "user": None, "team_id": "", "organization_id": "org", "agent_id": None, + "end_user": "", "request_tags": '["tag","tag"]', + } + monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) + client: Final = PrismaClient(os.environ["DATABASE_URL"], ProxyLogging(UserApiKeyCache())) + writer: Final = DBSpendUpdateWriter() + try: + await client.db.connect() + await _log(db, event) + await writer._enqueue_autorouter_turn_transaction(payload, client) + assert len(client.baseline_accounting_transactions) == 1 + queued: Final = client.baseline_accounting_transactions[0] + assert queued.daily is not None + assert [(target.entity, target.entity_id) for target in queued.daily.targets] == [ + ("user", None), ("team", ""), ("org", "org"), ("tag", "tag"), + ] + await writer.add_spend_log_transaction_to_daily_tag_transaction(payload, client) + actual_tags: Final = await writer.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + assert len(actual_tags) == 1 + assert next(iter(actual_tags.values()))["spend"] == 0.17 + durable: Final = BaselineAccountingStore.for_client(client) + anchor: Final = record("anchor", 9999.0) + await _log(db, anchor) + assert await durable.append(anchor) == "recorded" + assert await durable.append(queued) == "recorded" + assert await durable.append(queued) == "recorded" + assert await durable.project(queued.scope) == "published" + session: Final = await _session(db, queued) + assert session["turns"] == session["savings_estimated_turns"] == 2 + assert session["spend"] == session["savings_estimated_actual_spend"] == 0.34 + tag_rows: Final = await db.query_raw( + 'SELECT spend, api_requests, autorouter_savings_spend FROM "LiteLLM_DailyTagSpend" WHERE api_key=$1 AND tag=$2', + queued.api_key, "tag", + ) + assert session["saved_spend"] < 0 + assert len(tag_rows) == 1 + assert tag_rows[0]["autorouter_savings_spend"] == pytest.approx(session["saved_spend"]) + assert tag_rows[0]["spend"] == tag_rows[0]["api_requests"] == 0 + finally: + await client.db.disconnect() diff --git a/tests/proxy_migration_tests/test_autorouter_baseline_state.py b/tests/proxy_migration_tests/test_autorouter_baseline_state.py new file mode 100644 index 00000000000..d9021414bc4 --- /dev/null +++ b/tests/proxy_migration_tests/test_autorouter_baseline_state.py @@ -0,0 +1,103 @@ +"""Idempotent journal migration and primary transactional ownership.""" + +import asyncio +import os +import time +from collections.abc import AsyncGenerator, Iterator +from contextlib import asynccontextmanager +from datetime import timedelta +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import psycopg +import pytest +from psycopg import sql + +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db.baseline_accounting import BaselineAccountingStore +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper +from litellm.proxy.utils import PrismaClient, ProxyLogging + +_MIGRATION: Final = Path(__file__).parents[2] / ( + "litellm-proxy-extras/litellm_proxy_extras/migrations/20260915010000_add_autorouter_baseline_state/migration.sql" +) + + +@pytest.fixture +def database() -> Iterator[tuple[str, psycopg.Connection[tuple[object, ...]]]]: + base: Final = os.environ["DATABASE_URL"].split("?")[0] + schema: Final = f"baseline_{uuid4().hex}" + with psycopg.connect(base, autocommit=True) as connection: + connection.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema))) + connection.execute(sql.SQL("SET search_path TO {}").format(sql.Identifier(schema))) + try: + connection.execute(_MIGRATION.read_bytes()) + connection.execute(_MIGRATION.read_bytes()) + yield f"{base}?schema={schema}", connection + finally: + connection.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema))) + + +@asynccontextmanager +async def _client(env: pytest.MonkeyPatch, url: str, replica: str | None = None) -> AsyncGenerator[PrismaClient]: + with env.context() as context: + context.setenv("DATABASE_URL", url) + context.delenv("DATABASE_URL_READ_REPLICA", raising=False) + if replica is not None: + context.setenv("DATABASE_URL_READ_REPLICA", replica) + client: Final = PrismaClient(url, ProxyLogging(UserApiKeyCache())) + try: + await client.db.connect(timeout=timedelta(seconds=1)) + yield client + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_migration_and_projector_use_the_primary_across_clients( + database: tuple[str, psycopg.Connection[tuple[object, ...]]], monkeypatch: pytest.MonkeyPatch, +) -> None: + url, connection = database + connection.execute('CREATE TABLE "LiteLLM_SpendLogs" (request_id TEXT PRIMARY KEY)') + connection.execute('INSERT INTO "LiteLLM_AutoRouterBaselineComparison" ' + '(scope,api_key,session_id,router_name,initial_equivalent,revision) ' + "VALUES ('test','key','session','router',TRUE,1)") + async with _client(monkeypatch, url, url.split("?")[0]) as first: + assert await BaselineAccountingStore.for_client(first).project("test") == "published" + async with _client(monkeypatch, url) as restarted: + assert await BaselineAccountingStore.for_client(restarted).project("test") == "unchanged" + assert connection.execute('SELECT revision=published_revision FROM "LiteLLM_AutoRouterBaselineComparison"').fetchone() == (True,) + + +@pytest.mark.asyncio +async def test_primary_outage_and_missing_table_are_unavailable( + database: tuple[str, psycopg.Connection[tuple[object, ...]]], monkeypatch: pytest.MonkeyPatch, +) -> None: + url, connection = database + async with _client(monkeypatch, "postgresql://unused:unused@127.0.0.1:1/unreachable", url) as degraded: + assert isinstance(degraded.db, RoutingPrismaWrapper) and degraded.db.writer_unavailable + assert await BaselineAccountingStore.for_client(degraded).project("scope") == "unavailable" + connection.execute('DROP TABLE "LiteLLM_AutoRouterBaselineComparison"') + async with _client(monkeypatch, url) as missing: + assert await BaselineAccountingStore.for_client(missing).project("scope") == "unavailable" + + +@pytest.mark.asyncio +async def test_locked_projection_is_bounded_and_cancellation_propagates( + database: tuple[str, psycopg.Connection[tuple[object, ...]]], monkeypatch: pytest.MonkeyPatch, +) -> None: + url, connection = database + async with _client(monkeypatch, url) as client: + store: Final = BaselineAccountingStore.for_client(client) + with connection.transaction(): + connection.execute('LOCK TABLE "LiteLLM_AutoRouterBaselineComparison" IN ACCESS EXCLUSIVE MODE') + started: Final = time.monotonic() + assert await store.project("scope") == "unavailable" + assert time.monotonic() - started < 2 + pending: Final = asyncio.create_task(store.project("scope")) + await asyncio.sleep(0.01) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + assert await store.project("scope") == "unchanged" diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/proxy_unit_tests/test_update_spend.py index a28a78cc4a1..0b158c33c73 100644 --- a/tests/proxy_unit_tests/test_update_spend.py +++ b/tests/proxy_unit_tests/test_update_spend.py @@ -37,10 +37,15 @@ class MockPrismaClient: self.daily_user_spend_transactions = {} self.tool_usage_transactions = [] self.autorouter_turn_transactions = [] + self.baseline_accounting_transactions = [] + self.baseline_accounting_lock = asyncio.Lock() + self.spend_log_flush_requested = None + self.db.tx = MagicMock() + self.db.tx.return_value.__aenter__ = AsyncMock(return_value=self.db) + self.db.tx.return_value.__aexit__ = AsyncMock(return_value=None) + self.db.query_raw.return_value = [] # Add locks for the transaction queues (matches real PrismaClient) - import asyncio - self._spend_log_transactions_lock = asyncio.Lock() self._tool_usage_transactions_lock = asyncio.Lock() self._autorouter_turn_transactions_lock = asyncio.Lock() diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index b78d61c7bd4..7e4598c2e58 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -6,7 +6,7 @@ import sys from collections.abc import AsyncIterator from pathlib import Path from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import anyio import httpx2 @@ -18,12 +18,14 @@ from mcp.types import ( CONNECTION_CLOSED, INTERNAL_ERROR, REQUEST_TIMEOUT, + CallToolRequestParams, CallToolResult, ErrorData, Implementation, InitializeResult, JSONRPCError, JSONRPCMessage, + JSONRPCRequest, JSONRPCResponse, LoggingMessageNotificationParams, ServerCapabilities, @@ -61,8 +63,10 @@ class _MockTransportClient(MCPClient): super().__init__(**kwargs) self._respond = respond - def _create_transport_context(self): - http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(self._respond)) + def _create_transport_context(self) -> tuple[_TransportContext, httpx2.AsyncClient]: + http_client: Final = self._create_httpx_client_factory(transport=httpx2.MockTransport(self._respond))( + headers=self._get_auth_headers(), timeout=httpx2.Timeout(self.timeout) + ) return streamable_http_client(self.server_url, http_client=http_client), http_client @@ -1178,6 +1182,107 @@ def test_v1_static_headers_still_win_their_own_slot(): assert headers["Authorization"] == "Bearer static-upstream-mcp-token" +@pytest.mark.asyncio +async def test_sdk_same_origin_redirect_lists_and_calls_tools() -> None: + def respond(request: httpx2.Request) -> httpx2.Response: + if request.url.path == "/mcp": + return httpx2.Response(307, headers={"Location": "/final/mcp"}) + assert request.url == "https://upstream.example.com/final/mcp" + assert request.headers["x-upstream-token"] == "Bearer synthetic-token" + if request.method != "POST": + return httpx2.Response(405) + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + match payload.method: + case "initialize": + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": LATEST_HANDSHAKE_VERSION, + "capabilities": {"tools": {}}, + "serverInfo": {"name": "redirect-test", "version": "1"}, + }, + }, + ) + case "tools/list": + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": {"tools": [{"name": "add", "inputSchema": {"type": "object"}}]}, + }, + ) + case "tools/call": + assert payload.params is not None + assert payload.params["name"] == "add" + assert payload.params["arguments"] == {"a": 2, "b": 3} + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": {"content": [{"type": "text", "text": "5"}], "isError": False}, + }, + ) + case _: + pytest.fail(f"Unexpected MCP request: {payload.method}") + + responder: Final = Mock(side_effect=respond) + client: Final = _MockTransportClient( + responder, + server_url="https://upstream.example.com/mcp", + auth_type=MCPAuth.bearer_token, + auth_value="synthetic-token", + auth_header_name="x-upstream-token", + timeout=5, + ) + with anyio.fail_after(10): + tools: Final = await client.list_tools(raise_on_error=True) + result: Final = await client.call_tool( + CallToolRequestParams(name="add", arguments={"a": 2, "b": 3}), raise_on_error=True + ) + assert [tool.name for tool in tools] == ["add"] + assert result.is_error is False + assert len(result.content) == 1 + assert result.content[0].type == "text" + assert result.content[0].text == "5" + assert any(call.args[0].url.path == "/mcp" for call in responder.call_args_list) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ("list", "call")) +async def test_sdk_cross_origin_redirect_never_contacts_destination(operation: str) -> None: + responder: Final = Mock( + return_value=httpx2.Response(307, headers={"Location": "https://destination.example.com/mcp"}) + ) + client: Final = _MockTransportClient( + responder, + server_url="https://upstream.example.com/mcp", + auth_type=MCPAuth.bearer_token, + auth_value="synthetic-token", + auth_header_name="x-upstream-token", + timeout=5, + ) + pending_operation: Final = ( + client.list_tools(raise_on_error=True) + if operation == "list" + else client.call_tool(CallToolRequestParams(name="add", arguments={"a": 2, "b": 3}), raise_on_error=True) + ) + with anyio.fail_after(10), pytest.raises(MCPError): + await pending_operation + assert responder.call_count == 1 + request: Final = responder.call_args.args[0] + assert request.method == "POST" + assert request.url == "https://upstream.example.com/mcp" + assert request.headers["x-upstream-token"] == "Bearer synthetic-token" + assert all(call.args[0].url.host != "destination.example.com" for call in responder.call_args_list) + + @pytest.mark.asyncio async def test_a_custom_credential_header_is_stripped_when_a_redirect_crosses_origin(): """httpx drops Authorization across origins but keeps every other header, so a credential the diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 5e179f950a0..633dd1d9460 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -124,22 +124,32 @@ def test_calculate_usage_prefers_served_speed_from_response_usage(): assert no_response_speed.speed == "fast" -def test_streaming_iterator_persists_served_speed_across_usage_chunks(): +@pytest.mark.parametrize("input_update, expected_fresh", [({}, 1000), ({"input_tokens": 0}, 0), ({"input_tokens": 2000}, 2000)]) +def test_streaming_iterator_persists_cumulative_usage_across_partial_chunks(input_update, expected_fresh): """ - Only ``message_start`` usage carries the served speed; the final - ``message_delta`` usage does not. The iterator must remember the served - value so the last usage chunk, which wins in the stream chunk builder, does - not fall back to the requested speed. + Omitted input/cache/pricing fields retain their last cumulative values; + explicit input updates, including zero, replace them. """ from litellm.llms.anthropic.chat.handler import ModelResponseIterator iterator = ModelResponseIterator(None, sync_stream=True, speed="fast") - start_usage = iterator._handle_usage({"input_tokens": 12, "output_tokens": 1, "speed": "standard"}) - delta_usage = iterator._handle_usage({"output_tokens": 5}) + start_usage = iterator._handle_usage({ + "input_tokens": 1000, "output_tokens": 1, "speed": "standard", "inference_geo": "us", + "cache_creation_input_tokens": 3000, "cache_read_input_tokens": 2000, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 3000}, + }) + delta_usage = iterator._handle_usage({"output_tokens": 5, **input_update}) assert start_usage.speed == "standard" assert delta_usage.speed == "standard" + assert delta_usage.inference_geo == "us" + assert delta_usage.prompt_tokens == expected_fresh + 5000 + assert delta_usage.completion_tokens == 5 + details = delta_usage.prompt_tokens_details + assert (details.text_tokens, details.cached_tokens, details.cache_creation_tokens) == (expected_fresh, 2000, 3000) + assert details.cache_creation_token_details.ephemeral_1h_input_tokens == 3000 + assert start_usage.prompt_tokens_details.text_tokens == 1000 def test_calculate_usage_aggregates_cache_creation_split_across_iterations(): diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py b/tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py index 2b36866a1a0..62099f97b71 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py @@ -15,11 +15,17 @@ from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.anthropic.count_tokens import handler as count_handler from litellm.llms.anthropic.experimental_pass_through.messages.transformation import DEFAULT_ANTHROPIC_API_VERSION from litellm.llms.anthropic.prompt_cache_prediction import ( + CountedPromptCachePlan, NativePredictionTarget, + PromptCachePlan, + UnsupportedCachePlan, cache_scope, + count_cache_plan, count_prompt_tokens, + parse_cache_plan, parse_observed_cache, parse_prompt, + resolve_baseline_prediction_target, resolve_prediction_target, supported_prediction_headers, ) @@ -207,3 +213,232 @@ async def test_named_credential_is_explicitly_unsupported_before_count( assert arm.cache_state == "unknown" assert arm.reason == "unsupported_deployment_configuration" assert arm.estimate is None and arm.cold is None and arm.warm is None + + +def _cache_plan(body: Mapping[str, JsonValue]) -> PromptCachePlan: + plan: Final = parse_cache_plan(body) + assert isinstance(plan, PromptCachePlan) + return plan + + +def _text(text: str, ttl: str | None = None) -> dict[str, JsonValue]: + return {"type": "text", "text": text, + **({"cache_control": {"type": "ephemeral", "ttl": ttl}} if ttl else {})} + + +def _prompt(*blocks: dict[str, JsonValue], role: str = "user", **options: JsonValue) -> dict[str, JsonValue]: + return {**options, "messages": [{"role": role, "content": list(blocks)}]} + + +@pytest.mark.parametrize("text, supported", [("", False), (" \t", False), ("Context", True)]) +def test_public_predictor_preserves_string_message_policy(text: str, supported: bool) -> None: + body: Final = _body() + messages: Final = body["messages"] + assert isinstance(messages, list) + request: Final[dict[str, JsonValue]] = {**body, "messages": [{"role": "user", "content": text}, *messages]} + assert (parse_prompt(request) is not None) is supported + + +def test_cache_plan_preserves_hierarchical_prefixes_and_public_policy() -> None: + body: Final = _prompt( + _text("First turn", "5m"), system=[_text("Stable instructions", "1h")], + tools=[{"name": "lookup", "input_schema": {"type": "object"}, + "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + ) + plan: Final = _cache_plan(body) + changed: Final = _cache_plan({**body, "system": [_text("Changed instructions", "1h")]}) + assert tuple(marker.ttl_seconds for marker in plan.breakpoints) == (3600, 3600, 300) + assert plan.breakpoints[0].fingerprint == changed.breakpoints[0].fingerprint + assert all(left.fingerprint != right.fingerprint for left, right + in zip(plan.breakpoints[1:], changed.breakpoints[1:])) + assert plan.breakpoints[0].prefix_body == { + "tools": [{"name": "lookup", "input_schema": {"type": "object"}}], + "messages": [], + } + assert parse_prompt(body) is None + + +@pytest.mark.parametrize("kind, added, matches", [ + ("text", 19, True), ("text", 20, False), ("tool_use", 30, True), + ("tool_result", 30, True), +]) +def test_cache_plan_lookback_counts_native_positions( + kind: str, added: int, matches: bool, +) -> None: + previous: Final = _cache_plan(_body()) + appended: Final[list[dict[str, JsonValue]]] = [ + {"type": "tool_use", "id": f"tool_{index}", "name": "lookup", "input": {}} + if kind == "tool_use" else + {"type": "tool_result", "tool_use_id": f"tool_{index}", "content": "done"} + if kind == "tool_result" else + {"type": "text", "text": f"Added {index}"} + for index in range(added) + ] + current: Final = _cache_plan({**_body(), **_prompt( + _text("A cacheable prefix"), *appended[:-1], + {**appended[-1], "cache_control": {"type": "ephemeral"}}, + )}) + assert (previous.breakpoints[0].fingerprint + in current.breakpoints[0].lookback_fingerprints) is matches + + +@pytest.mark.parametrize("change, same_prefix, same_content", [ + ("tool_order", False, False), ("effort", False, False), + ("standard_speed", True, True), ("ttl", False, True), +]) +def test_cache_plan_identity_respects_settings_and_preserves_content( + change: str, same_prefix: bool, same_content: bool, +) -> None: + tool_input: Final[dict[str, JsonValue]] = {"a": 1, "b": 2, "cache_control": {"ttl": "user-data"}} + block: Final[dict[str, JsonValue]] = { + "type": "tool_use", "id": "tool_1", "name": "lookup", "input": tool_input, + "cache_control": {"type": "ephemeral", "ttl": "5m"}, + } + changed_block: Final = ( + {**block, "input": dict(reversed(tool_input.items()))} if change == "tool_order" else + {**block, "cache_control": {"type": "ephemeral", "ttl": "1h"}} if change == "ttl" else block + ) + before: Final = _cache_plan(_prompt(block, role="assistant", output_config={"effort": "low"})).breakpoints[0] + after: Final = _cache_plan(_prompt( + changed_block, role="assistant", output_config={"effort": "high" if change == "effort" else "low"}, + **({"speed": "standard"} if change == "standard_speed" else {}), + )).breakpoints[0] + assert (before.fingerprint == after.fingerprint) is same_prefix + assert (before.fingerprint in after.lookback_fingerprints) is same_prefix + assert (before.content_fingerprint == after.content_fingerprint) is same_content + assert (before.content_fingerprint in after.lookback_content_fingerprints) is same_content + assert "user-data" in json.dumps(dict(before.prefix_body)) + assert not supported_prediction_headers({"anthropic-beta": "fast-mode-2026-02-01"}) + + +def test_cache_plan_automatic_cache_and_thinking_use_last_cacheable_block() -> None: + body: Final = _prompt( + _text("A stable answer"), {"type": "thinking", "thinking": "Thinking", "signature": "signature"}, + role="assistant", thinking={"type": "adaptive"}, cache_control={"type": "ephemeral", "ttl": "1h"}, + ) + plan: Final = _cache_plan(body) + assert len(plan.breakpoints) == 1 + assert plan.breakpoints[0].ttl_seconds == 3600 + assert plan.breakpoints[0].prefix_body == { + "thinking": {"type": "adaptive"}, + "messages": [{"role": "assistant", "content": [ + {"type": "text", "text": "A stable answer"}, + ]}], + } + assert parse_prompt(body) is None + + +@pytest.mark.parametrize("body, reason", [ + (_prompt({"type": "image"}), "unsupported_prompt_shape"), + ({**_body(), "unknown_native_setting": True}, "unsupported_prompt_shape"), + ({**_body(), "cache_control": {"type": "ephemeral", "ttl": "1h"}}, + "conflicting_cache_ttl"), + (_prompt(_text("five", "5m"), _text("hour", "1h")), "invalid_cache_ttl_order"), +]) +def test_cache_plan_unsupported_is_explicit( + body: Mapping[str, JsonValue], reason: str, +) -> None: + result: Final = parse_cache_plan(body) + assert isinstance(result, UnsupportedCachePlan) + assert result.reason == reason + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model, counts, reason", [ + (None, (100, 150, 200), None), + (None, (100, 201, 200), "inconsistent_prefix_token_count"), + (None, (151, 150, 200), "inconsistent_prefix_token_count"), + (None, (None, 150, 200), "token_count_unavailable"), + ("claude-opus-5", (100, 150, 200), None), + ("claude-sonnet-5", (100, 150, 200), None), + ("declared-cache-model", (100, 150, 200), None), + ("unknown-cache-model", (100, 150, 200), "unsupported_thinking_cache_semantics"), + ("claude-haiku-4-5", (100, 150, 200), "unsupported_thinking_cache_semantics"), + ("claude-sonnet-4-5", (100, 150, 200), "unsupported_thinking_cache_semantics"), +]) +async def test_cache_plan_count_conserves_total_and_rejects_unknown( + model: str | None, counts: tuple[int | None, int | None, int], reason: str | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem(litellm.model_cost, "declared-cache-model", { + "litellm_provider": "anthropic", "mode": "chat", "supports_thinking_cache_preservation": True, + }) + plan: Final = _cache_plan(_prompt( + {"type": "thinking", "thinking": "Retained thought", "signature": "signature"} + if model else _text("first", "5m"), + _text("second", "5m"), _text("uncached"), role="assistant" if model else "user", + )) + + async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + assert reason != "unsupported_thinking_cache_semantics", "Unverified thinking retention must skip counting" + if body is plan.full_body: + return counts[2] + return counts[0] if body is plan.breakpoints[0].prefix_body else counts[1] + + result: Final = await count_cache_plan(model or _MODEL, _KEY, plan, count) + if reason is not None: + assert isinstance(result, UnsupportedCachePlan) + assert result.reason == reason + else: + assert isinstance(result, CountedPromptCachePlan) + assert result.total_tokens == 200 + assert tuple(marker.prefix_tokens for marker in result.breakpoints) == ((100,) if model else (100, 150)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("section", ["system", "tools"]) +@pytest.mark.parametrize("rejects_prefix", (False, True)) +async def test_native_count_preserves_settings_and_requires_every_prefix( + section: str, rejects_prefix: bool, monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + params: Final = LiteLLM_Params( + model=f"anthropic/{_MODEL}", api_key=_KEY, + api_base="https://gateway.example/v1/messages", + ) + target: Final = resolve_baseline_prediction_target(params) + assert isinstance(target, NativePredictionTarget) + assert target.api_base == params.api_base + assert not isinstance(resolve_prediction_target(params), NativePredictionTarget) + body: Final = _body() + marker: Final[dict[str, JsonValue]] = {"type": "ephemeral", "ttl": "1h"} + body[section] = ([_text("A cached system", "1h")] if section == "system" else [{ + "name": "lookup", "input_schema": {"type": "object"}, "cache_control": marker, + }]) + plan: Final = _cache_plan({**body, **_prompt( + _text("A later prefix", "5m"), _text("An uncached suffix"), + thinking={"type": "adaptive"}, tool_choice={"type": "auto"}, output_config={"effort": "high"}, + )}) + assert len(plan.breakpoints) == 2 + assert plan.breakpoints[0].prefix_body["messages"] == [] + + async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + return await count_prompt_tokens( + model, api_key, {**body, "max_tokens": 100}, api_base=target.api_base, + ) + + with respx.mock(assert_all_called=False) as upstream: + endpoint: Final = "https://gateway.example/v1/messages/count_tokens" + routes: Final = tuple( + upstream.post(endpoint, json={**body, "model": _MODEL}).respond( + 400 if rejects_prefix and index == 1 else 200, + json={"detail": {"error": "messages parameter is required"}} + if rejects_prefix and index == 1 else {"input_tokens": tokens}, + ) + for index, (body, tokens) in enumerate(( + (plan.full_body, 6000), (plan.breakpoints[0].prefix_body, 5000), + (plan.breakpoints[1].prefix_body, 5800), + )) + ) + unexpected: Final = upstream.post(endpoint).respond(200, json={"input_tokens": 1}) + result: Final = await count_cache_plan(target.model, target.api_key, plan, count) + + if rejects_prefix: + assert result == UnsupportedCachePlan("token_count_unavailable") + else: + assert isinstance(result, CountedPromptCachePlan) + assert result.total_tokens == 6000 + assert tuple(marker.prefix_tokens for marker in result.breakpoints) == (5000, 5800) + assert tuple(route.call_count for route in routes) == (1, 1, 1) + assert unexpected.call_count == 0 diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index fff1372f271..420adc9338e 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1552,7 +1552,7 @@ async def test_anthropic_post_uses_prebuilt_body_without_redumping(): provider_config = Mock() provider_config.max_retry_on_anthropic_messages_http_error = 2 - logging_obj = Mock() + logging_obj: Final = Mock(baseline_cache_context=None) logging_obj.model_call_details = {} out = await handler._async_post_anthropic_messages_with_http_error_retry( @@ -1592,7 +1592,7 @@ async def test_anthropic_post_falls_back_to_json_dumps_when_unsigned_none(): provider_config = Mock() provider_config.max_retry_on_anthropic_messages_http_error = 1 - logging_obj = Mock() + logging_obj: Final = Mock(baseline_cache_context=None) logging_obj.model_call_details = {} await handler._async_post_anthropic_messages_with_http_error_retry( @@ -1640,7 +1640,7 @@ async def test_anthropic_post_retry_reserializes_mutated_body(): # Re-sign returns no signed body (native anthropic path) -> must re-dump. provider_config.sign_request = Mock(return_value=({}, None)) - logging_obj = Mock() + logging_obj: Final = Mock(baseline_cache_context=None) logging_obj.model_call_details = {} await handler._async_post_anthropic_messages_with_http_error_retry( @@ -2579,7 +2579,7 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques posts.append({"headers": dict(headers), "data": data}) return invalid_signature_response if len(posts) == 1 else ok_response - logging_obj = Mock() + logging_obj: Final = Mock(baseline_cache_context=None) logging_obj.model_call_details = {} response = await handler._async_post_anthropic_messages_with_http_error_retry( diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py index 7774f6b543d..777b4a265ac 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/test_litellm/models/test_models.py @@ -629,7 +629,7 @@ class TestManagedTables: class TestAutoRouterSession: @staticmethod - def _row(baseline_models: dict) -> LiteLLM_AutoRouterSession: + def _row(estimated_baseline_models: dict[str, int]) -> LiteLLM_AutoRouterSession: return LiteLLM_AutoRouterSession( api_key="k", session_id="s", @@ -643,7 +643,9 @@ class TestAutoRouterSession: saved_spend=0.24, classifier_cost=0.0, tier_turns={}, - baseline_models=baseline_models, + baseline_models={"legacy-baseline": 100}, + savings_estimated_turns=sum(estimated_baseline_models.values()), + savings_estimated_baseline_models=estimated_baseline_models, ) def test_the_baseline_label_is_the_one_most_turns_were_priced_against(self): @@ -655,5 +657,5 @@ class TestAutoRouterSession: assert self._row({"b-model": 1, "a-model": 1}).baseline_model == "b-model" assert self._row({"a-model": 1, "b-model": 1}).baseline_model == "b-model" - def test_a_row_whose_turns_recorded_no_baseline_has_no_label(self): + def test_a_row_without_current_estimates_has_no_baseline_label(self) -> None: assert self._row({}).baseline_model is None diff --git a/tests/test_litellm/proxy/client/cli/test_statusline_script.py b/tests/test_litellm/proxy/client/cli/test_statusline_script.py index 39d0e24d7b0..0cbeec86ee8 100644 --- a/tests/test_litellm/proxy/client/cli/test_statusline_script.py +++ b/tests/test_litellm/proxy/client/cli/test_statusline_script.py @@ -325,9 +325,12 @@ class TestRender: use_color=False, ) - def test_a_session_that_cost_more_than_its_baseline_reads_as_a_plus(self, config_dir): - dearer = RECORDED._replace(spend=0.50, baseline_spend=0.40) - assert "+25% vs Claude Opus 5" in render("m", dearer, config_dir, use_color=False) + @pytest.mark.parametrize("spend,delta", ((0.50, "+25%"), (0.40, "0%"), (0.4001, "0%"), (0.3999, "0%"), (0.30, "-25%"))) + def test_rounded_cost_delta_uses_a_sign_only_for_nonzero_percentages( + self, config_dir: Path, spend: float, delta: str, + ) -> None: + session: Final = RECORDED._replace(spend=spend, baseline_spend=0.40) + assert render("m", session, config_dir, use_color=False).splitlines()[0] == f"Routed to: m {delta} vs Claude Opus 5" def test_without_a_baseline_only_the_routed_line_shows(self, config_dir): assert render("m", RECORDED._replace(baseline_model=None), config_dir, False) == "Routed to: m" @@ -339,6 +342,37 @@ class TestRender: class TestClaudeCodeMode: + @pytest.mark.parametrize("estimated_turns", (0, 1)) + def test_current_estimates_keep_the_routed_model_and_compare_only_covered_turns( + self, tmp_path: Path, transcript: Path, config_dir: Path, estimated_turns: int + ) -> None: + session: Final = statusline_script._session_from_payload( + { + **RECORDED._asdict(), + "spend": 10.0, + "baseline_spend": None, + "savings_estimated_baseline_spend": 1.5 if estimated_turns else None, + "turns": 3, + "savings_estimated_turns": estimated_turns, + "savings_estimated_actual_spend": 2.0 if estimated_turns else 0.0, + } + ) + assert session is not None + + def fetch(credentials: Credentials, session_id: str) -> Fetched: + return Fetched(session, True) + + first: Final = _run(_payload(transcript), _env(tmp_path, config_dir), fetch) + assert first == _run(_payload(transcript), _env(tmp_path, config_dir), fetch) + assert first.startswith("Routed to: claude-sonnet-5") + if estimated_turns: + assert "+33% vs Claude Opus 5 · 1 of 3 turns estimated" in first + assert "$2.00" in first and "$1.50" in first + assert "$10.00" not in first and "+567%" not in first + else: + assert "Savings unavailable" in first + assert "%" not in first and "$" not in first + @pytest.mark.parametrize("transcript_model", ("claude-auto", "anthropic/claude-opus-5")) def test_the_session_names_the_routed_model_even_when_the_transcript_differs( self, tmp_path: Path, config_dir: Path, transcript_model: str diff --git a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py index 271751a3ff8..acd3dc18b54 100644 --- a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py +++ b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py @@ -281,6 +281,9 @@ class TestFlush: 0, "medium", "anthropic/claude-opus-5", + 0, + 0.0, + 0.0, ) def test_a_connect_error_retries_the_same_statement(self): @@ -307,7 +310,20 @@ class TestFlush: class TestEnqueueSeam: @pytest.mark.asyncio @pytest.mark.parametrize("classifier_cost", [0.005, 0.0, None]) - async def test_update_database_seam_enqueues_only_auto_routed_success(self, classifier_cost: float | None): + @pytest.mark.parametrize("estimate, covered, saved", [ + ({"version": 1, "status": "estimated"}, 1, -0.003), + ({"version": 1, "status": "estimated"}, 1, 0.0), + ({"version": 2, "status": "estimated"}, 1, 0.0), + ({"version": 3, "status": "estimated"}, 1, -0.003), + ({"version": 1, "status": "unknown"}, 0, 0.0), + ({"version": 0, "status": "estimated"}, 0, 0.0), + ({"version": 4, "status": "estimated"}, 0, 0.0), + ({"version": True, "status": "estimated"}, 0, 0.0), + (None, 0, -0.003), + ]) + async def test_update_database_seam_enqueues_only_auto_routed_success( + self, classifier_cost: float | None, estimate: dict[str, object] | None, covered: int, saved: float, + ) -> None: from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter writer: Final = DBSpendUpdateWriter() @@ -315,7 +331,8 @@ class TestEnqueueSeam: _autorouter_turn_transactions_lock=asyncio.Lock(), autorouter_turn_transactions=[] ) metadata: Final = _metadata( - routing_decision={**ROUTING_DECISION, "classifier_cost": classifier_cost}, autorouter_savings=-0.003 + routing_decision={**ROUTING_DECISION, "classifier_cost": classifier_cost}, + autorouter_savings=saved if covered else -0.003, autorouter_savings_estimate=estimate, ) for payload in ( _payload(metadata=json.dumps(metadata)), @@ -330,7 +347,10 @@ class TestEnqueueSeam: assert transaction.router_name == "live-auto" assert transaction.spend == pytest.approx(0.01 + (classifier_cost or 0.0)) assert transaction.classifier_cost == (classifier_cost or 0.0) - assert transaction.saved_spend == -0.003 + assert transaction.saved_spend == saved + assert transaction.savings_estimated_turns == covered + assert transaction.savings_estimated_actual_spend == pytest.approx(transaction.spend if covered else 0.0) + assert transaction.savings_estimated_saved_spend == (saved if covered else 0.0) def test_every_drain_trigger_reads_the_one_queue_census_owner(): diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index d08ff77f364..2ed5f263775 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2758,7 +2758,18 @@ async def test_daily_transaction_carries_compression_saved_tokens(): @pytest.mark.asyncio -async def test_daily_transaction_compression_saved_tokens_zero_when_absent(): +@pytest.mark.parametrize("estimate, recorded_savings, expected", [ + pytest.param(None, None, -0.005, id="plain-classifier-cost"), + pytest.param({"version": 1, "status": "unknown"}, None, 0.0, id="unknown"), + pytest.param({"version": 2, "status": "unknown"}, None, 0.0, id="unknown-v2"), + pytest.param({"version": 1, "status": "unknown"}, -0.003, 0.0, id="unknown-stale-value"), + pytest.param({"version": 0, "status": "estimated"}, -0.003, 0.0, id="unsupported-version"), + pytest.param({"version": 1, "status": "estimated"}, -0.003, -0.003, id="estimated"), + pytest.param(None, -0.003, -0.003, id="legacy"), +]) +async def test_daily_transaction_compression_saved_tokens_zero_when_absent( + estimate: dict[str, object] | None, recorded_savings: float | None, expected: float, +) -> None: """Requests without any compression metadata produce a zero count.""" writer = DBSpendUpdateWriter() mock_prisma = MagicMock() @@ -2776,7 +2787,12 @@ async def test_daily_transaction_compression_saved_tokens_zero_when_absent(): "prompt_tokens": 100, "completion_tokens": 10, "spend": 0.01, - "metadata": json.dumps({"usage_object": {}}), + "metadata": json.dumps({ + "usage_object": {"prompt_tokens": 100, "completion_tokens": 10}, + "routing_decision": {"savings_baseline_model": "anthropic/claude-sonnet-5", "classifier_cost": 0.005}, + "autorouter_savings": recorded_savings, + "autorouter_savings_estimate": estimate, + }), } transaction = await writer._common_add_spend_log_transaction_to_daily_transaction( @@ -2789,6 +2805,8 @@ async def test_daily_transaction_compression_saved_tokens_zero_when_absent(): assert transaction["compression_saved_tokens"] == 0 assert transaction["compression_savings_spend"] == 0 assert transaction["prompt_caching_savings_spend"] == 0 + assert transaction["spend"] == 0.01 + assert transaction["autorouter_savings_spend"] == expected # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/hooks/test_autorouter_baseline_cache.py b/tests/test_litellm/proxy/hooks/test_autorouter_baseline_cache.py new file mode 100644 index 00000000000..c6bb7833310 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_autorouter_baseline_cache.py @@ -0,0 +1,326 @@ +import asyncio +import json +from collections.abc import AsyncIterator, Callable, Generator, Mapping +from contextlib import contextmanager +from datetime import datetime +from types import MappingProxyType +from typing import Final, cast +from uuid import uuid4 + +import httpx +import pytest +import respx +from pydantic import JsonValue, TypeAdapter +from typing_extensions import NotRequired, ReadOnly, TypedDict + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.anthropic.prompt_cache_prediction import NativePredictionTarget, TokenCounter +from litellm.proxy.hooks.autorouter_baseline_cache import AutoRouterBaselineCache, CapturedBaselineObservation +from litellm.router import Router +from litellm.types.router import RetryPolicy +from litellm.types.utils import CallTypes, StandardLoggingRoutingDecision + +pytestmark: Final = pytest.mark.asyncio + + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +_OBJECTS: Final = TypeAdapter(dict[str, object]) + + +_MESSAGES: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +_MESSAGES_JSON: Final = """[{"role":"user","content":[ + {"type":"text","text":"stable","cache_control":{"type":"ephemeral","ttl":"1h"}}, + {"type":"text","text":"question"}]}]""" + + +_MODELS: Final = _MESSAGES.validate_json("""[ + {"model_name":"test-router","litellm_params":{"model":"auto_router/complexity_router", + "complexity_router_config":{"tiers":{"SIMPLE":"sonnet","MEDIUM":"sonnet","COMPLEX":"sonnet", + "REASONING":"opus"},"session_affinity":false, + "keyword_tier_rules":[{"keywords":["USE_OPUS"],"tier":"REASONING"}]}}}, + {"model_name":"sonnet","litellm_params":{"model":"anthropic/claude-sonnet-5","api_key":"test-selected"}, + "model_info":{"id":"selected"}}, + {"model_name":"opus","litellm_params":{"model":"anthropic/claude-opus-5","api_key":"test-selected"}, + "model_info":{"id":"baseline"}}]""") + + +def _message(completed: bool, model: str) -> Mapping[str, JsonValue]: + return _JSON_OBJECT.validate_json(f"""{{ + "id":"msg_baseline_test","type":"message","role":"assistant","model":{json.dumps(model)}, + "content":{'[{"type":"text","text":"OK"}]' if completed else "[]"}, + "stop_reason":{'"end_turn"' if completed else "null"},"stop_sequence":null, + "usage":{{"input_tokens":1000,"output_tokens":{10 if completed else 0}, + "cache_creation_input_tokens":5000,"cache_read_input_tokens":0, + "cache_creation":{{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":5000}}}}}}""") + + +_EVENTS: Final = _MESSAGES.validate_json("""[ + {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}, + {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"OK"}}, + {"type":"content_block_stop","index":0}, + {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":10}}, + {"type":"message_stop"} +]""") + + +async def _count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int: + assert model == "claude-opus-5" + return 6000 if "question" in json.dumps(_JSON_OBJECT.validate_python(body)) else 5000 + + +class _CallContext(TypedDict): + litellm_logging_obj: NotRequired[ReadOnly[Logging]] + litellm_call_id: ReadOnly[str] + litellm_metadata: ReadOnly[Mapping[str, object]] + litellm_session_id: ReadOnly[str] + + +def _kwargs(logging_obj: Logging, trusted: bool = True, *, explicit_logging: bool = True) -> _CallContext: + context: Final = _OBJECTS.validate_json('{"litellm_metadata":{"user_api_key_hash":"test-caller-hash"}}') + Router._record_routing_decision( # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # production trusted stamp owner + context, + StandardLoggingRoutingDecision( + router_model_name="test-router", + router_type="complexity", + routed_model="sonnet", + cause="heuristic_scorer", + conversation_continuing=True, + savings_baseline_model="anthropic/claude-opus-5", + savings_baseline_deployment_id="baseline", + ), + ) + metadata: Final = _OBJECTS.validate_python(context["litellm_metadata"]) + if not trusted: + metadata["_autorouter_baseline_route"] = _JSON_OBJECT.validate_json( + '{"router_name":"test-router","baseline_model":"anthropic/claude-opus-5","baseline_deployment_id":"baseline"}' + ) + envelope: Final[_CallContext] = { + "litellm_call_id": logging_obj.litellm_call_id, + "litellm_session_id": "baseline-session", + "litellm_metadata": metadata, + } + supplied: Final[_CallContext] = {**envelope, "litellm_logging_obj": logging_obj} + return supplied if explicit_logging else envelope + + +def _stream(logging_obj: Logging) -> bool: + return logging_obj.stream is True # pyright: ignore[reportUnknownMemberType] # normalize the legacy Logging flag + + +def _sse(completed: bool = True, model: str = "claude-sonnet-5") -> tuple[bytes, ...]: + events: Final = ( + { # mutable-ok: json.dumps needs a concrete event dictionary + "type": "message_start", + "message": _message(False, model), + }, + *_EVENTS, + ) + return tuple( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + for event in (events if completed else events[:-1]) + ) + + +def _upstream(request: httpx.Request) -> httpx.Response: + body: Final = _JSON_OBJECT.validate_json(request.content) + model: Final = body.get("model") + assert isinstance(model, str) + stream: Final = body.get("stream") is True + content: Final = b"".join(_sse(model=model)) if stream else json.dumps(_message(True, model)).encode() + return httpx.Response(200, content=content, request=request, + headers=MappingProxyType({"content-type": "text/event-stream" if stream else "application/json"}), + ) + + +def _error(request: httpx.Request, code: int, message: str) -> httpx.Response: + return httpx.Response( + code, + text='{"type":"error","error":{"type":"rate_limit_error","message":' + json.dumps(message) + "}}", + headers=MappingProxyType({"retry-after": "0"}), + request=request, + ) + + +@contextmanager +def _transport(upstream: Callable[[httpx.Request], httpx.Response]) -> Generator[respx.Route]: + with respx.mock() as transport: + yield transport.post("https://api.anthropic.com/v1/messages").mock(side_effect=upstream) + + +class _NativeOptions(TypedDict): + api_key: NotRequired[ReadOnly[str]] + num_retries: NotRequired[ReadOnly[int]] + + +async def _call( + target: Router | None, + logging_obj: Logging, + *, + trusted: bool = True, + messages: str = _MESSAGES_JSON, + explicit_logging: bool = True, +) -> None: + invoke: Final = target.anthropic_messages if target else litellm.anthropic_messages # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # legacy native call signatures + options: Final = _NativeOptions() if target else _NativeOptions(api_key="test-selected", num_retries=0) + response: Final[object] = await invoke( # pyright: ignore[reportUnknownVariableType] # native Router returns an opaque SDK result + model="test-router" if target else "anthropic/claude-sonnet-5", + max_tokens=16, + stream=_stream(logging_obj), + messages=_MESSAGES.validate_json(messages), + **options, + **_kwargs(logging_obj, trusted, explicit_logging=explicit_logging), + ) + assert response is not None + if _stream(logging_obj): + assert isinstance(response, AsyncIterator) + stream: Final = cast(AsyncIterator[object], response) # cast-ok: iterator checked; all items satisfy object + assert tuple([chunk async for chunk in stream]) + +class _Capture(CustomLogger): + def __init__(self, call_id: str) -> None: + self.call_id: Final = call_id + self.payloads: Final[asyncio.Queue[Mapping[str, object]]] = asyncio.Queue() + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + payload: Final = _OBJECTS.validate_python(kwargs.get("standard_logging_object")) + if payload.get("litellm_call_id") == self.call_id: + self.payloads.put_nowait(payload) + + async def payload(self) -> Mapping[str, object]: + return await asyncio.wait_for(self.payloads.get(), timeout=20) + + +class _Rig: + def __init__(self, monkeypatch: pytest.MonkeyPatch, *, retries: int = 0, count: TokenCounter = _count) -> None: + self.router: Final = Router(model_list=_MODELS, num_retries=retries, + retry_policy=RetryPolicy(RateLimitErrorRetries=retries), disable_cooldowns=True) + + def router() -> Router: + return self.router + + self.hook: Final = AutoRouterBaselineCache(None, router=router, token_counter=count) + self.call_id: Final = uuid4().hex + self.capture: Final = _Capture(self.call_id) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + for name in ("ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "callbacks", [self.hook]) + for name in ("success_callback", "failure_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(litellm, "_async_success_callback", [self.capture]) + + def logging(self, stream: bool = False) -> Logging: + return Logging(model="anthropic/claude-sonnet-5", messages=_MESSAGES.validate_json(_MESSAGES_JSON), + stream=stream, call_type=CallTypes.anthropic_messages.value, start_time=datetime.now(), + litellm_call_id=self.call_id, function_id=self.call_id, kwargs={"litellm_session_id":"baseline-session"}) + + +def _observation(payload: Mapping[str, object]) -> CapturedBaselineObservation: + encoded: Final = payload["autorouter_baseline_observation"] + assert isinstance(encoded, str) + assert "test-selected" not in encoded and "stable" not in encoded and "x-api-key" not in encoded + return CapturedBaselineObservation.model_validate_json(encoded) + + +@pytest.mark.parametrize("stream,baseline", ((False, False), (True, False), (False, True), (True, True))) +async def test_native_logging_captures_usage_without_publishing_hypothetical_savings( + monkeypatch: pytest.MonkeyPatch, stream: bool, baseline: bool, +) -> None: + rig: Final = _Rig(monkeypatch) + messages: Final = _MESSAGES_JSON.replace("question", "question USE_OPUS") if baseline else _MESSAGES_JSON + with _transport(_upstream): + await _call(rig.router, rig.logging(stream), messages=messages) + payload: Final = await rig.capture.payload() + captured: Final = _observation(payload) + assert payload["autorouter_savings"] is None + assert _OBJECTS.validate_python(payload["autorouter_savings_estimate"])["reason"] == "pending_projection" + assert captured.observation.outcome == "complete" + assert captured.observation.baseline_equivalent == baseline + assert captured.observation.usage is not None and captured.observation.usage.completion_tokens == 10 + assert captured.observation.plan is not None and captured.observation.plan.total_tokens == 6000 + + +async def test_count_failure_preserves_initial_observed_equivalence(monkeypatch: pytest.MonkeyPatch) -> None: + async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + return None + + rig: Final = _Rig(monkeypatch, count=count) + with _transport(_upstream): + await _call(rig.router, rig.logging(), messages=_MESSAGES_JSON.replace("question", "question USE_OPUS")) + captured: Final = _observation(await rig.capture.payload()) + assert captured.observation.baseline_equivalent and captured.observation.usage is not None + assert captured.observation.plan is None and captured.observation.reason == "token_count_unavailable" + + +async def test_native_retry_is_uncertain_even_when_final_response_succeeds(monkeypatch: pytest.MonkeyPatch) -> None: + rig: Final = _Rig(monkeypatch, retries=1) + + def upstream(request: httpx.Request) -> httpx.Response: + return _upstream(request) if route.call_count else _error(request, 429, "retry") + + with _transport(upstream) as route: + await _call(rig.router, rig.logging()) + captured: Final = _observation(await rig.capture.payload()) + assert route.call_count == 2 + assert captured.observation.outcome == "uncertain" + assert captured.observation.reason == "retried_request" + + +async def test_caller_cannot_forge_an_observation_scope(monkeypatch: pytest.MonkeyPatch) -> None: + rig: Final = _Rig(monkeypatch) + with _transport(_upstream): + await _call(None, rig.logging(), trusted=False) + payload: Final = await rig.capture.payload() + assert payload["autorouter_baseline_observation"] is None + assert payload["autorouter_savings"] is None + + +@pytest.mark.parametrize("model,key,endpoint", ( + ("claude-sonnet-5", "test-first", None), + ("claude-opus-5", "test-second", None), + ("claude-opus-5", "test-first", "https://example.test"), +)) +async def test_count_memo_is_scoped_to_provider_recipient(model: str, key: str, endpoint: str | None) -> None: + counts: Final = iter((5000, 6000)) + + async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int: + return next(counts) + + collector: Final = AutoRouterBaselineCache(None, token_counter=count) + original: Final = NativePredictionTarget("claude-opus-5", "test-first") + other: Final = NativePredictionTarget(model, key, endpoint) + assert await collector._count(original, {}) == 5000 # pyright: ignore[reportPrivateUsage] + assert await collector._count(other, {}) == 6000 # pyright: ignore[reportPrivateUsage] + assert await collector._count(original, {}) == 5000 # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.parametrize("stream", (False, True)) +async def test_provider_counting_does_not_hold_the_inference_response( + monkeypatch: pytest.MonkeyPatch, stream: bool, +) -> None: + counting: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int: + counting.set() + await release.wait() + return await _count(model, api_key, body) + + rig: Final = _Rig(monkeypatch, count=count) + try: + with _transport(_upstream): + await asyncio.wait_for(_call(rig.router, rig.logging(stream)), timeout=2) + await asyncio.wait_for(counting.wait(), timeout=2) + assert rig.capture.payloads.empty() + release.set() + assert _observation(await rig.capture.payload()).observation.plan is not None + finally: + release.set() diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 067f30c2fd7..6ac053f4e15 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -546,6 +546,9 @@ class TestAutoRouterBenchmarks: total_tokens=4000, spend=10.0, saved_spend=30.0, + savings_estimated_turns=40, + savings_estimated_actual_spend=10.0, + savings_estimated_saved_spend=30.0, classifier_cost=0.4, classifier_cost_recorded_turns=40, session_seconds=400.0, @@ -582,12 +585,29 @@ class TestAutoRouterBenchmarks: def test_a_losing_router_reports_negative_savings(self): from litellm.proxy.management_endpoints.auto_router_endpoints import _benchmark_totals - losing = self.ROW.model_copy(update={"saved_spend": -5.0}) + losing = self.ROW.model_copy(update={"saved_spend": -5.0, "savings_estimated_saved_spend": -5.0}) totals = _benchmark_totals(losing) assert totals.baseline_spend == 5.0 assert totals.saved_pct == -100.0 assert totals.classifier_cost == 0.4 + @pytest.mark.parametrize("estimated_turns", [0, 4]) + def test_savings_compare_only_the_current_estimated_cohort(self, estimated_turns: int) -> None: + from litellm.proxy.management_endpoints.auto_router_endpoints import _benchmark_totals + + row: Final = self.ROW.model_copy(update={ + "savings_estimated_turns": estimated_turns, + "savings_estimated_actual_spend": 2.0 if estimated_turns else 0.0, + "savings_estimated_saved_spend": -0.5 if estimated_turns else 0.0, + }) + totals: Final = _benchmark_totals(row) + assert totals.spend == 10.0 + assert totals.savings_estimated_turns == estimated_turns + assert totals.saved_spend == (-0.5 if estimated_turns else None) + assert totals.baseline_spend == (1.5 if estimated_turns else None) + assert totals.saved_pct == (pytest.approx(-33.3) if estimated_turns else None) + assert totals.saved_per_session is None + def test_an_empty_window_folds_to_zeros(self): from litellm.proxy.management_endpoints.auto_router_endpoints import ( _benchmark_totals, @@ -607,7 +627,10 @@ class TestAutoRouterBenchmarks: _summed_agg_row, ) - other = self.ROW.model_copy(update={"router_name": "auto-2", "sessions": 1, "turns": 10, "spend": 0.0}) + other = self.ROW.model_copy(update={ + "router_name": "auto-2", "sessions": 1, "turns": 10, "spend": 0.0, + "savings_estimated_turns": 10, "savings_estimated_actual_spend": 0.0, + }) summed = _summed_agg_row([self.ROW, other]) totals = _benchmark_totals(summed) assert summed.sessions == 5 @@ -696,6 +719,9 @@ class TestAutoRouterBenchmarks: "turns": 10, "spend": 2.0, "saved_spend": -0.5, + "savings_estimated_turns": 10, + "savings_estimated_actual_spend": 2.0, + "savings_estimated_saved_spend": -0.5, "classifier_cost": recorded_turns * 0.02, "classifier_cost_recorded_turns": recorded_turns, } @@ -876,6 +902,10 @@ class TestAutoRouterSession: "last_model": "anthropic/claude-sonnet-5", "spend": 0.14, "saved_spend": 0.24, + "savings_estimated_turns": 3, + "savings_estimated_actual_spend": 0.14, + "savings_estimated_saved_spend": 0.24, + "savings_estimated_baseline_models": {"anthropic/claude-opus-5": 3}, "classifier_cost": 0.0, "tier_turns": {"simple": 1, "complex": 2}, "baseline_models": {"anthropic/claude-opus-5": 3}, @@ -899,25 +929,33 @@ class TestAutoRouterSession: return lookups @pytest.mark.asyncio + @pytest.mark.parametrize("turns, estimated", [(3, True), (10, True), (10, False)], ids=["full", "partial", "legacy"]) async def test_a_key_reads_its_own_session_with_the_baseline_its_turns_were_priced_against( - self, monkeypatch: pytest.MonkeyPatch - ): + self, monkeypatch: pytest.MonkeyPatch, turns: int, estimated: bool, + ) -> None: from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session caller = UserAPIKeyAuth(api_key="sk-caller") - self._rig(monkeypatch, [{**self.ROW, "api_key": caller.api_key, "session_id": "sess-1"}]) + row: Final = {key: value for key, value in self.ROW.items() if estimated or not key.startswith("savings_estimated_")} + spend: Final = 0.14 if turns == 3 else 10.0 + if estimated and turns != 3: + row["savings_estimated_saved_spend"] = -0.04 + self._rig(monkeypatch, [{**row, "api_key": caller.api_key, "session_id": "sess-1", "turns": turns, "spend": spend}]) response = await get_auto_router_session(user_api_key_dict=caller, session_id="sess-1") assert response.model_dump() == { "session_id": "sess-1", "router_name": "claude-auto", "router_type": "complexity", - "turns": 3, + "turns": turns, "last_model": "anthropic/claude-sonnet-5", - "spend": 0.14, - "saved_spend": 0.24, - "baseline_spend": pytest.approx(0.38), - "baseline_model": "anthropic/claude-opus-5", - "baseline_models": {"anthropic/claude-opus-5": 3}, + "spend": spend, + "saved_spend": (0.24 if turns == 3 else -0.04) if estimated else None, + "savings_estimated_turns": 3 if estimated else 0, + "savings_estimated_actual_spend": 0.14 if estimated else 0.0, + "baseline_spend": pytest.approx(0.38) if turns == 3 else None, + "savings_estimated_baseline_spend": pytest.approx(0.38 if turns == 3 else 0.1) if estimated else None, + "baseline_model": "anthropic/claude-opus-5" if estimated else None, + "baseline_models": {"anthropic/claude-opus-5": 3} if estimated else {}, } @pytest.mark.asyncio @@ -959,22 +997,14 @@ class TestAutoRouterSession: from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session priced = {"anthropic/claude-opus-5": 2, "anthropic/claude-sonnet-5": 1} - self._rig(monkeypatch, [{**self.ROW, "api_key": ADMIN.api_key, "session_id": "s", "baseline_models": priced}]) + self._rig(monkeypatch, [{ + **self.ROW, "api_key": ADMIN.api_key, "session_id": "s", + "baseline_models": {"old-baseline": 100}, "savings_estimated_baseline_models": priced, + }]) response = await get_auto_router_session(user_api_key_dict=ADMIN, session_id="s") assert response.baseline_model == "anthropic/claude-opus-5" assert response.baseline_models == priced - @pytest.mark.asyncio - async def test_a_session_whose_turns_recorded_no_baseline_reports_the_money_without_a_name( - self, monkeypatch: pytest.MonkeyPatch - ): - from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session - - self._rig(monkeypatch, [{**self.ROW, "api_key": ADMIN.api_key, "session_id": "s", "baseline_models": {}}]) - response = await get_auto_router_session(user_api_key_dict=ADMIN, session_id="s") - assert response.baseline_model is None - assert response.baseline_spend == pytest.approx(0.38) - @pytest.mark.asyncio async def test_an_oversized_client_session_id_is_bounded_like_the_writer_bounded_it( self, monkeypatch: pytest.MonkeyPatch diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index aaa3b205312..680dd4df0ae 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -11,12 +11,23 @@ from fastapi.testclient import TestClient from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.public_endpoints import router +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, ) from litellm.types.utils import LlmProviders +def test_fuse_presets_route_serves_the_shared_catalog_without_authentication() -> None: + app: Final = FastAPI() + app.include_router(router) + client: Final = TestClient(app) + response: Final = client.get("/public/complexity_router/fuse_presets") + assert response.status_code == 200 + assert response.json() == get_fuse_presets().model_dump(mode="json") + assert client.get("/public/complexity_router/fuse_presets").json() == response.json() + + def test_get_supported_providers_returns_enum_values(): app_instance = FastAPI() app_instance.include_router(router) diff --git a/tests/test_litellm/proxy/spend_tracking/test_baseline_accounting.py b/tests/test_litellm/proxy/spend_tracking/test_baseline_accounting.py new file mode 100644 index 00000000000..a188d65502d --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_baseline_accounting.py @@ -0,0 +1,199 @@ +from dataclasses import replace +from itertools import groupby +from typing import Final + +import pytest + +import litellm +from litellm.llms.anthropic.cost_calculation import cost_per_token +from litellm.llms.anthropic.prompt_cache_prediction import CountedBreakpoint, CountedPromptCachePlan +from litellm.proxy.spend_tracking.baseline_accounting import ( + BaselineEstimate, + BaselineHistory, + BaselineObservation, + CacheEntry, + advance_baseline_history, +) +from litellm.types.utils import CacheCreationTokenDetails, PromptTokensDetailsWrapper, Usage + + +def _usage() -> Usage: + return Usage( + prompt_tokens=6200, + completion_tokens=30, + total_tokens=6230, + cache_read_input_tokens=0, + cache_creation_input_tokens=6000, + speed="fast", + inference_geo="us", + completion_tokens_details={"reasoning_tokens": 20}, + server_tool_use={"web_search_requests": 1}, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=200, + cached_tokens=0, + cache_creation_tokens=6000, + cache_write_tokens=6000, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_5m_input_tokens=0, ephemeral_1h_input_tokens=6000 + ), + ), + ) + + +def _marker( + name: str = "prefix", ttl: int = 3600, tokens: int = 6000, previous: tuple[str, ...] = () +) -> CountedBreakpoint: + return CountedBreakpoint( + fingerprint=f"{name}:{ttl}", + ttl_seconds=ttl, + prefix_tokens=tokens, + lookback_fingerprints=(*(f"{item}:{ttl}" for item in previous), f"{name}:{ttl}"), + content_fingerprint=name, + lookback_content_fingerprints=(*previous, name), + ) + + +def _observation(request_id: str, started: float = 10000.0, **overrides: object) -> BaselineObservation: + return BaselineObservation.model_validate( + { + "request_id": request_id, + "started_at": started, + "available_at": started + 0.1, + "outcome": "complete", + "baseline_equivalent": False, + "usage": _usage(), + "plan": CountedPromptCachePlan(6200, (_marker(),)), + "minimum_cache_tokens": 4096, + **overrides, + } + ) + + +def _replay(*observations: BaselineObservation) -> tuple[BaselineEstimate, ...]: + history = BaselineHistory() + results: list[BaselineEstimate] = [] + for _, group in groupby(sorted(observations, key=lambda item: item.started_at), key=lambda item: item.started_at): + history, estimates = advance_baseline_history(history, tuple(group)) + results.extend(estimates) + return tuple(results) + + +def test_initial_identical_path_preserves_full_usage_without_counting_or_exclusive_owner() -> None: + initial: Final = _observation("main", baseline_equivalent=True, plan=None, reason="unsupported_request_headers") + background: Final = initial.model_copy(update={"request_id": "background"}) + later: Final = initial.model_copy(update={"request_id": "later", "started_at": 10001.0, "available_at": 10002.0}) + estimates: Final = _replay(initial, background, later) + assert all(item.provenance == "observed_identical" and item.usage == initial.usage for item in estimates) + assert all(item.usage is not initial.usage for item in estimates) + assert all(item.usage.prompt_tokens == 6200 for item in estimates if item.usage is not None) + + +def test_late_divergent_observation_replays_in_event_order_and_removes_initial_zero() -> None: + same: Final = _observation("same", 10001.0, baseline_equivalent=True) + early: Final = _observation("early") + assert _replay(same)[0].provenance == "observed_identical" + replayed: Final = _replay(same, early) + assert replayed == _replay(early, same) + assert replayed[0].usage is None + assert replayed[1].provenance == "modeled" + assert replayed[1].usage is not None and replayed[1].usage.prompt_tokens_details.cached_tokens == 6000 + + +@pytest.mark.parametrize("ttl", [300, 3600]) +def test_prefix_match_expiry_and_usage_pricing_fields(ttl: int) -> None: + plan: Final = CountedPromptCachePlan(6200, (_marker(ttl=ttl),)) + first: Final = _observation("first", baseline_equivalent=True, plan=plan) + # Each replay starts from the original observation, so warm does not refresh the expiry case. + warm: Final = _replay(first, _observation("warm", 10000.0 + ttl - 0.01, plan=plan))[-1] + cold: Final = _replay(first, _observation("cold", 10000.0 + ttl, plan=plan))[-1] + assert warm.reason == "cache_prefix_available" and cold.reason == "cache_prefix_expired" + assert warm.usage is not None and cold.usage is not None + assert warm.usage.prompt_tokens_details.cached_tokens == 6000 + assert cold.usage.prompt_tokens_details.cached_tokens == 0 + assert cold.usage.prompt_tokens_details.cache_creation_tokens == 6000 + unaffected: Final = {"prompt_tokens", "total_tokens", "prompt_tokens_details", "cache_read_input_tokens", "cache_creation_input_tokens"} + assert warm.usage.model_dump(exclude=unaffected) == first.usage.model_dump(exclude=unaffected) + assert cold.usage.model_dump(exclude=unaffected) == first.usage.model_dump(exclude=unaffected) + + +@pytest.mark.parametrize("warm_tail", (False, True)) +def test_growth_lookback_and_mixed_ttl_keep_distinct_read_write_buckets(warm_tail: bool) -> None: + first: Final = _observation("first", baseline_equivalent=True) + grown: Final = CountedPromptCachePlan(7100, (_marker("grown", 3600, 6500, ("prefix",)), _marker("tail", 300, 7000))) + # Initial unseen suffixes remain unknown within their potential pre-existing cache horizon. + second: Final = _replay(first, _observation("second", 10001.0, plan=grown))[-1] + assert second.reason == "history_unavailable" + history: Final = BaselineHistory( + first_at=1.0, last_at=10000.0, equivalent=False, uncertain_before=1.0, + entries=(CacheEntry("tail:300", "tail", 7000, 300, 10000.0, 10300.0),) if warm_tail else (), + ) + _, estimates = advance_baseline_history(history, (_observation("mixed", 10001.0, plan=grown),)) + usage: Final = estimates[0].usage + assert usage is not None + assert usage.prompt_tokens_details.text_tokens == 100 + # Anthropic billing locations: B is the highest 1h breakpoint AFTER the highest hit A. + # https://platform.claude.com/docs/en/build-with-claude/prompt-caching#mixing-different-ttls (2026-09-15) + assert usage.prompt_tokens_details.cached_tokens == (7000 if warm_tail else 0) + assert usage.prompt_tokens_details.cache_creation_token_details.ephemeral_1h_input_tokens == (0 if warm_tail else 6500) + assert usage.prompt_tokens_details.cache_creation_token_details.ephemeral_5m_input_tokens == (0 if warm_tail else 500) + + +@pytest.mark.parametrize("change", ["prefix", "ttl", "unavailable", "failed", "response_cache"]) +def test_uncertainty_and_replays_do_not_manufacture_hits(change: str) -> None: + first: Final = _observation("first", baseline_equivalent=True) + changes: Final = { + "prefix": {"plan": CountedPromptCachePlan(6200, (_marker("changed"),))}, + "ttl": {"plan": CountedPromptCachePlan(6200, (_marker(ttl=300),))}, + "unavailable": {"plan": None, "reason": "token_count_unavailable"}, + "failed": {"outcome": "uncertain", "reason": "incomplete_response"}, + "response_cache": {"outcome": "response_cache"}, + } + second: Final = _observation("second", 10001.0, **changes[change]) + third: Final = _observation("third", 10002.0) + middle, result = _replay(first, second, third)[1:] + assert middle.usage is None + if change in ("unavailable", "failed", "ttl"): + assert result.usage is None + else: + assert result.usage is not None and result.usage.prompt_tokens_details.cached_tokens == 6000 + + +def test_first_token_availability_and_simultaneous_divergence_are_conservative() -> None: + slow: Final = _observation("slow", available_at=10002.0, baseline_equivalent=True) + overlap: Final = _observation("overlap", 10001.0) + assert _replay(slow, overlap)[-1].usage is None + assert all(item.provenance != "observed_identical" for item in _replay(slow, _observation("tie"))) + + +def test_invalid_usage_and_invalid_count_plan_cannot_seed_cache() -> None: + bad: Final = _observation("bad", baseline_equivalent=True, usage=_usage().model_copy(update={"total_tokens": 1})) + assert all(item.usage is None for item in _replay(bad, _observation("next", 10001.0))) + broken: Final = CountedPromptCachePlan(6200, (replace(_marker(), prefix_tokens=7000),)) + assert _replay(_observation("bad", plan=broken))[0].usage is None + + +def test_overlapping_uncertain_request_cannot_be_warmed_by_a_later_callback() -> None: + uncertain: Final = _observation("incomplete", outcome="uncertain", available_at=10010.0) + overlap: Final = _observation("overlap", 10001.0) + during: Final = _observation("during", 10002.0) + after: Final = _observation("after", 10011.0) + warmed: Final = _observation("warmed", 10012.0) + estimates: Final = _replay(uncertain, overlap, during, after, warmed) + assert estimates[1].reason == estimates[2].reason == "concurrent_uncertainty" + assert estimates[3].usage is None + assert estimates[4].usage is not None and estimates[4].usage.prompt_tokens_details.cached_tokens == 6000 + + +def test_modeled_read_cannot_recharge_the_original_private_write_count() -> None: + warm: Final = _replay(_observation("initial", baseline_equivalent=True), _observation("warm", 10001.0))[-1] + assert warm.usage is not None + prices: Final = { + **litellm.get_model_info("claude-opus-5", custom_llm_provider="anthropic"), + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_read_input_token_cost": 1e-7, + "cache_creation_input_token_cost": 1.25e-6, + "provider_specific_entry": {"fast": 2.0, "us": 1.1}, + } + input_cost, output_cost = cost_per_token("claude-opus-5", warm.usage, model_info=prices) + assert input_cost + output_cost == pytest.approx((200 * 1e-6 + 6000 * 1e-7 + 30 * 2e-6) * 2.0 * 1.1) diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index 615938f2e33..aae966022e3 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -1,4 +1,4 @@ -from typing import Final +from typing import Final, Literal import pytest @@ -23,13 +23,13 @@ pytestmark = pytest.mark.usefixtures("local_model_cost_map") def test_baseline_preserves_anthropic_pricing_fields(modifier: dict[str, str], continuing: bool) -> None: usage: Final = _usage(1000, 0, 1000, 100).model_copy(update=modifier) expected: Final = (_usage(1000, 1000, 0, 100) if continuing else usage).model_copy(update=modifier) - normalized: Final = _baseline_usage(usage, continuing) + normalized: Final = _baseline_usage(expected) cache_fields: Final = {"prompt_tokens_details", "cache_read_input_tokens", "cache_creation_input_tokens"} assert normalized.model_dump(exclude=cache_fields) == usage.model_dump(exclude=cache_fields) assert usage.prompt_tokens_details.cached_tokens == 0 selected_cost: Final = 0.013 assert compute_autorouter_savings( - "claude-opus-5", "claude-sonnet-5", "anthropic", usage, conversation_continuing=continuing, + "claude-opus-5", "claude-sonnet-5", "anthropic", usage, baseline_usage=expected, cost_breakdown={"input_cost": 0.01, "output_cost": 0.003}, ) == pytest.approx(sum(anthropic_cost_per_token("claude-opus-5", expected)) - selected_cost) @@ -41,7 +41,7 @@ def test_anthropic_baseline_keeps_negotiated_prices_with_provider_multiplier() - } usage: Final = _usage(1000, 1000, 0, 100).model_copy(update={"speed": "fast"}) assert compute_autorouter_savings( - "claude-opus-5", "claude-sonnet-5", "anthropic", usage, baseline_info=info, + "claude-opus-5", "claude-sonnet-5", "anthropic", usage, baseline_info=info, baseline_usage=usage, cost_breakdown={"input_cost": 0.01, "output_cost": 0.003}, ) == pytest.approx(0.0015 * 2 - 0.013) @@ -405,146 +405,109 @@ def test_negative_token_counts_clamp_to_zero(): assert result.prompt_caching == 0.0 -def _usage(fresh: int, cached: int, written: int, out: int) -> Usage: +def _usage(fresh: int, cached: int, written: int, out: int, *, hour: bool = False, image: int = 0) -> Usage: """Usage as the spend log records it; `prompt_tokens` is the inclusive total.""" return Usage( prompt_tokens=fresh + cached + written, completion_tokens=out, total_tokens=fresh + cached + written + out, - prompt_tokens_details={"cached_tokens": cached, "cache_creation_tokens": written, "text_tokens": fresh}, + prompt_tokens_details={ + "cached_tokens": cached, "cache_creation_tokens": written, "text_tokens": fresh - image, "image_tokens": image, + "cache_creation_token_details": {"ephemeral_1h_input_tokens": written} if hour else None, + }, cache_read_input_tokens=cached, cache_creation_input_tokens=written, ) -def _savings(baseline: str, selected: str, usage: Usage, continuing: bool = True) -> float: - """Savings for a request, defaulting to a conversation already underway. - - `continuing=True` is the mid-conversation case, where the baseline had the prompt - cached and this request's write is what the switch cost. `continuing=False` is a - conversation's first turn, where nothing was cached for any model. - """ +def _savings(baseline: str, selected: str, usage: Usage, baseline_usage: Usage | None = None) -> float | None: return compute_autorouter_savings( baseline_model=baseline, selected_model=selected, selected_provider="anthropic", usage=usage, - conversation_continuing=continuing, + baseline_usage=baseline_usage, ) -def test_switching_models_mid_conversation_charges_the_cold_cache_write(): - """Staying on one model writes the cache once and reads it thereafter. Switching - leaves the new model cold, so it pays to write the whole prompt again; when that - charge outweighs the cheaper rates the route lost money and must report a loss. - - Pricing the baseline as if it too re-wrote the cache credits a charge it never - paid, which is how a losing switch used to read as the largest saving on the page. - """ - usage = _usage(fresh=3, cached=500, written=12304, out=500) - result = _savings("claude-sonnet-5", "claude-haiku-4-5", usage) - - sonnet = litellm.get_model_info("claude-sonnet-5", "anthropic") - haiku = litellm.get_model_info("claude-haiku-4-5", "anthropic") - warm_baseline = ( - 3 * sonnet["input_cost_per_token"] - + 12804 * sonnet["cache_read_input_token_cost"] - + 500 * sonnet["output_cost_per_token"] +@pytest.mark.parametrize("baseline, selected, actual, modeled, loses_money", [ + pytest.param("claude-sonnet-5", "claude-haiku-4-5", _usage(3, 500, 12304, 500), + _usage(3, 12804, 0, 500), True, id="warm-baseline-cold-route"), + pytest.param("claude-opus-5", "claude-opus-5", _usage(0, 0, 20000, 1000), + _usage(0, 20000, 0, 1000), True, id="same-model-cold-route"), + pytest.param("claude-opus-5", "claude-sonnet-5", _usage(0, 19000, 1000, 1000), + _usage(0, 19500, 500, 1000), False, id="partly-cached-growth"), + pytest.param("claude-opus-5", "claude-sonnet-5", _usage(0, 0, 100000, 1000, hour=True), + _usage(0, 0, 100000, 1000, hour=True), False, id="expired-one-hour"), + pytest.param("claude-opus-5", "claude-sonnet-5", _usage(0, 0, 100000, 1000, hour=True), + _usage(0, 100000, 0, 1000), True, id="invented-one-hour-hit"), + pytest.param("claude-opus-5", "claude-sonnet-5", _usage(4000, 0, 16000, 1000, hour=True, image=4000), + _usage(4000, 0, 16000, 1000, hour=True, image=4000), False, id="image-and-one-hour-write"), +]) +def test_supplied_baseline_usage_is_priced_independently( + baseline: str, selected: str, actual: Usage, modeled: Usage, loses_money: bool, +) -> None: + result: Final = _savings(baseline, selected, actual, modeled) + expected: Final = sum(generic_cost_per_token(model=baseline, usage=modeled, custom_llm_provider="anthropic")) - sum( + generic_cost_per_token(model=selected, usage=actual, custom_llm_provider="anthropic") ) - actually_paid = ( - 3 * haiku["input_cost_per_token"] - + 500 * haiku["cache_read_input_token_cost"] - + 12304 * haiku["cache_creation_input_token_cost"] - + 500 * haiku["output_cost_per_token"] - ) - assert result == pytest.approx(warm_baseline - actually_paid) - assert result < 0, "a cache-thrashing switch must report a loss, not a saving" - - phantom = 12304 * sonnet["cache_creation_input_token_cost"] - assert result != pytest.approx(warm_baseline + phantom - actually_paid) + assert result == pytest.approx(expected) + assert result is not None and (result < 0) is loses_money + assert _baseline_usage(modeled).prompt_tokens_details == modeled.prompt_tokens_details -def test_a_cold_switch_never_beats_turning_caching_off(): - """Switching to a cold model makes it write the whole prompt again. That write is a - real cost of switching, so the same traffic must look worse than if caching were off - entirely. - - The baseline is priced as a warm cache even though this request read nothing: a - switch reads nothing precisely because the new model's cache is empty, and staying - on one model would have had the prompt cached already. Gating the warm baseline on - a read charged the baseline a write it would never repeat, which made a cold switch - report a larger saving than no caching at all. - """ - cold_switch = _savings("anthropic/claude-opus-5", "claude-haiku-4-5", _usage(0, 0, 20_000, 1_000)) - caching_off = _savings("anthropic/claude-opus-5", "claude-haiku-4-5", _usage(20_000, 0, 0, 1_000)) - - assert cold_switch < caching_off - - opus = litellm.get_model_info("claude-opus-5", "anthropic") - haiku = litellm.get_model_info("claude-haiku-4-5", "anthropic") - warm_baseline = 20_000 * opus["cache_read_input_token_cost"] + 1_000 * opus["output_cost_per_token"] - actually_paid = 20_000 * haiku["cache_creation_input_token_cost"] + 1_000 * haiku["output_cost_per_token"] - assert cold_switch == pytest.approx(warm_baseline - actually_paid) +@pytest.mark.parametrize("modifier, multiplier", [({}, 1.0), ({"inference_geo": "us"}, 1.1), ({"speed": "fast"}, 2.0)]) +@pytest.mark.parametrize("negotiated", [False, True]) +@pytest.mark.parametrize("provenance", [None, "modeled", "observed_initial"]) +def test_observed_initial_uses_provider_billing_and_effective_rates( + modifier: dict[str, str], multiplier: float, negotiated: bool, + provenance: Literal["modeled", "observed_initial"] | None, +) -> None: + usage: Final = _usage(1000, 2000, 3000, 100).model_copy(update=modifier) + info: Final = litellm.get_model_info("claude-opus-5", "anthropic").copy() + if negotiated: + info["input_cost_per_token"] = 1e-6 + info["output_cost_per_token"] = 2e-6 + info["cache_read_input_token_cost"] = 3e-7 + info["cache_creation_input_token_cost"] = 4e-6 + billed: Final = anthropic_cost_per_token("claude-opus-5", usage, model_info=info) + if negotiated: + assert sum(billed) == pytest.approx(0.0138 * multiplier) + assert compute_autorouter_savings( + "anthropic/claude-opus-5", "claude-opus-5", "anthropic", usage, + selected_info=info, baseline_info=info, baseline_usage=usage, + baseline_deployment_id="same", selected_deployment_id="same", + cost_breakdown={"input_cost": billed[0], "output_cost": billed[1]}, + baseline_provenance=provenance, + ) == 0.0 -def test_moving_one_token_between_cache_buckets_does_not_move_the_answer(): - """A continuing conversation writes a few new tokens and reads the rest. Treating the - presence of a write as the signal for a switch made that ordinary increment flip the - result, so a request reading 19,999 and writing 1 landed somewhere entirely different - from one reading 20,000 and writing none. - """ - reads_nothing = _savings("anthropic/claude-opus-5", "claude-haiku-4-5", _usage(0, 0, 20_000, 1_000)) - reads_one = _savings("anthropic/claude-opus-5", "claude-haiku-4-5", _usage(0, 1, 19_999, 1_000)) - assert reads_one == pytest.approx(reads_nothing, abs=1e-4) - - -def test_multimodal_prompts_are_priced_on_the_baseline_too(): - """The baseline is this same request met by a warm cache, so every field it was - priced on has to survive. Rebuilding the details from the cache buckets alone - dropped the image and audio counts, which priced the baseline as a text-only - request that never ran and shrank the reported saving on multimodal traffic. - """ - details = {"cached_tokens": 0, "cache_creation_tokens": 16_000, "text_tokens": 0, "image_tokens": 4_000} - with_images = Usage( - prompt_tokens=20_000, - completion_tokens=1_000, - total_tokens=21_000, - prompt_tokens_details=details, - ) - baseline = _baseline_usage(with_images, conversation_continuing=True) - - assert baseline.prompt_tokens_details.image_tokens == 4_000, "image tokens must survive into the baseline" - - opus = litellm.get_model_info("claude-opus-5", "anthropic") - priced, _ = generic_cost_per_token(model="claude-opus-5", usage=baseline, custom_llm_provider="anthropic") - text_only = 20_000 * opus["cache_read_input_token_cost"] - assert priced > text_only, "dropping the image tokens undercharges the baseline and hides the saving" - - -def test_the_baseline_is_never_charged_a_cache_write(): - """Carrying the details through must not carry the 5m/1h creation breakdown with - them. `generic_cost_per_token` charges a creation cost whenever that breakdown is - present, even against a zeroed creation count, which would put the phantom write - back on the baseline for every long-cache request. - """ - long_cache = Usage( - prompt_tokens=20_000, - completion_tokens=1_000, - total_tokens=21_000, - prompt_tokens_details={ - "cached_tokens": 0, - "cache_creation_tokens": 20_000, - "text_tokens": 0, - "cache_creation_token_details": {"ephemeral_1h_input_tokens": 20_000}, - }, - ) - baseline = _baseline_usage(long_cache, conversation_continuing=True) - - opus = litellm.get_model_info("claude-opus-5", "anthropic") - priced, _ = generic_cost_per_token(model="claude-opus-5", usage=baseline, custom_llm_provider="anthropic") - assert priced == pytest.approx(20_000 * opus["cache_read_input_token_cost"]), ( - "the baseline reads a warm cache; it never pays to create one" - ) +@pytest.mark.parametrize("model, deployment, known, delta", [ + ("claude-sonnet-5", "same", "observed", 0.0), + ("claude-opus-5", "other", "observed", 0.0), + ("claude-opus-5", "", "observed", 0.0), + ("claude-opus-5", "same", "missing", 0.0), + ("claude-opus-5", "same", "different", 0.0), + ("claude-opus-5", "same", "observed", 0.01), + ("claude-opus-5", "same", "prices", 0.0), + ("claude-opus-5", "same", "unbilled", 0.0), +]) +def test_initial_provenance_cannot_override_mismatched_evidence( + model: str, deployment: str, known: Literal["observed", "missing", "different", "prices", "unbilled"], delta: float, +) -> None: + usage: Final = _usage(1000, 0, 1000, 100) + billed: Final = anthropic_cost_per_token("claude-opus-5", usage) + info: Final = litellm.get_model_info(model, "anthropic").copy() + if known == "prices": + info["cache_read_input_token_cost"] = 0.001 # No reads here: equal charge alone cannot establish equal rates. + assert compute_autorouter_savings( + "claude-opus-5", model, "anthropic", usage, + baseline_usage=(None if known == "missing" else _usage(1000, 1000, 0, 100) if known == "different" else usage), + selected_info=info, + baseline_provenance="observed_initial", + baseline_deployment_id="same", selected_deployment_id=deployment, + cost_breakdown=None if known == "unbilled" else {"input_cost": billed[0] + delta, "output_cost": billed[1]}, + ) is None def test_uncached_request_is_the_plain_rate_difference(): @@ -565,11 +528,12 @@ def test_escalation_reports_its_real_cost(): def test_autorouter_savings_zero_when_model_unchanged(): - assert _savings("claude-opus-5", "claude-opus-5", _usage(3, 500, 12304, 500)) == 0.0 + usage: Final = _usage(3, 500, 12304, 500) + assert _savings("claude-opus-5", "claude-opus-5", usage, usage) == 0.0 -def test_autorouter_savings_unknown_baseline_fails_open_to_zero(): - assert _savings("totally-made-up-model-xyz", "claude-haiku-4-5", _usage(3, 500, 12304, 500)) == 0.0 +def test_autorouter_savings_unknown_baseline_remains_unknown(): + assert _savings("totally-made-up-model-xyz", "claude-haiku-4-5", _usage(3, 500, 12304, 500)) is None def test_autorouter_savings_zero_without_baseline(): @@ -584,9 +548,7 @@ def test_autorouter_savings_zero_without_baseline(): assert result.autorouter == 0.0 -def test_compute_savings_spend_carries_a_losing_switch_through(): - """The signed value must survive into SavingsSpend; clamping it here would put the - dashboard back to only ever showing gains.""" +def test_compute_savings_spend_carries_a_recorded_losing_switch_through(): result = compute_savings_spend( model="claude-haiku-4-5", custom_llm_provider="anthropic", @@ -594,6 +556,7 @@ def test_compute_savings_spend_carries_a_losing_switch_through(): gateway_injected_cache=True, routing_decision={"conversation_continuing": True, "savings_baseline_model": "anthropic/claude-sonnet-5"}, usage_object=_cached_usage_object(), + recorded_autorouter_savings=-0.01, ) assert result.autorouter < 0 @@ -628,18 +591,10 @@ def test_malformed_usage_object_does_not_fail_the_spend_write(): assert result.compression > 0 -def test_the_same_deployment_spelled_two_ways_is_not_a_switch(): - """The spend log records a normalized model name while the baseline arrives as the - operator wrote it in config. Comparing the raw strings makes a request that never - changed model look like a switch, and prices one deployment against itself.""" - # Must be a cached request: the baseline arm is priced against a warm cache and the - # selected arm against what was actually paid, so treating one deployment as two - # charges it a cold-cache write it never took, inventing a loss on a request that - # never changed model. An uncached request prices identically either way and would - # make this assertion vacuous. - usage = _usage(fresh=3, cached=500, written=12304, out=500) - assert _savings("anthropic/claude-opus-5", "claude-opus-5", usage) == 0.0 - assert _savings("claude-opus-5", "anthropic/claude-opus-5", usage) == 0.0 +def test_equal_modeled_usage_is_zero_under_equivalent_model_names() -> None: + usage: Final = _usage(3, 500, 12304, 500) + assert _savings("anthropic/claude-opus-5", "claude-opus-5", usage, usage) == 0.0 + assert _savings("claude-opus-5", "anthropic/claude-opus-5", usage, usage) == 0.0 def test_baseline_is_priced_under_its_own_provider(): @@ -663,99 +618,9 @@ def test_baseline_is_priced_under_its_own_provider(): assert azure > 0 > deepseek -def test_unresolvable_baseline_fails_open_to_zero(): +def test_unresolvable_baseline_remains_unknown(): usage = _usage(fresh=2000, cached=0, written=0, out=500) - assert _savings("no-such-provider-xyz/no-such-model", "claude-haiku-4-5", usage) == 0.0 - - -def test_a_first_turn_is_the_rate_difference_not_a_switch_penalty(): - """Nothing was cached anywhere on a conversation's first turn, so the baseline would - have paid the same cache write. Charging it to the selected arm alone reported a - fraction of the real saving; on this shape roughly 4% of it. - """ - usage = _usage(fresh=0, cached=0, written=20_000, out=1_000) - first_turn = _savings("anthropic/claude-opus-5", "claude-haiku-4-5", usage, continuing=False) - - opus = litellm.get_model_info("claude-opus-5", "anthropic") - haiku = litellm.get_model_info("claude-haiku-4-5", "anthropic") - both_write = (20_000 * opus["cache_creation_input_token_cost"] + 1_000 * opus["output_cost_per_token"]) - ( - 20_000 * haiku["cache_creation_input_token_cost"] + 1_000 * haiku["output_cost_per_token"] - ) - assert first_turn == pytest.approx(both_write) - - mid_conversation = _savings("anthropic/claude-opus-5", "claude-haiku-4-5", usage) - assert first_turn > mid_conversation * 10, "a first turn must not be priced as a switch" - - -def test_a_first_turn_that_saves_money_never_reports_a_loss(): - """The write premium is fixed by prompt size while the saving grows with completion - length, so charging the write to a first turn made short answers over a large cached - prompt read as losses on requests that genuinely saved. That is the shape most likely - to be on the dashboard, and the sign has to be right. - """ - short_answer = _usage(fresh=0, cached=0, written=20_000, out=200) - assert _savings("anthropic/claude-opus-5", "claude-haiku-4-5", short_answer, continuing=False) > 0 - assert _savings("anthropic/claude-opus-5", "claude-haiku-4-5", short_answer) < 0 - - -def test_an_undetermined_conversation_shape_stays_conservative(): - """The default must charge the write. A caller that cannot be read, or a surface the - router never classified, has said nothing about whether the baseline was warm, and a - savings figure must not inflate on a guess. - """ - usage = _usage(fresh=0, cached=0, written=20_000, out=1_000) - defaulted = compute_autorouter_savings( - baseline_model="anthropic/claude-opus-5", - selected_model="claude-haiku-4-5", - selected_provider="anthropic", - usage=usage, - ) - assert defaulted == pytest.approx(_savings("anthropic/claude-opus-5", "claude-haiku-4-5", usage)) - assert defaulted < _savings("anthropic/claude-opus-5", "claude-haiku-4-5", usage, continuing=False) - - -def test_a_continuing_turn_on_the_same_model_writes_its_growth_on_both_arms(): - """A conversation that grew by a few tokens writes those on whatever model serves - it, and they are new to every model, so the baseline would have written them too. - Moving them into the baseline's read bucket forgives it a write it really owes and - shrinks the reported saving on ordinary steady-state traffic. - """ - usage = _usage(fresh=0, cached=19_900, written=100, out=1_000) - opus = litellm.get_model_info("claude-opus-5", "anthropic") - haiku = litellm.get_model_info("claude-haiku-4-5", "anthropic") - - def cost(info: dict) -> float: - return ( - 19_900 * info["cache_read_input_token_cost"] - + 100 * info["cache_creation_input_token_cost"] - + 1_000 * info["output_cost_per_token"] - ) - - both_write_the_growth = cost(opus) - cost(haiku) - assert _savings("anthropic/claude-opus-5", "claude-haiku-4-5", usage) == pytest.approx(both_write_the_growth) - - -def test_a_switch_onto_a_partly_cached_model_still_pays_for_the_write(): - """A model holding a small prefix of this prompt still has to write the rest, and - that write is the switch's cost. Keying the same-model case off reading *anything* - rather than reading *most of it* would hand this request the full rate gap and - inflate the saving by an order of magnitude. - """ - mostly_written = _usage(fresh=0, cached=500, written=19_500, out=1_000) - reported = _savings("anthropic/claude-opus-5", "claude-haiku-4-5", mostly_written) - - opus = litellm.get_model_info("claude-opus-5", "anthropic") - haiku = litellm.get_model_info("claude-haiku-4-5", "anthropic") - if_treated_as_same_model = ( - 500 * opus["cache_read_input_token_cost"] - + 19_500 * opus["cache_creation_input_token_cost"] - + 1_000 * opus["output_cost_per_token"] - ) - ( - 500 * haiku["cache_read_input_token_cost"] - + 19_500 * haiku["cache_creation_input_token_cost"] - + 1_000 * haiku["output_cost_per_token"] - ) - assert reported < if_treated_as_same_model / 10, "a mostly-cold switch must not be priced as a continuation" + assert _savings("no-such-provider-xyz/no-such-model", "claude-haiku-4-5", usage) is None def test_a_baseline_that_prices_caching_implicitly_still_pays_for_its_prompt(): @@ -771,7 +636,7 @@ def test_a_baseline_that_prices_caching_implicitly_still_pays_for_its_prompt(): selected_model="claude-haiku-4-5", selected_provider="anthropic", usage=first_turn, - conversation_continuing=False, + baseline_usage=first_turn, ) gpt5 = litellm.get_model_info("gpt-5", "openai") @@ -807,7 +672,7 @@ def _priced_chat_model_without_cache_read_rate() -> tuple[str, str, str]: usage=_usage(fresh=1_000, cached=0, written=0, out=100), conversation_continuing=True, ) - if priced == 0.0: + if priced is None or priced == 0.0: continue return key, key.removeprefix(f"{provider}/"), provider raise AssertionError("the bundled map has no per-token chat model without a cache-read rate") @@ -825,7 +690,7 @@ def test_a_baseline_with_no_cache_read_rate_is_charged_its_input_rate(): selected_model="claude-haiku-4-5", selected_provider="anthropic", usage=continuing, - conversation_continuing=True, + baseline_usage=_usage(0, 20000, 0, 1000), ) baseline = litellm.get_model_info(baseline_name, baseline_provider) @@ -966,7 +831,7 @@ def test_a_baseline_recorded_on_the_decision_turns_the_driver_on(): compression_saved_tokens=0, gateway_injected_cache=True, routing_decision={"conversation_continuing": True, "savings_baseline_model": "anthropic/claude-opus-5"}, - usage_object=_cached_usage_object(), + usage_object=_usage(12807, 0, 0, 500).model_dump(), ) assert result.autorouter != 0.0 @@ -981,13 +846,13 @@ def test_a_leftover_configured_baseline_does_not_override_the_recorded_one(monke compression_saved_tokens=0, gateway_injected_cache=True, routing_decision={"conversation_continuing": True, "savings_baseline_model": "anthropic/claude-opus-5"}, - usage_object=_cached_usage_object(), + usage_object=_usage(12807, 0, 0, 500).model_dump(), ) against_opus = compute_autorouter_savings( baseline_model="anthropic/claude-opus-5", selected_model="claude-haiku-4-5", selected_provider="anthropic", - usage=Usage(**_cached_usage_object()), + usage=_usage(12807, 0, 0, 500), ) assert result.autorouter == against_opus @@ -1054,12 +919,12 @@ def test_prompt_caching_prices_at_the_deployment_rate_not_the_public_one(): ("baseline", "selected", 2.0, None, 0.0, -0.015), ("baseline", "selected", 1.0, None, 0.0, 0.0), ("baseline", "selected", 0.1, 0.004, 0.001, 0.01), - ("baseline", "baseline", 0.1, 0.004, 0.001, -0.001), - (None, "selected", 0.1, None, 0.0, 0.0), - ("baseline", None, 0.1, None, 0.0, 0.0), + ("baseline", "baseline", 0.1, 0.004, 0.001, 0.01), + (None, "selected", 0.1, None, 0.0, 0.006), + ("baseline", None, 0.1, None, 0.0, 0.0075), (None, None, 0.1, None, 0.0, 0.0), - ("", "selected", 0.1, None, 0.0, 0.0), - ("baseline", "", 0.1, None, 0.0, 0.0), + ("", "selected", 0.1, None, 0.0, 0.006), + ("baseline", "", 0.1, None, 0.0, 0.0075), ], ) def test_autorouter_savings_distinguishes_priced_deployments( @@ -1172,7 +1037,7 @@ def test_a_recorded_baseline_deployment_prices_at_its_configured_rate(): compression_saved_tokens=0, gateway_injected_cache=True, routing_decision=decision, - usage_object=_cached_usage_object(), + usage_object=_usage(12807, 0, 0, 500).model_dump(), llm_router=lambda: router, ) at_public_rate = compute_savings_spend( @@ -1181,7 +1046,7 @@ def test_a_recorded_baseline_deployment_prices_at_its_configured_rate(): compression_saved_tokens=0, gateway_injected_cache=True, routing_decision={k: v for k, v in decision.items() if k != "savings_baseline_deployment_id"}, - usage_object=_cached_usage_object(), + usage_object=_usage(12807, 0, 0, 500).model_dump(), llm_router=lambda: router, ) assert with_deployment_rate.autorouter > at_public_rate.autorouter @@ -1234,9 +1099,8 @@ def test_a_boolean_is_not_a_recorded_savings_figure(): assert result.autorouter == 0.0 -def test_rows_written_before_the_field_shipped_recompute(): - """No recorded figure means the row predates the logging-path stamp; the writer - recomputes exactly what the one shared helper would have recorded.""" +@pytest.mark.parametrize("continuing", [False, True]) +def test_legacy_cache_rows_without_an_estimate_do_not_invent_a_new_figure(continuing: bool) -> None: from litellm.proxy.spend_tracking.savings import autorouter_savings_for_request recomputed = compute_savings_spend( @@ -1244,17 +1108,17 @@ def test_rows_written_before_the_field_shipped_recompute(): custom_llm_provider="anthropic", compression_saved_tokens=0, gateway_injected_cache=False, - routing_decision=_routed_decision(), + routing_decision={**_routed_decision(), "conversation_continuing": continuing}, usage_object=_cached_usage_object(), ) direct = autorouter_savings_for_request( model="claude-haiku-4-5", custom_llm_provider="anthropic", - routing_decision=_routed_decision(), + routing_decision={**_routed_decision(), "conversation_continuing": continuing}, usage_object=_cached_usage_object(), ) - assert direct is not None and direct != 0.0 - assert recomputed.autorouter == direct + assert direct is None + assert recomputed.autorouter == 0.0 def test_driver_off_is_none_not_zero_for_the_request_helper(): @@ -1294,7 +1158,7 @@ def test_logging_payload_never_stamps_internal_calls(): model="claude-haiku-4-5", custom_llm_provider="anthropic", model_id=None, - usage_object=_cached_usage_object(), + usage_object=_usage(12807, 0, 0, 500).model_dump(), cost_breakdown=None, ) assert stamped is not None and stamped != 0.0 @@ -1304,7 +1168,7 @@ def test_logging_payload_never_stamps_internal_calls(): model="claude-haiku-4-5", custom_llm_provider="anthropic", model_id=None, - usage_object=_cached_usage_object(), + usage_object=_usage(12807, 0, 0, 500).model_dump(), cost_breakdown=None, ) assert internal is None @@ -1320,13 +1184,13 @@ def test_savings_are_net_of_a_priced_classifier(): model="claude-haiku-4-5", custom_llm_provider="anthropic", routing_decision=_routed_decision(), - usage_object=_cached_usage_object(), + usage_object=_usage(12807, 0, 0, 500).model_dump(), ) net = autorouter_savings_for_request( model="claude-haiku-4-5", custom_llm_provider="anthropic", routing_decision={**_routed_decision(), "classifier_cost": 0.005}, - usage_object=_cached_usage_object(), + usage_object=_usage(12807, 0, 0, 500).model_dump(), ) assert gross is not None and net == pytest.approx(gross - 0.005) @@ -1339,13 +1203,13 @@ def test_an_unpriced_classifier_deducts_nothing(classifier_cost: object): model="claude-haiku-4-5", custom_llm_provider="anthropic", routing_decision=_routed_decision(), - usage_object=_cached_usage_object(), + usage_object=_usage(12807, 0, 0, 500).model_dump(), ) with_cost_field = autorouter_savings_for_request( model="claude-haiku-4-5", custom_llm_provider="anthropic", routing_decision={**_routed_decision(), "classifier_cost": classifier_cost}, - usage_object=_cached_usage_object(), + usage_object=_usage(12807, 0, 0, 500).model_dump(), ) assert with_cost_field == gross @@ -1454,4 +1318,41 @@ def test_marks_gateway_injection_credits_only_the_deployment_that_was_injected() assert marks_gateway_injection({"litellm_gateway_injected_cache": ""}, "dep-a") is True assert marks_gateway_injection({"litellm_gateway_injected_cache": ""}, None) is True assert marks_gateway_injection({"litellm_call_id": "c1"}, "dep-a") is False + + +@pytest.mark.parametrize("classifier", [0.0, 0.02]) +def test_observed_baseline_keeps_both_costs_and_classifier_overhead(classifier: float) -> None: + from litellm.proxy.spend_tracking.savings import BaselineCostSnapshot, price_baseline_comparison + + snapshot: Final = BaselineCostSnapshot( + model="baseline", provider="anthropic", prices=None, + actual_spend=0.17, classifier_cost=classifier, + ) + restored: Final = BaselineCostSnapshot.model_validate_json(snapshot.model_dump_json()) + result: Final = price_baseline_comparison(restored, Usage(prompt_tokens=100, completion_tokens=10), "observed_identical") + assert result is not None + assert result.baseline == snapshot.actual_spend + assert result.actual == snapshot.actual_spend + classifier + assert result.savings == pytest.approx(-classifier) + assert price_baseline_comparison(restored, None, None) is None + + +def test_modeled_baseline_uses_recorded_prices_and_preserves_other_actual_charges() -> None: + from litellm.proxy.spend_tracking.savings import BaselineCostSnapshot, price_baseline_comparison + + snapshot: Final = BaselineCostSnapshot( + model="claude-opus-5", provider="anthropic", + prices={ + **litellm.get_model_info("claude-opus-5", custom_llm_provider="anthropic"), + "input_cost_per_token": 0.001, "output_cost_per_token": 0.002, + }, + actual_token_cost=0.2, actual_spend=0.23, classifier_cost=0.01, + ) + usage: Final = Usage(prompt_tokens=100, completion_tokens=10) + result: Final = price_baseline_comparison(snapshot, usage, "modeled") + assert result is not None + assert result.actual == pytest.approx(0.24) + assert result.baseline == pytest.approx(0.12 + 0.03) + assert result.savings == pytest.approx(-0.09) + assert price_baseline_comparison(snapshot.model_copy(update={"prices": None}), usage, "modeled") is None assert marks_gateway_injection({"litellm_gateway_injected_cache": True}, "dep-a") is False diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 8d15fb094d5..9de6679472e 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3745,7 +3745,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "azure_spillover": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, @@ -3841,7 +3841,7 @@ class TestSpendLogsPayload: "model": "claude-4-sonnet-20250514", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "azure_spillover": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, @@ -3935,7 +3935,7 @@ class TestSpendLogsPayload: "model": "claude-4-sonnet-20250514", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "attempted_fallbacks": 0, "original_model_group": "my-anthropic-model-group", "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "azure_spillover": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "attempted_fallbacks": 0, "original_model_group": "my-anthropic-model-group", "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 0004711954a..f471e3f8fbb 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -3,6 +3,7 @@ import datetime import json from collections.abc import Callable, Mapping from datetime import timezone +from types import MappingProxyType from typing import Any, Final, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -5208,3 +5209,17 @@ def test_azure_spillover_absent_without_spillover_headers(): ) metadata = json.loads(payload["metadata"]) assert metadata["azure_spillover"] is None + + +def test_baseline_estimate_metadata_comes_from_the_logging_stamp() -> None: + supplied: Final = MappingProxyType({"version": 1, "status": "estimated", "reason": "caller_supplied"}) + recorded: Final = MappingProxyType({"version": 1, "status": "unknown", "reason": "history_unavailable"}) + result: Final = _get_spend_logs_metadata( + {"autorouter_savings": 999.0, "autorouter_savings_estimate": supplied}, # mutable-ok: legacy metadata helper accepts dicts + autorouter_savings=None, + autorouter_savings_estimate=recorded, + ) + assert result["autorouter_savings"] is None + assert result["autorouter_savings_estimate"] == recorded + absent: Final = _get_spend_logs_metadata({"autorouter_savings_estimate": supplied}) # mutable-ok: legacy metadata helper accepts dicts + assert absent["autorouter_savings_estimate"] is None diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 14d4929d27e..0f57af7f82c 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3324,6 +3324,38 @@ class TestWebSearchInterceptionSettingsEndpoints: assert resp.status_code == 200, resp.text assert resp.json()["values"]["enabled"] is True + def test_get_flags_a_pod_that_has_not_applied_the_stored_setting( + self, mock_proxy_config, mock_auth, monkeypatch + ): + import litellm + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", object()) + monkeypatch.setattr(litellm, "callbacks", []) + mock_proxy_config["config"]["litellm_settings"]["websearch_interception_params"] = {"enabled": True} + + resp = client.get("/get/websearch_interception_settings") + + assert resp.status_code == 200, resp.text + assert resp.json()["values"]["enabled"] is True + assert resp.json()["active_on_this_pod"] is False + + def test_get_reports_the_pod_as_active_once_the_callback_is_registered( + self, mock_proxy_config, mock_auth, monkeypatch + ): + import litellm + from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", object()) + monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger(search_tool_name="running")]) + mock_proxy_config["config"]["litellm_settings"]["websearch_interception_params"] = {"enabled": True} + + resp = client.get("/get/websearch_interception_settings") + + assert resp.status_code == 200, resp.text + assert resp.json()["active_on_this_pod"] is True + def test_get_reports_no_database_instead_of_empty_settings(self, mock_proxy_config, mock_auth, monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py index d7a6124dd97..a4bb7d63548 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio from datetime import datetime +from typing import Any, Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -13,7 +14,7 @@ from fastapi import HTTPException import litellm from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import AlertType, ProxyErrorTypes +from litellm.proxy._types import AlertType, ProxyErrorTypes, UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging @@ -156,6 +157,23 @@ async def test_post_call_failure_hook_non_http_exception_in_callback_swallowed( assert out is None +@pytest.mark.asyncio +@pytest.mark.parametrize("logging_value", (None, "caller-controlled", {"baseline_cache_context": "untrusted"})) # mutable-ok: emulate an untrusted JSON request field +async def test_terminal_baseline_cleanup_ignores_missing_or_untrusted_logging( + proxy_logging: ProxyLogging, monkeypatch: pytest.MonkeyPatch, logging_value: object +) -> None: + monkeypatch.setattr(litellm, "callbacks", ()) + proxy_logging.alert_types = [] # mutable-ok: disable optional alert sinks for this boundary test # rebind-ok: isolate the fixture-owned alert configuration + request_data: Final = {"litellm_call_id": "untrusted-logging", "litellm_logging_obj": logging_value} # mutable-ok: the production failure owner removes internal fields in place + result: Final = await proxy_logging.post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # exercise the existing proxy terminal owner with its legacy request dictionary contract + request_data=request_data, + original_exception=ValueError("original provider failure"), + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/messages"), + ) + assert result is None + assert "litellm_logging_obj" not in request_data + + # --------------------------------------------------------------------------- # _handle_logging_proxy_only_error # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index b6b4a8072fa..63fde9b2b8f 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -2269,6 +2269,10 @@ class TestAutoRouterSessionRepository: "classifier_cost": 0.01, "tier_turns": {"complex": 3}, "baseline_models": {"anthropic/claude-opus-5": 3}, + "savings_estimated_turns": 3, + "savings_estimated_actual_spend": 0.14, + "savings_estimated_saved_spend": 0.24, + "savings_estimated_baseline_models": {"anthropic/claude-opus-5": 3}, } @staticmethod @@ -2295,6 +2299,9 @@ class TestAutoRouterSessionRepository: assert (row.router_name, row.turns, row.spend, row.saved_spend) == ("claude-auto", 3, 0.14, 0.24) assert row.baseline_models == {"anthropic/claude-opus-5": 3} assert row.baseline_model == "anthropic/claude-opus-5" + assert row.savings_estimated_turns == 3 + assert row.savings_estimated_actual_spend == 0.14 + assert row.savings_estimated_saved_spend == 0.24 @pytest.mark.asyncio async def test_find_latest_for_key_is_none_when_the_key_wrote_no_such_session(self): diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 9b25c869f1c..ecd25ff654f 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -3721,7 +3721,11 @@ class TestLLMClassifier: assert outcome.score is not None @pytest.mark.asyncio - async def test_heuristic_v2_routes_directly_to_predicted_builtin_tier(self, mock_router_instance): + @pytest.mark.parametrize("redact", (False, True)) + async def test_heuristic_v2_routes_directly_to_predicted_builtin_tier( + self, mock_router_instance: MagicMock, redact: bool, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "turn_off_message_logging", redact) router = ComplexityRouter( model_name="tier-router", litellm_router_instance=mock_router_instance, @@ -3754,6 +3758,21 @@ class TestLLMClassifier: "tier-probability:complex=0.892157", "tier-probability:reasoning=0.980392", ] + redacted: Final = Router._redact_prompt_text_if_needed( + request_kwargs={}, routing_decision=response.routing_decision + ) + assert ("signals" in redacted) is not redact + assert redacted["heuristic_v2_forecast"] == { + "probabilities": { + "SIMPLE": 11 / 102, + "MEDIUM": 21 / 102, + "COMPLEX": 91 / 102, + "REASONING": 100 / 102, + }, + "threshold": 0.8, + "predicted_tier": "COMPLEX", + "request_type": "general", + } def test_heuristic_v2_needs_no_classifier_model(self): config = ComplexityRouterConfig(classifier_type="heuristic_v2") @@ -8879,13 +8898,29 @@ class TestRoutingDecisionSurvivesToSpendLogOnEveryMetadataShape: ], ) @pytest.mark.asyncio - async def test_decision_reaches_the_spend_log_payload(self, request_kwargs, expected_bucket): + @pytest.mark.parametrize("classifier_type", ("heuristic", "heuristic_v2")) + async def test_decision_reaches_the_spend_log_payload(self, request_kwargs, expected_bucket, classifier_type): import datetime import json from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload - router = Router(model_list=self.MODEL_LIST) + model_list: Final = [ + { + **row, + "litellm_params": { + **row["litellm_params"], + "complexity_router_config": { + **row["litellm_params"]["complexity_router_config"], + "classifier_type": classifier_type, + }, + }, + } + if row["model_name"] == "smart-router" + else row + for row in self.MODEL_LIST + ] + router = Router(model_list=model_list) response = await router.async_pre_routing_hook( model="smart-router", request_kwargs=request_kwargs, @@ -8915,6 +8950,15 @@ class TestRoutingDecisionSurvivesToSpendLogOnEveryMetadataShape: persisted = json.loads(payload["metadata"])["routing_decision"] assert persisted is not None, f"routing_decision dropped for {expected_bucket}" assert persisted["router_model_name"] == "smart-router" + if classifier_type == "heuristic_v2": + assert persisted["heuristic_v2_forecast"] == request_kwargs[expected_bucket]["routing_decision"][ + "heuristic_v2_forecast" + ] + assert set(persisted["heuristic_v2_forecast"]["probabilities"]) == { + "SIMPLE", "MEDIUM", "COMPLEX", "REASONING" + } + else: + assert "heuristic_v2_forecast" not in persisted class TestRoutingDecisionIsPerAttempt: @@ -9001,19 +9045,26 @@ class TestRecordRoutingDecision: Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) assert request_kwargs == {} - def test_clearing_the_decision_takes_the_savings_facts_with_it(self): + def test_clearing_the_decision_takes_the_savings_facts_with_it(self) -> None: """A fallback to a plain model group re-enters the hook with the same `request_kwargs`. The baseline and the conversation shape ride inside the decision rather than beside it, so one clear cannot leave either behind and attribute an auto-router saving to a deployment that never routed.""" - decision = { + from litellm.types.router import BaselineRouteStamp + + decision: Final = { "router_model_name": "smart-router", "router_type": "complexity", "routed_model": "gpt-4o-mini", "savings_baseline_model": "anthropic/claude-opus-5", + "savings_baseline_deployment_id": "opus-deployment", "conversation_continuing": False, } - request_kwargs: Dict = {"litellm_metadata": {"routing_decision": decision}} + request_kwargs: Final[dict[str, dict[str, object]]] = {"litellm_metadata": {}} + Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=decision) + stamp: Final = request_kwargs["litellm_metadata"]["_autorouter_baseline_route"] + assert isinstance(stamp, BaselineRouteStamp) + assert stamp.baseline_deployment_id == "opus-deployment" Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) assert request_kwargs["litellm_metadata"] == {} @@ -14104,6 +14155,33 @@ class TestModalityRouting: BASE_TIERS = {"SIMPLE": "text-cheap", "MEDIUM": "vision-mid", "COMPLEX": "vision-big"} BASE_VISION = {"text-cheap": False, "vision-mid": True, "vision-big": True, "vision-default": True} + @pytest.mark.asyncio + async def test_modality_escalation_preserves_the_original_heuristic_v2_forecast( + self, mock_router_instance: MagicMock + ) -> None: + router: Final = self._router( + mock_router_instance, + { + "classifier_type": "heuristic_v2", + "heuristic_v2_artifact": _heuristic_v2_artifact(), + "tiers": {"COMPLEX": "text-cheap", "REASONING": "vision-big"}, + "modality_routing": True, + }, + self.BASE_VISION, + ) + original: Final = await router.aclassify("What color is this?") + result: Final = await router.async_pre_routing_hook( + model="m", request_kwargs={}, messages=self.IMAGE_MESSAGE + ) + + assert original.heuristic_v2_forecast is not None + assert result is not None and result.routing_decision is not None + assert result.model == "vision-big" + assert result.routing_decision["cause"] == "modality_escalation" + assert result.routing_decision["tier"] == "REASONING" + assert result.routing_decision["heuristic_v2_forecast"] == original.heuristic_v2_forecast + assert result.routing_decision["heuristic_v2_forecast"]["predicted_tier"] == "COMPLEX" + @staticmethod def _router(mock_router_instance, config, vision_by_model): """vision_by_model: model name -> True/False (deployment model_info) or None (undeclared).""" @@ -14479,6 +14557,69 @@ class TestModalityRouting: @pytest.mark.usefixtures("local_model_cost_map") class TestHealthFallbackDispatch: + @pytest.mark.asyncio + @pytest.mark.parametrize("peer", (True, False), ids=("peer_failover", "default_fallback")) + async def test_health_rewrites_preserve_the_original_heuristic_v2_forecast(self, peer: bool) -> None: + router: Final = self._router( + config={ + "classifier_type": "heuristic_v2", + "heuristic_v2_artifact": _heuristic_v2_artifact(), + "tiers": {"COMPLEX": ["primary", "peer"] if peer else "primary"}, + } + ) + + def select_primary(models: Sequence[str]) -> str: + return max(models) + + with patch( # test-quality-ok: force initial classification onto the failing group in a mixed tier pool + "litellm.router_strategy.complexity_router.complexity_router.random.choice", + side_effect=select_primary, + ): + original: Final = await router.async_pre_routing_hook( + model="health-router", request_kwargs={}, messages=[{"role": "user", "content": "Hello!"}] + ) + self._unavailable(router, "primary-id", "cooldown") + result: Final = await router.async_pre_routing_hook( + model="health-router", request_kwargs={}, messages=[{"role": "user", "content": "Hello!"}] + ) + + assert original is not None and original.routing_decision is not None + assert original.model == "primary" + assert original.routing_decision["cause"] == "heuristic_v2" + assert result is not None and result.routing_decision is not None + assert result.model == ("peer" if peer else "fallback") + assert result.routing_decision["cause"] == ("health_failover" if peer else "health_default_fallback") + assert result.routing_decision["heuristic_v2_forecast"] == original.routing_decision["heuristic_v2_forecast"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("pinned", (False, True), ids=("keyword_bypass", "session_pin")) + async def test_heuristic_v2_bypasses_have_no_fabricated_forecast(self, pinned: bool) -> None: + router: Final = self._router( + session=pinned, + config={ + "classifier_type": "heuristic_v2", + "heuristic_v2_artifact": _heuristic_v2_artifact(), + "tiers": {"COMPLEX": "primary"}, + "keyword_tier_rules": [{"keywords": ["quick lookup"], "tier": "COMPLEX"}], + }, + ) + original: Final = await router.async_pre_routing_hook( + model="health-router", + request_kwargs={"metadata": {"session_id": "v2-forecast"}}, + messages=[{"role": "user", "content": "Hello!"}], + ) + result: Final = await router.async_pre_routing_hook( + model="health-router", + request_kwargs={"metadata": {"session_id": "v2-forecast"}}, + messages=[{"role": "user", "content": "quick lookup"}], + ) + + assert original is not None and original.routing_decision is not None + assert "heuristic_v2_forecast" in original.routing_decision + assert result is not None and result.routing_decision is not None + assert result.routing_decision["cause"] == ("session_affinity_pin" if pinned else "literal_keyword_match") + assert "heuristic_v2_forecast" not in result.routing_decision + @pytest.fixture(autouse=True) def httpx_transport(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) diff --git a/tests/test_litellm/router_strategy/test_fuse_presets.py b/tests/test_litellm/router_strategy/test_fuse_presets.py new file mode 100644 index 00000000000..0cd4b1f660b --- /dev/null +++ b/tests/test_litellm/router_strategy/test_fuse_presets.py @@ -0,0 +1,69 @@ +import json +from hashlib import sha256 +from importlib.resources import files +from typing import Final, Literal + +import pytest +from pydantic import ValidationError + +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets, resolve_fuse_profile + + +def test_catalog_is_loaded_once_and_preserves_bundled_content() -> None: + get_fuse_presets.cache_clear() + first: Final = get_fuse_presets() + second: Final = get_fuse_presets() + assert first is second + bundled: Final = json.loads( + files("litellm.router_strategy.complexity_router").joinpath("fuse_presets.json").read_text(encoding="utf-8") + ) + assert first.model_dump(mode="json") == bundled + entries: Final = (*first.models, *first.harnesses) + assert len({entry.id for entry in entries}) == len(entries) + assert all(entry.sources and all(source.startswith("https://") for source in entry.sources) for entry in entries) + + +@pytest.mark.parametrize( + ("kind", "preset_id", "expected_digest"), + ( + ("model", "gpt-6-astra-v1", "a9403b0c00ea64081b7b08b5b968850670f3a047d219a7e0668f2169146ae96e"), + ("model", "gpt-5.6-sol-v1", "2b91a6c43e0e93183aaaf9c355e1bbb8ed2e9817aab6b0c2f50148f53a23247b"), + ("model", "gpt-5.6-luna-v1", "fff94a9e01bf4519798d5be4e76a3f9d57b75a2d9966a59dc92cbfeb5cd08d07"), + ("model", "gpt-5.6-terra-v1", "75de040f3bea841fa4764885738303893ee7ac0804aed1e932cd3959185ff893"), + ("model", "claude-haiku-4-5-v1", "91c1920953073462b6b70ef810596a5325f08286b5e62630cff47938fc4157db"), + ("model", "claude-sonnet-5-v1", "133f4414c644a707cd8cf565a486153856f4836ca4e4f75ee0553f2b7a1e3663"), + ("model", "claude-opus-5-v1", "9cbfcae45d2e3a2575e44ce5adf618f56614abff4b3221d35900c647200b99ef"), + ("model", "claude-fable-5-v1", "25c275d7403f1572ffb4fe899d5feecd9a434ebdc37b4dd9ef601a8ecf4850fc"), + ("model", "claude-fable-5-1-v1", "37693107c878ab6266530395bbdc2d2813d676d179bf281d05e5a5ec1b9d4c60"), + ("harness", "unspecified-v1", "d9eb30b61509456f0c71ca805b33d821cab6605578d567a29ab421d8f602ce7b"), + ("harness", "claude-code-v1", "7ee8e9d50f1cf44a8a58461efff66d6182f245d25499702c144d1c642c101ed9"), + ("harness", "codex-cli-v1", "0678047e34562ef05b5e2fba099c1f9e5876304f7eaf3b0d8c3809e707eb3311"), + ("harness", "opencode-v1", "8b6cc240d90091ac2ef9b374b535f981a55abb91e25d4c04fdb9fc206eeb907e"), + ("harness", "mini-swe-agent-v1", "21e2dc4a8a2320a5a554a498b30516326dc3592ebf20b0f4db20ddb33e879a39"), + ), +) +def test_existing_preset_text_is_unchanged( + kind: Literal["model", "harness"], preset_id: str, expected_digest: str +) -> None: + text: Final = resolve_fuse_profile(None, preset_id, kind) + assert text is not None + assert sha256(text.encode("utf-8")).hexdigest() == expected_digest + + +def test_every_catalog_entry_resolves_without_changing_custom_ownership() -> None: + catalog: Final = get_fuse_presets() + for entry in catalog.models: + assert resolve_fuse_profile(None, entry.id, "model") == entry.text + assert resolve_fuse_profile("Custom text", entry.id, "model") == "Custom text" + for entry in catalog.harnesses: + assert resolve_fuse_profile(None, entry.id, "harness") == entry.text + assert resolve_fuse_profile("Custom text", entry.id, "harness") == "Custom text" + assert resolve_fuse_profile("Custom text", None, "model") == "Custom text" + assert resolve_fuse_profile("Custom text", None, "harness") == "Custom text" + + +def test_cached_catalog_and_records_cannot_be_modified() -> None: + catalog: Final = get_fuse_presets() + for record, field in ((catalog, "version"), (catalog.models[0], "text"), (catalog.harnesses[0], "text")): + with pytest.raises(ValidationError, match="frozen"): + setattr(record, field, "Changed") diff --git a/tests/test_litellm/router_strategy/test_llm_v2.py b/tests/test_litellm/router_strategy/test_llm_v2.py index 5447c8b43ce..27d31cbe640 100644 --- a/tests/test_litellm/router_strategy/test_llm_v2.py +++ b/tests/test_litellm/router_strategy/test_llm_v2.py @@ -11,8 +11,10 @@ from litellm import ModelResponse, Router from litellm.caching.dual_cache import DualCache from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, ComplexityTier +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets from litellm.router_strategy.complexity_router.llm_v2 import ( LLM_V2_PROMPT_VERSION, + LLM_V2_SYSTEM_PROMPT, LLMV2Calibration, LLMV2Config, LLMV2ProbabilityCalibration, @@ -174,6 +176,114 @@ def test_invalid_forecast_settings_are_rejected(overrides: dict[str, object]) -> LLMV2Config.model_validate({**base.model_dump(), **overrides}) +def _preset_config(**overrides: object) -> LLMV2Config: + catalog: Final = get_fuse_presets() + return LLMV2Config.model_validate( + { + "efficient_profile_preset": catalog.models[0].id, + "capable_profile_preset": catalog.models[-1].id, + "harness_preset": catalog.harnesses[-1].id, + "max_quality_gap": 0.05, + **overrides, + } + ) + + +def test_preset_roundtrip_keeps_references_without_materializing_text() -> None: + config: Final = _preset_config() + serialized: Final = config.model_dump(exclude_none=True) + assert serialized["efficient_profile_preset"] == config.efficient_profile_preset + assert serialized["capable_profile_preset"] == config.capable_profile_preset + assert serialized["harness_preset"] == config.harness_preset + assert not {"efficient_profile", "capable_profile", "harness"}.intersection(serialized) + assert LLMV2Config.model_validate(config.model_dump()) == config + assert LLMV2Config.model_validate_json(config.model_dump_json()) == config + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_preset_explicit_override_wins_and_survives_roundtrip(field: str) -> None: + config: Final = _preset_config(**{field: " Operator description "}) + roundtrip: Final = LLMV2Config.model_validate_json(config.model_dump_json()) + assert roundtrip.model_dump()[field] == "Operator description" + assert roundtrip.efficient_profile_preset == config.efficient_profile_preset + assert roundtrip.capable_profile_preset == config.capable_profile_preset + assert roundtrip.harness_preset == config.harness_preset + payload: Final = json.loads( + roundtrip.system_prompt("opaque-efficient", "opaque-capable").split("Configured solver profiles:\n")[1] + ) + if field == "harness": + assert payload["harness"] == "Operator description" + else: + assert payload[field.removesuffix("_profile")]["profile"] == "Operator description" + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +@pytest.mark.parametrize("invalid", ("", " \n\t", "x" * 4001)) +def test_preset_does_not_bypass_supplied_text_bounds(field: str, invalid: str) -> None: + with pytest.raises(ValidationError, match=field): + _preset_config(**{field: invalid}) + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +@pytest.mark.parametrize("override", (None, "Custom override")) +@pytest.mark.parametrize("invalid_id", ("missing-v1", "")) +def test_preset_unknown_reference_rejects_even_when_overridden( + field: str, override: str | None, invalid_id: str +) -> None: + with pytest.raises(ValidationError, match=f"{field}.*preset"): + _preset_config(**{field: override, f"{field}_preset": invalid_id}) + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_preset_missing_text_and_reference_rejects(field: str) -> None: + with pytest.raises(ValidationError, match=field): + _preset_config(**{f"{field}_preset": None}) + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_preset_reference_rejects_the_wrong_catalog_kind(field: str) -> None: + catalog: Final = get_fuse_presets() + wrong_id: Final = catalog.models[0].id if field == "harness" else catalog.harnesses[0].id + with pytest.raises(ValidationError, match=field): + _preset_config(**{f"{field}_preset": wrong_id}) + + +@pytest.mark.parametrize("mode", ("json_schema", "json_object")) +def test_custom_profile_prompt_bytes_are_unchanged(mode: str) -> None: + base: Final = _config().llm_v2_config + assert base is not None + config: Final = LLMV2Config.model_validate({**base.model_dump(), "response_format": mode}) + old_payload: Final = { + "prompt_version": LLM_V2_PROMPT_VERSION, + "harness": config.harness, + "efficient": {"model": "opaque-efficient", "profile": config.efficient_profile}, + "capable": {"model": "opaque-capable", "profile": config.capable_profile}, + } + schema: Final = ( + "\n\nResponse JSON schema:\n" + json.dumps(LLMV2Verdict.model_json_schema()) if mode == "json_object" else "" + ) + assert config.system_prompt("opaque-efficient", "opaque-capable") == ( + LLM_V2_SYSTEM_PROMPT + "\n\nConfigured solver profiles:\n" + json.dumps(old_payload) + schema + ) + + +@pytest.mark.asyncio +async def test_preset_router_passes_catalog_text_and_opaque_group_names_to_judge() -> None: + catalog: Final = get_fuse_presets() + config: Final = _config(llm_v2_config=_preset_config().model_dump()) + router, client = _router(_verdict().model_dump_json(), config) + outcome: Final = await router.aclassify("Complete the supplied task") + assert outcome.tier == ComplexityTier.SIMPLE + prompt: Final = client.acompletion.call_args.kwargs["messages"][0]["content"] + payload: Final = json.loads(prompt.split("Configured solver profiles:\n")[1]) + assert payload == { + "prompt_version": LLM_V2_PROMPT_VERSION, + "harness": catalog.harnesses[-1].text, + "efficient": {"model": "efficient", "profile": catalog.models[0].text}, + "capable": {"model": "capable", "profile": catalog.models[-1].text}, + } + + @pytest.mark.asyncio async def test_one_judge_fuses_whole_task_and_keeps_caller_text_out_of_system_prompt() -> None: router, client = _router(_verdict().model_dump_json()) diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py index 3dcb8d5af94..7d59a0590f2 100644 --- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py +++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py @@ -1,7 +1,10 @@ from collections.abc import Mapping +from typing import Final import pytest +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets + from litellm.router_utils.auto_router_model_naming import ( carries_complexity_router_settings, classify_strategy_router_model, @@ -171,6 +174,52 @@ def test_validate_accepts_loadable_complexity_config(complexity_router_config): assert validate_complexity_router_config_write(complexity_router_config=complexity_router_config) is None +def _fuse_write_config(profiles: Mapping[str, object]) -> Mapping[str, object]: + return { + "classifier_type": "llm_v2", + "classifier_llm_config": {"model": "judge"}, + "tiers": {"SIMPLE": ["opaque-efficient"], "REASONING": ["opaque-capable"]}, + "llm_v2_config": {"max_quality_gap": 0.05, **profiles}, + } + + +def test_fuse_write_accepts_presets_and_custom_text_with_the_same_entitlement() -> None: + catalog: Final = get_fuse_presets() + presets: Final = _fuse_write_config( + { + "efficient_profile_preset": catalog.models[0].id, + "capable_profile_preset": catalog.models[-1].id, + "harness_preset": catalog.harnesses[0].id, + } + ) + custom: Final = _fuse_write_config( + { + "efficient_profile": catalog.models[0].text, + "capable_profile": catalog.models[-1].text, + "harness": catalog.harnesses[0].text, + } + ) + assert validate_complexity_router_config_write(presets) is None + assert validate_complexity_router_config_write(custom) is None + assert claimed_capability(presets) is claimed_capability(custom) + assert claimed_capability(presets) is not None + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_fuse_write_rejects_unknown_preset_even_with_custom_text(field: str) -> None: + config: Final = _fuse_write_config( + { + "efficient_profile": "Custom efficient solver", + "capable_profile": "Custom capable solver", + "harness": "Custom runtime", + f"{field}_preset": "unknown-v1", + } + ) + violation: Final = validate_complexity_router_config_write(config) + assert violation is not None + assert f"{field}_preset" in violation + + def test_naming_check_ignores_the_config_entirely(): """The naming contract and the config's contents are separate questions with separate owners; a write may carry a config without naming a model, so neither can stand in for the other.""" diff --git a/tests/test_litellm/rust_bridge/test_legacy_callbacks.py b/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py similarity index 91% rename from tests/test_litellm/rust_bridge/test_legacy_callbacks.py rename to tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py index a0906c7c5be..1f6a214398a 100644 --- a/tests/test_litellm/rust_bridge/test_legacy_callbacks.py +++ b/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py @@ -10,8 +10,8 @@ from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.rust_bridge import legacy_callbacks as legacy -from litellm.rust_bridge.legacy_callbacks import check_limits, setup +from litellm.rust_bridge import callbacks_legacy_python as legacy +from litellm.rust_bridge.callbacks_legacy_python import check_limits, setup _OCR_KWARGS: Final = MappingProxyType( { @@ -81,7 +81,9 @@ def test_setup_builds_a_logger_when_none_is_supplied(call_type: str, kwargs: Map assert result.logger.litellm_call_id == result.kwargs["litellm_call_id"] -CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/callbacks-legacy/python_contract.json" +CONTRACT_PATH: Final = ( + Path(__file__).parents[3] / "litellm-rust/crates/callbacks-legacy-python/python_contract.json" +) def test_the_rust_contract_matches_the_shim_signatures() -> None: diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b40c10de428..3336ad6d33a 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -869,6 +869,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_pdf_input": {"type": "boolean"}, "prompt_cache_min_tokens": {"type": "number"}, "supports_prompt_cache_breakpoint": {"type": "boolean"}, + "supports_thinking_cache_preservation": {"type": "boolean"}, "supports_prompt_caching": {"type": "boolean"}, "supports_response_schema": {"type": "boolean"}, "supports_system_messages": {"type": "boolean"}, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx index 2820a9dce83..5c7453c1394 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx @@ -68,6 +68,8 @@ const totals = (overrides: Partial = {}): Totals => ({ avg_session_seconds: 7560, avg_tokens_per_session: 5_300_000, spend: 359.86, + savings_estimated_turns: overrides.turns ?? 3073, + savings_estimated_actual_spend: overrides.spend ?? 359.86, classifier_cost: 6.146, saved_spend: 2174.59, baseline_spend: 2534.45, @@ -100,6 +102,8 @@ const zeroTotals: Totals = { avg_session_seconds: 0, avg_tokens_per_session: 0, spend: 0, + savings_estimated_turns: 0, + savings_estimated_actual_spend: 0, classifier_cost: 0, saved_spend: 0, baseline_spend: 0, @@ -153,6 +157,39 @@ describe("AutoRouterBenchmarksTab", () => { mockAutoRouters(); }); + it.each([ + { estimatedTurns: 0, saved: null, pct: null }, + { estimatedTurns: 10, saved: -0.5, pct: -33.3 }, + { estimatedTurns: 10, saved: 0, pct: 0 }, + ])("preserves costs for $estimatedTurns estimated turns with savings $saved", ({ estimatedTurns, saved, pct }) => { + const cohort = { + savings_estimated_turns: estimatedTurns, + savings_estimated_actual_spend: estimatedTurns ? 2 : 0, + saved_spend: saved, + baseline_spend: estimatedTurns ? 2 + (saved ?? 0) : null, + saved_pct: pct, + saved_per_session: null, + }; + const partial = totals(cohort); + mockHook({ data: response([], partial) }); + renderTab(); + expect(screen.getByText("Estimated savings on covered turns")).toBeInTheDocument(); + expect(screen.getByText(`${estimatedTurns} of 3,073 turns estimated`)).toBeInTheDocument(); + expect(screen.getByText("$359.86")).toBeInTheDocument(); + expect(screen.getByText("Actual spend on covered turns")).toBeInTheDocument(); + expect(screen.getByText("Estimated baseline spend on covered turns")).toBeInTheDocument(); + expect(screen.getAllByText("Unavailable")).toHaveLength(estimatedTurns ? 1 : 3); + if (saved === 0) { + expect(screen.getByText("0%")).toBeInTheDocument(); + expect(screen.getAllByText("$2.00")).toHaveLength(2); + } else if (estimatedTurns) { + expect(screen.getByText("-$0.5000")).toBeInTheDocument(); + expect(screen.getByText("+33%")).toBeInTheDocument(); + } else { + expect(screen.queryByText("+0%")).not.toBeInTheDocument(); + } + }); + it("leads with total estimated savings, before the four session-shape metrics", () => { mockHook({ data: response([group(), group({ router_name: "gpt-auto" })]) }); renderTab(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx index ce5ab1c6776..ce55b633b60 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx @@ -73,26 +73,37 @@ const SpendRow: React.FC<{ label: string; value: string; hint?: string; subdued? const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => { const stats = view.stats; - const cheaper = stats.saved_spend >= 0; + const cheaper = stats.saved_spend != null && stats.saved_spend >= 0; + const completeCoverage = stats.savings_estimated_turns === stats.turns; return (

- Total estimated savings + {completeCoverage ? "Total estimated savings" : "Estimated savings on covered turns"}

- {usd(stats.saved_spend)} + {stats.saved_spend == null ? "Unavailable" : usd(stats.saved_spend)}

- - {stats.saved_spend !== 0 && (cheaper ? "-" : "+")} - {Math.abs(stats.saved_pct).toFixed(0)}% - + {stats.saved_pct != null && ( + + {stats.saved_spend !== 0 && (cheaper ? "-" : "+")} + {Math.abs(stats.saved_pct).toFixed(0)}% + + )}
+

+ {stats.savings_estimated_turns.toLocaleString()} of {stats.turns.toLocaleString()} turns estimated +

+ {!completeCoverage && ( +

+ Turns without a current estimate are excluded, including older estimates. +

+ )}
@@ -120,7 +131,15 @@ const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => {

)} - + {!completeCoverage && ( + + )} +
@@ -279,7 +298,7 @@ const BenchmarksBody: React.FC = ({ isPending, error, data,
@@ -288,12 +307,13 @@ const BenchmarksBody: React.FC = ({ isPending, error, data,

- Compares your actual routed spend with the estimated cost of using only the most expensive model configured in - the auto-router. It accounts for both the cache savings from staying on one model and the added cache costs from - switching models. Savings are net of recorded LLM classification cost, which is included in actual spend. - Classification cost per 1K turns is averaged over all auto-router turns, including those that skip - classification. The range counts whole sessions that overlap it, so totals can differ slightly from the Overall - tab, which buckets savings by UTC day. + Compares covered turns with the estimated cost of using the router's highest-tier baseline model. Estimates + use registered requests since tracking began, matching cache prefixes and expiry, and the actual response + length. Total actual spend includes every turn; savings and baseline spend include only turns with a current + estimate, including turns with zero savings. Savings are net of recorded LLM classification cost. Classification + cost per 1K turns is averaged over all auto-router turns, including those that skip classification. The range + counts whole sessions that overlap it, so totals can differ slightly from the Overall tab, which buckets savings + by UTC day.

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.test.tsx index da4af8baf29..e4417d77463 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.test.tsx @@ -21,6 +21,8 @@ const totalsOnly = { avg_session_seconds: 60, avg_tokens_per_session: 100, spend: 1, + savings_estimated_turns: 9, + savings_estimated_actual_spend: 1, saved_spend: 1, baseline_spend: 2, saved_pct: 50, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.test.ts index 22d6336e86f..0586163e77e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.test.ts @@ -37,6 +37,8 @@ const totals = (overrides: Partial = {}) => ({ avg_session_seconds: 7560, avg_tokens_per_session: 5_300_000, spend: 359.86, + savings_estimated_turns: overrides.turns ?? 3073, + savings_estimated_actual_spend: overrides.spend ?? 359.86, classifier_cost: 6.146, saved_spend: 2174.59, baseline_spend: 2534.45, diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/WebSearchInterceptionSettings/WebSearchInterceptionSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/WebSearchInterceptionSettings/WebSearchInterceptionSettings.test.tsx index ae981aa9767..a3d09c8f6d6 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/WebSearchInterceptionSettings/WebSearchInterceptionSettings.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/WebSearchInterceptionSettings/WebSearchInterceptionSettings.test.tsx @@ -127,6 +127,32 @@ describe("WebSearchInterceptionSettings", () => { expect(mockMutate.mock.calls[0][0]).toEqual(ENABLED_PAYLOAD); }); + it("warns when the cluster has it on but the serving pod has not applied it", async () => { + vi.mocked(useWebSearchInterceptionSettings).mockReturnValue({ + data: { ...storedSettings, values: { ...storedSettings.values, enabled: true }, active_on_this_pod: false }, + isLoading: false, + isError: false, + error: null, + } as any); + + await renderSettings(); + + expect(screen.getByText(/has not applied it/i)).toBeInTheDocument(); + }); + + it("stays quiet when the serving pod has applied the cluster setting", async () => { + vi.mocked(useWebSearchInterceptionSettings).mockReturnValue({ + data: { ...storedSettings, values: { ...storedSettings.values, enabled: true }, active_on_this_pod: true }, + isLoading: false, + isError: false, + error: null, + } as any); + + await renderSettings(); + + expect(screen.queryByText(/has not applied it/i)).not.toBeInTheDocument(); + }); + it("ignores stored values whose types do not match the field", async () => { vi.mocked(useWebSearchInterceptionSettings).mockReturnValue({ data: { diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/WebSearchInterceptionSettings/WebSearchInterceptionSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/WebSearchInterceptionSettings/WebSearchInterceptionSettings.tsx index 3e9e04e720b..93b9cbffd32 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/WebSearchInterceptionSettings/WebSearchInterceptionSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/WebSearchInterceptionSettings/WebSearchInterceptionSettings.tsx @@ -5,7 +5,7 @@ import { useUpdateWebSearchInterceptionSettings } from "@/app/(dashboard)/hooks/ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { toast } from "@/lib/toast"; import { Skeleton } from "@/components/ui/skeleton"; -import { CircleHelp, Info, Save } from "lucide-react"; +import { CircleHelp, Info, Save, TriangleAlert } from "lucide-react"; import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { useEffect, useState } from "react"; import { useForm } from "react-hook-form"; @@ -293,9 +293,22 @@ export default function WebSearchInterceptionSettings() { } const values: WebSearchInterceptionStoredValues = toStoredValues(data?.values ?? NO_STORED_VALUES); + const notAppliedHere = values.enabled === true && data?.active_on_this_pod === false; return (
+ {notAppliedHere && ( + + + Not running on the proxy that answered this page + + Interception is switched on for the cluster, but the proxy serving this page has not applied it. That is + expected for about 10 seconds after a change or a restart. If it persists, check that proxy's logs: + requests it handles are not being intercepted. + + + )} + Web Search Interception diff --git a/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx index 4a574ac736d..a03ccb11456 100644 --- a/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx @@ -1,7 +1,7 @@ import React, { useState } from "react"; -import { describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import userEvent from "@testing-library/user-event"; -import { fireEvent, renderWithProviders, screen } from "../../../tests/test-utils"; +import { act, fireEvent, renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; import ClassificationMethodConfig from "./ClassificationMethodConfig"; import AutoRouterClassifierTabs from "./AutoRouterClassifierTabs"; import ForecastClassifierConfig from "./ForecastClassifierConfig"; @@ -37,6 +37,53 @@ const fuseInitial: ComplexityRouterConfigValue = { }, }; const options = ["judge", "efficient", "capable"].map((model) => ({ value: model, label: model })); +const catalog = { + version: "catalog-v1", + models: [ + { + id: "efficient-v1", + label: "Efficient preset", + text: "Maintained efficient profile", + sources: ["https://example.com/efficient"], + model: "efficient-model", + }, + { + id: "capable-v1", + label: "Capable preset", + text: "Maintained capable profile", + sources: ["https://example.com/capable"], + model: "capable-model", + }, + ], + harnesses: [ + { + id: "runtime-v1", + label: "Runtime preset", + text: "Maintained runtime profile", + sources: ["https://example.com/runtime"], + }, + ], +}; +const presetConfig = { + efficient_profile_preset: catalog.models[0].id, + capable_profile_preset: catalog.models[1].id, + harness_preset: catalog.harnesses[0].id, + max_quality_gap: 0.05, +}; +const presetInitial = { ...fuseInitial, llm_v2_config: presetConfig }; + +beforeEach(() => { + testQueryClient.clear(); + vi.stubGlobal( + "fetch", + vi.fn().mockImplementation(async () => Response.json(catalog)), + ); +}); + +afterEach(() => { + testQueryClient.clear(); + vi.unstubAllGlobals(); +}); function Form({ initialValue = initial }: { initialValue?: ComplexityRouterConfigValue }) { const [value, setValue] = useState(initialValue); @@ -72,6 +119,193 @@ function Form({ initialValue = initial }: { initialValue?: ComplexityRouterConfi } describe("forecast classifier form", () => { + it("selects all three maintained presets, previews provenance, and saves only references", async () => { + const user = userEvent.setup(); + renderWithProviders(
); + await user.click(screen.getByRole("combobox", { name: "Efficient solver profile preset" })); + await user.click(await screen.findByRole("option", { name: /^Efficient preset/ })); + await user.click(screen.getByRole("combobox", { name: "Capable solver profile preset" })); + await user.click(screen.getByRole("option", { name: /^Capable preset/ })); + await user.click(screen.getByRole("combobox", { name: "Harness and budget preset" })); + await user.click(screen.getByRole("option", { name: /^Runtime preset/ })); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(catalog.models[0].text); + expect(screen.getByLabelText("Efficient solver profile")).toHaveAttribute("readonly"); + expect(screen.getByLabelText("Capable solver profile")).toHaveValue(catalog.models[1].text); + expect(screen.getByLabelText("Harness and budget")).toHaveValue(catalog.harnesses[0].text); + expect(screen.getAllByText(`Catalog version: ${catalog.version}`)).toHaveLength(3); + expect(screen.getByText(`Model: ${catalog.models[0].model}`)).toBeInTheDocument(); + expect(screen.getAllByRole("link", { name: "Source 1" }).map((link) => link.getAttribute("href"))).toEqual([ + catalog.models[0].sources[0], + catalog.models[1].sources[0], + catalog.harnesses[0].sources[0], + ]); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect(JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config).toEqual( + presetConfig, + ); + expect(fetch).toHaveBeenCalledTimes(1); + expect(fetch).toHaveBeenCalledWith( + expect.objectContaining({ url: expect.stringMatching(/\/public\/complexity_router\/fuse_presets$/) }), + ); + }); + + it.each([undefined, null, "Explicit override"])( + "copies effective text to Custom and clears only that reference, override=%s", + async (override) => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + const effectiveText = override ?? catalog.models[0].text; + await waitFor(() => expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(effectiveText)); + await user.click(screen.getByRole("combobox", { name: "Efficient solver profile preset" })); + await user.click(screen.getByRole("option", { name: "Custom" })); + expect(screen.getByLabelText("Efficient solver profile")).not.toHaveAttribute("readonly"); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(effectiveText); + fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Custom budget" } }); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + const { efficient_profile_preset: _preset, ...rest } = presetConfig; + expect( + JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config, + ).toEqual({ + ...rest, + efficient_profile: "Custom budget", + }); + }, + ); + + it.each([ + { ...fuseInitial.llm_v2_config!, efficient_profile: catalog.models[0].text }, + { + ...presetConfig, + efficient_profile: "Explicit override", + capable_profile: "Capable override", + harness: "Harness override", + }, + ])("keeps existing custom ownership and references on an unchanged save: %j", async (settings) => { + renderWithProviders(); + await waitFor(() => expect(screen.queryByText(/Loading profile presets/)).not.toBeInTheDocument()); + expect(screen.getByRole("combobox", { name: "Efficient solver profile preset" })).toHaveValue("Custom"); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(settings.efficient_profile); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect(JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config).toEqual( + settings, + ); + }); + + it.each([true, false])( + "keeps edits and stored IDs while the pending catalog settles, success=%s", + async (success) => { + let resolveCatalog: (response: Response) => void = () => {}; + vi.mocked(fetch).mockReturnValue( + new Promise((resolve) => { + resolveCatalog = resolve; + }), + ); + const settings = { ...presetConfig, efficient_profile: "Original override" }; + renderWithProviders(); + expect(screen.getByText(/Loading profile presets/)).toBeInTheDocument(); + fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Typed while loading" } }); + await act(async () => + resolveCatalog(success ? Response.json(catalog) : Response.json({ error: "unavailable" }, { status: 503 })), + ); + if (success) await screen.findAllByText(`Catalog version: ${catalog.version}`); + else expect(await screen.findByText(/Profile presets could not be loaded/)).toBeInTheDocument(); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue("Typed while loading"); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect( + JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config, + ).toEqual({ ...settings, efficient_profile: "Typed while loading" }); + }, + ); + + it.each([ + ["efficient_profile", "Efficient solver profile"], + ["capable_profile", "Capable solver profile"], + ["harness", "Harness and budget"], + ] as const)( + "preserves the saved %s reference during a catalog outage until Custom text replaces it", + async (field, label) => { + const user = userEvent.setup(); + vi.mocked(fetch).mockImplementation(async () => Response.json({ error: "unavailable" }, { status: 503 })); + renderWithProviders(); + expect(await screen.findByText(/Profile presets could not be loaded/)).toBeInTheDocument(); + const save = screen.getByRole("button", { name: "Save configuration" }); + const output = screen.getByRole("status", { name: "Saved configuration" }); + expect(save).toBeEnabled(); + await user.click(save); + expect(JSON.parse(output.textContent!).llm_v2_config).toEqual(presetConfig); + + await user.click(screen.getByRole("combobox", { name: `${label} preset` })); + await user.click(screen.getByRole("option", { name: "Custom" })); + expect(screen.getByLabelText(label)).toHaveValue(""); + expect(screen.getByLabelText(label)).not.toHaveAttribute("readonly"); + expect(screen.getByRole("combobox", { name: `${label} preset` })).toHaveValue("Custom"); + expect(save).toBeEnabled(); + await user.click(save); + expect(JSON.parse(output.textContent!).llm_v2_config).toEqual(presetConfig); + + fireEvent.change(screen.getByLabelText(label), { target: { value: " " } }); + expect(save).toBeDisabled(); + await user.click(screen.getByRole("button", { name: `Keep saved ${label.toLowerCase()} preset` })); + expect(screen.getByLabelText(label)).toHaveAttribute("readonly"); + expect(save).toBeEnabled(); + await user.click(save); + expect(JSON.parse(output.textContent!).llm_v2_config).toEqual(presetConfig); + + await user.click(screen.getByRole("combobox", { name: `${label} preset` })); + await user.click(screen.getByRole("option", { name: "Custom" })); + const replacement = "Manually authored replacement"; + fireEvent.change(screen.getByLabelText(label), { target: { value: replacement } }); + expect(save).toBeEnabled(); + await user.click(save); + const referenceKey = `${field}_preset` as const; + const { [referenceKey]: _reference, ...remaining } = presetConfig; + expect(JSON.parse(output.textContent!).llm_v2_config).toEqual({ ...remaining, [field]: replacement }); + }, + ); + + it.each([true, false])( + "keeps a reference selected as Custom while the catalog settles, success=%s", + async (success) => { + const user = userEvent.setup(); + const response = Promise.withResolvers(); + vi.mocked(fetch).mockReturnValue(response.promise); + renderWithProviders(); + await user.click(screen.getByRole("combobox", { name: "Efficient solver profile preset" })); + await user.click(screen.getByRole("option", { name: "Custom" })); + await act(async () => response.resolve(success ? Response.json(catalog) : Response.json({}, { status: 503 }))); + if (success) await screen.findAllByText(`Catalog version: ${catalog.version}`); + else await screen.findByText(/Profile presets could not be loaded/); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(success ? catalog.models[0].text : ""); + await user.click(screen.getByRole("button", { name: "Save configuration" })); + expect( + JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config, + ).toEqual(presetConfig); + fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Replacement" } }); + await user.click(screen.getByRole("button", { name: "Save configuration" })); + const { efficient_profile_preset: _reference, ...remaining } = presetConfig; + expect( + JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config, + ).toEqual({ + ...remaining, + efficient_profile: "Replacement", + }); + }, + ); + + it("keeps unknown saved IDs visible with unavailable previews rather than replacing them", async () => { + const settings = { ...presetConfig, efficient_profile_preset: "unavailable-v8" }; + renderWithProviders(); + await screen.findAllByText(`Catalog version: ${catalog.version}`); + expect(screen.getByRole("combobox", { name: "Efficient solver profile preset" })).toHaveValue("unavailable-v8"); + expect(screen.getByText("Preset preview unavailable. The saved reference is preserved")).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect(JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config).toEqual( + settings, + ); + }); + it("switches a populated standard router to Capability without saving hidden pools or their overrides", () => { renderWithProviders( { fireEvent.click(screen.getByRole("tab", { name: "Complexity" })); fireEvent.click(screen.getByRole("radio", { name: new RegExp(`^${target}`) })); await user.click(screen.getByRole("combobox", { name: "Classifier Model" })); - await user.click(screen.getByRole("option", { name: "judge", exact: true })); + await user.click(screen.getByRole("option", { name: "judge" })); fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); const output = screen.getByRole("status", { name: "Saved configuration" }); expect(output).toHaveTextContent('"classification_rubric":"agentic"'); diff --git a/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.tsx index b901a509435..c336423fe39 100644 --- a/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.tsx @@ -3,7 +3,7 @@ import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/component import { ChevronRight } from "lucide-react"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; -import { Textarea } from "@/components/ui/textarea"; +import FuseProfilePresets from "./FuseProfilePresets"; import { Switch } from "@/components/ui/switch"; import { SearchSelect } from "@/components/shared/SearchSelect"; import { MultiSelect } from "@/components/shared/MultiSelect"; @@ -239,35 +239,13 @@ const ForecastClassifierConfig = ({ value, onChange, modelOptions, effortOptions ) : ( <> - {(["efficient_profile", "capable_profile", "harness"] as const).map((field) => { - const label = { - efficient_profile: "Efficient solver profile", - capable_profile: "Capable solver profile", - harness: "Harness and budget", - }[field]; - return ( -
- -