scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge remote-tracking branch 'origin/main' into fix-responses-sse-framing

Browse source 
# Conflicts:
#	tests/unit/proxy/proxy_server/test_streaming_helpers.py
This commit is contained in:
mateo-berri 2026-11-02 10:49:38 -08:00
commit 263e6aa4db
216 changed files with 16631 additions and 4125 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
AGENTS.md
View file

@ -41,6 +41,8 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just drop the section
Run QA proxies with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --num_workers 2 --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`, and swap its `--config` for a copy of `litellm/proxy/dev_config.yaml` that points both `router_settings` (`redis_host`, `redis_port`) and `general_settings.coordination_redis` (`host`, `port`) at one dedicated Redis instance, with credentials when required, so the workers share router state and spend, budget, and rate-limit counters. Check that two worker PIDs are serving, and use the same worker count and Redis setup for the before and after proof. When the PR is about `--run_gunicorn`, `--run_hypercorn`, or `--run_granian`, which `--reload` does not support, drop `--reload`. When it is about behavior without Redis, drop Redis. When the proxy prints `"workers" flag is ignored when reloading is enabled`, that checkout or its uvicorn reloads only one worker, so drop the two-worker check for it. Say in the proof which part you dropped and why
Simple PRs may skip before/after e2e when the full diff is mechanical and focused checks completely prove it, such as #49030 restoring an existing helper name. Trace affected callers and run relevant tests/static checks at the tip. For a bug fix, establish that the regression catches the original defect. Size alone never qualifies. New behavior, API/provider semantics, auth, billing, persistence, concurrency or uncertain integration risk need live e2e. In Screenshots / Proof of Fix, record `E2E exception at <tip SHA>`, scope, rationale, commands and results. Tests count as proof only for this exception. Reassess changed diffs, retain CI/bot gates, and honor explicit e2e requests (2026-10-28, unnecessary live QA for #49030)
For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce

2
README.md
View file

@ -696,7 +696,7 @@ For companies that need better security, user management and professional suppor
[Talk to founders](https://enterprise.litellm.ai/demo)
This covers:
- ✅ **Features under the [LiteLLM Commercial License](https://docs.litellm.ai/docs/proxy/enterprise):**
- ✅ **Features under the [LiteLLM Commercial License](https://docs.litellm.ai/docs/enterprise):**
- ✅ **Feature Prioritization**
- ✅ **Custom Integrations**
- ✅ **Professional Support - Dedicated discord + slack**

293
litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs
View file

@ -124,12 +124,86 @@ pub enum AnthropicBeta {
ToolSearchTool20251019,
#[strum(serialize = "inline-tools-2026-09-15")]
InlineTools20260915,
#[strum(serialize = "afk-mode-2026-01-31")]
AfkMode20260131,
#[strum(serialize = "agent-memory-2026-07-22")]
AgentMemory20260722,
#[strum(serialize = "analytics-usage-cost-2026-08-24")]
AnalyticsUsageCost20260824,
#[strum(serialize = "cache-diagnosis-2026-04-07")]
CacheDiagnosis20260407,
#[strum(serialize = "ce-plugins-2026-09-01")]
CePlugins20260901,
#[strum(serialize = "ce-user-management-2026-07-13")]
CeUserManagement20260713,
#[strum(serialize = "code-execution-2025-05-22")]
CodeExecution20250522,
#[strum(serialize = "computer-use-2024-10-22")]
ComputerUse20241022,
#[strum(serialize = "dev-full-thinking-2025-05-14")]
DevFullThinking20250514,
#[strum(serialize = "dreaming-2026-04-21")]
Dreaming20260421,
#[strum(serialize = "experimental-claude-tag-2026-10-21")]
ExperimentalClaudeTag20261021,
#[strum(serialize = "extended-cache-ttl-2025-04-11")]
ExtendedCacheTtl20250411,
#[strum(serialize = "fallback-credit-2026-06-01")]
FallbackCredit20260601,
#[strum(serialize = "fallback-credit-2026-07-01")]
FallbackCredit20260701,
#[strum(serialize = "managed-agents-2026-04-01")]
ManagedAgents20260401,
#[strum(serialize = "mcp-client-2026-09-15")]
McpClient20260915,
#[strum(serialize = "mcp-tunnels-2026-06-22")]
McpTunnels20260622,
#[strum(serialize = "message-batches-2024-09-24")]
MessageBatches20240924,
#[strum(serialize = "message-threads-2026-08-12")]
MessageThreads20260812,
#[strum(serialize = "mid-conversation-system-clear-at-2026-08-21")]
MidConversationSystemClearAt20260821,
#[strum(serialize = "model-context-window-exceeded-2025-08-26")]
ModelContextWindowExceeded20250826,
#[strum(serialize = "output-300k-2026-03-24")]
Output300k20260324,
#[strum(serialize = "pdfs-2024-09-25")]
Pdfs20240925,
#[strum(serialize = "prompt-caching-2024-07-31")]
PromptCaching20240731,
#[strum(serialize = "redact-thinking-2026-02-12")]
RedactThinking20260212,
#[strum(serialize = "server-side-fallback-2026-06-01")]
ServerSideFallback20260601,
#[strum(serialize = "server-side-fallback-2026-07-01")]
ServerSideFallback20260701,
#[strum(serialize = "spend-limit-reads-2026-09-26")]
SpendLimitReads20260926,
#[strum(serialize = "structured-outputs-2025-12-15")]
StructuredOutputs20251215,
#[strum(serialize = "task-budgets-2026-03-13")]
TaskBudgets20260313,
#[strum(serialize = "telemetry-destinations-2026-08-11")]
TelemetryDestinations20260811,
#[strum(serialize = "thinking-token-count-2026-05-13")]
ThinkingTokenCount20260513,
#[strum(serialize = "token-counting-2024-11-01")]
TokenCounting20241101,
#[strum(serialize = "user-profiles-2026-03-24")]
UserProfiles20260324,
#[strum(serialize = "user-profiles-2026-08-18")]
UserProfiles20260818,
#[strum(serialize = "user-profiles-2026-09-04")]
UserProfiles20260904,
#[strum(serialize = "webhooks-2026-10-07")]
Webhooks20261007,
#[strum(default, transparent)]
Other(String),
}
impl AnthropicBeta {
pub const KNOWN: [Self; 40] = [
pub const KNOWN: [Self; 77] = [
Self::Oauth20250420,
Self::WebFetch20250910,
Self::WebSearch20250305,
@ -170,6 +244,43 @@ impl AnthropicBeta {
Self::ToolExamples20251029,
Self::ToolSearchTool20251019,
Self::InlineTools20260915,
Self::AfkMode20260131,
Self::AgentMemory20260722,
Self::AnalyticsUsageCost20260824,
Self::CacheDiagnosis20260407,
Self::CePlugins20260901,
Self::CeUserManagement20260713,
Self::CodeExecution20250522,
Self::ComputerUse20241022,
Self::DevFullThinking20250514,
Self::Dreaming20260421,
Self::ExperimentalClaudeTag20261021,
Self::ExtendedCacheTtl20250411,
Self::FallbackCredit20260601,
Self::FallbackCredit20260701,
Self::ManagedAgents20260401,
Self::McpClient20260915,
Self::McpTunnels20260622,
Self::MessageBatches20240924,
Self::MessageThreads20260812,
Self::MidConversationSystemClearAt20260821,
Self::ModelContextWindowExceeded20250826,
Self::Output300k20260324,
Self::Pdfs20240925,
Self::PromptCaching20240731,
Self::RedactThinking20260212,
Self::ServerSideFallback20260601,
Self::ServerSideFallback20260701,
Self::SpendLimitReads20260926,
Self::StructuredOutputs20251215,
Self::TaskBudgets20260313,
Self::TelemetryDestinations20260811,
Self::ThinkingTokenCount20260513,
Self::TokenCounting20241101,
Self::UserProfiles20260324,
Self::UserProfiles20260818,
Self::UserProfiles20260904,
Self::Webhooks20261007,
];
pub fn as_str(&self) -> &str {
@ -195,10 +306,9 @@ impl AnthropicBeta {
fn rejected_by(&self) -> &'static [BetaProvider] {
match self {
Self::AdvancedToolUse20251120 => {
Self::AdvancedToolUse20251120 | Self::ToolSearchTool20251019 => {
&[BetaProvider::BedrockConverse, BetaProvider::Databricks]
}
Self::ContextManagement20250627 => &[BetaProvider::BedrockConverse],
Self::AdvisorTool20260301 => &[
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
@ -207,98 +317,99 @@ impl AnthropicBeta {
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::Compact20260904 => &[
Self::AfkMode20260131
| Self::CePlugins20260901
| Self::CeUserManagement20260713
| Self::Compact20260904
| Self::Dreaming20260421
| Self::ExperimentalClaudeTag20261021
| Self::ExtendedCacheTtl20250411
| Self::FallbackCredit20260601
| Self::FallbackCredit20260701
| Self::InlineTools20260915
| Self::ManagedAgents20260401
| Self::McpTunnels20260622
| Self::MidConversationSystemClearAt20260821
| Self::ModelContextWindowExceeded20250826
| Self::Oauth20250420
| Self::Pdfs20240925
| Self::PromptCaching20240731
| Self::PromptCachingScope20260105
| Self::RedactThinking20260212
| Self::ServerSideFallback20260601
| Self::SpendLimitReads20260926
| Self::StructuredOutputs20251215
| Self::TaskBudgets20260313
| Self::TelemetryDestinations20260811
| Self::ThinkingTokenCount20260513
| Self::TokenCounting20241101
| Self::UserProfiles20260324
| Self::UserProfiles20260818
| Self::UserProfiles20260904 => &[
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::BedrockMantle,
BetaProvider::Databricks,
],
Self::Bash20241022
Self::AgentMemory20260722
| Self::Bash20241022
| Self::Bash20250124
| Self::DevFullThinking20250514
| Self::McpServers20251204
| Self::StructuredOutput20240301
| Self::TextEditor20241022
| Self::TextEditor20250124 => BetaProvider::VARIANTS,
Self::ClaudeCode20250219 | Self::ToolExamples20251029 => &[
BetaProvider::Anthropic,
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::CodeExecution20250825
Self::AnalyticsUsageCost20260824
| Self::CacheDiagnosis20260407
| Self::CodeExecution20250522
| Self::CodeExecution20250825
| Self::ComputerUse20241022
| Self::FastMode20260201
| Self::FilesApi20250414
| Self::McpClient20250404
| Self::McpClient20251120
| Self::PromptCachingScope20260105
| Self::McpClient20260915
| Self::MessageBatches20240924
| Self::MessageThreads20260812
| Self::Output300k20260324
| Self::ServerSideFallback20260701
| Self::Skills20251002
| Self::WebFetch20250910 => &[
| Self::WebFetch20250910
| Self::Webhooks20261007 => &[
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::BedrockMantle,
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse],
Self::ClaudeCode20250219
| Self::PerTurnControl20260701
| Self::ToolExamples20251029 => &[
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::Databricks,
],
Self::Compact20260112
| Self::ContextManagement20250627
| Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse],
Self::ComputerUse20250124
| Self::ComputerUse20251124
| Self::Context1m20250807
| Self::Effort20251124
| Self::FineGrainedToolStreaming20250514
| Self::ThinkingBindingControls20260801
| Self::ThinkingDisplayUpdates20260818 => &[],
Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse],
Self::Effort20251124 => &[BetaProvider::VertexAi],
Self::FastMode20260201 | Self::Oauth20250420 => &[
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::BedrockMantle,
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::FineGrainedToolStreaming20250514 => {
&[BetaProvider::AzureAi, BetaProvider::VertexAi]
}
Self::InterleavedThinking20250514 | Self::WebSearch20250305 => {
&[BetaProvider::BedrockConverse, BetaProvider::Bedrock]
}
Self::MidConversationOutputConfig20260701 => &[
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::BedrockMantle,
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::MidConversationToolChanges20260701 => &[
Self::InterleavedThinking20250514
| Self::Output128k20250219
| Self::TokenEfficientTools20250219
| Self::WebSearch20250305 => &[BetaProvider::BedrockConverse, BetaProvider::Bedrock],
Self::MidConversationOutputConfig20260701
| Self::MidConversationToolChanges20260701 => &[
BetaProvider::BedrockConverse,
BetaProvider::BedrockMantle,
BetaProvider::Databricks,
],
Self::Output128k20250219 | Self::TokenEfficientTools20250219 => &[
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::VertexAi,
],
Self::PerTurnControl20260701 => &[
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::Databricks,
],
Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock, BetaProvider::VertexAi],
Self::ToolSearchTool20251019 => &[
BetaProvider::Anthropic,
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::Databricks,
],
Self::InlineTools20260915 => &[
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::BedrockMantle,
BetaProvider::Databricks,
],
Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock],
Self::Other(_) => BetaProvider::VARIANTS,
}
}
@ -451,7 +562,44 @@ mod tests {
AnthropicBeta::TokenEfficientTools20250219,
AnthropicBeta::ToolExamples20251029,
AnthropicBeta::ToolSearchTool20251019,
AnthropicBeta::InlineTools20260915
AnthropicBeta::InlineTools20260915,
AnthropicBeta::AfkMode20260131,
AnthropicBeta::AgentMemory20260722,
AnthropicBeta::AnalyticsUsageCost20260824,
AnthropicBeta::CacheDiagnosis20260407,
AnthropicBeta::CePlugins20260901,
AnthropicBeta::CeUserManagement20260713,
AnthropicBeta::CodeExecution20250522,
AnthropicBeta::ComputerUse20241022,
AnthropicBeta::DevFullThinking20250514,
AnthropicBeta::Dreaming20260421,
AnthropicBeta::ExperimentalClaudeTag20261021,
AnthropicBeta::ExtendedCacheTtl20250411,
AnthropicBeta::FallbackCredit20260601,
AnthropicBeta::FallbackCredit20260701,
AnthropicBeta::ManagedAgents20260401,
AnthropicBeta::McpClient20260915,
AnthropicBeta::McpTunnels20260622,
AnthropicBeta::MessageBatches20240924,
AnthropicBeta::MessageThreads20260812,
AnthropicBeta::MidConversationSystemClearAt20260821,
AnthropicBeta::ModelContextWindowExceeded20250826,
AnthropicBeta::Output300k20260324,
AnthropicBeta::Pdfs20240925,
AnthropicBeta::PromptCaching20240731,
AnthropicBeta::RedactThinking20260212,
AnthropicBeta::ServerSideFallback20260601,
AnthropicBeta::ServerSideFallback20260701,
AnthropicBeta::SpendLimitReads20260926,
AnthropicBeta::StructuredOutputs20251215,
AnthropicBeta::TaskBudgets20260313,
AnthropicBeta::TelemetryDestinations20260811,
AnthropicBeta::ThinkingTokenCount20260513,
AnthropicBeta::TokenCounting20241101,
AnthropicBeta::UserProfiles20260324,
AnthropicBeta::UserProfiles20260818,
AnthropicBeta::UserProfiles20260904,
AnthropicBeta::Webhooks20261007
)]
beta: AnthropicBeta,
) {
@ -551,7 +699,7 @@ mod tests {
BetaProvider::Anthropic,
Some(AnthropicBeta::Oauth20250420)
)]
#[case::rejected("effort-2025-11-24", BetaProvider::VertexAi, None)]
#[case::rejected("skills-2025-10-02", BetaProvider::VertexAi, None)]
#[case::unknown("example-beta-2099-01-01", BetaProvider::Anthropic, None)]
fn on_resolves_known_spellings_held_as_other(
#[case] raw: &str,
@ -580,12 +728,11 @@ mod tests {
#[test]
fn unknown_values_are_kept_verbatim() {
let parsed: AnthropicBeta = "claude-code-20250219".parse().unwrap();
assert_eq!(
parsed,
AnthropicBeta::Other("claude-code-20250219".to_string())
let parsed: AnthropicBeta = "example-beta-2099-01-01".parse().unwrap();
assert!(
matches!(&parsed, AnthropicBeta::Other(value) if value == "example-beta-2099-01-01")
);
assert_eq!(parsed.to_string(), "claude-code-20250219");
assert_eq!(parsed.to_string(), "example-beta-2099-01-01");
}
#[test]

5
litellm-rust/crates/llms/src/anthropic/beta_headers.rs
View file

@ -96,7 +96,10 @@ mod tests {
}
#[rstest]
#[case::azure_ai(BetaProvider::AzureAi, "fast-mode-2026-02-01,example-beta-2099-01-01")]
#[case::azure_ai(
BetaProvider::AzureAi,
"advisor-tool-2026-03-01,example-beta-2099-01-01"
)]
#[case::anthropic(BetaProvider::Anthropic, "bash_20241022")]
#[case::vertex_ai(BetaProvider::VertexAi, " , ")]
fn filter_removes_header_when_nothing_survives(

19
litellm-rust/crates/model-catalog/src/wire/attributes.rs
View file

@ -11,7 +11,7 @@ use crate::error::{DateCause, Error, Repr};
use crate::vocabulary::{
AudioFormat, EffortLevel, Endpoint, InputModality, Mode, OutputModality, VertexAiAudioApi,
};
use crate::wire::lax::Count;
use crate::wire::lax::{Count, NumericCount};
macro_rules! attributes {
($($(#[$doc:meta])* $name:ident: $ty:ty $(as $via:ty)?),* $(,)?) => {
@ -104,9 +104,12 @@ attributes! {
/// Cloud regions the model is available in ('global' or region ids).
supported_regions: Vec<String>,
thinking_always_on: bool,
token_count_framing: u64 as Count,
token_count_framing: u64 as NumericCount,
token_count_model: TokenCountModel,
token_count_responses_framing: u64 as NumericCount,
token_count_responses_tool_framing: u64 as NumericCount,
token_count_tokenizer: TokenCountTokenizer,
token_count_tool_framing: u64 as NumericCount,
/// Provider default tokens-per-minute limit.
tpm: u64 as Count,
use_openai_responses_path: bool,
@ -302,6 +305,18 @@ mod tests {
#[case::reasoning_modes("reasoning_modes", json!(["none", "adaptive"]), true)]
#[case::reasoning_context_not_a_list("reasoning_contexts", json!("reasoning"), false)]
#[case::reasoning_mode_not_a_string("reasoning_modes", json!([1]), false)]
#[case::framing("token_count_framing", json!(12), true)]
#[case::framing_written_as_text("token_count_framing", json!("12"), false)]
#[case::framing_written_as_whole_float("token_count_framing", json!(12.0), true)]
#[case::fractional_tool_framing("token_count_tool_framing", json!(1.5), false)]
#[case::tool_framing("token_count_tool_framing", json!(214), true)]
#[case::tool_framing_written_as_text("token_count_tool_framing", json!("214"), false)]
#[case::tool_framing_as_a_flag("token_count_tool_framing", json!(true), false)]
#[case::negative_tool_framing("token_count_tool_framing", json!(-1), false)]
#[case::responses_framing("token_count_responses_framing", json!(38), true)]
#[case::responses_framing_as_a_flag("token_count_responses_framing", json!(false), false)]
#[case::responses_tool_framing("token_count_responses_tool_framing", json!(267), true)]
#[case::responses_tool_framing_written_as_text("token_count_responses_tool_framing", json!("267"), false)]
fn each_attribute_reads_its_own_type(
#[case] key: &str,
#[case] value: Value,

55
litellm-rust/crates/model-catalog/src/wire/lax.rs
View file

@ -8,6 +8,8 @@ pub(crate) struct LaxI64;
pub(crate) struct Count;
pub(crate) struct NumericCount;
impl<'de> DeserializeAs<'de, i64> for LaxI64 {
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<i64, D::Error> {
deserializer.deserialize_any(Self)
@ -52,6 +54,34 @@ impl<'de> DeserializeAs<'de, u64> for Count {
}
}
impl<'de> DeserializeAs<'de, u64> for NumericCount {
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<u64, D::Error> {
deserializer.deserialize_any(Self)
}
}
impl Visitor<'_> for NumericCount {
type Value = u64;
fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("a whole number of zero or more")
}
fn visit_u64<E: de::Error>(self, value: u64) -> Result<u64, E> {
Ok(value)
}
fn visit_i64<E: de::Error>(self, value: i64) -> Result<u64, E> {
u64::try_from(value).map_err(|_| E::invalid_value(de::Unexpected::Signed(value), &self))
}
fn visit_f64<E: de::Error>(self, value: f64) -> Result<u64, E> {
integral_float(value)
.and_then(|whole| u64::try_from(whole).ok())
.ok_or_else(|| E::invalid_value(de::Unexpected::Float(value), &self))
}
}
fn integer_string(value: &str) -> Option<i64> {
let integer = match value.split_once('.') {
Some((integer, fraction)) => {
@ -99,6 +129,12 @@ mod tests {
.map(DeserializeAsWrap::into_inner)
}
fn numeric_count(value: Value) -> Option<u64> {
DeserializeAsWrap::<u64, NumericCount>::deserialize(value)
.ok()
.map(DeserializeAsWrap::into_inner)
}
use serde::Deserialize;
#[rstest]
@ -125,4 +161,23 @@ mod tests {
) {
assert_eq!(count(written), expected);
}
#[rstest]
#[case::integer(json!(214), Some(214))]
#[case::zero(json!(0), Some(0))]
#[case::whole_float(json!(214.0), Some(214))]
#[case::negative_zero(json!(-0.0), Some(0))]
#[case::above_i64(json!(u64::MAX), Some(u64::MAX))]
#[case::negative(json!(-1), None)]
#[case::negative_whole_float(json!(-1.0), None)]
#[case::fraction(json!(1.5), None)]
#[case::string(json!("214"), None)]
#[case::boolean(json!(true), None)]
#[case::list(json!([214]), None)]
fn numeric_counts_accept_only_whole_json_numbers(
#[case] written: Value,
#[case] expected: Option<u64>,
) {
assert_eq!(numeric_count(written), expected);
}
}

5
litellm/__init__.py
View file

@ -1647,6 +1647,7 @@ def set_global_gitlab_config(config: Mapping[str, object]) -> None:
# Lazy loading system for heavy modules to reduce initial import time and memory usage
if TYPE_CHECKING:
from . import guardrails as guardrails
from litellm.types.utils import ModelInfo as _ModelInfoType
from litellm.types.utils import PriorityReservationSettings
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
@ -2499,10 +2500,10 @@ def __getattr__(name: str) -> Any:
handler_func: Final = registry[name]
return handler_func(name)
if name == "proxy":
if name in ("proxy", "guardrails"):
import importlib
return importlib.import_module("litellm.proxy")
return importlib.import_module(f"litellm.{name}")
if name == "harness" or name in _AGENT_EXPORTS:
import importlib

175
litellm/anthropic_beta_headers_config.json
View file

@ -37,7 +37,47 @@
"mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01",
"thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18",
"mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01",
"inline-tools-2026-09-15": "inline-tools-2026-09-15"
"inline-tools-2026-09-15": "inline-tools-2026-09-15",
"agent-memory-2026-07-22": null,
"analytics-usage-cost-2026-08-24": "analytics-usage-cost-2026-08-24",
"cache-diagnosis-2026-04-07": "cache-diagnosis-2026-04-07",
"ce-plugins-2026-09-01": "ce-plugins-2026-09-01",
"ce-user-management-2026-07-13": "ce-user-management-2026-07-13",
"claude-code-20250219": "claude-code-20250219",
"code-execution-2025-05-22": "code-execution-2025-05-22",
"computer-use-2024-10-22": "computer-use-2024-10-22",
"dev-full-thinking-2025-05-14": null,
"dreaming-2026-04-21": "dreaming-2026-04-21",
"experimental-claude-tag-2026-10-21": "experimental-claude-tag-2026-10-21",
"extended-cache-ttl-2025-04-11": "extended-cache-ttl-2025-04-11",
"fallback-credit-2026-06-01": "fallback-credit-2026-06-01",
"fallback-credit-2026-07-01": "fallback-credit-2026-07-01",
"managed-agents-2026-04-01": "managed-agents-2026-04-01",
"mcp-client-2026-09-15": "mcp-client-2026-09-15",
"mcp-tunnels-2026-06-22": "mcp-tunnels-2026-06-22",
"message-batches-2024-09-24": "message-batches-2024-09-24",
"message-threads-2026-08-12": "message-threads-2026-08-12",
"mid-conversation-system-clear-at-2026-08-21": "mid-conversation-system-clear-at-2026-08-21",
"model-context-window-exceeded-2025-08-26": "model-context-window-exceeded-2025-08-26",
"output-300k-2026-03-24": "output-300k-2026-03-24",
"pdfs-2024-09-25": "pdfs-2024-09-25",
"prompt-caching-2024-07-31": "prompt-caching-2024-07-31",
"server-side-fallback-2026-06-01": "server-side-fallback-2026-06-01",
"server-side-fallback-2026-07-01": "server-side-fallback-2026-07-01",
"spend-limit-reads-2026-09-26": "spend-limit-reads-2026-09-26",
"task-budgets-2026-03-13": "task-budgets-2026-03-13",
"telemetry-destinations-2026-08-11": "telemetry-destinations-2026-08-11",
"thinking-token-count-2026-05-13": "thinking-token-count-2026-05-13",
"token-counting-2024-11-01": "token-counting-2024-11-01",
"tool-examples-2025-10-29": "tool-examples-2025-10-29",
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
"user-profiles-2026-03-24": "user-profiles-2026-03-24",
"user-profiles-2026-08-18": "user-profiles-2026-08-18",
"user-profiles-2026-09-04": "user-profiles-2026-09-04",
"webhooks-2026-10-07": "webhooks-2026-10-07",
"afk-mode-2026-01-31": "afk-mode-2026-01-31",
"redact-thinking-2026-02-12": "redact-thinking-2026-02-12",
"structured-outputs-2025-12-15": "structured-outputs-2025-12-15"
},
"azure_ai": {
"advisor-tool-2026-03-01": null,
@ -45,23 +85,21 @@
"bash_20241022": null,
"bash_20250124": null,
"code-execution-2025-08-25": "code-execution-2025-08-25",
"compact-2026-01-12": null,
"compact-2026-09-04": "compact-2026-09-04",
"compact-2026-01-12": "compact-2026-01-12",
"computer-use-2025-01-24": "computer-use-2025-01-24",
"computer-use-2025-11-24": "computer-use-2025-11-24",
"context-1m-2025-08-07": "context-1m-2025-08-07",
"context-management-2025-06-27": "context-management-2025-06-27",
"dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03",
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": null,
"fast-mode-2026-02-01": "fast-mode-2026-02-01",
"files-api-2025-04-14": "files-api-2025-04-14",
"fine-grained-tool-streaming-2025-05-14": null,
"fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
"interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
"mcp-client-2025-11-20": "mcp-client-2025-11-20",
"mcp-client-2025-04-04": "mcp-client-2025-04-04",
"mcp-servers-2025-12-04": null,
"mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01",
"output-128k-2025-02-19": null,
"output-128k-2025-02-19": "output-128k-2025-02-19",
"structured-output-2024-03-01": null,
"per-turn-control-2026-07-01": "per-turn-control-2026-07-01",
"prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
@ -71,9 +109,54 @@
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18",
"token-efficient-tools-2025-02-19": null,
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
"web-fetch-2025-09-10": "web-fetch-2025-09-10",
"web-search-2025-03-05": "web-search-2025-03-05"
"web-search-2025-03-05": "web-search-2025-03-05",
"agent-memory-2026-07-22": null,
"analytics-usage-cost-2026-08-24": "analytics-usage-cost-2026-08-24",
"cache-diagnosis-2026-04-07": "cache-diagnosis-2026-04-07",
"ce-plugins-2026-09-01": "ce-plugins-2026-09-01",
"ce-user-management-2026-07-13": "ce-user-management-2026-07-13",
"claude-code-20250219": "claude-code-20250219",
"code-execution-2025-05-22": "code-execution-2025-05-22",
"compact-2026-09-04": "compact-2026-09-04",
"computer-use-2024-10-22": "computer-use-2024-10-22",
"dev-full-thinking-2025-05-14": null,
"dreaming-2026-04-21": "dreaming-2026-04-21",
"experimental-claude-tag-2026-10-21": "experimental-claude-tag-2026-10-21",
"extended-cache-ttl-2025-04-11": "extended-cache-ttl-2025-04-11",
"fallback-credit-2026-06-01": "fallback-credit-2026-06-01",
"fallback-credit-2026-07-01": "fallback-credit-2026-07-01",
"inline-tools-2026-09-15": "inline-tools-2026-09-15",
"managed-agents-2026-04-01": "managed-agents-2026-04-01",
"mcp-client-2026-09-15": "mcp-client-2026-09-15",
"mcp-tunnels-2026-06-22": "mcp-tunnels-2026-06-22",
"message-batches-2024-09-24": "message-batches-2024-09-24",
"message-threads-2026-08-12": "message-threads-2026-08-12",
"mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01",
"mid-conversation-system-clear-at-2026-08-21": "mid-conversation-system-clear-at-2026-08-21",
"mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01",
"model-context-window-exceeded-2025-08-26": "model-context-window-exceeded-2025-08-26",
"oauth-2025-04-20": "oauth-2025-04-20",
"output-300k-2026-03-24": "output-300k-2026-03-24",
"pdfs-2024-09-25": "pdfs-2024-09-25",
"prompt-caching-2024-07-31": "prompt-caching-2024-07-31",
"server-side-fallback-2026-06-01": "server-side-fallback-2026-06-01",
"server-side-fallback-2026-07-01": "server-side-fallback-2026-07-01",
"spend-limit-reads-2026-09-26": "spend-limit-reads-2026-09-26",
"task-budgets-2026-03-13": "task-budgets-2026-03-13",
"telemetry-destinations-2026-08-11": "telemetry-destinations-2026-08-11",
"thinking-token-count-2026-05-13": "thinking-token-count-2026-05-13",
"token-counting-2024-11-01": "token-counting-2024-11-01",
"tool-examples-2025-10-29": "tool-examples-2025-10-29",
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
"user-profiles-2026-03-24": "user-profiles-2026-03-24",
"user-profiles-2026-08-18": "user-profiles-2026-08-18",
"user-profiles-2026-09-04": "user-profiles-2026-09-04",
"webhooks-2026-10-07": "webhooks-2026-10-07",
"afk-mode-2026-01-31": "afk-mode-2026-01-31",
"redact-thinking-2026-02-12": "redact-thinking-2026-02-12",
"structured-outputs-2025-12-15": "structured-outputs-2025-12-15"
},
"bedrock_converse": {
"advisor-tool-2026-03-01": null,
@ -196,34 +279,92 @@
"context-1m-2025-08-07": "context-1m-2025-08-07",
"context-management-2025-06-27": "context-management-2025-06-27",
"dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03",
"effort-2025-11-24": null,
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": null,
"fine-grained-tool-streaming-2025-05-14": null,
"fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
"interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
"mcp-client-2025-11-20": null,
"mcp-client-2025-04-04": null,
"mcp-servers-2025-12-04": null,
"mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01",
"output-128k-2025-02-19": null,
"output-128k-2025-02-19": "output-128k-2025-02-19",
"structured-output-2024-03-01": null,
"per-turn-control-2026-07-01": "per-turn-control-2026-07-01",
"prompt-caching-scope-2026-01-05": null,
"skills-2025-10-02": null,
"structured-outputs-2025-11-13": null,
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18",
"token-efficient-tools-2025-02-19": null,
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
"web-fetch-2025-09-10": null,
"web-search-2025-03-05": "web-search-2025-03-05",
"inline-tools-2026-09-15": "inline-tools-2026-09-15"
"inline-tools-2026-09-15": "inline-tools-2026-09-15",
"agent-memory-2026-07-22": null,
"analytics-usage-cost-2026-08-24": null,
"cache-diagnosis-2026-04-07": null,
"ce-plugins-2026-09-01": null,
"ce-user-management-2026-07-13": null,
"claude-code-20250219": "claude-code-20250219",
"code-execution-2025-05-22": null,
"compact-2026-09-04": null,
"computer-use-2024-10-22": null,
"dev-full-thinking-2025-05-14": null,
"dreaming-2026-04-21": null,
"experimental-claude-tag-2026-10-21": null,
"extended-cache-ttl-2025-04-11": null,
"fallback-credit-2026-06-01": "fallback-credit-2026-06-01",
"fallback-credit-2026-07-01": "fallback-credit-2026-07-01",
"managed-agents-2026-04-01": null,
"mcp-client-2026-09-15": null,
"mcp-tunnels-2026-06-22": null,
"message-batches-2024-09-24": null,
"message-threads-2026-08-12": null,
"mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01",
"mid-conversation-system-clear-at-2026-08-21": "mid-conversation-system-clear-at-2026-08-21",
"mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01",
"model-context-window-exceeded-2025-08-26": "model-context-window-exceeded-2025-08-26",
"oauth-2025-04-20": null,
"output-300k-2026-03-24": null,
"pdfs-2024-09-25": "pdfs-2024-09-25",
"prompt-caching-2024-07-31": null,
"server-side-fallback-2026-06-01": "server-side-fallback-2026-06-01",
"server-side-fallback-2026-07-01": null,
"spend-limit-reads-2026-09-26": null,
"task-budgets-2026-03-13": "task-budgets-2026-03-13",
"telemetry-destinations-2026-08-11": null,
"thinking-token-count-2026-05-13": "thinking-token-count-2026-05-13",
"token-counting-2024-11-01": "token-counting-2024-11-01",
"tool-examples-2025-10-29": "tool-examples-2025-10-29",
"user-profiles-2026-03-24": null,
"user-profiles-2026-08-18": null,
"user-profiles-2026-09-04": null,
"webhooks-2026-10-07": null,
"afk-mode-2026-01-31": "afk-mode-2026-01-31",
"redact-thinking-2026-02-12": null,
"structured-outputs-2025-12-15": "structured-outputs-2025-12-15"
},
"vertex_ai_header_only": {
"ce-plugins-2026-09-01": "ce-plugins-2026-09-01",
"ce-user-management-2026-07-13": "ce-user-management-2026-07-13",
"compact-2026-09-04": "compact-2026-09-04",
"inline-tools-2026-09-15": "inline-tools-2026-09-15"
"dreaming-2026-04-21": "dreaming-2026-04-21",
"experimental-claude-tag-2026-10-21": "experimental-claude-tag-2026-10-21",
"extended-cache-ttl-2025-04-11": "extended-cache-ttl-2025-04-11",
"inline-tools-2026-09-15": "inline-tools-2026-09-15",
"managed-agents-2026-04-01": "managed-agents-2026-04-01",
"mcp-tunnels-2026-06-22": "mcp-tunnels-2026-06-22",
"oauth-2025-04-20": "oauth-2025-04-20",
"prompt-caching-2024-07-31": "prompt-caching-2024-07-31",
"prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
"redact-thinking-2026-02-12": "redact-thinking-2026-02-12",
"spend-limit-reads-2026-09-26": "spend-limit-reads-2026-09-26",
"telemetry-destinations-2026-08-11": "telemetry-destinations-2026-08-11",
"user-profiles-2026-03-24": "user-profiles-2026-03-24",
"user-profiles-2026-08-18": "user-profiles-2026-08-18",
"user-profiles-2026-09-04": "user-profiles-2026-09-04"
},
"databricks": {
"advisor-tool-2026-03-01": null,

47
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -34,10 +34,12 @@ from litellm.completion_extras.litellm_responses_transformation.history import (
requested_phase,
)
from litellm.integrations.anthropic_cache_control_hook import supports_openai_prompt_cache_breakpoint
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params
from litellm.litellm_core_utils.prompt_templates.common_utils import (
prompt_cache_breakpoint_for_wire,
prompt_cache_breakpoint_with_readable_mode,
responses_reasoning_items_from_thinking_blocks,
whole_prompt_cache_breakpoint_or_none,
with_prompt_cache_breakpoint,
)
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
@ -97,6 +99,12 @@ _RESPONSES_API_ONLY_FIELDS: Final = frozenset((*Response.model_fields, *Response
)
def _bridged_prompt_cache_breakpoint(marker: object, drop_params: bool, target_trims_breakpoints: bool) -> object:
if target_trims_breakpoints:
return prompt_cache_breakpoint_with_readable_mode(marker, drop_params)
return whole_prompt_cache_breakpoint_or_none(marker) if drop_params else marker
def _strip_prompt_cache_breakpoint_from_content_block(value: object) -> object:
if not isinstance(value, dict):
return value
@ -844,6 +852,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
messages: list["AllMessageValues"],
*,
drop_params: bool = False,
target_trims_breakpoints: bool = True,
keep_prompt_cache_breakpoints: bool = False,
replay_tool_search: bool = False,
include_namespace: bool = False,
@ -852,6 +861,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
converted_input_items, instructions = self._convert_chat_completion_messages_to_responses_input(
messages,
drop_params=drop_params,
target_trims_breakpoints=target_trims_breakpoints,
replay_tool_search=replay_tool_search,
include_namespace=include_namespace,
replay_openai_history=replay_openai_history,
@ -868,6 +878,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
messages: list["AllMessageValues"],
*,
drop_params: bool = False,
target_trims_breakpoints: bool = True,
replay_tool_search: bool = False,
include_namespace: bool = False,
replay_openai_history: bool = False,
@ -921,6 +932,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
content,
role,
drop_params=drop_params,
target_trims_breakpoints=target_trims_breakpoints,
),
}
)
@ -940,6 +952,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
content,
"user", # Use "user" role to get input_* types
drop_params=drop_params,
target_trims_breakpoints=target_trims_breakpoints,
)
else:
# Fallback: convert unexpected types to input_text
@ -968,7 +981,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
"type": "message",
"role": "assistant",
"content": self._convert_content_to_responses_format(
content, "assistant", drop_params=drop_params
content,
"assistant",
drop_params=drop_params,
target_trims_breakpoints=target_trims_breakpoints,
),
**_assistant_phase_field(msg, replay_openai_history, drop_params),
}
@ -1009,7 +1025,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
"type": "message",
"role": role,
"content": self._convert_content_to_responses_format(
content, cast(str, role), drop_params=drop_params
content,
cast(str, role),
drop_params=drop_params,
target_trims_breakpoints=target_trims_breakpoints,
),
**_assistant_phase_field(msg, replay_openai_history and role == "assistant", drop_params),
}
@ -1135,6 +1154,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
converted_input_items, converted_instructions = self.convert_chat_completion_messages_to_responses_api(
messages,
drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params),
target_trims_breakpoints=drops_unsupported_params(litellm_params.get("drop_params")),
keep_prompt_cache_breakpoints=supports_prompt_cache_breakpoint,
replay_tool_search=_requests_tool_search(cast(object, optional_params)), # cast-ok: upcast
include_namespace=True,
@ -1640,6 +1660,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
| None,
role: str,
drop_params: bool = False,
target_trims_breakpoints: bool = True,
) -> list[dict[str, object]]:
"""Convert chat completion content to responses API format"""
from litellm.types.llms.openai import ChatCompletionImageObject
@ -1666,7 +1687,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if original_type == "text":
converted = with_prompt_cache_breakpoint(
self._convert_content_str_to_input_text(item.get("text", ""), role),
prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params),
_bridged_prompt_cache_breakpoint(
item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints
),
)
result.append(converted)
verbose_logger.debug("Chat provider: text -> %s", converted)
@ -1679,14 +1702,18 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
cast(ChatCompletionImageObject, item), role
),
),
prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params),
_bridged_prompt_cache_breakpoint(
item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints
),
)
result.append(converted)
verbose_logger.debug("Chat provider: image_url -> %s", converted)
elif original_type == "input_audio":
converted = with_prompt_cache_breakpoint(
{"type": "input_audio", "input_audio": item.get("input_audio")},
prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params),
_bridged_prompt_cache_breakpoint(
item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints
),
)
result.append(converted)
verbose_logger.debug("Chat provider: input_audio -> %s", converted)
@ -1702,7 +1729,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
_input_file_from_file_value(
cast("ChatCompletionFileObject", item).get("file"), # cast-ok: type tag checked
),
prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params),
_bridged_prompt_cache_breakpoint(
item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints
),
)
result.append(converted)
verbose_logger.debug("Chat provider: file -> %s", converted)
@ -1726,7 +1755,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
# Default to input_text for unknown types
converted = with_prompt_cache_breakpoint(
self._convert_content_str_to_input_text(str(item.get("text", item)), role),
prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params),
_bridged_prompt_cache_breakpoint(
item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints
),
)
result.append(converted)
verbose_logger.debug("Chat provider: unknown(%s) -> %s", original_type, converted)

76
litellm/cost_calculator.py
View file

@ -6,6 +6,7 @@ import time
from collections.abc import Mapping, Sequence
from datetime import datetime
from functools import lru_cache
from itertools import accumulate
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast
@ -3850,44 +3851,65 @@ def handle_realtime_transcription_cost_calculation(
- {"type": "duration", "seconds": <float>} → priced via input_cost_per_second
- {"type": "tokens", "input_tokens": ...} → priced via input/audio token cost
"""
completed_events: Final = [
cast(dict, result) for result in results if result.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE
]
if not completed_events:
configured_models: Final = tuple(accumulate(map(_session_transcription_model, results), _newer_model))
clips: Final = tuple(
(as_str_mapping(result.get("usage")) or {}, configured)
for result, configured in zip(results, configured_models, strict=True)
if result.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE
)
if not clips:
return 0.0
model_name: Final = _get_transcription_model_name_from_results(results) or litellm_model_name
model_info: Final = _cost_map_model_info(model_name, custom_llm_provider)
fallback_model: Final = _get_transcription_model_name_from_results(results) or litellm_model_name
clip_models: Final = tuple((usage, configured or fallback_model) for usage, configured in clips)
model_infos: Final = MappingProxyType(
{model: _cost_map_model_info(model, custom_llm_provider) for model in frozenset(m for _, m in clip_models)}
)
override_info: Final = (
_cost_map_model_info(custom_pricing_model, custom_llm_provider) if custom_pricing_model is not None else None
)
return sum(
(_transcription_usage_cost(dict(usage), model_infos[model], override_info) for usage, model in clip_models),
0.0,
)
total_cost = 0.0
for event in completed_events:
usage = event.get("usage") or {}
total_cost += _transcription_usage_cost(usage, model_info, override_info)
return total_cost
_TRANSCRIPTION_SESSION_EVENT_TYPES: Final = frozenset(
{"transcription_session.created", "transcription_session.updated", "session.created", "session.updated"}
)
def _str_field(mapping: Mapping[str, object] | None, key: str) -> str | None:
value: Final = None if mapping is None else mapping.get(key)
return value if isinstance(value, str) and value else None
def _configured_transcription_model(session: Mapping[str, object]) -> str | None:
audio_input: Final = as_str_mapping((as_str_mapping(session.get("audio")) or {}).get("input"))
nested: Final = _str_field(as_str_mapping((audio_input or {}).get("transcription")), "model")
return nested or _str_field(as_str_mapping(session.get("input_audio_transcription")), "model")
def _session_of(result: Mapping[str, object]) -> Mapping[str, object] | None:
return as_str_mapping(result.get("session")) if result.get("type") in _TRANSCRIPTION_SESSION_EVENT_TYPES else None
def _session_transcription_model(result: Mapping[str, object]) -> str | None:
session: Final = _session_of(result)
return None if session is None else _configured_transcription_model(session)
def _newer_model(current: str | None, newer: str | None) -> str | None:
return newer or current
def _get_transcription_model_name_from_results(
results: OpenAIRealtimeStreamList,
) -> str | None:
"""Resolve the ASR model from a transcription_session.* / session.* event."""
for result in results:
if result.get("type") in (
"transcription_session.created",
"transcription_session.updated",
"session.created",
"session.updated",
):
session = cast(dict, result).get("session", {}) or {}
transcription = ((session.get("audio", {}) or {}).get("input", {}) or {}).get(
"transcription", {}
) or session.get("input_audio_transcription", {})
model = (transcription or {}).get("model") or session.get("model")
if model:
return model
return None
return next(
(model for result in results if (model := _str_field(_session_of(result), "model"))),
None,
)
_TranscriptionRateKey: TypeAlias = Literal[

66
litellm/guardrails.py Normal file
View file

@ -0,0 +1,66 @@
"""Invoke configured gateway guardrails without generating a completion."""
from collections.abc import Mapping, Sequence
from enum import Enum
from typing import Final
import httpx
from pydantic import TypeAdapter
import litellm
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared factory has untyped optional parameters
)
from litellm.types.guardrails import NativeGuardrailResponse
_NATIVE_RESPONSE: Final[TypeAdapter[NativeGuardrailResponse]] = TypeAdapter(NativeGuardrailResponse)
class _InputType(Enum):
REQUEST = "request"
RESPONSE = "response"
async def apply(
guardrail_name: str,
*,
input: str | None = None,
response: str | None = None,
messages: Sequence[Mapping[str, object]] | None = None,
metadata: Mapping[str, object] | None = None,
api_base: str | None = None,
api_key: str | None = None,
timeout: float | httpx.Timeout | None = None,
) -> NativeGuardrailResponse:
"""Return a configured guardrail's native structured findings.
Supply exactly one of ``input`` or ``response``. Grounding guardrails can
use ``messages`` as evidence. Connection settings default to
``litellm.api_base`` and ``litellm.api_key``.
"""
if (input is None) == (response is None):
raise ValueError("Supply exactly one of input or response.")
resolved_base: Final = api_base if api_base is not None else litellm.api_base
if not resolved_base:
raise ValueError("Set api_base or litellm.api_base to your LiteLLM gateway URL.")
resolved_key: Final = api_key if api_key is not None else litellm.api_key
input_type: Final = _InputType.RESPONSE if response is not None else _InputType.REQUEST
payload: Final = dict[str, object](
guardrail_name=guardrail_name,
text=response if response is not None else input,
input_type=input_type.value,
response_format="native",
)
if messages is not None:
payload["messages"] = [dict(message) for message in messages]
if metadata is not None:
payload["metadata"] = dict(metadata)
gateway_root: Final = resolved_base.rstrip("/").removesuffix("/v1")
client: Final = get_async_httpx_client(llm_provider="litellm_guardrails")
result: Final = await client.post( # pyright: ignore[reportUnknownMemberType] # shared handler has untyped optional parameters
url=f"{gateway_root}/guardrails/apply_guardrail",
json=payload,
headers={"Authorization": f"Bearer {resolved_key}"} if resolved_key else {},
timeout=timeout,
)
return _NATIVE_RESPONSE.validate_json(result.content)

30
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -17,7 +17,7 @@ from collections.abc import Callable, Container, Iterable, Iterator, Mapping, Se
from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar, cast
from urllib.parse import urlparse
from pydantic import TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
@ -73,7 +73,7 @@ MAX_CACHE_CONTROL_BLOCKS: Final = 4
CACHE_BREAKPOINT_KEYS: Final = ("cache_control", "prompt_cache_breakpoint")
OPENAI_PROMPT_CACHE_BREAKPOINT_MIN_GPT_VERSION: Final = (5, 6)
_GPT_VERSION_PATTERN: Final = re.compile(r"^gpt-(\d+)(?:\.(\d+))?")
_GPT_VERSION_PATTERN: Final = re.compile(r"^gpt-(\d+)(?:\.(\d+)|-(0|[1-9]\d?)(?![\da-z]))?")
_FINE_TUNED_BASE_PATTERN: Final = re.compile(r"^ft:([^:]+)")
_O_SERIES_PATTERN: Final = re.compile(r"^o\d+(?:-|$)")
OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset(
@ -335,6 +335,30 @@ def openai_prompt_cache_breakpoint_support(model: str, custom_llm_provider: str
return _named_model_breakpoint_support(model)
class _DeclaredBaseModels(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True)
base_model: object = None
model_info: object = None
def names(self) -> tuple[str, ...]:
model_info: Final = _validated_object_mapping(self.model_info) or {}
candidates: Final = (self.base_model, model_info.get("base_model"))
return tuple(name for name in candidates if isinstance(name, str) and name)
def deployment_prompt_cache_breakpoint_support(
model: str, custom_llm_provider: str, litellm_params: object
) -> bool | None:
"""``openai_prompt_cache_breakpoint_support`` for a deployment, read from its declared base model first."""
declared: Final = _DeclaredBaseModels.model_validate(litellm_params)
supports: Final = (
openai_prompt_cache_breakpoint_support(name, custom_llm_provider)
for name in dict.fromkeys((*declared.names(), model))
)
return next((support for support in supports if support is not None), None)
def _named_model_breakpoint_support(model: str) -> bool | None:
bare_model: Final = model.rsplit("/", 1)[-1].lower()
if _O_SERIES_PATTERN.match(bare_model) is not None:
@ -342,7 +366,7 @@ def _named_model_breakpoint_support(model: str) -> bool | None:
version_match: Final = _GPT_VERSION_PATTERN.match(bare_model)
if version_match is None:
return None
version: Final = (int(version_match.group(1)), int(version_match.group(2) or 0))
version: Final = (int(version_match.group(1)), int(version_match.group(2) or version_match.group(3) or 0))
return version >= OPENAI_PROMPT_CACHE_BREAKPOINT_MIN_GPT_VERSION

6
litellm/litellm_core_utils/core_helpers.py
View file

@ -4,7 +4,7 @@ import copy
import logging
import math
import re
from collections.abc import Collection, Iterable, Iterator, Mapping
from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable
@ -104,6 +104,10 @@ def as_str_mapping(value: object) -> Mapping[str, object] | None:
return value if isinstance(value, Mapping) else None
def as_object_list(value: object) -> Sequence[object] | None:
return value if isinstance(value, list) else None
def decoded_json_object(value: object) -> Mapping[str, object]:
mapping: Final = as_str_mapping(value)
if mapping is not None:

6
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -146,6 +146,10 @@ def handle_anthropic_text_model_custom_llm_provider(
return model, custom_llm_provider
def declared_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None:
return custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None)
def declared_authenticating_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None:
"""The authenticating provider this pair already names, or None.
@ -154,7 +158,7 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str
and for a declared pair the resolver's answer is the declaration itself, so metadata callers
adopt the declaration instead of resolving.
"""
declared: Final = custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None)
declared: Final = declared_provider(model, custom_llm_provider)
return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None

85
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -29,6 +29,7 @@ from litellm.types.llms.openai import (
ChatCompletionAudioObject,
ChatCompletionFileObject,
ChatCompletionImageObject,
ChatCompletionPromptCacheBreakpoint,
ChatCompletionReasoningItem,
ChatCompletionReasoningSummaryTextBlock,
ChatCompletionRedactedThinkingBlock,
@ -2193,16 +2194,62 @@ def with_prompt_cache_breakpoint(target: _MarkedT, marker: object) -> _MarkedT:
_PROMPT_CACHE_BREAKPOINT: Final = TypeAdapter(PromptCacheBreakpoint)
_CHAT_COMPLETION_PROMPT_CACHE_BREAKPOINT: Final = TypeAdapter(ChatCompletionPromptCacheBreakpoint)
_AcceptedBreakpointT: Final = TypeVar("_AcceptedBreakpointT", bound=Mapping[str, object])
def _accepted_prompt_cache_breakpoint(
accepted: TypeAdapter[_AcceptedBreakpointT], marker: object
) -> _AcceptedBreakpointT | None:
try:
return accepted.validate_python(marker)
except ValidationError:
return None
def prompt_cache_breakpoint_within(
accepted: TypeAdapter[_AcceptedBreakpointT], marker: object
) -> _AcceptedBreakpointT | None:
"""The client's marker cut down to the ``accepted`` shape: an unsupported ttl is dropped and the mode kept,
and a marker whose mode the target cannot read is dropped whole."""
whole: Final = _accepted_prompt_cache_breakpoint(accepted, marker)
if whole is not None:
return whole
fields: Final = _breakpoint_fields_or_none(marker)
without_ttl: Final = (
None
if fields is None
else _accepted_prompt_cache_breakpoint(accepted, {key: item for key, item in fields.items() if key != "ttl"})
)
verbose_logger.debug(
"Dropping malformed prompt_cache_breakpoint %r under drop_params"
if without_ttl is None
else "Dropping the unsupported ttl of prompt_cache_breakpoint %r under drop_params",
marker,
)
return without_ttl
def prompt_cache_breakpoint_for_wire(marker: object, drop_params: bool) -> object:
if marker is None or not drop_params:
return marker
try:
return _PROMPT_CACHE_BREAKPOINT.validate_python(marker)
except ValidationError:
verbose_logger.debug("Dropping malformed prompt_cache_breakpoint %r under drop_params", marker)
return None
return prompt_cache_breakpoint_within(_PROMPT_CACHE_BREAKPOINT, marker)
def prompt_cache_breakpoint_with_readable_mode(marker: object, drop_params: bool) -> object:
"""The client's marker as written, or None under drop_params when no target could read its mode, so the
target's own config decides which ttl and keys it keeps."""
return None if prompt_cache_breakpoint_for_wire(marker, drop_params) is None else marker
def whole_prompt_cache_breakpoint_or_none(marker: object) -> PromptCacheBreakpoint | None:
return _accepted_prompt_cache_breakpoint(_PROMPT_CACHE_BREAKPOINT, marker)
def chat_completion_prompt_cache_breakpoint_for_wire(marker: object, drop_params: bool) -> object:
if marker is None or not drop_params:
return marker
return prompt_cache_breakpoint_within(_CHAT_COMPLETION_PROMPT_CACHE_BREAKPOINT, marker)
def cache_control_as_prompt_cache_breakpoint(block: _MarkedT) -> _MarkedT:
@ -2220,6 +2267,13 @@ _BREAKPOINT_FIELDS: Final = TypeAdapter(Mapping[str, object])
_BREAKPOINT_BLOCKS: Final = TypeAdapter(list[object])
def _breakpoint_fields_or_none(value: object) -> Mapping[str, object] | None:
try:
return _BREAKPOINT_FIELDS.validate_python(value)
except ValidationError:
return None
def _without_prompt_cache_breakpoint(value: object) -> object:
if not isinstance(value, Mapping):
return value
@ -2242,6 +2296,27 @@ def without_prompt_cache_breakpoints(block: _MarkedT) -> _MarkedT:
return cast(_MarkedT, stripped) # cast-ok: same block shape as the input minus the marker keys
def _part_with_chat_completion_prompt_cache_breakpoint(part: object) -> object:
if not isinstance(part, Mapping):
return part
fields: Final = _BREAKPOINT_FIELDS.validate_python(part)
marker: Final = fields.get("prompt_cache_breakpoint")
if marker is None:
return fields
unmarked: Final = {key: item for key, item in fields.items() if key != "prompt_cache_breakpoint"}
return with_prompt_cache_breakpoint(unmarked, chat_completion_prompt_cache_breakpoint_for_wire(marker, True))
def chat_completion_prompt_cache_breakpoints_for_wire(message: _MarkedT, drop_params: bool) -> _MarkedT:
content: Final = message.get("content")
if not drop_params or not isinstance(content, list):
return message
parts: Final = [
_part_with_chat_completion_prompt_cache_breakpoint(part) for part in _BREAKPOINT_BLOCKS.validate_python(content)
]
return cast(_MarkedT, {**message, "content": parts}) # cast-ok: same message shape with its parts' markers trimmed
LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS: Final = frozenset(
{"thinking_blocks", "reasoning_items", "provider_specific_fields"}
)

14
litellm/litellm_core_utils/streaming_handler.py
View file

@ -905,13 +905,21 @@ class CustomStreamWrapper:
except Exception as e:
raise e
def _served_openrouter_model(self) -> str | None:
if self.custom_llm_provider == LlmProviders.OPENROUTER.value or (
self.custom_llm_provider == "cached_response"
and self.logging_obj.model_call_details.get("custom_llm_provider") == LlmProviders.OPENROUTER.value
):
return self._provider_response_model
return None
def model_response_creator(
self,
chunk: dict | None = None,
hidden_params: Mapping[str, object] | None = None,
delta: Delta | None = None,
) -> ModelResponseStream:
_model: Final = self._cached_model_name
_model: Final = self._served_openrouter_model() or self._cached_model_name
_logging_obj_llm_provider: Final = self._cached_logging_llm_provider
if chunk is None:
@ -1699,7 +1707,7 @@ class CustomStreamWrapper:
return dispatch_result.value
response_obj = dispatch_result.response_obj
model_response.model = self.model
model_response.model = self._served_openrouter_model() or self.model
## FUNCTION CALL PARSING
original_chunk: Final[_UpstreamChunkLike | None] = (
response_obj.get("original_chunk") if response_obj is not None else None
@ -2561,7 +2569,7 @@ class CustomStreamWrapper:
if partial is None or not isinstance(usage, Usage):
return None
if self.model:
partial.model = self.model
partial.model = self._served_openrouter_model() or self.model
backfill_missing_cache_usage_fields(usage)
return partial

20
litellm/llms/anthropic/common_utils.py
View file

@ -1297,7 +1297,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
return None
return anthropic_beta_header.split(",")
def get_computer_tool_beta_header(self, computer_tool_version: str) -> str:
def get_computer_tool_beta_header(self, computer_tool_version: str) -> str | None:
"""
Get the appropriate beta header for a given computer tool version.
@ -1305,16 +1305,14 @@ class AnthropicModelInfo(BaseLLMModelInfo):
computer_tool_version: The computer tool version (e.g., 'computer_20250124', 'computer_20241022')
Returns:
The corresponding beta header string
The corresponding beta header string, or None for a version with no known beta
"""
computer_tool_beta_mapping: Final = {
"computer_20251124": "computer-use-2025-11-24",
"computer_20250124": "computer-use-2025-01-24",
"computer_20241022": "computer-use-2024-10-22",
}
return computer_tool_beta_mapping.get(
computer_tool_version,
"computer-use-2024-10-22", # Default fallback
)
return computer_tool_beta_mapping.get(computer_tool_version)
def get_anthropic_beta_list(
self,
@ -1346,9 +1344,8 @@ class AnthropicModelInfo(BaseLLMModelInfo):
if effort_used:
betas.append(ANTHROPIC_EFFORT_BETA_HEADER) # effort-2025-11-24
if computer_tool_used:
beta_header: Final = self.get_computer_tool_beta_header(computer_tool_used)
betas.append(beta_header)
if computer_tool_used and (computer_tool_beta := self.get_computer_tool_beta_header(computer_tool_used)):
betas.append(computer_tool_beta)
# Anthropic no longer requires the prompt-caching beta header
# Prompt caching now works automatically when cache_control is used in messages
@ -1437,9 +1434,8 @@ class AnthropicModelInfo(BaseLLMModelInfo):
# Anthropic no longer requires the prompt-caching beta header
# Prompt caching now works automatically when cache_control is used in messages
# Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching
if computer_tool_used:
beta_header: Final = self.get_computer_tool_beta_header(computer_tool_used)
betas.add(beta_header)
if computer_tool_used and (computer_tool_beta := self.get_computer_tool_beta_header(computer_tool_used)):
betas.add(computer_tool_beta)
# if pdf_used:
# betas.add("pdfs-2024-09-25")
if file_id_used:

138
litellm/llms/anthropic/messages/context_management/inference_view.py Normal file
View file

@ -0,0 +1,138 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Final
from pydantic import JsonValue, TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm.types.llms.anthropic import AnthropicMessagesTextParam
from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat
@dataclass(frozen=True, slots=True)
class InferenceRequestView:
messages: Sequence[Mapping[str, object]]
system: object | None
@dataclass(frozen=True, slots=True)
class InferenceRequestViewError:
error_message: str
status_code: int
@dataclass(frozen=True, slots=True)
class CountContextEdits:
spec: JsonValue
model: str
tools: tuple[Mapping[str, object], ...] | None
drop_invalid_edits: bool = False
_CONTEXT_MANAGEMENT_SPEC: Final[TypeAdapter[Mapping[str, object] | tuple[Mapping[str, object], ...]]] = TypeAdapter(
Mapping[str, object] | tuple[Mapping[str, object], ...]
)
_SYSTEM_BLOCKS: Final = TypeAdapter(tuple[AnthropicMessagesTextParam, ...])
def count_context_edits(
context_management: CountContextManagement | None,
tools: Sequence[Mapping[str, object]] | None,
litellm_params: Mapping[str, object],
fallback_model: str,
) -> CountContextEdits:
model: Final = litellm_params.get("model")
return CountContextEdits(
spec=None if context_management is None else context_management.spec,
model=model if isinstance(model, str) else fallback_model,
tools=None if tools is None else tuple(tools),
drop_invalid_edits=context_management is not None and context_management.drop_invalid_edits,
)
def _typed_spec(spec: JsonValue) -> Mapping[str, object] | tuple[Mapping[str, object], ...] | None:
try:
return _CONTEXT_MANAGEMENT_SPEC.validate_python(spec)
except ValidationError:
return None
def _typed_system(system: object | None) -> str | tuple[AnthropicMessagesTextParam, ...] | None | ValidationError:
if system is None or isinstance(system, str):
return system
try:
return _SYSTEM_BLOCKS.validate_python(system)
except ValidationError as error:
return error
def _view_without_edits(
messages: Sequence[Mapping[str, object]],
system: object | None,
typed_system: str | Sequence[Mapping[str, object]] | None,
) -> InferenceRequestView:
from litellm.llms.anthropic.messages.context_management.editors.compact import (
apply_client_compaction_block_history,
)
history: Final = apply_client_compaction_block_history(
messages=[dict(message) for message in messages],
system=typed_system
if typed_system is None or isinstance(typed_system, str)
else [dict(block) for block in typed_system],
)
if history is None:
return InferenceRequestView(messages=messages, system=system)
return InferenceRequestView(messages=history.messages, system=history.system)
async def _anthropic_inference_view(
messages: Sequence[Mapping[str, object]], system: object | None, edits: CountContextEdits
) -> InferenceRequestView | InferenceRequestViewError:
from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request
from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError
from litellm.llms.anthropic.messages.native.utils import normalize_native_messages
typed_system: Final = _typed_system(system)
if isinstance(typed_system, ValidationError):
return InferenceRequestView(messages=messages, system=system)
normalized: Final = normalize_native_messages(messages)
spec: Final = _typed_spec(edits.spec)
try:
prepared: Final = await prepare_context_managed_request(
model=edits.model,
messages=normalized,
tools=None if edits.tools is None else [dict(tool) for tool in edits.tools],
system=[dict(block) for block in typed_system] if isinstance(typed_system, tuple) else typed_system,
context_management_spec=(
None if spec is None else dict(spec) if isinstance(spec, Mapping) else [dict(edit) for edit in spec]
),
litellm_metadata=None,
additional_drop_params=None,
llm_router=None,
count_only=True,
drop_invalid_edits=edits.drop_invalid_edits,
)
except AnthropicContextManagementError as error:
if error.status_code >= 500:
return _view_without_edits(normalized, system, typed_system)
return InferenceRequestViewError(error_message=error.message, status_code=error.status_code)
if prepared is None:
return InferenceRequestView(messages=normalized, system=system)
return InferenceRequestView(messages=prepared.messages, system=prepared.system)
async def inference_request_view(
messages: Sequence[Mapping[str, object]],
system: object | None,
messages_format: TokenCountMessagesFormat,
edits: CountContextEdits,
) -> InferenceRequestView | InferenceRequestViewError:
"""Apply the history rewrites and context edits /v1/messages runs before it converts a body for another provider."""
match messages_format:
case "openai":
return InferenceRequestView(messages=messages, system=system)
case "anthropic":
return await _anthropic_inference_view(messages, system, edits)
case _:
assert_never(messages_format)

100
litellm/llms/anthropic/messages/native/handler.py
View file

@ -8,6 +8,7 @@
import asyncio
import contextvars
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence
from dataclasses import dataclass
from functools import partial
from typing import Final, cast
@ -15,6 +16,7 @@ from pydantic import TypeAdapter
import litellm
from litellm.anthropic_beta_headers_manager import aload_beta_headers_config
from litellm.constants import PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.exception_mapping_utils import exception_type
@ -44,7 +46,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import CallTypes, LlmProviders
from litellm.types.utils import CallTypes, LlmProviders, LlmProvidersSet
from litellm.utils import ProviderConfigManager, client
from ..adapters.handler import LiteLLMMessagesToCompletionTransformationHandler, request_safeguards_evaluator
@ -94,6 +96,68 @@ def _responses_mode_is_lost_by_prefix_strip(
) and not _bridges_to_responses_api(resolved_model, custom_llm_provider, deployment_model_info)
@dataclass(frozen=True, slots=True)
class _BridgeTarget:
model: str
api_key: str | None
api_base: str | None
def _resolve_requested_provider(
model: str, custom_llm_provider: str | None, litellm_params: GenericLiteLLMParams
) -> tuple[str, str, str | None, str | None]:
resolved: Final = litellm.get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=litellm_params.api_base,
api_key=litellm_params.api_key,
litellm_params=litellm_params,
)
requested_provider: Final = model.partition("/")[0]
if (
custom_llm_provider in (None, requested_provider)
or requested_provider not in LlmProvidersSet
or requested_provider in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO
):
return resolved
own: Final = _resolve_by_model_prefix(model, litellm_params)
return own if own is not None and own[1] == custom_llm_provider else resolved
def _resolve_by_model_prefix(
model: str, litellm_params: GenericLiteLLMParams
) -> tuple[str, str, str | None, str | None] | None:
try:
return litellm.get_llm_provider(
model=model,
api_base=litellm_params.api_base,
api_key=litellm_params.api_key,
litellm_params=litellm_params.model_copy(update={"custom_llm_provider": None}),
)
except litellm.BadRequestError:
return None
def _sets_own_entra_credential(litellm_params: GenericLiteLLMParams) -> bool:
return bool(
litellm_params.azure_ad_token
or (litellm_params.model_extra or {}).get("azure_ad_token_provider") is not None
or (litellm_params.tenant_id and litellm_params.client_id and litellm_params.client_secret)
or (litellm_params.azure_username and litellm_params.azure_password and litellm_params.client_id)
)
def _bridge_target(requested: _BridgeTarget, resolved: _BridgeTarget, custom_llm_provider: str) -> _BridgeTarget:
requested_provider: Final = requested.model.partition("/")[0]
if requested_provider == custom_llm_provider or requested_provider not in LlmProvidersSet:
return requested
return _BridgeTarget(
model=f"{custom_llm_provider}/{resolved.model}",
api_key=requested.api_key if requested.api_base else requested.api_key or resolved.api_key,
api_base=resolved.api_base,
)
def _should_route_to_responses_api(
custom_llm_provider: str | None,
requested_model: str | None = None,
@ -624,13 +688,7 @@ def anthropic_messages_handler(
custom_llm_provider,
dynamic_api_key,
dynamic_api_base,
) = litellm.get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=litellm_params.api_base,
api_key=litellm_params.api_key,
litellm_params=litellm_params,
)
) = _resolve_requested_provider(model, custom_llm_provider, litellm_params)
# Store agentic loop params in logging object for agentic hooks
# This provides original request context needed for follow-up calls
@ -656,6 +714,16 @@ def anthropic_messages_handler(
mock_response=litellm_params.mock_response,
)
bridge: Final = _bridge_target(
_BridgeTarget(model=original_model, api_key=api_key, api_base=api_base),
_BridgeTarget(
model=model,
api_key=None if _sets_own_entra_credential(litellm_params) else dynamic_api_key,
api_base=dynamic_api_base,
),
custom_llm_provider,
)
# Expand litellm_proxy MCP references through the MCP gateway before dispatch, so every
# downstream path (native passthrough and both bridges) gets real tools rather than a
# reference the provider cannot resolve. Popped from kwargs so it never reaches the provider.
@ -682,8 +750,8 @@ def anthropic_messages_handler(
top_k=top_k,
top_p=top_p,
container=container,
api_key=api_key,
api_base=api_base,
api_key=bridge.api_key,
api_base=bridge.api_base,
client=client,
custom_llm_provider=custom_llm_provider,
**kwargs,
@ -700,7 +768,7 @@ def anthropic_messages_handler(
return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler(
max_tokens=max_tokens,
messages=messages,
model=original_model,
model=bridge.model,
metadata=metadata,
stop_sequences=stop_sequences,
stream=stream,
@ -716,8 +784,8 @@ def anthropic_messages_handler(
top_k=top_k,
top_p=top_p,
_is_async=is_async,
api_key=api_key,
api_base=api_base,
api_key=bridge.api_key,
api_base=bridge.api_base,
client=client,
custom_llm_provider=custom_llm_provider,
**kwargs,
@ -730,7 +798,7 @@ def anthropic_messages_handler(
return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler(
max_tokens=max_tokens,
messages=messages,
model=original_model,
model=bridge.model,
metadata=metadata,
stop_sequences=stop_sequences,
stream=stream,
@ -744,8 +812,8 @@ def anthropic_messages_handler(
top_k=top_k,
top_p=top_p,
_is_async=is_async,
api_key=api_key,
api_base=api_base,
api_key=bridge.api_key,
api_base=bridge.api_base,
client=client,
custom_llm_provider=custom_llm_provider,
**kwargs,

4
litellm/llms/anthropic/messages/responses_adapters/handler.py
View file

@ -21,7 +21,7 @@ from litellm.litellm_core_utils.asyncify import (
from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params
from litellm.litellm_core_utils.prompt_templates.common_utils import (
cache_control_as_prompt_cache_breakpoint,
prompt_cache_breakpoint_for_wire,
prompt_cache_breakpoint_with_readable_mode,
with_prompt_cache_breakpoint,
)
from litellm.responses.utils import auto_reasoning_summary
@ -207,7 +207,7 @@ def _without_prompt_cache_breakpoint(block: Mapping[str, object]) -> Mapping[str
def _cache_control_as_wire_safe_prompt_cache_breakpoint(block: Mapping[str, object]) -> Mapping[str, object]:
if "prompt_cache_breakpoint" not in block:
return cache_control_as_prompt_cache_breakpoint(block)
wire_marker: Final = prompt_cache_breakpoint_for_wire(block["prompt_cache_breakpoint"], drop_params=True)
wire_marker: Final = prompt_cache_breakpoint_with_readable_mode(block["prompt_cache_breakpoint"], drop_params=True)
return cache_control_as_prompt_cache_breakpoint(
with_prompt_cache_breakpoint(_without_prompt_cache_breakpoint(block), wire_marker)
)

4
litellm/llms/azure/chat/gpt_reasoning_transformation.py
View file

@ -1,6 +1,6 @@
"""Support for Azure OpenAI gpt-5 model family."""
from collections.abc import Collection
from collections.abc import Collection, Mapping
from typing import ClassVar, Final
import litellm
@ -105,7 +105,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
model: str,
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
litellm_params: Mapping[str, object],
headers: dict,
) -> dict:
model = model.replace(self.GPT5_SERIES_ROUTE, "")

21
litellm/llms/azure/chat/gpt_transformation.py
View file

@ -5,12 +5,16 @@ from typing import TYPE_CHECKING, Any, Final
from httpx._models import Headers, Response
import litellm
from litellm.integrations.anthropic_cache_control_hook import deployment_prompt_cache_breakpoint_support
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.litellm_core_utils.prompt_templates.common_utils import (
chat_completion_prompt_cache_breakpoints_for_wire,
drop_tool_reference_parts_from_tool_messages,
flatten_combinators_and_drop_non_python_regex_patterns,
hoist_images_from_tool_messages,
system_messages_first,
tool_with_sanitized_parameters,
without_prompt_cache_breakpoints,
)
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_azure_openai_messages,
@ -290,11 +294,16 @@ class AzureOpenAIConfig(BaseConfig):
model: str,
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
litellm_params: Mapping[str, object],
headers: dict,
) -> dict:
ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages
stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages)
wire_messages: Final = (
self._with_wire_safe_prompt_cache_breakpoints(model, ordered_messages, litellm_params)
if drops_unsupported_params(litellm_params.get("drop_params"))
else ordered_messages
)
stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(wire_messages)
azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages))
request_params: Final = MappingProxyType(
{
@ -310,6 +319,14 @@ class AzureOpenAIConfig(BaseConfig):
**sanitized_tools_update(optional_params),
}
@staticmethod
def _with_wire_safe_prompt_cache_breakpoints(
model: str, messages: list[AllMessageValues], litellm_params: Mapping[str, object]
) -> list[AllMessageValues]:
if deployment_prompt_cache_breakpoint_support(model, "azure", litellm_params) is False:
return [without_prompt_cache_breakpoints(message) for message in messages]
return [chat_completion_prompt_cache_breakpoints_for_wire(message, drop_params=True) for message in messages]
def transform_response(
self,
model: str,

12
litellm/llms/azure/responses/transformation.py
View file

@ -1,6 +1,7 @@
import re
from collections.abc import Mapping
from collections.abc import Callable, Mapping
from copy import deepcopy
from functools import partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal
from urllib.parse import urlsplit
@ -11,6 +12,7 @@ from pydantic import TypeAdapter
from litellm._logging import verbose_logger
from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION
from litellm.litellm_core_utils.prompt_templates.common_utils import prompt_cache_breakpoint_within
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config
from litellm.llms.azure.common_utils import BaseAzureLLM
@ -29,6 +31,7 @@ else:
LiteLLMLoggingObj = Any
_DECODED_JSON: Final = TypeAdapter(object)
_AZURE_RESPONSES_PROMPT_CACHE_BREAKPOINT: Final = TypeAdapter(AzureResponsesPromptCacheBreakpoint)
_DATED_GA_API_VERSION: Final = re.compile(r"\d{4}-\d{2}-\d{2}")
@ -81,6 +84,13 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
def _is_gpt_5_model(model: str) -> bool:
return AzureOpenAIGPT5Config.is_model_gpt_5_model(model)
def _prompt_cache_breakpoint_to_wire(
self, litellm_params: GenericLiteLLMParams, reads_breakpoints: bool, known_support: bool | None
) -> Callable[[object], object]:
if reads_breakpoints or known_support is True:
return partial(prompt_cache_breakpoint_within, _AZURE_RESPONSES_PROMPT_CACHE_BREAKPOINT)
return super()._prompt_cache_breakpoint_to_wire(litellm_params, reads_breakpoints, known_support)
@staticmethod
def _supports_reasoning_effort_none(model: str) -> bool:
return AzureOpenAIGPT5Config.supports_reasoning_effort_level(model, "none")

6
litellm/llms/azure_ai/anthropic/count_tokens/handler.py
View file

@ -12,6 +12,7 @@ from pydantic import JsonValue
import litellm
from litellm._logging import verbose_logger
from litellm.anthropic_beta_headers_manager import aload_beta_headers_config
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.llms.anthropic.common_utils import AnthropicError
@ -39,6 +40,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
tools: list[dict[str, JsonValue]] | None = None,
system: object = None,
optional_params: Mapping[str, JsonValue] | None = None,
context_management: JsonValue = None,
) -> dict[str, object]:
"""
Handle a CountTokens request using httpx with Azure authentication.
@ -70,6 +72,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
tools=tools,
system=system,
optional_params=optional_params,
context_management=context_management,
drop_params=drops_unsupported_params((litellm_params or {}).get("drop_params")),
)
@ -80,11 +83,14 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
verbose_logger.debug("Making request to: %s", endpoint_url)
if context_management is not None:
await aload_beta_headers_config()
# Get required headers with Azure authentication
headers: Final = self.get_required_headers(
api_key=api_key,
litellm_params=litellm_params,
messages=messages,
context_management=context_management,
)
if not {"x-api-key", "Authorization"} & headers.keys():
raise AnthropicError(

1
litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py
View file

@ -135,6 +135,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
tools=tools,
system=system,
optional_params=count_options,
context_management=None if context_management is None else context_management.spec,
)
if result is not None:

17
litellm/llms/azure_ai/anthropic/count_tokens/transformation.py
View file

@ -8,9 +8,16 @@ from collections.abc import Sequence
from typing import Final
from urllib.parse import urlsplit
from pydantic import JsonValue
from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.anthropic.common_utils import count_tokens_feature_betas, merge_anthropic_beta_headers
from litellm.llms.anthropic.common_utils import (
context_management_betas,
count_tokens_feature_betas,
merge_anthropic_beta_headers,
)
from litellm.llms.anthropic.count_tokens.transformation import (
AnthropicCountTokensConfig,
)
@ -31,6 +38,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
api_key: str | None,
litellm_params: dict[str, object] | None = None,
messages: Sequence[object] = (),
context_management: JsonValue = None,
) -> dict[str, str]:
"""
Get the required headers for the Azure AI Anthropic CountTokens API.
@ -42,6 +50,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
api_key: The Azure AI API key
litellm_params: Optional LiteLLM parameters for additional auth config
messages: The request messages, read for betas the counted features need
context_management: The edit spec the count forwards, read for the betas its edits need
Returns:
Dictionary of required headers with both x-api-key and Azure authentication
@ -51,7 +60,11 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
"Content-Type": "application/json",
"anthropic-version": "2023-06-01",
"anthropic-beta": merge_anthropic_beta_headers(
ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, count_tokens_feature_betas("azure_ai", {}, messages)
ANTHROPIC_TOKEN_COUNTING_BETA_VERSION,
(
*count_tokens_feature_betas("azure_ai", {}, messages),
*filter_and_transform_beta_headers(list(context_management_betas(context_management)), "azure_ai"),
),
),
**({"x-api-key": api_key} if api_key else {}),
}

7
litellm/llms/bedrock/common_utils.py
View file

@ -43,6 +43,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned
from litellm.secret_managers.main import get_secret, get_secret_str
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams
if TYPE_CHECKING:
@ -50,6 +51,12 @@ if TYPE_CHECKING:
from litellm.types.llms.openai import AllMessageValues
BEDROCK_INVOKE_TRACE_FIELD: Final = "amazon-bedrock-trace"
BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS: Final[Mapping[str, str]] = MappingProxyType(
{
"compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
"clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
}
)
_ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs"

68
litellm/llms/bedrock/count_tokens/bedrock_token_counter.py
View file

@ -13,9 +13,14 @@ from typing_extensions import assert_never
import litellm
from litellm._logging import verbose_logger
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.common_utils import context_management_betas, context_management_with_edits
from litellm.llms.anthropic.count_tokens.transformation import ChatCountTranslationError, anthropic_count_fields
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model
from litellm.llms.bedrock.common_utils import (
BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS,
BedrockError,
get_bedrock_base_model,
)
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler
from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model
@ -74,6 +79,16 @@ def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model:
)
def _mantle_retry_outcome(
runtime: CountTokensOutcome, mantle: CountTokensOutcome, carried_edits: bool
) -> CountTokensOutcome:
match runtime, mantle:
case CountTokensFailure(status_code=400), CountTokensFailure(status_code=404) if carried_edits:
return runtime
case _:
return mantle
def _request_data(
model: str, messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] | None, system: object
) -> Mapping[str, object]:
@ -85,6 +100,33 @@ def _request_data(
}
def _invoke_model_message(message: Mapping[str, object]) -> Mapping[str, object]:
content: Final = message.get("content")
return (
{**message, "content": [{"type": "text", "text": content}]} if isinstance(content, str) and content else message
)
def _invoke_supports_edit(edit: JsonValue) -> bool:
edit_type: Final = edit.get("type") if isinstance(edit, dict) else None
return isinstance(edit_type, str) and edit_type in BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS
def _invoke_context_management(context_management: CountContextManagement | None) -> Mapping[str, JsonValue]:
spec: Final = context_management_with_edits(None if context_management is None else context_management.spec)
edits: Final = spec.get("edits") if isinstance(spec, dict) else None
supported_edits: Final = tuple(filter(_invoke_supports_edit, edits if isinstance(edits, list) else ()))
if not isinstance(spec, dict) or not supported_edits:
return MappingProxyType({})
supported_spec: Final[JsonValue] = {**spec, "edits": list(supported_edits)}
return MappingProxyType(
{
"context_management": supported_spec,
"anthropic_beta": list(context_management_betas(supported_spec)),
}
)
class BedrockTokenCounter(BaseTokenCounter):
"""Counts on the configured Bedrock endpoint, with Mantle retry when Runtime rejects a Claude model"""
@ -167,11 +209,19 @@ class BedrockTokenCounter(BaseTokenCounter):
)
if isinstance(translated, ChatCountTranslationError):
return translated.response(request_model, model_to_use, primary_tokenizer)
runtime_request: Final = (
_request_data(model_to_use, messages, tools, system)
if translated is None
else _request_data(model_to_use, translated.messages, translated.tools, translated.system)
native_context_management: Final = _invoke_context_management(
context_management if is_mantle_claude_model(resolved_model) else None
)
counted_messages: Final = messages if translated is None else translated.messages
runtime_request: Final = {
**_request_data(
model_to_use,
tuple(map(_invoke_model_message, counted_messages)) if native_context_management else counted_messages,
tools if translated is None else translated.tools,
system if translated is None else translated.system,
),
**native_context_management,
}
mantle_request: Final = runtime_request if translated is None else {**runtime_request, **translated.options}
primary: Final = (
@ -184,8 +234,12 @@ class BedrockTokenCounter(BaseTokenCounter):
)
)
outcome: Final = (
await self._count_with(
self._mantle_handler, MANTLE_TOKENIZER_TYPE, mantle_request, litellm_params, resolved_model
_mantle_retry_outcome(
primary,
await self._count_with(
self._mantle_handler, MANTLE_TOKENIZER_TYPE, mantle_request, litellm_params, resolved_model
),
carried_edits=bool(native_context_management),
)
if not use_mantle and _runtime_rejected_claude_model(primary, resolved_model)
else primary

11
litellm/llms/bedrock/count_tokens/mantle_handler.py
View file

@ -8,7 +8,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.anthropic.common_utils import tools_without_custom_field
from litellm.llms.anthropic.common_utils import context_management_betas, tools_without_custom_field
from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url
@ -22,6 +22,7 @@ class MantleCountTokensRequest(TypedDict):
messages: ReadOnly[list[dict[str, JsonValue]]]
system: ReadOnly[NotRequired[JsonValue]]
tools: ReadOnly[NotRequired[list[dict[str, JsonValue]]]]
context_management: ReadOnly[NotRequired[JsonValue]]
_COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest)
@ -74,13 +75,19 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM):
optional_params=_COUNT_OPTIONS.validate_python(
{name: value for name, value in request_data.items() if name in COUNT_TOKEN_OPTION_NAMES}
),
context_management=request.get("context_management"),
)
betas: Final = context_management_betas(request.get("context_management"))
verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url)
api_key: Final = litellm_params.get("api_key")
signed_headers, signed_body = await run_aws_signing(
self._sign_request,
service_name="bedrock",
headers={"Content-Type": "application/json", "anthropic-version": MANTLE_ANTHROPIC_VERSION},
headers={
"Content-Type": "application/json",
"anthropic-version": MANTLE_ANTHROPIC_VERSION,
**({"anthropic-beta": ",".join(betas)} if betas else {}),
},
optional_params=litellm_params,
request_data=body,
api_base=endpoint_url,

8
litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py
View file

@ -37,6 +37,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation
bedrock_invoke_guardrail_headers,
)
from litellm.llms.bedrock.common_utils import (
BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS,
BedrockError,
apply_bedrock_invoke_structured_output,
bedrock_supports_tool_search,
@ -499,12 +500,7 @@ class AmazonAnthropicClaudeMessagesConfig(
# Bedrock InvokeModel DOES support ``clear_tool_uses_20250919`` under the
# ``context-management-2025-06-27`` beta. AWS docs:
# https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-tool-use.md
_BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = MappingProxyType(
{
"compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
"clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
}
)
_BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS
@classmethod
def _filter_context_management_for_bedrock_invoke(

15
litellm/llms/bedrock_mantle/responses/transformation.py
View file

@ -15,7 +15,8 @@ role / access key / profile / web identity), signed via the shared
BaseAWSLLM._sign_request after the request body is finalized.
"""
from collections.abc import Mapping, Sequence
from collections.abc import Callable, Mapping, Sequence
from functools import partial
from typing import Final, cast # noqa: TID251 # map_openai_params returns the filtered params as a bare dict
import httpx
@ -23,6 +24,7 @@ import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.litellm_core_utils.prompt_templates.common_utils import chat_completion_prompt_cache_breakpoint_for_wire
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.responses.codex_compat import drop_unsupported_tools, normalize_codex_input_items
@ -62,6 +64,10 @@ _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"})
_BEDROCK_MANTLE_OPENAI_PATH_SUPPORTED_REASONING_SUMMARIES: Final = frozenset({"auto"})
def _prompt_cache_breakpoint_as_sent_or_none(marker: object) -> object:
return marker if chat_completion_prompt_cache_breakpoint_for_wire(marker, drop_params=True) == marker else None
class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPIConfig):
def __init__(
self,
@ -113,6 +119,13 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
headers["OpenAI-Project"] = litellm_params.aws_bedrock_project_id
return headers
def _prompt_cache_breakpoint_to_wire(
self, litellm_params: GenericLiteLLMParams, reads_breakpoints: bool, known_support: bool | None
) -> Callable[[object], object]:
if reads_breakpoints or known_support is True:
return partial(chat_completion_prompt_cache_breakpoint_for_wire, drop_params=True)
return _prompt_cache_breakpoint_as_sent_or_none
def supports_native_file_search(self) -> bool:
return False

5
litellm/llms/custom_httpx/aiohttp_handler.py
View file

@ -14,6 +14,7 @@ from litellm.llms.base_llm.chat.transformation import BaseConfig
from litellm.llms.base_llm.image_variations.transformation import (
BaseImageVariationConfig,
)
from litellm.llms.custom_httpx.aiohttp_pool_diagnostics import build_pool_wait_trace_config
from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@ -90,7 +91,7 @@ class BaseLLMAIOHTTPHandler:
if self.transport is None:
connector: Final = self._get_connector()
if connector:
return aiohttp.ClientSession(connector=connector)
return aiohttp.ClientSession(connector=connector, trace_configs=[build_pool_wait_trace_config()])
transport: Final = self.transport or self._get_or_create_transport()
if transport is not None and hasattr(transport, "_get_valid_client_session"):
@ -99,7 +100,7 @@ class BaseLLMAIOHTTPHandler:
except RuntimeError:
pass
return aiohttp.ClientSession()
return aiohttp.ClientSession(trace_configs=[build_pool_wait_trace_config()])
def _get_async_client_session(self, dynamic_client_session: ClientSession | None = None) -> ClientSession:
if dynamic_client_session:

130
litellm/llms/custom_httpx/aiohttp_pool_diagnostics.py Normal file
View file

@ -0,0 +1,130 @@
import asyncio
from types import SimpleNamespace
from typing import Final
from aiohttp import (
ClientSession,
TraceConfig,
TraceConnectionQueuedEndParams,
TraceConnectionQueuedStartParams,
TraceRequestExceptionParams,
TraceRequestRedirectParams,
TraceRequestStartParams,
)
from yarl import URL
from litellm._logging import verbose_logger
_POOL_WAIT_WARNING_SECONDS: Final = 1.0
class _PoolWaitContext(SimpleNamespace):
def __init__(self, trace_request_ctx: object = None) -> None:
super().__init__()
self.host: str | None = None
self.port: int | None = None
self.started_at: float | None = None
self.waited: float = 0.0
self.warning: asyncio.TimerHandle | None = None
self.warned: bool = False
def set_destination(self, params: TraceRequestStartParams) -> None:
self.host = params.url.host
self.port = params.url.port
def redirect(self, params: TraceRequestRedirectParams) -> None:
location: Final = params.response.headers.get("Location") or params.response.headers.get("URI")
if location is None:
return
try:
destination: Final = params.url.join(URL(location))
host: Final = destination.host
port: Final = destination.port
except ValueError:
return
self.host = host
self.port = port
def start(self, session: ClientSession) -> None:
loop: Final = asyncio.get_running_loop()
self.started_at = loop.time()
if not self.warned:
self.warning = loop.call_later(max(0.0, _POOL_WAIT_WARNING_SECONDS - self.waited), self.warn, session)
def elapsed(self) -> float:
return self.waited + (
asyncio.get_running_loop().time() - self.started_at if self.started_at is not None else 0.0
)
def warn(self, session: ClientSession) -> None:
if self.warned or self.started_at is None:
return
self.warned = True
connector: Final = session.connector
if connector is None:
return
verbose_logger.warning(
"aiohttp connection pool saturated: host=%s port=%s waited=%.3fs for a connection slot "
"(acquired_total=%s, limit=%s, limit_per_host=%s). "
"Wait occurred before sending the request upstream.",
self.host,
self.port,
self.elapsed(),
len(connector._acquired), # pyright: ignore[reportPrivateUsage] # aiohttp exposes no public acquired count
connector.limit,
connector.limit_per_host,
)
def finish(self, session: ClientSession, failed: bool = False) -> None:
if self.warning is not None:
self.warning.cancel()
self.warning = None
if self.started_at is None:
return
if failed or self.elapsed() >= _POOL_WAIT_WARNING_SECONDS:
self.warn(session)
self.waited = self.elapsed()
self.started_at = None
async def _on_request_start(session: ClientSession, context: SimpleNamespace, params: TraceRequestStartParams) -> None:
if isinstance(context, _PoolWaitContext):
context.set_destination(params)
async def _on_queued_start(
session: ClientSession, context: SimpleNamespace, params: TraceConnectionQueuedStartParams
) -> None:
if isinstance(context, _PoolWaitContext):
context.start(session)
async def _on_queued_end(
session: ClientSession, context: SimpleNamespace, params: TraceConnectionQueuedEndParams
) -> None:
if isinstance(context, _PoolWaitContext):
context.finish(session)
async def _on_request_exception(
session: ClientSession, context: SimpleNamespace, params: TraceRequestExceptionParams
) -> None:
if isinstance(context, _PoolWaitContext):
context.finish(session, failed=True)
async def _on_request_redirect(
session: ClientSession, context: SimpleNamespace, params: TraceRequestRedirectParams
) -> None:
if isinstance(context, _PoolWaitContext):
context.redirect(params)
def build_pool_wait_trace_config() -> TraceConfig:
trace: Final = TraceConfig(trace_config_ctx_factory=_PoolWaitContext)
trace.on_request_start.append(_on_request_start)
trace.on_connection_queued_start.append(_on_queued_start)
trace.on_connection_queued_end.append(_on_queued_end)
trace.on_request_exception.append(_on_request_exception)
trace.on_request_redirect.append(_on_request_redirect)
return trace

2
litellm/llms/custom_httpx/http_handler.py
View file

@ -44,6 +44,7 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
from litellm.litellm_core_utils.request_timeout_resolver import (
get_configured_request_timeout,
)
from litellm.llms.custom_httpx.aiohttp_pool_diagnostics import build_pool_wait_trace_config
from litellm.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread
from litellm.types.llms.custom_http import *
@ -1456,6 +1457,7 @@ class AsyncHTTPHandler:
connector=TCPConnector(**transport_connector_kwargs),
cookie_jar=DummyCookieJar(),
trust_env=trust_env,
trace_configs=[build_pool_wait_trace_config()],
)
# Use shared session if provided and valid

16
litellm/llms/fireworks_ai/responses/transformation.py
View file

@ -29,6 +29,7 @@ from litellm.utils import (
declared_mid_conversation_system_support,
deployment_declared_mid_conversation_system_support,
deployment_model_info,
get_bundled_model_cost_map,
)
if TYPE_CHECKING:
@ -54,6 +55,15 @@ def _forwarded_user_params(litellm_params: GenericLiteLLMParams) -> Mapping[str,
)
def _bundled_cost_map_declares_mid_conversation_system(model: str) -> bool:
bundled_cost_map: Final[Mapping[str, object]] = get_bundled_model_cost_map()
names: Final = (model.removeprefix("fireworks_ai/"), resolve_fireworks_resource_name(model))
return any(
deployment_declared_mid_conversation_system_support(bundled_cost_map.get(f"fireworks_ai/{name}")) is True
for name in names
)
class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
@ -90,9 +100,13 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
validated_input: Final = self._validate_input_param(input)
model_info: Final = deployment_model_info(litellm_params)
declared: Final = declared_mid_conversation_system_support(model, self.custom_llm_provider.value, model_info)
keeps_developer_items_in_place: Final = declared is True and (
deployment_declared_mid_conversation_system_support(model_info) is True
or _bundled_cost_map_declares_mid_conversation_system(model)
)
folded: Final = (
None
if deployment_declared_mid_conversation_system_support(model_info) is True
if keeps_developer_items_in_place
else instruction_items_folded(
validated_input, response_api_optional_request_params, "user" if declared is False else "system"
)

9
litellm/llms/gemini/common_utils.py
View file

@ -537,13 +537,16 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter):
) -> TokenCountResponse | None:
import copy
from litellm.llms.anthropic.messages.context_management.inference_view import (
InferenceRequestViewError,
count_context_edits,
inference_request_view,
)
from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter
from litellm.llms.vertex_ai.count_tokens.context import (
VertexGeminiCountContextError,
count_context_edits,
count_context_error_response,
gemini_count_contents,
inference_request_view,
malformed_tool_result_text_count_error,
provider_error_response,
unconvertible_messages_error,
@ -556,7 +559,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter):
messages_format,
count_context_edits(context_management, tools, (deployment or {}).get("litellm_params", {}), model_to_use),
)
if isinstance(request_view, VertexGeminiCountContextError):
if isinstance(request_view, InferenceRequestViewError):
return count_context_error_response(request_view, request_model, model_to_use)
count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools, model=model_to_use)
if isinstance(count_context, VertexGeminiCountContextError):

18
litellm/llms/openai/chat/gpt_transformation.py
View file

@ -21,6 +21,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
should_convert_tool_call_to_json_mode,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
chat_completion_prompt_cache_breakpoints_for_wire,
drop_non_python_regex_patterns,
drop_tool_reference_parts_from_tool_messages,
flatten_combinators_and_drop_non_python_regex_patterns,
@ -556,9 +557,10 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
"""
A deployment that reads OpenAI's ``prompt_cache_breakpoint`` (gpt-5.6+ on
api.openai.com, or one opted in through ``prompt_cache_options`` or the cost
map) gets each Anthropic ``cache_control`` mark rewritten in place. Any other
endpoint keeps or strips the marks as before, and under drop_params also loses
every ``prompt_cache_breakpoint`` it cannot read.
map) gets each Anthropic ``cache_control`` mark rewritten in place, and under
drop_params each content part's ``prompt_cache_breakpoint`` cut down to what chat
completions accepts. Any other endpoint keeps or strips the marks as before, and
under drop_params also loses every ``prompt_cache_breakpoint`` it cannot read.
"""
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
@ -572,9 +574,15 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
_claude_messages_within_cache_control_cap(model, messages, tools, litellm_params, request_params),
tools,
)
drops_breakpoints: Final = not reads_breakpoints and _drops_params(litellm_params)
drop_params: Final = _drops_params(litellm_params)
drops_breakpoints: Final = not reads_breakpoints and drop_params
renamed: Final = (
[AnthropicCacheControlHook.cache_control_as_prompt_cache_breakpoints(message) for message in messages]
[
AnthropicCacheControlHook.cache_control_as_prompt_cache_breakpoints(
chat_completion_prompt_cache_breakpoints_for_wire(message, drop_params)
)
for message in messages
]
if reads_breakpoints
else [without_prompt_cache_breakpoints(message) for message in messages]
if drops_breakpoints

6
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -1989,7 +1989,7 @@ class OpenAIResponsesHandler(BaseTranslation):
"annotations": (),
"logprobs": None,
}
return (
events: Final = (
*_open_item_closing_events(responses_so_far),
OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
@ -2034,6 +2034,10 @@ class OpenAIResponsesHandler(BaseTranslation):
response=_blocked_response(exc, response_id=response_id, model=model, output_item=item),
),
)
return tuple(
event.model_copy(update={"sequence_number": number})
for number, event in enumerate(events, start=_next_stream_sequence_number(responses_so_far))
)
class _BlockedContentPart(TypedDict):

64
litellm/llms/openai/responses/transformation.py
View file

@ -1,5 +1,5 @@
from collections.abc import Callable, Collection, Iterable, Mapping, Sequence
from functools import lru_cache
from functools import lru_cache, partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints
@ -12,7 +12,7 @@ import litellm
from litellm._logging import verbose_logger
from litellm.integrations.anthropic_cache_control_hook import (
AnthropicCacheControlHook,
openai_prompt_cache_breakpoint_support,
deployment_prompt_cache_breakpoint_support,
targets_openai_api,
)
from litellm.litellm_core_utils.core_helpers import normalize_drop_params, process_response_headers
@ -25,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
drop_non_python_regex_patterns,
flatten_combinators_and_drop_non_python_regex_patterns,
prompt_cache_breakpoint_for_wire,
whole_prompt_cache_breakpoint_or_none,
with_prompt_cache_breakpoint,
)
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
@ -103,13 +104,6 @@ _INPUT_ITEM: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str,
_CONTENT_PARTS: Final[TypeAdapter[Sequence[Mapping[str, object]]]] = TypeAdapter(Sequence[Mapping[str, object]])
class _DeclaredModel(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True)
base_model: object = None
model_info: object = None
class _DeploymentHost(BaseModel):
"""`responses()` builds the params from raw kwargs, where callers may spell the host `base_url`."""
@ -241,12 +235,12 @@ def _json_array_or_empty(value: object) -> Sequence[object]:
return value if isinstance(value, (list, tuple)) else ()
def _item_with_wire_safe_prompt_cache_breakpoints(item: object, keeps_markers: bool) -> object:
def _item_with_wire_safe_prompt_cache_breakpoints(item: object, to_wire: Callable[[object], object] | None) -> object:
fields: Final = _json_object_or_none(item)
if fields is None or not _carries_prompt_cache_breakpoint(fields):
return item
wire_safe_fields: Final = {
key: _parts_with_wire_safe_prompt_cache_breakpoints(key, value, keeps_markers)
key: _parts_with_wire_safe_prompt_cache_breakpoints(key, value, to_wire)
for key, value in fields.items()
if key != "prompt_cache_breakpoint"
}
@ -270,21 +264,23 @@ def _parts_carry_prompt_cache_breakpoint(parts: object) -> bool:
return any(isinstance(part, Mapping) and "prompt_cache_breakpoint" in part for part in _json_array_or_empty(parts))
def _parts_with_wire_safe_prompt_cache_breakpoints(key: str, value: object, keeps_markers: bool) -> object:
def _parts_with_wire_safe_prompt_cache_breakpoints(
key: str, value: object, to_wire: Callable[[object], object] | None
) -> object:
if key not in _INPUT_PART_LIST_KEYS or not _parts_carry_prompt_cache_breakpoint(value):
return value
parts: Final = _json_array_or_empty(value)
wire_safe_parts: Final = tuple(_part_with_wire_safe_prompt_cache_breakpoint(part, keeps_markers) for part in parts)
wire_safe_parts: Final = tuple(_part_with_wire_safe_prompt_cache_breakpoint(part, to_wire) for part in parts)
return value if _same_objects(wire_safe_parts, parts) else list(wire_safe_parts)
def _part_with_wire_safe_prompt_cache_breakpoint(part: object, keeps_markers: bool) -> object:
def _part_with_wire_safe_prompt_cache_breakpoint(part: object, to_wire: Callable[[object], object] | None) -> object:
fields: Final = _json_object_or_none(part)
if fields is None or "prompt_cache_breakpoint" not in fields:
return part
marker: Final = fields["prompt_cache_breakpoint"]
wire_marker: Final = prompt_cache_breakpoint_for_wire(marker, drop_params=True) if keeps_markers else None
if keeps_markers and wire_marker == marker:
wire_marker: Final = None if to_wire is None else to_wire(marker)
if to_wire is not None and wire_marker == marker:
return part
unmarked: Final = {key: value for key, value in fields.items() if key != "prompt_cache_breakpoint"}
return with_prompt_cache_breakpoint(unmarked, wire_marker)
@ -681,24 +677,31 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
) -> str | ResponseInputParam:
if not isinstance(input, list) or not drops_params(litellm_params):
return input
keeps_markers: Final = reads_breakpoints or self._keeps_prompt_cache_breakpoints(
model=model, litellm_params=litellm_params
)
wire_safe_items: Final = tuple(
_item_with_wire_safe_prompt_cache_breakpoints(item, keeps_markers) for item in input
known_support: Final = self._prompt_cache_breakpoint_support(model=model, litellm_params=litellm_params)
to_wire: Final = (
self._prompt_cache_breakpoint_to_wire(litellm_params, reads_breakpoints, known_support)
if reads_breakpoints or known_support is not False
else None
)
wire_safe_items: Final = tuple(_item_with_wire_safe_prompt_cache_breakpoints(item, to_wire) for item in input)
if _same_objects(wire_safe_items, input):
return input
wire_safe_input: Final = list(wire_safe_items)
return cast("ResponseInputParam", wire_safe_input) # cast-ok: items keep their shape, minus a rejected marker
def _keeps_prompt_cache_breakpoints(self, model: str, litellm_params: GenericLiteLLMParams) -> bool:
provider: Final = _provider_name(self.custom_llm_provider)
base_model: Final = self._declared_base_model(model=model, litellm_params=litellm_params)
known_support: Final = (
openai_prompt_cache_breakpoint_support(name, provider) for name in dict.fromkeys((base_model, model))
def _prompt_cache_breakpoint_to_wire(
self, litellm_params: GenericLiteLLMParams, reads_breakpoints: bool, known_support: bool | None
) -> Callable[[object], object]:
opted_in_openai_host: Final = reads_breakpoints and self.custom_llm_provider == LlmProviders.OPENAI
known_openai_reader: Final = known_support is True and self._targets_openai_api(litellm_params)
if opted_in_openai_host or known_openai_reader:
return partial(prompt_cache_breakpoint_for_wire, drop_params=True)
return whole_prompt_cache_breakpoint_or_none
def _prompt_cache_breakpoint_support(self, model: str, litellm_params: GenericLiteLLMParams) -> bool | None:
return deployment_prompt_cache_breakpoint_support(
model, _provider_name(self.custom_llm_provider), litellm_params
)
return next((support for support in known_support if support is not None), None) is not False
def _targets_openai_api(self, litellm_params: GenericLiteLLMParams) -> bool:
host: Final = _DeploymentHost.model_validate(litellm_params)
@ -872,13 +875,6 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
base_model: Final[object] = model_info.get("base_model") if isinstance(model_info, dict) else None
return base_model if isinstance(base_model, str) and base_model else model
@staticmethod
def _declared_base_model(model: str, litellm_params: GenericLiteLLMParams) -> str:
declared: Final = _DeclaredModel.model_validate(litellm_params)
model_info: Final = _json_object_or_none(declared.model_info)
candidates: Final = (declared.base_model, None if model_info is None else model_info.get("base_model"))
return next((name for name in candidates if isinstance(name, str) and name), model)
@staticmethod
def _sanitized_tool_entry(
entry: Mapping[str, object],

13
litellm/llms/vertex_ai/common_utils.py
View file

@ -1458,6 +1458,7 @@ class VertexAITokenCounter(BaseTokenCounter):
) -> TokenCountResponse | None:
import copy
from litellm.llms.anthropic.common_utils import context_management_with_edits
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
VertexAIError as PartnerVertexAIError,
@ -1539,6 +1540,9 @@ class VertexAITokenCounter(BaseTokenCounter):
system=system,
tools=tools,
count_options=count_options,
context_management=(
None if context_management is None else context_management_with_edits(context_management.spec)
),
)
except (PartnerVertexAIError, httpx.HTTPStatusError) as e:
status_code: Final = e.response.status_code
@ -1565,12 +1569,15 @@ class VertexAITokenCounter(BaseTokenCounter):
original_response=result,
)
else:
from litellm.llms.anthropic.messages.context_management.inference_view import (
InferenceRequestViewError,
count_context_edits,
inference_request_view,
)
from litellm.llms.vertex_ai.count_tokens.context import (
VertexGeminiCountContextError,
count_context_edits,
count_context_error_response,
gemini_count_contents,
inference_request_view,
malformed_tool_result_text_count_error,
provider_error_response,
unconvertible_messages_error,
@ -1584,7 +1591,7 @@ class VertexAITokenCounter(BaseTokenCounter):
messages_format,
count_context_edits(context_management, tools, count_tokens_params_request, model_to_use),
)
if isinstance(request_view, VertexGeminiCountContextError):
if isinstance(request_view, InferenceRequestViewError):
return count_context_error_response(request_view, request_model, model_to_use)
count_context: Final = vertex_gemini_count_context(
system=request_view.system, tools=tools, model=model_to_use

129
litellm/llms/vertex_ai/count_tokens/context.py
View file

@ -11,10 +11,11 @@ from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_
from litellm.types.llms.anthropic import AnthropicMessagesTextParam
from litellm.types.llms.openai import AllMessageValues
from litellm.types.llms.vertex_ai import ContentType
from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
if TYPE_CHECKING:
from litellm.exceptions import APIError
from litellm.llms.anthropic.messages.context_management.inference_view import InferenceRequestViewError
@dataclass(frozen=True, slots=True)
@ -152,7 +153,7 @@ def malformed_tool_result_text_count_error(
def count_context_error_response(
error: VertexGeminiCountContextError, request_model: str, model_used: str
error: "VertexGeminiCountContextError | InferenceRequestViewError", request_model: str, model_used: str
) -> TokenCountResponse:
return TokenCountResponse(
total_tokens=0,
@ -178,130 +179,6 @@ def provider_error_response(error: "APIError", request_model: str, model_used: s
)
@dataclass(frozen=True, slots=True)
class InferenceRequestView:
messages: Sequence[Mapping[str, object]]
system: object | None
@dataclass(frozen=True, slots=True)
class CountContextEdits:
spec: JsonValue
model: str
tools: tuple[Mapping[str, object], ...] | None
drop_invalid_edits: bool = False
_CONTEXT_MANAGEMENT_SPEC: Final[TypeAdapter[Mapping[str, object] | tuple[Mapping[str, object], ...]]] = TypeAdapter(
Mapping[str, object] | tuple[Mapping[str, object], ...]
)
def count_context_edits(
context_management: CountContextManagement | None,
tools: Sequence[Mapping[str, object]] | None,
litellm_params: Mapping[str, object],
fallback_model: str,
) -> CountContextEdits:
model: Final = litellm_params.get("model")
return CountContextEdits(
spec=None if context_management is None else context_management.spec,
model=model if isinstance(model, str) else fallback_model,
tools=None if tools is None else tuple(tools),
drop_invalid_edits=context_management is not None and context_management.drop_invalid_edits,
)
def _typed_spec(spec: JsonValue) -> Mapping[str, object] | tuple[Mapping[str, object], ...] | None:
try:
return _CONTEXT_MANAGEMENT_SPEC.validate_python(spec)
except ValidationError:
return None
def _typed_system(
system: object | None,
) -> str | tuple[AnthropicMessagesTextParam, ...] | None | VertexGeminiCountContextError:
if system is None or isinstance(system, str):
return system
try:
return _SYSTEM_BLOCKS.validate_python(system)
except ValidationError:
return _INVALID_SYSTEM
def _view_without_edits(
messages: Sequence[Mapping[str, object]],
system: object | None,
typed_system: str | Sequence[Mapping[str, object]] | None,
) -> InferenceRequestView:
from litellm.llms.anthropic.messages.context_management.editors.compact import (
apply_client_compaction_block_history,
)
history: Final = apply_client_compaction_block_history(
messages=[dict(message) for message in messages],
system=typed_system
if typed_system is None or isinstance(typed_system, str)
else [dict(block) for block in typed_system],
)
if history is None:
return InferenceRequestView(messages=messages, system=system)
return InferenceRequestView(messages=history.messages, system=history.system)
async def _anthropic_inference_view(
messages: Sequence[Mapping[str, object]], system: object | None, edits: CountContextEdits
) -> InferenceRequestView | VertexGeminiCountContextError:
from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request
from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError
from litellm.llms.anthropic.messages.native.utils import normalize_native_messages
typed_system: Final = _typed_system(system)
if isinstance(typed_system, VertexGeminiCountContextError):
return typed_system
normalized: Final = normalize_native_messages(messages)
spec: Final = _typed_spec(edits.spec)
try:
prepared: Final = await prepare_context_managed_request(
model=edits.model,
messages=normalized,
tools=None if edits.tools is None else [dict(tool) for tool in edits.tools],
system=[dict(block) for block in typed_system] if isinstance(typed_system, tuple) else typed_system,
context_management_spec=(
None if spec is None else dict(spec) if isinstance(spec, Mapping) else [dict(edit) for edit in spec]
),
litellm_metadata=None,
additional_drop_params=None,
llm_router=None,
count_only=True,
drop_invalid_edits=edits.drop_invalid_edits,
)
except AnthropicContextManagementError as error:
if error.status_code >= 500:
return _view_without_edits(normalized, system, typed_system)
return VertexGeminiCountContextError(error_message=error.message, status_code=error.status_code)
if prepared is None:
return InferenceRequestView(messages=normalized, system=system)
return InferenceRequestView(messages=prepared.messages, system=prepared.system)
async def inference_request_view(
messages: Sequence[Mapping[str, object]],
system: object | None,
messages_format: TokenCountMessagesFormat,
edits: CountContextEdits,
) -> InferenceRequestView | VertexGeminiCountContextError:
"""Apply the history rewrites and context edits /v1/messages runs before it converts a body for Gemini."""
match messages_format:
case "openai":
return InferenceRequestView(messages=messages, system=system)
case "anthropic":
return await _anthropic_inference_view(messages, system, edits)
case _:
assert_never(messages_format)
def _chat_messages(
messages: Sequence[Mapping[str, object]],
messages_format: TokenCountMessagesFormat,

46
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -437,6 +437,8 @@ _NAMED_UNKNOWN_VOICE_KEYS: Final = 10
_NAMED_VOICE_KEY_CHARS: Final = 100
_NAMED_AUDIO_FORMAT_CHARS: Final = 100
def _gemini_speech_config_wrapper(value: Mapping[str, object]) -> Mapping[str, object] | None:
camel_config: Final = as_str_mapping(value.get("speechConfig"))
@ -473,6 +475,31 @@ def reject_unknown_gemini_tts_voice_keys(
)
def _named_audio_format(audio_format: object) -> str:
if isinstance(audio_format, str):
return repr(audio_format[:_NAMED_AUDIO_FORMAT_CHARS])
return repr(audio_format)[:_NAMED_AUDIO_FORMAT_CHARS]
def unsupported_gemini_tts_audio_format(
model: str, audio: Mapping[str, object], drop_params: bool, llm_provider: str | None
) -> UnsupportedParamsError | None:
nested_config: Final = _gemini_speech_config_wrapper(audio) or {}
audio_format: Final = next(
(item for item in (audio.get("format"), nested_config.get("format")) if item not in (None, "pcm16")), None
)
if audio_format is None or drops_unsupported_params(drop_params):
return None
return UnsupportedParamsError(
message=(
f"{model} does not take the audio format {_named_audio_format(audio_format)}. Set format "
"to 'pcm16' or leave it out, or set drop_params to drop it."
),
model=model,
llm_provider=llm_provider,
)
def normalize_gemini_speech_config(
value: Mapping[str, object], drop_unknown_fields: bool = False
) -> dict[str, object]: # mutable-ok: provider request serialization requires a concrete dict
@ -1555,15 +1582,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
languageCode: "en-US",
}
"""
# Validate audio format - Gemini TTS only supports pcm16
audio_format: Final = value.get("format")
if audio_format is not None and audio_format != "pcm16":
raise ValueError(
f"Unsupported audio format for Gemini TTS models: {audio_format}. "
f"Gemini TTS models only support 'pcm16' format as they return audio data in L16 PCM format. "
f"Please set audio format to 'pcm16'."
)
audio_format_error: Final = unsupported_gemini_tts_audio_format(
"Gemini TTS", value, drop_unknown_fields, self.custom_llm_provider
)
if audio_format_error is not None:
raise audio_format_error
camel_speech_config: Final = value.get("speechConfig")
snake_speech_config: Final = value.get("speech_config")
speech_config_value: Final = (
@ -1601,6 +1624,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
def _map_chat_audio_params(
self, model: str, value: Mapping[str, object], drop_params: bool
) -> dict[str, object]: # mutable-ok: the public forwarding interface returns a concrete request dictionary
audio_format_error: Final = unsupported_gemini_tts_audio_format(
model, value, drop_params, self.custom_llm_provider
)
if audio_format_error is not None:
raise audio_format_error
if _gemini_speech_config_wrapper(value) is not None:
reject_unknown_gemini_tts_voice_keys(model, value, drop_params, self.custom_llm_provider)
return self._map_audio_params(value, drop_unknown_fields=drops_unsupported_params(drop_params))

17
litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py
View file

@ -9,12 +9,14 @@ their respective publisher-specific count-tokens endpoints.
from collections.abc import Mapping, Sequence
from typing import Any, Final
from pydantic import TypeAdapter
from pydantic import JsonValue, TypeAdapter
from litellm.anthropic_beta_headers_manager import aload_beta_headers_config, filter_and_transform_beta_headers
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.anthropic.common_utils import (
OBJECT_LIST_ADAPTER,
context_management_betas,
count_tokens_feature_betas,
is_claude_model_target,
tools_without_custom_field,
@ -31,6 +33,7 @@ _COUNT_TOKENS_PUBLISHERS: Final = (
_REQUEST_MESSAGES: Final = TypeAdapter(Sequence[object])
_CONTEXT_MANAGEMENT: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
def count_tokens_request_body(
@ -55,8 +58,14 @@ def count_tokens_publisher(model: str) -> str | None:
def count_tokens_headers(access_token: str, request_data: Mapping[str, object]) -> Mapping[str, str]:
authorization: Final = {"Authorization": f"Bearer {access_token}"}
betas: Final = count_tokens_feature_betas(
"vertex_ai", request_data, _REQUEST_MESSAGES.validate_python(request_data["messages"])
betas: Final = (
*count_tokens_feature_betas(
"vertex_ai", request_data, _REQUEST_MESSAGES.validate_python(request_data["messages"])
),
*filter_and_transform_beta_headers(
list(context_management_betas(_CONTEXT_MANAGEMENT.validate_python(request_data.get("context_management")))),
"vertex_ai",
),
)
return {**authorization, "anthropic-beta": ",".join(betas)} if betas else authorization
@ -193,6 +202,8 @@ class VertexAIPartnerModelsTokenCounter(VertexBase):
api_base=litellm_params.get("api_base"),
)
if request_data.get("context_management") is not None:
await aload_beta_headers_config()
headers: Final = dict(count_tokens_headers(access_token, request_data))
# Get async HTTP client

11
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py
View file

@ -268,6 +268,7 @@ class VertexAIPartnerModels(VertexBase):
system: object | None = None,
tools: list[dict[str, object]] | None = None,
count_options: Mapping[str, object] = MappingProxyType({}),
context_management: object = None,
):
"""
Count tokens for Vertex AI partner models (Anthropic Claude, Mistral, etc.)
@ -302,7 +303,15 @@ class VertexAIPartnerModels(VertexBase):
"model": model,
"messages": messages,
**MappingProxyType(
{key: value for key, value in (("system", system), ("tools", tools)) if value is not None}
{
key: value
for key, value in (
("system", system),
("tools", tools),
("context_management", context_management),
)
if value is not None
}
),
**count_options,
}

42
litellm/messages/dispatch.py
View file

@ -75,9 +75,7 @@ def _public_request(
return native_call(legacy, args, _native_kwargs(fields, kwargs))
def _resolved(request: NativeCall) -> tuple[str, str | None]:
model: Final = str(request.resolved["model"])
provider: Final = optional_str(request.resolved.get("custom_llm_provider"))
def _resolved_model(model: str, provider: str | None) -> tuple[str, str | None]:
try:
resolved_model, resolved_provider, _, _ = get_llm_provider(model, provider)
except BadRequestError:
@ -85,21 +83,43 @@ def _resolved(request: NativeCall) -> tuple[str, str | None]:
return resolved_model, resolved_provider
def _target(request: NativeCall, model: str, provider: str | None) -> Target:
requested_model: Final = str(request.resolved["model"])
model_info: Final = request.resolved.get("model_info")
api_base: Final = request.resolved.get("api_base")
native_config: Final = main._native_messages_config( # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup
model, provider, model_info, api_base if isinstance(api_base, str) else None
)
def _resolved(request: NativeCall) -> tuple[str, str | None]:
return _resolved_model(str(request.resolved["model"]), optional_str(request.resolved.get("custom_llm_provider")))
def _route_target(
requested_model: str,
model: str,
provider: str | None,
model_info: object,
api_base: str | None,
compacting: bool,
) -> Target:
native_config: Final = main._native_messages_config(model, provider, model_info, api_base) # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup
if native_config is not None:
return Target.MESSAGES
compacting: Final = request.resolved.get("compaction") is not None
if not compacting and main._should_route_to_responses_api(provider, requested_model, model, model_info): # pyright: ignore[reportPrivateUsage] # the handler's own translation choice
return Target.RESPONSES
return Target.CHAT_COMPLETIONS
def _target(request: NativeCall, model: str, provider: str | None) -> Target:
api_base: Final = request.resolved.get("api_base")
return _route_target(
str(request.resolved["model"]),
model,
provider,
request.resolved.get("model_info"),
api_base if isinstance(api_base, str) else None,
compacting=request.resolved.get("compaction") is not None,
)
def messages_target(model: str, custom_llm_provider: str | None, model_info: object, api_base: str | None) -> Target:
resolved_model, resolved_provider = _resolved_model(model, custom_llm_provider)
return _route_target(model, resolved_model, resolved_provider, model_info, api_base, compacting=False)
_FEATURES: Final[detectors.Detectors] = MappingProxyType(
{
**detectors.COMMON,

82
litellm/model_prices_and_context_window_backup.json
View file

@ -15555,18 +15555,20 @@
"input_cost_per_token": 2.65e-06,
"litellm_provider": "bedrock",
"max_input_tokens": 8000,
"max_output_tokens": 2048,
"max_output_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 3.5e-06,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-70b-instruct.md",
"supports_pdf_input": true
},
"bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0": {
"input_cost_per_token": 3e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 8000,
"max_output_tokens": 2048,
"max_output_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-8b-instruct.md",
"supports_pdf_input": true
},
"bedrock/us-west-1/meta.llama3-70b-instruct-v1:0": {
@ -48955,6 +48957,7 @@
},
"together_ai/Qwen/Qwen3.8-2.4T-A95B": {
"cache_read_input_token_cost": 2.5e-07,
"deprecation_date": "2026-11-04",
"input_cost_per_token": 2e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 1010000,
@ -53954,23 +53957,23 @@
"supports_vision": true
},
"wandb/meta-llama/Llama-3.1-8B-Instruct": {
"max_input_tokens": 131000,
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"input_cost_per_token": 2.2e-07,
"output_cost_per_token": 2.2e-07,
"litellm_provider": "wandb",
"mode": "chat",
"source": "https://wandb.ai/site/pricing/tokens/"
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference"
},
"wandb/deepseek-ai/DeepSeek-V3.1": {
"supports_reasoning": true,
"max_input_tokens": 161000,
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"input_cost_per_token": 5.5e-07,
"output_cost_per_token": 1.65e-06,
"litellm_provider": "wandb",
"mode": "chat",
"source": "https://wandb.ai/site/pricing/tokens/"
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference"
},
"wandb/meta-llama/Llama-3.3-70B-Instruct": {
"max_input_tokens": 128000,
@ -59872,6 +59875,7 @@
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"mode": "responses",
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-cyber.html",
"use_openai_responses_path": true,
"supported_endpoints": [
@ -59893,7 +59897,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_xhigh_reasoning_effort": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol": {
"input_cost_per_token": 4.4e-06,
@ -59908,6 +59913,7 @@
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"mode": "responses",
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/responses"
@ -59927,7 +59933,8 @@
"supports_tool_choice": true,
"supports_xhigh_reasoning_effort": true,
"supports_vision": true,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html"
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html",
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.6-luna": {
"input_cost_per_token": 2.2e-07,
@ -60333,6 +60340,7 @@
"mode": "responses",
"output_cost_per_token_above_272k_tokens_ultrafast": 0.000495,
"output_cost_per_token_ultrafast": 0.00033,
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
@ -60356,7 +60364,8 @@
"supports_tool_choice": true,
"supports_xhigh_reasoning_effort": true,
"supports_vision": true,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html"
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html",
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-6-sol": {
"input_cost_per_token": 2.2e-06,
@ -60371,6 +60380,7 @@
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"mode": "responses",
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
@ -60394,7 +60404,8 @@
"supports_tool_choice": true,
"supports_xhigh_reasoning_effort": true,
"supports_vision": true,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html"
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html",
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-6-luna": {
"input_cost_per_token": 1.1e-07,
@ -60409,6 +60420,7 @@
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"mode": "responses",
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
@ -60432,7 +60444,8 @@
"supports_tool_choice": true,
"supports_xhigh_reasoning_effort": true,
"supports_vision": true,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html"
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html",
"supports_web_search": true
},
"us.openai.gpt-6-astra": {
"supports_bedrock_runtime_chat_completions_response_format": true,
@ -64026,6 +64039,7 @@
"output_cost_per_token_priority": 1.5e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -64045,6 +64059,7 @@
"output_cost_per_token_priority": 2.25e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -64087,6 +64102,7 @@
],
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64124,6 +64140,7 @@
"output_cost_per_token_priority": 1.5e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -64143,6 +64160,7 @@
"output_cost_per_token_priority": 2.25e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -64217,6 +64235,7 @@
],
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64237,6 +64256,7 @@
],
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64260,6 +64280,7 @@
],
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64335,6 +64356,7 @@
"output_cost_per_token_priority": 3e-06,
"source": "https://api.fireworks.ai/v1/serverless/models?format=nested",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64371,6 +64393,7 @@
"output_cost_per_token_priority": 6.375e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_vision": true
},
@ -64407,6 +64430,7 @@
"output_cost_per_token_priority": 3e-06,
"source": "https://api.fireworks.ai/v1/serverless/models?format=nested",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64477,6 +64501,7 @@
],
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64500,6 +64525,7 @@
],
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67293,6 +67319,7 @@
"output_cost_per_token_priority": 5.5e-06,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67311,6 +67338,7 @@
"output_cost_per_token_priority": 8.25e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67329,6 +67357,7 @@
"output_cost_per_token_priority": 5.5e-06,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67347,6 +67376,7 @@
"output_cost_per_token_priority": 8.25e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67362,6 +67392,7 @@
"output_cost_per_token": 6.6e-06,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67404,6 +67435,7 @@
"output_cost_per_token": 6.6e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67422,6 +67454,7 @@
"output_cost_per_token_priority": 6.25e-07,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
@ -67440,6 +67473,7 @@
"output_cost_per_token_priority": 9.375e-07,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
@ -67457,6 +67491,7 @@
"output_cost_per_token_priority": 6.25e-07,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
@ -67475,6 +67510,7 @@
"output_cost_per_token_priority": 9.375e-07,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
@ -67493,6 +67529,7 @@
"output_cost_per_token_priority": 4.05e-06,
"source": "https://api.fireworks.ai/v1/serverless/models?format=nested",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
@ -68630,7 +68667,9 @@
"cache_read_input_token_cost_above_272k_tokens": 5.4e-07,
"output_cost_per_token": 1.62e-05,
"output_cost_per_token_above_272k_tokens": 2.43e-05,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html"
"search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144},
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html",
"supports_web_search": true
},
"bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": {
"litellm_provider": "bedrock_mantle",
@ -68665,7 +68704,9 @@
"cache_read_input_token_cost_above_272k_tokens": 5.4e-08,
"output_cost_per_token": 1.62e-06,
"output_cost_per_token_above_272k_tokens": 2.43e-06,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html"
"search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144},
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html",
"supports_web_search": true
},
"bedrock_mantle/us-gov-west-1/openai.gpt-5.4": {
"litellm_provider": "bedrock_mantle",
@ -68694,7 +68735,9 @@
"input_cost_per_token": 3.375e-06,
"cache_read_input_token_cost": 3.375e-07,
"output_cost_per_token": 2.025e-05,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html"
"search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144},
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html",
"supports_web_search": true
},
"bedrock_mantle/us-gov-west-1/xai.grok-4.3": {
"use_openai_responses_path": true,
@ -73728,6 +73771,7 @@
"output_cost_per_token_priority": 1.875e-05,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -75309,6 +75353,7 @@
"output_cost_per_token_above_272k_tokens": 1.65e-05,
"output_cost_per_token_above_272k_tokens_ultrafast": 9.9e-05,
"output_cost_per_token_ultrafast": 6.6e-05,
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html",
"supported_endpoints": [
"/v1/chat/completions",
@ -75331,6 +75376,7 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true,
"supports_xhigh_reasoning_effort": true,
"use_openai_responses_path": true
},
@ -77228,6 +77274,7 @@
"output_cost_per_token_priority": 1.5e-06,
"source": "https://api.fireworks.ai/v1/serverless/models?format=nested",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -77386,6 +77433,7 @@
"output_cost_per_token_priority": 9e-07,
"source": "https://platform.minimax.io/docs/guides/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -79019,6 +79067,7 @@
"mode": "responses",
"output_cost_per_token": 5.5e-06,
"output_cost_per_token_above_272k_tokens": 8.25e-06,
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-terra.html",
"supported_endpoints": [
"/v1/chat/completions",
@ -79041,6 +79090,7 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true,
"supports_xhigh_reasoning_effort": true,
"use_openai_responses_path": true
},
@ -81161,6 +81211,7 @@
],
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -81177,6 +81228,7 @@
"output_cost_per_token": 3.08e-06,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -81193,6 +81245,7 @@
"output_cost_per_token": 8.4e-07,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -81209,6 +81262,7 @@
"output_cost_per_token": 8.4e-07,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,

4
litellm/proxy/_lazy_features.py
View file

@ -105,10 +105,10 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
LazyFeature(
name="a2a",
module_path="litellm.proxy.agent_endpoints.a2a_endpoints",
# ``/v1/a2a/{agent_id}/message/send`` is caught via the suffix so the
# ``/v1/a2a/{agent_id}/message/send`` and ``/message/stream`` are caught via suffixes so the
# ``/v1/a2a`` prefix doesn't subsume the discover prefix below.
path_prefixes=("/a2a",),
path_suffixes=("/message/send",),
path_suffixes=("/message/send", "/message/stream"),
),
LazyFeature(
name="a2a_registration",

359
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -54,6 +54,57 @@
],
"title": "ValidationError",
"type": "object"
},
"_A2AJsonRpcResponse": {
"additionalProperties": true,
"properties": {
"error": {
"anyOf": [
{
"additionalProperties": true,
"type": "object"
},
{
"type": "null"
}
],
"title": "Error"
},
"id": {
"anyOf": [
{
"type": "string"
},
{
"type": "integer"
},
{
"type": "null"
}
],
"title": "Id"
},
"jsonrpc": {
"const": "2.0",
"title": "Jsonrpc",
"type": "string"
},
"result": {
"anyOf": [
{},
{
"type": "null"
}
],
"title": "Result"
}
},
"required": [
"jsonrpc",
"id"
],
"title": "_A2AJsonRpcResponse",
"type": "object"
}
}
},
@ -242,6 +293,119 @@
]
}
},
"/a2a/{agent_id}/message/stream": {
"post": {
"description": "Invoke an agent using the A2A protocol (JSON-RPC 2.0).\n\nSupported methods:\n- message/send: Send a message and get a response\n- message/stream: Send a message and stream the response",
"operationId": "invoke_agent_a2a_a2a__agent_id__message_stream_post",
"parameters": [
{
"in": "path",
"name": "agent_id",
"required": true,
"schema": {
"title": "Agent Id",
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
"example": {
"id": "stream-1",
"jsonrpc": "2.0",
"method": "message/stream",
"params": {
"message": {
"kind": "message",
"messageId": "message-1",
"parts": [
{
"kind": "text",
"text": "Hello"
}
],
"role": "user"
}
}
},
"schema": {
"additionalProperties": true,
"properties": {
"id": {
"anyOf": [
{
"type": "string"
},
{
"type": "integer"
},
{
"type": "null"
}
],
"title": "Id"
},
"jsonrpc": {
"const": "2.0",
"title": "Jsonrpc",
"type": "string"
},
"method": {
"title": "Method",
"type": "string"
},
"params": {
"additionalProperties": true,
"title": "Params",
"type": "object"
}
},
"required": [
"jsonrpc",
"method",
"params"
],
"title": "_A2AJsonRpcRequest",
"type": "object"
}
}
},
"required": true
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/_A2AJsonRpcResponse"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Invoke Agent A2A",
"tags": [
"a2a"
]
}
},
"/v1/a2a/{agent_id}/message/send": {
"post": {
"description": "Invoke an agent using the A2A protocol (JSON-RPC 2.0).\n\nSupported methods:\n- message/send: Send a message and get a response\n- message/stream: Send a message and stream the response",
@ -287,6 +451,119 @@
"a2a"
]
}
},
"/v1/a2a/{agent_id}/message/stream": {
"post": {
"description": "Invoke an agent using the A2A protocol (JSON-RPC 2.0).\n\nSupported methods:\n- message/send: Send a message and get a response\n- message/stream: Send a message and stream the response",
"operationId": "invoke_agent_a2a_v1_a2a__agent_id__message_stream_post",
"parameters": [
{
"in": "path",
"name": "agent_id",
"required": true,
"schema": {
"title": "Agent Id",
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
"example": {
"id": "stream-1",
"jsonrpc": "2.0",
"method": "message/stream",
"params": {
"message": {
"kind": "message",
"messageId": "message-1",
"parts": [
{
"kind": "text",
"text": "Hello"
}
],
"role": "user"
}
}
},
"schema": {
"additionalProperties": true,
"properties": {
"id": {
"anyOf": [
{
"type": "string"
},
{
"type": "integer"
},
{
"type": "null"
}
],
"title": "Id"
},
"jsonrpc": {
"const": "2.0",
"title": "Jsonrpc",
"type": "string"
},
"method": {
"title": "Method",
"type": "string"
},
"params": {
"additionalProperties": true,
"title": "Params",
"type": "object"
}
},
"required": [
"jsonrpc",
"method",
"params"
],
"title": "_A2AJsonRpcRequest",
"type": "object"
}
}
},
"required": true
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/_A2AJsonRpcResponse"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Invoke Agent A2A",
"tags": [
"a2a"
]
}
}
}
},
@ -15326,6 +15603,10 @@
],
"title": "Metadata"
},
"response_format": {
"$ref": "#/components/schemas/GuardrailResponseFormat",
"default": "standard"
},
"text": {
"title": "Text",
"type": "string"
@ -16203,6 +16484,14 @@
"title": "BedrockChecksSensitiveInformationModel",
"type": "object"
},
"BedrockGuardrailOutputScope": {
"enum": [
"INTERVENTIONS",
"FULL"
],
"title": "BedrockGuardrailOutputScope",
"type": "string"
},
"BlockedWord": {
"description": "Represents a blocked word with its action and optional description",
"properties": {
@ -16746,6 +17035,14 @@
"title": "GuardrailInfoResponse",
"type": "object"
},
"GuardrailResponseFormat": {
"enum": [
"standard",
"native"
],
"title": "GuardrailResponseFormat",
"type": "string"
},
"GuardrailSubmissionItem": {
"properties": {
"created_at": {
@ -16998,6 +17295,7 @@
"title": "HTTPValidationError",
"type": "object"
},
"JsonValue": {},
"LakeraCategoryThresholds": {
"additionalProperties": true,
"properties": {
@ -17924,6 +18222,18 @@
"description": "Whether to fail the request if the guardrail encounters an error. Implemented by guardrail='model_armor', 'generic_guardrail_api' and 'crowdstrike_aidr'. True (default) raises the error. False logs a critical error and lets the request proceed, so only a valid guardrail response can block or modify it.",
"title": "Fail On Error"
},
"file_checkpoint_id": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "File-checkpoint ID for the Ovalix Tracker service (falls back to pre/post).",
"title": "File Checkpoint Id"
},
"file_guardrail_timeout": {
"anyOf": [
{
@ -18426,6 +18736,17 @@
],
"description": "Optional parameters for the guardrail"
},
"outputScope": {
"anyOf": [
{
"$ref": "#/components/schemas/BedrockGuardrailOutputScope"
},
{
"type": "null"
}
],
"description": "ApplyGuardrail response scope. FULL includes detected and non-detected findings. Omit to retain Bedrock's default of returning detected findings only."
},
"output_parse_pii": {
"anyOf": [
{
@ -20308,7 +20629,24 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/ApplyGuardrailResponse"
"anyOf": [
{
"$ref": "#/components/schemas/ApplyGuardrailResponse"
},
{
"additionalProperties": {
"$ref": "#/components/schemas/JsonValue"
},
"type": "object"
},
{
"items": {
"$ref": "#/components/schemas/JsonValue"
},
"type": "array"
}
],
"title": "Response Apply Guardrail Apply Guardrail Post"
}
}
},
@ -20400,7 +20738,24 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/ApplyGuardrailResponse"
"anyOf": [
{
"$ref": "#/components/schemas/ApplyGuardrailResponse"
},
{
"additionalProperties": {
"$ref": "#/components/schemas/JsonValue"
},
"type": "object"
},
{
"items": {
"$ref": "#/components/schemas/JsonValue"
},
"type": "array"
}
],
"title": "Response Apply Guardrail Guardrails Apply Guardrail Post"
}
}
},

1
litellm/proxy/_types.py
View file

@ -501,6 +501,7 @@ class LiteLLMRoutes(enum.Enum):
"/a2a/{agent_id}/message/send",
"/v1/a2a/{agent_id}/message/send",
"/a2a/{agent_id}/message/stream",
"/v1/a2a/{agent_id}/message/stream",
"/a2a/{agent_id}/.well-known/agent-card.json",
)

21
litellm/proxy/a2a/agent_card.py
View file

@ -12,6 +12,8 @@ from collections.abc import Mapping
from copy import deepcopy
from typing import Final, Literal
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items
SupportedA2AVersion = Literal["0.3", "1.0"]
# Protocol versions LiteLLM can serve to A2A clients. The admin pins one per agent;
@ -21,7 +23,6 @@ SUPPORTED_A2A_PROTOCOL_VERSIONS: Final[tuple[SupportedA2AVersion, ...]] = ("0.3"
# Default served version when the agent card does not pin one.
LITELLM_A2A_PROTOCOL_VERSION: Final = "1.0"
_PROTOCOL_VERSION_PATTERN: Final = re.compile(
r"^(\d+\.\d+)(?:\.\d+(?:-[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?(?:\+[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?)?$"
)
@ -50,6 +51,24 @@ def resolve_served_protocol_version(card: Mapping[str, object] | None) -> str:
return normalized if normalized is not None else LITELLM_A2A_PROTOCOL_VERSION
def point_agent_card_at_proxy(
agent_card: Mapping[str, object],
proxy_url: str,
) -> dict[str, object]: # mutable-ok: the public helper contract requires a dictionary
interfaces: Final = agent_card.get("supportedInterfaces")
interface_items: Final = (
tuple(iter_items(agent_card.get("supportedInterfaces"))) if isinstance(interfaces, list) else ()
)
first: Final[object] = interface_items[0] if interface_items else None
first_mapping: Final = as_str_mapping(first)
pointed_interfaces: Final = (
{"supportedInterfaces": [{**first_mapping, "url": proxy_url}, *interface_items[1:]]}
if isinstance(first, dict) and first_mapping is not None
else {}
)
return deepcopy({**agent_card, "url": proxy_url, **pointed_interfaces})
# Security scheme exposed by the LiteLLM-fronted agent card. Always replaces
# whatever upstream advertised — the client must authenticate to the proxy,
# not the upstream agent.

65
litellm/proxy/agent_endpoints/a2a_endpoints.py
View file

@ -12,9 +12,8 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM
import json
from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping
from copy import deepcopy
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
from urllib.parse import urlparse
from fastapi import APIRouter, Depends, HTTPException, Request, Response
@ -29,6 +28,7 @@ from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.a2a.agent_card import point_agent_card_at_proxy
from litellm.proxy.a2a.version_convert import (
A2AVersion,
normalize_agent_card,
@ -43,6 +43,7 @@ from litellm.proxy.agent_endpoints.databricks_oauth import (
)
from litellm.proxy.agent_endpoints.utils import merge_agent_headers
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.custom_openapi_spec import inline_request_body
from litellm.proxy.common_utils.openai_error_payload import status_typed_proxy_exception
from litellm.proxy.common_utils.sse_keepalive import (
SSE_COMMENT_PING,
@ -59,6 +60,42 @@ if TYPE_CHECKING:
router: Final = APIRouter()
class _A2AJsonRpcRequest(BaseModel):
model_config = ConfigDict(extra="allow", frozen=True)
jsonrpc: Literal["2.0"]
id: str | int | None = None
method: str
params: Mapping[str, object]
class _A2AJsonRpcResponse(BaseModel):
model_config = ConfigDict(extra="allow", frozen=True)
jsonrpc: Literal["2.0"]
id: str | int | None
result: object | None = None
error: Mapping[str, object] | None = None
_A2A_MESSAGE_STREAM_REQUEST_BODY: Final = inline_request_body(
_A2AJsonRpcRequest,
{
"jsonrpc": "2.0",
"id": "stream-1",
"method": "message/stream",
"params": {
"message": {
"messageId": "message-1",
"role": "user",
"parts": [{"kind": "text", "text": "Hello"}],
"kind": "message",
}
},
},
)
# Mirrors the native seam's own headers: a reverse proxy that batches the whole
# stream would swallow the keepalives this route sends to defeat idle timeouts.
_SSE_KEEPALIVE_HEADERS: Final[Mapping[str, str]] = MappingProxyType(
@ -713,16 +750,12 @@ async def get_agent_card(
)
proxy_url: Final = get_custom_url(str(request.base_url), route=f"a2a/{agent_id}")
agent_card = deepcopy(agent.agent_card_params)
agent_card["url"] = proxy_url
interfaces: Final = agent_card.get("supportedInterfaces")
if isinstance(interfaces, list) and interfaces:
interfaces[0]["url"] = proxy_url
agent_card: Final = point_agent_card_at_proxy(agent.agent_card_params, proxy_url)
served_version: Final = _served_version(agent, request)
agent_card = normalize_agent_card(agent_card, served_version)
normalized_agent_card: Final = normalize_agent_card(agent_card, served_version)
verbose_proxy_logger.debug("Returning agent card for '%s' with proxy URL: %s", agent_id, proxy_url)
return JSONResponse(content=agent_card)
return JSONResponse(content=normalized_agent_card)
except HTTPException:
raise
@ -746,6 +779,20 @@ async def get_agent_card(
tags=["[beta] A2A Agents"],
dependencies=[Depends(user_api_key_auth)],
)
@router.post(
"/a2a/{agent_id}/message/stream",
tags=["[beta] A2A Agents"],
dependencies=[Depends(user_api_key_auth)],
openapi_extra={"requestBody": _A2A_MESSAGE_STREAM_REQUEST_BODY},
responses={200: {"model": _A2AJsonRpcResponse}},
)
@router.post(
"/v1/a2a/{agent_id}/message/stream",
tags=["[beta] A2A Agents"],
dependencies=[Depends(user_api_key_auth)],
openapi_extra={"requestBody": _A2A_MESSAGE_STREAM_REQUEST_BODY},
responses={200: {"model": _A2AJsonRpcResponse}},
)
async def invoke_agent_a2a(
agent_id: str,
request: Request,

78
litellm/proxy/agent_endpoints/endpoints.py
View file

@ -22,6 +22,7 @@ from typing_extensions import ReadOnly, Required, assert_never
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS, AGENT_HEALTH_CHECK_TIMEOUT_SECONDS
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._types import (
CommonProxyErrors,
@ -33,6 +34,7 @@ from litellm.proxy.a2a.agent_card import (
SUPPORTED_A2A_PROTOCOL_VERSIONS,
merge_agent_card,
normalize_protocol_version,
point_agent_card_at_proxy,
)
from litellm.proxy.agent_endpoints.agent_registry import (
AgentIdWhere,
@ -180,10 +182,36 @@ def _redact_agent_litellm_params_dict(
return dict(parse_agent_litellm_params(redact_sensitive_agent_litellm_params(litellm_params)))
def _non_admin_agent_card(
card: Mapping[str, object], proxy_url: str
) -> dict[str, object]: # mutable-ok: AgentResponse.agent_card_params requires a mutable dictionary
pointed_card: Final = point_agent_card_at_proxy(card, proxy_url)
interfaces: Final = pointed_card.get("supportedInterfaces")
interface_items: Final = (
tuple(iter_items(pointed_card.get("supportedInterfaces"))) if isinstance(interfaces, list) else ()
)
first_interface: Final[object] = interface_items[0] if interface_items else None
first_mapping: Final = as_str_mapping(first_interface)
proxy_interface: Final = (
{"supportedInterfaces": [{**first_mapping, "url": proxy_url}]}
if isinstance(first_interface, dict) and first_mapping is not None and first_mapping.get("url") == proxy_url
else {}
)
return {
**{
key: value
for key, value in pointed_card.items()
if key not in ("additionalInterfaces", "supportedInterfaces")
},
**proxy_interface,
}
def _redact_sensitive_agent_fields(
agents: Sequence[AgentResponse],
*,
is_admin: bool,
request_base_url: str,
) -> list[AgentResponse]:
"""
Return copies of the given agents with credential-bearing litellm_params
@ -207,6 +235,11 @@ def _redact_sensitive_agent_fields(
copy.keys = None
copy.kill_switch = None
copy.identity = None
if copy.agent_card_params:
copy.agent_card_params = _non_admin_agent_card(
copy.agent_card_params,
get_custom_url(request_base_url, route=f"a2a/{copy.agent_id}"),
)
if copy.litellm_params:
copy.litellm_params = _redact_agent_litellm_params_dict(copy.litellm_params)
copy.kill_switch = redact_kill_switch(copy.kill_switch)
@ -461,14 +494,6 @@ async def get_agents(
global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups)
)
# litellm_params secrets are always redacted; keys/headers stay
# admin-only.
is_admin: Final = (
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
)
returned_agents = _redact_sensitive_agent_fields(returned_agents, is_admin=is_admin)
if health_check:
agents_with_url: Final = [agent for agent in returned_agents if (agent.agent_card_params or {}).get("url")]
agents_without_url = [agent for agent in returned_agents if not (agent.agent_card_params or {}).get("url")]
@ -493,6 +518,18 @@ async def get_agents(
healthy_ids: Final = {result["agent_id"] for result in health_results if result["healthy"]}
returned_agents = [agent for agent in agents_with_url if agent.agent_id in healthy_ids] + agents_without_url
# litellm_params secrets are always redacted; keys/headers stay
# admin-only.
is_admin: Final = (
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
)
returned_agents = _redact_sensitive_agent_fields(
returned_agents,
is_admin=is_admin,
request_base_url=str(request.base_url),
)
if query is None:
return returned_agents
return await _rank_agents_by_query(query, returned_agents, top_k, user_api_key_dict)
@ -691,7 +728,11 @@ async def create_agent(
# The caller is a proxy admin (enforced above); litellm_params
# secrets are still never echoed back in the response.
return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
return _redact_sensitive_agent_fields(
(result,),
is_admin=True,
request_base_url=str(http_request.base_url),
)[0]
except HTTPException:
raise
@ -708,6 +749,7 @@ async def create_agent(
)
async def get_agent_by_id(
agent_id: str,
request: Request,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
@ -770,7 +812,11 @@ async def get_agent_by_id(
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
)
agent = _redact_sensitive_agent_fields((agent,), is_admin=is_admin)[0]
agent = _redact_sensitive_agent_fields(
(agent,),
is_admin=is_admin,
request_base_url=str(request.base_url),
)[0]
return agent
except HTTPException:
@ -879,7 +925,11 @@ async def update_agent(
"Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id
)
return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
return _redact_sensitive_agent_fields(
(result,),
is_admin=True,
request_base_url=str(http_request.base_url),
)[0]
except HTTPException:
raise
except Exception as e:
@ -987,7 +1037,11 @@ async def patch_agent(
"Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id
)
return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
return _redact_sensitive_agent_fields(
(result,),
is_admin=True,
request_base_url=str(http_request.base_url),
)[0]
except HTTPException:
raise
except Exception as e:

73
litellm/proxy/auth/auth_checks.py
View file

@ -5378,79 +5378,6 @@ async def _granted_model_lists(
)
async def _user_object_or_none(
valid_token: UserAPIKeyAuth,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> LiteLLM_UserTable | None:
try:
return await get_user_object(
user_id=valid_token.user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
proxy_logging_obj=proxy_logging_obj,
)
except UserNotFoundError:
return None
async def enforced_model_allowlists(
valid_token: UserAPIKeyAuth,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> tuple[Sequence[str], ...]:
"""One model allowlist per level that ``common_checks`` enforces on a request from this identity."""
key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token)
if prisma_client is None:
return (key_models, tuple(valid_token.team_models or ()))
team_object: Final = (
None
if valid_token.team_id is None
else await get_team_object(
team_id=valid_token.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
)
user_object: Final = (
None
if team_object is not None
else await _user_object_or_none(
valid_token=valid_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
)
project_object: Final = (
None
if valid_token.project_id is None
else await get_project_object(
project_id=valid_token.project_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
)
return (
key_models,
team_object.models if team_object is not None else (),
await _team_member_granted_models(
valid_token=valid_token,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
),
user_object.models if user_object is not None else (),
project_object.models if project_object is not None else (),
)
async def collect_matched_model_access_groups(
model: str | Sequence[str] | None,
valid_token: UserAPIKeyAuth | None,

3
litellm/proxy/common_utils/realtime_utils.py
View file

@ -1,3 +1,4 @@
import json
from functools import lru_cache
from typing import Final
@ -10,7 +11,7 @@ def realtime_request_body(model: str | None) -> bytes:
Generate the realtime websocket request body. Cached with LRU semantics to avoid repeated
string formatting work while keeping memory usage bounded.
"""
return f'{{"model": "{model or ""}"}}'.encode()
return json.dumps({"model": model or ""}).encode()
_realtime_request_body: Final = realtime_request_body

172
litellm/proxy/dev_reload.py Normal file
View file

@ -0,0 +1,172 @@
import importlib
import logging
import multiprocessing
import os
import signal
import sys
import threading
from collections.abc import Callable, Mapping, Sequence
from functools import partial
from inspect import signature
from multiprocessing.process import BaseProcess
from socket import socket
from typing import Final, Literal
from pydantic import ConfigDict, ImportString, JsonValue, TypeAdapter, with_config
from typing_extensions import ReadOnly, Required, TypedDict
from uvicorn import Config, Server
from uvicorn import config as uvicorn_config
from uvicorn.supervisors import ChangeReload, Multiprocess
@with_config(ConfigDict(extra="allow"))
class ServerOptions(TypedDict, total=False):
host: ReadOnly[str]
port: ReadOnly[int]
server_header: ReadOnly[bool]
timeout_graceful_shutdown: ReadOnly[float]
log_config: ReadOnly[ # mutable-ok: uvicorn.Config requires a concrete logging dict
str | dict[str, JsonValue] | None
]
log_level: ReadOnly[str | int]
timeout_keep_alive: ReadOnly[int]
timeout_worker_healthcheck: ReadOnly[int]
limit_max_requests: ReadOnly[int]
limit_max_requests_jitter: ReadOnly[int]
limit_concurrency: ReadOnly[int]
ssl_keyfile: ReadOnly[str]
ssl_certfile: ReadOnly[str]
loop: ReadOnly[Literal["none", "auto", "asyncio", "uvloop"]]
@with_config(ConfigDict(extra="allow"))
class ReloadOptions(ServerOptions):
app: Required[ReadOnly[str]]
reload: ReadOnly[bool]
reload_dirs: ReadOnly[tuple[str, ...]]
reload_includes: ReadOnly[tuple[str, ...]]
RELOAD_OPTIONS: Final = TypeAdapter(ReloadOptions)
SERVER_OPTIONS: Final = TypeAdapter(ServerOptions)
ASGI_APP: Final = TypeAdapter(ImportString[Callable[..., object]])
LOGGER: Final = logging.getLogger("uvicorn.error")
RELOAD_WORKER_HEALTHCHECK_SECONDS: Final = 30
RELOADER_ONLY_KEYS: Final = frozenset(ReloadOptions.__annotations__) - frozenset(ServerOptions.__annotations__)
def multi_worker_reload_supported(platform: str = sys.platform) -> bool:
return platform != "win32" and hasattr(uvicorn_config, "STARTUP_FAILURE")
def server_options(options: ReloadOptions) -> ServerOptions:
return SERVER_OPTIONS.validate_python(
{key: value for key, value in options.items() if key not in RELOADER_ONLY_KEYS}
)
def load_app_or_exit(import_path: str, factory: bool = False) -> Callable[..., object]:
module_name: Final = import_path.partition(":")[0]
try:
importlib.import_module(module_name)
app: Final = ASGI_APP.validate_python(import_path)
return ASGI_APP.validate_python(app()) if factory else app
except Exception:
LOGGER.exception("Error loading ASGI app %s", import_path)
sys.exit(uvicorn_config.STARTUP_FAILURE)
def stop_worker_when_pool_exits(pool_parent: BaseProcess) -> None:
pool_parent.join()
LOGGER.info("Worker pool process [%s] exited, stopping worker [%s].", pool_parent.pid, os.getpid())
os.kill(os.getpid(), signal.SIGTERM)
def load_worker_app(import_path: str, worker_app_factory: bool) -> Callable[..., object]:
pool_parent: Final = multiprocessing.parent_process()
if pool_parent is not None:
threading.Thread(target=stop_worker_when_pool_exits, args=(pool_parent,), daemon=True).start()
return load_app_or_exit(import_path, worker_app_factory)
def pool_config(options: ReloadOptions, workers: int, worker_overrides: Mapping[str, object] | None = None) -> Config:
overrides: Final = worker_overrides or {}
worker_app: Final = TypeAdapter(str).validate_python(overrides.get("app", options["app"]))
worker_app_factory: Final = TypeAdapter(bool).validate_python(overrides.get("factory", False))
return Config(
partial(load_worker_app, worker_app, worker_app_factory),
factory=True,
workers=workers,
**SERVER_OPTIONS.validate_python(
{
"timeout_worker_healthcheck": RELOAD_WORKER_HEALTHCHECK_SECONDS,
**server_options(options),
**{key: value for key, value in overrides.items() if key not in ("app", "factory")},
}
), # pyright: ignore[reportArgumentType] # uvicorn types the float graceful shutdown timeout as int
)
def stop_pool_when_reloader_exits(reloader: BaseProcess, pool_should_exit: threading.Event) -> None:
reloader.join()
LOGGER.info("Reloader process [%s] exited, stopping the worker pool.", reloader.pid)
pool_should_exit.set()
class ReloadWorkerPool(Multiprocess):
quit_received: bool = False
def handle_quit(self) -> None:
LOGGER.info("Received SIGQUIT, exiting.")
self.quit_received = True
self.should_exit.set()
def terminate_all(self) -> None:
quit_signal: Final = getattr(signal, "SIGQUIT", None)
if quit_signal is not None and quit_signal in self.signal_queue:
self.handle_quit()
if not self.quit_received:
super().terminate_all()
return
for process in self.processes:
process.kill()
def _run_worker_pool(config: Config, sockets: Sequence[socket] | None) -> None:
shared_sockets: Final = list(sockets or ())
pool: Final = (
ReloadWorkerPool(
config=config,
target=Server(config=config).run, # pyright: ignore[reportCallIssue] # Older uvicorn versions require target
sockets=shared_sockets,
)
if "target" in signature(Multiprocess).parameters
else ReloadWorkerPool(config=config, sockets=shared_sockets)
)
reloader: Final = multiprocessing.parent_process()
if reloader is not None:
threading.Thread(target=stop_pool_when_reloader_exits, args=(reloader, pool.should_exit), daemon=True).start()
pool.run()
def run_multi_worker_reload(
options: ReloadOptions, workers: int, worker_overrides: Mapping[str, object] | None = None
) -> None:
worker_config: Final = pool_config(options, workers, worker_overrides)
reload_config: Final = Config(
options["app"],
workers=1,
reload=options["reload"],
reload_dirs=list(options["reload_dirs"]),
reload_includes=list(options["reload_includes"]),
**server_options(options), # pyright: ignore[reportArgumentType] # uvicorn types the float graceful shutdown timeout as int
)
shared_socket: Final = worker_config.bind_socket()
try:
ChangeReload(
config=reload_config,
target=partial(_run_worker_pool, worker_config),
sockets=[shared_socket],
).run()
finally:
shared_socket.close()

80
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -9,7 +9,7 @@ import os
from collections.abc import Awaitable, Callable, Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType, UnionType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, cast, get_args, get_origin
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, Union, cast, get_args, get_origin
from urllib.parse import urlparse
from fastapi import APIRouter, Depends, HTTPException, Request
@ -43,6 +43,11 @@ from litellm.proxy.guardrails.guardrail_registry import (
encrypt_guardrail_litellm_params,
parse_tolerant_litellm_params,
)
from litellm.proxy.guardrails.response_handlers import (
GUARDRAIL_RESPONSE_HANDLERS,
GuardrailEndpointResponse,
StandardGuardrailResponseHandler,
)
from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router
from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # legacy module exports
_user_has_admin_view, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
@ -53,7 +58,6 @@ from litellm.repositories.table_repositories import GuardrailsRepository
from litellm.types.guardrails import (
PII_ENTITY_CATEGORIES_MAP,
ApplyGuardrailRequest,
ApplyGuardrailResponse,
BaseLitellmParams,
BedrockGuardrailConfigModel,
BedrockGuardrailStreamingParams,
@ -2350,13 +2354,7 @@ async def _run_test_guardrail(execute_guardrail: Callable[[], object], timeout:
return await await_with_timeout(raw_result, remaining, label="test")
def _resolve_guardrail_input_type(active_guardrail: CustomGuardrail, input_type: str) -> Literal["request", "response"]:
"""Return the effective input_type, auto-upgrading to 'response' for post_call guardrails."""
if input_type == "request":
hook: Final = getattr(active_guardrail, "event_hook", None)
if hook == GuardrailEventHooks.post_call or hook == "post_call":
return "response"
return "response" if input_type == "response" else "request"
_RESPONSE: Final[TypeAdapter[GuardrailEndpointResponse]] = TypeAdapter(GuardrailEndpointResponse)
class _GuardrailLoggingObj(Protocol):
@ -2384,35 +2382,43 @@ def _patch_logging_obj_for_guardrail(litellm_logging_obj: _GuardrailLoggingObj,
)
async def _emit_guardrail_success_logs(
async def _guardrail_response_after_success_hook(
response_handler: StandardGuardrailResponseHandler,
proxy_logging_obj: _GuardrailProxyLogging,
litellm_logging_obj: _GuardrailLoggingObj | None,
data: dict,
data: Mapping[str, object],
user_api_key_dict: UserAPIKeyAuth,
response: ApplyGuardrailResponse,
start_time: datetime,
) -> ApplyGuardrailResponse:
"""Fire proxy and LiteLLM success hooks after a successful guardrail run.
Each hook is wrapped defensively so a callback failure never prevents the
caller from receiving the guardrail response. Returns the (possibly
hook-modified) response.
"""
response: GuardrailEndpointResponse,
) -> GuardrailEndpointResponse:
try:
modified: Final = await proxy_logging_obj.post_call_success_hook(
data=data,
user_api_key_dict=user_api_key_dict,
response=response,
)
if isinstance(modified, ApplyGuardrailResponse):
response = modified
return response_handler.accept_hook_result(response, modified)
except Exception:
verbose_proxy_logger.exception("apply_guardrail: post_call_success_hook failed")
return response
async def _emit_guardrail_success_logs(
response_handler: StandardGuardrailResponseHandler,
proxy_logging_obj: _GuardrailProxyLogging,
litellm_logging_obj: _GuardrailLoggingObj | None,
data: Mapping[str, object],
user_api_key_dict: UserAPIKeyAuth,
response: GuardrailEndpointResponse,
start_time: datetime,
) -> GuardrailEndpointResponse:
"""Apply the success hook and log the response returned to the caller."""
updated_response: Final = await _guardrail_response_after_success_hook(
response_handler, proxy_logging_obj, data, user_api_key_dict, response
)
# Build the logging payload after post_call_success_hook so that logged
# data matches what the caller actually receives if the hook modified
# the response.
response_for_logging: Final = {"response": response.model_dump(exclude_none=True)}
response_for_logging: Final = {"response": _RESPONSE.dump_python(updated_response, mode="json")}
if litellm_logging_obj is not None:
end_time: Final = datetime.now(timezone.utc)
@ -2426,11 +2432,11 @@ async def _emit_guardrail_success_logs(
except Exception:
verbose_proxy_logger.exception("apply_guardrail: async_success_handler failed")
return response
return updated_response
@router.post("/guardrails/apply_guardrail", response_model=ApplyGuardrailResponse)
@router.post("/apply_guardrail", response_model=ApplyGuardrailResponse)
@router.post("/guardrails/apply_guardrail", response_model=GuardrailEndpointResponse)
@router.post("/apply_guardrail", response_model=GuardrailEndpointResponse)
async def apply_guardrail(
fastapi_request: Request,
request: ApplyGuardrailRequest,
@ -2455,7 +2461,7 @@ async def apply_guardrail(
)
from litellm.proxy.utils import handle_exception_on_proxy
data: dict = {
data: dict[str, object] = {
"guardrail_name": request.guardrail_name,
"input": [request.text],
"messages": request.messages or [],
@ -2495,18 +2501,15 @@ async def apply_guardrail(
if litellm_logging_obj is not None:
_patch_logging_obj_for_guardrail(litellm_logging_obj, request)
request_data: Final[dict] = {
**({"messages": request.messages} if request.messages is not None else {}),
**({"metadata": request.metadata} if request.metadata is not None else {}),
}
_input_type: Final = _resolve_guardrail_input_type(active_guardrail, request.input_type)
response_handler: Final = GUARDRAIL_RESPONSE_HANDLERS[request.response_format]
request_data: Final = dict(response_handler.prepare_request_data(request))
input_type: Final = response_handler.resolve_input_type(request, active_guardrail)
guardrailed_inputs: Final = await active_guardrail.apply_guardrail(
inputs={"texts": [request.text]},
request_data=request_data,
input_type=_input_type,
input_type=input_type,
)
response_text: Final = guardrailed_inputs.get("texts", [])
response = ApplyGuardrailResponse(response_text=response_text[0] if response_text else request.text)
response: Final = response_handler.build_response(request, guardrailed_inputs, request_data)
except Exception as e:
if litellm_logging_obj is not None and not isinstance(e, HTTPException):
try:
@ -2537,7 +2540,8 @@ async def apply_guardrail(
raise handle_exception_on_proxy(status_typed_proxy_exception(e) if isinstance(e, HTTPException) else e)
# Success logging outside except so a hook error never triggers failure handlers.
response = await _emit_guardrail_success_logs(
logged_response: Final = await _emit_guardrail_success_logs(
response_handler=response_handler,
proxy_logging_obj=proxy_logging_obj,
litellm_logging_obj=litellm_logging_obj,
data=data,
@ -2545,7 +2549,7 @@ async def apply_guardrail(
response=response,
start_time=start_time,
)
return response
return logged_response
# Usage (dashboard) endpoints: overview, detail, logs

11
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -65,6 +65,7 @@ from litellm.proxy.guardrails.anthropic_sse import (
)
from litellm.types.guardrails import (
BedrockChecksConfigModel,
BedrockGuardrailOutputScope,
BedrockGuardrailStreamingParams,
GuardrailEventHooks,
LitellmParams,
@ -289,6 +290,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
streaming_buffer_release_on_scan: bool | None = None,
streaming_scan_in_background: bool | None = None,
async_handler: AsyncHTTPHandler | None = None,
outputScope: BedrockGuardrailOutputScope | None = None,
**kwargs,
):
self.async_handler = async_handler or get_async_httpx_client(
@ -309,6 +311,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
self.guardrailIdentifier = guardrailIdentifier
self.guardrailVersion = guardrailVersion
self.outputScope = outputScope
self.guardrail_provider = "bedrock"
self.chunk_budget_chars = chunk_budget_chars
self.contextual_grounding_from_messages = contextual_grounding_from_messages
@ -1280,10 +1283,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
AWS billed them to ``completed_chunk_usages``, and the attempt log sums those
with the blocking call's own usage.
"""
bedrock_request_data: Final = {
**base_request_data,
"content": content,
}
bedrock_request_data: Final = dict(base_request_data)
bedrock_request_data["content"] = content
if self.outputScope is not None and "outputScope" not in bedrock_request_data:
bedrock_request_data["outputScope"] = self.outputScope.value
prepared_request: Final = await run_aws_signing(
self._prepare_request,
credentials=credentials,

30
litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py
View file

@ -261,7 +261,8 @@ class ResponsesStreamRestorer:
own window, keyed by the family, the item id and the part index.
- Any `*.done` event closes the stream of the same family. Whatever its window still
holds goes out first, as a copy of that stream's last delta event -- so it carries
the stream's own ids, and repeats that event's `sequence_number`. Then every text
the stream's own ids. It takes the done event's `sequence_number`, and every later
event's number moves up by one, so the numbers stay unique and in order. Then every text
field on the done event is restored in full: its string fields other than
identifiers, plus any `part` or `item` it repeats.
- `response.completed` / `response.incomplete` repeat the whole reply, and are
@ -273,6 +274,8 @@ class ResponsesStreamRestorer:
self._rehydrate: Final = rehydrate
self._carries: Final[dict[ResponsesStreamKey, str]] = {} # mutable-ok: per-stream windows advanced in place.
self._last_deltas: Final[dict[ResponsesStreamKey, object]] = {} # mutable-ok: newest delta per stream.
self._inserted: int = 0
self._last_sequence_number: int | None = None
async def restore(self, event: object) -> tuple[object, ...]:
"""The events to emit in place of `event`: any flush, then the event itself."""
@ -281,7 +284,7 @@ class ResponsesStreamRestorer:
return (event,)
if kind.endswith(".delta") and kind not in RESPONSES_BINARY_DELTAS:
await self._restore_delta(event, kind)
return (event,)
return self._numbered((), event)
slots: Final[SlotSink] = []
flushed: Final = await self._flush(responses_stream_key(event, kind)) if kind.endswith(".done") else ()
if kind.endswith(".done"):
@ -294,12 +297,29 @@ class ResponsesStreamRestorer:
for item in read_list(read_field(event, "response"), "output"):
collect_response_item(item, slots)
await rehydrate_slots(slots, self._rehydrate)
return (*flushed, event)
return self._numbered(flushed, event)
async def finish(self) -> tuple[object, ...]:
"""Flushes every stream the provider never closed, e.g. a truncated reply."""
flushed: Final = tuple([await self._flush(key) for key in tuple(self._carries)])
return tuple(itertools.chain.from_iterable(flushed))
flushed: Final = tuple(itertools.chain.from_iterable([await self._flush(key) for key in tuple(self._carries)]))
last: Final = self._last_sequence_number
if last is not None:
for offset, flush in enumerate(flushed, start=1):
write_field(flush, "sequence_number", last + offset)
return flushed
def _numbered(self, flushed: tuple[object, ...], event: object) -> tuple[object, ...]:
"""Gives each flush the slot `event` arrived in and moves `event` and later events up."""
number: Final = read_field(event, "sequence_number")
if not isinstance(number, int):
return (*flushed, event)
for offset, flush in enumerate(flushed):
write_field(flush, "sequence_number", number + self._inserted + offset)
self._inserted += len(flushed)
self._last_sequence_number = number + self._inserted
if self._inserted:
write_field(event, "sequence_number", self._last_sequence_number)
return (*flushed, event)
async def _restore_delta(self, event: object, kind: str) -> None:
text: Final = read_field(event, "delta")

2
litellm/proxy/guardrails/guardrail_hooks/ovalix/__init__.py
View file

@ -19,6 +19,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
application_id: Final = getattr(litellm_params, "application_id", None)
pre_checkpoint_id: Final = getattr(litellm_params, "pre_checkpoint_id", None)
post_checkpoint_id: Final = getattr(litellm_params, "post_checkpoint_id", None)
file_checkpoint_id: Final = getattr(litellm_params, "file_checkpoint_id", None)
_ovalix_callback: Final = OvalixGuardrail(
guardrail_name=guardrail.get("guardrail_name", ""),
@ -27,6 +28,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
application_id=application_id,
pre_checkpoint_id=pre_checkpoint_id,
post_checkpoint_id=post_checkpoint_id,
file_checkpoint_id=file_checkpoint_id,
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
timeout=litellm_params.timeout,

443
litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py
View file

@ -4,13 +4,19 @@ Use Ovalix Guardrails for your LLM calls. Supports pre_call (user input) and
post_call (model output) checkpoints with optional correction/blocking.
"""
import asyncio
import base64
import datetime
import gzip
import hashlib
import mimetypes
import os
from typing import TYPE_CHECKING, Any, Final, Literal
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
import httpx
from typing_extensions import ReadOnly, TypedDict, Unpack
from typing_extensions import TypedDict, Unpack
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import GuardrailRaisedException
@ -18,10 +24,23 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix_extraction import (
FilePart,
extract_file_parts_from_images,
extract_file_parts_from_messages,
extract_tool_calls_from_messages,
extract_tool_results,
make_tool_data,
str_or_none,
tool_call_to_tool_data,
tool_data_key,
tool_result_text_indices,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import GenericGuardrailAPIInputs
@ -32,13 +51,34 @@ if TYPE_CHECKING:
BLOCKED_BY_OVALIX_FALLBACK_MESSAGE: Final = "This message was blocked by Ovalix"
BLOCKED_ACTION_TYPE: Final = "block"
_MODIFY_ACTION_TYPES: Final = ("anonymize", "sanitize")
_DEFAULT_FILE_SIZE_LIMIT: Final = 64 * 1024 * 1024
_NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
_FILE_BLOCK_ESCALATION_REASON: Final = (
"This message was blocked by Ovalix because file content anonymization isn't possible via LiteLLM"
)
_UNDECODABLE_FILE_BLOCK_REASON: Final = (
"This message was blocked by Ovalix because an attached file couldn't be decoded for inspection"
)
_TOOL_BLOCK_ESCALATION_REASON: Final = (
"This message was blocked by Ovalix because tool call anonymization isn't possible via LiteLLM"
)
_TOOL_RESULT_BLOCK_ESCALATION_REASON: Final = (
"This message was blocked by Ovalix because tool result anonymization isn't possible via LiteLLM"
)
def _encode_file_wire_format(raw: bytes) -> str:
return base64.b64encode(gzip.compress(raw)).decode()
def _mapping_or_empty(value: object) -> Mapping[str, object]:
return as_str_mapping(value) or _NO_METADATA
class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
"""Base-class constructor options this guardrail forwards untouched to CustomGuardrail."""
supported_event_hooks: ReadOnly[list[GuardrailEventHooks]]
class OvalixGuardrailMissingSecrets(Exception):
"""Raised when required Ovalix config (API base, key, application/checkpoint IDs) is missing."""
@ -87,6 +127,8 @@ class OvalixGuardrail(CustomGuardrail):
application_id: str | None = None,
pre_checkpoint_id: str | None = None,
post_checkpoint_id: str | None = None,
file_checkpoint_id: str | None = None,
supported_event_hooks: list[GuardrailEventHooks] | None = None,
**kwargs: Unpack[_CustomGuardrailOptions],
):
self._tracker_api_base = tracker_api_base or os.environ.get("OVALIX_TRACKER_API_BASE")
@ -94,23 +136,26 @@ class OvalixGuardrail(CustomGuardrail):
self._application_id = application_id or os.environ.get("OVALIX_APPLICATION_ID")
self._pre_checkpoint_id = pre_checkpoint_id or os.environ.get("OVALIX_PRE_CHECKPOINT_ID")
self._post_checkpoint_id = post_checkpoint_id or os.environ.get("OVALIX_POST_CHECKPOINT_ID")
self._file_checkpoint_id = file_checkpoint_id or os.environ.get("OVALIX_FILE_CHECKPOINT_ID")
supported_event_hooks: Final = kwargs.get("supported_event_hooks", [])
event_hooks: Final = self._validated_event_hooks(supported_event_hooks or ())
self._validate_config(supported_event_hooks)
self._tracker_headers = httpx.Headers(
{
"Authorization": f"Bearer {self._tracker_api_key}",
"Content-Type": "application/json",
},
encoding="utf-8",
self._tracker_headers = dict(
httpx.Headers(
MappingProxyType(
{
"Authorization": f"Bearer {self._tracker_api_key}",
"x-api-key": self._tracker_api_key or "",
"Content-Type": "application/json",
}
),
encoding="utf-8",
)
)
self._async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
forwarded: Final[_CustomGuardrailOptions] = {**kwargs, "supported_event_hooks": supported_event_hooks}
super().__init__(**forwarded)
super().__init__(supported_event_hooks=event_hooks, **kwargs)
verbose_proxy_logger.debug(
"Ovalix Guardrail initialized: tracker=%s, application_id=%s, pre_checkpoint_id=%s, post_checkpoint_id=%s",
self._tracker_api_base,
@ -119,44 +164,53 @@ class OvalixGuardrail(CustomGuardrail):
self._post_checkpoint_id,
)
def _validate_config(self, supported_event_hooks: list[GuardrailEventHooks]) -> None:
"""Ensure required secrets and checkpoint IDs are set; auto-add hooks when IDs are present."""
errors: Final[list[str]] = []
if not self._tracker_api_base:
errors.append("Tracker API base, set OVALIX_TRACKER_API_BASE or pass tracker_api_base")
if not self._tracker_api_key:
errors.append("Tracker API key, set OVALIX_TRACKER_API_KEY or pass tracker_api_key")
if not self._application_id:
errors.append("Application ID, set OVALIX_APPLICATION_ID or pass application_id")
if not self._pre_checkpoint_id and GuardrailEventHooks.pre_call in supported_event_hooks:
errors.append("Pre-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or pass pre_checkpoint_id")
if not self._post_checkpoint_id and GuardrailEventHooks.post_call in supported_event_hooks:
errors.append("Post-checkpoint ID, set OVALIX_POST_CHECKPOINT_ID or pass post_checkpoint_id")
if not self._pre_checkpoint_id and not self._post_checkpoint_id:
errors.append(
"Pre-checkpoint ID or Post-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or OVALIX_POST_CHECKPOINT_ID or pass pre_checkpoint_id or post_checkpoint_id"
def _validated_event_hooks(self, requested: Sequence[GuardrailEventHooks]) -> list[GuardrailEventHooks]:
"""Ensure required secrets and checkpoint IDs are set; return the requested hooks plus one per configured checkpoint."""
errors: Final = tuple(
message
for missing, message in (
(not self._tracker_api_base, "Tracker API base, set OVALIX_TRACKER_API_BASE or pass tracker_api_base"),
(not self._tracker_api_key, "Tracker API key, set OVALIX_TRACKER_API_KEY or pass tracker_api_key"),
(not self._application_id, "Application ID, set OVALIX_APPLICATION_ID or pass application_id"),
(
not self._pre_checkpoint_id and GuardrailEventHooks.pre_call in requested,
"Pre-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or pass pre_checkpoint_id",
),
(
not self._post_checkpoint_id and GuardrailEventHooks.post_call in requested,
"Post-checkpoint ID, set OVALIX_POST_CHECKPOINT_ID or pass post_checkpoint_id",
),
(
not self._pre_checkpoint_id and not self._post_checkpoint_id,
"Pre-checkpoint ID or Post-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or OVALIX_POST_CHECKPOINT_ID or pass pre_checkpoint_id or post_checkpoint_id",
),
)
if missing
)
if errors:
raise OvalixGuardrailMissingSecrets("Missing Ovalix guardrail configuration errors: " + ". ".join(errors))
# auto-add hooks when checkpoint IDs are present
if self._pre_checkpoint_id and GuardrailEventHooks.pre_call not in supported_event_hooks:
supported_event_hooks.append(GuardrailEventHooks.pre_call)
if self._post_checkpoint_id and GuardrailEventHooks.post_call not in supported_event_hooks:
supported_event_hooks.append(GuardrailEventHooks.post_call)
auto_added: Final = tuple(
hook
for configured, hook in (
(self._pre_checkpoint_id, GuardrailEventHooks.pre_call),
(self._post_checkpoint_id, GuardrailEventHooks.post_call),
)
if configured and hook not in requested
)
return [*requested, *auto_added]
def _get_actor(self, data: dict) -> str:
def _get_actor(self, data: Mapping[str, object]) -> str:
"""Return a stable actor identifier from request metadata (e.g. user email or id)."""
metadata: Final = data.get("metadata") or data.get("litellm_metadata") or {}
if metadata.get("user_api_key_user_email"):
return metadata["user_api_key_user_email"]
if metadata.get("user_api_key_user_id"):
return metadata["user_api_key_user_id"]
return "unknown"
metadata: Final = _mapping_or_empty(data.get("metadata") or data.get("litellm_metadata"))
return (
str_or_none(metadata.get("user_api_key_user_email"))
or str_or_none(metadata.get("user_api_key_user_id"))
or "unknown"
)
def _get_tracker_actor_id(self, data: dict) -> str:
def _get_tracker_actor_id(self, data: Mapping[str, object]) -> str:
"""Normalize the actor string into a short, stable id for Tracker API payloads."""
# NOTE: this hash is purely for normalization — it collapses an arbitrary actor
# string (email, user id, or "unknown") into a compact, fixed-length, consistent
@ -166,7 +220,7 @@ class OvalixGuardrail(CustomGuardrail):
normalized_actor_id: Final = hashlib.sha256(actor_id).hexdigest()[:8]
return normalized_actor_id
def _get_session_id(self, data: dict) -> str:
def _get_session_id(self, data: Mapping[str, object]) -> str:
"""Return a unique identifier for the chat/session (actor + date + application_id)."""
actor_hash: Final = self._get_tracker_actor_id(data)
today: Final = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d")
@ -174,106 +228,254 @@ class OvalixGuardrail(CustomGuardrail):
async def _call_checkpoint(
self,
content: str,
data_type: str,
data: Mapping[str, object],
checkpoint_id: str,
actor: str,
session_id: str,
) -> dict[str, Any]:
"""Call the Ovalix Tracker checkpoint API and return the JSON response."""
oversize: bool = False,
) -> Mapping[str, object]:
"""Call the Ovalix Tracker checkpoint API and return the JSON response.
Both routes live on the tracker's /beta litellm router, which accepts the api key this
guardrail already sends.
"""
application_id: Final = self._application_id
if not application_id or not checkpoint_id:
raise ValueError("Ovalix: application_id or checkpoint_id not resolved")
url: Final = f"{self._tracker_api_base}/tracking/custom_application/checkpoint"
headers: Final = dict(self._tracker_headers)
route: Final = "file_checkpoint" if data_type == "FILE" else "checkpoint"
payload: Final = {
"application_id": application_id,
"checkpoint_id": checkpoint_id,
"actor": actor,
"session_id": session_id,
"data_type": "TEXT",
"data": {"content": content},
"data_type": data_type,
"data": data,
"tool": "LiteLLM",
**(MappingProxyType({"oversize": True}) if oversize else _NO_METADATA),
}
response: Final = await self._async_handler.post(url, headers=headers, json=payload, timeout=self.timeout)
response: Final = await self._async_handler.post(
f"{self._tracker_api_base}/tracking/beta/{route}",
headers=self._tracker_headers,
json=payload,
timeout=self.timeout,
)
response.raise_for_status()
return response.json()
def _verdict(self, resp: Mapping[str, object]) -> tuple[str, str | None]:
return (str_or_none(resp.get("action_type")) or "").lower(), self._get_trackers_corrected_message(resp)
async def _block_reason_for_item(
self,
data_type: str,
data: Mapping[str, object],
checkpoint_id: str,
actor: str,
session_id: str,
escalation_reason: str,
oversize: bool = False,
) -> str | None:
try:
resp: Final = await self._call_checkpoint(
data_type, data, checkpoint_id, actor, session_id, oversize=oversize
)
except Exception as e:
verbose_proxy_logger.exception("Ovalix checkpoint call failed: %s", e)
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message=f"Ovalix guardrail error: {e}",
should_wrap_with_default_message=False,
) from e
action, corrected = self._verdict(resp)
if action == BLOCKED_ACTION_TYPE:
return corrected or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE
if action in _MODIFY_ACTION_TYPES:
return escalation_reason
return None
async def _check_items_block_only(
self,
items: Sequence[tuple[str, Mapping[str, object]]],
checkpoint_id: str,
actor: str,
session_id: str,
escalation_reason: str,
) -> str | None:
for data_type, data in items:
reason = await self._block_reason_for_item(
data_type, data, checkpoint_id, actor, session_id, escalation_reason
)
if reason is not None:
return reason
return None
async def _check_files_for_block(
self,
file_parts: Sequence[FilePart],
checkpoint_id: str,
actor: str,
session_id: str,
) -> str | None:
if any(part.undecodable for part in file_parts):
return _UNDECODABLE_FILE_BLOCK_REASON
for part in sorted(file_parts, key=lambda p: p.message_index, reverse=True):
data = await self._file_part_to_data(part)
reason = await self._block_reason_for_item(
"FILE",
data,
checkpoint_id,
actor,
session_id,
_FILE_BLOCK_ESCALATION_REASON,
oversize=part.oversize,
)
if reason is not None:
return reason
return None
@log_guardrail_information
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
request_data: Mapping[str, object],
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
"""
Apply Ovalix guardrail to the given inputs (request or response text).
Used by the unified guardrail flow and the /apply_guardrail API.
For "request", uses the pre-checkpoint; for "response", uses the post-checkpoint.
Args:
inputs: Guardrail API inputs (e.g. texts to check).
request_data: Full request payload (messages, metadata, response).
input_type: "request" (pre_call) or "response" (post_call).
logging_obj: Optional logging context.
Returns:
Updated inputs (e.g. with replaced/corrected texts, or unchanged).
"""
if not self._pre_checkpoint_id and not self._post_checkpoint_id:
return inputs
tracker_actor_id: Final = self._get_tracker_actor_id(request_data)
actor: Final = self._get_actor(request_data)
session_id: Final = self._get_session_id(request_data)
texts: Final = inputs.get("texts") or []
if not texts or not isinstance(texts, list):
is_response: Final = input_type == "response"
prompt_checkpoint: Final = self._post_checkpoint_id if is_response else self._pre_checkpoint_id
file_checkpoint: Final = self._file_checkpoint_id or prompt_checkpoint
if not file_checkpoint:
verbose_proxy_logger.debug(
"Ovalix guardrail: application %s has no %s checkpoint, leaving this direction uninspected",
self._application_id,
input_type,
)
return inputs
if input_type == "response":
if not self._post_checkpoint_id:
return inputs
corrected_llm_responses: Final = await self._generate_post_guardrail_llm_texts(
texts, tracker_actor_id, session_id, self._post_checkpoint_id
structured_messages: Final = inputs.get("structured_messages") or ()
file_parts: Final = (
extract_file_parts_from_images(inputs.get("images"), size_limit=_DEFAULT_FILE_SIZE_LIMIT)
if is_response
else extract_file_parts_from_messages(structured_messages, size_limit=_DEFAULT_FILE_SIZE_LIMIT)
)
file_block: Final = await self._check_files_for_block(file_parts, file_checkpoint, actor, session_id)
if file_block is not None:
self._block_current_message(file_block)
if not prompt_checkpoint:
verbose_proxy_logger.debug(
"Ovalix guardrail: application %s has only a %s file checkpoint, skipping text and tool inspection",
self._application_id,
input_type,
)
return {**inputs, "texts": corrected_llm_responses}
return inputs
if self._pre_checkpoint_id:
post_guardrail_texts: Final = await self._generate_post_guardrail_llm_texts(
texts, tracker_actor_id, session_id, self._pre_checkpoint_id
)
return {**inputs, "texts": post_guardrail_texts}
return inputs
tool_calls: Final = (
*(inputs.get("tool_calls") or ()),
*extract_tool_calls_from_messages(structured_messages),
)
unique_tool_data: Final = MappingProxyType(
{tool_data_key(data): data for data in (tool_call_to_tool_data(tc) for tc in tool_calls) if data}
)
tool_call_items: Final = tuple(("TOOL", data) for data in unique_tool_data.values())
tool_block: Final = await self._check_items_block_only(
tool_call_items,
prompt_checkpoint,
actor,
session_id,
_TOOL_BLOCK_ESCALATION_REASON,
)
if tool_block is not None:
self._block_current_message(tool_block)
async def _generate_post_guardrail_llm_texts(
self, texts: list[str], actor: str, session_id: str, checkpoint_id: str
) -> list[str]:
"""Generate post-guardrail LLM responses for the given LLM responses."""
post_guardrail_texts: Final[list[str]] = []
tool_results: Final = extract_tool_results(structured_messages)
tool_result_items: Final = tuple(("TOOL", make_tool_data(name, content)) for name, content, _ in tool_results)
tool_result_block: Final = await self._check_items_block_only(
tool_result_items,
prompt_checkpoint,
actor,
session_id,
_TOOL_RESULT_BLOCK_ESCALATION_REASON,
)
if tool_result_block is not None:
self._block_current_message(tool_result_block)
is_first_response = True
for llm_response in reversed(texts):
try:
resp = await self._call_checkpoint(llm_response, checkpoint_id, actor, session_id)
except Exception as e:
verbose_proxy_logger.exception("Ovalix apply_guardrail checkpoint call failed: %s", e)
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message=f"Ovalix guardrail error: {e}",
should_wrap_with_default_message=False,
) from e
texts: Final = inputs.get("texts") or ()
if not texts:
return inputs
output_texts: Final = await self._check_texts(
texts,
prompt_checkpoint,
actor,
session_id,
tool_result_text_indices(structured_messages, texts),
)
if output_texts is None:
return inputs
return {**inputs, "texts": output_texts}
action_type = (resp.get("action_type") or "").lower()
blocking_message = self._get_trackers_corrected_message(resp) or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE
if action_type == BLOCKED_ACTION_TYPE and is_first_response:
self._block_current_message(blocking_message)
elif action_type == BLOCKED_ACTION_TYPE:
post_guardrail_texts.insert(0, blocking_message)
else:
corrected_text = self._get_trackers_corrected_message(resp) or llm_response
post_guardrail_texts.insert(0, corrected_text)
is_first_response = False
return post_guardrail_texts
async def _file_part_to_data(self, part: FilePart) -> Mapping[str, object]:
extension: Final = mimetypes.guess_extension(part.mime_hint) if part.mime_hint else None
name: Final = part.name or (f"file{extension}" if extension else "file")
content: Final = (
await asyncio.get_event_loop().run_in_executor(None, _encode_file_wire_format, part.data)
if part.data
else None
)
return {"name": name, "content": content}
async def _check_texts(
self,
texts: Sequence[str],
checkpoint_id: str,
actor: str,
session_id: str,
skip_indices: frozenset[int],
) -> list[str] | None:
newest_first: Final = tuple(index for index in reversed(range(len(texts))) if index not in skip_indices)
checked: Final = MappingProxyType(
{
index: await self._check_text(
texts[index], checkpoint_id, actor, session_id, is_newest=index == len(texts) - 1
)
for index in newest_first
}
)
output: Final = [checked.get(index, text) for index, text in enumerate(texts)]
return output if output != list(texts) else None
async def _check_text(
self,
content: str,
checkpoint_id: str,
actor: str,
session_id: str,
is_newest: bool,
) -> str:
"""Return the text that replaces `content` after Ovalix checks it; a blocked newest text raises."""
try:
resp: Final = await self._call_checkpoint("TEXT", {"content": content}, checkpoint_id, actor, session_id)
except Exception as e:
verbose_proxy_logger.exception("Ovalix checkpoint call failed: %s", e)
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message=f"Ovalix guardrail error: {e}",
should_wrap_with_default_message=False,
) from e
action, corrected = self._verdict(resp)
if action == BLOCKED_ACTION_TYPE:
block_message: Final = corrected or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE
if is_newest:
self._block_current_message(block_message)
return block_message
if action in _MODIFY_ACTION_TYPES and corrected is not None:
return corrected
return content
def _block_current_message(self, blocking_message: str) -> None:
"""Raise OvalixGuardrailBlockedException with the given message (no default wrapper)."""
@ -283,12 +485,13 @@ class OvalixGuardrail(CustomGuardrail):
should_wrap_with_default_message=False,
)
def _get_trackers_corrected_message(self, resp: dict) -> str | None:
def _get_trackers_corrected_message(self, resp: Mapping[str, object]) -> str | None:
"""Extract corrected/blocking message content from Tracker checkpoint response."""
modified: Final = resp.get("modified_data")
if isinstance(modified, dict) and "content" in modified:
return modified["content"]
return None
modified: Final = as_str_mapping(resp.get("modified_data"))
if modified is None:
return None
content: Final = modified.get("content")
return content if isinstance(content, str) else None
@staticmethod
def get_config_model() -> type["GuardrailConfigModel"] | None:

396
litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py Normal file
View file

@ -0,0 +1,396 @@
import base64
import json
import posixpath
import re
from bisect import bisect_left
from collections.abc import Callable, Iterator, Mapping, Sequence
from itertools import chain, groupby
from operator import attrgetter
from types import MappingProxyType
from typing import Final, NamedTuple
from urllib.parse import unquote, urlparse
from pydantic import TypeAdapter, ValidationError
from litellm.litellm_core_utils.core_helpers import as_str_mapping
_TOOL_NAME_MAX_LENGTH: Final = 100
_DEFAULT_TOOL_RESULT_NAME: Final = "tool_result"
_NO_TOOL_INPUT: Final[Mapping[str, object]] = MappingProxyType({})
_DATA_URL_RE: Final = re.compile(r"^data:(?P<mime>[^;,]+)?(?P<params>(?:;[^;,]+)*?)(?P<b64>;base64)?,", re.IGNORECASE)
_URLSAFE_TO_STANDARD_B64: Final = str.maketrans("-_", "+/")
_OBJECT_LIST: Final = TypeAdapter(list[object])
def str_or_none(value: object) -> str | None:
return value if isinstance(value, str) and value else None
def _list_items(value: object) -> Sequence[object]:
try:
return _OBJECT_LIST.validate_python(value, strict=True)
except ValidationError:
return ()
class FilePart(NamedTuple):
name: str | None
data: bytes | None
mime_hint: str | None
inline: bool
oversize: bool
message_index: int
undecodable: bool = False
def _split_data_url(value: str) -> tuple[str | None, str | None]:
match: Final = _DATA_URL_RE.match(value)
if not match:
if value.lower().startswith("data:"):
return None, None
return None, value
mime: Final = match.group("mime") or None
if not match.group("b64"):
return mime, None
return mime, value[match.end() :]
def _b64decode_or_none(payload: str) -> bytes | None:
try:
return base64.b64decode(payload, validate=True)
except ValueError:
return None
def _decoded_base64(cleaned: str) -> bytes | None:
strict: Final = _b64decode_or_none(cleaned)
if strict is not None or ("-" not in cleaned and "_" not in cleaned):
return strict
return _b64decode_or_none(cleaned.translate(_URLSAFE_TO_STANDARD_B64))
def _inline_part(name: str | None, value: str, size_limit: int | None, message_index: int) -> FilePart | None:
mime_hint, payload = _split_data_url(value)
if payload is None:
return FilePart(name, None, mime_hint, True, False, message_index, undecodable=True)
cleaned: Final = "".join(payload.split())
if not cleaned:
return None
if size_limit is not None and (len(cleaned) * 3) // 4 - 2 > size_limit:
return FilePart(name, None, mime_hint, True, True, message_index)
data: Final = _decoded_base64(cleaned)
if data is None:
return FilePart(name, None, mime_hint, True, False, message_index, undecodable=True)
if size_limit is not None and len(data) > size_limit:
return FilePart(name, None, mime_hint, True, True, message_index)
return FilePart(name, data, mime_hint, True, False, message_index)
def _name_from_url(url: str) -> str | None:
try:
return unquote(posixpath.basename(urlparse(url).path)) or None
except ValueError:
return None
def _part_from_file_block(block: Mapping[str, object], size_limit: int | None, message_index: int) -> FilePart | None:
file_obj: Final = as_str_mapping(block.get("file"))
if file_obj is None:
return None
name: Final = str_or_none(file_obj.get("filename")) or str_or_none(file_obj.get("file_id"))
file_data: Final = str_or_none(file_obj.get("file_data"))
inline: Final = _inline_part(name, file_data, size_limit, message_index) if file_data else None
return inline or FilePart(name, None, None, False, False, message_index)
def _part_from_image_url_block(
block: Mapping[str, object], size_limit: int | None, message_index: int
) -> FilePart | None:
image_url: Final = block.get("image_url")
image_url_fields: Final = as_str_mapping(image_url)
url: Final = image_url_fields.get("url") if image_url_fields is not None else image_url
if not isinstance(url, str) or not url:
return None
if url.startswith("data:"):
return _inline_part(None, url, size_limit, message_index)
return FilePart(_name_from_url(url), None, None, False, False, message_index)
def _part_from_input_file_block(
block: Mapping[str, object], size_limit: int | None, message_index: int
) -> FilePart | None:
declared_name: Final = str_or_none(block.get("filename")) or str_or_none(block.get("file_id"))
file_data: Final = str_or_none(block.get("file_data"))
inline: Final = _inline_part(declared_name, file_data, size_limit, message_index) if file_data else None
if inline is not None:
return inline
file_url: Final = block.get("file_url")
name: Final = (
_name_from_url(file_url) if isinstance(file_url, str) and file_url and not declared_name else declared_name
)
return FilePart(name, None, None, False, False, message_index)
def _part_from_input_audio_block(
block: Mapping[str, object], size_limit: int | None, message_index: int
) -> FilePart | None:
audio: Final = as_str_mapping(block.get("input_audio"))
if audio is None:
return None
data_b64: Final = audio.get("data")
if not isinstance(data_b64, str) or not data_b64:
return None
name: Final = f"audio.{audio.get('format') or 'bin'}"
return _inline_part(name, data_b64, size_limit, message_index) or FilePart(
name, None, None, False, False, message_index
)
_BLOCK_PARSERS: Final[Mapping[str, Callable[[Mapping[str, object], int | None, int], FilePart | None]]] = (
MappingProxyType(
{
"file": _part_from_file_block,
"image_url": _part_from_image_url_block,
"input_image": _part_from_image_url_block,
"input_file": _part_from_input_file_block,
"input_audio": _part_from_input_audio_block,
}
)
)
def _file_parts_of_message(
message: Mapping[str, object], size_limit: int | None, message_index: int
) -> Iterator[FilePart]:
for block in map(as_str_mapping, _list_items(message.get("content"))):
if block is None:
continue
block_type = block.get("type")
if not isinstance(block_type, str):
continue
parser = _BLOCK_PARSERS.get(block_type)
if parser is None:
continue
try:
part = parser(block, size_limit, message_index)
except (TypeError, ValueError, AttributeError, KeyError):
continue
if part is not None and (part.inline or part.name):
yield part
def _file_parts_of_messages(structured_messages: Sequence[object], size_limit: int | None) -> Iterator[FilePart]:
for message_index, message in enumerate(map(as_str_mapping, structured_messages)):
if message is not None:
yield from _file_parts_of_message(message, size_limit, message_index)
def extract_file_parts_from_messages(
structured_messages: Sequence[object] | None, size_limit: int | None = None
) -> tuple[FilePart, ...]:
return tuple(_file_parts_of_messages(structured_messages or (), size_limit))
def _file_part_of_image(value: str, size_limit: int | None, index: int) -> FilePart | None:
if value.startswith(("http://", "https://")):
name: Final = _name_from_url(value)
return FilePart(name, None, None, False, False, index) if name else None
return _inline_part(None, value, size_limit, index)
def extract_file_parts_from_images(
images: Sequence[object] | None, size_limit: int | None = None
) -> tuple[FilePart, ...]:
candidates: Final = (
_file_part_of_image(value, size_limit, index)
for index, value in enumerate(images or ())
if isinstance(value, str) and value
)
return tuple(part for part in candidates if part is not None)
def make_tool_data(
name: str, content: str | None, tool_input: Mapping[str, object] | None = None
) -> Mapping[str, object]:
action_name: Final = str(name) if str(name).strip() else _DEFAULT_TOOL_RESULT_NAME
truncated: Final = action_name[:_TOOL_NAME_MAX_LENGTH]
tool_name: Final = truncated if truncated.strip() else _DEFAULT_TOOL_RESULT_NAME
return {
"content": content,
"tool_name": tool_name,
"action_name": action_name,
"tool_input": dict(tool_input or ()),
}
def _tool_call_field(tool_call: object, key: str) -> object:
fields: Final = as_str_mapping(tool_call)
return fields.get(key) if fields is not None else getattr(tool_call, key, None)
def _json_or_str(value: object) -> str:
try:
return json.dumps(value)
except (TypeError, ValueError):
return str(value)
def _parsed_tool_input(raw_arguments: str) -> Mapping[str, object]:
if not raw_arguments:
return _NO_TOOL_INPUT
try:
parsed: Final = json.loads(raw_arguments)
except (ValueError, TypeError):
return _NO_TOOL_INPUT
return as_str_mapping(parsed) or _NO_TOOL_INPUT
def _tool_content_and_input(raw_arguments: object) -> tuple[str, Mapping[str, object]]:
if isinstance(raw_arguments, str):
return raw_arguments, _parsed_tool_input(raw_arguments)
if raw_arguments is None:
return "", _NO_TOOL_INPUT
arguments: Final = as_str_mapping(raw_arguments)
return _json_or_str(raw_arguments), arguments if arguments is not None else _NO_TOOL_INPUT
def tool_call_to_tool_data(tool_call: object) -> Mapping[str, object] | None:
function: Final = _tool_call_field(tool_call, "function")
name: Final = _tool_call_field(function, "name")
if not name or not str(name).strip():
return None
content, tool_input = _tool_content_and_input(_tool_call_field(function, "arguments"))
return make_tool_data(str(name), content, tool_input)
def _message_tool_calls(message: object) -> Sequence[object]:
fields: Final = as_str_mapping(message)
return _list_items(fields.get("tool_calls")) if fields is not None else ()
def extract_tool_calls_from_messages(structured_messages: Sequence[object] | None) -> tuple[object, ...]:
"""Tool calls declared on the messages themselves.
Surfaces such as the Anthropic request path populate ``structured_messages`` but leave the
top-level ``tool_calls`` input empty, so calls made in prior assistant turns are only visible here.
"""
return tuple(chain.from_iterable(_message_tool_calls(message) for message in structured_messages or ()))
def tool_data_key(tool_data: Mapping[str, object]) -> str:
"""Stable identity for a tool payload, so a call reached from two sources is only scanned once."""
return json.dumps(tool_data, sort_keys=True, default=str)
def _tool_content_block_text(block: object) -> str | None:
if isinstance(block, str):
return block
fields: Final = as_str_mapping(block)
return str_or_none(fields.get("text")) if fields is not None else None
def _tool_content_text(content: object) -> str:
if isinstance(content, str):
return content
fields: Final = as_str_mapping(content)
if fields is not None:
return _json_or_str(fields)
texts: Final = (_tool_content_block_text(block) for block in _list_items(content))
return "\n".join(text for text in texts if text is not None)
def _extract_tool_content(content: object) -> str | None:
text: Final = _tool_content_text(content)
return text if text.strip() else None
class _Declaration(NamedTuple):
call_id: str
message_index: int
name: str
def _declarations_in(message_index: int, message: object) -> Iterator[_Declaration]:
fields: Final = as_str_mapping(message)
if fields is None or fields.get("role") != "assistant":
return
for tool_call in map(as_str_mapping, _list_items(fields.get("tool_calls"))):
if tool_call is None:
continue
call_id = tool_call.get("id")
function = as_str_mapping(tool_call.get("function"))
name = str_or_none(function.get("name")) if function is not None else None
if isinstance(call_id, str) and call_id and name and name.strip():
yield _Declaration(call_id, message_index, name)
def _declarations_by_call_id(messages: Sequence[object]) -> Mapping[str, tuple[_Declaration, ...]]:
"""Every named declaration of each tool call id, in message order, built in one pass."""
declarations: Final = chain.from_iterable(
_declarations_in(index, message) for index, message in enumerate(messages)
)
ordered: Final = sorted(declarations, key=attrgetter("call_id", "message_index"))
return MappingProxyType({call_id: tuple(group) for call_id, group in groupby(ordered, key=attrgetter("call_id"))})
def _resolve_tool_name(
declarations: Mapping[str, tuple[_Declaration, ...]], tool_index: int, tool_call_id: object
) -> str:
"""Name from the nearest declaration of ``tool_call_id`` before ``tool_index``."""
if not isinstance(tool_call_id, str) or not tool_call_id:
return _DEFAULT_TOOL_RESULT_NAME
declared: Final = declarations.get(tool_call_id, ())
earlier: Final = bisect_left(declared, tool_index, key=attrgetter("message_index"))
return declared[earlier - 1].name if earlier else _DEFAULT_TOOL_RESULT_NAME
def extract_tool_results(
structured_messages: Sequence[object] | None,
) -> tuple[tuple[str, str, object], ...]:
messages: Final = tuple(structured_messages or ())
declarations: Final = _declarations_by_call_id(messages)
def _results() -> Iterator[tuple[str, str, object]]:
for index, message in enumerate(map(as_str_mapping, messages)):
if message is None or message.get("role") != "tool":
continue
content = _extract_tool_content(message.get("content"))
if content is None:
continue
tool_call_id = message.get("tool_call_id")
yield _resolve_tool_name(declarations, index, tool_call_id), content, tool_call_id
return tuple(_results())
def _block_texts(content: object) -> Iterator[object]:
blocks: Final = (as_str_mapping(block) for block in _list_items(content))
return (block.get("text") for block in blocks if block is not None and block.get("text") is not None)
def _message_text_origins(structured_messages: Sequence[object] | None) -> Iterator[tuple[object, bool]]:
for message in map(as_str_mapping, structured_messages or ()):
if message is None:
continue
content = message.get("content")
from_tool_result = message.get("role") == "tool" and _extract_tool_content(content) is not None
if isinstance(content, str):
yield content, from_tool_result
else:
yield from ((text, from_tool_result) for text in _block_texts(content))
def tool_result_text_indices(structured_messages: Sequence[object] | None, texts: Sequence[str]) -> frozenset[int]:
"""Positions in ``texts`` that hold content already submitted under the TOOL policy.
The chat-completions guardrail flow builds ``texts`` and ``structured_messages`` from the
same message list, so tool-role content lands in both and would otherwise be checked twice.
Other surfaces (e.g. Anthropic messages) build ``texts`` from a differently shaped payload,
so the mapping is only trusted when replaying it reproduces ``texts`` exactly; anything else
falls back to checking every text.
"""
origins: Final = tuple(_message_text_origins(structured_messages))
if tuple(text for text, _ in origins) != tuple(texts):
return frozenset()
return frozenset(index for index, (_, from_tool_result) in enumerate(origins) if from_tool_result)

1
litellm/proxy/guardrails/guardrail_initializers.py
View file

@ -18,6 +18,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
event_hook=litellm_params.mode,
guardrailIdentifier=litellm_params.guardrailIdentifier,
guardrailVersion=litellm_params.guardrailVersion,
outputScope=litellm_params.outputScope,
checks=litellm_params.checks if isinstance(litellm_params.checks, BedrockChecksConfigModel) else None,
content_filter_threshold=litellm_params.content_filter_threshold,
prompt_attack_threshold=litellm_params.prompt_attack_threshold,

117
litellm/proxy/guardrails/response_handlers.py Normal file
View file

@ -0,0 +1,117 @@
"""Response policies for the standalone guardrail endpoint."""
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final, Literal, TypeAlias
from fastapi import HTTPException
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.types.guardrails import (
ApplyGuardrailRequest,
ApplyGuardrailResponse,
GuardrailEventHooks,
GuardrailResponseFormat,
NativeGuardrailResponse,
)
from litellm.types.utils import GenericGuardrailAPIInputs
GuardrailEndpointResponse: TypeAlias = ApplyGuardrailResponse | NativeGuardrailResponse
_NATIVE_RESPONSE: Final[TypeAdapter[NativeGuardrailResponse]] = TypeAdapter(NativeGuardrailResponse)
class StandardGuardrailResponseHandler:
"""Preserve the endpoint's existing text response and hook behavior."""
def prepare_request_data(self, request: ApplyGuardrailRequest) -> Mapping[str, object]:
return request.model_dump(include={"messages", "metadata"}, exclude_none=True)
def resolve_input_type(
self, request: ApplyGuardrailRequest, guardrail: CustomGuardrail
) -> Literal["request", "response"]:
if request.input_type == "request" and guardrail.event_hook in (GuardrailEventHooks.post_call, "post_call"):
return "response"
return "response" if request.input_type == "response" else "request"
def build_response(
self,
request: ApplyGuardrailRequest,
processed_inputs: GenericGuardrailAPIInputs,
request_data: Mapping[str, object],
) -> GuardrailEndpointResponse:
texts: Final = processed_inputs.get("texts", [])
return ApplyGuardrailResponse(response_text=texts[0] if texts else request.text)
def accept_hook_result(self, original: GuardrailEndpointResponse, modified: object) -> GuardrailEndpointResponse:
return modified if isinstance(modified, ApplyGuardrailResponse) else original
class _NativeGuardrailRecord(BaseModel):
model_config = ConfigDict(frozen=True)
guardrail_name: str | None = None
guardrail_response: JsonValue = None
class _NativeGuardrailMetadata(BaseModel):
model_config = ConfigDict(frozen=True)
standard_logging_guardrail_information: tuple[_NativeGuardrailRecord, ...] = ()
class NativeGuardrailResponseHandler(StandardGuardrailResponseHandler):
"""Return structured findings recorded by this guardrail invocation."""
def prepare_request_data(self, request: ApplyGuardrailRequest) -> Mapping[str, object]:
return request.model_dump(
include={"messages", "metadata"},
exclude={"metadata": {"standard_logging_guardrail_information"}},
exclude_none=True,
)
def resolve_input_type(
self, request: ApplyGuardrailRequest, guardrail: CustomGuardrail
) -> Literal["request", "response"]:
if "input_type" in request.model_fields_set:
return "response" if request.input_type == "response" else "request"
return super().resolve_input_type(request, guardrail)
def build_response(
self,
request: ApplyGuardrailRequest,
processed_inputs: GenericGuardrailAPIInputs,
request_data: Mapping[str, object],
) -> GuardrailEndpointResponse:
metadata: Final = request_data.get(get_metadata_variable_name_from_kwargs(request_data))
recorded: Final = _NativeGuardrailMetadata.model_validate_json(
safe_dumps(metadata if metadata is not None else {})
)
findings: Final = tuple(
record.guardrail_response
for record in recorded.standard_logging_guardrail_information
if record.guardrail_name == request.guardrail_name and isinstance(record.guardrail_response, (dict, list))
)
if not findings:
raise HTTPException(
status_code=422,
detail=f"Guardrail '{request.guardrail_name}' did not report structured native findings.",
)
return findings[0] if len(findings) == 1 else findings
def accept_hook_result(self, original: GuardrailEndpointResponse, modified: object) -> GuardrailEndpointResponse:
if modified is None:
return original
return _NATIVE_RESPONSE.validate_python(modified)
GUARDRAIL_RESPONSE_HANDLERS: Final[Mapping[GuardrailResponseFormat, StandardGuardrailResponseHandler]] = (
MappingProxyType(
{
GuardrailResponseFormat.STANDARD: StandardGuardrailResponseHandler(),
GuardrailResponseFormat.NATIVE: NativeGuardrailResponseHandler(),
}
)
)

2
litellm/proxy/hooks/dynamic_rate_limiter.py
View file

@ -118,7 +118,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
elif priority is not None and litellm.priority_reservation is not None:
if os.getenv("LITELLM_LICENSE", None) is None:
verbose_proxy_logger.error(
"PREMIUM FEATURE: Reserving tpm/rpm by priority is a premium feature. Please add a 'LITELLM_LICENSE' to your .env to enable this.\nGet a license: https://docs.litellm.ai/docs/proxy/enterprise."
"PREMIUM FEATURE: Reserving tpm/rpm by priority is a premium feature. Please add a 'LITELLM_LICENSE' to your .env to enable this.\nGet a license: https://docs.litellm.ai/docs/enterprise."
)
else:
value: Final = litellm.priority_reservation[priority]

2
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View file

@ -133,7 +133,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
elif priority is not None and litellm.priority_reservation is not None:
if os.getenv("LITELLM_LICENSE", None) is None:
verbose_proxy_logger.error(
"PREMIUM FEATURE: Reserving tpm/rpm by priority is a premium feature. Please add a 'LITELLM_LICENSE' to your .env to enable this.\nGet a license: https://docs.litellm.ai/docs/proxy/enterprise."
"PREMIUM FEATURE: Reserving tpm/rpm by priority is a premium feature. Please add a 'LITELLM_LICENSE' to your .env to enable this.\nGet a license: https://docs.litellm.ai/docs/enterprise."
)
else:
value: Final = litellm.priority_reservation[priority]

11
litellm/proxy/hooks/parallel_request_limiter.py
View file

@ -13,7 +13,7 @@ from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import RateLimitType
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs
from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_parent_otel_span_from_kwargs
from litellm.proxy._types import CommonProxyErrors, CurrentItemRateLimit, UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import (
get_key_model_rpm_limit,
@ -53,9 +53,16 @@ def _response_total_tokens(kwargs: Mapping[str, object], response_obj: object) -
_CONNECTION_SCOPED_CALL_TYPES: Final = frozenset({CallTypes.aresponses_websocket.value})
def _is_websocket_passthrough_row(kwargs: Mapping[str, object]) -> bool:
passthrough_payload: Final = as_str_mapping(kwargs.get("passthrough_logging_payload"))
return passthrough_payload is not None and passthrough_payload.get("request_method") == "WEBSOCKET"
def _slots_released(kwargs: Mapping[str, object]) -> int:
"""A websocket turn callback releases nothing; the socket's slot is released once at disconnect."""
return 0 if kwargs.get("call_type") in _CONNECTION_SCOPED_CALL_TYPES else 1
if kwargs.get("call_type") in _CONNECTION_SCOPED_CALL_TYPES or _is_websocket_passthrough_row(kwargs):
return 0
return 1
class CacheObject(TypedDict):

6
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -2405,7 +2405,7 @@ async def generate_key_fn(
- soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached.
- tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/enterprise#governance-and-cost)
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
- allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"]
- allowed_passthrough_routes: Optional[list] - List of allowed pass through endpoints for the key. Store the actual endpoint or store a wildcard pattern for a set of endpoints. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through endpoints the key can access, without specifying the routes. If allowed_routes is specified, allowed_pass_through_endpoints is ignored.
@ -2613,7 +2613,7 @@ async def generate_service_account_key_fn(
- tpd_limit: Optional[int] - Specify tpd limit for a given key (Tokens per day). Charged by batch submissions instead of tpm_limit/rpm_limit.
- soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached.
- tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/enterprise#governance-and-cost)
- allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"]
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
Examples:
@ -4045,7 +4045,7 @@ async def update_key_fn(
- models: Optional[list] - Model_name's a user is allowed to call
- tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/enterprise#governance-and-cost)
- spend: Optional[float] - Amount spent by key
- max_budget: Optional[float] - Max budget for key
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}

197
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -15,11 +15,11 @@ import os
import posixpath
import re
import sys
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from functools import partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, cast
import httpx
import openai
@ -51,6 +51,7 @@ from litellm.constants import (
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, attribute_or_none
from litellm.litellm_core_utils.fallback_generalizations import match_routing_generalization
from litellm.llms.anthropic.common_utils import (
AnthropicModelInfo,
@ -85,7 +86,12 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.llms.vllm.common_utils import endpoint_under_api_base
from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
from litellm.proxy.auth.auth_checks import (
can_agent_call_model,
can_key_call_resolved_model,
key_model_aliases_for_auth_check,
team_model_aliases_for_auth_check,
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports
@ -93,6 +99,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul
get_bearer_token,
is_no_auth_dev_mode,
user_api_key_auth,
user_api_key_auth_accepted_websocket_for_model,
user_api_key_auth_websocket,
user_api_key_auth_websocket_for_model,
)
@ -114,6 +121,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
safe_get_request_headers,
safe_set_request_parsed_body,
)
from litellm.proxy.common_utils.model_listing_utils import alias_map
from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER
from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
from litellm.proxy.common_utils.sse_keepalive import (
@ -129,16 +137,26 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import (
stream_usage_relay,
without_injected_usage_frames,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import (
openai_websocket_frame_kind,
)
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
STORED_PROMPT_PROBE_MODEL,
HttpPassThroughEndpointHelpers,
WebsocketFrameModelAllowed,
WebsocketFrameModelGate,
WebsocketFrameModelRefused,
WebsocketFrameModelVerdict,
WebsocketUpstreamFrameClassifier,
YieldsToConfiguredPassThroughRoute,
create_pass_through_route,
create_websocket_passthrough_route,
refuse_websocket,
truncated_close_reason,
websocket_passthrough_request,
)
from litellm.proxy.utils import ProxyLogging as ProxyLoggingType
from litellm.proxy.utils import is_known_model
from litellm.proxy.utils import is_known_model, normalize_route_for_root_path
from litellm.proxy.vector_store_endpoints.utils import (
assert_proxy_admin_for_vector_store_index_management,
assert_user_can_access_vector_store,
@ -198,7 +216,7 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li
"""
try:
model: Final = request_body.get("model")
return is_known_model(model, llm_router)
return isinstance(model, str) and is_known_model(model, llm_router)
except Exception:
return False
@ -228,9 +246,8 @@ def is_passthrough_request_streaming(request_body: object) -> bool:
A JSON body need not be an object, so a list or scalar can reach here; it
carries no streaming flag.
"""
if not isinstance(request_body, dict):
return False
return bool(request_body.get("stream", False))
body: Final = as_str_mapping(request_body)
return body is not None and bool(body.get("stream", False))
def _optional_str(value: object) -> str | None:
@ -3387,19 +3404,6 @@ def _join_url_paths(base_url: httpx.URL, path: str, custom_llm_provider: litellm
return joined_path_str
_OPENAI_WS_ALL_MODEL_ACCESS: Final = frozenset(
{
SpecialModelNames.all_proxy_models.value,
SpecialModelNames.all_team_models.value,
"*",
}
)
def _has_model_restrictions(model_allowlists: tuple[Sequence[str], ...]) -> bool:
return any(str(model) not in _OPENAI_WS_ALL_MODEL_ACCESS for allowlist in model_allowlists for model in allowlist)
@dataclass(frozen=True, slots=True)
class _OpenAIWebsocketRefusal:
close_reason: str
@ -3416,6 +3420,16 @@ class _OpenAIWebsocketErrorFrame(TypedDict):
error: ReadOnly[_OpenAIWebsocketErrorDetail]
_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL: Final = (
"This gateway could not authorize the model named in this frame, so the session was closed. Try again, "
"and ask a proxy admin to check the gateway logs if it keeps happening."
)
_OPENAI_WS_ALIASED_MODEL_REFUSAL: Final = (
"This key cannot use a model alias over OpenAI websocket passthrough. The passthrough sends each frame to "
"OpenAI unchanged, so name the model this key may use directly."
)
_OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal(
close_reason="OpenAI websocket passthrough is disabled",
message=(
@ -3424,14 +3438,6 @@ _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal(
),
)
_OPENAI_WS_MODEL_RESTRICTED_REFUSAL: Final = _OpenAIWebsocketRefusal(
close_reason="Keys with model restrictions cannot use OpenAI websocket passthrough",
message=(
"Keys with model restrictions cannot use OpenAI websocket passthrough, because this route "
"relays frames to the provider without reading which model they ask for."
),
)
def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, object]) -> bool:
setting: Final = general_settings.get("enable_openai_websocket_passthrough")
@ -3440,19 +3446,9 @@ def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, obje
return setting is True
class _OpenAIWebsocketModelAllowlists(Protocol):
async def __call__(self, valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: ...
async def _openai_websocket_refusal(
user_api_key_dict: UserAPIKeyAuth,
general_settings: Mapping[str, object],
model_allowlists: _OpenAIWebsocketModelAllowlists,
) -> _OpenAIWebsocketRefusal | None:
def _openai_websocket_refusal(general_settings: Mapping[str, object]) -> _OpenAIWebsocketRefusal | None:
if not _is_openai_websocket_passthrough_enabled(general_settings):
return _OPENAI_WS_DISABLED_REFUSAL
if _has_model_restrictions(await model_allowlists(user_api_key_dict)):
return _OPENAI_WS_MODEL_RESTRICTED_REFUSAL
return None
@ -3467,6 +3463,8 @@ class _WebsocketRelay(Protocol):
forward_headers: bool,
endpoint: str,
accept_websocket: bool,
client_frame_model_gate: WebsocketFrameModelGate | None = None,
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None,
meter: WebSocketSessionMeter | None = None,
) -> None: ...
@ -3475,18 +3473,95 @@ def _websocket_relay() -> _WebsocketRelay:
return websocket_passthrough_request
def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists:
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
def _proxy_error_type(denial: ProxyException) -> str:
return denial.type.value if isinstance(denial.type, ProxyErrorTypes) else denial.type
async def resolve(valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]:
return await enforced_model_allowlists(
valid_token=valid_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
FrameModelAuthorizer: TypeAlias = Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]]
def _frame_model_authorizer() -> FrameModelAuthorizer:
return user_api_key_auth_accepted_websocket_for_model
def _proxy_model_list() -> Sequence[object] | None:
from litellm.proxy import proxy_server
return as_object_list(attribute_or_none(proxy_server, "llm_model_list"))
def _gateway_remaps_model(model: str, frame_auth: UserAPIKeyAuth, llm_router: Router | None) -> bool:
alias_maps: Final = (
key_model_aliases_for_auth_check(frame_auth) or {},
team_model_aliases_for_auth_check(frame_auth) or {},
alias_map(litellm.model_alias_map),
)
targets: Final = (
*(aliases.get(model) for aliases in alias_maps),
llm_router.registry().resolve(model, kind="alias") if llm_router is not None else None,
)
return any(target not in (None, model) for target in targets)
def _aliased_model_verdict(
allowed: WebsocketFrameModelAllowed, probe: WebsocketFrameModelVerdict
) -> WebsocketFrameModelVerdict:
match probe:
case WebsocketFrameModelRefused(error_type=error_type) if error_type.endswith("_model_access_denied"):
return WebsocketFrameModelRefused(message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=error_type)
case WebsocketFrameModelRefused():
return probe
case WebsocketFrameModelAllowed():
return allowed
def _proxy_frame_model_gate(
authorize_frame_model: Annotated[FrameModelAuthorizer, Depends(_frame_model_authorizer)],
) -> WebsocketFrameModelGate:
from litellm.proxy import proxy_server
from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation
async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelVerdict:
connect_reservation: Final = attribute_or_none(websocket.state, "budget_reservation")
await release_or_invalidate_budget_reservation(
budget_reservation=connect_reservation if isinstance(connect_reservation, dict) else None
)
frame_auth: Final = await authorize_frame_model(websocket, model)
access_groups: Final = tuple(frame_auth.matched_model_access_groups or ())
llm_router: Final = proxy_server.llm_router
try:
await can_key_call_resolved_model(
model=model, llm_model_list=_proxy_model_list(), valid_token=frame_auth, llm_router=llm_router
)
await can_agent_call_model(
model=model,
valid_token=frame_auth,
llm_router=llm_router,
prisma_client=proxy_server.prisma_client,
user_api_key_cache=proxy_server.user_api_key_cache,
proxy_logging_obj=proxy_server.proxy_logging_obj,
)
finally:
await release_or_invalidate_budget_reservation(budget_reservation=frame_auth.budget_reservation)
allowed: Final = WebsocketFrameModelAllowed(access_groups=access_groups)
if not _gateway_remaps_model(model, frame_auth, llm_router):
return allowed
return _aliased_model_verdict(allowed, await gate(websocket, STORED_PROMPT_PROBE_MODEL))
return resolve
async def gate(websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict:
try:
return await authorize(websocket, model)
except ProxyException as denial:
return WebsocketFrameModelRefused(message=denial.message, error_type=_proxy_error_type(denial))
except HTTPException as denial:
return WebsocketFrameModelRefused(message=str(denial.detail), error_type="invalid_request_error")
except litellm.BudgetExceededError as exhausted:
return WebsocketFrameModelRefused(message=exhausted.message, error_type="budget_exceeded")
except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band
verbose_proxy_logger.exception("OpenAI websocket passthrough: the frame model gate failed on %s", model)
return WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error")
return gate
def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None:
@ -3498,6 +3573,14 @@ def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None:
return requested_subprotocols[0] if requested_subprotocols else None
async def _opening_model_refusal(
websocket: WebSocket, query_string: str, frame_model_gate: WebsocketFrameModelGate
) -> WebsocketFrameModelRefused | None:
opening_models: Final = dict.fromkeys(httpx.QueryParams(query_string).get_list("model"))
verdicts: Final = tuple([await frame_model_gate(websocket, model) for model in opening_models])
return next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None)
@router.websocket("/openai_passthrough/{endpoint:path}")
@router.websocket("/openai/{endpoint:path}")
async def openai_websocket_proxy_route(
@ -3506,12 +3589,12 @@ async def openai_websocket_proxy_route(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)],
general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)],
relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)],
model_allowlists: Annotated[_OpenAIWebsocketModelAllowlists, Depends(_proxy_model_allowlists)],
frame_model_gate: Annotated[WebsocketFrameModelGate, Depends(_proxy_frame_model_gate)],
) -> None:
"""WebSocket passthrough for OpenAI prefixes (realtime / responses.connect)."""
negotiated_subprotocol: Final = _negotiated_websocket_subprotocol(websocket)
refusal: Final = await _openai_websocket_refusal(user_api_key_dict, general_settings, model_allowlists)
refusal: Final = _openai_websocket_refusal(general_settings)
if refusal is not None:
await websocket.accept(subprotocol=negotiated_subprotocol)
error_frame: Final[_OpenAIWebsocketErrorFrame] = {
@ -3557,14 +3640,21 @@ async def openai_websocket_proxy_route(
await websocket.accept(subprotocol=negotiated_subprotocol)
opening_refusal: Final = await _opening_model_refusal(websocket, query_string, frame_model_gate)
if opening_refusal is not None:
await refuse_websocket(websocket, opening_refusal)
return
await relay(
websocket=websocket,
target=wss_target,
custom_headers=custom_headers,
user_api_key_dict=user_api_key_dict,
forward_headers=False,
endpoint=websocket.url.path,
endpoint=normalize_route_for_root_path(websocket.url.path) or websocket.url.path,
accept_websocket=False,
client_frame_model_gate=frame_model_gate,
upstream_frame_kind=openai_websocket_frame_kind,
)
@ -3882,7 +3972,8 @@ async def _tinyfish_json_object_field_names(request: Request) -> frozenset[str]
parsed: Final = _DECODED_JSON.validate_python(json.loads(raw_body))
except (json.JSONDecodeError, UnicodeDecodeError):
return None
return frozenset(parsed) if isinstance(parsed, dict) else None
fields: Final = as_str_mapping(parsed)
return None if fields is None else frozenset(fields)
def _tinyfish_route_timeout() -> float | None:

262
litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py Normal file
View file

@ -0,0 +1,262 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import Final, Literal, TypeAlias, cast # noqa: TID251 # cost helpers type raw frames as a TypedDict list
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.cost_calculator import RealtimeAPITokenUsageProcessor, handle_realtime_stream_cost_calculation
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponsesAPIResponse
from litellm.types.passthrough_endpoints.pass_through_endpoints import WebsocketUpstreamFrameKind
from litellm.types.utils import LlmProviders, ModelResponse, Usage
WebsocketMessages: TypeAlias = Sequence[Mapping[str, object]]
OPENAI_WEBSOCKET_PROVIDER: Final = LlmProviders.OPENAI.value
_SESSION_CREATED: Final = "session.created"
_SESSION_UPDATED: Final = "session.updated"
_RESPONSE_CREATED: Final = "response.created"
_ERROR: Final = "error"
_REALTIME_RESPONSE_DONE: Final = "response.done"
_TRANSCRIPTION_COMPLETED: Final = "conversation.item.input_audio_transcription.completed"
_RESPONSES_TURN_ENDS: Final = frozenset({"response.completed", "response.incomplete", "response.failed"})
_TURN_ENDS: Final = _RESPONSES_TURN_ENDS | {_REALTIME_RESPONSE_DONE}
_MODEL_NAMING_EVENTS: Final = _TURN_ENDS | {_SESSION_CREATED, _SESSION_UPDATED, _RESPONSE_CREATED}
_REALTIME_BILLED_EVENTS: Final = frozenset({_REALTIME_RESPONSE_DONE, _TRANSCRIPTION_COMPLETED})
_RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost"
_FRAME_KINDS: Final[Mapping[str, WebsocketUpstreamFrameKind]] = MappingProxyType(
{
_SESSION_CREATED: "opens_session",
_SESSION_UPDATED: "updates_session",
_RESPONSE_CREATED: "starts_turn",
_TRANSCRIPTION_COMPLETED: "bills_usage",
_ERROR: "refuses_turn",
}
)
class _NamedModel(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
model: str | None = None
class _AudioInput(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
transcription: _NamedModel | None = None
class _SessionAudio(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
input: _AudioInput | None = None
class _NamedSession(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
model: str | None = None
audio: _SessionAudio | None = None
class _ModelNamingEvent(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
session: _NamedSession | None = None
response: _NamedModel | None = None
class _ResponsesTurnEndEvent(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
type: Literal["response.completed", "response.incomplete", "response.failed"]
response: ResponsesAPIResponse
_MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent)
_RESPONSES_TURN_END_EVENT: Final = TypeAdapter(_ResponsesTurnEndEvent)
@dataclass(frozen=True, slots=True)
class _Billing:
usage: Usage
cost: float
@dataclass(frozen=True, slots=True)
class _CompletedTurn:
model: str
response: ResponsesAPIResponse
usage: Usage
def _event_model(event: Mapping[str, object]) -> str | None:
if event.get("type") not in _MODEL_NAMING_EVENTS:
return None
try:
naming: Final = _MODEL_NAMING_EVENT.validate_python(event)
except ValidationError:
return None
session: Final = naming.session or _NamedSession()
response: Final = naming.response or _NamedModel()
audio_input: Final = (session.audio or _SessionAudio()).input or _AudioInput()
transcription: Final = audio_input.transcription or _NamedModel()
return session.model or response.model or transcription.model or None
def observed_websocket_model(messages: WebsocketMessages) -> str | None:
return next((model for model in map(_event_model, messages) if model is not None), None)
def openai_websocket_frame_kind(frame: Mapping[str, object]) -> WebsocketUpstreamFrameKind:
frame_type: Final = frame.get("type")
if frame_type in _TURN_ENDS:
return "ends_turn"
if not isinstance(frame_type, str):
return "relayed"
return _FRAME_KINDS.get(frame_type, "relayed")
def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: LiteLLMLoggingObj) -> _Billing | None:
if not any(event.get("type") in _REALTIME_BILLED_EVENTS for event in messages):
return None
results: Final = cast( # cast-ok: the realtime cost helpers read raw provider frames through their TypedDict view
OpenAIRealtimeStreamList, tuple(event for event in messages if "type" in event)
)
usage: Final = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results)
cost: Final = handle_realtime_stream_cost_calculation(
results=results,
combined_usage_object=usage,
custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER,
litellm_model_name=model,
litellm_logging_obj=logging_obj,
)
return _Billing(usage=usage, cost=cost)
def _responses_turn(event: Mapping[str, object], session_model: str) -> _CompletedTurn | None:
if event.get("type") not in _RESPONSES_TURN_ENDS:
return None
try:
ended: Final = _RESPONSES_TURN_END_EVENT.validate_python(event)
except ValidationError:
verbose_proxy_logger.warning(
"OpenAI websocket passthrough: a %s frame carried no readable usage, so it is not billed", event.get("type")
)
return None
if ended.response.usage is None:
return None
return _CompletedTurn(
model=ended.response.model or session_model,
response=ended.response,
usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter
ended.response.usage
),
)
def _responses_billing(messages: WebsocketMessages, model: str, requested_model: str | None) -> _Billing | None:
turns: Final = tuple(turn for turn in (_responses_turn(event, model) for event in messages) if turn is not None)
if not turns:
return None
usage: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects([turn.usage for turn in turns])
return _Billing(
usage=usage,
cost=sum(_responses_cost(turn.response, (turn.model, requested_model)) for turn in turns),
)
def _model_cost(response: ResponsesAPIResponse, model: str) -> float | None:
try:
return litellm.completion_cost(
completion_response=response,
model=model,
custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER,
call_type="responses",
)
except Exception: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriced model, and the turn must still log
return None
def _responses_cost(response: ResponsesAPIResponse, price_as: Sequence[str | None]) -> float:
costs: Final = (_model_cost(response, model) for model in dict.fromkeys(price_as) if model is not None)
cost: Final = next((cost for cost in costs if cost), None)
if cost is None:
verbose_proxy_logger.warning(
"OpenAI websocket passthrough: could not price a Responses turn on %s, logging it at zero", price_as
)
return 0.0
return cost
def _requested_model(kwargs: Mapping[str, object]) -> str | None:
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params"))
metadata: Final = None if litellm_params is None else as_str_mapping(litellm_params.get("metadata"))
model_group: Final = None if metadata is None else metadata.get("model_group")
return model_group if isinstance(model_group, str) else None
def _priced_result(model: str, billing: _Billing, start_time: datetime, call_id: str) -> ModelResponse:
result: Final = ModelResponse(
id=call_id,
object="chat.completion",
created=int(start_time.timestamp()),
model=model,
usage=billing.usage,
)
hidden: Final = {"additional_headers": MappingProxyType({_RESPONSE_COST_HEADER: billing.cost})}
result._hidden_params = hidden # pyright: ignore[reportPrivateUsage] # the one settled-cost channel
return result
class OpenAIWebsocketPassthroughLoggingHandler:
def openai_websocket_passthrough_handler(
self,
websocket_messages: WebsocketMessages,
logging_obj: LiteLLMLoggingObj,
url_route: str,
start_time: datetime,
kwargs: Mapping[str, object],
) -> PassThroughEndpointLoggingTypedDict:
model: Final = observed_websocket_model(websocket_messages)
if model is None:
verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no model named, skipping cost", url_route)
unpriced: Final[PassThroughEndpointLoggingTypedDict] = {"result": None, "kwargs": {**kwargs}}
return unpriced
logging_obj.model = model # rebind-ok: the spend logger reads model and cost off the shared logging object
logging_obj.model_call_details.update(
MappingProxyType({"model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER})
)
billing: Final = _realtime_billing(websocket_messages, model, logging_obj) or _responses_billing(
websocket_messages, model, _requested_model(kwargs)
)
if billing is None:
verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no usage reported, skipping cost", url_route)
unbilled: Final[PassThroughEndpointLoggingTypedDict] = {
"result": None,
"kwargs": {**kwargs, "model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER},
}
return unbilled
logging_obj.model_call_details.update(MappingProxyType({"response_cost": billing.cost}))
priced: Final[PassThroughEndpointLoggingTypedDict] = {
"result": _priced_result(model, billing, start_time, logging_obj.litellm_call_id),
"kwargs": {
**kwargs,
"model": model,
"custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER,
"response_cost": billing.cost,
},
}
return priced

734
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -6,15 +6,17 @@ import posixpath
import re
import traceback
from base64 import b64encode
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, MutableMapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from itertools import chain, count, groupby
from functools import reduce
from itertools import chain, count, groupby, takewhile
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
from urllib.parse import urlencode, urlparse
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, TypeVar, cast
from urllib.parse import parse_qs, urlencode, urlparse
import httpx
import orjson
from fastapi import (
APIRouter,
Depends,
@ -24,6 +26,7 @@ from fastapi import (
Response,
UploadFile,
WebSocket,
WebSocketDisconnect,
status,
)
from fastapi.responses import StreamingResponse
@ -34,7 +37,7 @@ from starlette.routing import BaseRoute, Match, Route
from starlette.types import Scope
from starlette.websockets import WebSocketState
from typing_extensions import ReadOnly
from websockets.asyncio.client import connect
from websockets.asyncio.client import ClientConnection, connect
from websockets.exceptions import (
ConnectionClosedError,
ConnectionClosedOK,
@ -60,6 +63,8 @@ from litellm.integrations.custom_guardrail import (
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation
from litellm.litellm_core_utils.core_helpers import (
as_object_list,
as_str_mapping,
bind_budget_reservation_to_callbacks,
get_metadata_variable_name_from_kwargs,
get_or_create_metadata_bucket,
@ -69,6 +74,7 @@ from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import get_masked_values
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.litellm_core_utils.realtime_errors import realtime_error_event
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.anthropic.common_utils import (
@ -154,6 +160,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
WebSocketSessionMeter,
WebSocketSessionUsage,
WebsocketUpstreamFrameKind,
)
from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, Usage
@ -2574,11 +2581,16 @@ def _resolved_vertex_live_setup(
return {**setup_data, "model": setup_model_rewriter(setup_model)}
def _decoded_frame(frame: str | bytes) -> object:
try:
return _DECODED_JSON.validate_python(orjson.loads(frame))
except orjson.JSONDecodeError:
return _DECODED_JSON.validate_python(json.loads(frame if isinstance(frame, str) else frame.decode("utf-8")))
def _json_object_frame(frame: str | bytes) -> dict[str, object] | None:
try:
decoded: Final[object] = _DECODED_JSON.validate_python(
json.loads(frame if isinstance(frame, str) else frame.decode("utf-8"))
)
decoded: Final = _decoded_frame(frame)
except (json.JSONDecodeError, UnicodeDecodeError):
return None
return decoded if isinstance(decoded, dict) else None
@ -2637,6 +2649,359 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict
return inject_trace_context(headers, parent_span=parent_span)
@dataclass(frozen=True, slots=True)
class WebsocketFrameModelAllowed:
access_groups: tuple[str, ...]
@dataclass(frozen=True, slots=True)
class WebsocketFrameModelRefused:
message: str
error_type: str
WebsocketFrameModelVerdict: TypeAlias = WebsocketFrameModelAllowed | WebsocketFrameModelRefused
class WebsocketFrameModelGate(Protocol):
async def __call__(self, websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: ...
WebsocketUpstreamFrameClassifier: TypeAlias = Callable[[Mapping[str, object]], WebsocketUpstreamFrameKind]
_SESSION_FRAMES_EVERY_ROW_CARRIES: Final[frozenset[WebsocketUpstreamFrameKind]] = frozenset(
{"opens_session", "updates_session"}
)
_FrameT: Final = TypeVar("_FrameT", bound=Mapping[str, object])
@dataclass(frozen=True, slots=True)
class _TurnAuthorization:
requested_model: str | None
access_groups: tuple[str, ...]
@dataclass(frozen=True, slots=True)
class _ForwardClientFrame:
starts_turn: bool
turn: _TurnAuthorization | None
lane: str | None
@dataclass(frozen=True, slots=True)
class _RefuseClientFrame:
refusal: WebsocketFrameModelRefused
@dataclass(frozen=True, slots=True)
class _ClientFrameRequests:
starts_turn: bool
lane: str | None
turn_model: str | None
models: tuple[str, ...]
unverifiable_model_refusal: str | None
_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests(
starts_turn=False, lane=None, turn_model=None, models=(), unverifiable_model_refusal=None
)
_FORWARD_WITHOUT_TURN: Final = _ForwardClientFrame(starts_turn=False, turn=None, lane=None)
_UNREADABLE_CLIENT_FRAME: Final = _RefuseClientFrame(
refusal=WebsocketFrameModelRefused(
message="The gateway cannot parse this frame, so it cannot check which model it asks for.",
error_type="invalid_request_error",
)
)
_TURN_STARTING_FRAME: Final = "response.create"
_FRAME_TEXT_WORTH_PARSING: Final = (f'"{_TURN_STARTING_FRAME}"', '"model"', '"prompt"', '"tools"', "\\")
_TURN_MODEL_PLACEMENTS: Final = (("model",), ("response", "model"), ("session", "model"))
_TRANSCRIPTION_MODEL_PLACEMENTS: Final = (
("session", "input_audio_transcription", "model"),
("session", "audio", "input", "transcription", "model"),
)
_STORED_PROMPT_PLACEMENTS: Final = (("prompt",), ("response", "prompt"), ("session", "prompt"))
_TOOL_LIST_PLACEMENTS: Final = (("tools",), ("response", "tools"), ("session", "tools"))
_TOOL_THAT_RUNS_A_MODEL: Final = "image_generation"
STORED_PROMPT_PROBE_MODEL: Final = "litellm-unverifiable-stored-prompt-model"
_STORED_PROMPT_REFUSAL: Final = (
"This key is limited to specific models, so it cannot run a stored prompt. "
"The gateway cannot check which model a stored prompt runs."
)
_UNNAMED_TOOL_MODEL_REFUSAL: Final = (
"This key is limited to specific models, so an image_generation tool has to name its model. "
"The gateway cannot check which model an unnamed one runs."
)
_FIRST_SENTENCE: Final = re.compile(r".*?[.!?](?=\s|$)")
_LANE_NAME: Final = re.compile(r"[A-Za-z0-9_.-]{1,256}")
_UNROUTABLE_LANE_REFUSAL: Final = _RefuseClientFrame(
refusal=WebsocketFrameModelRefused(
message=(
"This stream_id is not a lane name the provider accepts. "
"The provider's error for it names no lane, so the gateway cannot tell which turn it refuses."
),
error_type="invalid_request_error",
)
)
def _may_name_a_model(text_data: str) -> bool:
return any(marker in text_data for marker in _FRAME_TEXT_WORTH_PARSING)
def _frame_lane(frame: Mapping[str, object]) -> str | None:
stream_id: Final = frame.get("stream_id")
return stream_id if isinstance(stream_id, str) else None
def _names_an_unroutable_lane(frame: Mapping[str, object]) -> bool:
stream_id: Final = frame.get("stream_id")
return stream_id is not None and not (isinstance(stream_id, str) and _LANE_NAME.fullmatch(stream_id) is not None)
def _frames_of_lane(lane: str | None) -> Callable[[Mapping[str, object]], bool]:
def in_lane(frame: Mapping[str, object]) -> bool:
return _frame_lane(frame) == lane
return in_lane
def _every_frame(_frame: Mapping[str, object]) -> bool:
return True
def _without_earlier_session_updates(
upstream_frame_kind: WebsocketUpstreamFrameClassifier,
frames: Sequence[_FrameT],
latest_update: _FrameT,
) -> tuple[_FrameT, ...]:
return tuple(
frame
for index, frame in enumerate(frames)
if frame is latest_update
or upstream_frame_kind(frame) != "updates_session"
or _unbilled_usage_ran_under(map(upstream_frame_kind, frames[index + 1 :]))
)
def _unbilled_usage_ran_under(later_kinds: Iterable[WebsocketUpstreamFrameKind]) -> bool:
return "bills_usage" in takewhile(lambda kind: kind != "updates_session", later_kinds)
def _nested_frame(frame: Mapping[str, object] | None, key: str) -> Mapping[str, object] | None:
return None if frame is None else as_str_mapping(frame.get(key))
def _placed_value(frame: Mapping[str, object], placement: Sequence[str]) -> object:
*parents, leaf = placement
container: Final = reduce(_nested_frame, parents, frame)
return None if container is None else container.get(leaf)
def _placed_models(frame: Mapping[str, object], placements: Iterable[Sequence[str]]) -> tuple[str, ...]:
placed: Final = (_placed_value(frame, placement) for placement in placements)
return tuple(value for value in placed if isinstance(value, str))
def _placed_tools(frame: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
tool_lists: Final = (_placed_value(frame, placement) for placement in _TOOL_LIST_PLACEMENTS)
tools: Final = chain.from_iterable(as_object_list(tool_list) or () for tool_list in tool_lists)
return tuple(tool for tool in map(as_str_mapping, tools) if tool is not None)
def _unverifiable_model_refusal(frame: Mapping[str, object], tools: Sequence[Mapping[str, object]]) -> str | None:
prompts: Final = (_placed_value(frame, placement) for placement in _STORED_PROMPT_PLACEMENTS)
if any(isinstance(prompt, Mapping) for prompt in prompts):
return _STORED_PROMPT_REFUSAL
if any(tool.get("type") == _TOOL_THAT_RUNS_A_MODEL and not isinstance(tool.get("model"), str) for tool in tools):
return _UNNAMED_TOOL_MODEL_REFUSAL
return None
def _client_frame_requests(frame_data: str | bytes) -> _ClientFrameRequests | _RefuseClientFrame:
if isinstance(frame_data, str) and not _may_name_a_model(frame_data):
return _NO_CLIENT_FRAME_REQUESTS
try:
frame: Final = _json_object_frame(frame_data)
except (RecursionError, ValueError):
return _UNREADABLE_CLIENT_FRAME
if frame is None:
return _NO_CLIENT_FRAME_REQUESTS
if _names_an_unroutable_lane(frame):
return _UNROUTABLE_LANE_REFUSAL
turn_models: Final = _placed_models(frame, _TURN_MODEL_PLACEMENTS)
tools: Final = _placed_tools(frame)
tool_models: Final = (model for model in (tool.get("model") for tool in tools) if isinstance(model, str))
return _ClientFrameRequests(
starts_turn=frame.get("type") == _TURN_STARTING_FRAME,
lane=_frame_lane(frame),
turn_model=next(iter(turn_models), None),
models=tuple(
dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS), *tool_models))
),
unverifiable_model_refusal=_unverifiable_model_refusal(frame, tools),
)
def _unverifiable_model_probe_refusal(probe: WebsocketFrameModelRefused, message: str) -> WebsocketFrameModelRefused:
if probe.error_type.endswith("_model_access_denied"):
return WebsocketFrameModelRefused(message=message, error_type=probe.error_type)
return probe
async def _client_frame_verdict(
websocket: WebSocket, frame_data: str | bytes, gate: WebsocketFrameModelGate | None
) -> _ForwardClientFrame | _RefuseClientFrame:
if gate is None:
return _FORWARD_WITHOUT_TURN
requests: Final = _client_frame_requests(frame_data)
if isinstance(requests, _RefuseClientFrame):
return requests
verdicts: Final = tuple([await gate(websocket, model) for model in requests.models])
refusal: Final = next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None)
if refusal is not None:
return _RefuseClientFrame(refusal=refusal)
if requests.unverifiable_model_refusal is not None and isinstance(
probe := await gate(websocket, STORED_PROMPT_PROBE_MODEL), WebsocketFrameModelRefused
):
return _RefuseClientFrame(refusal=_unverifiable_model_probe_refusal(probe, requests.unverifiable_model_refusal))
allowed: Final = tuple(verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelAllowed))
return _ForwardClientFrame(
starts_turn=requests.starts_turn,
lane=requests.lane,
turn=(
_TurnAuthorization(requested_model=requests.turn_model, access_groups=allowed[0].access_groups)
if requests.turn_model is not None and allowed
else None
),
)
def _refusal_close_reason(message: str) -> str:
first_sentence: Final = _FIRST_SENTENCE.match(message)
return truncated_close_reason(message if first_sentence is None else first_sentence.group(0))
async def refuse_websocket(websocket: WebSocket, refusal: WebsocketFrameModelRefused) -> None:
await websocket.send_text(realtime_error_event(refusal.message, error_type=refusal.error_type))
await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message))
def _stamp_billed_turn(kwargs: Mapping[str, object], turn: _TurnAuthorization) -> None:
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params"))
metadata: Final = None if litellm_params is None else litellm_params.get("metadata")
if not isinstance(metadata, MutableMapping):
return
metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(turn.access_groups)
if turn.requested_model is not None:
metadata["model_group"] = turn.requested_model
def _opening_model(target: str) -> str | None:
return next(iter(parse_qs(urlparse(target).query).get("model", ())), None)
class _WebsocketTurnLedger:
def __init__(self, opening: _TurnAuthorization, started: datetime) -> None:
self.opening = opening
self.queued_turns: Mapping[str | None, tuple[_TurnAuthorization | None, ...]] = MappingProxyType({})
self.session_opened = False
self.started = started
self.logged_turns = 0
self.unstarted_turns = 0
def current(self, lane: str | None) -> _TurnAuthorization:
oldest: Final = next(iter(self.queued_turns.get(lane, ())), None)
return self.opening if oldest is None else oldest
def queue_turn(self, lane: str | None, turn: _TurnAuthorization | None) -> None:
self.unstarted_turns += 1
if not self.session_opened:
self.queued_turns = MappingProxyType({**self.queued_turns, lane: (*self.queued_turns.get(lane, ()), turn)})
def open_session(self) -> None:
self.session_opened = True
self.queued_turns = MappingProxyType({})
def close_turn(self, lane: str | None, end_time: datetime) -> None:
self.logged_turns += 1
self.started = end_time
self._forget_turn(lane, 0)
def start_turn(self) -> None:
self.unstarted_turns = max(0, self.unstarted_turns - 1)
def drop_refused_turn(self, lane: str | None, lane_is_running_a_turn: bool) -> None:
self.start_turn()
self._forget_turn(lane, 1 if lane_is_running_a_turn else 0)
def _forget_turn(self, lane: str | None, position: int) -> None:
turns: Final = self.queued_turns.get(lane, ())
self.queued_turns = MappingProxyType({**self.queued_turns, lane: (*turns[:position], *turns[position + 1 :])})
def _session_end_needs_a_row(
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None,
unlogged_frames: Sequence[Mapping[str, object]],
ledger: _WebsocketTurnLedger,
) -> bool:
if upstream_frame_kind is None:
return True
kinds: Final = frozenset(map(upstream_frame_kind, unlogged_frames))
return "starts_turn" in kinds or (ledger.logged_turns == 0 and bool(kinds - {"opens_session"}))
def _turn_in_flight(
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None, unlogged_frames: Sequence[Mapping[str, object]]
) -> bool:
return upstream_frame_kind is not None and any(
upstream_frame_kind(frame) == "starts_turn" for frame in unlogged_frames
)
def _turn_owed(
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None,
unlogged_frames: Sequence[Mapping[str, object]],
ledger: _WebsocketTurnLedger,
) -> bool:
return _turn_in_flight(upstream_frame_kind, unlogged_frames) or (
upstream_frame_kind is not None and ledger.unstarted_turns > 0
)
def _lanes_in_flight(
upstream_frame_kind: WebsocketUpstreamFrameClassifier, unlogged_frames: Sequence[Mapping[str, object]]
) -> tuple[str | None, ...]:
started: Final = (frame for frame in unlogged_frames if upstream_frame_kind(frame) == "starts_turn")
return tuple(dict.fromkeys(map(_frame_lane, started)))
async def _relay_client_frame(
websocket: WebSocket,
upstream_ws: ClientConnection,
checked_frame: str | bytes,
forwarded_frame: str | bytes,
gate: WebsocketFrameModelGate | None,
ledger: _WebsocketTurnLedger,
) -> _ForwardClientFrame | _RefuseClientFrame:
verdict: Final = await _client_frame_verdict(websocket, checked_frame, gate)
match verdict:
case _RefuseClientFrame(refusal=refusal):
await refuse_websocket(websocket, refusal)
case _ForwardClientFrame(starts_turn=True, turn=turn, lane=lane):
ledger.queue_turn(lane, turn)
await upstream_ws.send(forwarded_frame)
case _ForwardClientFrame():
await upstream_ws.send(forwarded_frame)
return verdict
_EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"})
WALKOUT_TURN_DRAIN_SECONDS: Final = 600.0
def _metered_client_text(meter: WebSocketSessionMeter | None, client_text: str) -> int:
return 0 if meter is None else meter.measure(client_text)
def _log_session_usage(logging_obj: LiteLLMLoggingObj, usage: WebSocketSessionUsage) -> None:
logging_obj.model = usage.model # rebind-ok: the spend logger reads model and cost off the shared logging object
logging_obj.model_call_details["model"] = usage.model # rebind-ok: same shared logging object
@ -2654,6 +3019,9 @@ async def websocket_passthrough_request(
cost_per_request: float | None = None,
accept_websocket: bool = True,
setup_model_rewriter: Callable[[str], str] | None = None,
client_frame_model_gate: WebsocketFrameModelGate | None = None,
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None,
walkout_turn_drain_seconds: float = WALKOUT_TURN_DRAIN_SECONDS,
builtin_provider_route: bool = False,
model: str | None = None,
meter: WebSocketSessionMeter | None = None,
@ -2670,6 +3038,11 @@ async def websocket_passthrough_request(
endpoint: The endpoint path (for logging purposes)
cost_per_request: Optional field - cost per request to the target endpoint
setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream
client_frame_model_gate: Optional per-frame check of every model a client frame names
upstream_frame_kind: Optional sort of upstream frames into the one that ends a billable turn and the ones
that open the session, so each turn logs its own spend and keeps the session's provider-named model
walkout_turn_drain_seconds: How long to keep reading a turn the provider is still running after the client
left or was refused, so the provider's usage for it still reaches spend
builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider
model: The model the caller named before connecting, logged on a failure before the setup frame arrives
meter: Prices the session from the client text frames delivered upstream, for an upstream that reports no usage
@ -2709,26 +3082,31 @@ async def websocket_passthrough_request(
proxy_config=proxy_config,
route_description="websocket_passthrough",
)
logging_obj: Final = Logging(
model="unknown",
messages=[{"role": "user", "content": "WebSocket connection"}],
stream=True, # WebSockets are inherently streaming
call_type="pass_through_endpoint",
start_time=start_time,
litellm_call_id=litellm_call_id,
function_id="websocket_passthrough",
dynamic_success_callbacks=team_callbacks.success_callbacks,
dynamic_failure_callbacks=team_callbacks.failure_callbacks,
kwargs=team_callbacks.logging_kwargs,
)
# Create passthrough logging payload
passthrough_logging_payload: Final = PassthroughStandardLoggingPayload(
url=target,
request_body={}, # WebSocket doesn't have a traditional request body
request_method="WEBSOCKET",
cost_per_request=cost_per_request,
)
def new_call_logging(call_start: datetime, call_id: str) -> LiteLLMLoggingObj:
return Logging(
model="unknown",
messages=[{"role": "user", "content": "WebSocket connection"}],
stream=True,
call_type="pass_through_endpoint",
start_time=call_start,
litellm_call_id=call_id,
function_id="websocket_passthrough",
dynamic_success_callbacks=team_callbacks.success_callbacks,
dynamic_failure_callbacks=team_callbacks.failure_callbacks,
kwargs=team_callbacks.logging_kwargs,
)
def new_call_payload() -> PassthroughStandardLoggingPayload:
return PassthroughStandardLoggingPayload(
url=target,
request_body={},
request_method="WEBSOCKET",
cost_per_request=cost_per_request,
)
logging_obj: Final = new_call_logging(start_time, litellm_call_id)
passthrough_logging_payload: Final = new_call_payload()
# Create a dummy request object for WebSocket connections to maintain compatibility
# with the existing _init_kwargs_for_pass_through_endpoint function
@ -2750,38 +3128,136 @@ async def websocket_passthrough_request(
)
# Initialize kwargs for logging using the same pattern as HTTP passthrough
kwargs: Final = HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint(
user_api_key_dict=user_api_key_dict,
_parsed_body={}, # WebSocket doesn't have a traditional request body
passthrough_logging_payload=passthrough_logging_payload,
litellm_call_id=litellm_call_id,
request=dummy_request,
logging_obj=logging_obj,
def init_call_kwargs(
token: UserAPIKeyAuth,
call_payload: PassthroughStandardLoggingPayload,
call_id: str,
call_logging: LiteLLMLoggingObj,
) -> dict[str, object]: # mutable-ok: the Vertex relay stamps the model from setup frames into these kwargs
return HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint(
user_api_key_dict=token,
_parsed_body={}, # WebSocket doesn't have a traditional request body
passthrough_logging_payload=call_payload,
litellm_call_id=call_id,
request=dummy_request,
logging_obj=call_logging,
)
kwargs: Final = init_call_kwargs(user_api_key_dict, passthrough_logging_payload, litellm_call_id, logging_obj)
def start_call_logging(call_logging: LiteLLMLoggingObj, call_kwargs: Mapping[str, object]) -> None:
call_logging.update_environment_variables(
model="unknown",
user="unknown",
optional_params={},
litellm_params=dict(as_str_mapping(call_kwargs.get("litellm_params")) or {}),
call_type="pass_through_endpoint",
)
call_logging.pre_call(
input=[{"role": "user", "content": "WebSocket connection"}],
api_key="",
additional_args={
"complete_input_dict": {},
"api_base": target,
"headers": upstream_headers,
},
)
start_call_logging(logging_obj, kwargs)
class MockWebSocketRequest:
def __init__(self, target_url: str):
self.method = "WEBSOCKET"
self.url = target_url
class MockWebSocketResponse:
def __init__(self, target_url: str):
self.status_code = 200
self.text = "WebSocket connection successful"
self.headers: dict[str, str] = {}
self.request = MockWebSocketRequest(target_url)
ledger: Final = _WebsocketTurnLedger(
opening=_TurnAuthorization(
requested_model=_opening_model(target),
access_groups=tuple(user_api_key_dict.matched_model_access_groups or ()),
),
started=start_time,
)
# Update logging environment variables
logging_obj.update_environment_variables(
model="unknown",
user="unknown",
optional_params={},
litellm_params=dict(kwargs.get("litellm_params", {})),
call_type="pass_through_endpoint",
)
logging_obj.model_call_details["litellm_call_id"] = litellm_call_id
def log_frames(
call_logging: LiteLLMLoggingObj,
call_payload: PassthroughStandardLoggingPayload,
call_kwargs: Mapping[str, object],
end_time: datetime,
lane: str | None,
in_row: Callable[[Mapping[str, object]], bool],
) -> None:
all_frames: Final = list(websocket_messages)
frames: Final = [frame for frame in all_frames if in_row(frame)]
websocket_messages[:] = (
()
if upstream_frame_kind is None
else (
frame
for frame in all_frames
if not in_row(frame) or upstream_frame_kind(frame) in _SESSION_FRAMES_EVERY_ROW_CARRIES
)
)
logged_payload: Final[PassthroughStandardLoggingPayload] = {
**call_payload,
"response_body": frames,
"end_time": end_time,
}
if client_frame_model_gate is not None:
_stamp_billed_turn(call_kwargs, ledger.current(lane))
success_kwargs: Final = {
key: value for key, value in call_kwargs.items() if key not in _EXPLICIT_SUCCESS_KWARGS
}
custom_llm_provider: Final = call_kwargs.get("custom_llm_provider")
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=pass_through_endpoint_logging.pass_through_async_success_handler(
httpx_response=MockWebSocketResponse(target),
response_body=frames,
url_route=endpoint or "",
result="websocket_connection_successful",
start_time=ledger.started,
end_time=end_time,
logging_obj=call_logging,
cache_hit=False,
request_body={},
passthrough_logging_payload=logged_payload,
custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None,
finalize_before_dispatch=False,
**success_kwargs,
)
)
bind_budget_reservation_to_callbacks(call_logging.litellm_params)
# Pre-call logging
logging_obj.pre_call(
input=[{"role": "user", "content": "WebSocket connection"}],
api_key="",
additional_args={
"complete_input_dict": {},
"api_base": target,
"headers": upstream_headers,
},
)
def log_turn(lane: str | None, in_row: Callable[[Mapping[str, object]], bool]) -> None:
end_time: Final = datetime.now()
if upstream_frame_kind is None:
log_frames(logging_obj, passthrough_logging_payload, kwargs, end_time, lane, in_row)
else:
turn_call_id: Final = str(uuid.uuid4())
turn_logging: Final = new_call_logging(ledger.started, turn_call_id)
turn_payload: Final = new_call_payload()
turn_kwargs: Final = init_call_kwargs(
user_api_key_dict
if ledger.logged_turns == 0
else user_api_key_dict.model_copy(update={"budget_reservation": None}),
turn_payload,
turn_call_id,
turn_logging,
)
start_call_logging(turn_logging, turn_kwargs)
log_frames(turn_logging, turn_payload, turn_kwargs, end_time, lane, in_row)
ledger.close_turn(lane, end_time)
### CALL HOOKS ### - modify incoming data / reject request before calling the model
websocket_hook_data: Final = guardrail_request_data_with_streaming(MappingProxyType({}), is_streaming=True)
websocket_hook_data: Final = guardrail_request_data_with_streaming(
MappingProxyType({"litellm_call_id": litellm_call_id}), is_streaming=True
)
await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_dict,
data=websocket_hook_data,
@ -2808,7 +3284,6 @@ async def websocket_passthrough_request(
message = await websocket.receive()
message_type = message.get("type")
if message_type == "websocket.disconnect":
await upstream_ws.close()
break
text_data: str | None = message.get("text")
@ -2871,11 +3346,29 @@ async def websocket_passthrough_request(
)
# Not a JSON message or doesn't contain setup data
await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter))
if meter is not None:
metered_total += meter.measure(text_data)
if isinstance(
await _relay_client_frame(
websocket,
upstream_ws,
text_data,
_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter),
client_frame_model_gate,
ledger,
),
_RefuseClientFrame,
):
break
metered_total += _metered_client_text(meter, text_data)
elif bytes_data is not None:
await upstream_ws.send(bytes_data)
if isinstance(
await _relay_client_frame(
websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate, ledger
),
_RefuseClientFrame,
):
break
if not _turn_owed(upstream_frame_kind, websocket_messages, ledger):
await upstream_ws.close()
except asyncio.CancelledError:
raise
except Exception:
@ -2902,11 +3395,7 @@ async def websocket_passthrough_request(
is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint)
json_frame_ordinal: Final = count()
async def relay_upstream_frame(upstream_message: str | bytes) -> None:
if isinstance(upstream_message, bytes):
await websocket.send_bytes(upstream_message)
else:
await websocket.send_text(upstream_message)
def record_upstream_frame(upstream_message: str | bytes) -> None:
message_data: Final = _json_object_frame(upstream_message)
if message_data is None:
return
@ -2914,6 +3403,43 @@ async def websocket_passthrough_request(
_extract_vertex_live_model_from_setup_response(message_data)
return
websocket_messages.append(message_data)
match None if upstream_frame_kind is None else upstream_frame_kind(message_data):
case "opens_session":
ledger.open_session()
case "starts_turn":
ledger.start_turn()
case "updates_session" if upstream_frame_kind is not None:
websocket_messages[:] = _without_earlier_session_updates(
upstream_frame_kind, websocket_messages, message_data
)
case "ends_turn" if upstream_frame_kind is not None:
if _frame_lane(message_data) not in _lanes_in_flight(upstream_frame_kind, websocket_messages):
ledger.start_turn()
log_turn(_frame_lane(message_data), _frames_of_lane(_frame_lane(message_data)))
case "bills_usage" if not _turn_in_flight(upstream_frame_kind, websocket_messages):
log_turn(_frame_lane(message_data), _frames_of_lane(_frame_lane(message_data)))
case "refuses_turn" if upstream_frame_kind is not None:
lane: Final = _frame_lane(message_data)
ledger.drop_refused_turn(
lane, lane in _lanes_in_flight(upstream_frame_kind, websocket_messages)
)
case _:
pass
async def relay_upstream_frame(upstream_message: str | bytes) -> None:
record_upstream_frame(upstream_message)
if not _client_socket_is_open(websocket):
if not _turn_owed(upstream_frame_kind, websocket_messages, ledger):
await upstream_ws.close()
return
try:
if isinstance(upstream_message, bytes):
await websocket.send_bytes(upstream_message)
else:
await websocket.send_text(upstream_message)
except WebSocketDisconnect:
if not _turn_owed(upstream_frame_kind, websocket_messages, ledger):
raise
async def forward_upstream_to_client() -> Close | None:
try:
@ -2922,6 +3448,11 @@ async def websocket_passthrough_request(
except (ConnectionClosedOK, ConnectionClosedError) as e:
verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e)
return e.rcvd
except WebSocketDisconnect:
verbose_proxy_logger.debug(
"WebSocket passthrough (%s): client left while upstream was still sending", endpoint
)
return None
except asyncio.CancelledError:
verbose_proxy_logger.debug("asyncio.CancelledError in forward_upstream_to_client")
raise
@ -2933,12 +3464,15 @@ async def websocket_passthrough_request(
raise
# Create tasks for bidirectional message forwarding
tasks: Final = [
asyncio.create_task(forward_client_to_upstream()),
asyncio.create_task(forward_upstream_to_client()),
]
client_task: Final = asyncio.create_task(forward_client_to_upstream())
upstream_task: Final = asyncio.create_task(forward_upstream_to_client())
tasks: Final = (client_task, upstream_task)
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
if client_task.done() and _turn_owed(upstream_frame_kind, websocket_messages, ledger):
await asyncio.wait((upstream_task,), timeout=walkout_turn_drain_seconds)
done: Final = tuple(task for task in tasks if task.done())
pending: Final = tuple(task for task in tasks if not task.done())
# Cancel remaining tasks
for task in pending:
@ -2961,64 +3495,17 @@ async def websocket_passthrough_request(
reason=truncated_close_reason(upstream_close.reason),
)
end_time: Final = datetime.now()
# Update passthrough logging payload with response data
passthrough_logging_payload["response_body"] = websocket_messages
passthrough_logging_payload["end_time"] = end_time
# Remove logging_obj from kwargs to avoid duplicate keyword argument
success_kwargs: Final = kwargs.copy()
success_kwargs.pop("logging_obj", None)
# # Add user authentication context for database logging
# if user_api_key_dict:
# success_kwargs.setdefault('litellm_params', {})
# success_kwargs['litellm_params'].update({
# 'proxy_server_request': {
# 'body': {
# 'user': user_api_key_dict.user_id,
# 'team_id': user_api_key_dict.team_id,
# 'end_user_id': user_api_key_dict.end_user_id,
# }
# }
# })
# # Also add the user_api_key for direct access
# success_kwargs['user_api_key'] = user_api_key_dict.api_key
# Create a dummy httpx.Response for WebSocket connections
class MockWebSocketResponse:
def __init__(self, target_url: str):
self.status_code = 200
self.text = "WebSocket connection successful"
self.headers: dict[str, str] = {}
self.request = MockWebSocketRequest(target_url)
class MockWebSocketRequest:
def __init__(self, target_url: str):
self.method = "WEBSOCKET"
self.url = target_url
mock_response: Final = MockWebSocketResponse(target)
if meter is not None:
_log_session_usage(logging_obj, meter.price(metered_total))
# Use the same success handler as HTTP passthrough endpoints
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=pass_through_endpoint_logging.pass_through_async_success_handler(
httpx_response=mock_response,
response_body=websocket_messages,
url_route=endpoint or "",
result="websocket_connection_successful",
start_time=start_time,
end_time=end_time,
logging_obj=logging_obj,
cache_hit=False,
request_body={},
**success_kwargs,
if _session_end_needs_a_row(upstream_frame_kind, websocket_messages, ledger):
lanes_in_flight: Final = (
() if upstream_frame_kind is None else _lanes_in_flight(upstream_frame_kind, websocket_messages)
)
)
bind_budget_reservation_to_callbacks(logging_obj.litellm_params)
*earlier_lanes, last_lane = lanes_in_flight or (None,)
for lane in earlier_lanes:
log_turn(lane, _frames_of_lane(lane))
log_turn(last_lane, _every_frame)
# Call the proxy logging success hook
if proxy_logging_obj:
@ -3076,6 +3563,7 @@ async def websocket_passthrough_request(
if _client_socket_is_open(websocket):
await websocket.close(code=1011, reason="WebSocket passthrough error")
finally:
await proxy_logging_obj.arelease_responses_websocket_slot(user_api_key_dict)
if _client_socket_is_open(websocket):
await websocket.close()

70
litellm/proxy/pass_through_endpoints/success_handler.py
View file

@ -1,5 +1,6 @@
import asyncio
import json
from collections.abc import Mapping, Sequence
from contextlib import suppress
from datetime import datetime
from types import MappingProxyType
@ -13,7 +14,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.constants import AZURE_SPEECH_CUSTOM_LLM_PROVIDER
from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import PassThroughEndpointLoggingResultValues
from litellm.proxy._types import PassThroughEndpointLoggingResultValues, PassThroughEndpointLoggingTypedDict
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
)
@ -70,6 +71,14 @@ def _safe_response_text(httpx_response: httpx.Response) -> str:
return ""
def _websocket_frames(
response_body: Mapping[str, object] | Sequence[Mapping[str, object]] | None,
) -> tuple[Mapping[str, object], ...]:
return tuple(
message for message in (response_body if isinstance(response_body, list) else ()) if isinstance(message, dict)
)
class PassThroughEndpointLogging:
def __init__(
self,
@ -120,6 +129,7 @@ class PassThroughEndpointLogging:
# Vertex AI Live API WebSocket
self.TRACKED_VERTEX_AI_LIVE_ROUTES = ["/vertex_ai/live"]
self.TRACKED_OPENAI_WEBSOCKET_ROUTES: Final = ("/openai/", "/openai_passthrough/")
@property
def _log_dispatch(self) -> PassThroughLogDispatch:
@ -422,21 +432,18 @@ class PassThroughEndpointLogging:
standard_logging_response_object = vertex_ai_live_handler_result["result"]
kwargs = vertex_ai_live_handler_result["kwargs"]
elif DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route):
deepgram_handler_result: Final = (
DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
websocket_messages=tuple(
message
for message in (response_body if isinstance(response_body, list) else ())
if isinstance(message, dict)
),
logging_obj=logging_obj,
upstream_url=str(httpx_response.request.url),
kwargs=kwargs,
)
elif (
websocket_relay_result := self._websocket_relay_logging_result(
url_route=url_route,
websocket_frames=_websocket_frames(response_body),
logging_obj=logging_obj,
httpx_response=httpx_response,
start_time=start_time,
kwargs=kwargs,
)
standard_logging_response_object = deepgram_handler_result["result"] # rebind-ok: elif-chain
kwargs = deepgram_handler_result["kwargs"] # rebind-ok: elif-chain contract
) is not None:
standard_logging_response_object = websocket_relay_result["result"] # rebind-ok: elif-chain
kwargs = websocket_relay_result["kwargs"] # rebind-ok: elif-chain contract
elif FalAIPassthroughLoggingHandler.is_fal_ai_route(url_route, custom_llm_provider):
fal_ai_handler_result: Final = FalAIPassthroughLoggingHandler().fal_ai_passthrough_handler(
response_body=response_body if isinstance(response_body, dict) else MappingProxyType({}),
@ -452,6 +459,36 @@ class PassThroughEndpointLogging:
return_dict["kwargs"] = kwargs
return return_dict
def _websocket_relay_logging_result(
self,
url_route: str,
websocket_frames: tuple[Mapping[str, object], ...],
logging_obj: LiteLLMLoggingObj,
httpx_response: httpx.Response,
start_time: datetime,
kwargs: Mapping[str, object],
) -> PassThroughEndpointLoggingTypedDict | None:
if self.is_openai_websocket_route(url_route):
from .llm_provider_handlers.openai_websocket_passthrough_logging_handler import (
OpenAIWebsocketPassthroughLoggingHandler,
)
return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler(
websocket_messages=websocket_frames,
logging_obj=logging_obj,
url_route=url_route,
start_time=start_time,
kwargs=kwargs,
)
if DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route):
return DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
websocket_messages=websocket_frames,
logging_obj=logging_obj,
upstream_url=str(httpx_response.request.url),
kwargs=kwargs,
)
return None
async def pass_through_async_success_handler(
self,
httpx_response: httpx.Response,
@ -625,6 +662,9 @@ class PassThroughEndpointLogging:
return True
return False
def is_openai_websocket_route(self, url_route: str) -> bool:
return url_route.startswith(self.TRACKED_OPENAI_WEBSOCKET_ROUTES)
def is_cursor_route(self, url_route: str, custom_llm_provider: str | None = None):
"""Check if the URL route is a Cursor Cloud Agents API route."""
if custom_llm_provider == "cursor":

25
litellm/proxy/proxy_cli.py
View file

@ -1193,7 +1193,7 @@ class ProxyInitializationHelpers:
"--reload",
is_flag=True,
default=False,
help="Enable uvicorn hot reload (dev only). Also reloads when the --config YAML file changes. Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.",
help="Enable uvicorn hot reload (dev only), including every --num_workers worker on uvicorn>=0.50.0 outside Windows. Also reloads when the --config YAML or .env changes. Incompatible with --run_gunicorn, --run_hypercorn, and --run_granian.",
)
@click.option(
"--prometheus_metrics_port",
@ -1785,9 +1785,28 @@ def run_server(
if num_workers > 1:
start_query_engine_reaper()
if num_workers > 1 or reload:
uvicorn_args.update(
ProxyInitializationHelpers.child_process_uvicorn_overrides(uvicorn_args.get("log_config"))
worker_overrides: Final = ProxyInitializationHelpers.child_process_uvicorn_overrides(
uvicorn_args.get("log_config")
)
if reload and num_workers > 1:
from litellm.proxy.dev_reload import (
RELOAD_OPTIONS,
multi_worker_reload_supported,
run_multi_worker_reload,
)
if multi_worker_reload_supported():
run_multi_worker_reload(
RELOAD_OPTIONS.validate_python(uvicorn_args),
workers=TypeAdapter(int).validate_python(num_workers),
worker_overrides=worker_overrides,
)
return
print(
"\033[1;33mLiteLLM Proxy: --reload with --num_workers>1 requires uvicorn>=0.50.0 "
"and an OS other than Windows, so only a single worker reloads.\033[0m"
)
uvicorn_args.update(worker_overrides)
uvicorn.run(
**uvicorn_args,
workers=num_workers,

340
litellm/proxy/proxy_server.py
View file

@ -33,7 +33,7 @@ from collections.abc import (
from collections.abc import (
Set as AbstractSet,
)
from dataclasses import dataclass
from dataclasses import dataclass, replace
from datetime import datetime, timedelta, timezone
from itertools import chain
from types import MappingProxyType, TracebackType, UnionType
@ -61,6 +61,8 @@ import websockets
import websockets.exceptions
from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, ValidationError, with_config
from pydantic.fields import FieldInfo, PydanticUndefined
from pydantic.main import IncEx
from pydantic_core import PydanticSerializationError, to_jsonable_python
from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never
from litellm._uuid import uuid
@ -197,6 +199,7 @@ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import (
snapshot_tuning_baselines,
tuning_limit_violation,
)
from litellm.rust_bridge.catalog import Target
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import (
PRICING_OVERRIDES_KEY,
@ -387,15 +390,23 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
mask_sensitive_keys,
)
from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact
from litellm.llms.anthropic.messages.context_management.inference_view import (
InferenceRequestView,
InferenceRequestViewError,
count_context_edits,
inference_request_view,
)
from litellm.llms.anthropic.messages.convertible_blocks import (
deployment_drops_params,
drop_non_string_tool_result_texts,
malformed_tool_result_count_error,
)
from litellm.llms.custom_httpx.aiohttp_pool_diagnostics import build_pool_wait_trace_config
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.openai.live import parse_live_session_start
from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.messages.dispatch import messages_target
from litellm.proxy._lazy_features import attach_lazy_features, reserve_lazy_slot
from litellm.proxy._types import *
from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling
@ -1209,7 +1220,7 @@ ui_message += "\n\n💸 [```LiteLLM Model Cost Map```](https://models.litellm.ai
ui_message += f"\n\n🔎 [```LiteLLM Model Hub```]({model_hub_link}). See available models on the proxy. [**Docs**](https://docs.litellm.ai/docs/proxy/ai_hub)"
custom_swagger_message: Final = (
"[**Customize Swagger Docs**](https://docs.litellm.ai/docs/proxy/enterprise#swagger-docs---custom-routes--branding)"
"[**Customize Swagger Docs**](https://docs.litellm.ai/docs/enterprise#operations--branding)"
)
### CUSTOM BRANDING [ENTERPRISE FEATURE] ###
@ -1460,7 +1471,11 @@ async def initialize_shared_aiohttp_session() -> "ClientSession | None":
connector_kwargs["socket_factory"] = socket_factory
connector: Final = TCPConnector(**connector_kwargs)
session: Final = ClientSession(connector=connector, cookie_jar=DummyCookieJar())
session: Final = ClientSession(
connector=connector,
cookie_jar=DummyCookieJar(),
trace_configs=[build_pool_wait_trace_config()],
)
verbose_proxy_logger.info(
"SESSION REUSE: Created shared aiohttp session for connection pooling (ID: %s, limit=%s, limit_per_host=%s)",
@ -10181,6 +10196,8 @@ def _event_with_restamped_nested_model(event: Mapping[str, object], target_model
_SSE_DATA_FIELD: Final = b"data:"
_SSE_NESTED_RESPONSE_MARKER: Final = b'"response"'
_SSE_EVENT_OBJECT: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
_JSON_OBJECT: Final = TypeAdapter(dict[object, object])
_JSON_ARRAY: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...])
def _restamp_nested_response_model_in_sse_frame(frame: bytes, target_model: str) -> bytes:
@ -10206,91 +10223,64 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) -
return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending
def _fast_serialize_simple_model_response_stream(
chunk: ModelResponseStream,
) -> bytes | None:
"""
Serialize the common OpenAI text streaming chunk without the full Pydantic
serializer. Fall back for richer chunks so tool calls, logprobs, usage, and
provider-specific fields keep the canonical model_dump_json behavior.
"""
if (
getattr(chunk, "provider_specific_fields", None) is not None
or getattr(chunk, "system_fingerprint", None) is not None
or getattr(chunk, "usage", None) is not None
):
return None
choices: Final = getattr(chunk, "choices", None)
if not isinstance(choices, list) or len(choices) != 1:
return None
choice: Final = choices[0]
if getattr(choice, "logprobs", None) is not None or getattr(choice, "enhancements", None) is not None:
return None
delta: Final[object] = getattr(choice, "delta", None)
if delta is None:
return None
unsupported_delta_fields: Final = (
"function_call",
"tool_calls",
"audio",
"images",
"annotations",
"reasoning_content",
"thinking_blocks",
"provider_specific_fields",
"refusal",
)
if any(getattr(delta, field, None) is not None for field in unsupported_delta_fields):
return None
delta_dict: Final[dict] = {}
role: Final = getattr(delta, "role", None)
content: Final = getattr(delta, "content", None)
if role is not None:
delta_dict["role"] = role
if content is not None:
delta_dict["content"] = content
choice_dict: Final = {"index": getattr(choice, "index", 0), "delta": delta_dict}
finish_reason: Final = getattr(choice, "finish_reason", None)
if finish_reason is not None:
choice_dict["finish_reason"] = finish_reason
# Match the canonical ``model_dump_json(exclude_none=True)`` shape — if a
# field is None, omit it entirely rather than emitting ``"key": null``.
# Strict OpenAI-compatible clients reject ``null`` for optional fields like
# ``model``, so diverging here would surface as a client-side regression
# only on the fast path. Fall back to the slow path if a required-looking
# top-level identifier is missing.
model: Final = getattr(chunk, "model", None)
if model is None:
return None
payload: Final[dict] = {
"id": getattr(chunk, "id", None),
"object": getattr(chunk, "object", None),
"created": getattr(chunk, "created", None),
"model": model,
"service_tier": getattr(chunk, "service_tier", None),
"choices": [choice_dict],
}
for top_level_key in ("id", "object", "created", "service_tier"):
if payload[top_level_key] is None:
payload.pop(top_level_key)
return orjson.dumps(payload)
def _json_encodes(value: object) -> bool:
try:
to_jsonable_python(value)
except ValueError:
return False
return True
def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes:
def _unencodable_part(value: object, ancestors: frozenset[int]) -> IncEx | bool:
if id(value) in ancestors:
return True
lineage: Final = ancestors | {id(value)}
if isinstance(value, BaseModel):
return _unencodable_keys(iter(value), lineage) or True
if isinstance(value, dict):
return _unencodable_entries(_JSON_OBJECT.validate_python(value), lineage)
if isinstance(value, (list, tuple)):
return _unencodable_indexes(enumerate(_JSON_ARRAY.validate_python(value)), lineage) or True
return True
def _unencodable_entries(entries: Mapping[object, object], lineage: frozenset[int]) -> IncEx | bool:
by_name: Final = {key: item for key, item in entries.items() if isinstance(key, str)}
if len(by_name) == len(entries):
return _unencodable_keys(by_name.items(), lineage) or True
by_index: Final = {key: item for key, item in entries.items() if isinstance(key, int)}
if len(by_index) == len(entries):
return _unencodable_indexes(by_index.items(), lineage) or True
return True
def _unencodable_keys(items: Iterable[tuple[str, object]], lineage: frozenset[int]) -> Mapping[str, IncEx | bool]:
return {key: _unencodable_part(item, lineage) for key, item in items if not _json_encodes(item)}
def _unencodable_indexes(items: Iterable[tuple[int, object]], lineage: frozenset[int]) -> Mapping[int, IncEx | bool]:
return {index: _unencodable_part(item, lineage) for index, item in items if not _json_encodes(item)}
def _serialize_chat_chunk_leniently(chunk: ModelResponseStream) -> bytes:
exclude: Final = _unencodable_keys(iter(chunk), frozenset({id(chunk)}))
if exclude:
verbose_proxy_logger.warning("Dropping streamed chunk fields that JSON cannot encode: %s", exclude)
jsonable: Final = chunk.model_dump(mode="json", exclude_none=True, exclude_unset=True, exclude=exclude)
return json.dumps(jsonable, ensure_ascii=True, separators=(",", ":")).encode()
def _serialize_chat_chunk(chunk: ModelResponseStream) -> bytes:
try:
return chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode()
except PydanticSerializationError:
return _serialize_chat_chunk_leniently(chunk)
def _serialize_streaming_chunk(chunk: BaseModel) -> bytes:
if isinstance(chunk, ModelResponseStream):
serialized_chunk: Final = _fast_serialize_simple_model_response_stream(chunk)
if serialized_chunk is not None:
return serialized_chunk
return chunk.model_dump_json(exclude_none=True, exclude_unset=True)
return _serialize_chat_chunk(chunk)
return chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode()
async def _apply_streaming_chunk_hooks(
@ -14961,7 +14951,7 @@ async def _try_provider_token_count(
deployment: Mapping[str, object] | None,
request_model: str,
tools: list | None = None,
system: str | None = None,
system: object | None = None,
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
context_management: CountContextManagement | None = None,
@ -15031,6 +15021,7 @@ async def _provider_token_count(
request_model: str,
settings: TokenCountSettings,
local_estimate_allowed: bool,
forwarded_edits: bool = False,
) -> Union["TokenCountResponse", ProviderCountFailure, None]:
"""Count with the provider, using the token_count_model in its region when one is set.
@ -15054,7 +15045,7 @@ async def _provider_token_count(
pinned_count_model: Final = None if count_model is None or count_model.region is None else count_model
region_hint: Final = (
None
if pinned_count_model is None
if pinned_count_model is None or (forwarded_edits and counted.status_code == 400)
else f"This model counts tokens with {pinned_count_model.model} in {pinned_count_model.region}. "
"If that region is not enabled for the AWS account, enable it or set token_count_model in model_info"
)
@ -15316,6 +15307,9 @@ class _TokenCountDeployment(TypedDict, total=False):
_TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment)
_COUNT_TOOLS: Final[TypeAdapter[tuple[Mapping[str, object], ...] | None]] = TypeAdapter(
tuple[Mapping[str, object], ...] | None
)
_ROUTED_DEPLOYMENT: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
@ -15385,12 +15379,15 @@ def _count_context_management(
model, general_settings_view(), user_model if isinstance(user_model, str) else None
)
router_defaults: Final = _router_default_litellm_params(llm_router)
sent_params: Final = tuple(
{**deployment["litellm_params"], **router_defaults, **request_litellm_params}
deployment_params: Final = tuple(
deployment["litellm_params"]
for deployment in llm_router.registry().deployments(
inference_model if isinstance(inference_model, str) else model, source="routable"
)
)
sent_params: Final = tuple(
{**params, **router_defaults, **request_litellm_params} for params in deployment_params or ({},)
)
if any(context_management_explicitly_dropped(params.get("additional_drop_params")) for params in sent_params):
return None
return CountContextManagement(
@ -15400,6 +15397,90 @@ def _count_context_management(
)
_GATEWAY_EDIT_ROUTES: Final = frozenset({Target.CHAT_COMPLETIONS, Target.RESPONSES})
def _messages_route(deployment: _TokenCountDeployment, request_model: str) -> Target:
params: Final = deployment.get("litellm_params", {})
return messages_target(
_non_empty_str(params.get("model")) or request_model,
_non_empty_str(params.get("custom_llm_provider")),
deployment.get("model_info"),
_deployment_api_base(deployment),
)
def _rejected_view(status_code: int, message: str) -> ProxyException:
return ProviderCountFailure(
status_code=status_code, message=message, rejected_before_provider_call=True
).to_proxy_exception()
async def _bridged_request_view(
messages: Sequence[Mapping[str, object]],
system: object,
tools: Sequence[Mapping[str, object]] | None,
messages_format: TokenCountMessagesFormat,
edits: CountContextManagement | None,
deployment: _TokenCountDeployment,
fallback_model: str,
) -> InferenceRequestView | InferenceRequestViewError | None:
view: Final = await inference_request_view(
messages,
system,
messages_format,
count_context_edits(edits, tools, deployment.get("litellm_params", {}), fallback_model),
)
return None if isinstance(view, InferenceRequestView) and not view.messages else view
async def _responses_request_view(
messages: Sequence[Mapping[str, object]],
system: object,
tools: Sequence[Mapping[str, object]] | None,
edits: CountContextManagement,
deployment: _TokenCountDeployment,
fallback_model: str,
) -> InferenceRequestView | InferenceRequestViewError:
from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError
from litellm.llms.anthropic.messages.responses_adapters.handler import count_view_messages
try:
edited: Final = await count_view_messages(
model=_non_empty_str(deployment.get("litellm_params", {}).get("model")) or fallback_model,
messages=messages,
tools=tools,
system=system,
context_management=edits,
)
except AnthropicContextManagementError as error:
if error.status_code >= 500:
return InferenceRequestView(messages=messages, system=system)
return InferenceRequestViewError(error_message=error.message, status_code=error.status_code)
return InferenceRequestView(messages=edited, system=system)
async def _gateway_request_view(
route: Target,
messages: Sequence[Mapping[str, object]],
system: object,
tools: Sequence[Mapping[str, object]] | None,
messages_format: TokenCountMessagesFormat,
edits: CountContextManagement | None,
deployment: _TokenCountDeployment,
fallback_model: str,
) -> InferenceRequestView | InferenceRequestViewError | None:
match route:
case Target.CHAT_COMPLETIONS:
return await _bridged_request_view(
messages, system, tools, messages_format, edits, deployment, fallback_model
)
case Target.RESPONSES if edits is not None:
return await _responses_request_view(messages, system, tools, edits, deployment, fallback_model)
case _:
return None
def _responses_routing_kwargs(responses_request: Mapping[str, object] | None) -> Mapping[str, object]:
previous_response_id: Final = None if responses_request is None else responses_request.get("previous_response_id")
return {"previous_response_id": previous_response_id} if isinstance(previous_response_id, str) else {}
@ -15593,6 +15674,40 @@ async def _count_request_tokens(
)
prompt_only: Final = bool(prompt) and conversation_empty
request_messages: Final = _counted_request_messages(messages, messages_format, count_deployment)
local_count_error: Final = malformed_tool_result_count_error(request_messages, messages_format, count_deployment)
malformed_messages: Final = (
None
if local_count_error is None
else ProxyException(message=local_count_error, type="token_counting_error", param="messages", code=400)
)
context_edits: Final = _count_context_management(context_management, request.model, request_litellm_params)
route: Final = (
_messages_route(deployment_view, request.model)
if request_messages is not None and messages_format == "anthropic"
else None
)
if context_edits is not None and route in _GATEWAY_EDIT_ROUTES and malformed_messages is not None:
raise malformed_messages
edited_request: Final = (
None
if route is None or request_messages is None
else await _gateway_request_view(
route,
request_messages,
system,
_COUNT_TOOLS.validate_python(tools),
messages_format,
context_edits,
deployment_view,
model_to_use,
)
)
if isinstance(edited_request, InferenceRequestViewError):
raise _rejected_view(edited_request.status_code, edited_request.error_message)
inference_messages: Final = (
request_messages if edited_request is None else [dict(message) for message in edited_request.messages]
)
inference_system: Final = system if edited_request is None else edited_request.system
contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents))
prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES
routed_deployment: Final = _ROUTED_DEPLOYMENT.validate_python(deployment)
@ -15632,30 +15747,30 @@ async def _count_request_tokens(
if chat_count is not None
else [{"role": "user", "content": prompt}]
if prompt_as_message
else request_messages
else inference_messages
),
contents=None if prompt_as_message else contents,
deployment=count_deployment,
request_model=request.model,
tools=tools if chat_count is None else [dict(tool) for tool in chat_count.tools or ()] or None,
system=system,
system=inference_system,
messages_format=messages_format,
output_format=output_format,
context_management=_count_context_management(context_management, request.model, request_litellm_params),
context_management=None if edited_request is not None else context_edits,
chat_params=MappingProxyType({}) if chat_count is None else chat_count.params,
),
model_to_use=model_to_use,
request_model=request.model,
settings=settings,
local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable,
forwarded_edits=context_edits is not None and edited_request is None,
)
)
if isinstance(ladder, TokenCountResponse):
return ladder
local_count_error: Final = malformed_tool_result_count_error(request_messages, messages_format, count_deployment)
if local_count_error is not None:
raise ProxyException(message=local_count_error, type="token_counting_error", param="messages", code=400)
if malformed_messages is not None:
raise malformed_messages
# Check if token counter is disabled before fallback
if litellm.disable_token_counter is True:
@ -15666,6 +15781,27 @@ async def _count_request_tokens(
code=503,
)
local_view: Final = (
None
if route is not Target.MESSAGES or context_edits is None or request_messages is None
else await _bridged_request_view(
request_messages,
system,
_COUNT_TOOLS.validate_python(tools),
messages_format,
replace(context_edits, drop_invalid_edits=False),
deployment_view,
model_to_use,
)
)
if isinstance(local_view, InferenceRequestViewError):
raise _rejected_view(local_view.status_code, local_view.error_message)
counted_view: Final = edited_request if local_view is None else local_view
local_messages: Final = (
request_messages if counted_view is None else [dict(message) for message in counted_view.messages]
)
local_system: Final = system if counted_view is None else counted_view.system
# Default LiteLLM token counting
custom_tokenizer: CustomHuggingfaceTokenizer | None = None
if model_info is not None:
@ -15673,11 +15809,11 @@ async def _count_request_tokens(
CustomHuggingfaceTokenizer | None,
model_info.get("custom_tokenizer", None),
)
system_message: Final = _system_message(system)
system_message: Final = _system_message(local_system)
prompt_turn: Final = prompt_only and not (system_message is None and not tools)
typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes
Sequence[AllMessageValues] | None,
[{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else request_messages),
[{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else local_messages),
)
counted_messages: Final = (
typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages)
@ -15688,7 +15824,15 @@ async def _count_request_tokens(
framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0
kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None
exact: Final = (
await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, counted_tools)
await _kimi_k3_count(
request
if counted_view is None
else request.model_copy(update={"messages": local_messages, "system": local_system}),
messages_format,
model_to_use,
kimi_k3_messages,
counted_tools,
)
if kimi_k3_messages is not None and custom_tokenizer is None
else None
)

2
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -432,7 +432,7 @@ async def invalidate_budget_reservation_counters(
async def release_or_invalidate_budget_reservation(
budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict
budget_reservation: dict[str, object] | None, # mutable-ok: stamps finalized on the caller's reservation
) -> None:
"""Reconcile a still-open reservation on a terminal path that settles no cost.

88
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -16,6 +16,7 @@ from typing import (
from urllib.parse import urlparse
from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile
from prisma.errors import PrismaError
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model, field_validator
from pydantic.fields import FieldInfo, PydanticUndefined
from typing_extensions import NotRequired, ReadOnly, TypedDict
@ -26,6 +27,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
from litellm.proxy.common_utils.validation_error_body import public_validation_errors
from litellm.proxy.config_resolvers import FieldSource, SettingsStore, source_for
from litellm.proxy.config_resolvers.settings_store import ConfigOwnedKeyError
@ -1247,6 +1249,37 @@ async def get_sso_settings():
}
_ENVIRONMENT_VARIABLES_PARAM: Final = TypeAdapter(dict[str, object])
def _parse_environment_variables_param(param_value: str | Mapping[str, object] | None) -> Mapping[str, object]:
if isinstance(param_value, str):
return _ENVIRONMENT_VARIABLES_PARAM.validate_python(json.loads(param_value))
return param_value or {}
async def _load_environment_variables_param(
repo: ConfigRepository,
) -> Mapping[str, object] | None | PrismaError | ValueError:
try:
entry: Final = await _config_param_db(repo).find_unique(where={"param_name": "environment_variables"})
param_value: Final = None if entry is None else _parse_environment_variables_param(entry.param_value)
except (PrismaError, ValueError) as e:
return e
else:
return param_value
def _decrypted_sso_fields(environment_variables: Mapping[str, object]) -> Mapping[str, str | None]:
return MappingProxyType(
{
field_name: decrypt_value_helper(value=value, key=env_var, return_original_value=True)
for field_name, env_var in SSO_FIELD_ENV_VARS.items()
if isinstance(value := environment_variables.get(env_var), str)
}
)
@router.patch(
"/update/sso_settings",
tags=["SSO Settings"],
@ -1310,9 +1343,24 @@ async def update_sso_settings(
if "general_settings" not in config:
config["general_settings"] = {}
# Update environment variables in config and in memory
sso_data: Final = sso_config.model_dump()
for field_name, value in sso_data.items():
# PATCH semantics: omitted fields keep their stored value, an explicit null clears it. SSO values that
# still live only in the environment_variables row are carried over, since that row is stripped below
environment_variables_param: Final = await _load_environment_variables_param(
ConfigRepository(prisma_client, use_writer=True)
)
legacy_env_values: Final[Mapping[str, str | None]] = (
_decrypted_sso_fields(environment_variables_param)
if isinstance(environment_variables_param, Mapping)
else MappingProxyType({})
)
stored_values: Final = {k: v for k, v in (before_sso_data or {}).items() if v is not None}
fallback_values: Final = {**legacy_env_values, **stored_values}
submitted: Final = sso_config.model_dump(include=sso_config.model_fields_set)
sso_data: Final = {
field_name: submitted[field_name] if field_name in submitted else fallback_values.get(field_name)
for field_name in SSOConfig.model_fields
}
for field_name, value in submitted.items():
if field_name in SSO_FIELD_ENV_VARS:
env_var_name = SSO_FIELD_ENV_VARS[field_name]
if value:
@ -1348,25 +1396,17 @@ async def update_sso_settings(
)
)
# Remove SSO-related env vars from config.environment_variables
try:
env_var_entry: Final = await _config_param_db(ConfigRepository(prisma_client, use_writer=True)).find_unique(
where={"param_name": "environment_variables"}
if isinstance(environment_variables_param, (PrismaError, ValueError)):
raise HTTPException(
status_code=500,
detail={"error": f"Error updating environment_variables: {environment_variables_param}"},
)
# If no environment_variables entry exists, nothing to clean up
if env_var_entry is not None:
if env_var_entry.param_value is not None:
if isinstance(env_var_entry.param_value, str):
environment_variables: Mapping[str, object] = json.loads(env_var_entry.param_value)
else:
environment_variables = dict(env_var_entry.param_value)
else:
environment_variables = {}
if environment_variables_param is not None:
try:
env_vars_to_remove: Final = set(SSO_FIELD_ENV_VARS.values())
filtered_env_vars: Final = {
key: value for key, value in environment_variables.items() if key not in env_vars_to_remove
key: value for key, value in environment_variables_param.items() if key not in env_vars_to_remove
}
await _config_param_db(ConfigRepository(prisma_client)).update(
@ -1376,16 +1416,16 @@ async def update_sso_settings(
},
)
await invalidate_config_param("environment_variables")
except Exception as e:
raise HTTPException(
status_code=500,
detail={"error": f"Error updating environment_variables: {e}"},
)
except Exception as e:
raise HTTPException(
status_code=500,
detail={"error": f"Error updating environment_variables: {e}"},
)
return {
"message": "SSO settings updated successfully",
"status": "success",
"settings": sso_data,
"settings": mask_sensitive_keys(sso_data, set(SSO_SECRET_FIELDS)),
}

125
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -27,7 +27,11 @@ from litellm.responses.litellm_completion_transformation.tool_search import (
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, stamp_responses_usage_cost
from litellm.responses.streaming_iterator import (
SEQUENCE_NUMBER_STAMPED_ON_EMISSION,
ResponsesAPIStreamingIterator,
stamp_responses_usage_cost,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import (
PART_UNION_TYPES,
@ -239,14 +243,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
return execution
def _next_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> OutputItemAddedEvent:
self._sequence_number += 1
added: Final = OutputItemAddedEvent(
return OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=output_index,
item=item,
)
added.__dict__["sequence_number"] = self._sequence_number
return added
def _queue_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> None:
self._pending_tool_events.append(self._next_output_item_added(output_index, item))
@ -268,12 +269,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
item: Final = build_tool_search_call(
call_id, final_args or self._tool_args_by_call_id.get(call_id, ""), "completed", execution
)
self._sequence_number += 1
self._pending_tool_events.append(
OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=output_index,
sequence_number=self._sequence_number,
item=BaseLiteLLMOpenAIResponseObject(**item.model_dump()),
)
)
@ -346,12 +345,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
if streamed_item is None:
return
self._sent_compaction_item_done = True
self._sequence_number += 1
self._pending_response_events.append(
OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=self._compaction_output_index,
sequence_number=self._sequence_number,
item=streamed_item,
)
)
@ -494,15 +491,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
chunk_size = 10 # Match typical OpenAI delta size
for i in range(0, len(fn_args_delta), chunk_size):
delta_chunk = fn_args_delta[i : i + chunk_size]
self._sequence_number += 1
delta_event: BaseLiteLLMOpenAIResponseObject = FunctionCallArgumentsDeltaEvent(
type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
output_index=output_index,
delta=delta_chunk,
)
# Add sequence_number as extra field (BaseLiteLLMOpenAIResponseObject allows extra fields)
delta_event.__dict__["sequence_number"] = self._sequence_number
self._pending_tool_events.append(delta_event)
def _queue_final_tool_call_done_events(self, litellm_complete_object: ModelResponse) -> None:
@ -567,33 +561,27 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
chunk_size = 10 # Match typical OpenAI delta size
for i in range(0, len(remaining_args), chunk_size):
delta_chunk = remaining_args[i : i + chunk_size]
self._sequence_number += 1
delta_event = FunctionCallArgumentsDeltaEvent(
type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
output_index=output_index,
delta=delta_chunk,
)
delta_event.__dict__["sequence_number"] = self._sequence_number
self._pending_tool_events.append(delta_event)
self._sequence_number += 1
done_event = FunctionCallArgumentsDoneEvent(
type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE,
item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
output_index=output_index,
arguments=final_args,
)
done_event.__dict__["sequence_number"] = self._sequence_number
self._pending_tool_events.append(done_event)
self._sequence_number += 1
item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, final_args, "completed")
item_kwargs["id"] = self._tool_item_id_by_call_id.setdefault(call_id, item_kwargs["id"])
item_done_event = OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=output_index,
sequence_number=self._sequence_number,
item=BaseLiteLLMOpenAIResponseObject(
**item_kwargs,
**(_tool_call_provider_specific_fields(tc) if self.tool_call_provider_fields else {}),
@ -625,16 +613,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
_WEB_SEARCH_PROGRESS_EVENTS[:1] if item.status == "in_progress" else _WEB_SEARCH_PROGRESS_EVENTS
)
for event_type, event_class in progress_events:
self._sequence_number += 1
event = event_class(type=event_type, output_index=output_index, item_id=item.id)
event.__dict__["sequence_number"] = self._sequence_number
self._pending_tool_events.append(event)
self._sequence_number += 1
self._pending_tool_events.append(event_class(type=event_type, output_index=output_index, item_id=item.id))
self._pending_tool_events.append(
OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=output_index,
sequence_number=self._sequence_number,
item=BaseLiteLLMOpenAIResponseObject(**item.model_dump()),
)
)
@ -770,31 +753,24 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
"""
response_created_event_data: Final = self._default_response_created_event_data()
self._sequence_number += 1
event: Final = ResponseCreatedEvent(
return ResponseCreatedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_CREATED,
response=self._with_encoded_response_id(ResponsesAPIResponse(**response_created_event_data)),
)
event.__dict__["sequence_number"] = self._sequence_number
return event
def create_response_in_progress_event(self) -> ResponseInProgressEvent:
response_in_progress_event_data: Final = self._default_response_created_event_data()
response_in_progress_event_data["status"] = "in_progress"
self._sequence_number += 1
event: Final = ResponseInProgressEvent(
return ResponseInProgressEvent(
type=ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS,
response=self._with_encoded_response_id(ResponsesAPIResponse(**response_in_progress_event_data)),
)
event.__dict__["sequence_number"] = self._sequence_number
return event
def create_output_item_added_event(self) -> OutputItemAddedEvent:
if self._cached_item_id is None:
self._cached_item_id = f"msg_{uuid.uuid4()}"
self._sequence_number += 1
event: Final = OutputItemAddedEvent(
return OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=self._message_index(),
item=BaseLiteLLMOpenAIResponseObject(
@ -807,23 +783,18 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
}
),
)
event.__dict__["sequence_number"] = self._sequence_number
return event
def create_content_part_added_event(self) -> ContentPartAddedEvent:
if self._cached_item_id is None:
self._cached_item_id = f"msg_{uuid.uuid4()}"
self._sequence_number += 1
event: Final = ContentPartAddedEvent(
return ContentPartAddedEvent(
type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED,
item_id=self._cached_item_id,
output_index=self._message_index(),
content_index=0,
part=BaseLiteLLMOpenAIResponseObject(**{"type": "output_text", "text": "", "annotations": []}),
)
event.__dict__["sequence_number"] = self._sequence_number
return event
def _queue_message_item_added_events(self) -> None:
if not self.sent_message_item_added_event:
@ -834,7 +805,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
def _queue_message_item_added_event(self) -> None:
item_id: Final = self._message_item_id()
self.sent_message_item_added_event = True
self._sequence_number += 1
event: Final = OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=self._message_index(),
@ -848,7 +818,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
}
),
)
event.__dict__["sequence_number"] = self._sequence_number
self._pending_response_events.append(event)
def _queue_refusal_events(self, refusal: str) -> None:
@ -885,12 +854,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
part=ContentPartDonePartRefusal(type="refusal", refusal=refusal),
),
)
first_sequence_number: Final = self._sequence_number + 1
self._sequence_number += len(events)
self._pending_response_events.extend(
event.model_copy(update={"sequence_number": sequence_number})
for sequence_number, event in enumerate(events, start=first_sequence_number)
)
self._pending_response_events.extend(events)
def _merge_provider_specific_fields(self, src: dict) -> None:
"""Merge provider_specific_fields using last-value-wins for lists.
@ -944,7 +908,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self,
reasoning_item_id: str,
reasoning_content: str,
sequence_number: int,
) -> ReasoningSummaryTextDoneEvent:
"""
Create response.reasoning_summary_text.done event.
@ -954,7 +917,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
"type": "response.reasoning_summary_text.done",
"item_id": "rs_0c5dae30e53172980069708ba2f59c8197b71ca9820edad07c",
"output_index": 0,
"sequence_number": 97,
"summary_index": 0,
"text": "**Clarifying the first humans**\n\nThe I'm addressing the user's specific interest."
}
@ -963,7 +925,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
type=ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE,
item_id=reasoning_item_id,
output_index=self._reasoning_index(),
sequence_number=sequence_number,
sequence_number=SEQUENCE_NUMBER_STAMPED_ON_EMISSION,
summary_index=0,
text=reasoning_content,
)
@ -972,7 +934,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self,
reasoning_item_id: str,
reasoning_content: str,
sequence_number: int,
) -> ReasoningSummaryPartDoneEvent:
"""
Create response.reasoning_summary_part.done event.
@ -986,7 +947,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
"type": "summary_text",
"text": "**Clarifying the first humans**\n\nThe earlier hominins. It feels important to ensure I'm addressing the user's specific interest."
},
"sequence_number": 98,
"summary_index": 0
}
"""
@ -994,7 +954,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
type=ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE,
item_id=reasoning_item_id,
output_index=self._reasoning_index(),
sequence_number=sequence_number,
sequence_number=SEQUENCE_NUMBER_STAMPED_ON_EMISSION,
summary_index=0,
part=BaseLiteLLMOpenAIResponseObject(
**{
@ -1066,7 +1026,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
return OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=self._message_index(),
sequence_number=1,
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": self._cached_item_id,
@ -1082,7 +1041,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self,
reasoning_item_id: str,
reasoning_content: str,
sequence_number: int,
) -> OutputItemDoneEvent:
"""
Create response.output_item.done event for reasoning items.
@ -1091,7 +1049,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
{
"type": "response.output_item.done",
"output_index": 0,
"sequence_number": 99,
"item": {
"id": "rs_0c5dae30e53172980069708ba2f59c8197b71ca9820edad07c",
"type": "reasoning",
@ -1107,7 +1064,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
return OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=self._reasoning_index(),
sequence_number=sequence_number,
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": reasoning_item_id,
@ -1223,7 +1179,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
if chunk.choices[0].finish_reason is None and not _delta_carries_output(delta):
return
self._sequence_number += 1
self.sent_output_item_added_event = True
# Reasoning-first
@ -1233,20 +1188,20 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._cached_reasoning_item_id = mint_reasoning_item_id()
self._reasoning_item_id = self._cached_reasoning_item_id
event = OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=self._reasoning_index(),
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": self._cached_reasoning_item_id,
"type": "reasoning",
"status": "in_progress",
"summary": None,
}
),
self._pending_response_events.append(
OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=self._reasoning_index(),
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": self._cached_reasoning_item_id,
"type": "reasoning",
"status": "in_progress",
"summary": None,
}
),
)
)
event.__dict__["sequence_number"] = self._sequence_number
self._pending_response_events.append(event)
return
# Tool-first
@ -1279,30 +1234,34 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
)
reasoning_item_id: Final = self._cached_reasoning_item_id
self._queue_compaction_item_done()
self._sequence_number += 1
text_done_event: Final = self.create_reasoning_summary_text_done_event(
reasoning_item_id=reasoning_item_id,
reasoning_content=reasoning_content,
sequence_number=self._sequence_number,
)
self._sequence_number += 1
part_done_event: Final = self.create_reasoning_summary_part_done_event(
reasoning_item_id=reasoning_item_id,
reasoning_content=reasoning_content,
sequence_number=self._sequence_number,
)
self._sequence_number += 1
reasoning_output_item_done_event: Final = self.create_reasoning_output_item_done_event(
reasoning_item_id=reasoning_item_id,
reasoning_content=reasoning_content,
sequence_number=self._sequence_number,
)
self._pending_response_events.extend([text_done_event, part_done_event, reasoning_output_item_done_event])
self._reasoning_done_emitted = True
self._reasoning_active = False
def _stamp_sequence_number(self, event: BaseLiteLLMOpenAIResponseObject) -> BaseLiteLLMOpenAIResponseObject:
stamped: Final = event.model_copy(update={"sequence_number": self._sequence_number})
self._sequence_number += 1
return stamped
async def __anext__(
self,
) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject:
return self._stamp_sequence_number(await self._anext_unstamped_event())
async def _anext_unstamped_event(
self,
) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject:
try:
while True:
@ -1363,6 +1322,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
def __next__(
self,
) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject:
return self._stamp_sequence_number(self._next_unstamped_event())
def _next_unstamped_event(
self,
) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject:
try:
while True:
@ -1504,16 +1468,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
if not self.sent_content_part_added_event:
self._queue_message_item_added_events()
self._hold_tool_calls_behind_text(getattr(chunk.choices[0].delta, "tool_calls", None))
self._sequence_number += 1
text_delta_event: Final = OutputTextDeltaEvent(
return OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id=item_id,
output_index=self._message_index(),
content_index=0,
delta=delta_content,
)
text_delta_event.__dict__["sequence_number"] = self._sequence_number
return text_delta_event
# Priority 3: Handle tool call deltas (if any) -> queue events and emit them
# For each tool call delta, we emit events one at a time to match OpenAI's streaming behavior

12
litellm/responses/streaming_iterator.py
View file

@ -379,6 +379,7 @@ def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool:
PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: Final = frozenset({"response.created", "response.in_progress", "response.queued"})
SEQUENCE_NUMBER_STAMPED_ON_EMISSION: Final = -1
class BaseResponsesAPIStreamingIterator:
@ -1933,7 +1934,6 @@ def build_synthetic_response_events(
_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed),
]
sequence_number = 0
output_items: Final[Sequence[object]] = getattr(transformed, "output", []) or []
for output_index, output_item in enumerate(output_items):
output_item_payload = _dump_response_object(output_item)
@ -2043,35 +2043,31 @@ def build_synthetic_response_events(
delta=summary_text[i : i + chunk_size],
)
)
sequence_number += 1
events.append(
openai_types.ReasoningSummaryTextDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE,
item_id=item_id,
output_index=output_index,
sequence_number=sequence_number,
sequence_number=SEQUENCE_NUMBER_STAMPED_ON_EMISSION,
summary_index=summary_index,
text=summary_text,
)
)
sequence_number += 1
events.append(
openai_types.ReasoningSummaryPartDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE,
item_id=item_id,
output_index=output_index,
sequence_number=sequence_number,
sequence_number=SEQUENCE_NUMBER_STAMPED_ON_EMISSION,
summary_index=summary_index,
part=openai_types.BaseLiteLLMOpenAIResponseObject(**summary_payload),
)
)
sequence_number += 1
events.append(
openai_types.OutputItemDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=output_index,
sequence_number=sequence_number,
item=openai_types.BaseLiteLLMOpenAIResponseObject(**output_item_payload),
)
)
@ -2082,7 +2078,7 @@ def build_synthetic_response_events(
response=transformed,
)
)
return events
return [event.model_copy(update={"sequence_number": number}) for number, event in enumerate(events)]
# ---------------------------------------------------------------------------

15
litellm/responses/utils.py
View file

@ -678,20 +678,25 @@ class ResponsesAPIRequestUtils:
return
items: Final = cast(list[object], request_input) # cast-ok: untyped client json
stripped: Final = tuple(
ResponsesAPIRequestUtils._without_encrypted_reasoning(item)
if should_strip is None or (isinstance(item, Mapping) and should_strip(cast(Mapping[str, object], item)))
ResponsesAPIRequestUtils._without_encrypted_reasoning(reasoning)
if (reasoning := ResponsesAPIRequestUtils._encrypted_reasoning_item(item)) is not None
and (should_strip is None or should_strip(reasoning))
else item
for item in items
)
items[:] = (item for item in stripped if item is not None)
@staticmethod
def _without_encrypted_reasoning(item: object) -> object | None:
def _encrypted_reasoning_item(item: object) -> Mapping[str, object] | None:
if not isinstance(item, dict):
return item
return None
reasoning: Final = cast(Mapping[str, object], item) # cast-ok: untyped client json
if reasoning.get("type") != "reasoning" or not reasoning.get("encrypted_content"):
return reasoning
return None
return reasoning
@staticmethod
def _without_encrypted_reasoning(reasoning: Mapping[str, object]) -> object | None:
readable: Final = any(
ResponsesAPIRequestUtils._has_readable_text(reasoning.get(key)) for key in ("summary", "content")
)

19
litellm/router_python_core.py
View file

@ -544,6 +544,7 @@ if TYPE_CHECKING:
ResponseAPIUsage,
ResponseIncompleteEvent,
ResponsesAPIResponse,
ResponsesAPIStreamingResponse,
)
Span = _Span
@ -4385,12 +4386,13 @@ class PythonCore(DeprecatedRouterMembers):
fallback_headers_adopted: bool = False
def __init__(self, async_generator: AsyncGenerator):
def __init__(self, async_generator: "AsyncGenerator[ResponsesAPIStreamingResponse, None]"):
import time
from datetime import datetime
self._async_generator = async_generator
self._billing_source: object = source_iterator
self._last_sequence_number = -1
# Mirror every attribute BaseResponsesAPIStreamingIterator.__init__
# would have set. The wrapper bypasses super().__init__ (it has no
# httpx.Response of its own and no provider config to drive), so
@ -4497,7 +4499,16 @@ class PythonCore(DeprecatedRouterMembers):
if self.completed_response is None and getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES:
self.completed_response = chunk
self._yielded_first_chunk = True
return chunk
sequence_number: Final = getattr(chunk, "sequence_number", None)
if not isinstance(sequence_number, int):
return chunk
if not self.fallback_headers_adopted:
self._last_sequence_number = max(self._last_sequence_number, sequence_number)
return chunk
self._last_sequence_number = max(self._last_sequence_number + 1, sequence_number)
if self._last_sequence_number == sequence_number:
return chunk
return chunk.model_copy(update={"sequence_number": self._last_sequence_number})
async def aclose(self):
# async generators always expose aclose — no defensive check needed.
@ -4559,7 +4570,9 @@ class PythonCore(DeprecatedRouterMembers):
exc,
)
wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks())
wrapper: Final = FallbackResponsesStreamWrapper(
stream_with_fallbacks() # pyright: ignore[reportArgumentType] # relays response events typed as object
)
return wrapper
def _aresponses_background_stream_can_release(

285
litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py
View file

@ -36,31 +36,75 @@ Safe to enable globally:
- No cache required.
"""
from collections.abc import Iterator, Mapping, Sequence
from functools import cache
import hashlib
from collections.abc import Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass, field
from functools import cache, lru_cache, partial
from time import monotonic
from typing import TYPE_CHECKING, Final, Optional, cast
from pydantic import TypeAdapter
from pydantic import BaseModel, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_router_logger
from litellm.exceptions import AuthenticationError, BadRequestError
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.get_llm_provider_logic import declared_provider, get_llm_provider
from litellm.litellm_core_utils.litellm_logging import has_sensitive_keyword
from litellm.litellm_core_utils.prompt_templates.common_utils import (
anthropic_content_lists,
encrypted_content_of_block,
strip_encrypted_reasoning_from_messages,
)
from litellm.litellm_core_utils.routing_common_utils import as_excluded_deployment_ids, as_retry_skipped_deployment_ids
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param
from litellm.secret_managers.main import secret_manager_would_be_consulted
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import Deployment
from litellm.utils import get_order_filtered_deployments
from litellm.types.router import Deployment, GenericLiteLLMParams
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager, get_order_filtered_deployments, with_named_credential
if TYPE_CHECKING:
from litellm.router_api import RoutingContext
_REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_HEADERS_ADAPTER: Final = TypeAdapter(Mapping[object, object])
_ACCOUNT_SCOPE_HEADERS: Final = frozenset({"openai-organization", "openai-project", "cookie"})
_SINGLE_VENDOR_PROVIDERS: Final = frozenset({LlmProviders.OPENAI.value})
_DEFAULT_ROUTE_SECRETS: Final = ("USE_LITELLM_PROXY", "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENAI_API_BASE")
_SECRET_MANAGER_REFRESH_SECONDS: Final = 60
@dataclass(frozen=True, slots=True)
class _ProviderRoute:
model: str
provider: str
api_base: str | None
@dataclass(frozen=True, slots=True)
class _EncryptionBoundary:
api_base: object
api_key: object
model: Callable[[], str | None] = field(compare=False)
account_scope: frozenset[tuple[str, object]]
explicit_pair: bool
def is_shared_with(self, other: "_EncryptionBoundary") -> bool:
if (self.api_base, self.api_key) != (other.api_base, other.api_key):
return False
if self.explicit_pair and other.explicit_pair:
return True
return self.account_scope == other.account_scope and self.model() == other.model()
@dataclass(frozen=True, slots=True)
class _ParamsInRefreshWindow:
digest: str
window: int
params: Mapping[str, object] = field(compare=False)
class EncryptedContentAffinityCheck(CustomLogger):
@ -83,9 +127,11 @@ class EncryptedContentAffinityCheck(CustomLogger):
ctx: Optional["RoutingContext"] = None,
enable_global_affinity: bool = True,
model_group_affinity_config: dict[str, list[str]] | None = None,
clock: Callable[[], float] = monotonic,
) -> None:
super().__init__()
self.ctx = ctx
self._clock = clock
self.enable_global_affinity = enable_global_affinity
self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {}
@ -209,44 +255,156 @@ class EncryptedContentAffinityCheck(CustomLogger):
@staticmethod
def _encryption_boundary_key(
litellm_params: object,
) -> tuple[object, object] | None:
clock: Callable[[], float] = monotonic,
) -> _EncryptionBoundary | None:
"""
``(api_base, api_key)`` identifies an upstream encryption boundary.
The values are resolved from the deployment and its named credential
without modifying the deployment.
The upstream encryption boundary a deployment's encrypted reasoning belongs to. The values
are resolved the way the request resolves them, the deployment's own values winning over its
named credential's, without modifying the deployment. Two explicit ``(api_base, api_key)``
pairs match on those values alone. On a provider that serves one vendor's models, a deployment
with no ``api_key`` takes the provider's default key, one with no ``api_base`` takes the endpoint
its Responses URL builder resolves, and a match involving either also needs the same model and
the same headers that pick an account.
Accepts any object exposing dict-style ``.get(key, default)``: plain
dicts (the common case in ``healthy_deployments``) as well as
``LiteLLM_Params``-style Pydantic instances, which define a custom
``.get()``. A stricter ``isinstance(dict)`` guard would silently drop
the latter from boundary matching and fall back to the full pool —
i.e. trigger the exact ``invalid_encrypted_content`` failure this
check exists to prevent.
Accepts plain dicts (the common case in ``healthy_deployments``) as
well as ``LiteLLM_Params`` Pydantic instances. A stricter
``isinstance(dict)`` guard would silently drop the latter from
boundary matching and fall back to the full pool — i.e. trigger the
exact ``invalid_encrypted_content`` failure this check exists to prevent.
When a secret manager would be asked for the provider defaults, the boundary of one set of
deployment and named credential values is reused for up to a minute, so a run of pin misses
does not ask it again for every candidate, while a changed credential resolves at once.
"""
getter: Final = getattr(litellm_params, "get", None)
if not callable(getter):
params: Final = EncryptedContentAffinityCheck._params_mapping(litellm_params)
if params is None:
return None
api_base: Final = getter("api_base")
api_key: Final = getter("api_key")
credential_name: Final = getter("litellm_credential_name")
credential_values: Final[Mapping[str, object] | None] = (
CredentialAccessor.get_credential_values(credential_name)
if isinstance(credential_name, str) and credential_name
else None
effective_params: Final = with_named_credential(params)
if not any(secret_manager_would_be_consulted(name) for name in _DEFAULT_ROUTE_SECRETS):
return EncryptedContentAffinityCheck._resolve_encryption_boundary(effective_params)
return EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window(
_ParamsInRefreshWindow(
hashlib.sha256(repr(effective_params).encode()).hexdigest(),
int(clock() // _SECRET_MANAGER_REFRESH_SECONDS),
effective_params,
)
)
effective_api_base: Final = (
credential_values.get("api_base")
if credential_values is not None and "api_base" in credential_values
else api_base
)
effective_api_key: Final = (
credential_values.get("api_key")
if credential_values is not None and "api_key" in credential_values
else api_key
)
if not effective_api_base or not effective_api_key:
@staticmethod
@lru_cache(maxsize=1024)
def _encryption_boundary_in_refresh_window(key: _ParamsInRefreshWindow) -> _EncryptionBoundary | None:
return EncryptedContentAffinityCheck._resolve_encryption_boundary(key.params)
@staticmethod
def _resolve_encryption_boundary(effective_params: Mapping[str, object]) -> _EncryptionBoundary | None:
explicit_api_key: Final = effective_params.get("api_key")
explicit_api_base: Final = effective_params.get("api_base")
account_scope: Final = EncryptedContentAffinityCheck._account_scope(effective_params)
if explicit_api_key and explicit_api_base:
return _EncryptionBoundary(
explicit_api_base,
explicit_api_key,
cache(partial(EncryptedContentAffinityCheck._single_vendor_model, effective_params)),
account_scope,
explicit_pair=True,
)
route: Final = EncryptedContentAffinityCheck._single_vendor_route(effective_params)
if route is None:
return None
api_key: Final = explicit_api_key or EncryptedContentAffinityCheck._default_api_key(route)
if not api_key:
return None
api_base: Final = explicit_api_base or EncryptedContentAffinityCheck._default_responses_endpoint(
route, effective_params
)
if not api_base:
return None
return _EncryptionBoundary(api_base, api_key, lambda: route.model, account_scope, explicit_pair=False)
@staticmethod
def _params_mapping(litellm_params: object) -> Mapping[str, object] | None:
if isinstance(litellm_params, BaseModel):
return litellm_params.model_dump(exclude_none=True)
if isinstance(litellm_params, Mapping):
return _REQUEST_KWARGS_ADAPTER.validate_python(litellm_params)
return None
@staticmethod
def _account_scope(litellm_params: Mapping[str, object]) -> frozenset[tuple[str, object]]:
return frozenset(
(name.lower(), str(value))
for name, value in EncryptedContentAffinityCheck._header_items(litellm_params)
if name.lower() in _ACCOUNT_SCOPE_HEADERS or has_sensitive_keyword(name)
)
@staticmethod
def _header_items(litellm_params: Mapping[str, object]) -> Iterator[tuple[str, object]]:
headers: Final = EncryptedContentAffinityCheck._header_mapping(litellm_params.get("headers"))
extra_headers: Final = EncryptedContentAffinityCheck._header_mapping(litellm_params.get("extra_headers"))
overridden: Final = frozenset(name.lower() for name in extra_headers)
yield from ((name, value) for name, value in headers.items() if name.lower() not in overridden)
yield from extra_headers.items()
@staticmethod
def _header_mapping(headers: object) -> Mapping[str, object]:
if isinstance(headers, Mapping):
return {str(name): value for name, value in _HEADERS_ADAPTER.validate_python(headers).items()}
return {}
@staticmethod
def _single_vendor_route(litellm_params: Mapping[str, object]) -> _ProviderRoute | None:
route: Final = EncryptedContentAffinityCheck._provider_route(litellm_params)
return route if route is not None and route.provider in _SINGLE_VENDOR_PROVIDERS else None
@staticmethod
def _single_vendor_model(litellm_params: Mapping[str, object]) -> str | None:
route: Final = EncryptedContentAffinityCheck._single_vendor_route(litellm_params)
return route.model if route is not None else None
@staticmethod
def _provider_route(litellm_params: Mapping[str, object]) -> _ProviderRoute | None:
model: Final = litellm_params.get("model")
if not isinstance(model, str):
return None
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
declared: Final = declared_provider(
model, custom_llm_provider if isinstance(custom_llm_provider, str) else None
)
if declared in litellm.provider_list and declared not in _SINGLE_VENDOR_PROVIDERS:
return None
try:
provider_model, provider, _, provider_api_base = get_llm_provider(
model=model, litellm_params=GenericLiteLLMParams.model_validate(dict(litellm_params))
)
except (BadRequestError, ValidationError):
return None
return _ProviderRoute(provider_model, provider, provider_api_base)
@staticmethod
def _default_api_key(route: _ProviderRoute) -> str | None:
model_info: Final = ProviderConfigManager.get_provider_model_info(
model=route.model, provider=LlmProviders(route.provider)
)
if model_info is None:
return None
try:
return model_info.get_api_key()
except NotImplementedError:
return None
@staticmethod
def _default_responses_endpoint(route: _ProviderRoute, litellm_params: Mapping[str, object]) -> str | None:
responses_config: Final = ProviderConfigManager.get_provider_responses_api_config(
provider=route.provider, model=route.model
)
if responses_config is None:
return None
try:
return responses_config.get_complete_url( # pyright: ignore[reportUnknownMemberType] # the base types litellm_params as a bare dict
api_base=route.api_base, litellm_params=dict(litellm_params)
)
except (ValueError, BaseLLMException, AuthenticationError):
return None
return (effective_api_base, effective_api_key)
def _find_deployments_on_same_encryption_boundary(
self,
@ -255,7 +413,7 @@ class EncryptedContentAffinityCheck(CustomLogger):
) -> tuple[list[dict], Deployment | None]:
"""
Deployments in ``healthy_deployments`` sharing the originating
deployment's ``(api_base, api_key)``, alongside the originating
deployment's encryption boundary, alongside the originating
deployment object (or ``None`` if it was removed / router unavailable).
"""
if self.ctx is None:
@ -263,11 +421,16 @@ class EncryptedContentAffinityCheck(CustomLogger):
originating: Final = self.ctx.registry().deployment(model_id)
if originating is None:
return [], None
boundary: Final = self._encryption_boundary_key(originating.litellm_params.model_dump(exclude_none=True))
boundary: Final = self._encryption_boundary_key(
originating.litellm_params.model_dump(exclude_none=True), self._clock
)
if boundary is None:
return [], originating
matches: Final = [
d for d in healthy_deployments if self._encryption_boundary_key(d.get("litellm_params", {})) == boundary
d
for d in healthy_deployments
if (candidate := self._encryption_boundary_key(d.get("litellm_params", {}), self._clock)) is not None
and boundary.is_shared_with(candidate)
]
return matches, originating
@ -280,6 +443,7 @@ class EncryptedContentAffinityCheck(CustomLogger):
target_deployments: Sequence[Mapping[str, object]],
*,
unmarked_origin: str | None,
clock: Callable[[], float] = monotonic,
) -> None:
"""
Drop the encrypted reasoning that none of ``target_deployments`` minted or shares an
@ -293,12 +457,19 @@ class EncryptedContentAffinityCheck(CustomLogger):
for target in target_deployments
if isinstance((model_info := target.get("model_info")), Mapping) and model_info.get("id") is not None
)
target_boundaries: Final = frozenset(
boundary
for target in target_deployments
if (boundary := EncryptedContentAffinityCheck._encryption_boundary_key(target.get("litellm_params")))
is not None
)
@cache
def target_boundaries() -> tuple[_EncryptionBoundary, ...]:
return tuple(
boundary
for target in target_deployments
if (
boundary := EncryptedContentAffinityCheck._encryption_boundary_key(
target.get("litellm_params"), clock
)
)
is not None
)
@cache
def target_can_decrypt(marked_origin: str | None) -> bool:
@ -312,12 +483,14 @@ class EncryptedContentAffinityCheck(CustomLogger):
origin: Final = ctx.registry().deployment(origin_model_id)
origin_boundary: Final = (
EncryptedContentAffinityCheck._encryption_boundary_key(
origin.litellm_params.model_dump(exclude_none=True)
origin.litellm_params.model_dump(exclude_none=True), clock
)
if origin is not None
else None
)
return origin_boundary is not None and origin_boundary in target_boundaries
return origin_boundary is not None and any(
origin_boundary.is_shared_with(target_boundary) for target_boundary in target_boundaries()
)
def should_strip_input_item(item: Mapping[str, object]) -> bool:
return not target_can_decrypt(EncryptedContentAffinityCheck._model_id_of_input_item(item))
@ -424,7 +597,12 @@ class EncryptedContentAffinityCheck(CustomLogger):
)
request_kwargs["_encrypted_content_affinity_pinned"] = True
self.strip_reasoning_the_targets_cannot_decrypt(
self.ctx, request_input, anthropic_messages, (deployment,), unmarked_origin=model_id
self.ctx,
request_input,
anthropic_messages,
(deployment,),
unmarked_origin=model_id,
clock=self._clock,
)
return [deployment]
@ -442,7 +620,12 @@ class EncryptedContentAffinityCheck(CustomLogger):
)
request_kwargs["_encrypted_content_affinity_pinned"] = True
self.strip_reasoning_the_targets_cannot_decrypt(
self.ctx, request_input, anthropic_messages, boundary_matches, unmarked_origin=model_id
self.ctx,
request_input,
anthropic_messages,
boundary_matches,
unmarked_origin=model_id,
clock=self._clock,
)
return boundary_matches

23
litellm/types/guardrails.py
View file

@ -1,10 +1,10 @@
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from datetime import datetime
from enum import Enum
from types import MappingProxyType
from typing import Annotated, Final, Literal, TypeAlias, cast
from pydantic import ConfigDict, Field, field_validator, model_validator
from pydantic import ConfigDict, Field, JsonValue, field_validator, model_validator
from pydantic.fields import FieldInfo
from typing_extensions import ReadOnly, Required, TypedDict
@ -604,6 +604,11 @@ class BedrockChecksConfigModel(LiteLLMBaseModel):
return self
class BedrockGuardrailOutputScope(Enum):
INTERVENTIONS = "INTERVENTIONS"
FULL = "FULL"
class BedrockGuardrailConfigModel(LiteLLMBaseModel):
"""Configuration parameters for the AWS Bedrock guardrail"""
@ -612,6 +617,11 @@ class BedrockGuardrailConfigModel(LiteLLMBaseModel):
default=None,
description="The version of your Bedrock guardrail (e.g., DRAFT or version number)",
)
outputScope: BedrockGuardrailOutputScope | None = Field(
default=None,
description="ApplyGuardrail response scope. FULL includes detected and non-detected findings. "
"Omit to retain Bedrock's default of returning detected findings only.",
)
disable_exception_on_block: bool | None = Field(
default=False,
description="If True, will not raise an exception when the guardrail is blocked. Useful for OpenWebUI where exceptions can end the chat flow.",
@ -1540,6 +1550,14 @@ class PresidioPerRequestConfig(LiteLLMBaseModel):
entities: list[PiiEntityType] | None = None
NativeGuardrailResponse: TypeAlias = Mapping[str, JsonValue] | Sequence[JsonValue]
class GuardrailResponseFormat(Enum):
STANDARD = "standard"
NATIVE = "native"
class ApplyGuardrailRequest(LiteLLMBaseModel):
guardrail_name: str
text: str
@ -1548,6 +1566,7 @@ class ApplyGuardrailRequest(LiteLLMBaseModel):
input_type: str = "request"
messages: list[dict[str, object]] | None = None
metadata: dict[str, object] | None = None
response_format: GuardrailResponseFormat = GuardrailResponseFormat.STANDARD
class ApplyGuardrailResponse(LiteLLMBaseModel):

109
litellm/types/llms/openai.py
View file

@ -1,8 +1,9 @@
import builtins
import math
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import ItemsView, Iterable, Mapping, Sequence
from enum import Enum
from os import PathLike
from types import MappingProxyType
from typing import IO, TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeAlias, Union
import httpx
@ -692,6 +693,15 @@ class PromptCacheBreakpoint(TypedDict):
ttl: NotRequired[ReadOnly[Literal["30m", "12h"]]]
class ChatCompletionPromptCacheBreakpoint(TypedDict):
mode: ReadOnly[Literal["explicit"]]
class AzureResponsesPromptCacheBreakpoint(TypedDict):
mode: ReadOnly[Literal["explicit"]]
ttl: NotRequired[ReadOnly[Literal["30m"]]]
class PromptCacheOptions(TypedDict, total=False):
mode: ReadOnly[Literal["implicit", "explicit"]]
ttl: ReadOnly[Literal["30m", "12h"]]
@ -1754,17 +1764,40 @@ class ResponsesAPIStreamEvents(str, Enum):
ERROR = "error"
class ResponseCreatedEvent(BaseLiteLLMOpenAIResponseObject):
_NO_EXTRAS: Final[Mapping[str, object]] = MappingProxyType({})
class BaseResponsesStreamEvent(BaseLiteLLMOpenAIResponseObject): # frozen-ok: guardrails rewrite events in place
def __getitem__(self, key: str) -> object:
return self._mapping_holding(key)[key]
def get(self, key: str, default: object = None) -> object:
return self._mapping_holding(key).get(key, default)
def __contains__(self, key: object) -> bool:
return key in self.__dict__ or key in self._extras()
def items(self) -> ItemsView[str, object]:
return MappingProxyType({**self.__dict__, **self._extras()}).items()
def _extras(self) -> Mapping[str, object]:
return self.__pydantic_extra__ or _NO_EXTRAS
def _mapping_holding(self, key: str) -> Mapping[str, object]:
return self.__dict__ if key in self.__dict__ else self._extras()
class ResponseCreatedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.RESPONSE_CREATED]
response: ResponsesAPIResponse
class ResponseInProgressEvent(BaseLiteLLMOpenAIResponseObject):
class ResponseInProgressEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS]
response: ResponsesAPIResponse
class ResponseCompletedEvent(BaseLiteLLMOpenAIResponseObject):
class ResponseCompletedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.RESPONSE_COMPLETED]
response: ResponsesAPIResponse
_hidden_params: dict = PrivateAttr(default_factory=dict)
@ -1778,12 +1811,12 @@ class ResponseCompletedEvent(BaseLiteLLMOpenAIResponseObject):
self._hidden_params = hidden_params
class ResponseFailedEvent(BaseLiteLLMOpenAIResponseObject):
class ResponseFailedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.RESPONSE_FAILED]
response: ResponsesAPIResponse
class ResponseIncompleteEvent(BaseLiteLLMOpenAIResponseObject):
class ResponseIncompleteEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE]
response: ResponsesAPIResponse
@ -1791,7 +1824,7 @@ class ResponseIncompleteEvent(BaseLiteLLMOpenAIResponseObject):
ResponsesTerminalEvent: TypeAlias = ResponseCompletedEvent | ResponseIncompleteEvent | ResponseFailedEvent
class ResponsePartAddedEvent(BaseLiteLLMOpenAIResponseObject):
class ResponsePartAddedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.RESPONSE_PART_ADDED]
item_id: str
output_index: int
@ -1799,7 +1832,7 @@ class ResponsePartAddedEvent(BaseLiteLLMOpenAIResponseObject):
part: dict
class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject):
class ReasoningSummaryTextDeltaEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA]
item_id: str
output_index: int
@ -1807,7 +1840,7 @@ class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject):
delta: str
class ReasoningSummaryTextDoneEvent(BaseLiteLLMOpenAIResponseObject):
class ReasoningSummaryTextDoneEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE]
item_id: str
output_index: int
@ -1816,7 +1849,7 @@ class ReasoningSummaryTextDoneEvent(BaseLiteLLMOpenAIResponseObject):
text: str
class ReasoningSummaryPartDoneEvent(BaseLiteLLMOpenAIResponseObject):
class ReasoningSummaryPartDoneEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE]
item_id: str
output_index: int
@ -1825,13 +1858,13 @@ class ReasoningSummaryPartDoneEvent(BaseLiteLLMOpenAIResponseObject):
part: BaseLiteLLMOpenAIResponseObject
class OutputItemAddedEvent(BaseLiteLLMOpenAIResponseObject):
class OutputItemAddedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED]
output_index: int
item: BaseLiteLLMOpenAIResponseObject | None
class OutputItemDoneEvent(BaseLiteLLMOpenAIResponseObject):
class OutputItemDoneEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE]
output_index: int
sequence_number: int = 1
@ -1851,7 +1884,7 @@ class OpenAIChatCompletionLogprobsContent(TypedDict, total=False):
top_logprobs: list[OpenAIChatCompletionLogprobsContentTopLogprobs]
class ContentPartAddedEvent(BaseLiteLLMOpenAIResponseObject):
class ContentPartAddedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.CONTENT_PART_ADDED]
item_id: str
output_index: int
@ -1879,7 +1912,7 @@ class ContentPartDonePartReasoningText(BaseLiteLLMOpenAIResponseObject):
PART_UNION_TYPES = ContentPartDonePartOutputText | ContentPartDonePartRefusal | ContentPartDonePartReasoningText
class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject):
class ContentPartDoneEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.CONTENT_PART_DONE]
item_id: str
output_index: int
@ -1887,7 +1920,7 @@ class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject):
part: PART_UNION_TYPES
class OutputTextDeltaEvent(BaseLiteLLMOpenAIResponseObject):
class OutputTextDeltaEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA]
item_id: str
output_index: int
@ -1895,7 +1928,7 @@ class OutputTextDeltaEvent(BaseLiteLLMOpenAIResponseObject):
delta: str
class OutputTextAnnotationAddedEvent(BaseLiteLLMOpenAIResponseObject):
class OutputTextAnnotationAddedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED]
item_id: str
output_index: int
@ -1904,7 +1937,7 @@ class OutputTextAnnotationAddedEvent(BaseLiteLLMOpenAIResponseObject):
annotation: dict
class OutputTextDoneEvent(BaseLiteLLMOpenAIResponseObject):
class OutputTextDoneEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE]
item_id: str
output_index: int
@ -1912,7 +1945,7 @@ class OutputTextDoneEvent(BaseLiteLLMOpenAIResponseObject):
text: str
class RefusalDeltaEvent(BaseLiteLLMOpenAIResponseObject):
class RefusalDeltaEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.REFUSAL_DELTA]
item_id: str
output_index: int
@ -1920,7 +1953,7 @@ class RefusalDeltaEvent(BaseLiteLLMOpenAIResponseObject):
delta: str
class RefusalDoneEvent(BaseLiteLLMOpenAIResponseObject):
class RefusalDoneEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.REFUSAL_DONE]
item_id: str
output_index: int
@ -1928,72 +1961,72 @@ class RefusalDoneEvent(BaseLiteLLMOpenAIResponseObject):
refusal: str
class FunctionCallArgumentsDeltaEvent(BaseLiteLLMOpenAIResponseObject):
class FunctionCallArgumentsDeltaEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA]
item_id: str
output_index: int
delta: str
class FunctionCallArgumentsDoneEvent(BaseLiteLLMOpenAIResponseObject):
class FunctionCallArgumentsDoneEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE]
item_id: str
output_index: int
arguments: str
class FileSearchCallInProgressEvent(BaseLiteLLMOpenAIResponseObject):
class FileSearchCallInProgressEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.FILE_SEARCH_CALL_IN_PROGRESS]
output_index: int
item_id: str
class FileSearchCallSearchingEvent(BaseLiteLLMOpenAIResponseObject):
class FileSearchCallSearchingEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.FILE_SEARCH_CALL_SEARCHING]
output_index: int
item_id: str
class FileSearchCallCompletedEvent(BaseLiteLLMOpenAIResponseObject):
class FileSearchCallCompletedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.FILE_SEARCH_CALL_COMPLETED]
output_index: int
item_id: str
class WebSearchCallInProgressEvent(BaseLiteLLMOpenAIResponseObject):
class WebSearchCallInProgressEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.WEB_SEARCH_CALL_IN_PROGRESS]
output_index: int
item_id: str
class WebSearchCallSearchingEvent(BaseLiteLLMOpenAIResponseObject):
class WebSearchCallSearchingEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.WEB_SEARCH_CALL_SEARCHING]
output_index: int
item_id: str
class WebSearchCallCompletedEvent(BaseLiteLLMOpenAIResponseObject):
class WebSearchCallCompletedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.WEB_SEARCH_CALL_COMPLETED]
output_index: int
item_id: str
# MCP List Tools Events
class MCPListToolsInProgressEvent(BaseLiteLLMOpenAIResponseObject):
class MCPListToolsInProgressEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.MCP_LIST_TOOLS_IN_PROGRESS]
sequence_number: int
output_index: int
item_id: str
class MCPListToolsCompletedEvent(BaseLiteLLMOpenAIResponseObject):
class MCPListToolsCompletedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED]
sequence_number: int
output_index: int
item_id: str
class MCPListToolsFailedEvent(BaseLiteLLMOpenAIResponseObject):
class MCPListToolsFailedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.MCP_LIST_TOOLS_FAILED]
sequence_number: int
output_index: int
@ -2001,14 +2034,14 @@ class MCPListToolsFailedEvent(BaseLiteLLMOpenAIResponseObject):
# MCP Call Events
class MCPCallInProgressEvent(BaseLiteLLMOpenAIResponseObject):
class MCPCallInProgressEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.MCP_CALL_IN_PROGRESS]
sequence_number: int
output_index: int
item_id: str
class MCPCallArgumentsDeltaEvent(BaseLiteLLMOpenAIResponseObject):
class MCPCallArgumentsDeltaEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA]
output_index: int
item_id: str
@ -2016,7 +2049,7 @@ class MCPCallArgumentsDeltaEvent(BaseLiteLLMOpenAIResponseObject):
sequence_number: int
class MCPCallArgumentsDoneEvent(BaseLiteLLMOpenAIResponseObject):
class MCPCallArgumentsDoneEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DONE]
output_index: int
item_id: str
@ -2024,21 +2057,21 @@ class MCPCallArgumentsDoneEvent(BaseLiteLLMOpenAIResponseObject):
sequence_number: int
class MCPCallCompletedEvent(BaseLiteLLMOpenAIResponseObject):
class MCPCallCompletedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.MCP_CALL_COMPLETED]
sequence_number: int
item_id: str
output_index: int
class MCPCallFailedEvent(BaseLiteLLMOpenAIResponseObject):
class MCPCallFailedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.MCP_CALL_FAILED]
sequence_number: int
item_id: str
output_index: int
class ImageGenerationPartialImageEvent(BaseLiteLLMOpenAIResponseObject):
class ImageGenerationPartialImageEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.IMAGE_GENERATION_PARTIAL_IMAGE]
partial_image_index: int
b64_json: str
@ -2053,7 +2086,7 @@ class ErrorEventError(BaseLiteLLMOpenAIResponseObject):
param: str | dict[str, object] | None = None
class ErrorEvent(BaseLiteLLMOpenAIResponseObject):
class ErrorEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.ERROR]
sequence_number: int
error: ErrorEventError | None = None
@ -2062,7 +2095,7 @@ class ErrorEvent(BaseLiteLLMOpenAIResponseObject):
param: str | Mapping[str, object] | None = None
class GenericEvent(BaseLiteLLMOpenAIResponseObject):
class GenericEvent(BaseResponsesStreamEvent):
type: str
model_config = ConfigDict(extra="allow", protected_namespaces=())

14
litellm/types/passthrough_endpoints/pass_through_endpoints.py
View file

@ -1,9 +1,11 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Final, Protocol
from typing import Final, Literal, Protocol, TypeAlias
from pydantic import Field
from typing_extensions import TypedDict
from typing_extensions import ReadOnly, TypedDict
from litellm.types.llms.base import LiteLLMPydanticObjectBase
@ -28,6 +30,10 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__
SETTLES_BACKGROUND_RESPONSE_KWARG: Final = "settles_background_response"
WebsocketUpstreamFrameKind: TypeAlias = Literal[
"starts_turn", "ends_turn", "opens_session", "updates_session", "bills_usage", "refuses_turn", "relayed"
]
class BackgroundResponseSettlement(str, Enum):
SETTLES = "settles"
@ -63,7 +69,7 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False):
"""
The body of the request
"""
response_body: dict | None # only tracked for non-streaming responses
response_body: dict | Sequence[Mapping[str, object]] | None # only tracked for non-streaming responses
"""
The body of the response
"""
@ -75,6 +81,8 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False):
Optional field, we use this for cost tracking only if it's set.
"""
end_time: ReadOnly[datetime]
@dataclass(frozen=True, slots=True)
class WebSocketSessionUsage:

4
litellm/types/proxy/guardrails/guardrail_hooks/ovalix.py
View file

@ -28,6 +28,10 @@ class OvalixGuardrailConfigModel(GuardrailConfigModel):
default=None,
description="Post-checkpoint ID for the Ovalix Tracker service.",
)
file_checkpoint_id: str | None = Field(
default=None,
description="File-checkpoint ID for the Ovalix Tracker service (falls back to pre/post).",
)
@staticmethod
def ui_friendly_name() -> str:

2
litellm/types/utils.py
View file

@ -1151,6 +1151,8 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = {
# A2A (Agent-to-Agent)
"/a2a/{agent_id}": [CallTypes.asend_message, CallTypes.send_message],
"/a2a/{agent_id}/message/send": [CallTypes.asend_message, CallTypes.send_message],
"/a2a/{agent_id}/message/stream": [CallTypes.asend_message, CallTypes.send_message],
"/v1/a2a/{agent_id}/message/stream": [CallTypes.asend_message, CallTypes.send_message],
# Passthrough endpoints
"/llm_passthrough": [
CallTypes.allm_passthrough_route,

82
model_prices_and_context_window.json
View file

@ -15555,18 +15555,20 @@
"input_cost_per_token": 2.65e-06,
"litellm_provider": "bedrock",
"max_input_tokens": 8000,
"max_output_tokens": 2048,
"max_output_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 3.5e-06,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-70b-instruct.md",
"supports_pdf_input": true
},
"bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0": {
"input_cost_per_token": 3e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 8000,
"max_output_tokens": 2048,
"max_output_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-8b-instruct.md",
"supports_pdf_input": true
},
"bedrock/us-west-1/meta.llama3-70b-instruct-v1:0": {
@ -48955,6 +48957,7 @@
},
"together_ai/Qwen/Qwen3.8-2.4T-A95B": {
"cache_read_input_token_cost": 2.5e-07,
"deprecation_date": "2026-11-04",
"input_cost_per_token": 2e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 1010000,
@ -53954,23 +53957,23 @@
"supports_vision": true
},
"wandb/meta-llama/Llama-3.1-8B-Instruct": {
"max_input_tokens": 131000,
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"input_cost_per_token": 2.2e-07,
"output_cost_per_token": 2.2e-07,
"litellm_provider": "wandb",
"mode": "chat",
"source": "https://wandb.ai/site/pricing/tokens/"
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference"
},
"wandb/deepseek-ai/DeepSeek-V3.1": {
"supports_reasoning": true,
"max_input_tokens": 161000,
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"input_cost_per_token": 5.5e-07,
"output_cost_per_token": 1.65e-06,
"litellm_provider": "wandb",
"mode": "chat",
"source": "https://wandb.ai/site/pricing/tokens/"
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference"
},
"wandb/meta-llama/Llama-3.3-70B-Instruct": {
"max_input_tokens": 128000,
@ -59872,6 +59875,7 @@
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"mode": "responses",
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-cyber.html",
"use_openai_responses_path": true,
"supported_endpoints": [
@ -59893,7 +59897,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_xhigh_reasoning_effort": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol": {
"input_cost_per_token": 4.4e-06,
@ -59908,6 +59913,7 @@
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"mode": "responses",
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/responses"
@ -59927,7 +59933,8 @@
"supports_tool_choice": true,
"supports_xhigh_reasoning_effort": true,
"supports_vision": true,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html"
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html",
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.6-luna": {
"input_cost_per_token": 2.2e-07,
@ -60333,6 +60340,7 @@
"mode": "responses",
"output_cost_per_token_above_272k_tokens_ultrafast": 0.000495,
"output_cost_per_token_ultrafast": 0.00033,
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
@ -60356,7 +60364,8 @@
"supports_tool_choice": true,
"supports_xhigh_reasoning_effort": true,
"supports_vision": true,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html"
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html",
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-6-sol": {
"input_cost_per_token": 2.2e-06,
@ -60371,6 +60380,7 @@
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"mode": "responses",
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
@ -60394,7 +60404,8 @@
"supports_tool_choice": true,
"supports_xhigh_reasoning_effort": true,
"supports_vision": true,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html"
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html",
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-6-luna": {
"input_cost_per_token": 1.1e-07,
@ -60409,6 +60420,7 @@
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"mode": "responses",
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
@ -60432,7 +60444,8 @@
"supports_tool_choice": true,
"supports_xhigh_reasoning_effort": true,
"supports_vision": true,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html"
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html",
"supports_web_search": true
},
"us.openai.gpt-6-astra": {
"supports_bedrock_runtime_chat_completions_response_format": true,
@ -64026,6 +64039,7 @@
"output_cost_per_token_priority": 1.5e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -64045,6 +64059,7 @@
"output_cost_per_token_priority": 2.25e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -64087,6 +64102,7 @@
],
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64124,6 +64140,7 @@
"output_cost_per_token_priority": 1.5e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -64143,6 +64160,7 @@
"output_cost_per_token_priority": 2.25e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -64217,6 +64235,7 @@
],
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64237,6 +64256,7 @@
],
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64260,6 +64280,7 @@
],
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64335,6 +64356,7 @@
"output_cost_per_token_priority": 3e-06,
"source": "https://api.fireworks.ai/v1/serverless/models?format=nested",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64371,6 +64393,7 @@
"output_cost_per_token_priority": 6.375e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_vision": true
},
@ -64407,6 +64430,7 @@
"output_cost_per_token_priority": 3e-06,
"source": "https://api.fireworks.ai/v1/serverless/models?format=nested",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64477,6 +64501,7 @@
],
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -64500,6 +64525,7 @@
],
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67293,6 +67319,7 @@
"output_cost_per_token_priority": 5.5e-06,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67311,6 +67338,7 @@
"output_cost_per_token_priority": 8.25e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67329,6 +67357,7 @@
"output_cost_per_token_priority": 5.5e-06,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67347,6 +67376,7 @@
"output_cost_per_token_priority": 8.25e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67362,6 +67392,7 @@
"output_cost_per_token": 6.6e-06,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67404,6 +67435,7 @@
"output_cost_per_token": 6.6e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -67422,6 +67454,7 @@
"output_cost_per_token_priority": 6.25e-07,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
@ -67440,6 +67473,7 @@
"output_cost_per_token_priority": 9.375e-07,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
@ -67457,6 +67491,7 @@
"output_cost_per_token_priority": 6.25e-07,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
@ -67475,6 +67510,7 @@
"output_cost_per_token_priority": 9.375e-07,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
@ -67493,6 +67529,7 @@
"output_cost_per_token_priority": 4.05e-06,
"source": "https://api.fireworks.ai/v1/serverless/models?format=nested",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
@ -68630,7 +68667,9 @@
"cache_read_input_token_cost_above_272k_tokens": 5.4e-07,
"output_cost_per_token": 1.62e-05,
"output_cost_per_token_above_272k_tokens": 2.43e-05,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html"
"search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144},
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html",
"supports_web_search": true
},
"bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": {
"litellm_provider": "bedrock_mantle",
@ -68665,7 +68704,9 @@
"cache_read_input_token_cost_above_272k_tokens": 5.4e-08,
"output_cost_per_token": 1.62e-06,
"output_cost_per_token_above_272k_tokens": 2.43e-06,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html"
"search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144},
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html",
"supports_web_search": true
},
"bedrock_mantle/us-gov-west-1/openai.gpt-5.4": {
"litellm_provider": "bedrock_mantle",
@ -68694,7 +68735,9 @@
"input_cost_per_token": 3.375e-06,
"cache_read_input_token_cost": 3.375e-07,
"output_cost_per_token": 2.025e-05,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html"
"search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144},
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html",
"supports_web_search": true
},
"bedrock_mantle/us-gov-west-1/xai.grok-4.3": {
"use_openai_responses_path": true,
@ -73728,6 +73771,7 @@
"output_cost_per_token_priority": 1.875e-05,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -75309,6 +75353,7 @@
"output_cost_per_token_above_272k_tokens": 1.65e-05,
"output_cost_per_token_above_272k_tokens_ultrafast": 9.9e-05,
"output_cost_per_token_ultrafast": 6.6e-05,
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html",
"supported_endpoints": [
"/v1/chat/completions",
@ -75331,6 +75376,7 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true,
"supports_xhigh_reasoning_effort": true,
"use_openai_responses_path": true
},
@ -77228,6 +77274,7 @@
"output_cost_per_token_priority": 1.5e-06,
"source": "https://api.fireworks.ai/v1/serverless/models?format=nested",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -77386,6 +77433,7 @@
"output_cost_per_token_priority": 9e-07,
"source": "https://platform.minimax.io/docs/guides/pricing",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -79019,6 +79067,7 @@
"mode": "responses",
"output_cost_per_token": 5.5e-06,
"output_cost_per_token_above_272k_tokens": 8.25e-06,
"search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012},
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-terra.html",
"supported_endpoints": [
"/v1/chat/completions",
@ -79041,6 +79090,7 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true,
"supports_xhigh_reasoning_effort": true,
"use_openai_responses_path": true
},
@ -81161,6 +81211,7 @@
],
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -81177,6 +81228,7 @@
"output_cost_per_token": 3.08e-06,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -81193,6 +81245,7 @@
"output_cost_per_token": 8.4e-07,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -81209,6 +81262,7 @@
"output_cost_per_token": 8.4e-07,
"source": "https://api.fireworks.ai/v1/serverless/models",
"supports_function_calling": true,
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,

13
tests/e2e/CONTRIBUTING.md
View file

@ -90,18 +90,7 @@ The suites run against a live proxy, so bring one up first by running the litell
They also need a proxy whose bundled UI contains the change under test, so run the proxy from your branch (an editable install serves the UI your checkout builds)
Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy. The presidio guardrail tests need a running Presidio analyzer and anonymizer the proxy can reach, addressed by `PRESIDIO_ANALYZER_API_BASE` / `PRESIDIO_ANONYMIZER_API_BASE`
The Presidio spend-log audit test also requires prompt storage on the proxy, because its assertions inspect the detected entities and their scores. Add the following to the proxy config before starting it:
```yaml
general_settings:
store_prompts_in_spend_logs: true
```
Alternatively, start the proxy with `STORE_PROMPTS_IN_SPEND_LOGS=true litellm --config <your-e2e-config>.yml --port 4000`. Setting the variable only on the pytest process does not configure the proxy. With prompt storage disabled, the proxy correctly redacts `guardrail_response`, so that configuration cannot exercise this test's entity-detail assertions. Enable this only on a test stack using synthetic prompts
A couple of logging destinations are configured on the proxy rather than by the test. The Weave tests scope their callback to the key they create, but litellm builds the `weave_otel` logger from `WANDB_API_KEY` and `WANDB_PROJECT_ID` before it applies the per-key vars, so the proxy needs both in its own environment or the key-scoped callback never initializes and nothing ships
Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and start the proxy with `LITELLM_OTEL_V2=true`, the `otel` callback, and `OTEL_EXPORTER_OTLP_ENDPOINT` pointing at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy.
### Secret manager lanes

128
tests/e2e/a2a/a2a_client.py
View file

@ -18,7 +18,7 @@ from dataclasses import dataclass
from pydantic import BaseModel, ConfigDict, Field
from e2e_config import settle_propagation
from e2e_http import NoBody, Result, Success, get_external, is_ok
from e2e_http import NoBody, Result, Success, is_ok
from e2e_metadata import STEP_FRAMES, step
from proxy_client import ProxyClient
@ -36,11 +36,6 @@ class A2ASkill(BaseModel):
examples: list[str] | None = None
class A2AProvider(BaseModel):
organization: str
url: str
class AgentCardParams(BaseModel):
"""The upstream agent card an admin registers. `protocolVersion` is the field the
proxy validates against SUPPORTED_A2A_PROTOCOL_VERSIONS on registration."""
@ -57,29 +52,6 @@ class AgentCardParams(BaseModel):
preferred_transport: str | None = Field(default=None, serialization_alias="preferredTransport")
class UpstreamAgentCard(BaseModel):
"""A real published agent card parsed from a public /.well-known endpoint. Keys on
the A2A wire aliases so `model_validate_json` reads the served JSON and
`model_dump(by_alias=True)` re-emits it unchanged for verbatim registration; it is
only ever fetched-and-validated, never hand-constructed, so aliasing on the wire
names does not affect any call site."""
model_config = ConfigDict(populate_by_name=True)
protocol_version: str = Field(alias="protocolVersion")
name: str
description: str
version: str
url: str
provider: A2AProvider | None = None
documentation_url: str | None = Field(default=None, alias="documentationUrl")
capabilities: A2ACapabilities = A2ACapabilities()
skills: list[A2ASkill]
default_input_modes: list[str] = Field(default=["text"], alias="defaultInputModes")
default_output_modes: list[str] = Field(default=["text"], alias="defaultOutputModes")
preferred_transport: str | None = Field(default=None, alias="preferredTransport")
class A2ABridgeParams(BaseModel):
"""litellm_params that route the agent through the completion bridge: an A2A
message/send is transformed into a litellm.acompletion against this provider."""
@ -93,7 +65,7 @@ class A2ABridgeParams(BaseModel):
class AgentRegisterBody(BaseModel):
agent_name: str
agent_card_params: AgentCardParams | UpstreamAgentCard
agent_card_params: AgentCardParams
litellm_params: A2ABridgeParams | None = None
@ -136,32 +108,9 @@ class A2ATextPart(BaseModel):
text: str
class A2ASearchPropertiesParams(BaseModel):
"""The strict param schema of the published property agent's `search_properties`
skill (unknown keys are rejected upstream), so a natural-language query like
"properties for sale in SF under $2M" is expressed as typed fields."""
un_locode: str | None = None
service_type: str | None = None
property_type: str | None = None
bedrooms_min: int | None = None
asking_price_max: float | None = None
limit: int | None = None
class A2ASkillInvocation(BaseModel):
skill: str
params: A2ASearchPropertiesParams
class A2ADataPart(BaseModel):
kind: str = "data"
data: A2ASkillInvocation
class A2AOutboundMessage(BaseModel):
role: str = "user"
parts: list[A2ATextPart | A2ADataPart]
parts: list[A2ATextPart]
message_id: str = Field(serialization_alias="messageId")
@ -189,57 +138,10 @@ class A2AResponseMessage(BaseModel):
parts: list[A2AResponsePart] = []
class A2ATaskStatus(BaseModel):
state: str | None = None
message: A2AResponseMessage | None = None
class A2AListingLocation(BaseModel):
"""Only the location fields a test reads back off a returned listing."""
un_locode: str | None = None
class A2AListing(BaseModel):
"""A single property card from the agent's `search_results` artifact; only the
identity/location fields a test asserts on are modelled."""
raia_id: str
property_type: str | None = None
service_type: str | None = None
location: A2AListingLocation = A2AListingLocation()
class A2ASearchResults(BaseModel):
"""The DataPart payload the property agent's `search_properties` skill returns:
the run count plus the listing cards themselves. Proof the tool actually ran and
matched, not just that the task completed with some text."""
total: int
count: int
listings: list[A2AListing] = []
class A2AArtifactPart(BaseModel):
kind: str | None = None
data: A2ASearchResults | None = None
class A2AArtifact(BaseModel):
model_config = ConfigDict(populate_by_name=True)
artifact_id: str | None = Field(default=None, alias="artifactId")
name: str | None = None
parts: list[A2AArtifactPart] = []
class A2AResult(BaseModel):
"""A message/send result. In 0.3 the message fields sit directly on the result
(`kind`/`role`/`parts`); in 1.0 they are nested under `message`; a real agent that
runs a task replies with a `task` whose agent text lives on `status.message` and
whose tool output lives on `artifacts`. `text` reads the agent's reply from
whichever shape the served version produced; `search_results` reads the tool's
structured output when the agent ran a skill."""
(`kind`/`role`/`parts`); in 1.0 they are nested under `message`. `text` reads the
agent's reply from whichever shape the served version produced."""
model_config = ConfigDict(populate_by_name=True)
@ -248,8 +150,6 @@ class A2AResult(BaseModel):
message_id: str | None = Field(default=None, alias="messageId")
parts: list[A2AResponsePart] = []
message: A2AResponseMessage | None = None
status: A2ATaskStatus | None = None
artifacts: list[A2AArtifact] = []
@property
def text(self) -> str:
@ -257,8 +157,6 @@ class A2AResult(BaseModel):
parts = self.message.parts
elif self.parts:
parts = self.parts
elif self.status is not None and self.status.message is not None:
parts = self.status.message.parts
else:
parts = []
return "".join(part.text or "" for part in parts)
@ -267,14 +165,6 @@ class A2AResult(BaseModel):
def is_nested_v1_shape(self) -> bool:
return self.message is not None
@property
def search_results(self) -> A2ASearchResults | None:
for artifact in self.artifacts:
for part in artifact.parts:
if part.data is not None:
return part.data
return None
class A2AError(BaseModel):
code: int
@ -380,11 +270,3 @@ class A2AClient:
def build_a2a_client(proxy: ProxyClient) -> A2AClient:
return A2AClient(proxy=proxy)
@step("Fetch a published A2A agent card from its /.well-known endpoint")
def fetch_agent_card(url: str, *, timeout: float = 20.0) -> Result[UpstreamAgentCard]:
"""Fetch a live A2A agent card from its /.well-known endpoint and parse it into the
registration model, so a test can register a real published card verbatim rather
than a hand-rolled one."""
return get_external(url, response_type=UpstreamAgentCard, timeout=timeout)

49
tests/e2e/a2a/test_a2a_agent_e2e.py
View file

@ -17,18 +17,14 @@ import pytest
from a2a_client import (
A2ABridgeParams,
A2AClient,
A2ADataPart,
A2AJsonRpcRequest,
A2AMessageSendParams,
A2AOutboundMessage,
A2ASearchPropertiesParams,
A2ASkill,
A2ASkillInvocation,
A2ATextPart,
AgentCardParams,
AgentRegisterBody,
AgentResponse,
fetch_agent_card,
)
from e2e_config import unique_marker
from e2e_http import Result, UnknownApiError, unwrap
@ -50,9 +46,6 @@ BRIDGE = A2ABridgeParams(
model=BRIDGE_MODEL,
)
MOVEHOME_AGENT_CARD_URL = "https://movehome.org/.well-known/agent.json"
MOVEHOME_ORIGIN = "https://movehome.org"
pytestmark = pytest.mark.e2e
@ -135,48 +128,6 @@ class TestA2AAgentLifecycle:
assert result is not None
assert result.text != ""
@pytest.mark.covers("other.a2a.message_send.real_world_agent_replies")
@meta(
Subject(
domain=Domain.AGENTS_API,
route=Route.A2A,
)
)
def test_real_world_agent_replies_to_property_query(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None:
upstream = unwrap(fetch_agent_card(MOVEHOME_AGENT_CARD_URL)).model_copy(update={"url": MOVEHOME_ORIGIN})
assert upstream.protocol_version == "0.3.0"
marker = unique_marker()
body = AgentRegisterBody(agent_name=f"e2e-a2a-real-{marker}", agent_card_params=upstream)
agent = unwrap(client.register_agent(body))
resources.defer(lambda: client.delete_agent(agent.agent_id))
assert agent.agent_card_params.protocol_version == "0.3"
location = "GBLON"
request = A2AJsonRpcRequest(
id=f"e2e-{unique_marker()}",
params=A2AMessageSendParams(
message=A2AOutboundMessage(
parts=[
A2ADataPart(
data=A2ASkillInvocation(
skill="search_properties",
params=A2ASearchPropertiesParams(un_locode=location, service_type="long_term", limit=3),
)
)
],
message_id=unique_marker(),
)
),
)
response = unwrap(client.send_message(agent.agent_id, scoped_key, request))
assert response.error is None
assert response.result is not None
results = response.result.search_results
assert results is not None, "agent returned no search_results artifact; skill did not run"
assert results.total > 0
assert results.listings, "search_properties matched nothing; agent returned no property cards"
assert all(listing.raia_id for listing in results.listings)
assert all(listing.location.un_locode == location for listing in results.listings)
@pytest.mark.covers("other.a2a.discovery.proxy_fronted_card")
@meta(
Subject(

5
tests/e2e/coverage_registry/guardrail.yaml
View file

@ -1,11 +1,6 @@
# Guardrail enforcement (behavior features). Grounded in litellm/proxy/guardrails/guardrail_hooks/.
# Rolls up into the "Logging & Guardrails" dashboard module together with logging.*
- {id: guardrail.presidio.pre_call.masks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "PII masking pre-call; data-leak blast radius"}
- {id: guardrail.presidio.post_call.masks, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Mask PII in model output"}
- {id: guardrail.presidio.post_call.masks_generated_output, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, chat_completions_stream, anthropic_messages_stream], source: "guardrail_hooks/presidio.py", rationale: "Mask model-generated credit-card output with the UI default scope"}
- {id: guardrail.presidio.post_call.spend_log_stores_masked_output, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, chat_completions_stream, messages, anthropic_messages_stream, responses], source: "guardrail_hooks/presidio.py", fail_before_fix: proven, rationale: "When an output guardrail masks the response, the spend log stores the masked text the caller received rather than the raw model output, on every endpoint and both stream modes (LIT-8325)"}
- {id: guardrail.presidio.logging_only.masks, module: guardrail, tier: P0, hook_point: logging_only, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Redact in logs without blocking"}
- {id: guardrail.presidio.pre_call.logs_masked_entities, module: guardrail, tier: P0, hook_point: pre_call, assertions: [logs_masked_entities], exercised_on: [chat_completions], source: "guardrail_hooks/presidio.py", rationale: "A masking run must record itself on the spend log: the dashboard's guardrail panel renders the masked-entity counts and per-entity scores straight off metadata.guardrail_information, so a run that masks but records nothing leaves an operator unable to audit it"}
- {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions, messages, responses], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"}
- {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"}
- {id: guardrail.litellm_content_filter.pre_call.blocks_video, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [videos], source: "test_key_guardrail_video_e2e.py", fail_before_fix: proven, rationale: "A content-filter guardrail attached to a key (metadata.guardrails) blocks a banned prompt on POST /v1/videos before the provider is called; before the fix the route's call type was unknown to the unified guardrail hook and the prompt went to the provider unscanned (LIT-6685)"}

5
tests/e2e/coverage_registry/llm_conversational.yaml
View file

@ -139,11 +139,6 @@
- {id: llm.responses.ollama.basic.stream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses streams text deltas, usage and a terminal event"}
- {id: llm.responses.ollama.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses returns one addressable get_weather call"}
- {id: llm.responses.ollama.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses tool result round trip reaches the model"}
- {id: llm.chat_completions.sail.service_tier.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: sail, capability: service_tier, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_sail_e2e.py", rationale: "service_tier balanced and auto map to Sail completion windows and bill the matching price columns; Sail serves flex only to background responses and Batch"}
- {id: llm.chat_completions.sail.service_tier.nonstream.rejects_unknown_tier, module: llm, tier: P1, subject_endpoint: chat_completions, route: sail, capability: service_tier, streaming: nonstream, assertions: [rejects_unknown_tier], source: "llm_translation/test_sail_e2e.py", rationale: "A service_tier Sail has no completion window for is a 400 without drop_params"}
- {id: llm.chat_completions.sail.service_tier.nonstream.drops_unknown_tier_and_bills_asap, module: llm, tier: P1, subject_endpoint: chat_completions, route: sail, capability: service_tier, streaming: nonstream, assertions: [drops_unknown_tier_and_bills_asap], source: "llm_translation/test_sail_e2e.py", rationale: "An unknown service_tier under drop_params is dropped and billed at asap in both the cost header and spend log"}
- {id: llm.responses.sail.service_tier.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: responses, route: sail, capability: service_tier, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_sail_e2e.py", rationale: "A caller metadata.completion_window of balanced on /v1/responses bills Sail balanced rates"}
- {id: llm.messages.sail.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: sail, capability: basic, streaming: nonstream, assertions: [works], source: "llm_translation/test_sail_e2e.py", rationale: "Sail over /v1/messages"}
- {id: llm.chat_completions.anthropic.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: chat_completions, route: anthropic, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic over /chat/completions: cost header and spend row agree"}
- {id: llm.chat_completions.anthropic.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic tool result round trip over /chat/completions"}
- {id: llm.messages.anthropic.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic tool result round trip over /v1/messages"}

1
tests/e2e/coverage_registry/other.yaml
View file

@ -55,7 +55,6 @@
- {id: other.a2a.register.persists, module: other, tier: P1, area: a2a, assertions: [persists], source: "agent_endpoints/endpoints.py:325-443", rationale: "POST /v1/agents registers an agent card; GET /v1/agents/{id} reads it back"}
- {id: other.a2a.register.unsupported_version_rejected, module: other, tier: P1, area: a2a, assertions: [unsupported_version_rejected], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a protocolVersion outside SUPPORTED_A2A_PROTOCOL_VERSIONS is refused with 400"}
- {id: other.a2a.register.semver_version_accepted, module: other, tier: P1, area: a2a, assertions: [semver_version_accepted], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a patch-level semver like 0.3.0 (what the Google A2A SDK emits) registers, stores and serves the canonical 0.3 rather than 400ing; regression guard for the v1.92 report"}
- {id: other.a2a.message_send.real_world_agent_replies, module: other, tier: P1, area: a2a, assertions: [real_world_agent_replies], source: "agent_endpoints/a2a_endpoints.py asend_message", rationale: "A real published a2a agent fetched live from a public /.well-known endpoint (pinning the full semver 0.3.0 the a2a-sdk emits) registers, serves the canonical 0.3, and a message/send skill invocation proxies to the live upstream and returns the agent's reply"}
- {id: other.a2a.register.malformed_version_rejected, module: other, tier: P1, area: a2a, assertions: [malformed_version_rejected], source: "a2a/agent_card.py normalize_protocol_version", rationale: "A malformed protocolVersion like 0.3.garbage fails full-string semver validation and is refused with 400 instead of truncating to a supported family"}
- {id: other.a2a.discovery.proxy_fronted_card, module: other, tier: P1, area: a2a, assertions: [proxy_fronted_card], source: "agent_endpoints/a2a_endpoints.py get_agent_card", rationale: "/.well-known/agent-card.json serves the proxy url + supportedInterfaces and the LiteLLM virtual-key bearer scheme, not the upstream"}
- {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"}

8
tests/e2e/gateway/stage_mirror_ci_config.yml
View file

@ -32,7 +32,7 @@ litellm_settings:
- host: 127.0.0.1
port: 6379
ssl: true
callbacks: ["arize_phoenix", "datadog", "smtp_email", "prometheus", "otel", "websearch_interception"]
callbacks: ["datadog", "smtp_email", "prometheus", "otel", "websearch_interception"]
websearch_interception_params:
enabled_providers: ["bedrock"]
search_tool_name: e2e-search
@ -97,9 +97,3 @@ files_settings:
finetune_settings:
- custom_llm_provider: openai
api_key: os.environ/OPENAI_API_KEY
mcp_servers:
devin:
url: "https://mcp.devin.ai/mcp"
auth_type: api_key
auth_value: os.environ/DEVIN_API_KEY

19
tests/e2e/guardrails/guardrails_client.py
View file

@ -41,8 +41,6 @@ from pydantic import BaseModel, Field
GUARDRAIL_BACKEND: Final = "gemini/gemini-2.5-flash"
GuardrailMode = Literal["pre_call", "post_call", "during_call", "logging_only"]
PiiEntity = Literal["EMAIL_ADDRESS", "PHONE_NUMBER", "PERSON", "CREDIT_CARD", "US_SSN"]
PiiAction = Literal["MASK", "BLOCK"]
BlockedWordAction = Literal["BLOCK", "MASK"]
@ -80,22 +78,6 @@ class BlockCodeExecutionParamsBody(GuardrailParamsBase):
guardrail: Literal["block_code_execution"] = "block_code_execution"
class PresidioParamsBody(GuardrailParamsBase):
"""Presidio PII guardrail params. `presidio_filter_scope="input"` keeps the
registration to a single callback on the configured mode; the default
("both") also registers a second post_call output-masking callback, which a
pre_call- or logging_only-scoped test must not drag in. `output_parse_pii`
stays unset/False: True would unmask the response back to the caller."""
guardrail: Literal["presidio"] = "presidio"
presidio_analyzer_api_base: str
presidio_anonymizer_api_base: str
presidio_filter_scope: Literal["input", "output", "both"] | None = None
presidio_language: str | None = None
output_parse_pii: bool | None = None
pii_entities_config: dict[PiiEntity, PiiAction] | None = None
class ToolPermissionRuleBody(BaseModel):
"""One tool_permission rule: a decision for the tool named by `tool_name`."""
@ -130,7 +112,6 @@ GuardrailParamsBody = (
| BedrockGuardrailParamsBody
| OpenAIModerationParamsBody
| BlockCodeExecutionParamsBody
| PresidioParamsBody
| ToolPermissionParamsBody
| CustomCodeParamsBody
)

857
tests/e2e/guardrails/test_presidio_masking_e2e.py
View file

@ -1,857 +0,0 @@
"""Live e2e: the Presidio PII guardrail masks, per its configured hook point.
pre_call: the guardrail calls the Presidio analyzer/anonymizer on the request
messages BEFORE the model runs, so the model only ever sees placeholders like
<EMAIL_ADDRESS>. A prompt asking the model to repeat a fake email + phone back
must come back with the placeholders echoed and the raw PII absent, on
/chat/completions and on /v1/messages (Anthropic format).
post_call: the mirror hook. The request reaches the model unmasked and the
MODEL OUTPUT is what gets anonymized, so the caller never receives raw PII the
model repeated back. The two hooks are told apart behaviorally rather than by
configuration: the post_call prompt asks for a value derived from the raw email
(its local part, which is not itself an entity Presidio masks) alongside the
address itself, so the answer proves the model saw the raw address while the
address in the same response comes back as <EMAIL_ADDRESS>.
The analyzer/anonymizer endpoints come from PRESIDIO_ANALYZER_API_BASE /
PRESIDIO_ANONYMIZER_API_BASE; missing env is a hard failure, never a skip.
Each guardrail registers with an explicit presidio_filter_scope so only the
configured hook's callback exists (the default "both" registers input masking
AND a post_call output masker), and is deleted on teardown.
The spend-log audit test requires general_settings.store_prompts_in_spend_logs:
true in the proxy config, or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy
process environment before startup. Setting it only on pytest has no effect.
Without this opt-in, redacting guardrail_response is expected proxy behavior;
this suite deliberately requires the detected-entity details to remain visible.
"""
from __future__ import annotations
import os
import re
import time
from collections.abc import Callable, Iterator
from typing import Final, Literal
import pytest
from pydantic import BaseModel, JsonValue, TypeAdapter
from e2e_config import unique_marker
from e2e_http import Result, StreamingResponse, Success
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
from guardrails_client import GuardrailMode, GuardrailsClient, PiiAction, PiiEntity, PresidioParamsBody
from lifecycle import ResourceManager
from models import (
AnthropicMessagesResponse,
ChatResponse,
GuardrailEntityMatch,
GuardrailRunRecord,
SpendLogRow,
)
pytestmark = pytest.mark.e2e
MODEL = "gemini-2.5-flash"
# A guardrail created via POST /guardrails reaches the worker that served the
# create immediately, but every other worker only picks it up on its next
# periodic DB sync (~30s), so the first requests can be served unguarded.
GUARDRAIL_PROPAGATION_DEADLINE_SECONDS = 40.0
GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS = 5.0
# Presidio's anonymizer replaces a detected entity with its unnumbered type
# placeholder, e.g. <EMAIL_ADDRESS>. The pre_call assertions match on the bare
# token because the model is echoing the masked prompt and may not preserve the
# angle brackets; the logged payload keeps the placeholder verbatim.
MASKED_EMAIL_TOKEN = "EMAIL_ADDRESS"
MASKED_PHONE_TOKEN = "PHONE_NUMBER"
MASKED_CREDIT_CARD_TOKEN = "CREDIT_CARD"
# Fictional NANP 555 number; a standard format Presidio's phone recognizer detects.
FAKE_PHONE = "+1 415-555-0134"
FAKE_VISA_TEST_CARD = "4111 1111 1111 1111"
_CARD_DIGIT_RUN: Final = re.compile(r"(?:\d[ -]?){13,19}")
_CONTENT_KEYS: Final = frozenset({"content", "text"})
def _presidio_bases() -> tuple[str, str]:
analyzer = os.environ.get("PRESIDIO_ANALYZER_API_BASE", "").strip()
anonymizer = os.environ.get("PRESIDIO_ANONYMIZER_API_BASE", "").strip()
if not analyzer or not anonymizer:
pytest.fail(
"Presidio e2e requires PRESIDIO_ANALYZER_API_BASE and PRESIDIO_ANONYMIZER_API_BASE "
"(the running Presidio analyzer/anonymizer services); missing env is a hard failure, not a skip"
)
return analyzer, anonymizer
def _register_presidio(
client: GuardrailsClient,
resources: ResourceManager,
*,
name: str,
mode: GuardrailMode | list[GuardrailMode] = "pre_call",
filter_scope: Literal["input", "output", "both"] | None = "input",
entities: dict[PiiEntity, PiiAction] | None = None,
) -> None:
analyzer, anonymizer = _presidio_bases()
guardrail_id = client.register(
name,
PresidioParamsBody(
mode=mode,
default_on=False,
presidio_analyzer_api_base=analyzer,
presidio_anonymizer_api_base=anonymizer,
presidio_filter_scope=filter_scope,
pii_entities_config=entities,
),
)
resources.defer(lambda: client.delete_guardrail(guardrail_id))
def _fake_email() -> str:
return f"jane.doe.{unique_marker()}@example.com"
def _pii_prompt(marker: str, email: str) -> str:
return (
f"{marker} Repeat this sentence back to me exactly, word for word: "
f"My email address is {email} and my phone number is {FAKE_PHONE}."
)
def _first_content(response: ChatResponse) -> str:
if not response.choices:
return ""
message = response.choices[0].message
return (message.content if message else None) or ""
class _StreamDelta(BaseModel):
content: str | None = None
class _StreamChoice(BaseModel):
delta: _StreamDelta
class _StreamChunk(BaseModel):
choices: tuple[_StreamChoice, ...] = ()
class _AnthropicStreamDelta(BaseModel):
type: str | None = None
text: str | None = None
class _AnthropicStreamEvent(BaseModel):
type: str
delta: _AnthropicStreamDelta | None = None
def _credit_card_prompt(marker: str) -> str:
return (
f"{marker} Reply with only the well known Visa sandbox test card number that starts with 4111, "
"the 16 digits grouped in fours separated by spaces, and nothing else."
)
def _passes_luhn(digits: str) -> bool:
checksum = sum(
digit if position % 2 == 0 else (digit * 2 - 9 if digit * 2 > 9 else digit * 2)
for position, digit in enumerate(int(char) for char in reversed(digits))
)
return checksum % 10 == 0
def _contains_card_number(text: str) -> bool:
"""Presidio's CREDIT_CARD recognizer only reports Luhn-valid digit runs, so a
Luhn-invalid number the model hallucinates is not something masking can catch."""
return any(
13 <= len(digits) <= 19 and _passes_luhn(digits)
for digits in (re.sub(r"[ -]", "", match.group()) for match in _CARD_DIGIT_RUN.finditer(text))
)
def _stream_content(result: StreamingResponse) -> str:
return "".join(
choice.delta.content
for event in result.stream_events
if event != "[DONE]"
for choice in _StreamChunk.model_validate_json(event).choices[:1]
if choice.delta.content
)
def _anthropic_stream_content(result: StreamingResponse) -> str:
return "".join(
event.delta.text
for payload in result.stream_events
for event in [_AnthropicStreamEvent.model_validate_json(payload)]
if event.type == "content_block_delta"
and event.delta is not None
and event.delta.type == "text_delta"
and event.delta.text
)
def _messages_text(response: AnthropicMessagesResponse) -> str:
"""The text of a /v1/messages answer, whichever shape the proxy produced
(Anthropic-native content blocks or OpenAI-normalized choices)."""
parts: list[str] = []
for block in response.content or []:
if block.text:
parts.append(block.text)
for choice in response.choices or []:
if choice.message and choice.message.content:
parts.append(choice.message.content)
return "\n".join(parts)
def _assert_eventually_masked[R: BaseModel](
fetch: Callable[[], Result[R]], extract: Callable[[R], str], *, email: str
) -> None:
"""Retry the call until the response comes back masked, to the propagation
deadline. An unmasked early response is in-flight guardrail propagation, not
a failure, and neither is a transient non-Success (a replica that has not
reloaded the guardrail answers 404, the live model can rate-limit) - only a
response that still carries the raw PII at the deadline is."""
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
last: str = "<no successful response yet>"
while True:
result = fetch()
match result:
case Success(data=data):
content = extract(data)
last = content
masked = MASKED_EMAIL_TOKEN in content and MASKED_PHONE_TOKEN in content and email not in content
if masked:
assert FAKE_PHONE not in content, (
f"the raw phone number must be masked before the model sees it, but the "
f"response echoed it: {content[:300]!r}"
)
return
case _:
last = f"<non-Success result: {result}>"
if time.monotonic() >= deadline:
pytest.fail(
f"presidio pre_call guardrail never masked the PII within "
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
)
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
class TestPresidioPreCallMasking:
@pytest.mark.covers(
"guardrail.presidio.pre_call.masks",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_pre_call_masks_pii_on_chat_completions(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-pre-chat-{unique_marker()}"
_register_presidio(client, resources, name=name)
email = _fake_email()
prompt = _pii_prompt(unique_marker(), email)
_assert_eventually_masked(
lambda: client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128),
_first_content,
email=email,
)
@pytest.mark.covers(
"guardrail.presidio.pre_call.masks",
exercised_on=["messages"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.MESSAGES,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_pre_call_masks_pii_on_messages(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-pre-msg-{unique_marker()}"
_register_presidio(client, resources, name=name)
email = _fake_email()
prompt = _pii_prompt(unique_marker(), email)
_assert_eventually_masked(
lambda: client.messages(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128),
_messages_text,
email=email,
)
#: Room for the model's reasoning tokens plus the three-line answer; a lower cap
#: truncates the response before the address it is supposed to mask.
_POST_CALL_MAX_TOKENS = 512
#: The post_call scenario masks these two entities and nothing else. Left
#: unscoped, Presidio's broader recognizers claim the local part too (a random
#: marker reads as an NRP), which would erase the very token that tells output
#: masking apart from input masking.
_POST_CALL_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"}
def _post_call_prompt(marker: str, local_part: str) -> str:
"""Ask for the local part and the full address in one answer. Presidio masks
an EMAIL_ADDRESS entity and a bare local part is not one, so the two land
differently in the same response and pin the hook point behaviorally."""
return (
f"{marker} My email address is {local_part}@example.com and my phone number is {FAKE_PHONE}. "
"Reply with exactly three lines and nothing else. "
"Line 1: the part of the email address before the @ sign. "
"Line 2: the full email address. "
"Line 3: the phone number."
)
class TestPresidioPostCallMasking:
@pytest.mark.covers(
"guardrail.presidio.post_call.masks",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_post_call_masks_pii_in_model_output(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
"""A guardrail scoped to the output must anonymize the PII the model
repeats back, so a caller (or a downstream log of the response) never
receives it, while the request itself reaches the model untouched.
Both facts are asserted from one response: the local part comes back raw,
which is only possible if the model saw the real address, and the address
itself comes back as <EMAIL_ADDRESS> in the same answer.
"""
name = f"e2e-presidio-post-chat-{unique_marker()}"
_register_presidio(
client,
resources,
name=name,
mode="post_call",
filter_scope="output",
entities=_POST_CALL_ENTITIES,
)
local_part = f"e2euser{unique_marker()}"
email = f"{local_part}@example.com"
prompt = _post_call_prompt(unique_marker(), local_part)
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
last = "<no successful response yet>"
while True:
result = client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=_POST_CALL_MAX_TOKENS)
match result:
case Success(data=data):
last = _first_content(data)
if MASKED_EMAIL_TOKEN in last and email not in last:
assert local_part in last, (
"the model must have seen the RAW address (it is asked for the local "
"part, which Presidio does not mask); the local part is missing, so "
f"this response cannot tell post_call masking from pre_call: {last[:300]!r}"
)
assert MASKED_PHONE_TOKEN in last and FAKE_PHONE not in last, (
f"the phone number in the model's answer must be masked too, got: {last[:300]!r}"
)
return
case _:
last = f"<non-Success result: {result}>"
if time.monotonic() >= deadline:
pytest.fail(
f"presidio post_call guardrail never masked the model's output within "
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
)
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
def _assert_eventually_masks_generated_card(fetch: Callable[[], str | None]) -> None:
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
last: str = "<no successful response yet>"
while True:
content = fetch()
if content is not None:
last = content
if _contains_card_number(content):
pytest.fail(
"the post_call output masking let a card number through: "
f"{content[:300]!r}"
)
if MASKED_CREDIT_CARD_TOKEN in content:
return
if time.monotonic() >= deadline:
pytest.fail(
"presidio post_call output masking never masked the generated card within "
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
)
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
class TestPresidioCreditCardOutputMasking:
"""Proves the UI-default Presidio scope masks model-generated card output."""
@pytest.mark.covers(
"guardrail.presidio.post_call.masks_generated_output",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_ui_default_scope_masks_a_card_number_the_model_generates_on_chat_completions(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-card-chat-{unique_marker()}"
_register_presidio(
client,
resources,
name=name,
mode=["pre_call", "post_call"],
filter_scope=None,
entities={"CREDIT_CARD": "MASK"},
)
prompt: Final = _credit_card_prompt(unique_marker())
def fetch() -> str | None:
result: Final = client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=512)
match result:
case Success(data=data):
return _first_content(data)
case _:
return None
_assert_eventually_masks_generated_card(fetch)
@pytest.mark.covers(
"guardrail.presidio.post_call.masks_generated_output",
exercised_on=["chat_completions_stream"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.STREAM,
)
)
def test_ui_default_scope_masks_a_card_number_the_model_generates_on_streaming_chat_completions(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-card-stream-{unique_marker()}"
_register_presidio(
client,
resources,
name=name,
mode=["pre_call", "post_call"],
filter_scope=None,
entities={"CREDIT_CARD": "MASK"},
)
prompt: Final = _credit_card_prompt(unique_marker())
def fetch() -> str | None:
result: Final = client.chat_stream_raw(
scoped_key,
MODEL,
prompt,
guardrails=[name],
max_tokens=512,
)
if not result.ok or result.stream_error:
return None
return _stream_content(result)
_assert_eventually_masks_generated_card(fetch)
@pytest.mark.covers(
"guardrail.presidio.post_call.masks_generated_output",
exercised_on=["anthropic_messages_stream"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.MESSAGES,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.STREAM,
)
)
def test_ui_default_scope_masks_a_card_number_the_model_generates_on_streaming_anthropic_messages(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-card-messages-stream-{unique_marker()}"
_register_presidio(
client,
resources,
name=name,
mode=["pre_call", "post_call"],
filter_scope=None,
entities={"CREDIT_CARD": "MASK"},
)
prompt: Final = _credit_card_prompt(unique_marker())
def fetch() -> str | None:
result: Final = client.messages_stream_raw(
scoped_key,
MODEL,
prompt,
guardrails=[name],
max_tokens=512,
)
if not result.ok or result.stream_error:
return None
return _anthropic_stream_content(result)
_assert_eventually_masks_generated_card(fetch)
def _wire_text(outcome: StreamingResponse) -> str:
return "\n".join(outcome.stream_events) if outcome.is_streaming else outcome.body
def _poll_until_generated_card_masked(fetch: Callable[[], StreamingResponse]) -> StreamingResponse:
"""The raw HTTP outcome once the output masker is in effect on the serving
worker: whichever wire shape the endpoint speaks, a masked body carries the
CREDIT_CARD placeholder and no Luhn-valid card run. A raw card is a worker
that has not loaded the guardrail yet, so it is polled through like any
other unmasked answer."""
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
last: str = "<no successful response yet>"
while True:
outcome = fetch()
if outcome.ok and not outcome.stream_error:
wire = _wire_text(outcome)
last = wire
if MASKED_CREDIT_CARD_TOKEN in wire and not _contains_card_number(wire):
return outcome
if time.monotonic() >= deadline:
pytest.fail(
"presidio post_call output masking never masked the generated card within "
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
)
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
def _spend_log_response_text(client: GuardrailsClient, key: str, call_id: str) -> str:
rows = client.proxy.poll_logs_for_key(
key,
predicate=lambda logged: any(row.litellm_call_id == call_id for row in logged),
)
row = next((row for row in rows if row.litellm_call_id == call_id), None)
assert row is not None, f"no spend log row ever appeared for x-litellm-call-id {call_id}"
return "\n".join(_stored_content(row.response))
def _stored_content(node: JsonValue, key: str | None = None) -> Iterator[str]:
match node:
case str() if key in _CONTENT_KEYS:
yield node
case dict():
for child_key, child in node.items():
yield from _stored_content(child, child_key)
case list():
for item in node:
yield from _stored_content(item, key)
case _:
return
class TestPresidioSpendLogStoresMaskedOutput:
"""The spend log stores the response the caller received, on every endpoint
and both stream modes, when an output-only post_call Presidio guardrail masks
a card number the model generated."""
_CELL: Final = "guardrail.presidio.post_call.spend_log_stores_masked_output"
def _assert_spend_log_is_masked(
self,
client: GuardrailsClient,
resources: ResourceManager,
key: str,
*,
name: str,
fetch: Callable[[str, str], StreamingResponse],
) -> None:
_register_presidio(
client,
resources,
name=name,
mode="post_call",
filter_scope="output",
entities={"CREDIT_CARD": "MASK"},
)
prompt: Final = _credit_card_prompt(unique_marker())
outcome = _poll_until_generated_card_masked(lambda: fetch(prompt, name))
assert outcome.call_id, f"the served response must carry x-litellm-call-id: {dict(outcome.headers)}"
logged = _spend_log_response_text(client, key, outcome.call_id)
assert not _contains_card_number(logged), (
"the caller got the masked response but the spend log stored the raw model output: "
f"{logged[:400]!r}"
)
assert MASKED_CREDIT_CARD_TOKEN in logged, (
f"the spend log response carries neither the card nor the placeholder: {logged[:400]!r}"
)
@pytest.mark.covers(_CELL, exercised_on=["chat_completions"])
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_spend_log_stores_masked_output_on_chat_completions(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
self._assert_spend_log_is_masked(
client,
resources,
scoped_key,
name=f"e2e-presidio-log-card-chat-{unique_marker()}",
fetch=lambda prompt, guardrail: client.chat_raw(
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
),
)
@pytest.mark.covers(_CELL, exercised_on=["chat_completions_stream"])
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.STREAM,
)
)
def test_spend_log_stores_masked_output_on_streaming_chat_completions(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
self._assert_spend_log_is_masked(
client,
resources,
scoped_key,
name=f"e2e-presidio-log-card-chat-stream-{unique_marker()}",
fetch=lambda prompt, guardrail: client.chat_stream_raw(
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
),
)
@pytest.mark.covers(_CELL, exercised_on=["messages"])
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.MESSAGES,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_spend_log_stores_masked_output_on_anthropic_messages(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
self._assert_spend_log_is_masked(
client,
resources,
scoped_key,
name=f"e2e-presidio-log-card-messages-{unique_marker()}",
fetch=lambda prompt, guardrail: client.messages_raw(
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
),
)
@pytest.mark.covers(_CELL, exercised_on=["anthropic_messages_stream"])
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.MESSAGES,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.STREAM,
)
)
def test_spend_log_stores_masked_output_on_streaming_anthropic_messages(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
self._assert_spend_log_is_masked(
client,
resources,
scoped_key,
name=f"e2e-presidio-log-card-messages-stream-{unique_marker()}",
fetch=lambda prompt, guardrail: client.messages_stream_raw(
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
),
)
@pytest.mark.covers(_CELL, exercised_on=["responses"])
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.RESPONSES,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_spend_log_stores_masked_output_on_responses(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
self._assert_spend_log_is_masked(
client,
resources,
scoped_key,
name=f"e2e-presidio-log-card-responses-{unique_marker()}",
fetch=lambda prompt, guardrail: client.responses(scoped_key, MODEL, prompt, guardrails=[guardrail]),
)
_LOGGED_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"}
_ENTITY_LIST_ADAPTER: Final = TypeAdapter(list[GuardrailEntityMatch])
def _applied_guardrails(outcome: StreamingResponse) -> str:
return outcome.headers.get("x-litellm-applied-guardrails", "")
def _guardrail_records(row: SpendLogRow) -> tuple[GuardrailRunRecord, ...]:
metadata = row.metadata
if metadata is None:
return ()
return tuple(metadata.guardrail_information or ())
def _poll_until_guardrail_applied(
client: GuardrailsClient, key: str, guardrail_name: str, prompt: str
) -> StreamingResponse:
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128)
while time.monotonic() < deadline:
if last.ok and guardrail_name in _applied_guardrails(last):
return last
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128)
return last
class TestPresidioSpendLogRecord:
@pytest.mark.covers(
"guardrail.presidio.pre_call.logs_masked_entities",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_masking_run_is_recorded_on_the_spend_log(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-log-{unique_marker()}"
_register_presidio(client, resources, name=name, entities=_LOGGED_ENTITIES)
email = _fake_email()
prompt = _pii_prompt(unique_marker(), email)
outcome = _poll_until_guardrail_applied(client, scoped_key, name, prompt)
assert outcome.ok, f"the guarded call must be served, got {outcome.status_code}: {outcome.body[:400]}"
assert name in _applied_guardrails(outcome), (
"the response must carry x-litellm-applied-guardrails naming the guardrail; without it "
f"the 200 only proves the guardrail never attached. Got {_applied_guardrails(outcome)!r}"
)
request_id = ChatResponse.model_validate_json(outcome.body).id
assert request_id, f"the served response must carry an id to look the spend log up by: {outcome.body[:400]}"
rows = client.proxy.poll_logs_for_request_id(
request_id,
predicate=lambda logged: bool(_guardrail_records(logged[0])),
)
assert rows, f"no spend log row ever appeared for request {request_id}"
records = _guardrail_records(rows[0])
assert records, (
f"the spend log for {request_id} carries no guardrail_information, so the dashboard's "
"guardrail panel would render nothing for a request the guardrail demonstrably ran on"
)
record = next(
(entry for entry in records if entry.guardrail_name == name and entry.guardrail_mode == "pre_call"),
None,
)
assert record is not None, (
"guardrail_information carries no pre_call record for this guardrail, only "
f"{[(entry.guardrail_name, entry.guardrail_mode) for entry in records]}"
)
assert record.guardrail_status == "success", (
f"the recorded status must be success for a run that masked and served, got {record.guardrail_status!r}"
)
assert record.guardrail_provider == "presidio", (
f"the record must attribute the run to presidio so the dashboard picks the right "
f"renderer, got {record.guardrail_provider!r}"
)
counts = record.masked_entity_count or {}
assert _LOGGED_ENTITIES.keys() <= counts.keys(), (
f"every entity the guardrail was configured to mask must appear in masked_entity_count, got {counts}"
)
assert all(counts[entity] >= 1 for entity in _LOGGED_ENTITIES), (
f"each masked entity must be counted at least once, got {counts}"
)
assert not isinstance(record.guardrail_response, str), (
"This audit test requires general_settings.store_prompts_in_spend_logs: true in the proxy config "
"or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy process environment before startup "
"(not just the pytest environment). Default prompt redaction is valid proxy behavior, "
f"but prevents entity-detail assertions; got guardrail_response={record.guardrail_response!r}"
)
entities = _ENTITY_LIST_ADAPTER.validate_python(record.guardrail_response)
assert entities, (
"guardrail_response must carry the detected entities; the dashboard's Detected Entities "
"list and its per-entity scores are rendered from exactly this array"
)
assert {entity.entity_type for entity in entities} >= _LOGGED_ENTITIES.keys(), (
f"the detected entities must cover what was masked, got "
f"{sorted(entity.entity_type for entity in entities)}"
)
for entity in entities:
assert 0.0 < entity.score <= 1.0, f"{entity.entity_type} carries an out-of-range score: {entity.score}"
assert 0 <= entity.start < entity.end, (
f"{entity.entity_type} carries a degenerate span: {entity.start}-{entity.end}"
)

74
tests/e2e/llm_translation/test_messages_beta_headers_e2e.py
View file

@ -15,48 +15,85 @@ pytestmark = pytest.mark.e2e
_BETA_CONFIG: Final = files("litellm").joinpath("anthropic_beta_headers_config.json")
_ANTHROPIC_MODEL: Final = "anthropic/claude-sonnet-5-5"
_VERTEX_MODEL: Final = "vertex_ai/claude-sonnet-5-5"
_VERTEX_LOCATION: Final = "global"
_AZURE_AI_MODEL: Final = "azure_ai/claude-fable-5-1"
_BEDROCK_MODEL_ID: Final = "us.anthropic.claude-fable-5-1"
_BETA_CONFIG_FILE: Final = TypeAdapter(dict[str, JsonValue])
_BETA_MAPPING: Final = TypeAdapter(dict[str, str | None])
_CODE_EXECUTION: Final[ToolUnionParam] = {"type": "code_execution_20250825", "name": "code_execution"}
# A live Vertex probe on 2026-10-31: each Claude 4.5 and 4.6 model takes one of these and Claude 5.5 takes neither
_VERTEX_PER_MODEL_BETAS: Final = frozenset({"computer-use-2025-01-24", "computer-use-2025-11-24"})
_NO_PER_MODEL_BETAS: Final[frozenset[str]] = frozenset()
def _mapped_beta_headers(provider: str) -> str:
mapping: Final = _BETA_MAPPING.validate_python(_BETA_CONFIG_FILE.validate_json(_BETA_CONFIG.read_text())[provider])
return ",".join(name for name, value in mapping.items() if value is not None)
def _mapped_beta_headers(provider: str, per_model: frozenset[str]) -> str:
config: Final = _BETA_CONFIG_FILE.validate_json(_BETA_CONFIG.read_text())
mapping: Final = _BETA_MAPPING.validate_python(config[provider])
header_only: Final = _BETA_MAPPING.validate_python(config.get(f"{provider}_header_only", {}))
forwarded: Final = {name for name, value in {**mapping, **header_only}.items() if value is not None}
assert per_model <= forwarded, per_model - forwarded
return ",".join(sorted(forwarded - per_model))
def _deployment(backend: str) -> LiteLLMParamsBody:
if backend.startswith("anthropic/"):
return LiteLLMParamsBody(model=backend, api_key="os.environ/ANTHROPIC_API_KEY")
return LiteLLMParamsBody(
model=backend,
aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID",
aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY",
aws_region_name="us-east-1",
)
match backend.split("/", 1)[0]:
case "anthropic":
return LiteLLMParamsBody(model=backend, api_key="os.environ/ANTHROPIC_API_KEY")
case "vertex_ai":
return LiteLLMParamsBody(
model=backend,
vertex_project="os.environ/VERTEXAI_PROJECT",
vertex_location=_VERTEX_LOCATION,
vertex_credentials="os.environ/VERTEXAI_CREDENTIALS",
)
case "azure_ai":
return LiteLLMParamsBody(
model=backend,
api_base="os.environ/AZURE_AI_API_BASE",
api_key="os.environ/AZURE_AI_API_KEY",
)
case _:
return LiteLLMParamsBody(
model=backend,
aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID",
aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY",
aws_region_name="us-east-1",
)
class TestAnthropicBetaHeaders:
@pytest.mark.parametrize(
("backend", "beta_provider", "tools"),
("backend", "beta_provider", "tools", "per_model"),
[
pytest.param(_ANTHROPIC_MODEL, "anthropic", [_CODE_EXECUTION], id="anthropic"),
pytest.param(f"bedrock/invoke/{_BEDROCK_MODEL_ID}", "bedrock", [], id="bedrock-invoke"),
pytest.param(f"bedrock/converse/{_BEDROCK_MODEL_ID}", "bedrock_converse", [], id="bedrock-converse"),
pytest.param(_ANTHROPIC_MODEL, "anthropic", [_CODE_EXECUTION], _NO_PER_MODEL_BETAS, id="anthropic"),
pytest.param(_VERTEX_MODEL, "vertex_ai", [], _VERTEX_PER_MODEL_BETAS, id="vertex-ai"),
pytest.param(_AZURE_AI_MODEL, "azure_ai", [], _NO_PER_MODEL_BETAS, id="azure-ai"),
pytest.param(
f"bedrock/invoke/{_BEDROCK_MODEL_ID}", "bedrock", [], _NO_PER_MODEL_BETAS, id="bedrock-invoke"
),
pytest.param(
f"bedrock/converse/{_BEDROCK_MODEL_ID}",
"bedrock_converse",
[],
_NO_PER_MODEL_BETAS,
id="bedrock-converse",
),
],
)
@pytest.mark.covers(
"llm.messages.anthropic.basic.nonstream.works",
"llm.messages.bedrock_invoke.basic.nonstream.works",
"llm.messages.bedrock_converse.basic.nonstream.works",
"llm.messages.vertex.basic.nonstream.works",
"llm.messages.azure_foundry.basic.nonstream.works",
)
@meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.MESSAGES,
providers=(Provider.ANTHROPIC, Provider.BEDROCK),
models=(_ANTHROPIC_MODEL, _BEDROCK_MODEL_ID),
providers=(Provider.ANTHROPIC, Provider.BEDROCK, Provider.VERTEX_AI, Provider.AZURE_AI),
models=(_ANTHROPIC_MODEL, _BEDROCK_MODEL_ID, _VERTEX_MODEL, _AZURE_AI_MODEL),
mode=Mode.NONSTREAM,
)
)
@ -68,6 +105,7 @@ class TestAnthropicBetaHeaders:
backend: str,
beta_provider: str,
tools: list[ToolUnionParam],
per_model: frozenset[str],
) -> None:
model: Final = f"e2e-messages-beta-{unique_marker()}"
model_id: Final = proxy.create_model(model, _deployment(backend))
@ -78,7 +116,7 @@ class TestAnthropicBetaHeaders:
max_tokens=64,
messages=[{"role": "user", "content": "Say 'hello' and nothing else"}],
tools=tools,
extra_headers={"anthropic-beta": _mapped_beta_headers(beta_provider)},
extra_headers={"anthropic-beta": _mapped_beta_headers(beta_provider, per_model)},
extra_body=NO_PROXY_CACHE,
)

256
tests/e2e/llm_translation/test_sail_e2e.py
View file

@ -1,256 +0,0 @@
"""Live e2e: Sail through the gateway, where LiteLLM turns ``service_tier`` into Sail's
``metadata.completion_window`` and bills the price columns of the window it sent.
The deployment carries its own base, balanced and flex rates, each distinct, so a bill at
the wrong tier cannot pass. They are registered on the deployment instead of read from the
proxy's cost map, because a stack that loads the published map has no ``sail/`` rows until
this provider ships. Requires SAIL_API_KEY on the proxy; no skip gate.
"""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final, Literal
import pytest
from e2e_config import SLOW_PROVIDER_TIMEOUT_SECONDS, unique_marker
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
from lifecycle import ResourceManager
from models import LiteLLMParamsBody, SpendLogRow
from openai import OpenAI
from proxy_client import ProxyClient
from sdk_clients import NO_PROXY_CACHE, SdkClients, response_header
pytestmark = pytest.mark.e2e
BACKEND: Final = "sail/zai-org/GLM-5.3"
PricedTier = Literal["base", "balanced", "flex"]
PRICED_TIERS: Final[tuple[PricedTier, ...]] = ("base", "balanced", "flex")
PROMPT: Final = "Reply with one word."
MAX_TOKENS: Final = 512
@dataclass(frozen=True, slots=True)
class _Rates:
input: float
output: float
cache_read: float
RATES: Final[Mapping[PricedTier, _Rates]] = {
"base": _Rates(input=3e-06, output=9e-06, cache_read=1e-06),
"balanced": _Rates(input=2e-06, output=6e-06, cache_read=7e-07),
"flex": _Rates(input=1e-06, output=3e-06, cache_read=4e-07),
}
@dataclass(frozen=True, slots=True)
class _Tokens:
prompt: int
cached: int
completion: int
def _approx_equal(actual: float, expected: float) -> bool:
return abs(actual - expected) <= max(1e-12, abs(expected) * 1e-2)
def _cost(rates: _Rates, tokens: _Tokens) -> float:
return (
(tokens.prompt - tokens.cached) * rates.input
+ tokens.cached * rates.cache_read
+ tokens.completion * rates.output
)
def _register(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]:
model: Final = f"e2e-sail-{unique_marker()}"
model_id: Final = proxy.create_model(
model,
LiteLLMParamsBody(
model=BACKEND,
api_key="os.environ/SAIL_API_KEY",
input_cost_per_token=RATES["base"].input,
output_cost_per_token=RATES["base"].output,
cache_read_input_token_cost=RATES["base"].cache_read,
input_cost_per_token_balanced=RATES["balanced"].input,
output_cost_per_token_balanced=RATES["balanced"].output,
cache_read_input_token_cost_balanced=RATES["balanced"].cache_read,
input_cost_per_token_flex=RATES["flex"].input,
output_cost_per_token_flex=RATES["flex"].output,
cache_read_input_token_cost_flex=RATES["flex"].cache_read,
),
)
resources.defer(lambda: proxy.delete_model(model_id))
return model, resources.key()
def _openai(sdk: SdkClients, key: str) -> OpenAI:
return sdk.openai(key).with_options(timeout=SLOW_PROVIDER_TIMEOUT_SECONDS)
def _assert_billed_at(tier: PricedTier, tokens: _Tokens, header_cost: str | None) -> float:
assert tokens.prompt > 0 and tokens.completion > 0, f"Sail reported no usage, so no cost is real: {tokens}"
assert header_cost is not None, "x-litellm-response-cost header missing"
costs: Final = {priced: _cost(rates, tokens) for priced, rates in RATES.items()}
assert not any(_approx_equal(costs[other], costs[tier]) for other in PRICED_TIERS if other != tier), (
f"{BACKEND} tier rates too close together to tell {tier} apart at {tokens}: {costs}"
)
assert _approx_equal(float(header_cost), costs[tier]), (
f"header cost {header_cost} is not the {tier} price at {tokens}: expected {costs[tier]}, all tiers {costs}"
)
return float(header_cost)
def _assert_spend_row_matches(proxy: ProxyClient, key: str, header_cost: float) -> None:
def priced(rows: list[SpendLogRow]) -> bool:
return any((row.spend or 0) > 0 for row in rows)
rows: Final = [row for row in proxy.poll_logs_for_key(key, predicate=priced) if (row.spend or 0) > 0]
assert rows, f"no priced spend row landed for key {key}"
assert rows[0].custom_llm_provider == "sail", f"spend row misattributed: {rows[0]}"
assert rows[0].spend is not None and _approx_equal(rows[0].spend, header_cost), (
f"logged spend {rows[0].spend} disagrees with the x-litellm-response-cost header {header_cost}"
)
class TestSailChatCompletions:
@pytest.mark.covers("llm.chat_completions.sail.service_tier.nonstream.cost_logged")
@meta(
Subject(
domain=Domain.SPEND_BUDGETS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.SAIL,),
models=(BACKEND,),
mode=Mode.NONSTREAM,
)
)
@pytest.mark.parametrize(
("service_tier", "billed_tier"), [("balanced", "balanced"), ("auto", "base")]
)
def test_service_tier_bills_the_matching_completion_window(
self,
proxy: ProxyClient,
resources: ResourceManager,
sdk: SdkClients,
service_tier: str,
billed_tier: PricedTier,
) -> None:
model, key = _register(proxy, resources)
raw: Final = _openai(sdk, key).chat.completions.with_raw_response.create(
model=model,
messages=[{"role": "user", "content": f"{PROMPT} {unique_marker()}"}],
max_completion_tokens=MAX_TOKENS,
extra_body={**NO_PROXY_CACHE, "service_tier": service_tier},
)
usage: Final = raw.parse().usage
assert usage is not None, "chat response carries no usage"
details: Final = usage.prompt_tokens_details
tokens: Final = _Tokens(
prompt=usage.prompt_tokens,
cached=(details.cached_tokens or 0) if details else 0,
completion=usage.completion_tokens,
)
header_cost: Final = _assert_billed_at(
billed_tier, tokens, response_header(raw.headers, "x-litellm-response-cost")
)
_assert_spend_row_matches(proxy, key, header_cost)
@pytest.mark.covers("llm.chat_completions.sail.service_tier.nonstream.drops_unknown_tier_and_bills_asap")
@meta(
Subject(
domain=Domain.SPEND_BUDGETS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.SAIL,),
models=(BACKEND,),
mode=Mode.NONSTREAM,
)
)
@pytest.mark.parametrize("service_tier", ["bogus", 5])
def test_unknown_service_tier_is_dropped_and_billed_asap(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, service_tier: str | int
) -> None:
model, key = _register(proxy, resources)
raw: Final = _openai(sdk, key).chat.completions.with_raw_response.create(
model=model,
messages=[{"role": "user", "content": f"{PROMPT} {unique_marker()}"}],
max_completion_tokens=MAX_TOKENS,
extra_body={**NO_PROXY_CACHE, "service_tier": service_tier, "drop_params": True},
)
usage: Final = raw.parse().usage
assert usage is not None, "chat response carries no usage"
details: Final = usage.prompt_tokens_details
tokens: Final = _Tokens(
prompt=usage.prompt_tokens,
cached=(details.cached_tokens or 0) if details else 0,
completion=usage.completion_tokens,
)
header_cost: Final = _assert_billed_at("base", tokens, response_header(raw.headers, "x-litellm-response-cost"))
_assert_spend_row_matches(proxy, key, header_cost)
class TestSailResponses:
@pytest.mark.covers("llm.responses.sail.service_tier.nonstream.cost_logged")
@meta(
Subject(
domain=Domain.SPEND_BUDGETS,
route=Route.RESPONSES,
providers=(Provider.SAIL,),
models=(BACKEND,),
mode=Mode.NONSTREAM,
)
)
def test_caller_completion_window_bills_its_rates(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model, key = _register(proxy, resources)
raw: Final = _openai(sdk, key).responses.with_raw_response.create(
model=model,
input=f"{PROMPT} {unique_marker()}",
max_output_tokens=MAX_TOKENS,
metadata={"completion_window": "balanced"},
extra_body=NO_PROXY_CACHE,
)
usage: Final = raw.parse().usage
assert usage is not None, "responses answer carries no usage"
tokens: Final = _Tokens(
prompt=usage.input_tokens,
cached=usage.input_tokens_details.cached_tokens,
completion=usage.output_tokens,
)
header_cost: Final = _assert_billed_at(
"balanced", tokens, response_header(raw.headers, "x-litellm-response-cost")
)
_assert_spend_row_matches(proxy, key, header_cost)
class TestSailMessages:
@pytest.mark.covers("llm.messages.sail.basic.nonstream.works")
@meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.MESSAGES,
providers=(Provider.SAIL,),
models=(BACKEND,),
mode=Mode.NONSTREAM,
)
)
def test_plain_call_returns_a_message(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model, key = _register(proxy, resources)
message: Final = sdk.anthropic(key).messages.create(
model=model,
max_tokens=MAX_TOKENS,
messages=[{"role": "user", "content": PROMPT}],
extra_body=NO_PROXY_CACHE,
)
assert message.role == "assistant" and message.content, f"/v1/messages returned no content: {message}"
assert message.usage.output_tokens > 0, f"/v1/messages reported no output usage: {message.usage}"

39
tests/e2e/logging/logging_client.py
View file

@ -213,45 +213,6 @@ class LangfuseCreds:
)
@dataclass(frozen=True, slots=True)
class WeaveCreds:
"""Weights & Biases Weave credentials for a key-scoped ``weave_otel`` callback.
The proxy still needs WANDB_API_KEY / WANDB_PROJECT_ID in its own environment:
the weave_otel logger is constructed from those before the per-key vars are
applied, so a key-scoped callback on a proxy without them never initializes.
The per-key vars are what direct THIS key's spans at this project.
"""
api_key: str
project_id: str
def key_logging_metadata(self) -> KeyMetadata:
return KeyMetadata(
logging=[
KeyLoggingCallback(
callback_name="weave_otel",
callback_type="success_and_failure",
callback_vars=KeyLoggingCallbackVars(
wandb_api_key=self.api_key,
weave_project_id=self.project_id,
),
)
]
)
def load_weave_creds() -> WeaveCreds:
api_key = os.getenv("WANDB_API_KEY")
project_id = (os.getenv("WEAVE_PROJECT_ID") or os.getenv("WANDB_PROJECT_ID") or "").strip()
if not (api_key and project_id):
pytest.fail(
"Weave e2e requires WANDB_API_KEY and WEAVE_PROJECT_ID (or WANDB_PROJECT_ID, "
"format <entity>/<project>); missing credentials is a hard failure, not a skip"
)
return WeaveCreds(api_key=api_key, project_id=project_id)
def load_langfuse_creds() -> LangfuseCreds:
public_key = os.getenv("LANGFUSE_PUBLIC_KEY")
secret_key = os.getenv("LANGFUSE_SECRET_KEY")

2
tests/e2e/logging/test_otel_trace_e2e.py
View file

@ -54,7 +54,7 @@ def _assert_otel_destination_configured(client: LoggingClient) -> None:
details = _ReadinessDetails.model_validate_json(readiness_details_body(client))
assert OTEL_V2_LOGGER_NAME in details.success_callbacks, (
f"the proxy must report the {OTEL_V2_LOGGER_NAME} callback active "
f"(LITELLM_OTEL_V2 + arize_phoenix preset in the compose config); got: {details.success_callbacks}"
f"(LITELLM_OTEL_V2=true + the otel callback in the stack config); got: {details.success_callbacks}"
)

212
tests/e2e/logging/test_weave_log_e2e.py
View file

@ -1,212 +0,0 @@
"""Live e2e: key-scoped Weave (Weights & Biases) delivery, success and failure.
Covers the two `logging.niche_integrations.*.logs_spend` cells with a real member
of that cohort. A key carrying a `weave_otel` callback in its logging metadata
must deliver its calls to the real Weave project, and each call must arrive
exactly once, carrying the same cost the response header reported:
- success: one `litellm_request` call, OTEL status OK, `llm.response.cost` equal
to `x-litellm-response-cost`, and non-zero tokens
- failure: a provider-rejected call arrives too, as one call with OTEL status
ERROR naming the provider exception, and with no cost - a failed call that
silently never reaches the destination is an invisible outage, and a billed
one is worse
Both halves assert the recorded state (the key's callback registration answers
success and the destination holds the call) and the enforced behavior (the
delivered payload's status and cost). Delivery is read back through Weave's own
query API; nothing is mocked.
"""
from __future__ import annotations
import time
from typing import Final
import pytest
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
from e2e_http import StreamingResponse
from e2e_metadata import Domain, Mode, Provider, Subject, meta
from lifecycle import ResourceManager
from logging_client import (
INVALID_UPSTREAM_API_KEY,
LoggingClient,
WeaveCreds,
costs_agree,
first_ok,
load_weave_creds,
)
from models import LiteLLMParamsBody
from weave_reader import WeaveCall, WeaveReader, build_weave_reader
pytestmark = pytest.mark.e2e
UNREACHABLE_ANTHROPIC_BACKEND: Final = "anthropic/claude-haiku-4-5"
@pytest.fixture(scope="session")
def weave_creds() -> WeaveCreds:
return load_weave_creds()
@pytest.fixture(scope="session")
def weave_reader() -> WeaveReader:
return build_weave_reader()
#: How far before the request the Weave read-back window opens, to absorb clock
#: skew between this host and Weave. Without it a host running slightly fast
#: would filter out its own call.
_WINDOW_SKEW_SECONDS = 120.0
def _window_start() -> float:
return time.time() - _WINDOW_SKEW_SECONDS
def _exactly_one(calls: tuple[WeaveCall, ...], *, marker: str, what: str) -> WeaveCall:
assert calls, f"no Weave call for the {what} (marker {marker}) reached the project within the deadline"
assert len(calls) == 1, (
f"expected exactly ONE Weave call for the {what} (marker {marker}), got {len(calls)}: "
f"{[call.id for call in calls]} - more than one call for one request is the "
"duplicate-delivery bug"
)
return calls[0]
WEAVE_STAGE_RED_REASON = (
"stage red: product gap, key-scoped weave_otel spans are not delivered when the OTEL v2 callback is active"
)
class TestWeaveLogDelivery:
@pytest.mark.skip(reason=WEAVE_STAGE_RED_REASON)
@pytest.mark.covers("logging.niche_integrations.success.logs_spend", exercised_on=["chat_completions"])
@meta(
Subject(
domain=Domain.OBSERVABILITY,
providers=(Provider.ANTHROPIC,),
models=(CHEAP_ANTHROPIC_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_chat_completions_delivers_one_call_with_spend(
self,
client: LoggingClient,
weave_creds: WeaveCreds,
weave_reader: WeaveReader,
resources: ResourceManager,
) -> None:
alias = f"weave-key-{unique_marker()}"
key = client.key_with_alias(
alias,
models=[CHEAP_ANTHROPIC_MODEL],
metadata=weave_creds.key_logging_metadata(),
)
resources.defer(lambda: client.delete_key(key))
marker = unique_marker()
since = _window_start()
outcome = first_ok(
client,
lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=64),
)
assert outcome.response_cost is not None and outcome.response_cost > 0, (
f"the response must report x-litellm-response-cost, got {outcome.response_cost!r}"
)
call = _exactly_one(
weave_reader.poll_calls_matching(marker, since=since), marker=marker, what="successful call"
)
assert call.status_code == "OK", f"a successful call must land at OK span status, got {call.status_code!r}"
cost = call.response_cost
assert cost is not None and costs_agree(outcome.response_cost, cost), (
f"the Weave call's llm.response.cost {cost!r} must agree with the header cost "
f"{outcome.response_cost} - a delivered span with the wrong cost is a silent "
"billing-attribution bug"
)
assert call.total_tokens is not None and call.total_tokens > 0, (
f"the delivered call must carry token usage, got {call.total_tokens!r}"
)
@pytest.mark.skip(reason=WEAVE_STAGE_RED_REASON)
@pytest.mark.covers("logging.niche_integrations.failure.logs_spend", exercised_on=["chat_completions"])
@meta(
Subject(
domain=Domain.OBSERVABILITY,
providers=(Provider.ANTHROPIC,),
models=(UNREACHABLE_ANTHROPIC_BACKEND,),
mode=Mode.NONSTREAM,
)
)
def test_failed_chat_completions_delivers_one_error_call(
self,
client: LoggingClient,
weave_creds: WeaveCreds,
weave_reader: WeaveReader,
resources: ResourceManager,
) -> None:
"""A deployment with an invalid upstream key passes proxy auth and fails
at the provider, so exactly one provider failure exists for it. Proxy-side
401s during key propagation never reach the provider and ship no payload,
which is what the retry loop below relies on."""
model_name = f"weave-err-{unique_marker()}"
model_id = client.create_model(
model_name,
LiteLLMParamsBody(model=UNREACHABLE_ANTHROPIC_BACKEND, api_key=INVALID_UPSTREAM_API_KEY),
)
resources.defer(lambda: client.delete_model(model_id))
key = client.key_with_alias(
f"weave-err-key-{unique_marker()}",
models=[model_name],
metadata=weave_creds.key_logging_metadata(),
)
resources.defer(lambda: client.delete_key(key))
marker = unique_marker()
since = _window_start()
outcome = _provoke_provider_failure(client, key, model_name, marker)
call = _exactly_one(weave_reader.poll_calls_matching(marker, since=since), marker=marker, what="failed call")
assert call.status_code == "ERROR", (
f"a failed call must land at ERROR span status, got {call.status_code!r} - "
"Weave's own summary.weave.status reads success either way, which is exactly "
"why the span status is what this asserts on"
)
error = call.error
assert error is not None and error.message is not None and "AnthropicException" in error.message, (
f"the delivered call must carry the provider error, got {error!r}"
)
assert not call.response_cost, f"a failed call must not be billed, got llm.response.cost={call.response_cost!r}"
assert outcome.status_code == 401, (
f"an upstream auth failure must map to 401, got {outcome.status_code}: {outcome.body[:200]}"
)
def _provoke_provider_failure(client: LoggingClient, key: str, model_name: str, marker: str) -> StreamingResponse:
"""Send until the provider (not the proxy) is the one rejecting the call.
A network failure between the test and the proxy is NOT retried: the request
may have been served, and a retry would double-log the failure payload and
falsely trip the exactly-one assertion.
"""
deadline = time.monotonic() + client.proxy.poll_timeout
while True:
outcome = client.chat_raw(key, model_name, f"trigger an upstream auth failure {marker}", max_tokens=16)
assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid"
assert outcome.status_code != -1, (
"network failure between the test and the proxy while provoking the provider failure; "
"retrying now could double-log the failure payload and falsely trip the exactly-one "
f"assertion - fix the rig connectivity first: {outcome.body[:200]}"
)
if "AnthropicException" in outcome.body or time.monotonic() >= deadline:
break
time.sleep(client.proxy.poll_interval)
assert "AnthropicException" in outcome.body, (
"never saw the upstream provider failure before the deadline; the key may still be "
f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}"
)
return outcome

Some files were not shown because too many files have changed in this diff Show more