scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

Merge remote-tracking branch 'origin/main' into litellm-providers/price-sync-databricks

Browse source
This commit is contained in:
Devin AI 2026-11-02 19:09:57 +00:00
commit ecc53f58c2
119 changed files with 9633 additions and 1827 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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[

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

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

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

@ -1988,7 +1988,7 @@ class OpenAIResponsesHandler(BaseTranslation):
"annotations": (),
"logprobs": None,
}
return (
events: Final = (
*_open_item_closing_events(responses_so_far),
OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
@ -2033,6 +2033,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,

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,

67
litellm/model_prices_and_context_window_backup.json
View file

@ -59905,6 +59905,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": [
@ -59926,7 +59927,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,
@ -59941,6 +59943,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"
@ -59960,7 +59963,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,
@ -60366,6 +60370,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",
@ -60389,7 +60394,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,
@ -60404,6 +60410,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",
@ -60427,7 +60434,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,
@ -60442,6 +60450,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",
@ -60465,7 +60474,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,
@ -64059,6 +64069,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,
@ -64078,6 +64089,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,
@ -64120,6 +64132,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,
@ -64157,6 +64170,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,
@ -64176,6 +64190,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,
@ -64250,6 +64265,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,
@ -64270,6 +64286,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,
@ -64293,6 +64310,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,
@ -64368,6 +64386,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,
@ -64404,6 +64423,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
},
@ -64440,6 +64460,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,
@ -64510,6 +64531,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,
@ -64533,6 +64555,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,
@ -67326,6 +67349,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,
@ -67344,6 +67368,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 +67387,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,
@ -67380,6 +67406,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,
@ -67395,6 +67422,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,
@ -67437,6 +67465,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,
@ -67455,6 +67484,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
@ -67473,6 +67503,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
@ -67490,6 +67521,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
@ -67508,6 +67540,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
@ -67526,6 +67559,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
@ -68663,7 +68697,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",
@ -68698,7 +68734,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",
@ -68727,7 +68765,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,
@ -73777,6 +73817,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,
@ -75361,6 +75402,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",
@ -75383,6 +75425,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
},
@ -77284,6 +77327,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,
@ -77442,6 +77486,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,
@ -79075,6 +79120,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",
@ -79097,6 +79143,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
},
@ -81217,6 +81264,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,
@ -81233,6 +81281,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,
@ -81249,6 +81298,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,
@ -81265,6 +81315,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,

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

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")

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):

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":

331
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,6 +390,12 @@ 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,
@ -397,6 +406,7 @@ 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
@ -10186,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:
@ -10211,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(
@ -14950,7 +14935,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,
@ -15020,6 +15005,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.
@ -15043,7 +15029,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"
)
@ -15305,6 +15291,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)
@ -15374,12 +15363,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(
@ -15389,6 +15381,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 {}
@ -15582,6 +15658,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)
@ -15621,30 +15731,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:
@ -15655,6 +15765,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:
@ -15662,11 +15793,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)
@ -15677,7 +15808,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

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:

67
model_prices_and_context_window.json
View file

@ -59905,6 +59905,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": [
@ -59926,7 +59927,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,
@ -59941,6 +59943,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"
@ -59960,7 +59963,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,
@ -60366,6 +60370,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",
@ -60389,7 +60394,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,
@ -60404,6 +60410,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",
@ -60427,7 +60434,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,
@ -60442,6 +60450,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",
@ -60465,7 +60474,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,
@ -64059,6 +64069,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,
@ -64078,6 +64089,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,
@ -64120,6 +64132,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,
@ -64157,6 +64170,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,
@ -64176,6 +64190,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,
@ -64250,6 +64265,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,
@ -64270,6 +64286,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,
@ -64293,6 +64310,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,
@ -64368,6 +64386,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,
@ -64404,6 +64423,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
},
@ -64440,6 +64460,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,
@ -64510,6 +64531,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,
@ -64533,6 +64555,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,
@ -67326,6 +67349,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,
@ -67344,6 +67368,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 +67387,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,
@ -67380,6 +67406,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,
@ -67395,6 +67422,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,
@ -67437,6 +67465,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,
@ -67455,6 +67484,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
@ -67473,6 +67503,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
@ -67490,6 +67521,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
@ -67508,6 +67540,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
@ -67526,6 +67559,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
@ -68663,7 +68697,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",
@ -68698,7 +68734,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",
@ -68727,7 +68765,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,
@ -73777,6 +73817,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,
@ -75361,6 +75402,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",
@ -75383,6 +75425,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
},
@ -77284,6 +77327,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,
@ -77442,6 +77486,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,
@ -79075,6 +79120,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",
@ -79097,6 +79143,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
},
@ -81217,6 +81264,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,
@ -81233,6 +81281,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,
@ -81249,6 +81298,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,
@ -81265,6 +81315,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,

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,
)

2
tests/integration/_support/asgi.py
View file

@ -15,7 +15,7 @@ from starlette.types import ASGIApp
@contextmanager
def asgi_server(
app: ASGIApp, *, before_stop: Callable[[], None] | None = None, port: int = 0, keepalive_seconds: int = 1
app: ASGIApp, *, before_stop: Callable[[], None] | None = None, port: int = 0, keepalive_seconds: int = 30
) -> Iterator[str]:
with socket.socket() as listener:
listener.bind(("127.0.0.1", port))

3
tests/integration/database/test_prisma_migration_lock.py
View file

@ -19,7 +19,8 @@ def _can_take_prismas_session_lock(connection: psycopg.Connection[tuple[object,
def _granted_advisory_keys(connection: psycopg.Connection[tuple[object, ...]]) -> frozenset[int]:
rows: Final = connection.execute(
"SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted"
"SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted "
"AND database = (SELECT oid FROM pg_database WHERE datname = current_database())"
).fetchall()
return frozenset(int(row[0]) for row in rows)

6
tests/integration/management/test_user_updates_wedged_coordination_redis.py
View file

@ -193,8 +193,10 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged(
timings["s6_customer_update"] = _timed_post(
candidate, "/customer/update", {"user_id": customer_id, "max_budget": 6.0}
)
assert timings["s6_customer_update"] < _HANDLER_BUDGET_SECONDS, (
f"/customer/update took {timings['s6_customer_update']:.3f}s with a wedged coordination Redis"
assert timings["s6_customer_update"] < _SHORTEST_REDIS_WAIT_SECONDS, (
f"/customer/update took {timings['s6_customer_update']:.3f}s, reaching the "
f"{_SHORTEST_REDIS_WAIT_SECONDS}s shortest wait the proxy puts on Redis, so it waited on the "
"wedged coordination Redis"
)
timings["s7_reset_spend"] = _timed_post(candidate, f"/key/{key}/reset_spend", {"reset_to": 0})
assert timings["s7_reset_spend"] < _HANDLER_BUDGET_SECONDS, (

27
tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py
View file

@ -3,7 +3,7 @@ import json
import re
import time
import uuid
from collections.abc import Callable, Mapping
from collections.abc import Callable, Iterator, Mapping
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
@ -13,7 +13,7 @@ import anthropic
import httpx
import openai
import pytest
from integration._support.client import Gateway, eventually, object_value
from integration._support.client import Gateway, eventually, gateway_from_environment, object_value
from integration._support.database import read_rows
from integration._support.process import owned_proxy
from integration._support.vertex import service_account_json
@ -417,14 +417,31 @@ def test_a_deployment_dropping_context_management_sends_neither_the_field_nor_it
_spend_row(f"msg_{marker}")
@pytest.fixture(scope="module")
def local_beta_map_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
with (
gateway_from_environment() as shared,
owned_proxy(
shared,
tmp_path_factory.mktemp("safeguards-local-beta-map"),
{"LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS": "true"},
) as owned,
):
yield owned
@pytest.mark.parametrize("row", tuple(_FOUNDRY_SHAPES))
def test_azure_ai_messages_forward_the_mapped_client_betas_beside_the_derived_one(gateway: Gateway, row: str) -> None:
def test_azure_ai_messages_forward_the_mapped_client_betas_beside_the_derived_one(
local_beta_map_gateway: Gateway, row: str
) -> None:
case: Final = _FOUNDRY_SHAPES[row]
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
with wire_server(anthropic_peer) as wire, local_beta_map_gateway.scenario() as scenario:
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
response: Final = _post(
gateway, messages_body(model, marker, safeguards=case.sent, stream=case.stream), beta=case.client_beta
local_beta_map_gateway,
messages_body(model, marker, safeguards=case.sent, stream=case.stream),
beta=case.client_beta,
)
if case.stream:
_assert_streamed(response, marker)

19
tests/integration/messages_endpoint/test_beta_gated_fields_wire.py
View file

@ -15,8 +15,8 @@ from pydantic import JsonValue, TypeAdapter
pytestmark: Final = pytest.mark.shares_worker_state
ANTHROPIC_MODEL: Final = "beta-gated-anthropic"
FOUNDRY_MODEL: Final = "beta-gated-foundry"
FOUNDRY_DROP_PARAMS_MODEL: Final = "beta-gated-foundry-drop-params"
DATABRICKS_MODEL: Final = "beta-gated-databricks"
DATABRICKS_DROP_PARAMS_MODEL: Final = "beta-gated-databricks-drop-params"
CLEAR_AT_BETA: Final = "mid-conversation-system-clear-at-2026-08-21"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_NONCE: Final = re.compile(rb"nonce-([0-9a-f]{32})")
@ -71,9 +71,16 @@ def write_config(directory: Path, wire_url: str) -> Path:
{
"model_list": [
deployment(ANTHROPIC_MODEL, "anthropic/claude-opus-4-8", f"{wire_url}/anthropic"),
deployment(FOUNDRY_MODEL, "azure_ai/claude-opus-4-8", f"{wire_url}/foundry"),
deployment(
FOUNDRY_DROP_PARAMS_MODEL, "azure_ai/claude-opus-4-8", f"{wire_url}/foundry", drop_params=True
DATABRICKS_MODEL,
"databricks/databricks-claude-opus-4-8",
f"{wire_url}/databricks/serving-endpoints",
),
deployment(
DATABRICKS_DROP_PARAMS_MODEL,
"databricks/databricks-claude-opus-4-8",
f"{wire_url}/databricks/serving-endpoints",
drop_params=True,
),
],
"general_settings": {
@ -150,7 +157,7 @@ def system_message(body: bytes) -> dict[str, JsonValue]:
def test_a_drop_params_leg_whose_row_lacks_the_beta_sends_neither_the_beta_nor_clear_at(rig: Rig) -> None:
request_nonce: Final = nonce()
status, raw = rig.send(FOUNDRY_DROP_PARAMS_MODEL, request_nonce)
status, raw = rig.send(DATABRICKS_DROP_PARAMS_MODEL, request_nonce)
assert reply_text(status, raw) == f"answered {request_nonce.removeprefix('nonce-')}"
upstream: Final = rig.received(request_nonce)
assert CLEAR_AT_BETA not in beta_header(upstream.headers), upstream.headers
@ -159,7 +166,7 @@ def test_a_drop_params_leg_whose_row_lacks_the_beta_sends_neither_the_beta_nor_c
def test_a_leg_whose_row_lacks_the_beta_forwards_clear_at_as_sent_without_drop_params(rig: Rig) -> None:
request_nonce: Final = nonce()
status, raw = rig.send(FOUNDRY_MODEL, request_nonce)
status, raw = rig.send(DATABRICKS_MODEL, request_nonce)
assert reply_text(status, raw) == f"answered {request_nonce.removeprefix('nonce-')}"
upstream: Final = rig.received(request_nonce)
assert CLEAR_AT_BETA not in beta_header(upstream.headers), upstream.headers

13
tests/integration/observability/_logging_only_scope_support.py
View file

@ -22,7 +22,7 @@ import yaml
from anthropic import Anthropic, AsyncAnthropic
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from integration._support.client import Gateway, Scenario, eventually, object_value
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
from integration._support.database import read_rows, write_rows
from integration._support.upstream import delete_scenario, register_scenario
from integration._support.wire import Reply, Request, Wire
@ -798,6 +798,7 @@ def _insert_database_guardrail(
@contextmanager
def _database_guardrail(
gateway: Gateway,
identity: str,
policy_url: str,
scope: str | None,
@ -809,11 +810,15 @@ def _database_guardrail(
try:
yield
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
def _delete_database_guardrail(identity: str) -> None:
write_rows('DELETE FROM "LiteLLM_GuardrailsTable" WHERE guardrail_name=%s', (identity,))
def _delete_database_guardrail(gateway: Gateway, identity: str) -> None:
rows: Final = read_rows('SELECT guardrail_id FROM "LiteLLM_GuardrailsTable" WHERE guardrail_name=%s', (identity,))
responses: Final = tuple(
gateway.request("DELETE", f"/guardrails/{string_value(row['guardrail_id'])}") for row in rows
)
assert all(response.status_code == 200 for response in responses), [response.text for response in responses]
def _post_guardrail_body(

30
tests/integration/observability/test_logging_only_scope_config.py
View file

@ -115,7 +115,7 @@ def test_G7_database_invalid_scope_reads_preserve_pre_call_blocking(gateway: Gat
assert blocked_response.status_code == 400, blocked_response.text
assert _upstream_requests(gateway.upstream_url) == ()
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
finally:
delete_scenario(upstream_handle)
@ -678,7 +678,7 @@ def test_database_guardrail_load_keeps_pre_call_blocking(
for entry in entries
) == ((identity, "pre_call", "guardrail_intervened"),), entries
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
finally:
delete_scenario(upstream_handle)
@ -799,7 +799,7 @@ def test_G6_database_guardrail_polling_normalizes_invalid_scope_without_reinitia
) == ((identity, "pre_call", "guardrail_intervened"),), (call_id, entries)
assert _upstream_requests(gateway.upstream_url) == ()
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
finally:
delete_scenario(upstream_handle)
@ -952,7 +952,7 @@ def test_management_post_validates_logging_only_scope(
assert isinstance(after_list, list), after_list
assert identity not in {object_value(item)["guardrail_name"] for item in after_list}, after_list
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
def test_H5_management_put_rejection_preserves_database_and_runtime(gateway: Gateway, tmp_path: Path) -> None:
@ -1074,7 +1074,7 @@ def test_H5_management_put_rejection_preserves_database_and_runtime(gateway: Gat
for entry in entries
) == ((identity, "pre_call", "guardrail_intervened"),), (call_id, entries)
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
delete_scenario(upstream_handle)
@ -1153,7 +1153,7 @@ def test_H6_management_patch_clears_scope_when_switching_to_blocking_only_mode(
for entry in entries
) == ((identity, ("pre_call",), "guardrail_intervened"),), entries
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
delete_scenario(upstream_handle)
@ -1240,7 +1240,7 @@ def test_H7_management_patch_rejection_preserves_pre_call_guardrail(gateway: Gat
for entry in entries
) == ((identity, "pre_call", "guardrail_intervened"),), (call_id, entries)
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
delete_scenario(upstream_handle)
@ -1331,7 +1331,7 @@ def test_management_patch_default_on_heals_stored_scope(
for entry in entries
) == ((identity, "pre_call", "guardrail_intervened"),), entries
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
finally:
delete_scenario(upstream_handle)
@ -1429,7 +1429,7 @@ def test_H10_management_patch_null_scope_restores_both_directions(gateway: Gatew
(identity, "logging_only", "guardrail_intervened"),
), entries
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
delete_scenario(upstream_handle)
@ -1517,7 +1517,7 @@ def test_H11_management_patch_same_scope_is_idempotent_and_output_only(gateway:
(entry["guardrail_name"], entry["guardrail_mode"], entry["guardrail_status"]) for entry in entries
) == tuple((identity, "logging_only", "guardrail_intervened") for _ in expected_directions), entries
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
delete_scenario(upstream_handle)
@ -1550,7 +1550,7 @@ def test_H12_management_reads_expose_typed_logging_only_scope(gateway: Gateway,
assert object_value(info["litellm_params"])["logging_only_scope"] == "input", info
assert object_value(matches[0]["litellm_params"])["logging_only_scope"] == "input", matches
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
def test_H13_unauthenticated_management_and_chat_requests_do_not_scan(gateway: Gateway, tmp_path: Path) -> None:
@ -1629,7 +1629,7 @@ def test_K1_post_rejects_continue_flag_outside_logging_only(gateway: Gateway, tm
assert after_list == before_list, (before_list, after_list)
assert identity not in {object_value(item)["guardrail_name"] for item in after_list}, after_list
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
def test_K3_patch_flag_rejection_preserves_raw_stored_row(gateway: Gateway, tmp_path: Path) -> None:
@ -1675,7 +1675,7 @@ def test_K3_patch_flag_rejection_preserves_raw_stored_row(gateway: Gateway, tmp_
} == {key: after_params[key] for key in after_params if key != "api_key"}, (before_params, after_params)
assert str(after_params["api_key"]).startswith("litellm_enc::"), after_params
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
def test_K4_patch_dropping_logging_only_mode_clears_stored_flag(gateway: Gateway, tmp_path: Path) -> None:
@ -1713,7 +1713,7 @@ def test_K4_patch_dropping_logging_only_mode_clears_stored_flag(gateway: Gateway
assert params.get("logging_only_continue_on_input_failure") is None, persisted
assert params["mode"] == ["pre_call"], persisted
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
def test_K5_yaml_flag_on_blocking_guardrail_boots_and_keeps_blocking(gateway: Gateway, tmp_path: Path) -> None:
@ -1759,5 +1759,5 @@ def test_K5_yaml_flag_on_blocking_guardrail_boots_and_keeps_blocking(gateway: Ga
log_text: Final = owned.log.read_text()
assert "Ignoring logging_only_continue_on_input_failure" in log_text, log_text
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
delete_scenario(upstream_handle)

2
tests/integration/observability/test_logging_only_scope_runtime.py
View file

@ -1098,7 +1098,7 @@ def test_X2_scope_patch_toggles_during_concurrent_requests_keep_one_scan_per_res
with wire_server(policy) as edge:
config: Final = _empty_proxy_configuration(tmp_path, identity)
with (
_database_guardrail(identity, edge.url, "output", mode="logging_only"),
_database_guardrail(gateway, identity, edge.url, "output", mode="logging_only"),
owned_proxy(gateway, tmp_path, {}, config=config, workers=1) as candidate,
):
guardrails: Final = candidate.get("/v2/guardrails/list")["guardrails"]

29
tests/integration/providers/test_azure_ai_claude_messages_wire.py
View file

@ -1,10 +1,12 @@
from collections.abc import Iterator
from typing import Final
import anthropic
import pytest
from anthropic.types import Message
from integration._support import claude_code as cc
from integration._support.client import Gateway
from integration._support.client import Gateway, gateway_from_environment
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
@ -144,8 +146,23 @@ def test_azure_ai_claude_messages_preserves_target_auth_beta_and_typed_response(
assert _JSON_OBJECT.validate_json(provider_request.body) == expected_provider_body, provider_request.body
@pytest.fixture(scope="module")
def local_beta_map_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
with (
gateway_from_environment() as shared,
owned_proxy(
shared,
tmp_path_factory.mktemp("azure-ai-local-beta-map"),
{"LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS": "true"},
) as owned,
):
yield owned
@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"])
def test_azure_ai_claude_messages_filters_auto_injected_betas_by_azure_support(gateway: Gateway, stream: bool) -> None:
def test_azure_ai_claude_messages_forwards_the_auto_injected_betas_azure_accepts(
local_beta_map_gateway: Gateway, stream: bool
) -> None:
identity: Final = f"msg_azure_claude_beta_filter_{stream}"
api_path: Final = "/anthropic/v1/messages"
reply_body: Final = cc.message_reply(identity, _BACKEND, _CONTENT, _USAGE)
@ -172,7 +189,7 @@ def test_azure_ai_claude_messages_filters_auto_injected_betas_by_azure_support(g
"context_management": context_management,
}
expected_provider_body: Final = {**request_body_template, "model": _BACKEND}
expected_beta: Final = "structured-outputs-2025-11-13"
expected_beta: Final = "compact-2026-01-12,structured-outputs-2025-11-13"
def respond(request: Request) -> Reply:
assert request.method == "POST" and request.target == api_path, request.target
@ -185,7 +202,7 @@ def test_azure_ai_claude_messages_filters_auto_injected_betas_by_azure_support(g
return Reply(content_type="text/event-stream", chunks=stream_reply)
return Reply(body=reply_body)
with wire_server(respond) as wire, gateway.scenario() as scenario:
with wire_server(respond) as wire, local_beta_map_gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"azure_ai/{_BACKEND}",
api_base=wire.url,
@ -193,8 +210,8 @@ def test_azure_ai_claude_messages_filters_auto_injected_betas_by_azure_support(g
num_retries=0,
)
with anthropic.Anthropic(
api_key=gateway.key,
base_url=str(gateway.client.base_url),
api_key=local_beta_map_gateway.key,
base_url=str(local_beta_map_gateway.client.base_url),
max_retries=0,
) as client:
if stream:

57
tests/integration/providers/test_openai_chat_prompt_cache_breakpoint_wire.py
View file

@ -41,6 +41,7 @@ class _Proxy:
openai: Wire
gpt_5_6_with_prompt_cache_options: str
gpt_5_6_without_prompt_cache_options: str
gpt_5_6_dropping_params: str
def openai_receives(
self, model: str, messages: Sequence[JsonValue], *, stream: bool = False, tools: Sequence[JsonValue] = ()
@ -79,6 +80,9 @@ def proxy() -> Iterator[_Proxy]:
model="openai/gpt-5.6", api_base=api_base, prompt_cache_options={"mode": "explicit"}
),
gpt_5_6_without_prompt_cache_options=scenario.model(model="openai/gpt-5.6", api_base=api_base),
gpt_5_6_dropping_params=scenario.model(
model="openai/gpt-5.6", api_base=api_base, prompt_cache_options={"mode": "explicit"}, drop_params=True
),
)
@ -202,3 +206,56 @@ def test_custom_api_base_without_prompt_cache_options_forwards_cache_control_unc
assert received["messages"] == _SYSTEM_AND_USER_WITH_CACHE_CONTROL
assert "prompt_cache_options" not in received
assert "prompt_cache_breakpoint" not in str(received)
_CLIENT_BREAKPOINTS_ON_PARTS: Final[list[JsonValue]] = [
{"type": "text", "text": "Stable rules", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}},
{"type": "text", "text": "More rules", "prompt_cache_breakpoint": {"mode": "implicit"}},
{"type": "text", "text": "Describe this", "prompt_cache_breakpoint": {"mode": "explicit"}},
]
_CLIENT_BREAKPOINTS_AS_CHAT_COMPLETIONS_TAKES_THEM: Final[list[JsonValue]] = [
{"type": "text", "text": "Stable rules", "prompt_cache_breakpoint": {"mode": "explicit"}},
{"type": "text", "text": "More rules"},
{"type": "text", "text": "Describe this", "prompt_cache_breakpoint": {"mode": "explicit"}},
]
_CHAT_COMPLETIONS_REJECTS_TTL_AND_IMPLICIT: Final = (
"live gpt-6.1-sol Chat Completions answers 400 on a part breakpoint with any ttl or mode implicit (2026-10-28)"
)
@pytest.mark.parametrize("stream", [False, True], ids=["unary", "stream"])
def test_drop_params_cuts_client_breakpoints_down_to_what_chat_completions_takes(proxy: _Proxy, stream: bool) -> None:
received: Final = proxy.openai_receives(
proxy.gpt_5_6_dropping_params, [{"role": "user", "content": _CLIENT_BREAKPOINTS_ON_PARTS}], stream=stream
)
assert received["messages"] == [
{"role": "user", "content": _CLIENT_BREAKPOINTS_AS_CHAT_COMPLETIONS_TAKES_THEM}
], _CHAT_COMPLETIONS_REJECTS_TTL_AND_IMPLICIT
def test_without_drop_params_client_breakpoints_reach_openai_unchanged(proxy: _Proxy) -> None:
received: Final = proxy.openai_receives(
proxy.gpt_5_6_with_prompt_cache_options, [{"role": "user", "content": _CLIENT_BREAKPOINTS_ON_PARTS}]
)
assert received["messages"] == [{"role": "user", "content": _CLIENT_BREAKPOINTS_ON_PARTS}]
def test_drop_params_cuts_client_breakpoints_on_v1_messages_routed_to_chat_completions(proxy: _Proxy) -> None:
_ = proxy.openai.drain()
response: Final = proxy.gateway.request(
"POST",
"/v1/messages",
{
"model": proxy.gpt_5_6_dropping_params,
"max_tokens": 16,
"stop_sequences": ["END"],
"messages": [{"role": "user", "content": _CLIENT_BREAKPOINTS_ON_PARTS}],
},
)
assert response.status_code == 200, response.text
assert proxy.received_by_openai()["messages"] == [
{"role": "user", "content": _CLIENT_BREAKPOINTS_AS_CHAT_COMPLETIONS_TAKES_THEM}
], _CHAT_COMPLETIONS_REJECTS_TTL_AND_IMPLICIT

24
tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py
View file

@ -27,8 +27,9 @@ _MALFORMED: Final[tuple[tuple[str, JsonValue], ...]] = (
("5kb-string", "x" * 5000),
("empty-object", {}),
("bogus-mode", {"mode": "bogus"}),
("bad-ttl", {"mode": "explicit", "ttl": "1h"}),
("bogus-mode-with-ttl", {"mode": "bogus", "ttl": "30m"}),
)
_UNSUPPORTED_TTL: Final[Mapping[str, JsonValue]] = {"mode": "explicit", "ttl": "1h"}
_MALFORMED_IDS: Final = tuple(name for name, _ in _MALFORMED)
_MALFORMED_VALUES: Final = tuple(value for _, value in _MALFORMED)
_CASES: Final[tuple[tuple[str, Mode, JsonValue, JsonValue], ...]] = (
@ -192,6 +193,27 @@ def test_marker_with_an_unknown_key(bridge: _Bridge, mode: Mode, expected: JsonV
bridge.spend.landed(bridge.model(mode), call_id, marker)
@pytest.mark.parametrize(
("mode", "expected"), (("on", None), ("off", _UNSUPPORTED_TTL)), ids=("dropped-on", "verbatim-off")
)
def test_marker_with_an_unsupported_ttl(bridge: _Bridge, mode: Mode, expected: JsonValue) -> None:
marker: Final = uuid.uuid4().hex
call_id: Final = _completion(_chat(bridge, bridge.model(mode), [_user(marker, _UNSUPPORTED_TTL)]), marker)
pcb.assert_marker(_user_block_on_wire(bridge, marker), expected)
bridge.spend.landed(bridge.model(mode), call_id, marker)
def test_marker_with_an_unsupported_ttl_keeps_its_mode_on_a_host_opted_in_with_prompt_cache_options(
bridge: _Bridge,
) -> None:
marker: Final = uuid.uuid4().hex
with bridge.gateway.scenario() as scenario:
model: Final = scenario.model(model=pcb.MODEL, api_base=bridge.api_base, drop_params=True, **_OPTED_IN)
call_id: Final = _completion(_chat(bridge, model, [_user(marker, _UNSUPPORTED_TTL)]), marker)
pcb.assert_marker(_user_block_on_wire(bridge, marker), pcb.EXPLICIT)
bridge.spend.landed(model, call_id, marker)
def _second_block_on_wire(bridge: _Bridge, marker: str) -> dict[str, JsonValue]:
request: Final = pcb.posted(bridge.wire, marker)
_wire_body(request)

2
tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py
View file

@ -25,7 +25,7 @@ from pydantic import JsonValue, TypeAdapter
_INLINE_TOOLS: Final = "inline-tools-2026-09-15"
_THINKING: Final = "interleaved-thinking-2025-05-14"
_VERTEX_UNSUPPORTED: Final = "effort-2025-11-24"
_VERTEX_UNSUPPORTED: Final = "skills-2025-10-02"
_UNKNOWN: Final = "not-a-beta-2099-01-01"
_BACKEND: Final = "claude-opus-5-5"
_PROJECT: Final = "scripted-project"

36
tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py
View file

@ -303,7 +303,7 @@ def scenario_api_base(scenario: Scenario, scenario_id: str, response: StoredResp
def deployment(
scenario: Scenario,
group: str,
api_base: str,
api_base: str | None,
*,
order: int | None,
model: str = RESPONSES_MODEL,
@ -316,7 +316,7 @@ def deployment(
"litellm_params": {
"model": model,
"api_key": PROVIDER_KEY,
"api_base": api_base,
**({} if api_base is None else {"api_base": api_base}),
**({} if order is None else {"order": order}),
**extra_params,
},
@ -863,6 +863,38 @@ def test_hop_inside_one_encryption_boundary_keeps_the_reasoning(gateway: Gateway
assert [body.get("input") for body in bodies] == [history(mark), history(mark)], bodies
def test_hop_between_same_key_deployments_on_the_providers_default_endpoint_keeps_the_reasoning(
gateway: Gateway, tmp_path: Path
) -> None:
shared: Final = f"default-{uuid.uuid4().hex[:12]}"
with (
scratch_database() as database,
owned_proxy_process(
gateway,
tmp_path,
{"DATABASE_URL": database, "OPENAI_BASE_URL": f"{gateway.upstream_url}/{shared}"},
remove_environment=("DIRECT_URL", "DATABASE_URL_READ_REPLICA", "OPENAI_API_BASE"),
) as owned,
owned.gateway.scenario() as scenario,
):
owned.gateway.post("/config/update", {"router_settings": {"num_retries": 0}})
one: Final = scenario_api_base(scenario, f"{shared}-o1", failing()).rsplit("/", 1)[-1]
two: Final = scenario_api_base(scenario, f"{shared}-o2", healthy_json()).rsplit("/", 1)[-1]
order_one: Final = deployment(
scenario, shared, None, order=1, extra_params={"extra_headers": {"x-scripted-scenario": one}}
)
order_two: Final = deployment(
scenario, shared, None, order=2, extra_params={"extra_headers": {"x-scripted-scenario": two}}
)
mark: Final = marker()
answer: Final = responses_answer(
owned.gateway, HTTPX, Target(shared, frozenset({order_one, order_two})), history(mark)
)
assert_served_by(answer, order_two)
bodies: Final = bodies_for(observed(owned.gateway), shared)
assert [body.get("input") for body in bodies] == [history(mark), history(mark)], bodies
def assert_cross_group_hop_stripped(gateway: Gateway, source_scenario: str, target_scenario: str, mark: str) -> None:
records: Final = observed(gateway)
assert [body.get("input") for body in bodies_for(records, source_scenario)] == [history(mark)], records

47
tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
View file

@ -5884,9 +5884,9 @@ _HAND_WRITTEN_PROMPT_CACHE_BREAKPOINT_BLOCKS: Final = (
)
def test_convert_chat_completion_messages_to_responses_api_drops_malformed_prompt_cache_breakpoint_under_drop_params():
"""OpenAI's Responses API answered "Supported values are: '30m'" for a 1h breakpoint ttl on 2026-10-07,
so an unsupported ttl drops the marker as a unit while an unknown key is dropped from a valid one."""
def test_convert_chat_completion_messages_to_responses_api_drops_only_an_unreadable_prompt_cache_breakpoint_under_drop_params():
"""The target's Responses config decides which ttl and keys its API takes, so the bridge keeps every marker
whose mode some target reads."""
handler = LiteLLMResponsesTransformationHandler()
messages = [{"role": "user", "content": list(_HAND_WRITTEN_PROMPT_CACHE_BREAKPOINT_BLOCKS)}]
@ -5899,12 +5899,12 @@ def test_convert_chat_completion_messages_to_responses_api_drops_malformed_promp
None,
None,
None,
None,
{"mode": "explicit"},
{"mode": "explicit", "ttl": "1h"},
{"mode": "explicit", "scope": "all"},
{"mode": "explicit", "ttl": "30m"},
{"mode": "explicit"},
]
assert all("prompt_cache_breakpoint" not in block for block in content[:4])
assert all("prompt_cache_breakpoint" not in block for block in content[:3])
def test_convert_chat_completion_messages_to_responses_api_keeps_malformed_prompt_cache_breakpoint_by_default():
@ -5937,6 +5937,41 @@ def test_transform_request_drop_params_in_litellm_params_gates_the_prompt_cache_
assert "prompt_cache_breakpoint" not in result["input"][0]["content"][0]
_TTL_MARKED_CONTENT: Final = (
{"type": "text", "text": "unsupported ttl", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}},
{"type": "text", "text": "well formed", "prompt_cache_breakpoint": {"mode": "explicit"}},
)
@pytest.mark.parametrize(
"global_drop_params,deployment_drop_params,expected",
[
("true", None, [None, {"mode": "explicit"}]),
(False, True, [{"mode": "explicit", "ttl": "1h"}, {"mode": "explicit"}]),
(True, None, [{"mode": "explicit", "ttl": "1h"}, {"mode": "explicit"}]),
],
ids=("string-global-target-keeps-raw-markers", "deployment-drop-params", "global-drop-params"),
)
def test_transform_request_hands_a_marker_to_the_target_whole_only_when_the_target_trims_it(
monkeypatch: pytest.MonkeyPatch,
global_drop_params: object,
deployment_drop_params: bool | None,
expected: list[object],
) -> None:
monkeypatch.setattr(litellm, "drop_params", global_drop_params)
request: Final = LiteLLMResponsesTransformationHandler().transform_request(
model="gpt-6.1-sol",
messages=[{"role": "user", "content": list(_TTL_MARKED_CONTENT)}],
optional_params={},
litellm_params={"custom_llm_provider": "openai", "drop_params": deployment_drop_params},
headers={},
litellm_logging_obj=Mock(),
)
(item,) = cast(list[dict[str, object]], request["input"]) # cast-ok: the bridge emits message item mappings
parts: Final = cast(list[dict[str, object]], item["content"]) # cast-ok: message items carry content parts
assert [part.get("prompt_cache_breakpoint") for part in parts] == expected
def test_convert_response_output_folds_two_messages_and_function_call_into_one_choice() -> None:
first: Final = ResponseOutputMessage.model_validate(
{

11
tests/unit/integrations/test_anthropic_cache_control_hook.py
View file

@ -3485,12 +3485,23 @@ class TestOpenAIPromptCacheBreakpoint:
("gpt-5.6-sol", True),
("gpt-5.6-luna", True),
("gpt-5.7", True),
("gpt-5.6o", True),
("gpt-5.10", True),
("gpt-5.06", True),
("gpt-6", True),
("GPT-5.6", True),
("gpt-5.5", False),
("gpt-5", False),
("gpt-5-chat-latest", False),
("gpt-4.1", False),
("gpt-5-6-sol", True),
("gpt-6-1", True),
("gpt-5-4-mini", False),
("gpt-5-32k", False),
("gpt-5-6b", False),
("gpt-5-06", False),
("gpt-5-2026-03-01", False),
("gpt-4-1106-preview", False),
("o3", False),
("claude-sonnet-4-5", False),
],

18
tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py
View file

@ -2343,7 +2343,8 @@ _MALFORMED_PROMPT_CACHE_BREAKPOINTS: Final = (
"",
{},
{"mode": "bogus"},
{"mode": "explicit", "ttl": "1h"},
{"mode": "bogus", "ttl": "30m"},
{1: "explicit"},
)
_MALFORMED_PROMPT_CACHE_BREAKPOINT_IDS: Final = (
"string",
@ -2352,7 +2353,8 @@ _MALFORMED_PROMPT_CACHE_BREAKPOINT_IDS: Final = (
"empty-string",
"empty-object",
"bogus-mode",
"bad-ttl",
"bogus-mode-with-ttl",
"non-string-key",
)
_VALID_PROMPT_CACHE_BREAKPOINTS: Final = ({"mode": "explicit"}, {"mode": "explicit", "ttl": "30m"})
@ -2382,6 +2384,18 @@ def test_prompt_cache_breakpoint_for_wire_strips_an_unknown_key_under_drop_param
assert prompt_cache_breakpoint_for_wire({"mode": "explicit", "note": "kept"}, True) == {"mode": "explicit"}
@pytest.mark.parametrize(
"marker",
({"mode": "explicit", "ttl": "1h"}, {"mode": "explicit", "ttl": "1h", "note": "kept"}),
ids=("bad-ttl", "bad-ttl-and-unknown-key"),
)
def test_prompt_cache_breakpoint_for_wire_drops_an_unsupported_ttl_and_keeps_the_mode_under_drop_params(
marker: dict[str, str],
) -> None:
assert prompt_cache_breakpoint_for_wire(marker, True) == {"mode": "explicit"}
assert prompt_cache_breakpoint_for_wire(marker, False) is marker
def test_cache_control_is_replaced_by_an_explicit_prompt_cache_breakpoint_on_the_same_block() -> None:
marked: Final = {"type": "text", "text": "rules", "cache_control": {"type": "ephemeral", "ttl": "1h"}}

102
tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py
View file

@ -6788,6 +6788,34 @@ def test_eager_input_streaming_not_set_on_computer_use_tool():
assert "eager_input_streaming" not in mapped_tool
def _computer_tool_with_defer_loading_and_allowed_callers(tool_type: str) -> dict[str, object]:
return {
"type": tool_type,
"function": {"name": "computer", "parameters": {"display_width_px": 1024, "display_height_px": 768}},
"defer_loading": True,
"allowed_callers": ["code_execution_20250825"],
}
@pytest.mark.parametrize("tool_type", ["computer_20241022", "computer_20250124"])
def test_legacy_computer_tools_drop_defer_loading_and_allowed_callers(tool_type: str):
mapped_tool, _ = AnthropicConfig().map_tool_helper(_computer_tool_with_defer_loading_and_allowed_callers(tool_type))
assert mapped_tool["type"] == tool_type
assert "defer_loading" not in mapped_tool
assert "allowed_callers" not in mapped_tool
def test_computer_20251124_keeps_defer_loading_and_allowed_callers():
mapped_tool, _ = AnthropicConfig().map_tool_helper(
_computer_tool_with_defer_loading_and_allowed_callers("computer_20251124")
)
assert mapped_tool["type"] == "computer_20251124"
assert mapped_tool["defer_loading"] is True
assert mapped_tool["allowed_callers"] == ["code_execution_20250825"]
def test_eager_input_streaming_reaches_anthropic_request_tools():
result = AnthropicConfig().map_openai_params(
non_default_params={"tools": [_eager_chat_tool(eager_input_streaming=True)], "stream": True},
@ -7520,10 +7548,10 @@ def test_process_anthropic_headers_with_no_matching_headers():
@pytest.mark.parametrize(
"computer_tool_used, prompt_caching_set, expected_beta_header",
[
(True, False, True),
(False, True, False),
(True, True, True),
(False, False, False),
("computer_20250124", False, True),
(None, True, False),
("computer_20250124", True, True),
(None, False, False),
],
)
def test_anthropic_beta_header(computer_tool_used, prompt_caching_set, expected_beta_header):
@ -9233,6 +9261,7 @@ async def test_litellm_anthropic_prompt_caching_system(
}
@pytest.mark.usefixtures("local_beta_headers_config")
def test_completion_forwards_caller_anthropic_beta_header_with_computer_tool(respx_mock: respx.MockRouter):
route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=httpx.Response(
@ -9271,7 +9300,7 @@ def test_completion_forwards_caller_anthropic_beta_header_with_computer_tool(res
sent_betas: Final = frozenset(
beta.strip() for beta in route.calls.last.request.headers["anthropic-beta"].split(",")
)
assert sent_betas == frozenset({"computer-use-2025-01-24"})
assert sent_betas == frozenset({"computer-use-2024-10-22", "computer-use-2025-01-24"})
assert json.loads(route.calls.last.request.content)["tools"] == [
{
"type": "computer_20241022",
@ -9283,6 +9312,69 @@ def test_completion_forwards_caller_anthropic_beta_header_with_computer_tool(res
]
# Tool-to-beta pairs from Anthropic's computer use docs; api.anthropic.com answered 400 to
# computer_20251124 without computer-use-2025-11-24 and 200 with it on 2026-10-31
@pytest.mark.usefixtures("local_beta_headers_config")
@pytest.mark.parametrize(
("tool_type", "expected_beta"),
[
("computer_20241022", "computer-use-2024-10-22"),
("computer_20250124", "computer-use-2025-01-24"),
("computer_20251124", "computer-use-2025-11-24"),
],
)
def test_completion_adds_the_beta_matching_the_computer_tool_version(
respx_mock: respx.MockRouter, tool_type: str, expected_beta: str
):
route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=httpx.Response(200, json=_anthropic_message_body([{"type": "text", "text": "hi"}], "end_turn"))
)
litellm.completion(
model="anthropic/claude-opus-5",
api_key="mock_api_key",
messages=[{"role": "user", "content": "Say hi"}],
tools=[
{
"type": tool_type,
"function": {
"name": "computer",
"parameters": {"display_height_px": 768, "display_width_px": 1024},
},
}
],
)
sent_betas: Final = frozenset(
beta.strip() for beta in route.calls.last.request.headers["anthropic-beta"].split(",")
)
assert sent_betas == frozenset({expected_beta})
@pytest.mark.usefixtures("local_beta_headers_config")
def test_completion_sends_no_computer_use_beta_for_an_unknown_computer_tool_version(respx_mock: respx.MockRouter):
route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=httpx.Response(200, json=_anthropic_message_body([{"type": "text", "text": "hi"}], "end_turn"))
)
litellm.completion(
model="anthropic/claude-opus-5",
api_key="mock_api_key",
messages=[{"role": "user", "content": "Say hi"}],
tools=[
{
"type": "computer_20991231",
"function": {
"name": "computer",
"parameters": {"display_height_px": 768, "display_width_px": 1024},
},
}
],
)
assert "anthropic-beta" not in route.calls.last.request.headers
def _anthropic_message_body(content: list[dict[str, object]], stop_reason: str) -> dict[str, object]:
return {
"id": "msg_01XrAv7gc5tQNDuoADra7vB4",

239
tests/unit/llms/anthropic/messages/context_management/test_inference_view.py Normal file
View file

@ -0,0 +1,239 @@
import logging
from typing import Final
from unittest.mock import patch
import pytest
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.llms.anthropic.messages.context_management.inference_view import (
CountContextEdits,
InferenceRequestView,
InferenceRequestViewError,
inference_request_view,
)
_NO_EDITS: Final = CountContextEdits(spec=None, model="gemini/gemini-2.5-flash", tools=None)
@pytest.mark.parametrize(
"system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system")
)
async def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None:
view: Final = await inference_request_view(
(
{"role": "user", "content": "first question " * 50},
{"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]},
{"role": "user", "content": "latest question"},
),
system,
"anthropic",
_NO_EDITS,
)
assert isinstance(view, InferenceRequestView)
assert view.messages == [{"role": "user", "content": "latest question"}]
assert "Summary of the earlier turns." in str(view.system)
assert "Be brief." in str(view.system)
async def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None:
snippet: Final = "LiteLLM ships a weekly release candidate and a stable line."
view: Final = await inference_request_view(
(
{"role": "user", "content": "What is new?"},
{
"role": "assistant",
"content": [
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}},
{
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_01",
"content": [
{
"type": "web_search_result",
"url": "https://docs.litellm.ai/release_notes",
"title": "Release notes",
"encrypted_content": "",
"snippet": snippet,
}
],
},
],
},
),
None,
"anthropic",
_NO_EDITS,
)
assert isinstance(view, InferenceRequestView)
assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"])
assert [block["type"] for block in assistant_blocks] == ["text"]
assert snippet in assistant_blocks[0]["text"]
async def test_openai_view_is_the_request_unchanged() -> None:
messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},)
view: Final = await inference_request_view(messages, "Be brief.", "openai", _NO_EDITS)
assert isinstance(view, InferenceRequestView)
assert view.messages == messages
assert view.system == "Be brief."
_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180
_READ_FILE_TOOL: Final[dict[str, object]] = {
"name": "read_file",
"description": "Read a file",
"input_schema": {"type": "object", "properties": {"path": {"type": "string"}}},
}
_CLEAR_EDIT: Final[JsonValue] = {
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 500},
"keep": {"type": "tool_uses", "value": 1},
}
_CLEAR_OLD_TOOL_RESULTS: Final[JsonValue] = {"edits": [_CLEAR_EDIT]}
def _agent_turns(rounds: int, first_round: int = 0) -> tuple[dict[str, JsonValue], ...]:
return (
*(
turn
for index in range(first_round, first_round + rounds)
for turn in (
{"role": "user", "content": f"Read file {index}."},
{
"role": "assistant",
"content": [
{"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}}
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _TOOL_OUTPUT}],
},
)
), # comprehension-ok: flattens each round into its three turns
{"role": "user", "content": "Summarize what you read."},
)
def _tool_outputs_kept(view: InferenceRequestView) -> tuple[bool, ...]:
return tuple(
_TOOL_OUTPUT in str(message["content"]) for message in view.messages if "tool_result" in str(message["content"])
)
def _edits(spec: JsonValue) -> CountContextEdits:
return CountContextEdits(spec=spec, model="gemini/gemini-2.5-flash", tools=(_READ_FILE_TOOL,))
async def test_anthropic_view_clears_the_tool_results_inference_clears() -> None:
view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (False, False, True)
async def test_anthropic_view_applies_clear_edits_without_calling_the_compaction_summary_model() -> None:
spec: Final[JsonValue] = {
"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50000}}, _CLEAR_EDIT]
}
turns: Final = (
{"role": "user", "content": "first question " * 30000},
{"role": "assistant", "content": "first answer"},
*_agent_turns(3),
)
with (
patch(
"litellm.llms.anthropic.messages.context_management.editors.compact._read_summary_model_setting",
return_value="gemini/gemini-2.5-flash",
),
patch.object(litellm, "acompletion", side_effect=AssertionError("count_tokens must not summarize")) as summary,
):
view: Final = await inference_request_view(turns, None, "anthropic", _edits(spec))
assert summary.call_count == 0
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (False, False, True)
async def test_anthropic_view_counts_the_unedited_body_without_an_error_log_when_an_edit_fails(
caplog: pytest.LogCaptureFixture,
) -> None:
with (
patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")),
caplog.at_level(logging.DEBUG, logger="LiteLLM"),
):
view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True, True)
assert any("skipping edits" in record.getMessage() for record in caplog.records)
assert not any(record.levelno >= logging.ERROR for record in caplog.records)
async def test_anthropic_view_counts_the_turns_after_a_compaction_block_when_an_edit_fails() -> None:
turns: Final = (
*_agent_turns(3)[:-1],
{"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]},
*_agent_turns(2, first_round=3),
)
with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")):
view: Final = await inference_request_view(turns, None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True)
assert "Read files 0 to 2." in str(view.system)
async def test_anthropic_view_clears_only_the_turns_after_a_compaction_block_when_clear_runs_before_compact() -> None:
clear_before_compact: Final[JsonValue] = {
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "tool_uses", "value": 3},
"keep": {"type": "tool_uses", "value": 1},
},
{"type": "compact_20260112"},
]
}
turns: Final = (
*_agent_turns(3)[:-1],
{"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]},
*_agent_turns(2, first_round=3),
)
with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")):
view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True)
assert "Read files 0 to 2." in str(view.system)
async def test_anthropic_view_leaves_a_system_that_is_not_text_blocks_and_its_turns_to_the_provider_check() -> None:
image_system: Final[JsonValue] = [{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}]
view: Final = await inference_request_view(
_agent_turns(3), image_system, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)
)
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True, True)
assert view.system == image_system
async def test_anthropic_view_rejects_a_compact_trigger_inference_rejects() -> None:
low_trigger: Final[JsonValue] = {
"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}]
}
view: Final = await inference_request_view(_agent_turns(1), None, "anthropic", _edits(low_trigger))
assert isinstance(view, InferenceRequestViewError)
assert view.status_code == 400

257
tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py
View file

@ -3,10 +3,13 @@ import copy
import json
import os
import uuid
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Final, List
import httpx
import openai
import responses
import respx
import pytest
@ -1136,6 +1139,260 @@ async def test_messages_strips_provider_prefix_exactly_once(requested_model, exp
assert captured["url"] == expected_url
_GET_WEATHER_TOOL: Final = {
"name": "get_weather",
"description": "Get the current weather for a city",
"input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"tools, expected_path",
[
([_GET_WEATHER_TOOL], "/openai/v1/responses"),
(None, "/openai/deployments/gpt-6-luna/chat/completions"),
],
ids=["function_tool_bridged_to_responses", "no_tools_on_chat"],
)
@pytest.mark.parametrize("credentials_from_env", [False, True], ids=["credentials_in_params", "credentials_from_env"])
async def test_azure_ai_on_azure_openai_host_streams_natively_with_the_bare_deployment(
tools: list[dict[str, object]] | None,
expected_path: str,
credentials_from_env: bool,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
monkeypatch.delenv("AZURE_API_BASE", raising=False)
monkeypatch.delenv("AZURE_API_KEY", raising=False)
monkeypatch.setenv("AZURE_AI_API_BASE", "https://qa-resource.cognitiveservices.azure.com")
monkeypatch.setenv("AZURE_AI_API_KEY", "env-api-key")
credentials: Final = (
{}
if credentials_from_env
else {"api_base": "https://qa-resource.cognitiveservices.azure.com", "api_key": "params-api-key"}
)
with respx.mock(assert_all_called=True) as respx_mock:
upstream: Final = respx_mock.route(host="qa-resource.cognitiveservices.azure.com").mock(
side_effect=httpx.ConnectError("cut at the wire")
)
with pytest.raises(openai.APIError):
await litellm.anthropic.messages.acreate(
max_tokens=100,
messages=[{"role": "user", "content": "What is the weather in Berlin?"}],
model="azure_ai/gpt-6-luna",
stream=True,
tools=tools,
**credentials,
)
bodies: Final = [json.loads(call.request.content) for call in upstream.calls]
assert {call.request.url.path for call in upstream.calls} == {expected_path}
assert {call.request.headers["api-key"] for call in upstream.calls} == {credentials.get("api_key", "env-api-key")}
assert {body["model"] for body in bodies} == {"gpt-6-luna"}
assert all(body["stream"] is True for body in bodies)
@pytest.mark.asyncio
async def test_bridged_messages_keep_the_region_of_a_regional_bedrock_mantle_model(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
with respx.mock(assert_all_called=True) as respx_mock:
upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire"))
with pytest.raises(litellm.exceptions.ServiceUnavailableError):
await litellm.anthropic.messages.acreate(
max_tokens=100,
messages=[{"role": "user", "content": "ping"}],
model="bedrock_mantle/eu-west-1/openai.gpt-oss-120b",
api_key="test-api-key",
)
assert {call.request.url.host for call in upstream.calls} == {"bedrock-mantle.eu-west-1.api.aws"}
@pytest.mark.asyncio
@pytest.mark.parametrize("model", ["anthropic/claude-sonnet-5-5", "ragflow/gpt-4o"])
async def test_bridged_messages_keep_the_callers_explicit_provider_over_the_model_prefix(
model: str, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
with respx.mock(assert_all_called=True) as respx_mock:
upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire"))
with pytest.raises(openai.APIError):
await litellm.anthropic.messages.acreate(
max_tokens=100,
messages=[{"role": "user", "content": "ping"}],
model=model,
custom_llm_provider="openai",
api_base="https://gateway.example.com/v1",
api_key="test-api-key",
)
assert {call.request.url.host for call in upstream.calls} == {"gateway.example.com"}
assert {json.loads(call.request.content)["model"] for call in upstream.calls} == {model}
@pytest.mark.parametrize("model", ["chatgpt/gpt-5", "github_copilot/gpt-5"])
def test_bridged_messages_reach_the_callers_provider_without_signing_in_to_the_prefix_provider(
model: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("OPENAI_API_KEY", "server-openai-key")
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
litellm.in_memory_llm_clients_cache.flush_cache()
with respx.mock(assert_all_called=True) as respx_mock:
upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire"))
with pytest.raises(openai.APIError):
litellm.anthropic.messages.create(
max_tokens=100,
messages=[{"role": "user", "content": "ping"}],
model=model,
custom_llm_provider="openai",
)
assert {call.request.url.host for call in upstream.calls} == {"api.openai.com"}
assert {call.request.headers["authorization"] for call in upstream.calls} == {"Bearer server-openai-key"}
@pytest.mark.asyncio
async def test_bridged_messages_send_the_deployments_key_for_azure_models_served_as_openai(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("OPENAI_API_KEY", "server-openai-key")
litellm.in_memory_llm_clients_cache.flush_cache()
with respx.mock(assert_all_called=True) as respx_mock:
upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire"))
with pytest.raises(openai.APIError):
await litellm.anthropic.messages.acreate(
max_tokens=100,
messages=[{"role": "user", "content": "ping"}],
model="azure/mistral-large-latest",
api_base="https://example-resource.services.ai.azure.com/models",
api_key="deployment-key",
)
assert {call.request.headers["authorization"] for call in upstream.calls} == {"Bearer deployment-key"}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"entra_credential",
[
{"azure_ad_token": "entra-token"},
{"azure_ad_token_provider": lambda: "entra-token"},
{"tenant_id": "tenant", "client_id": "client", "client_secret": "secret"},
{"azure_username": "user", "azure_password": "password", "client_id": "client"},
],
)
@pytest.mark.parametrize("api_base_from_env", [False, True], ids=["api_base_in_params", "api_base_from_env"])
async def test_bridged_messages_keep_the_deployments_entra_credential_over_the_env_key(
monkeypatch: pytest.MonkeyPatch, entra_credential: Mapping[str, object], api_base_from_env: bool
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.cognitiveservices.azure.com")
monkeypatch.setenv("AZURE_AI_API_KEY", "azure-env-key")
litellm.in_memory_llm_clients_cache.flush_cache()
api_base: Final = {} if api_base_from_env else {"api_base": "https://example-resource.cognitiveservices.azure.com"}
with (
patch("litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", return_value=lambda: "entra-token"), # test-quality-ok: stands in for the Entra token fetch so the test can see which credential reaches Azure
patch(
"litellm.llms.azure.common_utils.get_azure_ad_token_from_username_password",
return_value=lambda: "entra-token",
), # test-quality-ok: stands in for the Entra token fetch so the test can see which credential reaches Azure
respx.mock(assert_all_called=True) as respx_mock,
):
upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire"))
with pytest.raises(openai.APIError):
await litellm.anthropic.messages.acreate(
max_tokens=100,
messages=[{"role": "user", "content": "ping"}],
model="azure_ai/gpt-6-luna",
tools=[{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}],
**api_base,
**entra_credential,
)
assert {call.request.headers.get("authorization") for call in upstream.calls} == {"Bearer entra-token"}
assert {call.request.headers.get("api-key") for call in upstream.calls} == {None}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"azure_env_credential, expected_header, expected_value",
[
({"AZURE_API_KEY": "azure-env-key"}, "api-key", "azure-env-key"),
(
{"AZURE_TENANT_ID": "tenant", "AZURE_CLIENT_ID": "client", "AZURE_CLIENT_SECRET": "secret"},
"authorization",
"Bearer entra-token",
),
],
ids=["azure_env_key", "azure_env_entra"],
)
async def test_bridged_messages_keep_the_azure_env_credential_when_the_deployment_sets_its_own_api_base(
azure_env_credential: Mapping[str, str], expected_header: str, expected_value: str, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
for name in ("AZURE_API_KEY", "AZURE_OPENAI_API_KEY", "AZURE_AD_TOKEN", "AZURE_TENANT_ID", "AZURE_CLIENT_ID"):
monkeypatch.delenv(name, raising=False)
monkeypatch.setenv("AZURE_AI_API_KEY", "key-for-another-resource")
for name, value in azure_env_credential.items():
monkeypatch.setenv(name, value)
litellm.in_memory_llm_clients_cache.flush_cache()
with (
patch("litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", return_value=lambda: "entra-token"), # test-quality-ok: stands in for the Entra token fetch so the test can see which credential reaches Azure
respx.mock(assert_all_called=True) as respx_mock,
):
upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire"))
with pytest.raises(openai.APIError):
await litellm.anthropic.messages.acreate(
max_tokens=100,
messages=[{"role": "user", "content": "ping"}],
model="azure_ai/gpt-6-luna",
api_base="https://example-resource.cognitiveservices.azure.com",
tools=[{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}],
)
assert {call.request.headers.get(expected_header) for call in upstream.calls} == {expected_value}
@pytest.mark.asyncio
async def test_bridged_messages_still_route_through_the_litellm_proxy_when_asked(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("LITELLM_PROXY_API_BASE", "https://proxy.example.com")
monkeypatch.setenv("LITELLM_PROXY_API_KEY", "proxy-key")
monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.cognitiveservices.azure.com")
monkeypatch.setenv("AZURE_AI_API_KEY", "azure-env-key")
litellm.in_memory_llm_clients_cache.flush_cache()
with respx.mock(assert_all_called=True) as respx_mock:
upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire"))
with pytest.raises(openai.APIError):
await litellm.anthropic.messages.acreate(
max_tokens=100,
messages=[{"role": "user", "content": "ping"}],
model="azure_ai/gpt-6-luna",
use_litellm_proxy=True,
)
assert {call.request.url.host for call in upstream.calls} == {"proxy.example.com"}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model_info, expected_url",

28
tests/unit/llms/anthropic/messages/native/test_mcp_handler.py
View file

@ -2,8 +2,12 @@ from importlib import import_module
from typing import Final
from unittest.mock import AsyncMock, patch
import httpx
import openai
import pytest
import respx
import litellm
from litellm.llms.anthropic.messages.adapters.transformation import (
LiteLLMAnthropicMessagesAdapter,
)
@ -56,6 +60,30 @@ def test_anthropic_messages_handler_routes_litellm_proxy_mcp_to_the_gateway():
assert result is not None
@pytest.mark.asyncio
async def test_anthropic_messages_with_mcp_reaches_an_env_only_deployment_with_its_env_credentials(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.cognitiveservices.azure.com")
monkeypatch.setenv("AZURE_AI_API_KEY", "azure-env-key")
litellm.in_memory_llm_clients_cache.flush_cache()
with respx.mock(assert_all_called=True) as respx_mock:
upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire"))
with pytest.raises(openai.APIError):
await litellm.anthropic.messages.acreate(
max_tokens=100,
messages=[{"role": "user", "content": "ping"}],
model="azure_ai/gpt-6-luna",
tools=[{**MCP_REFERENCE, "server_url": "litellm_proxy"}],
)
assert {call.request.url.host for call in upstream.calls} == {"example-resource.cognitiveservices.azure.com"}
assert {call.request.headers.get("api-key") for call in upstream.calls} == {"azure-env-key"}
assert not any(b"litellm_proxy" in call.request.content for call in upstream.calls)
def test_anthropic_messages_handler_skips_the_gateway_on_recursion():
"""The gateway's own follow-up call must not re-enter the gateway."""
with patch(

86
tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py
View file

@ -806,3 +806,89 @@ def test_azure_keeps_a_later_developer_message_in_place_as_system(config: AzureO
)
assert [message["role"] for message in messages] == ["system", "user", "assistant", "system", "user"]
assert messages[3]["content"] == "Answer with exactly one word."
_MARKED_AZURE_PARTS: Final = (
{"type": "text", "text": "Stable prefix", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}},
{"type": "text", "text": "Implicit", "prompt_cache_breakpoint": {"mode": "implicit"}},
)
_TRIMMED_AZURE_PARTS: Final = [
{"type": "text", "text": "Stable prefix", "prompt_cache_breakpoint": {"mode": "explicit"}},
{"type": "text", "text": "Implicit"},
]
@pytest.fixture
def registered_azure_breakpoint_reader(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setitem(
litellm.model_cost,
"azure/unit-breakpoint-reader",
{"litellm_provider": "azure", "mode": "chat", "supports_prompt_cache_breakpoint": True},
)
monkeypatch.setitem(
litellm.model_cost,
"azure/unit-breakpoint-ignorer",
{"litellm_provider": "azure", "mode": "chat", "supports_prompt_cache_breakpoint": False},
)
def _azure_request_parts(
config: AzureOpenAIConfig, model: str, litellm_params: dict[str, object]
) -> object:
request: Final = config.transform_request(
model=model,
messages=[{"role": "user", "content": list(_MARKED_AZURE_PARTS)}],
optional_params={},
litellm_params=litellm_params,
headers={},
)
return request["messages"][0]["content"]
@pytest.mark.parametrize("config", [AzureOpenAIConfig(), AzureOpenAIGPT5Config()])
@pytest.mark.parametrize(
"model, litellm_params",
[
("unit-breakpoint-reader", {"drop_params": True}),
("my-azure-alias", {"drop_params": True, "base_model": "azure/unit-breakpoint-reader"}),
("my-azure-alias", {"drop_params": True, "model_info": {"base_model": "azure/unit-breakpoint-reader"}}),
("my-azure-alias", {"drop_params": True}),
],
ids=("deployment-name", "litellm-params-base-model", "model-info-base-model", "unknown-deployment"),
)
def test_azure_drop_params_cuts_part_breakpoints_down_to_what_azure_chat_accepts(
registered_azure_breakpoint_reader: None,
config: AzureOpenAIConfig,
model: str,
litellm_params: dict[str, object],
) -> None:
"""Azure OpenAI chat answered 400 for a content part's breakpoint ttl and for implicit mode on 2026-10-29,
the same as OpenAI chat completions."""
assert _azure_request_parts(config, model, litellm_params) == _TRIMMED_AZURE_PARTS
@pytest.mark.parametrize("litellm_params", [{}, {"drop_params": False}], ids=("no-drop-params", "drop-params-false"))
def test_azure_sends_part_breakpoints_as_written_without_drop_params(
registered_azure_breakpoint_reader: None, litellm_params: dict[str, object]
) -> None:
assert _azure_request_parts(AzureOpenAIConfig(), "unit-breakpoint-reader", litellm_params) == list(
_MARKED_AZURE_PARTS
)
@pytest.mark.parametrize(
"model, litellm_params",
[
("unit-breakpoint-ignorer", {"drop_params": True}),
("unit-breakpoint-reader", {"drop_params": True, "base_model": "azure/unit-breakpoint-ignorer"}),
("my-azure-alias", {"drop_params": True, "model_info": {"base_model": "azure/unit-breakpoint-ignorer"}}),
],
ids=("known-non-reader", "base-model-known-non-reader", "model-info-base-model-known-non-reader"),
)
def test_azure_drop_params_strips_part_breakpoints_from_a_known_non_reader(
registered_azure_breakpoint_reader: None, model: str, litellm_params: dict[str, object]
) -> None:
assert _azure_request_parts(AzureOpenAIConfig(), model, litellm_params) == [
{"type": "text", "text": "Stable prefix"},
{"type": "text", "text": "Implicit"},
]

94
tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py
View file

@ -1,5 +1,7 @@
import json
import operator
import threading
from collections.abc import Mapping
from typing import Final
import httpx
@ -8,8 +10,9 @@ import respx
from pydantic import ValidationError
import litellm
from litellm import anthropic_beta_headers_manager
from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAnthropicTokenCounter
from litellm.types.utils import TokenCountResponse
from litellm.types.utils import CountContextManagement, TokenCountResponse
API_BASE: Final = "https://my-resource.services.ai.azure.com"
COUNT_URL: Final = f"{API_BASE}/anthropic/v1/messages/count_tokens"
@ -268,3 +271,92 @@ async def test_count_sends_the_tool_change_beta_when_the_history_adds_a_tool(mon
assert result is not None
assert result.total_tokens == 541
assert route.calls.last.request.headers["anthropic-beta"] == "mid-conversation-tool-changes-2026-07-01,token-counting-2024-11-01"
async def _betas_azure_inference_sends(context_management: Mapping[str, object]) -> frozenset[str]:
with respx.mock:
route: Final = respx.post(f"{API_BASE}/anthropic/v1/messages").mock(
return_value=httpx.Response(
200,
json={
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5-5",
"content": [{"type": "text", "text": "ok"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 12, "output_tokens": 2},
},
)
)
await litellm.anthropic.messages.acreate(
max_tokens=64,
messages=[{"role": "user", "content": "Summarize the thread."}],
model="azure_ai/claude-sonnet-5-5",
api_key="azure-key",
api_base=API_BASE,
context_management=context_management,
)
return frozenset(filter(None, route.calls.last.request.headers.get("anthropic-beta", "").split(",")))
@pytest.mark.asyncio
@pytest.mark.usefixtures("local_beta_headers_config")
@pytest.mark.parametrize("edit_type", ["clear_tool_uses_20250919", "compact_20260112"])
async def test_count_forwards_the_context_edits_with_the_betas_azure_inference_sends(
monkeypatch: pytest.MonkeyPatch, edit_type: str
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
spec: Final = {"edits": [{"type": edit_type}]}
inference_betas: Final = await _betas_azure_inference_sends(spec)
with respx.mock:
route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 88}))
result: Final = await AzureAIAnthropicTokenCounter().count_tokens(
model_to_use="claude-sonnet-5-5",
messages=[{"role": "user", "content": "Summarize the thread."}],
contents=None,
deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}},
request_model="azure_ai/claude-sonnet-5-5",
context_management=CountContextManagement(spec=spec, drop_invalid_edits=False),
)
assert result is not None
assert result.total_tokens == 88
assert json.loads(route.calls.last.request.content)["context_management"] == spec
assert frozenset(route.calls.last.request.headers["anthropic-beta"].split(",")) == inference_betas | {
"token-counting-2024-11-01"
}
class _FetchThreads:
seen: tuple[int, ...] = ()
def __call__(self, request: httpx.Request) -> httpx.Response:
self.seen = (*self.seen, threading.get_ident())
return httpx.Response(200, json={})
@pytest.mark.asyncio
async def test_count_with_an_edit_fetches_the_remote_beta_config_off_the_event_loop(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False)
monkeypatch.setattr(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None)
fetch_threads: Final = _FetchThreads()
with respx.mock:
respx.get(litellm.anthropic_beta_headers_url).mock(side_effect=fetch_threads)
respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 88}))
result: Final = await AzureAIAnthropicTokenCounter().count_tokens(
model_to_use="claude-sonnet-5-5",
messages=[{"role": "user", "content": "Summarize the thread."}],
contents=None,
deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}},
request_model="azure_ai/claude-sonnet-5-5",
context_management=CountContextManagement(
spec={"edits": [{"type": "clear_tool_uses_20250919"}]}, drop_invalid_edits=False
),
)
assert result is not None
assert len(fetch_threads.seen) == 1
assert threading.get_ident() not in fetch_threads.seen

6
tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py
View file

@ -1271,6 +1271,12 @@ def test_azure_ai_falls_back_to_the_model_name_when_the_base_model_row_does_not_
{"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}},
),
({"type": "ephemeral"}, {"drop_params": True}, {"type": "text", "text": "stable"}),
({"mode": "implicit"}, {"drop_params": True}, {"type": "text", "text": "stable"}),
(
{"mode": "explicit", "ttl": "1h"},
{"drop_params": True},
{"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}},
),
(
{"mode": "explicit", "ttl": "30m"},
{},

18
tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py
View file

@ -1,6 +1,8 @@
import base64
import json
import pytest
from litellm.llms.bedrock.count_tokens.transformation import (
DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS,
BedrockCountTokensConfig,
@ -53,6 +55,22 @@ def test_detect_input_type_converse_blocks_route_to_converse():
assert config._detect_input_type(request) == "converse"
@pytest.mark.parametrize("content", [[{"text": "hi"}], "hi"])
@pytest.mark.parametrize(
"context_management", [None, {"edits": [{"type": "clear_tool_uses_20250919"}]}]
)
def test_a_pass_through_body_with_context_management_keeps_its_converse_input(content, context_management):
request = {
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
"messages": [{"role": "user", "content": content}],
"context_management": context_management,
}
transformed = BedrockCountTokensConfig().transform_anthropic_to_bedrock_count_tokens(request)
assert transformed["input"]["converse"]["messages"] == [{"role": "user", "content": [{"text": "hi"}]}]
def test_transform_to_invoke_model_format_base64_encodes_body():
"""The CountTokens API expects invokeModel.body as a base64-encoded blob;
Anthropic Messages bodies additionally need anthropic_version/max_tokens

99
tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py
View file

@ -9,6 +9,7 @@ from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.types.utils import CountContextManagement
RUNTIME_HOST: Final = "bedrock-runtime.us-east-1.amazonaws.com"
MANTLE_HOST: Final = "bedrock-mantle.us-east-1.api.aws"
@ -144,6 +145,104 @@ async def test_mantle_failure_is_reported_with_its_status() -> None:
assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
@pytest.mark.asyncio
async def test_runtime_rejection_of_an_edit_stands_when_mantle_does_not_serve_the_model() -> None:
rejection: Final = "'claude-haiku-4-5-20251001' does not support the 'compact_20260112' context management strategy"
bedrock: Final = _Bedrock(
runtime=(400, {"message": rejection}),
mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}),
)
result: Final = await _counter(bedrock).count_tokens(
model_to_use="us.anthropic.claude-haiku-4-5-20251001-v1:0",
messages=MESSAGES,
contents=None,
deployment=DEPLOYMENT,
request_model="claude-haiku-4-5",
context_management=CountContextManagement(
spec={"edits": [{"type": "compact_20260112"}]}, drop_invalid_edits=False
),
)
assert result is not None
assert result.error is True
assert result.status_code == 400
assert result.tokenizer_type == "bedrock_api"
assert rejection in (result.error_message or "")
assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
@pytest.mark.asyncio
async def test_runtime_400_without_an_edit_still_reports_the_mantle_failure() -> None:
bedrock: Final = _Bedrock(
runtime=(400, {"message": "The provided model identifier is invalid."}),
mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}),
)
result: Final = await _counter(bedrock).count_tokens(
model_to_use="us.anthropic.claude-haiku-4-5-20251001-v1:0",
messages=MESSAGES,
contents=None,
deployment=DEPLOYMENT,
request_model="claude-haiku-4-5",
)
assert result is not None
assert result.status_code == 404
assert result.tokenizer_type == "bedrock_mantle_api"
assert "does not exist" in (result.error_message or "")
assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
@pytest.mark.asyncio
async def test_plain_text_history_with_an_edit_is_counted_through_invoke_carrying_the_edit() -> None:
bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 57}), mantle=(200, {"input_tokens": 1}))
spec: Final[dict[str, JsonValue]] = {"edits": [{"type": "clear_tool_uses_20250919"}]}
result: Final = await _counter(bedrock).count_tokens(
model_to_use="global.anthropic.claude-opus-4-8",
messages=MESSAGES,
contents=None,
deployment=DEPLOYMENT,
request_model="claude-opus-4-8",
context_management=CountContextManagement(spec=spec, drop_invalid_edits=False),
)
assert result is not None
assert result.total_tokens == 57
assert bedrock.posted_hosts() == (RUNTIME_HOST,)
invoke_body: Final = _invoke_body(bedrock.requests[0])
assert invoke_body["context_management"] == spec
assert invoke_body["anthropic_beta"] == ["context-management-2025-06-27"]
@pytest.mark.asyncio
async def test_an_empty_string_turn_with_an_edit_is_sent_as_an_empty_string_not_an_empty_text_block() -> None:
bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 16}), mantle=(200, {"input_tokens": 1}))
spec: Final[dict[str, JsonValue]] = {"edits": [{"type": "clear_tool_uses_20250919"}]}
result: Final = await _counter(bedrock).count_tokens(
model_to_use="global.anthropic.claude-opus-4-8",
messages=[
{"role": "user", "content": "first"},
{"role": "assistant", "content": ""},
{"role": "user", "content": "third question here"},
],
contents=None,
deployment=DEPLOYMENT,
request_model="claude-opus-4-8",
context_management=CountContextManagement(spec=spec, drop_invalid_edits=False),
)
assert result is not None
assert result.total_tokens == 16
assert _invoke_body(bedrock.requests[0])["messages"] == [
{"role": "user", "content": [{"type": "text", "text": "first"}]},
{"role": "assistant", "content": ""},
{"role": "user", "content": [{"type": "text", "text": "third question here"}]},
]
@pytest.mark.asyncio
async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> None:
bedrock: Final = _Bedrock(runtime=(404, UNSUPPORTED), mantle=(200, {"input_tokens": 2177}))

65
tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
View file

@ -1922,3 +1922,68 @@ class TestBedrockMantleResponsesPricing:
def test_models_registered(self, local_cost_map):
assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models
assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models
def _mantle_marked_part(
litellm_params: GenericLiteLLMParams, marker: object, model: str = "openai.gpt-6.1-sol"
) -> object:
body: Final = BedrockMantleResponsesAPIConfig().transform_responses_api_request(
model=model,
input=[
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "Stable prefix", "prompt_cache_breakpoint": marker}],
}
],
response_api_optional_request_params={},
litellm_params=litellm_params,
headers={},
)
return body["input"][0]["content"][0]
@pytest.mark.parametrize(
"marker, wire_part",
[
(
{"mode": "explicit", "ttl": "30m"},
{"type": "input_text", "text": "Stable prefix", "prompt_cache_breakpoint": {"mode": "explicit"}},
),
({"mode": "implicit"}, {"type": "input_text", "text": "Stable prefix"}),
],
ids=("ttl-dropped", "implicit-dropped"),
)
def test_mantle_drop_params_cuts_a_responses_breakpoint_down_to_its_mode(
monkeypatch: pytest.MonkeyPatch, marker: dict[str, str], wire_part: dict[str, object]
) -> None:
"""Bedrock Mantle Responses answered 400 for a 30m breakpoint ttl and 200 for a bare explicit mode on 2026-10-29."""
monkeypatch.setattr(litellm, "drop_params", False)
assert _mantle_marked_part(GenericLiteLLMParams(drop_params=True), marker) == wire_part
assert _mantle_marked_part(GenericLiteLLMParams(), marker) == {
"type": "input_text",
"text": "Stable prefix",
"prompt_cache_breakpoint": marker,
}
@pytest.mark.parametrize(
"marker, wire_part",
[
({"mode": "explicit", "ttl": "1h"}, {"type": "input_text", "text": "Stable prefix"}),
({"mode": "explicit", "ttl": "30m"}, {"type": "input_text", "text": "Stable prefix"}),
(
{"mode": "explicit"},
{"type": "input_text", "text": "Stable prefix", "prompt_cache_breakpoint": {"mode": "explicit"}},
),
],
ids=("1h-dropped", "30m-dropped", "explicit-kept"),
)
def test_mantle_drop_params_sends_an_unlisted_model_only_a_bare_explicit_breakpoint(
monkeypatch: pytest.MonkeyPatch, marker: dict[str, str], wire_part: dict[str, object]
) -> None:
monkeypatch.setattr(litellm, "drop_params", False)
assert (
_mantle_marked_part(GenericLiteLLMParams(drop_params=True), marker, model="openai.not-in-the-cost-map")
== wire_part
)

2
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -1668,7 +1668,7 @@ _BOTH_FOUNDRY_BETAS: Final = ("context-1m-2025-08-07,context-management-2025-06-
pytest.param(
None,
{"anthropic-beta": "context-1m-2025-08-07"},
"claude-code-20250219",
"mcp-servers-2025-12-04",
("context-1m-2025-08-07",),
id="unsupported-client-beta-is-dropped-and-deployment-beta-kept",
),

78
tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py
View file

@ -172,10 +172,10 @@ def test_responses_call_forwards_previous_response_id_and_store() -> None:
def test_responses_call_folds_developer_items_into_instructions() -> None:
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b"))
with patch(HTTPX_CLIENT_FACTORY, return_value=client):
litellm.responses(
model="fireworks_ai/accounts/fireworks/models/kimi-k3",
model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b",
input=[
{"role": "user", "content": "Hi there"},
{"role": "developer", "content": "Answer with exactly one word."},
@ -214,6 +214,68 @@ def test_responses_call_keeps_developer_items_in_place_for_a_deployment_that_acc
)
@respx.mock
@pytest.mark.parametrize(
"model",
[
"fireworks_ai/accounts/fireworks/models/kimi-k3",
"fireworks_ai/accounts/fireworks/models/glm-5p3",
"fireworks_ai/accounts/fireworks/routers/kimi-k3-fast",
"fireworks_ai/glm-5p3-us",
"fireworks_ai/models/kimi-k3",
"fireworks_ai/routers/kimi-k3-fast",
],
)
def test_responses_call_keeps_developer_items_in_place_for_a_model_the_cost_map_declares_accepts_them(
model: str,
) -> None:
route: Final = _fireworks_responses_route(model.removeprefix("fireworks_ai/"))
litellm.responses(
model=model,
instructions="You are Codex",
input=[
{"role": "user", "content": "Hi there"},
{"role": "developer", "content": "Answer with exactly one word."},
{"role": "user", "content": "What is the capital of France?"},
],
api_key="fw-test-key",
)
body: Final = _routed_body(route)
assert body["instructions"] == "You are Codex"
assert tuple(body["input"]) == (
{"role": "user", "content": "Hi there"},
{"role": "system", "content": "Answer with exactly one word.", "type": "message"},
{"role": "user", "content": "What is the capital of France?"},
)
@respx.mock
def test_responses_call_folds_developer_items_when_the_loaded_cost_map_turns_the_flag_off(
declare_mid_conversation_system: Callable[[str, bool], None],
) -> None:
declare_mid_conversation_system("fireworks_ai/accounts/fireworks/models/kimi-k3", False)
route: Final = _fireworks_responses_route("accounts/fireworks/models/kimi-k3")
litellm.responses(
model="fireworks_ai/accounts/fireworks/models/kimi-k3",
input=[
{"role": "user", "content": "Hi there"},
{"role": "developer", "content": "Answer with exactly one word."},
{"role": "user", "content": "What is the capital of France?"},
{"type": "message", "role": "assistant", "content": "Paris"},
{"role": "developer", "content": "Reply with PONG."},
],
api_key="fw-test-key",
)
body: Final = _routed_body(route)
assert body["instructions"] == "Answer with exactly one word."
assert tuple(body["input"]) == (
{"role": "user", "content": "Hi there"},
{"role": "user", "content": "What is the capital of France?"},
{"type": "message", "role": "assistant", "content": "Paris"},
{"role": "user", "content": "Reply with PONG.", "type": "message"},
)
@respx.mock
@pytest.mark.parametrize(
("model_info", "closing_role"),
@ -222,9 +284,9 @@ def test_responses_call_keeps_developer_items_in_place_for_a_deployment_that_acc
def test_responses_call_sends_the_closing_developer_item_as_a_user_turn_only_for_a_declared_system_first_deployment(
model_info: Mapping[str, object] | None, closing_role: str
) -> None:
route: Final = _fireworks_responses_route("accounts/fireworks/models/kimi-k3")
route: Final = _fireworks_responses_route("accounts/fireworks/models/qwen3p8-2p4t-a95b")
litellm.responses(
model="fireworks_ai/accounts/fireworks/models/kimi-k3",
model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b",
input=[
{"role": "user", "content": "Hi there"},
{"role": "developer", "content": "Answer with exactly one word."},
@ -402,7 +464,7 @@ def test_transform_request_forwards_non_string_instructions_and_input_untouched(
developer_item: Final = {"role": "developer", "content": "Answer with exactly one word."}
user_item: Final = {"role": "user", "content": "What is the capital of France?"}
request: Final = FireworksAIResponsesAPIConfig().transform_responses_api_request(
model="accounts/fireworks/models/kimi-k3",
model="accounts/fireworks/models/qwen3p8-2p4t-a95b",
input=cast(ResponseInputParam, [developer_item, user_item]),
response_api_optional_request_params={"instructions": ["not", "a", "string"]},
litellm_params=GenericLiteLLMParams(),
@ -416,7 +478,7 @@ def test_transform_request_forwards_non_string_instructions_and_input_untouched(
def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_output_items() -> None:
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b"))
pydantic_input: Final = cast(
ResponseInputParam,
[
@ -435,7 +497,9 @@ def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_outpu
)
with patch(HTTPX_CLIENT_FACTORY, return_value=client):
litellm.responses(
model="fireworks_ai/accounts/fireworks/models/kimi-k3", input=pydantic_input, api_key="fw-test-key"
model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b",
input=pydantic_input,
api_key="fw-test-key",
)
_, _, body = _sent_request(client)
assert body["instructions"] == "Answer with exactly one word."

119
tests/unit/llms/openai/chat/test_openai_gpt_transformation.py
View file

@ -1659,6 +1659,125 @@ class TestOpenAIPromptCacheBreakpointChatPath:
]
_CLIENT_BREAKPOINTS_AND_THEIR_CHAT_WIRE_FORM: Final = (
({"mode": "explicit", "ttl": "1h"}, {"mode": "explicit"}),
({"mode": "explicit", "ttl": "30m"}, {"mode": "explicit"}),
({"mode": "explicit"}, {"mode": "explicit"}),
({"mode": "implicit"}, None),
({}, None),
("explicit", None),
)
_CLIENT_BREAKPOINT_IDS: Final = ("ttl-1h", "ttl-30m", "explicit", "implicit", "empty", "string")
class TestChatCompletionsClientBreakpointUnderDropParams:
"""
Live OpenAI Chat Completions (gpt-6.1-sol, gpt-6.1-terra, gpt-6-sol, gpt-5.6, 2026-10-28) answers 400 on a
content part's prompt_cache_breakpoint carrying any ttl or a mode other than "explicit".
"""
@staticmethod
def _request(marker: object, drop_params: bool) -> dict[str, object]:
return OpenAIGPTConfig().transform_request(
model="gpt-5.6",
messages=[
{"role": "system", "content": "Answer briefly"},
{
"role": "user",
"content": [
{"type": "text", "text": "Describe this", "prompt_cache_breakpoint": marker},
{"type": "image_url", "image_url": {"url": "https://example.com/a.png"}},
{"type": "text", "text": "and nothing else"},
],
},
],
optional_params={},
litellm_params={"custom_llm_provider": "openai", "drop_params": drop_params},
headers={},
)
@pytest.mark.parametrize(
("marker", "wire_marker"), _CLIENT_BREAKPOINTS_AND_THEIR_CHAT_WIRE_FORM, ids=_CLIENT_BREAKPOINT_IDS
)
def test_drop_params_cuts_the_breakpoint_down_to_what_chat_completions_accepts(
self, marker: object, wire_marker: dict[str, str] | None
):
text_part: Final = {"type": "text", "text": "Describe this"}
assert self._request(marker, drop_params=True)["messages"] == [
{"role": "system", "content": "Answer briefly"},
{
"role": "user",
"content": [
text_part if wire_marker is None else {**text_part, "prompt_cache_breakpoint": wire_marker},
{"type": "image_url", "image_url": {"url": "https://example.com/a.png"}},
{"type": "text", "text": "and nothing else"},
],
},
]
@pytest.mark.parametrize(
"marker", [marker for marker, _ in _CLIENT_BREAKPOINTS_AND_THEIR_CHAT_WIRE_FORM], ids=_CLIENT_BREAKPOINT_IDS
)
def test_without_drop_params_the_client_breakpoint_reaches_openai_unchanged(self, marker: object):
assert self._request(marker, drop_params=False)["messages"][1]["content"][0] == {
"type": "text",
"text": "Describe this",
"prompt_cache_breakpoint": marker,
}
def test_cache_control_still_marks_a_part_whose_client_breakpoint_was_dropped(self):
request: Final = OpenAIGPTConfig().transform_request(
model="gpt-5.6",
messages=[
{
"role": "user",
"content": [
{
"type": "text",
"text": "Describe this",
"cache_control": {"type": "ephemeral"},
"prompt_cache_breakpoint": {"mode": "implicit"},
}
],
}
],
optional_params={},
litellm_params={"custom_llm_provider": "openai", "drop_params": True},
headers={},
)
assert request["messages"] == [
{
"role": "user",
"content": [{"type": "text", "text": "Describe this", "prompt_cache_breakpoint": {"mode": "explicit"}}],
}
]
@pytest.mark.parametrize("stray_part", ["Answer briefly", None], ids=["string", "null"])
def test_drop_params_leaves_a_part_that_is_not_an_object_for_openai_to_reject(self, stray_part: object):
assert OpenAIGPTConfig().transform_request(
model="gpt-5.6",
messages=[
{
"role": "system",
"content": [
stray_part,
{"type": "text", "text": "and nothing else", "prompt_cache_breakpoint": {"mode": "implicit"}},
],
},
{"role": "user", "content": "Describe this"},
],
optional_params={},
litellm_params={"custom_llm_provider": "openai", "drop_params": True},
headers={},
)["messages"] == [
{"role": "system", "content": [stray_part, {"type": "text", "text": "and nothing else"}]},
{"role": "user", "content": "Describe this"},
]
_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$'

28
tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py
View file

@ -4547,6 +4547,34 @@ class TestBuildBlockSseChunks:
assert completed["output"][0]["content"][0]["text"] == "Blocked by policy."
assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28}
def test_continuation_numbers_every_event_after_the_last_one_relayed(self):
yielded: Final = [
{"type": "response.created", "sequence_number": 0, "response": {"id": "resp_live", "output": []}},
{
"type": "response.output_item.added",
"sequence_number": 1,
"output_index": 0,
"item": {"id": "msg_orig", "type": "message", "role": "assistant"},
},
{
"type": "response.output_text.delta",
"sequence_number": 2,
"item_id": "msg_orig",
"output_index": 0,
"content_index": 0,
"delta": "partial",
},
]
payloads: Final = self._payloads(
OpenAIResponsesHandler().build_block_sse_chunks(
self._exc(original_response=yielded), stream_started=True, responses_so_far=yielded
)
)
assert payloads[-1]["type"] == "response.completed"
assert [payload.get("sequence_number") for payload in payloads] == list(range(3, 3 + len(payloads)))
def test_continuation_reads_usage_from_typed_completed_event(self):
from litellm.types.llms.openai import (
ResponseCompletedEvent,

112
tests/unit/llms/openai/responses/test_openai_responses_transformation.py
View file

@ -4393,8 +4393,9 @@ async def test_extra_body_merges_with_request_data(extra_body_mock_response_data
assert request_body["custom_field"] == "custom_value"
_MALFORMED_MARKERS: Final = ("yes", 1, {"mode": "bogus"}, {"mode": "explicit", "ttl": "1h"})
_MALFORMED_MARKER_IDS: Final = ("string", "int", "bogus-mode", "bad-ttl")
_MALFORMED_MARKERS: Final = ("yes", 1, {"mode": "bogus"}, {"mode": "bogus", "ttl": "30m"})
_MALFORMED_MARKER_IDS: Final = ("string", "int", "bogus-mode", "bogus-mode-with-ttl")
_UNSUPPORTED_TTL_MARKER: Final = {"mode": "explicit", "ttl": "1h"}
_VALID_MARKERS: Final = ({"mode": "explicit"}, {"mode": "explicit", "ttl": "30m"})
_VALID_MARKER_IDS: Final = ("explicit", "explicit-30m")
_PARTS: Final[Mapping[str, Mapping[str, object]]] = {
@ -4480,6 +4481,76 @@ def test_valid_prompt_cache_breakpoint_reaches_the_wire_under_both_settings(
assert _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params=drop_params)) == input
@pytest.mark.parametrize("config", _CONFIGS, ids=_CONFIG_IDS)
@pytest.mark.parametrize("part_type", tuple(_PARTS))
def test_an_unsupported_ttl_is_dropped_and_the_explicit_mode_kept_under_drop_params(
monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig, part_type: str
) -> None:
monkeypatch.setattr(litellm, "drop_params", False)
input: Final = _marked_message(part_type, _UNSUPPORTED_TTL_MARKER)
trimmed: Final = _transformed_input(config, input, GenericLiteLLMParams(drop_params=True))
assert _single_part(trimmed) == {**_PARTS[part_type], "prompt_cache_breakpoint": {"mode": "explicit"}}
assert _transformed_input(config, input, GenericLiteLLMParams(drop_params=False)) == input
@pytest.mark.parametrize(
"config, marker, wire_marker",
[
(OpenAIResponsesAPIConfig(), {"mode": "explicit", "ttl": "12h"}, {"mode": "explicit", "ttl": "12h"}),
(AzureOpenAIResponsesAPIConfig(), {"mode": "explicit", "ttl": "30m"}, {"mode": "explicit", "ttl": "30m"}),
(AzureOpenAIResponsesAPIConfig(), {"mode": "explicit", "ttl": "12h"}, {"mode": "explicit"}),
(AzureAIResponsesAPIConfig(), {"mode": "explicit", "ttl": "12h"}, {"mode": "explicit"}),
],
ids=("openai-keeps-12h", "azure-keeps-30m", "azure-drops-12h", "azure-ai-drops-12h"),
)
def test_drop_params_keeps_only_a_ttl_the_target_responses_api_accepts(
monkeypatch: pytest.MonkeyPatch,
config: BaseResponsesAPIConfig,
marker: Mapping[str, str],
wire_marker: Mapping[str, str],
) -> None:
"""Azure OpenAI Responses answered 400 for a 12h breakpoint ttl and 200 for 30m on 2026-10-29, and Azure AI
Foundry Responses did the same on 2026-10-30."""
monkeypatch.setattr(litellm, "drop_params", False)
input: Final = _marked_message("input_text", marker)
trimmed: Final = _transformed_input(config, input, GenericLiteLLMParams(drop_params=True))
assert _single_part(trimmed) == {**_PARTS["input_text"], "prompt_cache_breakpoint": wire_marker}
assert _transformed_input(config, input, GenericLiteLLMParams(drop_params=False)) == input
@pytest.mark.parametrize(
"config, litellm_params",
[
(XAIResponsesAPIConfig(), GenericLiteLLMParams(drop_params=True)),
(HostedVLLMResponsesAPIConfig(), GenericLiteLLMParams(drop_params=True, api_base="http://localhost:8000/v1")),
(OpenAIResponsesAPIConfig(), GenericLiteLLMParams(drop_params=True, api_base="http://localhost:8000/v1")),
],
ids=("xai", "hosted-vllm", "openai-compatible-host"),
)
def test_a_responses_api_not_known_to_take_breakpoints_gets_a_marker_whole_or_not_at_all(
monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig, litellm_params: GenericLiteLLMParams
) -> None:
monkeypatch.setattr(litellm, "drop_params", False)
unsupported_ttl: Final = _marked_message("input_text", _UNSUPPORTED_TTL_MARKER)
assert _single_part(_transformed_input(config, unsupported_ttl, litellm_params)) == _PARTS["input_text"]
openai_ttl: Final = _marked_message("input_text", {"mode": "explicit", "ttl": "12h"})
assert _transformed_input(config, openai_ttl, litellm_params) == openai_ttl
def test_a_custom_host_opted_in_with_prompt_cache_options_keeps_the_mode_of_a_marker_with_an_unsupported_ttl(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "drop_params", False)
request: Final = OpenAIResponsesAPIConfig().transform_responses_api_request(
model="gpt-6.1-sol",
input=cast("ResponseInputParam", _marked_message("input_text", _UNSUPPORTED_TTL_MARKER)),
response_api_optional_request_params={"prompt_cache_options": {"mode": "explicit"}},
litellm_params=GenericLiteLLMParams(drop_params=True, api_base="https://gateway.example/v1"),
headers={},
)
assert _single_part(request["input"]) == {**_PARTS["input_text"], "prompt_cache_breakpoint": {"mode": "explicit"}}
def test_prompt_cache_breakpoint_unknown_key_is_trimmed_under_drop_params(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "drop_params", False)
input: Final = _marked_message("input_text", {"mode": "explicit", "note": "kept"})
@ -4569,12 +4640,43 @@ def test_a_valid_marker_survives_drop_params_on_a_deployment_name_with_unknown_b
monkeypatch.setattr(litellm, "drop_params", False)
litellm_params: Final = GenericLiteLLMParams(drop_params=True)
valid: Final = _marked_message("input_text", {"mode": "explicit"})
malformed: Final = _marked_message("input_text", {"mode": "explicit", "ttl": "1h"})
malformed: Final = _marked_message("input_text", {"mode": "bogus"})
assert _transformed_input_for("litellm-qa-sol-alias", valid, litellm_params, config) == valid
dropped: Final = _transformed_input_for("litellm-qa-sol-alias", malformed, litellm_params, config)
assert _single_part(dropped) == _PARTS["input_text"]
@pytest.mark.parametrize("ttl", ("30m", "12h"))
@pytest.mark.parametrize(
"config",
(OpenAIResponsesAPIConfig(), AzureOpenAIResponsesAPIConfig(), AzureAIResponsesAPIConfig()),
ids=("openai", "azure", "azure_ai"),
)
def test_a_deployment_with_unknown_breakpoint_support_gets_a_marker_whole_or_not_at_all(
monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig, ttl: str
) -> None:
monkeypatch.setattr(litellm, "drop_params", False)
litellm_params: Final = GenericLiteLLMParams(drop_params=True)
supported_ttl: Final = _marked_message("input_text", {"mode": "explicit", "ttl": ttl})
unsupported_ttl: Final = _marked_message("input_text", _UNSUPPORTED_TTL_MARKER)
assert _transformed_input_for("litellm-qa-sol-alias", supported_ttl, litellm_params, config) == supported_ttl
dropped: Final = _transformed_input_for("litellm-qa-sol-alias", unsupported_ttl, litellm_params, config)
assert _single_part(dropped) == _PARTS["input_text"]
@pytest.mark.parametrize(
"config", (AzureOpenAIResponsesAPIConfig(), AzureAIResponsesAPIConfig()), ids=("azure", "azure_ai")
)
def test_an_azure_deployment_declaring_a_breakpoint_reader_keeps_the_mode_of_a_marker_with_an_unsupported_ttl(
monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig
) -> None:
monkeypatch.setattr(litellm, "drop_params", False)
litellm_params: Final = GenericLiteLLMParams(drop_params=True, base_model="gpt-6.1-sol")
unsupported_ttl: Final = _marked_message("input_text", _UNSUPPORTED_TTL_MARKER)
trimmed: Final = _transformed_input_for("litellm-qa-sol-alias", unsupported_ttl, litellm_params, config)
assert _single_part(trimmed) == {**_PARTS["input_text"], "prompt_cache_breakpoint": {"mode": "explicit"}}
def _declaring_base_model(where: str, base_model: str) -> GenericLiteLLMParams:
if where == "litellm_params":
return GenericLiteLLMParams(drop_params=True, base_model=base_model)
@ -4615,10 +4717,10 @@ def test_a_model_that_takes_markers_sends_kept_parts_byte_for_byte_and_drops_onl
{"prompt_cache_breakpoint": {"mode": "explicit"}, "type": "input_text", "text": "b"},
]
rejected: Final = {"type": "input_text", "text": "c"}
bad_ttl: Final = {"prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}}
bogus_mode: Final = {"prompt_cache_breakpoint": {"mode": "bogus"}}
input: Final = [
{"type": "message", "role": "user", "content": kept},
{"type": "message", "role": "user", "content": [{**bad_ttl, **rejected}]},
{"type": "message", "role": "user", "content": [{**bogus_mode, **rejected}]},
]
transformed: Final = _transformed_input_for("gpt-6.1-sol", input, GenericLiteLLMParams(drop_params=True))
assert json.dumps(transformed) == json.dumps(

230
tests/unit/llms/vertex_ai/count_tokens/test_context.py
View file

@ -1,6 +1,5 @@
import asyncio
import base64
import logging
import uuid
from collections.abc import Generator
from contextlib import contextmanager
@ -11,20 +10,22 @@ from typing import Final, Literal
from unittest.mock import patch
import pytest
from pydantic import JsonValue, TypeAdapter
from pydantic import JsonValue
import litellm
from litellm.llms.anthropic.messages.adapters.transformation import (
LiteLLMAnthropicMessagesAdapter,
truncate_tool_name,
)
from litellm.llms.vertex_ai.count_tokens.context import (
from litellm.llms.anthropic.messages.context_management.inference_view import (
CountContextEdits,
InferenceRequestView,
inference_request_view,
)
from litellm.llms.vertex_ai.count_tokens.context import (
VertexGeminiCountContext,
VertexGeminiCountContextError,
gemini_count_contents,
inference_request_view,
unconvertible_messages_error,
vertex_gemini_count_context,
)
@ -424,73 +425,6 @@ async def test_image_that_cannot_be_fetched_returns_an_error_value(
assert isinstance(contents, VertexGeminiCountContextError)
@pytest.mark.parametrize(
"system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system")
)
async def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None:
view: Final = await inference_request_view(
(
{"role": "user", "content": "first question " * 50},
{"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]},
{"role": "user", "content": "latest question"},
),
system,
"anthropic",
_NO_EDITS,
)
assert isinstance(view, InferenceRequestView)
assert view.messages == [{"role": "user", "content": "latest question"}]
assert "Summary of the earlier turns." in str(view.system)
assert "Be brief." in str(view.system)
async def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None:
snippet: Final = "LiteLLM ships a weekly release candidate and a stable line."
view: Final = await inference_request_view(
(
{"role": "user", "content": "What is new?"},
{
"role": "assistant",
"content": [
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}},
{
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_01",
"content": [
{
"type": "web_search_result",
"url": "https://docs.litellm.ai/release_notes",
"title": "Release notes",
"encrypted_content": "",
"snippet": snippet,
}
],
},
],
},
),
None,
"anthropic",
_NO_EDITS,
)
assert isinstance(view, InferenceRequestView)
assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"])
assert [block["type"] for block in assistant_blocks] == ["text"]
assert snippet in assistant_blocks[0]["text"]
async def test_openai_view_is_the_request_unchanged() -> None:
messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},)
view: Final = await inference_request_view(messages, "Be brief.", "openai", _NO_EDITS)
assert isinstance(view, InferenceRequestView)
assert view.messages == messages
assert view.system == "Be brief."
async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None:
with patch(
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
@ -556,138 +490,30 @@ async def test_anthropic_replayed_web_search_result_is_counted_as_its_text(
assert snippet in str(contents)
_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180
_READ_FILE_TOOL: Final[dict[str, object]] = {
"name": "read_file",
"description": "Read a file",
"input_schema": {"type": "object", "properties": {"path": {"type": "string"}}},
}
_CLEAR_EDIT: Final[JsonValue] = {
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 500},
"keep": {"type": "tool_uses", "value": 1},
}
_CLEAR_OLD_TOOL_RESULTS: Final[JsonValue] = {"edits": [_CLEAR_EDIT]}
def _agent_turns(rounds: int, first_round: int = 0) -> tuple[dict[str, JsonValue], ...]:
return (
*(
turn
for index in range(first_round, first_round + rounds)
for turn in (
{"role": "user", "content": f"Read file {index}."},
{
"role": "assistant",
"content": [
{"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}}
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _TOOL_OUTPUT}],
},
)
), # comprehension-ok: flattens each round into its three turns
{"role": "user", "content": "Summarize what you read."},
)
def _tool_outputs_kept(view: InferenceRequestView) -> tuple[bool, ...]:
return tuple(
_TOOL_OUTPUT in str(message["content"]) for message in view.messages if "tool_result" in str(message["content"])
)
def _edits(spec: JsonValue) -> CountContextEdits:
return CountContextEdits(spec=spec, model="gemini/gemini-2.5-flash", tools=(_READ_FILE_TOOL,))
async def test_anthropic_view_clears_the_tool_results_inference_clears() -> None:
view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (False, False, True)
async def test_anthropic_view_applies_clear_edits_without_calling_the_compaction_summary_model() -> None:
spec: Final[JsonValue] = {
"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50000}}, _CLEAR_EDIT]
}
turns: Final = (
{"role": "user", "content": "first question " * 30000},
{"role": "assistant", "content": "first answer"},
*_agent_turns(3),
)
with (
patch(
"litellm.llms.anthropic.messages.context_management.editors.compact._read_summary_model_setting",
return_value="gemini/gemini-2.5-flash",
),
patch.object(litellm, "acompletion", side_effect=AssertionError("count_tokens must not summarize")) as summary,
):
view: Final = await inference_request_view(turns, None, "anthropic", _edits(spec))
assert summary.call_count == 0
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (False, False, True)
async def test_anthropic_view_counts_the_unedited_body_without_an_error_log_when_an_edit_fails(
caplog: pytest.LogCaptureFixture,
) -> None:
with (
patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")),
caplog.at_level(logging.DEBUG, logger="LiteLLM"),
):
view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True, True)
assert any("skipping edits" in record.getMessage() for record in caplog.records)
assert not any(record.levelno >= logging.ERROR for record in caplog.records)
async def test_anthropic_view_clears_only_the_turns_after_a_compaction_block_when_clear_runs_before_compact() -> None:
clear_before_compact: Final[JsonValue] = {
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "tool_uses", "value": 3},
"keep": {"type": "tool_uses", "value": 1},
},
{"type": "compact_20260112"},
]
}
turns: Final = (
*_agent_turns(3)[:-1],
{"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]},
*_agent_turns(2, first_round=3),
)
with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")):
view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True)
assert "Read files 0 to 2." in str(view.system)
async def test_anthropic_view_rejects_a_compact_trigger_inference_rejects() -> None:
low_trigger: Final[JsonValue] = {
"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}]
}
view: Final = await inference_request_view(_agent_turns(1), None, "anthropic", _edits(low_trigger))
assert isinstance(view, VertexGeminiCountContextError)
assert view.status_code == 400
async def test_anthropic_view_names_the_block_gemini_cannot_count_when_an_edit_trigger_cannot_tokenize_it() -> None:
clear_old_tool_results: Final = CountContextEdits(
spec={
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 500},
"keep": {"type": "tool_uses", "value": 1},
}
]
},
model="gemini/gemini-2.5-flash",
tools=({"name": "read_file", "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}},),
)
turns: Final = (
*_agent_turns(1)[:-1],
{"role": "user", "content": "Read file 0."},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "toolu_0", "name": "read_file", "input": {"path": "0"}}],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "toolu_0", "content": "line of a long file. " * 180}],
},
{
"role": "user",
"content": [
@ -701,7 +527,7 @@ async def test_anthropic_view_names_the_block_gemini_cannot_count_when_an_edit_t
},
)
view: Final = await inference_request_view(turns, None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS))
view: Final = await inference_request_view(turns, None, "anthropic", clear_old_tool_results)
assert isinstance(view, InferenceRequestView)
blocks_error: Final = unconvertible_messages_error(view.messages, "anthropic")

99
tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
View file

@ -1,16 +1,19 @@
import copy
import threading
import time
from collections.abc import Iterator
from contextlib import contextmanager
from enum import Enum
from typing import Final
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm import anthropic_beta_headers_manager
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler
@ -1826,6 +1829,100 @@ async def test_vertex_ai_token_counter_forwards_system_and_tools_to_partner_requ
assert without_optional_fields.total_tokens == 37
@pytest.mark.asyncio
@pytest.mark.parametrize(
("edit_type", "expected_beta"),
[("clear_tool_uses_20250919", "context-management-2025-06-27"), ("compact_20260112", "compact-2026-01-12")],
)
async def test_vertex_claude_count_forwards_the_context_edits_and_beta_vertex_inference_sends(
edit_type: str, expected_beta: str, monkeypatch: pytest.MonkeyPatch
) -> None:
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler
class FakeResponse:
status_code = 200
def json(self) -> dict[str, int]:
return {"input_tokens": 21}
class FakeHttpClient:
posted: tuple[tuple[dict[str, str], dict[str, object]], ...] = ()
async def post(
self, url: str, headers: dict[str, str], json: dict[str, object], timeout: float
) -> FakeResponse:
self.posted = (*self.posted, (headers, json))
return FakeResponse()
fake_http_client: Final = FakeHttpClient()
spec: Final[dict[str, JsonValue]] = {"edits": [{"type": edit_type}]}
application_default_credentials: Final = MagicMock(token="fake-token", valid=True, expiry=None)
monkeypatch.setattr("google.auth.default", lambda *args, **kwargs: (application_default_credentials, "test-project"))
with patch.object(handler, "get_async_httpx_client", return_value=fake_http_client):
result: Final = await VertexAITokenCounter().count_tokens(
model_to_use="claude-opus-4-8",
messages=[{"role": "user", "content": "Hello"}],
contents=None,
deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}},
messages_format="anthropic",
context_management=CountContextManagement(spec=spec, drop_invalid_edits=False),
)
assert result is not None
assert result.total_tokens == 21
((headers, body),) = fake_http_client.posted
assert body["context_management"] == spec
assert headers["anthropic-beta"] == expected_beta
@pytest.mark.asyncio
async def test_vertex_claude_count_with_an_edit_fetches_the_remote_beta_config_off_the_event_loop(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler
class FakeResponse:
status_code = 200
def json(self) -> dict[str, int]:
return {"input_tokens": 21}
class FakeHttpClient:
async def post(
self, url: str, headers: dict[str, str], json: dict[str, object], timeout: float
) -> FakeResponse:
return FakeResponse()
class FetchThreads:
seen: tuple[int, ...] = ()
def __call__(self, request: httpx.Request) -> httpx.Response:
self.seen = (*self.seen, threading.get_ident())
return httpx.Response(200, json={})
monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False)
monkeypatch.setattr(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None)
fetch_threads: Final = FetchThreads()
application_default_credentials: Final = MagicMock(token="fake-token", valid=True, expiry=None)
monkeypatch.setattr("google.auth.default", lambda *args, **kwargs: (application_default_credentials, "test-project"))
with patch.object(handler, "get_async_httpx_client", return_value=FakeHttpClient()), respx.mock:
respx.get(litellm.anthropic_beta_headers_url).mock(side_effect=fetch_threads)
result: Final = await VertexAITokenCounter().count_tokens(
model_to_use="claude-opus-4-8",
messages=[{"role": "user", "content": "Hello"}],
contents=None,
deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}},
messages_format="anthropic",
context_management=CountContextManagement(
spec={"edits": [{"type": "clear_tool_uses_20250919"}]}, drop_invalid_edits=False
),
)
assert result is not None
assert len(fetch_threads.seen) == 1
assert threading.get_ident() not in fetch_threads.seen
_PARTNER_CHAT_TOOL_TURNS: Final[list[dict[str, object]]] = [
{"role": "user", "content": "Weather in SF?"},
{

38
tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py
View file

@ -640,36 +640,16 @@ async def test_vertex_ai_anthropic_prompt_caching_does_not_add_beta_header(
assert "anthropic_beta" not in json.loads(request.content)
def test_vertex_ai_partner_models_anthropic_remove_prompt_caching_scope_beta_header():
"""
Test that remove_unsupported_beta correctly filters out prompt-caching-scope-2026-01-05
from the anthropic-beta headers.
"""
def test_vertex_ai_partner_models_anthropic_forwards_prompt_caching_scope_beta_header(
local_beta_headers_config: None,
) -> None:
"""A live Vertex rawPredict probe on 2026-10-28 accepted prompt-caching-scope-2026-01-05, so it goes upstream
while a beta missing from the vertex_ai column is still dropped."""
headers: Final = update_headers_with_filtered_beta(
{"anthropic-beta": "other-feature,prompt-caching-scope-2026-01-05,web-search-2025-03-05"}, "vertex_ai"
)
# This beta header should be removed
PROMPT_CACHING_BETA_HEADER = "prompt-caching-scope-2026-01-05"
headers = {
"anthropic-beta": f"other-feature,{PROMPT_CACHING_BETA_HEADER},web-search-2025-03-05"
}
headers = update_headers_with_filtered_beta(headers, "vertex_ai")
beta_header = headers.get("anthropic-beta")
assert PROMPT_CACHING_BETA_HEADER not in (
beta_header or ""
), f"{PROMPT_CACHING_BETA_HEADER} should be filtered out"
assert "other-feature" not in (
beta_header or ""
), "Other non-excluded beta headers should remain"
assert "web-search-2025-03-05" in (
beta_header or ""
), "Other non-excluded beta headers should remain"
# If prompt-caching was the only value, header should be removed completely
headers2 = {"anthropic-beta": PROMPT_CACHING_BETA_HEADER}
headers2 = update_headers_with_filtered_beta(headers2, "vertex_ai")
assert (
"anthropic-beta" not in headers2
), "Header should be removed if no supported values remain"
assert headers == {"anthropic-beta": "prompt-caching-scope-2026-01-05,web-search-2025-03-05"}
def test_vertex_ai_anthropic_output_config_effort_only_forwarded():

133
tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
View file

@ -10998,139 +10998,6 @@ async def test_key_budget_error_keeps_the_masked_key_name(key_name):
assert f"Key=prod-key ({key_name}) Current cost" in message
class _UntouchedPrisma:
def __getattr__(self, name: str) -> object:
raise AssertionError(f"database reached through {name}")
class _MissingUserPrisma:
class db:
class litellm_usertable:
@staticmethod
async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None:
return None
@pytest.mark.asyncio
async def test_enforced_model_allowlists_treats_a_missing_user_row_as_unrestricted():
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import ProxyLogging
cache = UserApiKeyCache()
scopes = await enforced_model_allowlists(
valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="default_user_id"),
prisma_client=_MissingUserPrisma(),
user_api_key_cache=cache,
proxy_logging_obj=ProxyLogging(user_api_key_cache=cache),
)
assert [list(scope) for scope in scopes] == [[], [], [], [], []]
class _UnreachableUserPrisma:
class db:
class litellm_usertable:
@staticmethod
async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None:
raise RuntimeError("database gone")
@pytest.mark.asyncio
async def test_enforced_model_allowlists_surfaces_a_failed_user_lookup():
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import ProxyLogging
cache = UserApiKeyCache()
with pytest.raises(ValueError, match="database gone"):
await enforced_model_allowlists(
valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"),
prisma_client=_UnreachableUserPrisma(),
user_api_key_cache=cache,
proxy_logging_obj=ProxyLogging(user_api_key_cache=cache),
)
@pytest.mark.asyncio
async def test_enforced_model_allowlists_reads_every_level_from_cache():
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_ProjectTableCachedObj,
LiteLLM_TeamMembership,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
)
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
team_membership_reservation_cache_key,
)
from litellm.proxy.utils import ProxyLogging
cache = UserApiKeyCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=cache)
await cache.async_set_cache(
key="team_id:team-fake", value=LiteLLM_TeamTableCachedObj(team_id="team-fake", models=["gpt-4o"])
)
await cache.async_set_cache(
key=team_membership_reservation_cache_key(user_id="user-fake", team_id="team-fake"),
value=LiteLLM_TeamMembership(
user_id="user-fake",
team_id="team-fake",
litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["gpt-4o-mini"]),
),
)
await cache.async_set_cache(
key="project_id:project-fake",
value=LiteLLM_ProjectTableCachedObj(project_id="project-fake", models=["gpt-4.1"]),
)
await cache.async_set_cache(key="user-fake", value=LiteLLM_UserTable(user_id="user-fake", models=["o3"]))
prisma_client = _UntouchedPrisma()
team_scoped = await enforced_model_allowlists(
valid_token=UserAPIKeyAuth(
token="hashed-fake",
models=["all-team-models"],
team_models=["gpt-4o", "gpt-4o-mini"],
user_id="user-fake",
team_id="team-fake",
project_id="project-fake",
),
prisma_client=prisma_client,
user_api_key_cache=cache,
proxy_logging_obj=proxy_logging_obj,
)
personal = await enforced_model_allowlists(
valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"),
prisma_client=prisma_client,
user_api_key_cache=cache,
proxy_logging_obj=proxy_logging_obj,
)
without_database = await enforced_model_allowlists(
valid_token=UserAPIKeyAuth(
token="hashed-fake",
models=["gpt-4o"],
team_models=["gpt-4o-mini"],
user_id="user-fake",
team_id="team-fake",
),
prisma_client=None,
user_api_key_cache=cache,
proxy_logging_obj=proxy_logging_obj,
)
assert [list(scope) for scope in team_scoped] == [
["gpt-4o", "gpt-4o-mini"],
["gpt-4o"],
["gpt-4o-mini"],
[],
["gpt-4.1"],
]
assert [list(scope) for scope in personal] == [[], [], [], ["o3"], []]
assert [list(scope) for scope in without_database] == [["gpt-4o"], ["gpt-4o-mini"]]
@pytest.mark.asyncio
@pytest.mark.parametrize("channel", ["team", "key"])
async def test_access_group_model_fallback_uses_the_injected_database(channel: str) -> None:

9
tests/unit/proxy/common_utils/test_realtime_cache.py
View file

@ -1,4 +1,5 @@
from typing import Any, cast
import json
from typing import Any, Final, cast
import pytest
@ -40,6 +41,12 @@ def test_realtime_request_body_caches_each_model_separately():
assert gpt4o_body_first is not gpt4o_mini_body
def test_realtime_request_body_keeps_a_model_name_with_quotes_inside_its_own_field():
body: Final = realtime_request_body('gpt-4o", "user": "someone-else')
assert json.loads(body) == {"model": 'gpt-4o", "user": "someone-else'}
def test_realtime_query_params_template_caches_each_pair_separately():
params_with_intent_first = _realtime_query_params_template("gpt-4o", "intent-a")
params_with_intent_second = _realtime_query_params_template("gpt-4o", "intent-a")

42
tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py
View file

@ -1694,6 +1694,48 @@ class TestResponsesStreamRestoration:
assert "".join(e.delta for e in out) == "see a@example.com a["
@pytest.mark.asyncio
async def test_a_flushed_delta_takes_the_done_events_number_and_later_events_move_up(self):
guardrail, _ = _shielded(self.VAULT)
done = OutputTextDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
item_id="msg_1",
output_index=0,
content_index=0,
text="Mail [EMAIL_1] x[",
sequence_number=3,
)
out = await _restore_stream(
guardrail,
[
self._text_delta("Mail [EMA", 1),
self._text_delta("IL_1] x[", 2),
done,
self._text_delta("next", 4, content_index=1),
],
)
assert [(e.type, e.sequence_number) for e in out] == [
(ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, 1),
(ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, 2),
(ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, 3),
(ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, 4),
(ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, 5),
]
assert out[2].delta == "["
@pytest.mark.asyncio
async def test_a_truncated_stream_flush_continues_after_the_last_number(self):
guardrail, _ = _shielded(self.VAULT)
out = await _restore_stream(
guardrail, [self._text_delta("see [EMAIL_1] a[", 7), self._text_delta(" b", 8, content_index=1)]
)
assert [e.sequence_number for e in out] == [7, 8, 9]
assert out[-1].delta == "["
@pytest.mark.asyncio
async def test_completed_response_is_restored(self):
"""The terminal event repeats the whole reply, and clients read it as the answer."""

520
tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py Normal file
View file

@ -0,0 +1,520 @@
import uuid
from collections.abc import Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
import litellm
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
OpenAIPassthroughLoggingHandler,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import (
OpenAIWebsocketPassthroughLoggingHandler,
openai_websocket_frame_kind,
)
from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id
from litellm.types.utils import CallTypes
REALTIME_MODEL: Final = "gpt-realtime-2.1"
RESPONSES_MODEL: Final = "gpt-5.4-mini"
REASONING_MODEL: Final = "gpt-5.4"
TRANSCRIPTION_MODEL: Final = "gpt-4o-mini-transcribe"
SESSION_START: Final = datetime(2026, 1, 1, 12, 0, 0)
def _logging_obj() -> MagicMock:
logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id=str(uuid.uuid4()))
return logging_obj
def _session_created(model: str) -> dict[str, object]:
return {"type": "session.created", "session": {"id": "sess_fake", "model": model}}
def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_out: int) -> dict[str, object]:
return {
"type": "response.done",
"response": {
"id": "resp_fake",
"status": "completed",
"usage": {
"input_tokens": text_in + audio_in,
"output_tokens": text_out + audio_out,
"total_tokens": text_in + audio_in + text_out + audio_out,
"input_token_details": {"text_tokens": text_in, "audio_tokens": audio_in, "cached_tokens": 0},
"output_token_details": {"text_tokens": text_out, "audio_tokens": audio_out},
},
},
}
def _responses_completed(
model: str,
input_tokens: int,
output_tokens: int,
event: str = "response.completed",
served: Mapping[str, object] = MappingProxyType({}),
) -> dict[str, object]:
return {
"type": event,
"response": {
"id": "resp_fake",
"object": "response",
"created_at": 1761700000,
"model": model,
"status": "completed",
"output": [],
**served,
"usage": {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": input_tokens + output_tokens,
"input_tokens_details": {"cached_tokens": 0},
"output_tokens_details": {"reasoning_tokens": 0},
},
},
}
def _handle(
messages: Sequence[Mapping[str, object]],
kwargs: Mapping[str, object] = MappingProxyType({}),
logging_obj: MagicMock | None = None,
url_route: str = "/openai/v1/realtime",
):
return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler(
websocket_messages=messages,
logging_obj=logging_obj if logging_obj is not None else _logging_obj(),
url_route=url_route,
start_time=SESSION_START,
kwargs=kwargs,
)
def _realtime_cost(model: str, text_in: int, audio_in: int, text_out: int, audio_out: int) -> float:
prices: Final = litellm.model_cost[model]
return (
text_in * prices["input_cost_per_token"]
+ audio_in * prices["input_cost_per_audio_token"]
+ text_out * prices["output_cost_per_token"]
+ audio_out * prices["output_cost_per_audio_token"]
)
def test_realtime_session_is_priced_on_the_model_the_session_reported():
logging_obj: Final = _logging_obj()
messages: Final = [
_session_created(REALTIME_MODEL),
{"type": "response.created", "response": {"id": "resp_fake"}},
_realtime_response_done(text_in=33, audio_in=120, text_out=16, audio_out=240),
]
handled: Final = _handle(messages, logging_obj=logging_obj)
expected_cost: Final = _realtime_cost(REALTIME_MODEL, 33, 120, 16, 240)
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert handled["kwargs"]["model"] == REALTIME_MODEL
assert handled["kwargs"]["custom_llm_provider"] == "openai"
assert handled["result"].model == REALTIME_MODEL
assert handled["result"].usage.total_tokens == 409
assert logging_obj.model_call_details["response_cost"] == pytest.approx(expected_cost)
assert logging_obj.model_call_details["model"] == REALTIME_MODEL
def test_realtime_cost_covers_every_turn_of_the_session():
one_turn: Final = _handle([_session_created(REALTIME_MODEL), _realtime_response_done(10, 0, 5, 0)])
two_turns: Final = _handle(
[
_session_created(REALTIME_MODEL),
_realtime_response_done(10, 0, 5, 0),
_realtime_response_done(10, 0, 5, 0),
]
)
assert two_turns["kwargs"]["response_cost"] == pytest.approx(2 * one_turn["kwargs"]["response_cost"])
assert two_turns["result"].usage.total_tokens == 30
def test_responses_turn_is_priced_on_the_model_the_response_reported():
handled: Final = _handle(
[_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)],
url_route="/openai_passthrough/v1/responses",
)
prices: Final = litellm.model_cost[RESPONSES_MODEL]
expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert handled["kwargs"]["model"] == RESPONSES_MODEL
assert handled["result"].usage.prompt_tokens == 13
assert handled["result"].usage.completion_tokens == 5
def test_responses_cost_covers_every_turn_of_the_connection():
handled: Final = _handle(
[
_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5),
_responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7),
],
url_route="/openai_passthrough/v1/responses",
)
prices: Final = litellm.model_cost[RESPONSES_MODEL]
expected_cost: Final = 33 * prices["input_cost_per_token"] + 12 * prices["output_cost_per_token"]
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert handled["result"].usage.total_tokens == 45
def test_a_model_the_caller_asked_for_never_prices_a_session_the_provider_never_named():
"""
The caller picks what to ask for, and the provider picks what to run. Only
the provider's own frames may set the price of a session, so a model the
caller put in the request cannot stand in for one the provider never named.
"""
logging_obj: Final = _logging_obj()
handled: Final = _handle(
[_realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0)],
kwargs=MappingProxyType({"model": REALTIME_MODEL}),
logging_obj=logging_obj,
)
assert handled["result"] is None
assert handled["kwargs"] == {"model": REALTIME_MODEL}
assert logging_obj.model_call_details == {}
def test_each_responses_turn_is_priced_on_the_model_that_turn_ran():
"""
One socket can run several models one after another, so a connection that
starts on a big model and finishes on a small one has to pay the big
model's price for the turn that used it.
"""
handled: Final = _handle(
[
_responses_completed(REASONING_MODEL, input_tokens=13, output_tokens=5),
_responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7),
],
url_route="/openai_passthrough/v1/responses",
)
first: Final = litellm.model_cost[REASONING_MODEL]
second: Final = litellm.model_cost[RESPONSES_MODEL]
expected_cost: Final = (
13 * first["input_cost_per_token"]
+ 5 * first["output_cost_per_token"]
+ 20 * second["input_cost_per_token"]
+ 7 * second["output_cost_per_token"]
)
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert handled["result"].usage.total_tokens == 45
def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model():
handled: Final = _handle(
[
_session_created(RESPONSES_MODEL),
_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served={"model": None}),
],
url_route="/openai_passthrough/v1/responses",
)
prices: Final = litellm.model_cost[RESPONSES_MODEL]
expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
def test_a_session_nobody_named_a_model_for_is_left_unpriced():
logging_obj: Final = _logging_obj()
handled: Final = _handle(
[{"type": "response.done", "response": {"usage": {"input_tokens": 10, "output_tokens": 5}}}],
kwargs=MappingProxyType({"litellm_call_id": "call-fake"}),
logging_obj=logging_obj,
)
assert handled["result"] is None
assert handled["kwargs"] == {"litellm_call_id": "call-fake"}
assert logging_obj.model_call_details == {}
def test_a_session_that_reported_no_usage_still_records_its_model():
logging_obj: Final = _logging_obj()
handled: Final = _handle([_session_created(REALTIME_MODEL)], logging_obj=logging_obj)
assert handled["result"] is None
assert handled["kwargs"] == {"model": REALTIME_MODEL, "custom_llm_provider": "openai"}
assert logging_obj.model_call_details["model"] == REALTIME_MODEL
assert logging_obj.model_call_details["custom_llm_provider"] == "openai"
def test_a_responses_turn_without_readable_usage_is_not_billed():
handled: Final = _handle(
[{"type": "response.completed", "response": {"model": RESPONSES_MODEL, "status": "completed"}}],
url_route="/openai_passthrough/v1/responses",
)
assert handled["result"] is None
assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"}
def test_the_price_of_a_multi_model_connection_reaches_the_spend_row():
"""
The logging layer prices whatever result it is handed, and one result can
only name one model, so a connection that ran two models has to carry the
price it already worked out or the cheaper model pays for both turns. The
row it lands in names the model the connection opened on.
"""
handled: Final = _handle(
[
_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5),
_responses_completed(REASONING_MODEL, input_tokens=20, output_tokens=7),
],
url_route="/openai_passthrough/v1/responses",
)
logged_cost: Final = litellm.response_cost_calculator(
response_object=handled["result"],
model=handled["kwargs"]["model"],
custom_llm_provider="openai",
call_type="completion",
optional_params={},
)
session_prices: Final = litellm.model_cost[RESPONSES_MODEL]
priced_on_the_session_model: Final = (
33 * session_prices["input_cost_per_token"] + 12 * session_prices["output_cost_per_token"]
)
assert logged_cost == pytest.approx(handled["kwargs"]["response_cost"])
assert logged_cost > priced_on_the_session_model
assert handled["kwargs"]["model"] == RESPONSES_MODEL
def test_turns_that_start_in_the_same_instant_each_keep_their_own_spend_row():
"""
Spend logs key each row by the id its result carries and skip a repeated id
without a word, so two turns that start in the same instant, on one socket
or on two, each need an id of their own or one of them is never billed.
"""
turns: Final = [
_handle(
[_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)],
url_route="/openai_passthrough/v1/responses",
)
for _ in range(2)
]
spend_log_ids: Final = {
get_spend_logs_id(CallTypes.pass_through.value, turn["result"].model_dump(), turn["kwargs"]) for turn in turns
}
assert len(spend_log_ids) == 2
assert None not in spend_log_ids
def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran():
"""
A client that closes before response.completed leaves no usage behind to
bill, but response.created already named the model, so the row says which
model answered instead of falling back to unknown.
"""
logging_obj: Final = _logging_obj()
handled: Final = _handle(
[{"type": "response.created", "response": {"id": "resp_fake", "model": REASONING_MODEL}}],
logging_obj=logging_obj,
url_route="/openai_passthrough/v1/responses",
)
assert handled["result"] is None
assert handled["kwargs"] == {"model": REASONING_MODEL, "custom_llm_provider": "openai"}
assert logging_obj.model_call_details["model"] == REASONING_MODEL
def _transcription_session(model: str) -> tuple[dict[str, object], ...]:
return (
{"type": "session.created", "session": {"type": "transcription", "audio": {"input": {"transcription": None}}}},
{
"type": "session.updated",
"session": {"type": "transcription", "audio": {"input": {"transcription": {"model": model}}}},
},
)
def _transcription_completed(usage: Mapping[str, object]) -> dict[str, object]:
return {
"type": "conversation.item.input_audio_transcription.completed",
"item_id": "item_fake",
"transcript": "hello",
"usage": usage,
}
@pytest.mark.parametrize(
("usage", "billed_units"),
[
({"type": "duration", "seconds": 3}, ((3, "input_cost_per_second"),)),
(
{
"type": "tokens",
"input_tokens": 30,
"output_tokens": 11,
"input_token_details": {"text_tokens": 0, "audio_tokens": 30},
},
((30, "input_cost_per_audio_token"), (11, "output_cost_per_token")),
),
],
)
def test_a_transcription_session_is_billed_for_every_clip_it_transcribed(usage, billed_units):
handled: Final = _handle(
[
*_transcription_session(TRANSCRIPTION_MODEL),
_transcription_completed(usage),
_transcription_completed(usage),
],
url_route="/openai_passthrough/v1/realtime",
)
logged_cost: Final = litellm.response_cost_calculator(
response_object=handled["result"],
model=handled["kwargs"]["model"],
custom_llm_provider="openai",
call_type="completion",
optional_params={},
)
prices: Final = litellm.model_cost[TRANSCRIPTION_MODEL]
expected_cost: Final = 2 * sum(units * prices[rate] for units, rate in billed_units)
assert expected_cost > 0
assert handled["kwargs"]["model"] == TRANSCRIPTION_MODEL
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert logged_cost == pytest.approx(expected_cost)
def test_a_conversation_that_also_transcribes_the_caller_is_logged_on_its_own_model():
handled: Final = _handle(
[
{
"type": "session.created",
"session": {
"model": REALTIME_MODEL,
"audio": {"input": {"transcription": {"model": TRANSCRIPTION_MODEL}}},
},
},
_realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0),
]
)
assert handled["kwargs"]["model"] == REALTIME_MODEL
@pytest.mark.parametrize("event", ["response.incomplete", "response.failed"])
def test_a_responses_turn_that_ended_early_is_billed_for_the_tokens_it_used(event):
"""
A turn cut short by max_output_tokens or a provider failure still reports
the tokens it burned, and OpenAI bills them, so the row has to as well.
"""
handled: Final = _handle(
[_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, event=event)],
url_route="/openai_passthrough/v1/responses",
)
prices: Final = litellm.model_cost[RESPONSES_MODEL]
expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert handled["result"].usage.total_tokens == 18
def test_a_failed_responses_turn_that_reported_no_usage_is_not_billed():
handled: Final = _handle(
[{"type": "response.failed", "response": {"model": RESPONSES_MODEL, "status": "failed", "usage": None}}],
url_route="/openai_passthrough/v1/responses",
)
assert handled["result"] is None
assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"}
@pytest.mark.parametrize(
("turn_model", "requested_model", "priced_as"),
[
("gpt-unreleased-qa-snapshot", RESPONSES_MODEL, RESPONSES_MODEL),
("gpt-5.4-mini-qa-preview", RESPONSES_MODEL, RESPONSES_MODEL),
(RESPONSES_MODEL, REASONING_MODEL, RESPONSES_MODEL),
],
ids=["unmapped_model", "model_the_cost_map_prices_at_zero", "priced_turn_model"],
)
def test_a_responses_turn_falls_back_to_the_requested_model_only_when_its_own_has_no_price(
turn_model: str, requested_model: str, priced_as: str
):
"""
The provider can answer under a model name the cost map does not price
yet. That turn still ran the model the key was allowed to ask for, so it is
billed at that model's price instead of $0, while a turn whose own model
has a price keeps it.
"""
handled: Final = _handle(
[_responses_completed(turn_model, input_tokens=13, output_tokens=5)],
kwargs=MappingProxyType({"litellm_params": {"metadata": {"model_group": requested_model}}}),
url_route="/openai_passthrough/v1/responses",
)
prices: Final = litellm.model_cost[priced_as]
expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
assert expected_cost > 0
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert handled["kwargs"]["model"] == turn_model
def _http_passthrough_cost(response: Mapping[str, object]) -> float:
http_response: Final = httpx.Response(
200, json=response, request=httpx.Request("POST", "https://api.openai.com/v1/responses")
)
_, cost = OpenAIPassthroughLoggingHandler._build_responses_api_response_and_cost(
RESPONSES_MODEL, http_response, _logging_obj(), "openai"
)
return cost
@pytest.mark.parametrize(
"served",
[
{"service_tier": "priority"},
{"output": [{"type": "web_search_call", "id": "ws_fake", "status": "completed", "action": {"type": "search"}}]},
],
ids=["priority_tier", "hosted_web_search"],
)
def test_a_responses_turn_costs_what_the_http_passthrough_charges_for_the_same_response(served):
"""
The provider bills a turn by the tier that served it and the hosted tools it
called, not only its tokens, so a websocket turn has to cost exactly what the
same response costs over HTTP, never the cheaper token-only price.
"""
turn: Final = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served=served)
plain_turn: Final = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)
handled: Final = _handle([turn], url_route="/openai_passthrough/v1/responses")
assert handled["kwargs"]["response_cost"] == pytest.approx(_http_passthrough_cost(turn["response"]))
assert handled["kwargs"]["response_cost"] > _http_passthrough_cost(plain_turn["response"])
@pytest.mark.parametrize(
("frame_type", "kind"),
[
("response.completed", "ends_turn"),
("response.incomplete", "ends_turn"),
("response.failed", "ends_turn"),
("response.done", "ends_turn"),
("session.created", "opens_session"),
("session.updated", "updates_session"),
("conversation.item.input_audio_transcription.completed", "bills_usage"),
("response.created", "starts_turn"),
("response.output_text.delta", "relayed"),
("error", "refuses_turn"),
],
)
def test_each_frame_is_classified_by_what_it_does_to_a_turn(frame_type, kind):
assert openai_websocket_frame_kind({"type": frame_type}) == kind

1535
tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

File diff suppressed because it is too large Load diff

53
tests/unit/proxy/pass_through_endpoints/test_success_handler.py
View file

@ -652,6 +652,59 @@ def test_vertex_ai_live_route_without_usage_frames_yields_no_logging_response(
assert normalized["kwargs"] == {"model": _LIVE_MODEL}
REALTIME_FRAMES: Final = [
{"type": "session.created", "session": {"id": "sess_fake", "model": "gpt-realtime-2.1"}},
{
"type": "response.done",
"response": {
"usage": {
"input_tokens": 33,
"output_tokens": 16,
"total_tokens": 49,
"input_token_details": {"text_tokens": 33, "audio_tokens": 0, "cached_tokens": 0},
"output_token_details": {"text_tokens": 16, "audio_tokens": 0},
}
},
},
]
def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, object]]) -> dict:
logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id="call-fake")
return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload(
httpx_response=MagicMock(spec=httpx.Response),
response_body=response_body,
request_body={},
logging_obj=logging_obj,
url_route=url_route,
result="websocket_connection_successful",
start_time=_START,
end_time=_END,
cache_hit=False,
litellm_call_id="call-fake",
)
@pytest.mark.parametrize("url_route", ["/openai/v1/realtime", "/openai_passthrough/v1/realtime"])
def test_openai_websocket_frames_reach_the_websocket_cost_handler(url_route):
normalized: Final = _normalize(url_route, REALTIME_FRAMES)
assert normalized["kwargs"]["model"] == "gpt-realtime-2.1"
assert normalized["kwargs"]["custom_llm_provider"] == "openai"
prices: Final = litellm.model_cost["gpt-realtime-2.1"]
assert normalized["kwargs"]["response_cost"] == pytest.approx(
33 * prices["input_cost_per_token"] + 16 * prices["output_cost_per_token"]
)
assert normalized["standard_logging_response_object"].usage.total_tokens == 49
def test_openai_http_passthrough_bodies_are_left_to_the_http_handlers():
normalized: Final = _normalize("/openai/v1/responses", {"model": "gpt-5.4-mini", "usage": {"input_tokens": 13}})
assert normalized["standard_logging_response_object"] is None
assert normalized["kwargs"] == {"litellm_call_id": "call-fake"}
@pytest.mark.asyncio
@pytest.mark.parametrize("cancel_mode", ["once", "twice", "scope", "none"])
@pytest.mark.parametrize("cancel_phase", ["allocation", "lock"])

286
tests/unit/proxy/proxy_server/test_streaming_helpers.py
View file

@ -5,7 +5,6 @@ Pins covered:
- ``async_assistants_data_generator``
- ``_get_client_requested_model_for_streaming``
- ``_restamp_streaming_chunk_model``
- ``_fast_serialize_simple_model_response_stream``
- ``_serialize_streaming_chunk``
- ``_apply_streaming_chunk_hooks``
- ``_format_streaming_sse_chunk``
@ -16,8 +15,11 @@ Pins covered:
from __future__ import annotations
import asyncio
import hashlib
import json
import logging
from collections.abc import AsyncIterator, Callable
from types import MappingProxyType
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock
@ -36,7 +38,6 @@ from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.common_request_processing import create_response
from litellm.proxy.proxy_server import (
_apply_streaming_chunk_hooks,
_fast_serialize_simple_model_response_stream,
_format_fallback_metadata_sse_event,
_format_streaming_sse_chunk,
_get_client_requested_model_for_streaming,
@ -55,7 +56,14 @@ from litellm.types.llms.openai import (
ResponseFailedEvent,
ResponsesAPIResponse,
)
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
Delta,
Function,
ModelResponseStream,
StreamingChoices,
Usage,
)
from tests._support.clock_loops import AutojumpClockLoop
from .conftest import normalize, route_entry_points_to_stubs
@ -81,6 +89,10 @@ def _simple_chunk(model: str = "gpt-4", content: str = "hi") -> ModelResponseStr
)
def _sse_text(frame: str | bytes) -> str:
return frame.decode() if isinstance(frame, bytes) else frame
async def _async_iter(items):
for it in items:
yield it
@ -520,48 +532,13 @@ def test_format_fallback_metadata_sse_event():
assert isinstance(payload["created"], int)
# ---------------------------------------------------------------------------
# _fast_serialize_simple_model_response_stream
# ---------------------------------------------------------------------------
def test_fast_serialize_simple_model_response_stream_returns_bytes_payload():
chunk = _simple_chunk()
result = _fast_serialize_simple_model_response_stream(chunk)
assert isinstance(result, bytes)
payload = json.loads(result)
assert normalize(payload) == {
"id": "<VOLATILE>",
"object": "chat.completion.chunk",
"created": "<VOLATILE>",
"model": "gpt-4",
"choices": [
{
"index": 0,
"delta": {"role": "assistant", "content": "hi"},
}
],
}
def test_fast_serialize_simple_model_response_stream_with_usage_returns_none_invalid():
"""Fast path bails (returns None) when ``usage`` is populated — the slow
path is required to preserve usage fields. Returning None here is the
"I cannot handle this" sentinel, not a hard error."""
chunk = _simple_chunk()
chunk.usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2)
assert _fast_serialize_simple_model_response_stream(chunk) is None
# ---------------------------------------------------------------------------
# _serialize_streaming_chunk
# ---------------------------------------------------------------------------
def test_serialize_streaming_chunk_simple_uses_fast_path_bytes():
result = _serialize_streaming_chunk(_simple_chunk())
assert isinstance(result, bytes)
payload = json.loads(result)
def test_serialize_streaming_chunk_simple_payload():
payload = json.loads(_serialize_streaming_chunk(_simple_chunk()))
assert normalize(payload) == {
"id": "<VOLATILE>",
"object": "chat.completion.chunk",
@ -700,10 +677,7 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch):
out.append(line)
assert out[-1] == "data: [DONE]\n\n"
# First chunk is bytes (fast path) wrapped via _format_streaming_sse_chunk.
first = out[0]
assert isinstance(first, bytes)
payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n"))
payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
assert normalize(payload) == {
"id": "<VOLATILE>",
"object": "chat.completion.chunk",
@ -718,6 +692,205 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch):
}
_AZURE_CONTENT_FILTER_RESULTS: Final = {
"hate": {"filtered": False, "severity": "safe"},
"protected_material_text": {"detected": False, "filtered": False},
}
def _azure_chunk(choice_fields: dict[str, object], chunk_fields: dict[str, object]) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-azure",
choices=[StreamingChoices(index=0, delta=Delta(content="Red"), **choice_fields)],
created=0,
model="gpt-4",
object="chat.completion.chunk",
**chunk_fields,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("chunk", "read_field", "sent"),
[
pytest.param(
_azure_chunk({"content_filter_results": _AZURE_CONTENT_FILTER_RESULTS}, {}),
lambda payload: payload["choices"][0]["content_filter_results"],
_AZURE_CONTENT_FILTER_RESULTS,
id="choice-content-filter-results",
),
pytest.param(
_azure_chunk({}, {"obfuscation": "Ab3"}),
lambda payload: payload["obfuscation"],
"Ab3",
id="chunk-obfuscation",
),
pytest.param(
ModelResponseStream(
id="chatcmpl-or",
choices=[
StreamingChoices(
index=0,
delta=Delta(content="Red", reasoning_details=[{"type": "reasoning.text", "text": "colors"}]),
)
],
created=0,
model="gpt-4",
object="chat.completion.chunk",
),
lambda payload: payload["choices"][0]["delta"]["reasoning_details"],
[{"type": "reasoning.text", "text": "colors"}],
id="delta-reasoning-details",
),
pytest.param(
ModelResponseStream(
id="chatcmpl-resp",
choices=[
StreamingChoices(
index=0,
delta=Delta(content="Red", reasoning_items=[{"type": "reasoning", "id": "rs_1", "summary": []}]),
)
],
created=0,
model="gpt-4",
object="chat.completion.chunk",
),
lambda payload: payload["choices"][0]["delta"]["reasoning_items"],
[{"type": "reasoning", "id": "rs_1", "summary": []}],
id="delta-reasoning-items",
),
],
)
async def test_async_data_generator_sse_frame_keeps_provider_fields_on_a_text_chunk(
monkeypatch, chunk: ModelResponseStream, read_field: Callable[[dict[str, object]], object], sent: object
):
_patch_logging_flags(monkeypatch)
out = [
line
async for line in async_data_generator(
response=_async_iter([chunk]),
user_api_key_dict=_user_auth(),
request_data={"model": "gpt-4"},
)
]
assert read_field(json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))) == sent
def _circular_dict() -> dict[str, object]:
looped: Final[dict[str, object]] = {"keep": 1} # mutable-ok: a circular dict has to be inserted into itself
looped["self"] = looped
return looped
@pytest.mark.asyncio
@pytest.mark.parametrize(
"unencodable",
[
pytest.param(object(), id="unknown-type"),
pytest.param(hashlib.sha256(b"trace").digest(), id="invalid-utf8-bytes"),
],
)
async def test_async_data_generator_streams_a_text_chunk_a_hook_tagged_with_a_value_json_cannot_encode(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, unencodable: object
) -> None:
_patch_logging_flags(monkeypatch)
chunk: Final = _azure_chunk({"content_filter_results": _AZURE_CONTENT_FILTER_RESULTS}, {"obfuscation": "Ab3"})
setattr(chunk, "trace_chunk", unencodable)
setattr(chunk.choices[0], "trace_choice", unencodable)
setattr(chunk.choices[0].delta, "trace_delta", unencodable)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
out: Final = [
line
async for line in async_data_generator(
response=_async_iter([chunk]),
user_api_key_dict=_user_auth(),
request_data={"model": "gpt-4"},
)
]
payload: Final = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
assert payload["obfuscation"] == "Ab3"
assert payload["choices"][0]["content_filter_results"] == _AZURE_CONTENT_FILTER_RESULTS
assert payload["choices"][0]["delta"] == {"content": "Red"}
assert "trace_chunk" not in payload
assert "trace_choice" not in payload["choices"][0]
assert all(name in caplog.text for name in ("trace_chunk", "trace_choice", "trace_delta"))
def test_serialize_streaming_chunk_keeps_lone_surrogates_a_provider_split_across_frames(
caplog: pytest.LogCaptureFixture,
) -> None:
tool_call: Final = ChatCompletionDeltaToolCall(
id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Par\udc00is"}')
)
chunk: Final = ModelResponseStream(
id="chatcmpl-split",
choices=[StreamingChoices(index=0, delta=Delta(content="Hi \ud83d", tool_calls=[tool_call]))],
created=0,
model="gpt-4",
object="chat.completion.chunk",
)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
delta: Final = json.loads(_serialize_streaming_chunk(chunk))["choices"][0]["delta"]
assert caplog.text == ""
assert delta["content"] == "Hi \ud83d"
assert delta["tool_calls"][0]["id"] == "call_1"
assert delta["tool_calls"][0]["function"] == {"name": "get_weather", "arguments": '{"city": "Par\udc00is"}'}
def test_serialize_streaming_chunk_drops_only_the_nested_value_json_cannot_encode() -> None:
digest: Final = hashlib.sha256(b"trace").digest()
chunk: Final = ModelResponseStream(
id="chatcmpl-nested",
choices=[
StreamingChoices(index=0, delta=Delta(content="a")),
StreamingChoices(
index=1, delta=Delta(content="b", provider_specific_fields={"citations": ["ok"], "raw": digest})
),
],
created=0,
model="gpt-4",
object="chat.completion.chunk",
usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5),
)
setattr(chunk.usage, "raw", digest)
setattr(chunk, "trace_chunk", _circular_dict())
payload: Final = json.loads(_serialize_streaming_chunk(chunk))
assert payload["choices"][0]["delta"] == {"content": "a"}
assert payload["choices"][1]["delta"] == {"content": "b", "provider_specific_fields": {"citations": ["ok"]}}
assert payload["usage"] == {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}
assert payload["trace_chunk"] == {"keep": 1}
def test_serialize_streaming_chunk_keeps_streaming_when_unencodable_values_sit_under_non_string_keys() -> None:
chunk: Final = ModelResponseStream(
id="chatcmpl-keys",
choices=[StreamingChoices(index=0, delta=Delta(content="still streaming"))],
created=0,
model="gpt-4",
object="chat.completion.chunk",
)
setattr(chunk, "by_index", {0: "kept", 1: object()})
setattr(chunk, "mixed_keys", {"name": "x", 2: object()})
setattr(chunk, "bytes_keys", {b"raw": object()})
setattr(chunk, "read_only_view", MappingProxyType({"name": object()}))
payload: Final = json.loads(_serialize_streaming_chunk(chunk))
assert payload["choices"][0]["delta"] == {"content": "still streaming"}
assert payload["by_index"] == {"0": "kept"}
assert "mixed_keys" not in payload
assert "bytes_keys" not in payload
assert "read_only_view" not in payload
@pytest.mark.asyncio
async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch):
_patch_logging_flags(monkeypatch)
@ -739,9 +912,7 @@ async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch)
):
out.append(line)
first = out[0]
assert isinstance(first, bytes)
payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n"))
payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
assert payload["model"] == "fallback-model"
@ -764,9 +935,7 @@ async def test_async_data_generator_uses_chunk_fallback_metadata(monkeypatch):
):
out.append(line)
first = out[0]
assert isinstance(first, bytes)
payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n"))
payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
assert payload["model"] == "fallback-model"
@ -793,8 +962,8 @@ async def test_async_data_generator_switches_model_mid_stream_on_fallback(monkey
):
out.append(line)
first_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n"))
second_payload = json.loads(out[1].removeprefix(b"data: ").removesuffix(b"\n\n"))
first_payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
second_payload = json.loads(_sse_text(out[1]).removeprefix("data: ").removesuffix("\n\n"))
assert first_payload["model"] == "primary-model"
assert second_payload["model"] == "fallback-model"
@ -830,8 +999,7 @@ async def test_async_data_generator_emits_fallback_error_metadata_event(monkeypa
):
out.append(line)
assert isinstance(out[0], bytes)
chunk_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n"))
chunk_payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
assert chunk_payload["model"] == "fallback-model"
assert isinstance(out[1], str)
assert out[1].startswith("data: ")
@ -879,8 +1047,7 @@ async def test_async_data_generator_skips_fallback_error_event_without_opt_in(
):
out.append(line)
assert isinstance(out[0], bytes)
payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n"))
payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
assert payload["model"] == "fallback-model"
@ -2434,14 +2601,11 @@ async def test_queue_request_stream_is_untouched_while_keepalives_are_unconfigur
assert chunks[-1] == b"data: [DONE]\n\n"
def test_fast_serialize_simple_model_response_stream_keeps_served_service_tier():
def test_serialize_streaming_chunk_keeps_served_service_tier():
chunk = _simple_chunk()
chunk.service_tier = "priority"
result = _fast_serialize_simple_model_response_stream(chunk)
assert result is not None
assert json.loads(result)["service_tier"] == "priority"
assert json.loads(_serialize_streaming_chunk(chunk))["service_tier"] == "priority"
@pytest.mark.asyncio

439
tests/unit/proxy/test_openai_ws_passthrough_routes.py
View file

@ -1,31 +1,41 @@
"""OpenAI passthrough WebSocket route: registration, opt-in gating, and refusals."""
"""OpenAI passthrough WebSocket route: registration, opt-in gating, refusals, and per-frame model checks."""
import json
from collections.abc import Mapping, Sequence
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType, SimpleNamespace
from typing import Final
from unittest.mock import patch
from unittest.mock import AsyncMock, patch
import httpx
import pytest
import respx
from fastapi import HTTPException
from starlette.routing import WebSocketRoute
from starlette.websockets import WebSocket
import litellm
from litellm.llms.openai.workload_identity import _workload_identity_auth
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
_OPENAI_WS_ALIASED_MODEL_REFUSAL,
_OPENAI_WS_DISABLED_REFUSAL,
_OPENAI_WS_MODEL_RESTRICTED_REFUSAL,
_has_model_restrictions,
_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL,
_openai_websocket_refusal,
_proxy_model_allowlists,
_proxy_frame_model_gate,
openai_websocket_proxy_route,
router,
)
Scopes = tuple[Sequence[str], ...]
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 (
WebsocketFrameModelAllowed,
WebsocketFrameModelGate,
WebsocketFrameModelRefused,
WebsocketFrameModelVerdict,
WebsocketUpstreamFrameClassifier,
)
ENABLED: Final = MappingProxyType({"enable_openai_websocket_passthrough": True})
DISABLED_SETTINGS: Final = (
@ -76,6 +86,8 @@ class _RelayCall:
forward_headers: bool
endpoint: str
accept_websocket: bool
client_frame_model_gate: WebsocketFrameModelGate | None
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None
class _FakeRelay:
@ -92,6 +104,8 @@ class _FakeRelay:
forward_headers: bool,
endpoint: str,
accept_websocket: bool,
client_frame_model_gate: WebsocketFrameModelGate | None = None,
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None,
) -> None:
self.calls.append(
_RelayCall(
@ -100,24 +114,26 @@ class _FakeRelay:
forward_headers=forward_headers,
endpoint=endpoint,
accept_websocket=accept_websocket,
client_frame_model_gate=client_frame_model_gate,
upstream_frame_kind=upstream_frame_kind,
)
)
class _FakeModelAllowlists:
def __init__(self, scopes: Scopes) -> None:
self.scopes = scopes
self.calls: list[UserAPIKeyAuth] = []
class _FakeFrameModelGate:
def __init__(self, refusals: Mapping[str, WebsocketFrameModelRefused] = MappingProxyType({})) -> None:
self.refusals = refusals
self.checked: list[str] = []
async def __call__(self, valid_token: UserAPIKeyAuth, /) -> Scopes:
self.calls.append(valid_token)
return self.scopes
async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict:
self.checked.append(model)
return self.refusals.get(model) or WebsocketFrameModelAllowed(access_groups=())
@dataclass(frozen=True, slots=True)
class _Served:
relay: _FakeRelay
allowlists: _FakeModelAllowlists
frame_model_gate: _FakeFrameModelGate
async def _serve(
@ -125,16 +141,16 @@ async def _serve(
endpoint: str,
user_api_key_dict: UserAPIKeyAuth,
general_settings: Mapping[str, object],
scopes: Scopes = (),
frame_model_gate: _FakeFrameModelGate | None = None,
) -> _Served:
served = _Served(relay=_FakeRelay(), allowlists=_FakeModelAllowlists(scopes))
served: Final = _Served(relay=_FakeRelay(), frame_model_gate=frame_model_gate or _FakeFrameModelGate())
await openai_websocket_proxy_route(
websocket=websocket,
endpoint=endpoint,
user_api_key_dict=user_api_key_dict,
general_settings=general_settings,
relay=served.relay,
model_allowlists=served.allowlists,
frame_model_gate=served.frame_model_gate,
)
return served
@ -143,18 +159,20 @@ async def _serve(
@pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"])
async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, monkeypatch):
monkeypatch.delenv("OPENAI_API_BASE", raising=False)
websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview")
websocket: Final = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1")
with patch(GET_CREDENTIALS, return_value="sk-provider"):
served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED)
assert served.relay.calls == [
_RelayCall(
target="wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview",
target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1",
custom_headers=MappingProxyType({"Authorization": "Bearer sk-provider"}),
forward_headers=False,
endpoint=f"/{prefix}/v1/realtime",
accept_websocket=False,
client_frame_model_gate=served.frame_model_gate,
upstream_frame_kind=openai_websocket_frame_kind,
)
]
assert websocket.accepts == [None]
@ -162,11 +180,23 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m
assert websocket.closed is None
@pytest.mark.asyncio
async def test_openai_websocket_reports_its_route_without_the_proxy_root_path(monkeypatch):
"""Route allowlists and spend rows key on the route the proxy serves, which a root path prefix is not part of"""
monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm")
websocket: Final = _FakeWebSocket("/litellm/openai/v1/responses", "")
with patch(GET_CREDENTIALS, return_value="sk-provider"):
served: Final = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED)
assert [call.endpoint for call in served.relay.calls] == ["/openai/v1/responses"]
@pytest.mark.asyncio
async def test_openai_websocket_accepts_first_client_subprotocol():
websocket = _FakeWebSocket(
"/openai/v1/realtime",
"model=gpt-4o-realtime-preview",
"model=gpt-realtime-2.1",
subprotocols="realtime, openai-insecure-api-key.sk-abc, openai-beta.realtime-v1",
)
@ -239,7 +269,7 @@ async def test_openai_websocket_closes_cleanly_when_workload_identity_exchange_f
@pytest.mark.asyncio
async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing():
websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview")
websocket: Final = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1")
with patch(GET_CREDENTIALS, return_value=None):
served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED)
@ -255,7 +285,7 @@ async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing
@pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"])
@pytest.mark.parametrize("general_settings", DISABLED_SETTINGS)
async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, general_settings):
websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview")
websocket: Final = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1")
served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings)
@ -265,25 +295,21 @@ async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, genera
assert served.relay.calls == []
@pytest.mark.asyncio
@pytest.mark.parametrize("general_settings", DISABLED_SETTINGS)
async def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings):
refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), general_settings, _FakeModelAllowlists(()))
assert refusal is _OPENAI_WS_DISABLED_REFUSAL
def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings):
assert _openai_websocket_refusal(general_settings) is _OPENAI_WS_DISABLED_REFUSAL
@pytest.mark.asyncio
@pytest.mark.parametrize("value", [True, "true", "True"])
async def test_openai_websocket_refusal_is_none_for_truthy_settings(value):
settings = MappingProxyType({"enable_openai_websocket_passthrough": value})
assert await _openai_websocket_refusal(UserAPIKeyAuth(), settings, _FakeModelAllowlists(())) is None
def test_openai_websocket_refusal_is_none_for_truthy_settings(value):
assert _openai_websocket_refusal(MappingProxyType({"enable_openai_websocket_passthrough": value})) is None
@pytest.mark.asyncio
async def test_openai_websocket_refusal_echoes_requested_subprotocol():
websocket = _FakeWebSocket(
"/openai_passthrough/v1/realtime",
"model=gpt-4o-realtime-preview",
"model=gpt-realtime-2.1",
subprotocols="realtime, openai-beta.realtime-v1",
)
@ -294,68 +320,303 @@ async def test_openai_websocket_refusal_echoes_requested_subprotocol():
assert served.relay.calls == []
RESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = (
(("gpt-4o",),),
((), ("gpt-4o-realtime-preview",)),
(("all-team-models",), ("gpt-4o",)),
((), ("all-proxy-models",), ("gpt-4o",)),
((), (), (), ("gpt-4o",)),
(("*",), (), (), (), ("gpt-4o",)),
@pytest.mark.asyncio
@pytest.mark.parametrize(
"key_models",
[["gpt-5.4-mini"], ["gpt-realtime-2.1", "gpt-5.4"], []],
)
UNRESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = (
(),
((),),
(("all-proxy-models",),),
(("*",),),
(("all-team-models",), ("all-proxy-models",)),
((), (), (), (), ()),
(("*",), ("all-proxy-models",), ("all-team-models",), (), ()),
)
@pytest.mark.asyncio
@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES)
async def test_openai_websocket_rejects_model_restricted_identities(scopes):
websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview")
user_api_key_dict = UserAPIKeyAuth(token="hashed-fake", user_id="user-fake", team_id="team-fake")
served = await _serve(websocket, "v1/realtime", user_api_key_dict, ENABLED, scopes)
assert "model restrictions" in websocket.error_message()
assert websocket.closed == (1008, _OPENAI_WS_MODEL_RESTRICTED_REFUSAL.close_reason)
assert served.relay.calls == []
assert served.allowlists.calls == [user_api_key_dict]
@pytest.mark.asyncio
@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES)
async def test_openai_websocket_disabled_refusal_skips_allowlist_lookups(scopes):
allowlists = _FakeModelAllowlists(scopes)
refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), MappingProxyType({}), allowlists)
assert refusal is _OPENAI_WS_DISABLED_REFUSAL
assert allowlists.calls == []
@pytest.mark.asyncio
@pytest.mark.parametrize("scopes", UNRESTRICTED_SCOPES)
async def test_openai_websocket_allows_unrestricted_identities(scopes):
websocket = _FakeWebSocket("/openai/v1/responses", "")
async def test_openai_websocket_connects_model_restricted_keys(key_models):
websocket: Final = _FakeWebSocket("/openai/v1/responses", "")
user_api_key_dict: Final = UserAPIKeyAuth(models=key_models, token="hashed-fake", user_id="user-fake")
with patch(GET_CREDENTIALS, return_value="sk-provider"):
served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED, scopes)
served: Final = await _serve(websocket, "v1/responses", user_api_key_dict, ENABLED)
assert len(served.relay.calls) == 1
assert served.relay.calls[0].client_frame_model_gate is served.frame_model_gate
assert websocket.sent == []
assert websocket.closed is None
@pytest.mark.asyncio
async def test_proxy_model_allowlists_reads_the_token_scopes_without_a_database():
token: Final = UserAPIKeyAuth(models=[], team_id="team-fake", team_models=["gpt-4o"])
with patch("litellm.proxy.proxy_server.prisma_client", None):
scopes = await _proxy_model_allowlists()(token)
ALIAS_REFUSAL: Final = WebsocketFrameModelRefused(
message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=ProxyErrorTypes.key_model_access_denied.value
)
assert tuple(tuple(scope) for scope in scopes) == ((), ("gpt-4o",))
assert _has_model_restrictions(scopes)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"query",
[
"model=gpt-realtime-2.1",
"model=gpt-realtime-2.1&model=gpt-realtime-2.1-mini",
"model=gpt-realtime-2.1-mini&model=gpt-realtime-2.1",
],
ids=["one_model", "refused_model_first", "refused_model_last"],
)
async def test_openai_websocket_refuses_an_opening_model_the_frame_gate_refuses(query):
"""
The relay sends the connection URL to OpenAI unchanged, so every model it names has to pass the
same gate as a frame, or an alias the upgrade resolves to an allowed model opens the session on the
alias name itself
"""
websocket: Final = _FakeWebSocket("/openai/v1/realtime", query)
gate: Final = _FakeFrameModelGate(MappingProxyType({"gpt-realtime-2.1": ALIAS_REFUSAL}))
with patch(GET_CREDENTIALS, return_value="sk-provider"):
served: Final = await _serve(
websocket, "v1/realtime", UserAPIKeyAuth(models=["gpt-realtime-2.1-mini"]), ENABLED, gate
)
assert websocket.error_message() == _OPENAI_WS_ALIASED_MODEL_REFUSAL
assert websocket.closed == (1008, "This key cannot use a model alias over OpenAI websocket passthrough.")
assert served.relay.calls == []
assert "gpt-realtime-2.1" in gate.checked
@pytest.mark.asyncio
async def test_openai_websocket_relays_a_session_once_every_opening_model_passes_the_frame_gate():
websocket: Final = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1-mini")
with patch(GET_CREDENTIALS, return_value="sk-provider"):
served: Final = await _serve(
websocket, "v1/realtime", UserAPIKeyAuth(models=["gpt-realtime-2.1-mini"]), ENABLED
)
assert served.frame_model_gate.checked == ["gpt-realtime-2.1-mini"]
assert len(served.relay.calls) == 1
assert websocket.closed is None
class _FrameAuthorizer:
"""Stands in for the full key, team, user and budget auth the gate re-runs for every model a frame names."""
def __init__(self, token: UserAPIKeyAuth | None = None, denial: Exception | None = None) -> None:
self.token = token or UserAPIKeyAuth()
self.denial = denial
self.authorized: list[str] = []
async def __call__(self, websocket: object, model: str) -> UserAPIKeyAuth:
self.authorized.append(model)
if self.denial is not None:
raise self.denial
return self.token
SOCKET: Final = WebSocket({"type": "websocket"}, receive=AsyncMock(), send=AsyncMock())
async def _gate_verdict(authorizer: _FrameAuthorizer, model: str) -> WebsocketFrameModelVerdict:
with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the model checks need no database
return await _proxy_frame_model_gate(authorizer)(SOCKET, model)
def _refusal(verdict: WebsocketFrameModelVerdict) -> WebsocketFrameModelRefused:
assert isinstance(verdict, WebsocketFrameModelRefused)
return verdict
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_runs_full_auth_for_the_model_the_frame_names():
authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"]))
verdict: Final = await _gate_verdict(authorizer, "gpt-realtime-2.1")
assert verdict == WebsocketFrameModelAllowed(access_groups=())
assert authorizer.authorized == ["gpt-realtime-2.1"]
@pytest.mark.asyncio
@pytest.mark.parametrize("requested_model", ["gpt-5.4", "gpt-realtime-2.1-mini"])
async def test_proxy_frame_model_gate_refuses_a_model_the_key_lacks(requested_model):
authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"]))
refusal: Final = _refusal(await _gate_verdict(authorizer, requested_model))
assert requested_model in refusal.message
assert refusal.error_type == ProxyErrorTypes.key_model_access_denied.value
@pytest.mark.asyncio
@pytest.mark.parametrize(
"granted_models, refusal_names",
[(["gpt-5.4-mini"], "gpt-5.4"), ([], "no model grants")],
ids=["model_outside_the_grant", "no_grants"],
)
async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted(granted_models, refusal_names):
from litellm.types.agents import AgentResponse
token: Final = UserAPIKeyAuth(models=[], agent_id="agent-fake")
token.managed_agent_policy = AgentResponse(
agent_id="agent-fake",
agent_name="Agent",
agent_card_params={},
object_permission={"object_permission_id": "grant-fake", "models": granted_models},
)
refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4"))
assert refusal_names in refusal.message
@pytest.mark.asyncio
@pytest.mark.parametrize(
"token",
[
UserAPIKeyAuth(models=["gpt-5.4-mini"], aliases={"gpt-5.4": "gpt-5.4-mini"}),
UserAPIKeyAuth(models=["gpt-5.4-mini"], team_model_aliases={"gpt-5.4": "gpt-5.4-mini"}),
],
ids=["key_alias", "team_alias"],
)
async def test_proxy_frame_model_gate_refuses_an_alias_that_would_reach_openai_unchanged(token):
"""
The passthrough sends each frame to OpenAI as written, so OpenAI runs the
alias name itself. An alias pointing at an allowed model must not let a key
limited to that model run the model the alias is named after.
"""
refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4"))
assert refusal == WebsocketFrameModelRefused(
message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=ProxyErrorTypes.key_model_access_denied.value
)
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_refuses_a_proxy_wide_alias_for_a_key_limited_to_its_target(monkeypatch):
monkeypatch.setattr(litellm, "model_alias_map", {"gpt-5.4": "gpt-5.4-mini"})
refusal: Final = _refusal(
await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=["gpt-5.4-mini"])), "gpt-5.4")
)
assert refusal.message == _OPENAI_WS_ALIASED_MODEL_REFUSAL
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_lets_a_key_that_may_use_any_model_send_an_alias():
token: Final = UserAPIKeyAuth(models=[], aliases={"fast": "gpt-5.4-mini"})
verdict: Final = await _gate_verdict(_FrameAuthorizer(token), "fast")
assert verdict == WebsocketFrameModelAllowed(access_groups=())
@pytest.mark.asyncio
@pytest.mark.parametrize("unrestricted_models", [[], ["all-proxy-models"], ["*"]])
async def test_proxy_frame_model_gate_allows_every_model_for_unrestricted_keys(unrestricted_models):
verdict: Final = await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=unrestricted_models)), "gpt-5.4")
assert verdict == WebsocketFrameModelAllowed(access_groups=())
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_hands_back_the_access_groups_that_authorized_the_model():
"""The turn a frame starts is billed to the groups that let its model through, so the gate has to report them."""
token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"], matched_model_access_groups=["tier-a"])
verdict: Final = await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4-mini")
assert verdict == WebsocketFrameModelAllowed(access_groups=("tier-a",))
@pytest.mark.asyncio
@pytest.mark.parametrize(
"requested_model",
["gpt-5.4-mini", "gpt-5.4"],
ids=["allowed", "refused_by_the_model_check"],
)
async def test_proxy_frame_model_gate_releases_the_budget_it_reserved_for_the_check(requested_model):
"""
Full auth reserves budget for the request it authorizes, but a frame check is
not a request, so a reservation the gate leaves open would hold the key's
spend counter above its real spend and refuse later calls on every route.
"""
reservation: Final[dict[str, object]] = {"reserved_cost": 0.5, "entries": []}
token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"], budget_reservation=reservation)
await _gate_verdict(_FrameAuthorizer(token), requested_model)
assert token.budget_reservation == {**reservation, "finalized": True}
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_does_not_count_the_sessions_own_reservation_as_key_spend():
"""
The upgrade reserves up to the key's remaining budget for the session, and the key budget check
counts open reservations as spend, so a frame checked while that reservation is open is refused
on a key with budget left, where an HTTP call with the same key goes through.
"""
connect_reservation: Final[dict[str, object]] = {"reserved_cost": 0.0001, "entries": []}
socket: Final = WebSocket(
{"type": "websocket", "state": {"budget_reservation": connect_reservation}}, receive=AsyncMock(), send=AsyncMock()
)
async def key_budget_check(websocket: object, model: str) -> UserAPIKeyAuth:
if connect_reservation.get("finalized") is not True:
raise litellm.BudgetExceededError(current_cost=0.0001, max_budget=0.0001)
return UserAPIKeyAuth(models=["gpt-5.4-mini"])
with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the model checks need no database
verdict: Final = await _proxy_frame_model_gate(key_budget_check)(socket, "gpt-5.4-mini")
assert verdict == WebsocketFrameModelAllowed(access_groups=())
@pytest.mark.asyncio
@pytest.mark.parametrize(
"denial, error_type, message",
[
(
ProxyException(
message="Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0",
type=ProxyErrorTypes.budget_exceeded,
param=None,
code=400,
),
"budget_exceeded",
"Budget has been exceeded! Key=hashed",
),
(
litellm.BudgetExceededError(
current_cost=2.0,
max_budget=1.0,
message="Budget has been exceeded! Model access group=tier-a Current cost: 2.0, Max budget: 1.0",
entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value,
entity_id="tier-a",
),
"budget_exceeded",
"Model access group=tier-a",
),
(
ProxyException(
message="team not allowed to access model. This team can only access models=['gpt-5.4-mini']",
type=ProxyErrorTypes.team_model_access_denied,
param="model",
code=401,
),
"team_model_access_denied",
"team not allowed to access model",
),
(HTTPException(status_code=403, detail="Invalid Authorization header format"), "invalid_request_error", ""),
],
ids=["key_budget", "group_budget", "team_allowlist", "malformed_auth"],
)
async def test_proxy_frame_model_gate_refuses_what_full_auth_denies_with_its_reason(denial, error_type, message):
"""A frame is denied for the same reasons a request is, and the caller reads which reason in the error event."""
refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(denial=denial), "gpt-5.4-mini"))
assert refusal.error_type == error_type
assert message in refusal.message
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_refuses_out_loud_when_auth_cannot_answer():
"""
A gate that cannot answer has to close the session with a reason the caller can read.
Letting the failure escape drops the socket with no frame at all, which the caller
cannot tell apart from the network going away.
"""
refusal: Final = _refusal(
await _gate_verdict(_FrameAuthorizer(denial=ConnectionError("connection pool is exhausted")), "gpt-5.4-mini")
)
assert refusal == WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error")

638
tests/unit/proxy/test_proxy_token_counter.py
View file

@ -25,7 +25,7 @@ import pytest
import respx
import tiktoken
from dotenv import load_dotenv
from pydantic import BaseModel, ConfigDict, JsonValue
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
from tokenizers import Tokenizer
from tokenizers.models import WordLevel
from tokenizers.pre_tokenizers import Whitespace
@ -82,6 +82,7 @@ from litellm.proxy.proxy_server import (
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
from litellm.types.utils import (
ADDRESSED_RESPONSE_ID_FIELD,
@ -1741,7 +1742,10 @@ _OPUS_IN_US_EAST_1: Final = TokenCountModel(model="us.anthropic.claude-opus-5-5"
async def _ladder(
counter: _ScriptedCounter, settings: TokenCountSettings, local_estimate_allowed: bool = True
counter: _ScriptedCounter,
settings: TokenCountSettings,
local_estimate_allowed: bool = True,
forwarded_edits: bool = False,
) -> TokenCountResponse | None:
return await _provider_token_count(
count_with=partial(
@ -1759,6 +1763,7 @@ async def _ladder(
request_model=_COUNT_REQUEST.model,
settings=settings,
local_estimate_allowed=local_estimate_allowed,
forwarded_edits=forwarded_edits,
)
@ -1986,6 +1991,26 @@ async def test_configured_token_count_model_error_stands_even_with_framing(
assert raised.value.message.count("counts tokens with us.anthropic.claude-opus-5-5 in us-east-1") == region_hints
@pytest.mark.parametrize(("status_code", "region_hints"), [(400, 0), (403, 1)])
@pytest.mark.asyncio
async def test_a_pinned_token_count_model_rejecting_forwarded_edits_answers_without_the_region_hint(
status_code: int, region_hints: int
):
counter: Final = _ScriptedCounter(
{"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)}
)
with pytest.raises(ProxyException) as raised:
await _ladder(
counter,
TokenCountSettings(framing=None, count_model=replace(_OPUS_IN_US_EAST_1, configured=True)),
forwarded_edits=True,
)
assert raised.value.code == str(status_code)
assert raised.value.message.count("counts tokens with us.anthropic.claude-opus-5-5 in us-east-1") == region_hints
@pytest.mark.parametrize(
("count_model", "region_hints"),
[(_OPUS_IN_US_EAST_1, frozenset({True})), (replace(_OPUS_IN_US_EAST_1, region=None), frozenset())],
@ -3172,15 +3197,14 @@ _READ_FILE_TOOL: Final[dict[str, JsonValue]] = {
"name": "read_file",
"input_schema": {"type": "object", "properties": {"path": {"type": "string"}}},
}
_CLEAR_TOOL_USES: Final[dict[str, JsonValue]] = {
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 500},
"keep": {"type": "tool_uses", "value": 1},
}
]
}
_CLEAR_TOOL_USES_EDITS: Final[list[JsonValue]] = [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 500},
"keep": {"type": "tool_uses", "value": 1},
}
]
_CLEAR_TOOL_USES: Final[dict[str, JsonValue]] = {"edits": _CLEAR_TOOL_USES_EDITS}
_INVALID_COMPACT_EDIT: Final[dict[str, JsonValue]] = {
"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}]
}
@ -3516,6 +3540,28 @@ async def test_gemini_messages_count_tokens_counts_the_unedited_body_when_drop_p
assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == 3
@pytest.mark.asyncio
@pytest.mark.parametrize(
("router_defaults", "request_options", "expected_status"),
(({}, {"drop_params": True}, None), ({"drop_params": True}, {}, None), ({}, {}, 400)),
ids=("request-drop-params", "router-default-drop-params", "no-drop-params"),
)
async def test_messages_count_tokens_for_a_model_with_no_deployment_drops_an_invalid_edit_only_under_drop_params(
monkeypatch: pytest.MonkeyPatch,
router_defaults: dict[str, JsonValue],
request_options: dict[str, JsonValue],
expected_status: int | None,
) -> None:
monkeypatch.setenv("GEMINI_API_KEY", "test-key")
with patch.object(litellm, "drop_params", False):
response, _ = await _count_through_the_messages_endpoint(
Router(model_list=[_gemini_deployment("gemini-agent", {})], default_litellm_params=router_defaults),
{**request_options, "model": "gemini/gemini-3.8-pro", "context_management": _INVALID_COMPACT_EDIT},
)
assert (response.status_code if isinstance(response, JSONResponse) else None) == expected_status
@pytest.mark.asyncio
async def test_gemini_messages_count_tokens_rejects_an_invalid_edit_without_drop_params() -> None:
with patch.object(litellm, "drop_params", False):
@ -3529,6 +3575,576 @@ async def test_gemini_messages_count_tokens_rejects_an_invalid_edit_without_drop
assert sent_bodies == ()
_MISTRAL_MODEL: Final = "mistral/mistral-large-latest"
def _mistral_router(
litellm_params: Mapping[str, JsonValue] = MappingProxyType({}),
router_defaults: Mapping[str, JsonValue] = MappingProxyType({}),
) -> Router:
return Router(
model_list=[
{
"model_name": "mistral-agent",
"litellm_params": {"model": _MISTRAL_MODEL, "api_key": "test-key", **litellm_params},
}
],
default_litellm_params=dict(router_defaults),
)
async def _messages_endpoint_count(router: Router, request_body: Mapping[str, JsonValue]) -> object:
import litellm.proxy.anthropic_endpoints.endpoints as anthropic_endpoints
import litellm.proxy.proxy_server as proxy_server
mock_request: Final = MagicMock(spec=Request)
mock_request.headers = {}
with (
patch.object(proxy_server, "llm_router", router),
patch.object(anthropic_endpoints, "read_request_body", AsyncMock(return_value=dict(request_body))),
):
return await anthropic_count_tokens(mock_request, MagicMock())
class _MessagesCountReply(BaseModel):
model_config = ConfigDict(frozen=True)
input_tokens: int
async def _input_tokens(router: Router, request_body: Mapping[str, JsonValue]) -> int:
return _MessagesCountReply.model_validate(await _messages_endpoint_count(router, request_body)).input_tokens
async def _messages_the_inference_polyfill_sends(model: str) -> list[dict[str, JsonValue]]:
from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request
prepared: Final = await prepare_context_managed_request(
model=model,
messages=[dict(message) for message in _agent_turns_with_tool_results(3)],
tools=[dict(_READ_FILE_TOOL)],
system=None,
context_management_spec=dict(_CLEAR_TOOL_USES),
litellm_metadata=None,
additional_drop_params=None,
llm_router=None,
)
assert prepared is not None
return TypeAdapter(list[dict[str, JsonValue]]).validate_python(prepared.messages)
@pytest.mark.asyncio
async def test_messages_count_tokens_counts_the_body_the_inference_polyfill_sends_on_a_bridged_provider() -> None:
router: Final = _mistral_router()
unedited_body: Final[dict[str, JsonValue]] = {
"model": "mistral-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
}
edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES})
unedited: Final = await _input_tokens(router, unedited_body)
polyfilled: Final = await _input_tokens(
router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(_MISTRAL_MODEL)}
)
assert edited < unedited
assert edited == polyfilled
@pytest.mark.asyncio
@pytest.mark.parametrize(
("deployment_params", "router_defaults", "request_options"),
(
({"additional_drop_params": ["context_management"]}, {}, {}),
({}, {"additional_drop_params": ["context_management"]}, {}),
({}, {}, {"additional_drop_params": ["context_management"]}),
),
ids=("deployment-drops", "router-default-drops", "request-drops"),
)
async def test_messages_count_tokens_on_a_bridged_provider_keeps_the_unedited_count_when_context_management_is_dropped(
deployment_params: dict[str, JsonValue],
router_defaults: dict[str, JsonValue],
request_options: dict[str, JsonValue],
) -> None:
router: Final = _mistral_router(deployment_params, router_defaults)
unedited_body: Final[dict[str, JsonValue]] = {
"model": "mistral-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
**request_options,
}
edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES})
assert edited == await _input_tokens(router, unedited_body)
@pytest.mark.asyncio
async def test_messages_count_tokens_on_a_bridged_provider_rejects_an_invalid_edit_without_drop_params() -> None:
with patch.object(litellm, "drop_params", False):
response: Final = await _messages_endpoint_count(
_mistral_router(),
{
"model": "mistral-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": _INVALID_COMPACT_EDIT,
},
)
assert isinstance(response, JSONResponse)
assert response.status_code == 400
@pytest.mark.asyncio
async def test_messages_count_tokens_on_a_bridged_provider_counts_without_an_invalid_edit_under_drop_params() -> None:
router: Final = _mistral_router()
unedited_body: Final[dict[str, JsonValue]] = {
"model": "mistral-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"drop_params": True,
}
with patch.object(litellm, "drop_params", False):
edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _INVALID_COMPACT_EDIT})
unedited: Final = await _input_tokens(router, unedited_body)
assert edited == unedited
def _turns_compacted_by_the_client() -> list[dict[str, JsonValue]]:
return [
*_agent_turns_with_tool_results(3)[:-1],
{"role": "assistant", "content": [{"type": "compaction", "content": "The agent read files 0, 1 and 2."}]},
{"role": "user", "content": "Summarize what you read."},
]
@pytest.mark.asyncio
async def test_messages_count_tokens_on_a_bridged_provider_counts_only_what_follows_a_client_compaction_block() -> None:
from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request
router: Final = _mistral_router()
sent: Final = await prepare_context_managed_request(
model=_MISTRAL_MODEL,
messages=[dict(message) for message in _turns_compacted_by_the_client()],
tools=[dict(_READ_FILE_TOOL)],
system=None,
context_management_spec=None,
litellm_metadata=None,
additional_drop_params=None,
llm_router=None,
)
assert sent is not None
compacted: Final = await _input_tokens(
router, {"model": "mistral-agent", "messages": _turns_compacted_by_the_client(), "tools": [_READ_FILE_TOOL]}
)
uncompacted: Final = await _input_tokens(
router, {"model": "mistral-agent", "messages": _agent_turns_with_tool_results(3), "tools": [_READ_FILE_TOOL]}
)
sent_body_count: Final = await _input_tokens(
router,
{
"model": "mistral-agent",
"messages": TypeAdapter(list[dict[str, JsonValue]]).validate_python(sent.messages),
"system": TypeAdapter(JsonValue).validate_python(sent.system),
"tools": [_READ_FILE_TOOL],
},
)
assert compacted < uncompacted
assert compacted == sent_body_count
@pytest.mark.asyncio
async def test_messages_count_tokens_applies_the_edits_a_responses_api_bridge_polyfills_for_a_non_openai_provider() -> (
None
):
perplexity_model: Final = "perplexity/perplexity/sonar"
router: Final = Router(
model_list=[{"model_name": "sonar-agent", "litellm_params": {"model": perplexity_model, "api_key": "test-key"}}]
)
unedited_body: Final[dict[str, JsonValue]] = {
"model": "sonar-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
}
edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES})
unedited: Final = await _input_tokens(router, unedited_body)
polyfilled: Final = await _input_tokens(
router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(perplexity_model)}
)
assert edited < unedited
assert edited == polyfilled
@pytest.mark.asyncio
async def test_messages_count_tokens_on_a_responses_api_bridge_counts_the_unedited_body_when_an_edit_fails() -> None:
from litellm.llms.anthropic.messages.responses_adapters import handler as responses_adapters_handler
router: Final = Router(
model_list=[
{
"model_name": "sonar-agent",
"litellm_params": {"model": "perplexity/perplexity/sonar", "api_key": "test-key"},
}
]
)
unedited_body: Final[dict[str, JsonValue]] = {
"model": "sonar-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
}
with patch.object(
responses_adapters_handler, "apply_context_management", AsyncMock(side_effect=RuntimeError("editor bug"))
):
edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES})
assert edited == await _input_tokens(router, unedited_body)
@pytest.mark.asyncio
async def test_messages_count_tokens_on_a_responses_api_bridge_rejects_an_invalid_edit_without_drop_params() -> None:
router: Final = Router(
model_list=[
{
"model_name": "sonar-agent",
"litellm_params": {"model": "perplexity/perplexity/sonar", "api_key": "test-key"},
}
]
)
invalid_clear_edit: Final[dict[str, JsonValue]] = {
"edits": [{"type": "clear_tool_uses_20250919", "clear_at_least": "all of them"}]
}
with patch.object(litellm, "drop_params", False):
response: Final = await _messages_endpoint_count(
router,
{
"model": "sonar-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": invalid_clear_edit,
},
)
assert isinstance(response, JSONResponse)
assert response.status_code == 400
message: Final = json.loads(bytes(response.body))["error"]["message"]
assert "context_management.edits.0.clear_tool_uses_20250919.clear_at_least" in message
class _CountTokensPeer:
def __init__(self, reply: Mapping[str, JsonValue]) -> None:
self._reply: Final = reply
self.requests: tuple[httpx.Request, ...] = ()
def __call__(self, request: httpx.Request) -> httpx.Response:
self.requests = (*self.requests, request)
return httpx.Response(200, json=dict(self._reply), request=request)
class _BedrockInvokeModel(BaseModel):
model_config = ConfigDict(frozen=True)
body: str
class _BedrockCountTokensInput(BaseModel):
model_config = ConfigDict(frozen=True)
invoke_model: _BedrockInvokeModel = Field(alias="invokeModel")
class _BedrockCountTokensBody(BaseModel):
model_config = ConfigDict(frozen=True)
input: _BedrockCountTokensInput
def _bedrock_counted_request(request_body: bytes) -> str:
invoke_model: Final = _BedrockCountTokensBody.model_validate_json(request_body).input.invoke_model
return base64.b64decode(invoke_model.body).decode()
async def _count_on_a_bridged_provider_count_api(
request_body: Mapping[str, JsonValue],
monkeypatch: pytest.MonkeyPatch,
bedrock_model: str = "bedrock/us.amazon.nova-2-lite-v1:0",
peer_reply: Mapping[str, JsonValue] = MappingProxyType({"inputTokens": 57}),
model_info: Mapping[str, JsonValue] = MappingProxyType({}),
) -> tuple[object, tuple[httpx.Request, ...]]:
from litellm.llms.bedrock.count_tokens import handler as bedrock_count_tokens_handler
from litellm.llms.bedrock.count_tokens import mantle_handler as bedrock_mantle_count_tokens_handler
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
peer: Final = _CountTokensPeer(peer_reply)
router: Final = Router(
model_list=[
{
"model_name": "nova-agent",
"litellm_params": {
"model": bedrock_model,
"aws_access_key_id": "AKIATESTACCESSKEY",
"aws_secret_access_key": "test-secret",
"aws_region_name": "us-east-1",
},
"model_info": dict(model_info),
}
]
)
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer))
with (
patch.object(bedrock_count_tokens_handler, "get_async_httpx_client", return_value=client),
patch.object(bedrock_mantle_count_tokens_handler, "get_async_httpx_client", return_value=client),
):
response: Final = await _messages_endpoint_count(router, {**request_body, "model": "nova-agent"})
return response, peer.requests
@pytest.mark.asyncio
async def test_messages_count_tokens_sends_the_edited_messages_to_a_bridged_provider_count_api(
monkeypatch: pytest.MonkeyPatch,
) -> None:
response, sent_bodies = await _count_on_a_bridged_provider_count_api(
{
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": _CLEAR_TOOL_USES,
},
monkeypatch,
)
assert response == {"input_tokens": 57, "litellm_count_source": "provider"}
assert len(sent_bodies) == 1
assert _bedrock_counted_request(sent_bodies[0].content).count(_GEMINI_TOOL_OUTPUT) == 1
@pytest.mark.asyncio
async def test_messages_count_tokens_sends_a_body_with_only_empty_text_to_the_bridged_provider_count_api(
monkeypatch: pytest.MonkeyPatch,
) -> None:
empty_turn: Final[dict[str, JsonValue]] = {"role": "user", "content": [{"type": "text", "text": ""}]}
response, sent_bodies = await _count_on_a_bridged_provider_count_api({"messages": [empty_turn]}, monkeypatch)
assert response == {"input_tokens": 57, "litellm_count_source": "provider"}
assert len(sent_bodies) == 1
_BEDROCK_CLAUDE_MODEL: Final = "bedrock/us.anthropic.claude-sonnet-5-5"
_CLEAR_THINKING_EDIT: Final[dict[str, JsonValue]] = {"type": "clear_thinking_20251015"}
_COUNTED_ON_MANTLE: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"token_count_model": {"model": "us.anthropic.claude-sonnet-5-5", "litellm_provider": "bedrock_mantle"}}
)
_COUNTED_ON_RUNTIME: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"token_count_model": {"model": "us.anthropic.claude-sonnet-5-5", "litellm_provider": "bedrock"}}
)
def _bedrock_claude_count_request(edits: list[JsonValue]) -> Mapping[str, JsonValue]:
return MappingProxyType(
{
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": {"edits": edits},
}
)
@pytest.mark.asyncio
async def test_messages_count_tokens_forwards_bedrock_supported_edits_and_their_beta_to_the_mantle_count_api(
monkeypatch: pytest.MonkeyPatch,
) -> None:
response, sent = await _count_on_a_bridged_provider_count_api(
_bedrock_claude_count_request([*_CLEAR_TOOL_USES_EDITS, _CLEAR_THINKING_EDIT]),
monkeypatch,
_BEDROCK_CLAUDE_MODEL,
MappingProxyType({"input_tokens": 57, "context_management": {"original_input_tokens": 4100}}),
_COUNTED_ON_MANTLE,
)
assert isinstance(response, dict)
assert (response["input_tokens"], response["context_management"]) == (57, {"original_input_tokens": 4100})
assert len(sent) == 1
assert json.loads(sent[0].content)["context_management"] == _CLEAR_TOOL_USES
assert sent[0].headers["anthropic-beta"] == ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value
@pytest.mark.asyncio
async def test_messages_count_tokens_forwards_bedrock_supported_edits_and_their_beta_to_the_runtime_count_api(
monkeypatch: pytest.MonkeyPatch,
) -> None:
response, sent = await _count_on_a_bridged_provider_count_api(
_bedrock_claude_count_request([*_CLEAR_TOOL_USES_EDITS, _CLEAR_THINKING_EDIT]),
monkeypatch,
_BEDROCK_CLAUDE_MODEL,
model_info=_COUNTED_ON_RUNTIME,
)
assert isinstance(response, dict)
assert response["input_tokens"] == 57
assert len(sent) == 1
invoke_body: Final = json.loads(_bedrock_counted_request(sent[0].content))
assert (invoke_body["context_management"], invoke_body["anthropic_beta"]) == (
_CLEAR_TOOL_USES,
[ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value],
)
@pytest.mark.asyncio
async def test_messages_count_tokens_sends_no_context_management_to_bedrock_when_no_edit_is_bedrock_supported(
monkeypatch: pytest.MonkeyPatch,
) -> None:
response, sent = await _count_on_a_bridged_provider_count_api(
_bedrock_claude_count_request([_CLEAR_THINKING_EDIT]),
monkeypatch,
_BEDROCK_CLAUDE_MODEL,
MappingProxyType({"input_tokens": 57}),
_COUNTED_ON_MANTLE,
)
assert isinstance(response, dict)
assert response["input_tokens"] == 57
assert len(sent) == 1
assert "context_management" not in json.loads(sent[0].content)
assert "anthropic-beta" not in sent[0].headers
@pytest.mark.asyncio
async def test_messages_count_tokens_rejects_a_non_string_tool_result_text_before_a_bridged_provider_applies_edits(
monkeypatch: pytest.MonkeyPatch,
) -> None:
response, sent_bodies = await _count_on_a_bridged_provider_count_api(
{"messages": _build_log_messages(5), "context_management": _CLEAR_TOOL_USES}, monkeypatch
)
assert isinstance(response, JSONResponse)
assert response.status_code == 400
assert "tool_result toolu_log content.1.text" in json.loads(bytes(response.body))["error"]["message"]
assert sent_bodies == ()
@pytest.mark.asyncio
async def test_messages_count_tokens_applies_a_responses_api_deployment_edits_once() -> None:
peer: Final = _CountTokensPeer({"object": "response.input_tokens", "input_tokens": 71})
router: Final = Router(
model_list=[{"model_name": "gpt-agent", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "test-key"}}]
)
with patch.object(
openai_count_tokens_http,
"get_async_httpx_client",
return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)),
):
response: Final = await _messages_endpoint_count(
router,
{
"model": "gpt-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": _CLEAR_TOOL_USES,
},
)
assert response == {"input_tokens": 71, "litellm_count_source": "provider"}
assert len(peer.requests) == 1
assert peer.requests[0].content.decode().count(_GEMINI_TOOL_OUTPUT) == 1
_DATABRICKS_CLAUDE_MODEL: Final = "databricks/databricks-claude-sonnet-4-6"
def _databricks_claude_router() -> Router:
return Router(
model_list=[
{
"model_name": "databricks-claude-agent",
"litellm_params": {
"model": _DATABRICKS_CLAUDE_MODEL,
"api_key": "test-key",
"api_base": "https://example.cloud.databricks.com/serving-endpoints",
},
}
]
)
@pytest.mark.asyncio
async def test_messages_count_tokens_counts_the_edited_body_on_a_native_messages_deployment_without_a_count_api() -> (
None
):
router: Final = _databricks_claude_router()
unedited_body: Final[dict[str, JsonValue]] = {
"model": "databricks-claude-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
}
edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES})
unedited: Final = await _input_tokens(router, unedited_body)
polyfilled: Final = await _input_tokens(
router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(_DATABRICKS_CLAUDE_MODEL)}
)
assert edited < unedited
assert edited == polyfilled
@pytest.mark.asyncio
@pytest.mark.parametrize("drop_params", [False, True])
async def test_messages_count_tokens_on_a_native_messages_deployment_without_a_count_api_rejects_an_invalid_edit(
drop_params: bool,
) -> None:
with patch.object(litellm, "drop_params", drop_params):
response: Final = await _messages_endpoint_count(
_databricks_claude_router(),
{
"model": "databricks-claude-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": _INVALID_COMPACT_EDIT,
},
)
assert isinstance(response, JSONResponse)
assert response.status_code == 400
@pytest.mark.asyncio
async def test_messages_count_tokens_leaves_the_edits_to_a_native_messages_deployment() -> None:
peer: Final = _CountTokensPeer({"input_tokens": 63})
router: Final = Router(
model_list=[
{
"model_name": "claude-agent",
"litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "test-key"},
}
]
)
with patch.object(
anthropic_count_tokens_handler,
"get_async_httpx_client",
return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)),
):
response: Final = await _messages_endpoint_count(
router,
{
"model": "claude-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": _CLEAR_TOOL_USES,
},
)
assert response == {"input_tokens": 63, "litellm_count_source": "provider"}
assert len(peer.requests) == 1
assert peer.requests[0].content.decode().count(_GEMINI_TOOL_OUTPUT) == 3
@pytest.mark.asyncio
@pytest.mark.parametrize(
"text", ["hello there", "The quick brown fox jumps over the lazy dog. " * 50], ids=["short", "long"]

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