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 origin/main into litellm_responses_bridge_logprobs

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo 2026-11-02 17:36:32 +00:00
commit 94a46a07ce
126 changed files with 8581 additions and 2870 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
AGENTS.md
View file

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

2
README.md
View file

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

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

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

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

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

5
litellm/__init__.py
View file

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

175
litellm/anthropic_beta_headers_config.json
View file

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

66
litellm/guardrails.py Normal file
View file

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

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

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:

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,

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

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

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

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

15
litellm/model_prices_and_context_window_backup.json
View file

@ -15555,18 +15555,20 @@
"input_cost_per_token": 2.65e-06,
"litellm_provider": "bedrock",
"max_input_tokens": 8000,
"max_output_tokens": 2048,
"max_output_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 3.5e-06,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-70b-instruct.md",
"supports_pdf_input": true
},
"bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0": {
"input_cost_per_token": 3e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 8000,
"max_output_tokens": 2048,
"max_output_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-8b-instruct.md",
"supports_pdf_input": true
},
"bedrock/us-west-1/meta.llama3-70b-instruct-v1:0": {
@ -48955,6 +48957,7 @@
},
"together_ai/Qwen/Qwen3.8-2.4T-A95B": {
"cache_read_input_token_cost": 2.5e-07,
"deprecation_date": "2026-11-04",
"input_cost_per_token": 2e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 1010000,
@ -53954,23 +53957,23 @@
"supports_vision": true
},
"wandb/meta-llama/Llama-3.1-8B-Instruct": {
"max_input_tokens": 131000,
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"input_cost_per_token": 2.2e-07,
"output_cost_per_token": 2.2e-07,
"litellm_provider": "wandb",
"mode": "chat",
"source": "https://wandb.ai/site/pricing/tokens/"
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference"
},
"wandb/deepseek-ai/DeepSeek-V3.1": {
"supports_reasoning": true,
"max_input_tokens": 161000,
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"input_cost_per_token": 5.5e-07,
"output_cost_per_token": 1.65e-06,
"litellm_provider": "wandb",
"mode": "chat",
"source": "https://wandb.ai/site/pricing/tokens/"
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference"
},
"wandb/meta-llama/Llama-3.3-70B-Instruct": {
"max_input_tokens": 128000,

82
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -15603,6 +15603,10 @@
],
"title": "Metadata"
},
"response_format": {
"$ref": "#/components/schemas/GuardrailResponseFormat",
"default": "standard"
},
"text": {
"title": "Text",
"type": "string"
@ -16480,6 +16484,14 @@
"title": "BedrockChecksSensitiveInformationModel",
"type": "object"
},
"BedrockGuardrailOutputScope": {
"enum": [
"INTERVENTIONS",
"FULL"
],
"title": "BedrockGuardrailOutputScope",
"type": "string"
},
"BlockedWord": {
"description": "Represents a blocked word with its action and optional description",
"properties": {
@ -17023,6 +17035,14 @@
"title": "GuardrailInfoResponse",
"type": "object"
},
"GuardrailResponseFormat": {
"enum": [
"standard",
"native"
],
"title": "GuardrailResponseFormat",
"type": "string"
},
"GuardrailSubmissionItem": {
"properties": {
"created_at": {
@ -17275,6 +17295,7 @@
"title": "HTTPValidationError",
"type": "object"
},
"JsonValue": {},
"LakeraCategoryThresholds": {
"additionalProperties": true,
"properties": {
@ -18201,6 +18222,18 @@
"description": "Whether to fail the request if the guardrail encounters an error. Implemented by guardrail='model_armor', 'generic_guardrail_api' and 'crowdstrike_aidr'. True (default) raises the error. False logs a critical error and lets the request proceed, so only a valid guardrail response can block or modify it.",
"title": "Fail On Error"
},
"file_checkpoint_id": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "File-checkpoint ID for the Ovalix Tracker service (falls back to pre/post).",
"title": "File Checkpoint Id"
},
"file_guardrail_timeout": {
"anyOf": [
{
@ -18703,6 +18736,17 @@
],
"description": "Optional parameters for the guardrail"
},
"outputScope": {
"anyOf": [
{
"$ref": "#/components/schemas/BedrockGuardrailOutputScope"
},
{
"type": "null"
}
],
"description": "ApplyGuardrail response scope. FULL includes detected and non-detected findings. Omit to retain Bedrock's default of returning detected findings only."
},
"output_parse_pii": {
"anyOf": [
{
@ -20585,7 +20629,24 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/ApplyGuardrailResponse"
"anyOf": [
{
"$ref": "#/components/schemas/ApplyGuardrailResponse"
},
{
"additionalProperties": {
"$ref": "#/components/schemas/JsonValue"
},
"type": "object"
},
{
"items": {
"$ref": "#/components/schemas/JsonValue"
},
"type": "array"
}
],
"title": "Response Apply Guardrail Apply Guardrail Post"
}
}
},
@ -20677,7 +20738,24 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/ApplyGuardrailResponse"
"anyOf": [
{
"$ref": "#/components/schemas/ApplyGuardrailResponse"
},
{
"additionalProperties": {
"$ref": "#/components/schemas/JsonValue"
},
"type": "object"
},
{
"items": {
"$ref": "#/components/schemas/JsonValue"
},
"type": "array"
}
],
"title": "Response Apply Guardrail Guardrails Apply Guardrail Post"
}
}
},

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

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

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

@ -12,7 +12,6 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM
import json
from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping
from copy import deepcopy
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
from urllib.parse import urlparse
@ -29,6 +28,7 @@ from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.a2a.agent_card import point_agent_card_at_proxy
from litellm.proxy.a2a.version_convert import (
A2AVersion,
normalize_agent_card,
@ -750,16 +750,12 @@ async def get_agent_card(
)
proxy_url: Final = get_custom_url(str(request.base_url), route=f"a2a/{agent_id}")
agent_card = deepcopy(agent.agent_card_params)
agent_card["url"] = proxy_url
interfaces: Final = agent_card.get("supportedInterfaces")
if isinstance(interfaces, list) and interfaces:
interfaces[0]["url"] = proxy_url
agent_card: Final = point_agent_card_at_proxy(agent.agent_card_params, proxy_url)
served_version: Final = _served_version(agent, request)
agent_card = normalize_agent_card(agent_card, served_version)
normalized_agent_card: Final = normalize_agent_card(agent_card, served_version)
verbose_proxy_logger.debug("Returning agent card for '%s' with proxy URL: %s", agent_id, proxy_url)
return JSONResponse(content=agent_card)
return JSONResponse(content=normalized_agent_card)
except HTTPException:
raise

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

25
litellm/proxy/proxy_cli.py
View file

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

2
litellm/proxy/proxy_server.py
View file

@ -1210,7 +1210,7 @@ ui_message += "\n\n💸 [```LiteLLM Model Cost Map```](https://models.litellm.ai
ui_message += f"\n\n🔎 [```LiteLLM Model Hub```]({model_hub_link}). See available models on the proxy. [**Docs**](https://docs.litellm.ai/docs/proxy/ai_hub)"
custom_swagger_message: Final = (
"[**Customize Swagger Docs**](https://docs.litellm.ai/docs/proxy/enterprise#swagger-docs---custom-routes--branding)"
"[**Customize Swagger Docs**](https://docs.litellm.ai/docs/enterprise#operations--branding)"
)
### CUSTOM BRANDING [ENTERPRISE FEATURE] ###

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,
@ -247,14 +251,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))
@ -276,12 +277,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()),
)
)
@ -354,12 +353,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,
)
)
@ -502,15 +499,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:
@ -575,33 +569,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 {}),
@ -633,16 +621,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()),
)
)
@ -778,31 +761,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(
@ -815,23 +791,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:
@ -842,7 +813,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(),
@ -856,7 +826,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:
@ -893,12 +862,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.
@ -952,7 +916,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self,
reasoning_item_id: str,
reasoning_content: str,
sequence_number: int,
) -> ReasoningSummaryTextDoneEvent:
"""
Create response.reasoning_summary_text.done event.
@ -962,7 +925,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."
}
@ -971,7 +933,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,
)
@ -980,7 +942,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self,
reasoning_item_id: str,
reasoning_content: str,
sequence_number: int,
) -> ReasoningSummaryPartDoneEvent:
"""
Create response.reasoning_summary_part.done event.
@ -994,7 +955,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
}
"""
@ -1002,7 +962,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(
**{
@ -1094,7 +1054,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,
@ -1110,7 +1069,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.
@ -1119,7 +1077,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
{
"type": "response.output_item.done",
"output_index": 0,
"sequence_number": 99,
"item": {
"id": "rs_0c5dae30e53172980069708ba2f59c8197b71ca9820edad07c",
"type": "reasoning",
@ -1135,7 +1092,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,
@ -1251,7 +1207,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
@ -1261,20 +1216,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
@ -1307,30 +1262,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:
@ -1391,6 +1350,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:
@ -1532,11 +1496,10 @@ 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
delta_logprobs: Final = output_text_logprobs(getattr(chunk.choices[0], "logprobs", None))
if delta_logprobs is not None:
self._text_logprob_chunks.append(delta_logprobs)
text_delta_event: Final = OutputTextDeltaEvent(
return OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id=item_id,
output_index=self._message_index(),
@ -1544,8 +1507,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
delta=delta_content,
logprobs=delta_logprobs,
)
text_delta_event.__dict__["sequence_number"] = self._sequence_number
return text_delta_event
# Priority 3: Handle tool call deltas (if any) -> queue events and emit them
# For each tool call delta, we emit events one at a time to match OpenAI's streaming behavior

12
litellm/responses/streaming_iterator.py
View file

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

15
litellm/responses/utils.py
View file

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

19
litellm/router_python_core.py
View file

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

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

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

23
litellm/types/guardrails.py
View file

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

100
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
@ -1759,17 +1760,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)
@ -1783,12 +1807,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
@ -1796,7 +1820,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
@ -1804,7 +1828,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
@ -1812,7 +1836,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
@ -1821,7 +1845,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
@ -1830,13 +1854,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
@ -1856,7 +1880,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
@ -1884,7 +1908,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
@ -1892,7 +1916,7 @@ class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject):
part: PART_UNION_TYPES
class OutputTextDeltaEvent(CarriesOutputTextLogprobs):
class OutputTextDeltaEvent(CarriesOutputTextLogprobs, BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA]
item_id: str
output_index: int
@ -1900,7 +1924,7 @@ class OutputTextDeltaEvent(CarriesOutputTextLogprobs):
delta: str
class OutputTextAnnotationAddedEvent(BaseLiteLLMOpenAIResponseObject):
class OutputTextAnnotationAddedEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED]
item_id: str
output_index: int
@ -1909,7 +1933,7 @@ class OutputTextAnnotationAddedEvent(BaseLiteLLMOpenAIResponseObject):
annotation: dict
class OutputTextDoneEvent(CarriesOutputTextLogprobs):
class OutputTextDoneEvent(CarriesOutputTextLogprobs, BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE]
item_id: str
output_index: int
@ -1917,7 +1941,7 @@ class OutputTextDoneEvent(CarriesOutputTextLogprobs):
text: str
class RefusalDeltaEvent(BaseLiteLLMOpenAIResponseObject):
class RefusalDeltaEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.REFUSAL_DELTA]
item_id: str
output_index: int
@ -1925,7 +1949,7 @@ class RefusalDeltaEvent(BaseLiteLLMOpenAIResponseObject):
delta: str
class RefusalDoneEvent(BaseLiteLLMOpenAIResponseObject):
class RefusalDoneEvent(BaseResponsesStreamEvent):
type: Literal[ResponsesAPIStreamEvents.REFUSAL_DONE]
item_id: str
output_index: int
@ -1933,72 +1957,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
@ -2006,14 +2030,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
@ -2021,7 +2045,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
@ -2029,21 +2053,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
@ -2058,7 +2082,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
@ -2067,7 +2091,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=())

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

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

15
model_prices_and_context_window.json
View file

@ -15555,18 +15555,20 @@
"input_cost_per_token": 2.65e-06,
"litellm_provider": "bedrock",
"max_input_tokens": 8000,
"max_output_tokens": 2048,
"max_output_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 3.5e-06,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-70b-instruct.md",
"supports_pdf_input": true
},
"bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0": {
"input_cost_per_token": 3e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 8000,
"max_output_tokens": 2048,
"max_output_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-8b-instruct.md",
"supports_pdf_input": true
},
"bedrock/us-west-1/meta.llama3-70b-instruct-v1:0": {
@ -48955,6 +48957,7 @@
},
"together_ai/Qwen/Qwen3.8-2.4T-A95B": {
"cache_read_input_token_cost": 2.5e-07,
"deprecation_date": "2026-11-04",
"input_cost_per_token": 2e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 1010000,
@ -53954,23 +53957,23 @@
"supports_vision": true
},
"wandb/meta-llama/Llama-3.1-8B-Instruct": {
"max_input_tokens": 131000,
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"input_cost_per_token": 2.2e-07,
"output_cost_per_token": 2.2e-07,
"litellm_provider": "wandb",
"mode": "chat",
"source": "https://wandb.ai/site/pricing/tokens/"
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference"
},
"wandb/deepseek-ai/DeepSeek-V3.1": {
"supports_reasoning": true,
"max_input_tokens": 161000,
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"input_cost_per_token": 5.5e-07,
"output_cost_per_token": 1.65e-06,
"litellm_provider": "wandb",
"mode": "chat",
"source": "https://wandb.ai/site/pricing/tokens/"
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference"
},
"wandb/meta-llama/Llama-3.3-70B-Instruct": {
"max_input_tokens": 128000,

13
tests/e2e/CONTRIBUTING.md
View file

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

285
tests/e2e/logging/weave_reader.py
View file

@ -1,285 +0,0 @@
"""Read-back for the Weave (Weights & Biases) logging tests against the real
Weave project.
The proxy ships OTEL spans to https://trace.wandb.ai/otel/v1/traces with the
``weave_otel`` callback, and the tests read the ingested calls back through
Weave's own query API (``POST /calls/stream_query``), which answers JSON Lines:
one JSON object per call, so the body is parsed line by line rather than as one
document.
The project is shared with other traffic, so the read never relies on the target
being among the newest N calls: the query is scoped server-side to the
``litellm_request`` op and to calls that started after the test's own request,
and pages with ``offset`` until the window is exhausted.
Weave's own ``summary.weave.status`` is a rollup that reads "success" even for a
span the exporter marked failed, so status comes from the OTEL span itself
(``attributes.otel_span.status.code``), and the shipped cost from
``attributes.otel_span.attributes.llm.response.cost`` - the StandardLogging
``response_cost``, which is what makes this a spend assertion rather than a
delivery ping.
Missing configuration is a hard failure, never a skip.
"""
from __future__ import annotations
import base64
import json
import os
import time
from dataclasses import dataclass, field
from itertools import count, takewhile
from typing import Final
import pytest
from pydantic import BaseModel, ConfigDict, Field
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT
from e2e_http import URL, AuthHeaders, send
from e2e_metadata import step
_WEAVE_TRACE_API: Final = "https://trace.wandb.ai"
#: The op every litellm LLM call lands under. The proxy also exports a root
#: server span ("Received Proxy Server Request") and management spans; only the
#: LLM call carries the usage and cost this suite asserts on.
LITELLM_REQUEST_OP: Final = "litellm_request"
#: How long to keep re-reading after the first matching call before trusting the
#: exactly-one assertion. The OTEL batch exporter flushes on its own schedule, so
#: a duplicate export can surface well after the first one, and a duplicate IS
#: the bug being guarded against.
WEAVE_SETTLE_SECONDS: Final = 45.0
#: Rows per page. The query is already scoped to this run's time window, so this
#: only bounds one round trip, not what the read can see.
_PAGE_SIZE: Final = 500
class _WeaveSortBy(BaseModel):
field: str
direction: str
class _WeaveOpFilter(BaseModel):
op_names: list[str]
class _WeaveGetField(BaseModel):
get_field: str = Field(serialization_alias="$getField")
class _WeaveLiteral(BaseModel):
literal: float = Field(serialization_alias="$literal")
class _WeaveGreaterThan(BaseModel):
gt: tuple[_WeaveGetField, _WeaveLiteral] = Field(serialization_alias="$gt")
class _WeaveQuery(BaseModel):
expr: _WeaveGreaterThan = Field(serialization_alias="$expr")
class _WeaveQueryBody(BaseModel):
project_id: str
filter: _WeaveOpFilter
query: _WeaveQuery
limit: int = _PAGE_SIZE
offset: int = 0
sort_by: list[_WeaveSortBy] = [_WeaveSortBy(field="started_at", direction="asc")]
class _OtelStatus(BaseModel):
model_config = ConfigDict(extra="ignore")
code: str | None = None
message: str | None = None
class _OtelError(BaseModel):
model_config = ConfigDict(extra="ignore")
code: str | None = None
type: str | None = None
message: str | None = None
class _LlmResponse(BaseModel):
model_config = ConfigDict(extra="ignore")
cost: float | None = None
class _LlmAttributes(BaseModel):
model_config = ConfigDict(extra="ignore")
response: _LlmResponse | None = None
class _OtelSpanAttributes(BaseModel):
model_config = ConfigDict(extra="ignore")
llm: _LlmAttributes | None = None
error: _OtelError | None = None
class _OtelSpan(BaseModel):
model_config = ConfigDict(extra="ignore")
name: str | None = None
status: _OtelStatus | None = None
attributes: _OtelSpanAttributes | None = None
class _CallAttributes(BaseModel):
model_config = ConfigDict(extra="ignore")
otel_span: _OtelSpan | None = None
class _Usage(BaseModel):
model_config = ConfigDict(extra="ignore")
total_tokens: int | None = None
class _WeaveSummary(BaseModel):
model_config = ConfigDict(extra="ignore")
usage: dict[str, _Usage] = {}
class WeaveCall(BaseModel):
"""One ingested Weave call, reduced to what the scenarios assert on."""
model_config = ConfigDict(extra="ignore")
id: str
op_name: str
started_at: str | None = None
inputs: dict[str, object] = {}
attributes: _CallAttributes | None = None
summary: _WeaveSummary | None = Field(default=None)
@property
def op(self) -> str:
"""The bare op name out of ``weave:///<entity>/<project>/op/<op>:<digest>``."""
return self.op_name.split("/op/")[-1].split(":")[0]
@property
def status_code(self) -> str | None:
"""The OTEL span status, not Weave's own rollup (which reads "success"
even for a span the exporter marked ERROR)."""
span = self.attributes.otel_span if self.attributes else None
return span.status.code if span and span.status else None
@property
def error(self) -> _OtelError | None:
span = self.attributes.otel_span if self.attributes else None
return span.attributes.error if span and span.attributes else None
@property
def response_cost(self) -> float | None:
span = self.attributes.otel_span if self.attributes else None
llm = span.attributes.llm if span and span.attributes else None
return llm.response.cost if llm and llm.response else None
@property
def total_tokens(self) -> int | None:
"""Weave keys usage by model, so the total is summed across whatever
models the call reported."""
if not self.summary or not self.summary.usage:
return None
totals = [usage.total_tokens for usage in self.summary.usage.values() if usage.total_tokens is not None]
return sum(totals) if totals else None
def mentions(self, needle: str) -> bool:
return needle in json.dumps(self.inputs, default=str)
@dataclass(frozen=True, slots=True)
class WeaveReader:
project_id: str
api_key: str = field(repr=False)
@property
def _headers(self) -> AuthHeaders:
"""Weave authenticates with HTTP Basic as the fixed user ``api``."""
token = base64.b64encode(f"api:{self.api_key}".encode()).decode()
return AuthHeaders(authorization=f"Basic {token}")
def _query_body(self, *, since: float, offset: int, op: str) -> _WeaveQueryBody:
return _WeaveQueryBody(
project_id=self.project_id,
filter=_WeaveOpFilter(op_names=[f"weave:///{self.project_id}/op/{op}:*"]),
query=_WeaveQuery(
expr=_WeaveGreaterThan(gt=(_WeaveGetField(get_field="started_at"), _WeaveLiteral(literal=since)))
),
offset=offset,
)
def _page(self, *, since: float, offset: int, op: str) -> tuple[WeaveCall, ...]:
outcome = send(
URL(f"{_WEAVE_TRACE_API}/calls/stream_query"),
headers=self._headers,
json=self._query_body(since=since, offset=offset, op=op),
)
if not outcome.ok:
pytest.fail(
f"Weave calls query for project {self.project_id!r} failed "
f"({outcome.status_code}): {outcome.body[:300]}"
)
return tuple(WeaveCall.model_validate_json(line) for line in outcome.body.splitlines() if line.strip())
@step("Read the Weave {op} calls carrying the marker {marker}")
def calls_matching(self, marker: str, *, since: float, op: str = LITELLM_REQUEST_OP) -> tuple[WeaveCall, ...]:
"""Every call under ``op`` started after ``since`` whose inputs carry
``marker``, paging until the window is exhausted.
More than one is the duplicate-delivery bug, so this never collapses to a
single call.
"""
pages = tuple(
takewhile(
bool,
(self._page(since=since, offset=offset, op=op) for offset in count(0, _PAGE_SIZE)),
)
)
return tuple(call for page in pages for call in page if call.mentions(marker))
@step("Wait for Weave to ingest a {op} call carrying the marker {marker}, then watch for duplicates")
def poll_calls_matching(self, marker: str, *, since: float, op: str = LITELLM_REQUEST_OP) -> tuple[WeaveCall, ...]:
"""Poll until the call is readable, then keep re-reading for
WEAVE_SETTLE_SECONDS so a duplicate exported by a later batch flush
cannot hide from the exactly-one assertion. A duplicate ends the settle
early, because more waiting cannot clear it."""
deadline = time.monotonic() + POLL_TIMEOUT
while time.monotonic() < deadline:
calls = self.calls_matching(marker, since=since, op=op)
if calls:
return self._settled(marker, since=since, op=op, first=calls)
time.sleep(POLL_INTERVAL)
return ()
def _settled(self, marker: str, *, since: float, op: str, first: tuple[WeaveCall, ...]) -> tuple[WeaveCall, ...]:
"""A transiently empty re-read never downgrades what was already seen."""
settle_deadline = time.monotonic() + WEAVE_SETTLE_SECONDS
latest = first # rebind-ok: one settle window, re-read per poll interval
while time.monotonic() < settle_deadline and len(latest) <= 1:
time.sleep(POLL_INTERVAL)
latest = self.calls_matching(marker, since=since, op=op) or latest
return latest
def build_weave_reader() -> WeaveReader:
project_id = (os.environ.get("WEAVE_PROJECT_ID") or os.environ.get("WANDB_PROJECT_ID") or "").strip()
api_key = os.environ.get("WANDB_API_KEY", "").strip()
if not project_id or not api_key:
pytest.fail(
"Weave e2e requires WANDB_API_KEY and WEAVE_PROJECT_ID (or WANDB_PROJECT_ID, "
"format <entity>/<project>): the test reads the proxy's weave_otel delivery "
"back from the real Weave project; missing credentials is a hard failure, not a skip"
)
return WeaveReader(project_id=project_id, api_key=api_key)

2
tests/e2e/models.py
View file

@ -45,8 +45,6 @@ class KeyLoggingCallbackVars(BaseModel):
langfuse_public_key: str | None = Field(default=None, repr=False)
langfuse_secret_key: str | None = Field(default=None, repr=False)
langfuse_host: str | None = Field(default=None, repr=False)
wandb_api_key: str | None = Field(default=None, repr=False)
weave_project_id: str | None = None
class KeyLoggingCallback(BaseModel):

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

42
tests/integration/_support/mcp.py
View file

@ -10,6 +10,7 @@ from dataclasses import dataclass, field
from pathlib import Path
from typing import Final, Literal
import anyio
import httpx
from integration._support.asgi import asgi_server
from integration._support.client import Gateway, Scenario
@ -63,7 +64,9 @@ class Confirmation(BaseModel):
confirmed: bool
def math_service(name: str = "integration-math", *, rich: bool = False) -> MCPServer:
def math_service(
name: str = "integration-math", *, rich: bool = False, on_slow: Callable[[], None] | None = None
) -> MCPServer:
service: Final = MCPServer(name)
service.add_tool(add)
service.add_tool(multiply)
@ -77,6 +80,8 @@ def math_service(name: str = "integration-math", *, rich: bool = False) -> MCPSe
@service.tool()
async def slow(seconds: float) -> str:
if on_slow is not None:
on_slow()
await asyncio.sleep(seconds)
return "slept"
@ -114,6 +119,29 @@ def math_service(name: str = "integration-math", *, rich: bool = False) -> MCPSe
return service
def _cancel_on_disconnect(app: ASGIApp) -> ASGIApp:
async def connected(scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
await app(scope, receive, send)
return
sender, reader = anyio.create_memory_object_stream[Message]()
async with sender, reader, anyio.create_task_group() as tasks:
async def relay() -> None:
while True:
message: Final = await receive()
if message["type"] == "http.disconnect":
tasks.cancel_scope.cancel()
return
await sender.send(message)
tasks.start_soon(relay)
await app(scope, reader.receive, send)
tasks.cancel_scope.cancel()
return connected
def _capturing(app: Callable[[Scope, Receive, Send], object], observed: queue.Queue[dict[str, object]]):
async def capture(scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
@ -247,8 +275,10 @@ def _draining_sse_watcher(app: Callable[[Scope, Receive, Send], object]):
@contextmanager
def mcp_peer(transport: Literal["http", "sse"] = "http", *, rich: bool = False) -> Iterator[McpPeer]:
service: Final = math_service(rich=rich)
def mcp_peer(
transport: Literal["http", "sse"] = "http", *, rich: bool = False, on_slow: Callable[[], None] | None = None
) -> Iterator[McpPeer]:
service: Final = math_service(rich=rich, on_slow=on_slow)
security: Final = TransportSecuritySettings(enable_dns_rebinding_protection=False)
app: Final = (
_draining_sse_watcher(service.sse_app(transport_security=security))
@ -256,7 +286,11 @@ def mcp_peer(transport: Literal["http", "sse"] = "http", *, rich: bool = False)
else service.streamable_http_app(stateless_http=True, json_response=True, transport_security=security)
)
observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
with asgi_server(_capturing(app, observed), before_stop=_drain_sse_streams if transport == "sse" else None) as url:
captured: Final = _capturing(app, observed)
with asgi_server(
_cancel_on_disconnect(captured) if transport == "http" else captured,
before_stop=_drain_sse_streams if transport == "sse" else None,
) as url:
yield McpPeer(url + ("/sse" if transport == "sse" else "/mcp"), observed, transport)

85
tests/integration/authorization/test_agent_access.py
View file

@ -66,6 +66,13 @@ def _peer_card(url: str, name: str) -> dict[str, JsonValue]:
"description": "Synthetic peer",
"version": "1.0.0",
"url": url + "/",
"supportedInterfaces": [
{
"url": url + "/",
"protocolBinding": "JSONRPC",
"protocolVersion": "0.3",
}
],
"capabilities": {"streaming": True},
"defaultInputModes": ["text"],
"defaultOutputModes": ["text"],
@ -73,10 +80,15 @@ def _peer_card(url: str, name: str) -> dict[str, JsonValue]:
}
def _register_agent(gateway: Gateway, scenario: Scenario, name: str, url: str) -> str:
created: Final = gateway.request(
"POST", "/v1/agents", {"agent_name": name, "agent_card_params": _peer_card(url, name)}
)
def _register_agent(
gateway: Gateway,
scenario: Scenario,
name: str,
url: str,
agent_card: dict[str, JsonValue] | None = None,
) -> str:
card: Final = _peer_card(url, name) if agent_card is None else agent_card
created: Final = gateway.request("POST", "/v1/agents", {"agent_name": name, "agent_card_params": card})
assert created.status_code == 200, created.text
identity: Final = created.json()["agent_id"]
@ -416,3 +428,68 @@ def test_legacy_agent_json_card_path_applies_agent_permissions_to_non_admin_keys
)
assert (denied_legacy.status_code, denied_legacy.text) == (403, refusal(denied)), denied_legacy.text
assert wire.drain() == ()
def test_agent_list_and_detail_cards_use_proxy_urls_for_scoped_keys(gateway: Gateway) -> None:
marker: Final = "a2acardurl" + uuid.uuid4().hex[:12]
def upstream(request: Request) -> Reply:
assert request.method == "GET", request.method
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
with wire_server(upstream) as wire, gateway.scenario() as scenario:
backend_url: Final = wire.url + "/"
peer_card: Final = _peer_card(wire.url, marker)
peer_interfaces: Final = peer_card["supportedInterfaces"]
assert isinstance(peer_interfaces, list)
agent_card: Final = {
**peer_card,
"additionalInterfaces": [{"url": backend_url, "transport": "GRPC"}],
"supportedInterfaces": [
*peer_interfaces,
{
"url": backend_url,
"protocolBinding": "GRPC",
"protocolVersion": "0.3",
},
],
}
agent_id: Final = _register_agent(gateway, scenario, marker, wire.url, agent_card)
user_id: Final = scenario.user(user_role="internal_user")
scoped_key: Final = scenario.key(user_id=user_id, object_permission={"agents": [agent_id]})
public_card: Final = gateway.request(
"GET",
f"/a2a/{agent_id}/.well-known/agent-card.json",
key=scoped_key,
)
listed: Final = gateway.request("GET", "/v1/agents", key=scoped_key)
detailed: Final = gateway.request("GET", f"/v1/agents/{agent_id}", key=scoped_key)
assert public_card.status_code == 200, public_card.text
assert listed.status_code == 200, listed.text
assert detailed.status_code == 200, detailed.text
assert wire.url not in listed.text
assert wire.url not in detailed.text
assert "additionalInterfaces" not in listed.text
assert "additionalInterfaces" not in detailed.text
public_urls: Final = (
public_card.json()["url"],
public_card.json()["supportedInterfaces"][0]["url"],
)
listed_agent: Final = next(item for item in listed.json() if item["agent_id"] == agent_id)
for body in (listed_agent, detailed.json()):
card: Final = body["agent_card_params"]
assert (card["url"], card["supportedInterfaces"][0]["url"]) == public_urls
assert card["supportedInterfaces"] == [
{
"url": public_urls[1],
"protocolBinding": "JSONRPC",
"protocolVersion": "0.3",
}
]
admin_listed: Final = gateway.request("GET", "/v1/agents")
assert admin_listed.status_code == 200, admin_listed.text
admin_agent: Final = next(item for item in admin_listed.json() if item["agent_id"] == agent_id)
assert admin_agent["agent_card_params"]["url"] == backend_url

53
tests/integration/compatibility/test_a2a_wire_versions.py
View file

@ -859,6 +859,59 @@ def test_a2a_message_send_aliases_and_agent_name_reach_the_same_agent_with_the_s
), posts
def test_a2a_data_part_skill_call_and_data_artifact_reach_each_side_unchanged(gateway: Gateway) -> None:
marker: Final = "a2adata" + uuid.uuid4().hex
invocation: Final[dict[str, JsonValue]] = {"skill": "search_items", "params": {"region": "emea", "limit": 3}}
found: Final[dict[str, JsonValue]] = {
"total": 2,
"count": 2,
"listings": [{"id": marker + "-1", "region": "emea"}, {"id": marker + "-2", "region": "emea"}],
}
task: Final[dict[str, JsonValue]] = {
"kind": "task",
"id": marker + "-task",
"contextId": marker + "-context",
"status": {"state": "completed"},
"artifacts": [
{"artifactId": marker + "-artifact", "name": "search_results", "parts": [{"kind": "data", "data": found}]}
],
}
def upstream(request: Request) -> Reply:
if request.method == "GET":
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
body: Final = json.loads(request.body)
return Reply(body=json.dumps({"jsonrpc": "2.0", "id": body["id"], "result": task}).encode())
with wire_server(upstream) as wire, gateway.scenario() as scenario:
identity: Final = _register_agent(
gateway, scenario, {"agent_name": marker, "agent_card_params": _peer_card(wire.url, marker)}
)
response: Final = gateway.client.post(
f"/a2a/{identity}",
headers={"Authorization": f"Bearer {gateway.key}"},
json={
"jsonrpc": "2.0",
"id": marker,
"method": "message/send",
"params": {
"message": {
"kind": "message",
"role": "user",
"messageId": marker + "-in",
"parts": [{"kind": "data", "data": invocation}],
}
},
},
)
assert response.status_code == 200, response.text
assert response.json() == {"jsonrpc": "2.0", "id": marker, "result": task}, response.text
posts: Final = tuple(json.loads(item.body) for item in wire.drain() if item.method == "POST")
assert len(posts) == 1, posts
assert posts[0]["method"] == "message/send", posts
assert posts[0]["params"]["message"]["parts"] == [{"kind": "data", "data": invocation}], posts
def test_a2a_task_methods_pass_through_verbatim_and_push_callbacks_are_validated_before_forwarding(
gateway: Gateway,
) -> None:

95
tests/integration/configuration/test_multi_worker_reload.py Normal file
View file

@ -0,0 +1,95 @@
import json
import os
import re
from pathlib import Path
from typing import Final
import psutil
import pytest
from tests.integration._support.client import Gateway, eventually, object_value
from tests.integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
def _worker_pids(owned: OwnedProxy) -> frozenset[int]:
return frozenset(int(match.group(1)) for match in STARTED_WORKER.finditer(owned.log.read_text()))
def _launched_config(owned: OwnedProxy) -> Path:
arguments: Final = psutil.Process(owned.process.pid).cmdline()
return Path(arguments[arguments.index("--config") + 1])
def _config(model: str) -> str:
return json.dumps(
{
"model_list": [{"model_name": model, "litellm_params": {"model": "openai/reload-test"}}],
"general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY"},
"router_settings": {"redis_host": "os.environ/REDIS_HOST", "redis_port": "os.environ/REDIS_PORT"},
}
)
def _assert_replaced(owned: OwnedProxy, previous: frozenset[int], workers: int) -> frozenset[int]:
started: Final = eventually(
lambda: _worker_pids(owned),
lambda pids: len(pids - previous) == workers,
seconds=graceful_stop_seconds(),
)
assert all(not psutil.pid_exists(pid) for pid in previous), owned.log.read_text()
assert owned.process.poll() is None, owned.log.read_text()
assert len(started) == len(previous) + workers, owned.log.read_text()
return started
def _assert_models(gateway: Gateway) -> None:
data: Final = gateway.get("/v1/models")["data"]
assert isinstance(data, list), data
assert tuple(object_value(model)["id"] for model in data) == ("after-reload",)
@pytest.mark.parametrize("workers", (1, 2))
@pytest.mark.timeout(6 * graceful_stop_seconds())
def test_reload_replaces_every_worker_and_shutdown_leaves_none(
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, workers: int
) -> None:
root: Final = Path(__file__).resolve().parents[3]
monkeypatch.setenv("INTEGRATION_PROXY_ROOT", str(tmp_path))
config: Final = tmp_path / "config.yaml"
source: Final = tmp_path / "reload_source.py"
env: Final = tmp_path / ".env"
config.write_text(_config("before-reload"))
source.write_text("VALUE = 1\n")
env.write_text("RELOAD_TEST_VALUE=before\n")
with owned_proxy_process(
gateway,
tmp_path,
{
"PYTHONPATH": os.pathsep.join((str(root), str(root / "tests"), os.environ.get("PYTHONPATH", ""))),
"STORE_MODEL_IN_DB": "False",
},
config=config,
remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "PYTHON_DOTENV_DISABLED"),
workers=workers,
database_setup=(),
extra_arguments=("--reload",),
) as owned:
initial: Final = eventually(
lambda: _worker_pids(owned), lambda pids: len(pids) == workers, seconds=graceful_stop_seconds()
)
_launched_config(owned).write_text(_config("after-reload"))
after_config: Final = _assert_replaced(owned, initial, workers)
for _ in range(2 * workers):
_assert_models(owned.gateway)
source.write_text("VALUE = 2\n")
after_source: Final = _assert_replaced(owned, after_config, workers)
env.write_text("RELOAD_TEST_VALUE=after\n")
after_env: Final = _assert_replaced(owned, after_source, workers)
response: Final = owned.gateway.request("GET", "/health/readiness")
assert response.status_code == 200, response.text
owned.begin_shutdown()
owned.process.wait(timeout=graceful_stop_seconds())
assert all(not psutil.pid_exists(pid) for pid in after_env), owned.log.read_text()
assert owned.process.returncode == 0, owned.log.read_text()

246
tests/integration/configuration/test_reload_worker_pool.py Normal file
View file

@ -0,0 +1,246 @@
import multiprocessing
import re
import signal
import socket
import subprocess
import sys
import threading
import time
from collections.abc import Callable
from pathlib import Path
from typing import Final
import httpx
import psutil
import pytest
from litellm.proxy.dev_reload import multi_worker_reload_supported, stop_pool_when_reloader_exits
from tests.integration._support.client import eventually
from tests.integration._support.ports import allocate_port
from tests.integration._support.process import group_members, kill_stray_process, signal_group, stop_processes
POOL_WORKERS: Final = 2
POOL_TIMEOUT_SECONDS: Final = 120
WORKER_HEALTHCHECK_SECONDS: Final = 30
WORKER_STARTED: Final = re.compile(r"Started server process \[(\d+)\]")
POOL_STOPPED: Final = "failed to start, stopping the parent process"
APP_LOAD_ERROR: Final = "Error loading ASGI app"
PID_APP: Final = """import os
async def app(scope, receive, send):
if scope["type"] != "http":
return
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": str(os.getpid()).encode()})
"""
REQUIRES_UVICORN_STARTUP_FAILURE: Final = pytest.mark.skipif(
not multi_worker_reload_supported(), reason="uvicorn<0.50.0 has no startup failure exit code"
)
CORE_DUMP_SECONDS: Final = 15
SLOW_CORE_DUMP_ON_QUIT: Final = f"""
import os
import signal
import subprocess
def freeze_until_killed(signum, frame):
subprocess.Popen(("sh", "-c", f"sleep {CORE_DUMP_SECONDS}; kill -9 {{os.getpid()}}"))
os.kill(os.getpid(), signal.SIGSTOP)
signal.signal(signal.SIGQUIT, freeze_until_killed)
"""
POOL_RUNNER: Final = """import sys
sys.path.insert(0, sys.argv[1])
from litellm.proxy.dev_reload import run_multi_worker_reload
run_multi_worker_reload(
{
"app": "reload_pool_app:app",
"host": "127.0.0.1",
"port": int(sys.argv[2]),
"timeout_worker_healthcheck": int(sys.argv[3]),
"reload": True,
"reload_dirs": (sys.argv[1],),
"reload_includes": ("*.py",),
},
workers=int(sys.argv[4]),
)
"""
def _started_workers(log_text: str) -> frozenset[int]:
return frozenset(int(match.group(1)) for match in WORKER_STARTED.finditer(log_text))
def _until(done: Callable[[], bool], log: Path, seconds: float = POOL_TIMEOUT_SECONDS) -> None:
eventually(lambda: (done(), log.read_text()), lambda observed: observed[0], seconds=seconds)
def _wait_for_log(log: Path, process: subprocess.Popen[bytes], seen: Callable[[str], bool]) -> str:
log_text, _ = eventually(
lambda: (log.read_text(), process.poll()),
lambda observed: seen(observed[0]) or observed[1] is not None,
seconds=POOL_TIMEOUT_SECONDS,
)
assert seen(log_text), log_text
return log_text
def _start_pool(app_dir: Path, port: int, log: Path, runner: str = POOL_RUNNER) -> subprocess.Popen[bytes]:
with log.open("wb") as output:
return subprocess.Popen(
(
sys.executable,
"-I",
"-c",
runner,
str(app_dir),
str(port),
str(WORKER_HEALTHCHECK_SECONDS),
str(POOL_WORKERS),
),
stdout=output,
stderr=subprocess.STDOUT,
start_new_session=True,
)
def _port_refuses_connections(port: int) -> bool:
with socket.socket() as sock:
return sock.connect_ex(("127.0.0.1", port)) != 0
def _next_pool(log: Path, process: subprocess.Popen[bytes], port: int, previous: frozenset[int]) -> frozenset[int]:
log_text: Final = _wait_for_log(log, process, lambda text: len(_started_workers(text) - previous) == POOL_WORKERS)
pool: Final = _started_workers(log_text) - previous
assert not any(psutil.pid_exists(pid) for pid in previous), log_text
served: Final = httpx.get(f"http://127.0.0.1:{port}/", timeout=POOL_TIMEOUT_SECONDS)
assert int(served.text) in pool, log_text
return pool
def test_the_pool_is_told_to_stop_once_the_reloader_process_exits() -> None:
reloader: Final = multiprocessing.get_context("spawn").Process(target=time.sleep, args=(POOL_TIMEOUT_SECONDS,))
reloader.start()
pool_should_exit: Final = threading.Event()
threading.Thread(target=stop_pool_when_reloader_exits, args=(reloader, pool_should_exit), daemon=True).start()
try:
assert not pool_should_exit.wait(0.5)
kill_stray_process(reloader.pid)
assert pool_should_exit.wait(POOL_TIMEOUT_SECONDS)
finally:
if reloader.is_alive():
kill_stray_process(reloader.pid)
reloader.join()
@REQUIRES_UVICORN_STARTUP_FAILURE
@pytest.mark.timeout(5 * POOL_TIMEOUT_SECONDS)
def test_a_save_replaces_every_worker_and_a_broken_save_stops_the_pool_until_the_next_good_one(
tmp_path: Path,
) -> None:
app: Final = tmp_path / "reload_pool_app.py"
app.write_text(PID_APP)
log: Final = tmp_path / "reload.log"
port: Final = allocate_port()
process: Final = _start_pool(tmp_path, port, log)
try:
first_pool: Final = _next_pool(log, process, port, frozenset())
app.write_text("async def app(\n")
_wait_for_log(log, process, lambda text: POOL_STOPPED in text)
assert not any(psutil.pid_exists(pid) for pid in first_pool), log.read_text()
app.write_text(PID_APP)
last_pool: Final = _next_pool(log, process, port, first_pool)
signal_group(process.pid, signal.SIGINT)
assert process.wait(timeout=POOL_TIMEOUT_SECONDS) == 0, log.read_text()
assert not any(psutil.pid_exists(pid) for pid in last_pool), log.read_text()
assert 1 <= log.read_text().count(APP_LOAD_ERROR) <= POOL_WORKERS, log.read_text()
finally:
stop_processes(group_members(process.pid))
process.wait()
@REQUIRES_UVICORN_STARTUP_FAILURE
@pytest.mark.timeout(3 * POOL_TIMEOUT_SECONDS)
@pytest.mark.parametrize(
("kill_reloader", "slow_core_dump", "shutdown_timeout"),
(
pytest.param(kill_stray_process, False, POOL_TIMEOUT_SECONDS, id="sigkill-reloader"),
pytest.param(
lambda reloader: signal_group(reloader, signal.SIGHUP), False, POOL_TIMEOUT_SECONDS, id="terminal-closed"
),
pytest.param(lambda reloader: signal_group(reloader, signal.SIGQUIT), True, 10, id="ctrl-backslash"),
),
)
def test_the_pool_stops_and_frees_the_port_when_the_reloader_dies(
tmp_path: Path, kill_reloader: Callable[[int], None], slow_core_dump: bool, shutdown_timeout: float
) -> None:
app_source: Final = PID_APP + (SLOW_CORE_DUMP_ON_QUIT if slow_core_dump else "")
(tmp_path / "reload_pool_app.py").write_text(app_source)
log: Final = tmp_path / "reload.log"
port: Final = allocate_port()
process: Final = _start_pool(tmp_path, port, log)
try:
workers: Final = _next_pool(log, process, port, frozenset())
pool_parent: Final = psutil.Process(min(workers)).ppid()
assert psutil.Process(pool_parent).ppid() == process.pid, log.read_text()
kill_reloader(process.pid)
process.wait(timeout=POOL_TIMEOUT_SECONDS)
_until(
lambda: not psutil.pid_exists(pool_parent) and _port_refuses_connections(port),
log,
seconds=shutdown_timeout,
)
finally:
stop_processes(group_members(process.pid))
process.wait()
@REQUIRES_UVICORN_STARTUP_FAILURE
@pytest.mark.timeout(3 * POOL_TIMEOUT_SECONDS)
def test_a_killed_pool_parent_takes_its_workers_with_it_so_the_next_save_serves_only_new_workers(
tmp_path: Path,
) -> None:
app: Final = tmp_path / "reload_pool_app.py"
app.write_text(PID_APP)
log: Final = tmp_path / "reload.log"
port: Final = allocate_port()
process: Final = _start_pool(tmp_path, port, log)
try:
orphans: Final = _next_pool(log, process, port, frozenset())
kill_stray_process(psutil.Process(min(orphans)).ppid())
_until(lambda: not any(psutil.pid_exists(pid) for pid in orphans), log)
app.write_text(f"{PID_APP}\n")
_next_pool(log, process, port, orphans)
finally:
stop_processes(group_members(process.pid))
process.wait()
@REQUIRES_UVICORN_STARTUP_FAILURE
@pytest.mark.timeout(3 * POOL_TIMEOUT_SECONDS)
def test_ctrl_backslash_frees_the_port_without_waiting_out_the_worker_healthcheck_when_core_dumps_are_slow(
tmp_path: Path,
) -> None:
(tmp_path / "reload_pool_app.py").write_text(PID_APP + SLOW_CORE_DUMP_ON_QUIT)
log: Final = tmp_path / "reload.log"
port: Final = allocate_port()
process: Final = _start_pool(tmp_path, port, log)
try:
workers: Final = _next_pool(log, process, port, frozenset())
pool_parent: Final = psutil.Process(min(workers)).ppid()
signal_group(process.pid, signal.SIGQUIT)
_until(
lambda: (
process.poll() is not None and not psutil.pid_exists(pool_parent) and _port_refuses_connections(port)
),
log,
seconds=WORKER_HEALTHCHECK_SECONDS / 2,
)
finally:
stop_processes(group_members(process.pid))
process.wait()

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

30
tests/integration/mcp/test_mcp_lifecycle.py
View file

@ -6,6 +6,7 @@ from collections.abc import Mapping
from contextlib import ExitStack
from hashlib import sha256
from pathlib import Path
from threading import Event
from typing import Final, cast
import httpx
@ -24,6 +25,7 @@ from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_req
from integration._support.mcp import (
JsonRpc,
McpCaller,
McpPeer,
Outcome,
ScriptedTool,
call_tool,
@ -46,6 +48,34 @@ _OBJECTS: Final = TypeAdapter(Mapping[str, object])
_STRINGS: Final = TypeAdapter(Mapping[str, str])
@pytest.mark.timeout(30)
def test_owned_http_peer_cleans_up_a_disconnected_tool_call() -> None:
started: Final = Event()
async def hang_up(peer: McpPeer) -> None:
async with (
streamable_http_client(peer.url, terminate_on_close=False) as streams,
ClientSession(streams[0], streams[1]) as session,
):
await session.initialize()
peer.drain()
async with asyncio.TaskGroup() as tasks:
running: Final = tasks.create_task(session.call_tool("slow", {"seconds": 60}))
received: Final = await asyncio.to_thread(
eventually, lambda: tool_calls(peer.drain()), bool, seconds=30
)
assert len(received) == 1, received
params: Final = _OBJECTS.validate_python(_OBJECTS.validate_python(received[0]["body"])["params"])
assert (params["name"], params["arguments"]) == ("slow", {"seconds": 60}), received
assert await asyncio.to_thread(started.wait, 5), "The MCP tool handler never started"
assert not running.done(), running.result().model_dump_json()
running.cancel()
assert running.cancelled()
with mcp_peer(rich=True, on_slow=started.set) as peer:
asyncio.run(hang_up(peer))
@pytest.mark.covers("mcp.call_tool.saved_headers.reach_actual_transport")
def test_saved_headers_reach_real_mcp_tool_and_survive_unrelated_edit(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:

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

338
tests/integration/observability/test_ovalix_guardrail.py Normal file
View file

@ -0,0 +1,338 @@
"""Ovalix guardrail driven through a real proxy.
The Ovalix tracker is the only guardrail double: an owned wire peer that answers the `/tracking/beta` routes
the way the live tracker does, including allowing a checkpoint call whose `actor` is empty without inspecting
it. The provider is a second owned peer. The proxy, Postgres and Redis run for real with two workers.
"""
import json
import threading
import uuid
from collections.abc import Iterator
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from itertools import product
from pathlib import Path
from typing import Final
from urllib.parse import urlsplit
import httpx
import pytest
import yaml
from integration._support.client import Gateway, JsonValue, gateway_from_environment
from integration._support.process import owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
pytestmark: Final = pytest.mark.shares_worker_state
API_KEY: Final = "synthetic-ovalix-key"
APP_ID: Final = "synthetic-ovalix-app"
OTHER_APP: Final = "Synthetic Other App"
PRE_CHECKPOINT: Final = "synthetic-pre-checkpoint"
POST_CHECKPOINT: Final = "synthetic-post-checkpoint"
KEYWORD: Final = "Project Falcon"
NAME: Final = "David"
DENIAL: Final = "synthetic ovalix denial"
REPLY: Final = "synthetic provider reply"
PINNED: Final = "ovalix-pinned"
ENDPOINTS: Final = ("/v1/chat/completions", "/v1/messages", "/v1/responses")
CHECKPOINT: Final = "/tracking/beta/checkpoint"
OUTAGE_AFTER: Final = 6
class _Outage:
"""Takes the tracker double down once it has answered a set number of checkpoint calls."""
def __init__(self) -> None:
self._lock: Final = threading.Lock()
self._remaining: int | None = None
def arm(self, answered: int) -> None:
with self._lock:
self._remaining = answered
def clear(self) -> None:
with self._lock:
self._remaining = None
def admits(self, path: str) -> bool:
with self._lock:
if self._remaining is None:
return True
if self._remaining == 0:
return False
if path == CHECKPOINT:
self._remaining -= 1
return True
OUTAGE: Final = _Outage()
def _checkpoint_verdict(body: dict[str, object]) -> dict[str, object]:
data: Final = body["data"]
content: Final = str(data.get("content", "")) if isinstance(data, dict) else ""
if not body["actor"]:
return {"action_type": "allow"}
if KEYWORD in content:
return {"action_type": "block", "modified_data": {"content": DENIAL}}
if NAME in content:
return {"action_type": "anonymize", "modified_data": {"content": content.replace(NAME, "{Name}")}}
return {"action_type": "allow"}
def _tracker(request: Request) -> Reply:
if request.headers.get("x-api-key") != API_KEY:
return Reply(status=401, body=b'{"detail": "missing api key"}')
path: Final = urlsplit(request.target).path
if not OUTAGE.admits(path):
return Reply(status=503, body=b'{"detail": "tracker unavailable"}')
if path != CHECKPOINT:
return Reply(status=404, body=b'{"detail": "not found"}')
return Reply(body=json.dumps(_checkpoint_verdict(json.loads(request.body))).encode())
def _responses_reply(marker: str) -> dict[str, object]:
return {
"id": "resp_" + marker,
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": [
{
"type": "message",
"id": "msg_" + marker,
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": REPLY, "annotations": []}],
}
],
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
}
def _chat_reply(marker: str) -> dict[str, object]:
return {
"id": "chatcmpl-" + marker,
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": REPLY}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
}
def _provider(request: Request) -> Reply:
marker: Final = uuid.uuid4().hex
is_responses: Final = urlsplit(request.target).path.endswith("/responses")
return Reply(body=json.dumps(_responses_reply(marker) if is_responses else _chat_reply(marker)).encode())
def _alias() -> str:
return f"[{OTHER_APP}] integration key {uuid.uuid4().hex}"
def _ovalix(name: str, tracker: str, **routing: str) -> dict[str, object]:
return {
"guardrail_name": name,
"litellm_params": {
"guardrail": "ovalix",
"mode": "pre_call",
"default_on": False,
"tracker_api_base": tracker,
"tracker_api_key": API_KEY,
**routing,
},
}
@dataclass(frozen=True, slots=True)
class Keys:
no_user: str
user_id: str
with_user_id: str
email: str
with_email: str
@dataclass(frozen=True, slots=True)
class Rig:
proxy: Gateway
tracker: Wire
provider: Wire
model: str
keys: Keys
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
root: Final = tmp_path_factory.mktemp("ovalix")
with (
gateway_from_environment() as gateway,
wire_server(_tracker) as tracker,
wire_server(_provider) as provider,
):
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"]["cache"] = False
config["guardrails"] = [
_ovalix(
PINNED,
tracker.url,
application_id=APP_ID,
pre_checkpoint_id=PRE_CHECKPOINT,
post_checkpoint_id=POST_CHECKPOINT,
),
]
path: Final = root / "ovalix.yaml"
path.write_text(yaml.safe_dump(config))
with (
owned_proxy_process(gateway, root, {}, config=path, workers=2) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = scenario.model(
model="openai/gpt-4o-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key"
)
user_id: Final = scenario.user()
email: Final = f"ovalix-{uuid.uuid4().hex}@example.com"
keys: Final = Keys(
no_user=scenario.key(key_alias=_alias()),
user_id=user_id,
with_user_id=scenario.key(key_alias=_alias(), user_id=user_id),
email=email,
with_email=scenario.key(key_alias=_alias(), user_id=scenario.user(user_email=email)),
)
yield Rig(owned.gateway, tracker, provider, model, keys)
def _post(rig: Rig, endpoint: str, key: str, text: str, extra: dict[str, JsonValue] | None = None) -> httpx.Response:
prompt: Final[dict[str, JsonValue]] = (
{"input": text}
if endpoint == "/v1/responses"
else {"max_tokens": 64, "messages": [{"role": "user", "content": text}]}
if endpoint == "/v1/messages"
else {"messages": [{"role": "user", "content": text}]}
)
return rig.proxy.request(
"POST", endpoint, {"model": rig.model, "guardrails": [PINNED], **prompt, **(extra or {})}, key=key
)
def _send(rig: Rig, endpoint: str, key: str, text: str, extra: dict[str, JsonValue] | None = None) -> httpx.Response:
rig.tracker.drain()
rig.provider.drain()
return _post(rig, endpoint, key, text, extra)
def _checkpoint_actors(rig: Rig) -> list[object]:
return [
json.loads(request.body)["actor"]
for request in rig.tracker.drain()
if urlsplit(request.target).path == CHECKPOINT
]
@pytest.mark.parametrize("stream", (False, True))
@pytest.mark.parametrize("endpoint", ENDPOINTS)
def test_a_prompt_from_a_key_with_no_user_is_inspected(rig: Rig, endpoint: str, stream: bool) -> None:
response: Final = _send(
rig, endpoint, rig.keys.no_user, f"Tell me about {KEYWORD} {uuid.uuid4()}", {"stream": stream}
)
assert response.status_code == 400 and DENIAL in response.text, response.text
assert rig.provider.drain() == ()
assert _checkpoint_actors(rig) == ["unknown"]
@pytest.mark.parametrize(
("endpoint", "client_metadata"),
(("/v1/messages", {"user_id": "client-chosen"}), ("/v1/responses", {"client": "chosen"})),
)
def test_a_client_metadata_object_does_not_switch_inspection_off(
rig: Rig, endpoint: str, client_metadata: dict[str, JsonValue]
) -> None:
response: Final = _send(
rig,
endpoint,
rig.keys.with_user_id,
f"Tell me about {KEYWORD} {uuid.uuid4()}",
{"metadata": client_metadata},
)
assert response.status_code == 400 and DENIAL in response.text, response.text
assert rig.provider.drain() == ()
assert _checkpoint_actors(rig) == ["unknown"]
def test_a_name_from_a_key_with_no_user_reaches_the_provider_anonymized(rig: Rig) -> None:
marker: Final = uuid.uuid4().hex
response: Final = _send(rig, "/v1/chat/completions", rig.keys.no_user, f"My name is {NAME} {marker}")
assert response.status_code == 200, response.text
sent: Final = [json.loads(request.body)["messages"][-1]["content"] for request in rig.provider.drain()]
assert sent == [f"My name is {{Name}} {marker}"]
def test_the_actor_is_the_key_users_email_then_their_id(rig: Rig) -> None:
by_id: Final = _send(rig, "/v1/chat/completions", rig.keys.with_user_id, f"hello {uuid.uuid4()}")
assert by_id.status_code == 200, by_id.text
assert _checkpoint_actors(rig) == [rig.keys.user_id]
by_email: Final = _send(rig, "/v1/chat/completions", rig.keys.with_email, f"hello {uuid.uuid4()}")
assert by_email.status_code == 200, by_email.text
assert _checkpoint_actors(rig) == [rig.keys.email]
@pytest.mark.parametrize("endpoint", ENDPOINTS)
def test_a_key_alias_naming_another_app_does_not_change_the_application_checked(rig: Rig, endpoint: str) -> None:
response: Final = _send(rig, endpoint, rig.keys.with_email, f"Tell me about {KEYWORD} {uuid.uuid4()}")
assert response.status_code == 400 and DENIAL in response.text, response.text
assert rig.provider.drain() == ()
calls: Final = [(urlsplit(request.target).path, json.loads(request.body)) for request in rig.tracker.drain()]
assert [
(path, body["application_id"], body["checkpoint_id"], "application_name" in body) for path, body in calls
] == [(CHECKPOINT, APP_ID, PRE_CHECKPOINT, False)]
def _burst(rig: Rig) -> tuple[httpx.Response, ...]:
plan: Final = tuple(product(ENDPOINTS, (False, True), range(4)))
with ThreadPoolExecutor(max_workers=12) as pool:
return tuple(
pool.map(
lambda cell: _post(
rig,
cell[0],
rig.keys.no_user,
f"Tell me about {KEYWORD} {uuid.uuid4()}",
{"stream": cell[1]},
),
plan,
)
)
def test_a_tracker_outage_mid_burst_fails_closed_and_recovers(rig: Rig) -> None:
rig.tracker.drain()
rig.provider.drain()
OUTAGE.arm(OUTAGE_AFTER)
try:
during: Final = _burst(rig)
finally:
OUTAGE.clear()
denied: Final = [response for response in during if response.status_code == 400 and DENIAL in response.text]
failed: Final = [response for response in during if response.status_code >= 400 and DENIAL not in response.text]
assert len(denied) == OUTAGE_AFTER, [response.text for response in during]
assert len(failed) == len(during) - OUTAGE_AFTER, [response.text for response in during]
assert all("Ovalix guardrail error" in response.text for response in failed), [response.text for response in failed]
assert rig.provider.drain() == ()
assert set(_checkpoint_actors(rig)) == {"unknown"}
after: Final = _burst(rig)
assert all(response.status_code == 400 and DENIAL in response.text for response in after), [
response.text for response in after
]
assert rig.provider.drain() == ()
assert _checkpoint_actors(rig) == ["unknown"] * len(after)

592
tests/integration/observability/test_presidio_masking_endpoints.py Normal file
View file

@ -0,0 +1,592 @@
import json
import re
import uuid
from collections.abc import Callable, Iterator
from contextlib import ExitStack
from dataclasses import dataclass
from itertools import accumulate, chain
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway, eventually, gateway_from_environment, list_value, object_value
from integration._support.process import owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
CARD: Final = "4111 1111 1111 1111"
PHONE: Final = "+1 415-555-0134"
MASKED_CARD: Final = "<CREDIT_CARD>"
MASKED_EMAIL: Final = "<EMAIL_ADDRESS>"
MASKED_PHONE: Final = "<PHONE_NUMBER>"
REPEAT_REQUEST: Final = "Repeat this sentence back to me exactly, word for word: "
THREE_LINE_REQUEST: Final = (
"Reply with exactly three lines and nothing else. "
"Line 1: the part of the email address before the @ sign. "
"Line 2: the full email address. "
"Line 3: the phone number."
)
CARD_REQUEST: Final = "Reply with only the well known Visa sandbox test card number that starts with 4111."
EMAIL_PATTERN: Final = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9-]+(?:\.[A-Za-z0-9-]+)*\.[A-Za-z]{2,}")
PHONE_PATTERN: Final = re.compile(r"(?:\+1 )?\d{3}-\d{3}-\d{4}")
CARD_PATTERN: Final = re.compile(r"\b(?:4\d{3}|5[0-5]\d{2}|6\d{3}|1\d{3}|3\d{3})[- ]?\d{3,4}[- ]?\d{3,4}[- ]?\d{3,5}\b")
RECOGNIZERS: Final = {
"EMAIL_ADDRESS": (EMAIL_PATTERN, 1.0),
"PHONE_NUMBER": (PHONE_PATTERN, 0.4),
"CREDIT_CARD": (CARD_PATTERN, 1.0),
}
JSON: Final = TypeAdapter(JsonValue)
class _Frozen(BaseModel):
model_config = ConfigDict(frozen=True)
class _AnalyzeRequest(_Frozen):
text: str = ""
language: str = ""
entities: tuple[str, ...] | None = None
class _Span(_Frozen):
entity_type: str
start: int
end: int
score: float
class _AnonymizeRequest(_Frozen):
text: str
analyzer_results: tuple[_Span, ...]
class _Part(_Frozen):
text: str
class _Message(_Frozen):
content: str | tuple[_Part, ...]
class _ModelRequest(_Frozen):
messages: tuple[_Message, ...] = ()
input: str | tuple[_Message, ...] | None = None
stream: bool = False
@property
def prompt(self) -> str:
turns: Final = self.messages if self.input is None else self.input
if isinstance(turns, str):
return turns
content: Final = turns[-1].content
return content if isinstance(content, str) else "".join(part.text for part in content)
class _StreamDelta(_Frozen):
content: str | None = None
class _StreamChoice(_Frozen):
delta: _StreamDelta
class _ChatChunk(_Frozen):
choices: tuple[_StreamChoice, ...] = ()
class _AnthropicDelta(_Frozen):
type: str | None = None
text: str | None = None
class _AnthropicEvent(_Frozen):
type: str
delta: _AnthropicDelta | None = None
class _LoggedEntity(_Frozen):
entity_type: str
start: int
end: int
score: float
def _passes_luhn(digits: str) -> bool:
checksum: Final = sum(
digit if position % 2 == 0 else (digit * 2 - 9 if digit * 2 > 9 else digit * 2)
for position, digit in enumerate(int(char) for char in reversed(digits))
)
return checksum % 10 == 0
def _detections(entity: str, text: str) -> tuple[_Span, ...]:
pattern, score = RECOGNIZERS[entity]
return tuple(
_Span(entity_type=entity, start=match.start(), end=match.end(), score=score)
for match in pattern.finditer(text)
if entity != "CREDIT_CARD" or _passes_luhn(re.sub(r"\D", "", match.group()))
)
def _analyze(request: Request) -> Reply:
assert request.target == "/analyze", request.target
body: Final = _AnalyzeRequest.model_validate_json(request.body)
if not body.text or not body.language:
return Reply(status=400, body=json.dumps({"error": "No text provided"}).encode())
requested: Final = body.entities or tuple(RECOGNIZERS)
findings: Final = chain.from_iterable(
_detections(entity, body.text) for entity in requested if entity in RECOGNIZERS
)
return Reply(body=json.dumps([span.model_dump() for span in sorted(findings, key=lambda s: s.start)]).encode())
def _anonymize(request: Request) -> Reply:
assert request.target == "/anonymize", request.target
body: Final = _AnonymizeRequest.model_validate_json(request.body)
spans: Final = sorted(body.analyzer_results, key=lambda span: span.start)
masks: Final = [f"<{span.entity_type}>" for span in spans]
shifts: Final = tuple(
accumulate((len(mask) - (span.end - span.start) for mask, span in zip(masks, spans)), initial=0)
)
items: Final = [
{
"start": span.start + shift,
"end": span.start + shift + len(mask),
"entity_type": span.entity_type,
"text": mask,
"operator": "replace",
}
for span, mask, shift in zip(spans, masks, shifts)
]
pieces: Final = [
body.text[(spans[index - 1].end if index else 0) : span.start] + masks[index]
for index, span in enumerate(spans)
]
tail: Final = body.text[spans[-1].end :] if spans else body.text
return Reply(body=json.dumps({"text": "".join(pieces) + tail, "items": items}).encode())
def _model_answer(prompt: str) -> str:
if REPEAT_REQUEST in prompt:
return prompt.split(REPEAT_REQUEST, 1)[1]
if THREE_LINE_REQUEST in prompt:
email: Final = EMAIL_PATTERN.search(prompt)
phone: Final = PHONE_PATTERN.search(prompt)
assert email is not None and phone is not None, prompt
return "\n".join((email.group().split("@")[0], email.group(), phone.group()))
assert CARD_REQUEST in prompt, prompt
return CARD
def _chat_completion(answer: str) -> dict[str, JsonValue]:
return {
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": answer}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 11, "completion_tokens": 5, "total_tokens": 16},
}
def _chat_frames(answer: str) -> tuple[bytes, ...]:
identity: Final = f"chatcmpl-{uuid.uuid4().hex}"
def frame(delta: dict[str, JsonValue], finish: str | None = None) -> bytes:
payload: Final = {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
}
return b"data: " + json.dumps(payload).encode() + b"\n\n"
return (
frame({"role": "assistant", "content": ""}),
*(frame({"content": word}) for word in re.findall(r"\S+\s*", answer)),
frame({}, "stop"),
b"data: [DONE]\n\n",
)
def _response_object(answer: str, identity: str, item: str, *, completed: bool = True) -> dict[str, JsonValue]:
return {
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed" if completed else "in_progress",
"model": "gpt-4o-mini",
"output": [_message_item(answer, item)] if completed else [],
"usage": {
"input_tokens": 11,
"output_tokens": 5,
"total_tokens": 16,
"input_tokens_details": {"cached_tokens": 0},
"output_tokens_details": {"reasoning_tokens": 0},
}
if completed
else None,
}
def _message_item(answer: str, item: str, *, completed: bool = True) -> dict[str, JsonValue]:
return {
"type": "message",
"id": item,
"status": "completed" if completed else "in_progress",
"role": "assistant",
"content": [_output_text(answer)] if completed else [],
}
def _output_text(text: str) -> dict[str, JsonValue]:
return {"type": "output_text", "text": text, "annotations": []}
def _response_frames(answer: str) -> tuple[bytes, ...]:
identity: Final = f"resp_{uuid.uuid4().hex}"
item: Final = f"msg_{uuid.uuid4().hex}"
location: Final[dict[str, JsonValue]] = {"item_id": item, "output_index": 0, "content_index": 0}
events: Final[tuple[dict[str, JsonValue], ...]] = (
{"type": "response.created", "response": _response_object("", identity, item, completed=False)},
{
"type": "response.output_item.added",
"output_index": 0,
"item": _message_item("", item, completed=False),
},
{"type": "response.content_part.added", **location, "part": _output_text("")},
*({"type": "response.output_text.delta", **location, "delta": word} for word in re.findall(r"\S+\s*", answer)),
{"type": "response.output_text.done", **location, "text": answer},
{"type": "response.content_part.done", **location, "part": _output_text(answer)},
{"type": "response.output_item.done", "output_index": 0, "item": _message_item(answer, item)},
{"type": "response.completed", "response": _response_object(answer, identity, item)},
)
return tuple(
f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode()
for number, event in enumerate(events)
)
def _model(request: Request) -> Reply:
body: Final = _ModelRequest.model_validate_json(request.body)
answer: Final = _model_answer(body.prompt)
match request.target:
case "/v1/chat/completions" if body.stream:
return Reply(content_type="text/event-stream", chunks=_chat_frames(answer))
case "/v1/chat/completions":
return Reply(body=json.dumps(_chat_completion(answer)).encode())
case "/v1/responses" if body.stream:
return Reply(content_type="text/event-stream", chunks=_response_frames(answer))
case "/v1/responses":
return Reply(
body=json.dumps(
_response_object(answer, f"resp_{uuid.uuid4().hex}", f"msg_{uuid.uuid4().hex}")
).encode()
)
case _:
raise AssertionError(request.target)
def _presidio_guardrail(
name: str,
analyzer: Wire,
anonymizer: Wire,
*,
mode: str | list[str],
scope: str | None,
entities: dict[str, JsonValue] | None,
) -> dict[str, JsonValue]:
return {
"guardrail_name": name,
"litellm_params": {
"guardrail": "presidio",
"mode": mode,
"default_on": False,
"presidio_analyzer_api_base": analyzer.url,
"presidio_anonymizer_api_base": anonymizer.url,
**({} if scope is None else {"presidio_filter_scope": scope}),
**({} if entities is None else {"pii_entities_config": entities}),
},
}
@dataclass(frozen=True, slots=True)
class Guardrails:
input_scope: str
logged: str
output_scope: str
default_scope: str
card_output_scope: str
@dataclass(frozen=True, slots=True)
class Rig:
gateway: Gateway
model: Wire
guardrails: Guardrails
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
directory: Final = tmp_path_factory.mktemp("presidio-masking-endpoints")
suffix: Final = uuid.uuid4().hex[:8]
names: Final = Guardrails(
input_scope=f"presidio-input-{suffix}",
logged=f"presidio-logged-{suffix}",
output_scope=f"presidio-output-{suffix}",
default_scope=f"presidio-default-{suffix}",
card_output_scope=f"presidio-card-output-{suffix}",
)
contact_entities: Final[dict[str, JsonValue]] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"}
card_entities: Final[dict[str, JsonValue]] = {"CREDIT_CARD": "MASK"}
with ExitStack() as stack:
model: Final = stack.enter_context(wire_server(_model))
analyzer: Final = stack.enter_context(wire_server(_analyze))
anonymizer: Final = stack.enter_context(wire_server(_anonymize))
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["general_settings"]["store_prompts_in_spend_logs"] = True
config["guardrails"] = [
_presidio_guardrail(names.input_scope, analyzer, anonymizer, mode="pre_call", scope="input", entities=None),
_presidio_guardrail(
names.logged, analyzer, anonymizer, mode="pre_call", scope="input", entities=contact_entities
),
_presidio_guardrail(
names.output_scope, analyzer, anonymizer, mode="post_call", scope="output", entities=contact_entities
),
_presidio_guardrail(
names.default_scope,
analyzer,
anonymizer,
mode=["pre_call", "post_call"],
scope=None,
entities=card_entities,
),
_presidio_guardrail(
names.card_output_scope, analyzer, anonymizer, mode="post_call", scope="output", entities=card_entities
),
]
path: Final = directory / "presidio-masking.yaml"
path.write_text(yaml.safe_dump(config))
shared: Final = stack.enter_context(gateway_from_environment())
owned: Final = stack.enter_context(owned_proxy_process(shared, directory, {}, config=path))
yield Rig(owned.gateway, model, names)
@pytest.fixture
def model(rig: Rig) -> Iterator[str]:
with rig.gateway.scenario() as scenario:
yield scenario.model(api_base=rig.model.url + "/v1")
def _content_strings(node: JsonValue, key: str | None = None) -> Iterator[str]:
match node:
case str() if key in {"content", "text"}:
yield node
case dict():
for child_key, child in node.items():
yield from _content_strings(child, child_key)
case list():
for item in node:
yield from _content_strings(item, key)
case _:
return
def _json_text(raw: str) -> str:
return "\n".join(_content_strings(JSON.validate_json(raw)))
def _data_events(raw: str) -> tuple[str, ...]:
return tuple(
line.removeprefix("data: ") for line in raw.splitlines() if line.startswith("data: ") and line != "data: [DONE]"
)
def _chat_stream_text(raw: str) -> str:
return "".join(
choice.delta.content
for event in _data_events(raw)
for choice in _ChatChunk.model_validate_json(event).choices[:1]
if choice.delta.content
)
def _messages_stream_text(raw: str) -> str:
return "".join(
event.delta.text
for payload in _data_events(raw)
for event in [_AnthropicEvent.model_validate_json(payload)]
if event.type == "content_block_delta"
and event.delta is not None
and event.delta.type == "text_delta"
and event.delta.text
)
@dataclass(frozen=True, slots=True)
class Surface:
path: str
stream: bool
takes_input: bool
read: Callable[[str], str]
CHAT: Final = Surface("/v1/chat/completions", False, False, _json_text)
CHAT_STREAM: Final = Surface("/v1/chat/completions", True, False, _chat_stream_text)
MESSAGES: Final = Surface("/v1/messages", False, False, _json_text)
MESSAGES_STREAM: Final = Surface("/v1/messages", True, False, _messages_stream_text)
RESPONSES: Final = Surface("/v1/responses", False, True, _json_text)
@dataclass(frozen=True, slots=True)
class Answer:
text: str
wire: str
applied_guardrails: str
sent_to_model: tuple[str, ...]
def _request_body(surface: Surface, model: str, guardrail: str, prompt: str) -> dict[str, JsonValue]:
common: Final[dict[str, JsonValue]] = {"model": model, "guardrails": [guardrail], "stream": surface.stream}
if surface.takes_input:
return {**common, "input": prompt}
return {**common, "max_tokens": 512, "messages": [{"role": "user", "content": prompt}]}
def _ask(rig: Rig, model: str, surface: Surface, guardrail: str, prompt: str) -> Answer:
rig.model.drain()
with rig.gateway.client.stream(
"POST",
surface.path,
json=_request_body(surface, model, guardrail, prompt),
headers={"Authorization": f"Bearer {rig.gateway.key}"},
) as response:
wire: Final = response.read().decode()
status: Final = response.status_code
applied: Final = response.headers.get("x-litellm-applied-guardrails", "")
assert status == 200, wire
return Answer(
surface.read(wire),
wire,
applied,
tuple(_ModelRequest.model_validate_json(request.body).prompt for request in rig.model.drain()),
)
def _spend_row(rig: Rig, model: str) -> dict[str, JsonValue]:
rows: Final = eventually(
lambda: rig.gateway.rows('SELECT response, metadata FROM "LiteLLM_SpendLogs" WHERE model_group = %s', (model,)),
lambda found: len(found) == 1,
seconds=70,
)
return rows[0]
def _marker() -> str:
return uuid.uuid4().hex[:8]
def _repeat_prompt(marker: str, email: str) -> str:
return f"{marker} {REPEAT_REQUEST}My email address is {email} and my phone number is {PHONE}."
def _three_line_prompt(marker: str, local_part: str) -> str:
return f"{marker} My email address is {local_part}@example.com and my phone number is {PHONE}. {THREE_LINE_REQUEST}"
def _card_prompt(marker: str) -> str:
return f"{marker} {CARD_REQUEST}"
@pytest.mark.parametrize("surface", (CHAT, MESSAGES), ids=("chat", "messages"))
def test_pre_call_masks_contact_details_before_the_model_sees_them(rig: Rig, model: str, surface: Surface) -> None:
marker: Final = _marker()
email: Final = f"jane.doe.{marker}@example.com"
answer: Final = _ask(rig, model, surface, rig.guardrails.input_scope, _repeat_prompt(marker, email))
masked_sentence: Final = f"My email address is {MASKED_EMAIL} and my phone number is {MASKED_PHONE}."
assert answer.sent_to_model == (f"{marker} {REPEAT_REQUEST}{masked_sentence}",), answer.sent_to_model
assert answer.text == masked_sentence, answer.wire
def test_post_call_masks_the_model_output_while_the_model_sees_the_raw_request(rig: Rig, model: str) -> None:
marker: Final = _marker()
local_part: Final = f"e2euser{marker}"
prompt: Final = _three_line_prompt(marker, local_part)
answer: Final = _ask(rig, model, CHAT, rig.guardrails.output_scope, prompt)
assert answer.sent_to_model == (prompt,), answer.sent_to_model
assert answer.text == f"{local_part}\n{MASKED_EMAIL}\n{MASKED_PHONE}", answer.wire
@pytest.mark.parametrize(
"surface", (CHAT, CHAT_STREAM, MESSAGES_STREAM), ids=("chat", "chat_stream", "messages_stream")
)
def test_default_scope_masks_a_card_number_the_model_generates(rig: Rig, model: str, surface: Surface) -> None:
answer: Final = _ask(rig, model, surface, rig.guardrails.default_scope, _card_prompt(_marker()))
assert answer.text == MASKED_CARD, answer.wire
assert "4111" not in answer.wire, answer.wire
@pytest.mark.parametrize(
"surface",
(CHAT, CHAT_STREAM, MESSAGES, MESSAGES_STREAM, RESPONSES),
ids=("chat", "chat_stream", "messages", "messages_stream", "responses"),
)
def test_spend_log_stores_the_masked_output_the_caller_received(rig: Rig, model: str, surface: Surface) -> None:
answer: Final = _ask(rig, model, surface, rig.guardrails.card_output_scope, _card_prompt(_marker()))
assert answer.text == MASKED_CARD, answer.wire
stored: Final = _spend_row(rig, model)["response"]
assert "\n".join(_content_strings(stored)) == MASKED_CARD, stored
assert "4111" not in json.dumps(stored), stored
def test_masking_run_is_recorded_on_the_spend_log_with_the_detected_entities(rig: Rig, model: str) -> None:
marker: Final = _marker()
email: Final = f"jane.doe.{marker}@example.com"
prompt: Final = _repeat_prompt(marker, email)
answer: Final = _ask(rig, model, CHAT, rig.guardrails.logged, prompt)
assert rig.guardrails.logged in answer.applied_guardrails, answer.applied_guardrails
metadata: Final = object_value(_spend_row(rig, model)["metadata"])
records: Final = tuple(
object_value(record)
for record in list_value(metadata["guardrail_information"])
if object_value(record)["guardrail_name"] == rig.guardrails.logged
)
assert len(records) == 1, metadata["guardrail_information"]
record: Final = records[0]
assert {
key: record[key]
for key in (
"guardrail_mode",
"guardrail_status",
"guardrail_provider",
"masked_entity_count",
)
} == {
"guardrail_mode": "pre_call",
"guardrail_status": "success",
"guardrail_provider": "presidio",
"masked_entity_count": {"EMAIL_ADDRESS": 1, "PHONE_NUMBER": 1},
}, record
detected: Final = sorted(
(_LoggedEntity.model_validate(entity) for entity in list_value(record["guardrail_response"])),
key=lambda entity: entity.start,
)
email_start: Final = prompt.index(email)
phone_start: Final = prompt.index(PHONE)
assert [(entity.entity_type, entity.start, entity.end, entity.score) for entity in detected] == [
("EMAIL_ADDRESS", email_start, email_start + len(email), 1.0),
("PHONE_NUMBER", phone_start, phone_start + len(PHONE), 0.4),
], record["guardrail_response"]

223
tests/integration/observability/test_weave_key_scoped_delivery.py Normal file
View file

@ -0,0 +1,223 @@
import base64
import json
import uuid
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
from integration._support.process import owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest
from opentelemetry.proto.common.v1.common_pb2 import AnyValue
from opentelemetry.proto.trace.v1.trace_pb2 import Status
from pydantic import JsonValue
INPUT_RATE: Final = 0.001
OUTPUT_RATE: Final = 0.002
PROMPT_TOKENS: Final = 7
COMPLETION_TOKENS: Final = 2
KEY_WANDB_API_KEY: Final = "tenant-wandb-key"
KEY_WEAVE_PROJECT: Final = "tenant/key-project"
OPERATOR_WEAVE_PROJECT: Final = "operator/project"
FAILING_PREFIX: Final = "weave-fail-"
UPSTREAM_AUTH_ERROR: Final = "Incorrect API key provided for the scripted upstream"
LATE_DUPLICATE_WINDOW_SECONDS: Final = 3
@dataclass(frozen=True, slots=True)
class WeaveSpan:
status_code: int
attributes: Mapping[str, str]
project_id: str
authorization: str
@dataclass(frozen=True, slots=True)
class Rig:
proxy: Gateway
scenario: Scenario
model: str
weave: Wire
def _attribute_text(value: AnyValue) -> str:
match value.WhichOneof("value"):
case "string_value":
return value.string_value
case "int_value":
return str(value.int_value)
case "double_value":
return str(value.double_value)
case "bool_value":
return str(value.bool_value)
case _:
return ""
def _spans(request: Request) -> Iterator[WeaveSpan]:
if request.headers.get("content-type") != "application/x-protobuf":
return
batch: Final = ExportTraceServiceRequest.FromString(request.body)
for resource_spans in batch.resource_spans:
for scope_spans in resource_spans.scope_spans:
for span in scope_spans.spans:
yield WeaveSpan(
status_code=span.status.code,
attributes={attribute.key: _attribute_text(attribute.value) for attribute in span.attributes},
project_id=request.headers.get("project_id", ""),
authorization=request.headers.get("authorization", ""),
)
def _upstream_reply(request: Request) -> Reply:
prompt: Final = json.loads(request.body)["messages"][0]["content"]
if prompt.startswith(FAILING_PREFIX):
return Reply(
status=401,
body=json.dumps(
{
"error": {
"message": UPSTREAM_AUTH_ERROR,
"type": "invalid_request_error",
"param": None,
"code": "invalid_api_key",
}
}
).encode(),
)
return Reply(
body=json.dumps(
{
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}],
"usage": {
"prompt_tokens": PROMPT_TOKENS,
"completion_tokens": COMPLETION_TOKENS,
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
},
}
).encode()
)
def _collector(_request: Request) -> Reply:
return Reply(body=b"", content_type="application/x-protobuf")
def _weave_logging_for_the_key() -> list[JsonValue]:
return [
{
"callback_name": "weave_otel",
"callback_type": "success_and_failure",
"callback_vars": {"wandb_api_key": KEY_WANDB_API_KEY, "weave_project_id": KEY_WEAVE_PROJECT},
}
]
def _matching(weave: Wire, prompt: str) -> tuple[WeaveSpan, ...]:
return tuple(span for request in weave.drain() for span in _spans(request) if prompt in span.attributes.values())
def _delivered_once(weave: Wire, prompt: str) -> WeaveSpan:
first: Final = eventually(lambda: _matching(weave, prompt), bool, seconds=30)
late: Final = eventually(
lambda: _matching(weave, prompt),
bool,
seconds=LATE_DUPLICATE_WINDOW_SECONDS,
return_last_on_timeout=True,
)
assert len(first) + len(late) == 1, (first, late)
return first[0]
@contextmanager
def _started(directory: Path, callbacks: list[str]) -> Iterator[Rig]:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
path: Final = directory / "weave.yaml"
path.write_text(
yaml.safe_dump({**config, "litellm_settings": {**config["litellm_settings"], "callbacks": callbacks}})
)
with (
wire_server(_upstream_reply) as provider,
wire_server(_collector) as weave,
gateway_from_environment() as gateway,
owned_proxy_process(
gateway,
directory,
{
"LITELLM_OTEL_V2": "1",
"OTEL_BSP_SCHEDULE_DELAY": "100",
"OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT": "span_only",
"WANDB_API_KEY": "operator-wandb-key",
"WANDB_PROJECT_ID": OPERATOR_WEAVE_PROJECT,
"WANDB_HOST": weave.url,
},
config=path,
) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = scenario.model(
api_base=provider.url + "/v1", input_cost_per_token=INPUT_RATE, output_cost_per_token=OUTPUT_RATE
)
yield Rig(owned.gateway, scenario, model, weave)
@pytest.fixture(scope="module", params=("weave_otel", "otel"))
def rig(request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
if request.param == "otel":
pytest.skip(
"BUG: a key-scoped weave_otel callback ships nothing when the OTEL v2 otel callback is the only one configured"
)
with _started(tmp_path_factory.mktemp("weave-key-scoped"), [request.param]) as started:
yield started
def _chat(rig: Rig, key: str, prompt: str) -> tuple[int, str, str]:
response: Final = rig.proxy.request(
"POST",
"/v1/chat/completions",
{"model": rig.model, "messages": [{"role": "user", "content": prompt}]},
key=key,
)
return response.status_code, response.headers.get("x-litellm-response-cost", ""), response.text
def test_key_scoped_weave_gets_one_span_with_the_request_cost_in_the_keys_project(rig: Rig) -> None:
prompt: Final = f"weave-ok-{uuid.uuid4().hex}"
key: Final = rig.scenario.key(models=[rig.model], metadata={"logging": _weave_logging_for_the_key()})
status, header_cost, text = _chat(rig, key, prompt)
assert status == 200, text
span: Final = _delivered_once(rig.weave, prompt)
assert span.project_id == KEY_WEAVE_PROJECT, span
assert span.authorization == "Basic " + base64.b64encode(f"api:{KEY_WANDB_API_KEY}".encode()).decode(), span
assert span.status_code != Status.STATUS_CODE_ERROR, span
expected_cost: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE
assert float(header_cost) == pytest.approx(expected_cost), header_cost
assert float(span.attributes["litellm.cost.total"]) == pytest.approx(expected_cost), span
assert span.attributes["gen_ai.usage.input_tokens"] == str(PROMPT_TOKENS), span
assert span.attributes["gen_ai.usage.output_tokens"] == str(COMPLETION_TOKENS), span
def test_key_scoped_weave_gets_one_error_span_for_a_provider_failure_without_cost(rig: Rig) -> None:
prompt: Final = f"{FAILING_PREFIX}{uuid.uuid4().hex}"
key: Final = rig.scenario.key(models=[rig.model], metadata={"logging": _weave_logging_for_the_key()})
status, _, text = _chat(rig, key, prompt)
assert status == 401, text
span: Final = _delivered_once(rig.weave, prompt)
assert span.project_id == KEY_WEAVE_PROJECT, span
assert span.status_code == Status.STATUS_CODE_ERROR, span
assert span.attributes["error.type"] == "AuthenticationError", span
assert UPSTREAM_AUTH_ERROR in span.attributes["error.message"], span
assert float(span.attributes["litellm.cost.total"]) == 0, span

229
tests/integration/pricing/test_sail_service_tier_pricing.py Normal file
View file

@ -0,0 +1,229 @@
import json
import uuid
from dataclasses import dataclass
from typing import Final
import pytest
from pydantic import JsonValue
from tests.integration._support.client import Gateway, Scenario, eventually, object_value, string_value
from tests.integration._support.database import read_rows
from tests.integration._support.wire import Reply, Request, Wire, wire_server
BACKEND: Final = "zai-org/GLM-5.3"
PROMPT_TOKENS: Final = 1000
CACHED_TOKENS: Final = 400
COMPLETION_TOKENS: Final = 200
PRICES: Final[dict[str, JsonValue]] = {
"input_cost_per_token": 3e-06,
"output_cost_per_token": 9e-06,
"cache_read_input_token_cost": 1e-06,
"input_cost_per_token_balanced": 2e-06,
"output_cost_per_token_balanced": 6e-06,
"cache_read_input_token_cost_balanced": 7e-07,
"input_cost_per_token_flex": 1e-06,
"output_cost_per_token_flex": 3e-06,
"cache_read_input_token_cost_flex": 4e-07,
}
BASE_COST: Final = 600 * 3e-06 + 400 * 1e-06 + 200 * 9e-06
BALANCED_COST: Final = 600 * 2e-06 + 400 * 7e-07 + 200 * 6e-06
FLEX_COST: Final = 600 * 1e-06 + 400 * 4e-07 + 200 * 3e-06
def _chat_reply(_request: Request) -> Reply:
return Reply(
body=json.dumps(
{
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": BACKEND,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}],
"usage": {
"prompt_tokens": PROMPT_TOKENS,
"completion_tokens": COMPLETION_TOKENS,
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
"prompt_tokens_details": {"cached_tokens": CACHED_TOKENS},
},
}
).encode()
)
def _responses_reply(_request: Request) -> Reply:
return Reply(
body=json.dumps(
{
"id": f"resp_{uuid.uuid4().hex}",
"object": "response",
"created_at": 1,
"status": "completed",
"model": BACKEND,
"output": [
{
"type": "message",
"id": f"msg_{uuid.uuid4().hex}",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "OK", "annotations": []}],
}
],
"usage": {
"input_tokens": PROMPT_TOKENS,
"output_tokens": COMPLETION_TOKENS,
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
"input_tokens_details": {"cached_tokens": CACHED_TOKENS},
"output_tokens_details": {"reasoning_tokens": 0},
},
}
).encode()
)
def _messages_reply(_request: Request) -> Reply:
return Reply(
body=json.dumps(
{
"id": f"msg_{uuid.uuid4().hex}",
"type": "message",
"role": "assistant",
"model": BACKEND,
"content": [{"type": "text", "text": "OK"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": PROMPT_TOKENS, "output_tokens": COMPLETION_TOKENS},
}
).encode()
)
def _upstream_body(request: Request) -> dict[str, JsonValue]:
return object_value(json.loads(request.body))
def _sail_model(scenario: Scenario, wire: Wire) -> str:
return scenario.model(model=f"sail/{BACKEND}", api_base=wire.url, api_key="synthetic-sail-key", **PRICES)
def _assert_spend_row(gateway: Gateway, response_id: str, expected_cost: float) -> None:
rows: Final = eventually(
lambda: read_rows(
'SELECT spend, custom_llm_provider, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" '
"WHERE request_id = %s",
(response_id,),
),
lambda values: len(values) == 1,
seconds=70,
)
assert rows[0]["custom_llm_provider"] == "sail", rows
assert rows[0]["prompt_tokens"] == PROMPT_TOKENS, rows
assert rows[0]["completion_tokens"] == COMPLETION_TOKENS, rows
assert float(rows[0]["spend"]) == pytest.approx(expected_cost, rel=1e-6), rows
@dataclass(frozen=True, slots=True)
class Served:
upstream_body: dict[str, JsonValue]
header_cost: float
response_id: str
def _chat_through_sail(gateway: Gateway, extra: dict[str, JsonValue]) -> Served:
with wire_server(_chat_reply) as wire, gateway.scenario() as scenario:
model: Final = _sail_model(scenario, wire)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "Reply with one word."}], **extra},
)
assert response.status_code == 200, response.text
received: Final = wire.drain()
assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
return Served(
upstream_body=_upstream_body(received[0]),
header_cost=float(response.headers["x-litellm-response-cost"]),
response_id=string_value(object_value(response.json())["id"]),
)
@pytest.mark.parametrize(
("service_tier", "completion_window", "expected_cost"),
(
("balanced", "balanced", BALANCED_COST),
("flex", "flex", FLEX_COST),
("auto", None, BASE_COST),
),
ids=("balanced", "flex", "auto"),
)
def test_chat_service_tier_is_sent_as_completion_window_and_billed_at_that_window(
gateway: Gateway, service_tier: str, completion_window: str | None, expected_cost: float
) -> None:
served: Final = _chat_through_sail(gateway, {"service_tier": service_tier})
assert served.upstream_body == {
"model": BACKEND,
"messages": [{"role": "user", "content": "Reply with one word."}],
**({} if completion_window is None else {"metadata": {"completion_window": completion_window}}),
}, served.upstream_body
assert served.header_cost == pytest.approx(expected_cost, rel=1e-6)
_assert_spend_row(gateway, served.response_id, expected_cost)
@pytest.mark.parametrize("service_tier", ("bogus", 5), ids=("string", "integer"))
def test_chat_unknown_service_tier_under_drop_params_is_dropped_and_billed_at_the_base_rate(
gateway: Gateway, service_tier: str | int
) -> None:
served: Final = _chat_through_sail(gateway, {"service_tier": service_tier, "drop_params": True})
assert served.upstream_body == {
"model": BACKEND,
"messages": [{"role": "user", "content": "Reply with one word."}],
}, served.upstream_body
assert served.header_cost == pytest.approx(BASE_COST, rel=1e-6)
_assert_spend_row(gateway, served.response_id, BASE_COST)
def test_responses_caller_completion_window_reaches_sail_and_is_billed_at_that_window(gateway: Gateway) -> None:
with wire_server(_responses_reply) as wire, gateway.scenario() as scenario:
model: Final = _sail_model(scenario, wire)
response: Final = gateway.request(
"POST",
"/v1/responses",
{"model": model, "input": "Reply with one word.", "metadata": {"completion_window": "balanced"}},
)
assert response.status_code == 200, response.text
received: Final = wire.drain()
assert [(request.method, request.target) for request in received] == [("POST", "/responses")]
assert _upstream_body(received[0]) == {
"model": BACKEND,
"input": "Reply with one word.",
"metadata": {"completion_window": "balanced"},
}, received[0].body
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(BALANCED_COST, rel=1e-6), response.text
_assert_spend_row(gateway, string_value(object_value(response.json())["id"]), BALANCED_COST)
def test_messages_call_to_sail_returns_an_anthropic_message(gateway: Gateway) -> None:
with wire_server(_messages_reply) as wire, gateway.scenario() as scenario:
model: Final = _sail_model(scenario, wire)
response: Final = gateway.request(
"POST",
"/v1/messages",
{"model": model, "max_tokens": 64, "messages": [{"role": "user", "content": "Reply with one word."}]},
)
assert response.status_code == 200, response.text
received: Final = wire.drain()
assert [(request.method, request.target) for request in received] == [("POST", "/v1/messages")], received
assert _upstream_body(received[0]) == {
"model": BACKEND,
"max_tokens": 64,
"messages": [{"role": "user", "content": "Reply with one word."}],
"stream": False,
}, received[0].body
payload: Final = object_value(response.json())
assert payload["type"] == "message" and payload["role"] == "assistant", response.text
assert payload["content"] == [{"type": "text", "text": "OK"}], response.text
assert object_value(payload["usage"])["output_tokens"] == COMPLETION_TOKENS, response.text

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:

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

33
tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py
View file

@ -8,7 +8,7 @@ import threading
import uuid
from collections.abc import AsyncIterable, Awaitable, Callable, Iterable
from concurrent.futures import ThreadPoolExecutor
from types import FrameType
from types import FrameType, MappingProxyType
from typing import Final, TypeVar
import pytest
@ -28,12 +28,23 @@ _BURST_THREADS: Final = 8
_BURST_CALLS: Final = 24
_CHURN_UPSERTS: Final = 120
_UNREPORTED_MODEL_ID: Final = frozenset({"aresponses-stream", "aanthropic-messages"})
_ROUTER_FILE: Final = Router.__init__.__code__.co_filename
_INDEX_LOCK_LINES: Final = frozenset(
number
for number, line in enumerate(linecache.getlines(_ROUTER_FILE), start=1)
if "with self._deployment_index_lock" in line
)
_REINDEX: Final = Router._update_deployment_indices_after_removal.__code__
def _router_source_files() -> frozenset[str]:
facade_files: Final = (sys.modules[cls.__module__].__file__ for cls in Router.__mro__ if cls is not object)
return frozenset({_REINDEX.co_filename, *(path for path in facade_files if path is not None)})
def _lock_lines(path: str) -> frozenset[int]:
return frozenset(
number
for number, line in enumerate(linecache.getlines(path), start=1)
if "with self._deployment_index_lock" in line
)
_INDEX_LOCK_LINES: Final = MappingProxyType({path: _lock_lines(path) for path in _router_source_files()})
T = TypeVar("T")
@ -80,7 +91,6 @@ def _read_during_held_write(read: Callable[[], T], write: Callable[[], object])
"""Run `read` on its own thread while `write` waits where the list has shrunk but the index maps have not caught
up. The write resumes once `read` returns or reaches the router's index lock, so a read that skips the lock meets
the half-updated list and a read that takes the lock waits for the write"""
reindex: Final = Router._update_deployment_indices_after_removal.__code__
held: Final = threading.Event()
release: Final = threading.Event()
settled: Final = threading.Event()
@ -89,7 +99,7 @@ def _read_during_held_write(read: Callable[[], T], write: Callable[[], object])
failure: Final[list[BaseException]] = []
def hold_at_reindex(frame: FrameType, event: str, _arg: object) -> None:
if event == "call" and frame.f_code is reindex and not held.is_set():
if event == "call" and frame.f_code is _REINDEX and not held.is_set():
held.set()
released_in_time.append(release.wait(timeout=_SYNC_SECONDS))
@ -102,9 +112,10 @@ def _read_during_held_write(read: Callable[[], T], write: Callable[[], object])
sys.setprofile(previous)
def watch_read(frame: FrameType, event: str, _arg: object) -> Callable[..., object] | None:
if frame.f_code.co_filename != _ROUTER_FILE:
lock_lines: Final = _INDEX_LOCK_LINES.get(frame.f_code.co_filename)
if lock_lines is None:
return None
if event == "line" and frame.f_lineno in _INDEX_LOCK_LINES:
if event == "line" and frame.f_lineno in lock_lines:
settled.set()
return watch_read

6
tests/unit/enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py
View file

@ -29,7 +29,7 @@ async def test_apply_guardrail_endpoint_returns_correct_response(
mock_proxy_logging_ctx(),
):
# Create a mock guardrail
mock_guardrail = Mock(spec=CustomGuardrail)
mock_guardrail = Mock(spec=CustomGuardrail, event_hook=None)
# Apply guardrail returns GenericGuardrailAPIInputs (dict with texts key)
mock_guardrail.apply_guardrail = AsyncMock(
return_value={"texts": ["Redacted text: [REDACTED] and [REDACTED]"]}
@ -116,7 +116,7 @@ async def test_apply_guardrail_endpoint_with_presidio_guardrail(mock_proxy_loggi
mock_proxy_logging_ctx(),
):
# Create a mock guardrail that simulates Presidio behavior
mock_guardrail = Mock(spec=CustomGuardrail)
mock_guardrail = Mock(spec=CustomGuardrail, event_hook=None)
# Simulate masking PII entities - returns GenericGuardrailAPIInputs (dict with texts key)
mock_guardrail.apply_guardrail = AsyncMock(
return_value={
@ -168,7 +168,7 @@ async def test_apply_guardrail_endpoint_without_optional_params(mock_proxy_loggi
mock_proxy_logging_ctx(),
):
# Create a mock guardrail
mock_guardrail = Mock(spec=CustomGuardrail)
mock_guardrail = Mock(spec=CustomGuardrail, event_hook=None)
# Returns GenericGuardrailAPIInputs (dict with texts key)
mock_guardrail.apply_guardrail = AsyncMock(
return_value={"texts": ["Processed text"]}

33
tests/unit/integration_support/test_upstream.py
View file

@ -112,6 +112,39 @@ async def test_mcp_notification_ordering_preserves_other_responses(
assert outgoing == list(messages)
@pytest.mark.asyncio
async def test_mcp_http_disconnect_cancels_an_in_flight_handler(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
from integration._support.mcp import _cancel_on_disconnect
from starlette.types import Receive, Send
started: Final = asyncio.Event()
stopped: Final = asyncio.Event()
body: Final[Message] = {"type": "http.request", "body": b"tool-call", "more_body": False}
pending: Final = iter((body,))
async def receive() -> Message:
message: Final = next(pending, None)
if message is not None:
return message
await started.wait()
return {"type": "http.disconnect"}
async def send(message: Message) -> None:
pytest.fail(f"Disconnected handler sent a response: {message}")
async def app(scope: Scope, receive: Receive, send: Send) -> None:
assert await receive() == body
started.set()
try:
await asyncio.Future[None]()
finally:
stopped.set()
await asyncio.wait_for(_cancel_on_disconnect(app)({"type": "http"}, receive, send), 2)
assert stopped.is_set(), "The disconnected HTTP handler was left running"
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", (RuntimeError("dispatch failed"), asyncio.CancelledError()))
async def test_mcp_notification_dispatch_failure_is_not_acknowledged(

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

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(

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

4
tests/unit/llms/gemini/test_gemini_tts.py
View file

@ -298,11 +298,11 @@ class TestGeminiTTSTransformation:
# Test invalid format
non_default_params = {
"audio": {"voice": "Kore", "format": "wav"} # Invalid format
"audio": {"voice": "Kore", "format": "mp3"} # Invalid format
}
optional_params = {}
with pytest.raises(ValueError, match="Unsupported audio format for Gemini TTS models"):
with pytest.raises(litellm.UnsupportedParamsError, match="audio format 'mp3'"):
config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,

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

@ -4552,6 +4552,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/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -328,7 +328,9 @@ def test_chat_audio_drop_params_keeps_every_voice_config_field_gemini_accepts(sp
def test_chat_audio_reads_snake_case_wrapper_beside_null_camel_case_wrapper():
mapped: Final = VertexGeminiConfig().map_openai_params(
non_default_params={"audio": {"format": "pcm16", "speechConfig": None, "speech_config": {"voice_name": "Kore"}}},
non_default_params={
"audio": {"format": "pcm16", "speechConfig": None, "speech_config": {"voice_name": "Kore"}}
},
optional_params={},
model="gemini-2.5-flash-preview-tts",
drop_params=False,
@ -3468,9 +3470,7 @@ async def test_vertex_ai_streaming_bad_request_is_not_wrapped():
async def dispatch_failure_handlers(self, *args, **kwargs):
return None
async def dispatch_attempt_failure_handlers(
self, exception: Exception, traceback_exception: str
) -> None:
async def dispatch_attempt_failure_handlers(self, exception: Exception, traceback_exception: str) -> None:
return None
def attempt_scoped_copy(self) -> "DummyLogging":
@ -3600,9 +3600,7 @@ def test_grounding_annotations_preserve_all_web_sources(indices: list[int], expe
]
metadata: Final = {
"groundingChunks": chunks,
"groundingSupports": [
{"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": indices}
],
"groundingSupports": [{"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": indices}],
}
annotations: Final = VertexGeminiConfig._convert_grounding_metadata_to_annotations([metadata], "The passage")
expected: Final = [
@ -10425,6 +10423,106 @@ def test_map_openai_params_audio_voice_reaches_speech_config_mapped() -> None:
assert optional_params["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore"
def _top_level_format_audio(audio_format: str) -> Mapping[str, object]:
return {"voice": "Kore", "format": audio_format}
def _format_beside_speech_config_audio(audio_format: str) -> Mapping[str, object]:
return {"speech_config": {"voice_name": "Kore"}, "format": audio_format}
def _format_inside_speech_config_audio(audio_format: str) -> Mapping[str, object]:
return {"speech_config": {"voice_name": "Kore", "format": audio_format}}
_FORMAT_PLACEMENTS: Final = (
_top_level_format_audio,
_format_beside_speech_config_audio,
_format_inside_speech_config_audio,
)
@pytest.mark.parametrize("audio_format", ("mp3", "opus", "flac", "PCM16"))
@pytest.mark.parametrize("audio", _FORMAT_PLACEMENTS)
@pytest.mark.parametrize("model", ("gemini/gemini-3.8-flash-tts", "vertex_ai/gemini-3.8-flash-tts"))
def test_chat_audio_unsupported_format_is_a_400_naming_the_model_and_format(
model: str, audio: Callable[[str], Mapping[str, object]], audio_format: str
) -> None:
with respx.mock, pytest.raises(litellm.BadRequestError) as raised:
completion(
model=model,
messages=[{"role": "user", "content": "Say hello"}],
modalities=["audio"],
audio=audio(audio_format),
api_key="test-key",
vertex_project="test-project",
vertex_location="us-central1",
)
assert raised.value.status_code == 400
assert f"gemini-3.8-flash-tts does not take the audio format '{audio_format}'" in str(raised.value)
@pytest.mark.parametrize("audio", _FORMAT_PLACEMENTS)
@pytest.mark.parametrize("request_drop_params, global_drop_params", ((True, False), (False, True)))
def test_chat_audio_unsupported_format_is_dropped_under_drop_params(
monkeypatch: pytest.MonkeyPatch,
request_drop_params: bool,
global_drop_params: bool,
audio: Callable[[str], Mapping[str, object]],
) -> None:
monkeypatch.setattr(litellm, "drop_params", global_drop_params)
optional_params: Final = VertexGeminiConfig().map_openai_params(
non_default_params={"modalities": ["audio"], "audio": audio("mp3")},
optional_params={},
model="gemini-3.8-flash-tts",
drop_params=request_drop_params,
)
assert optional_params == {
"responseModalities": ["AUDIO"],
"speechConfig": {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}},
}
def test_chat_audio_unsupported_format_error_caps_the_echoed_format() -> None:
with pytest.raises(litellm.UnsupportedParamsError) as raised:
VertexGeminiConfig().map_openai_params(
non_default_params={"modalities": ["audio"], "audio": {"voice": "Kore", "format": "x" * 300}},
optional_params={},
model="gemini-3.8-flash-tts",
drop_params=False,
)
assert f"audio format '{'x' * 100}'" in str(raised.value)
assert "x" * 101 not in str(raised.value)
@pytest.mark.parametrize("audio_format", [5, {"codec": "mp3"}, ["x"] * 300])
def test_chat_audio_non_string_format_error_names_the_value_as_sent(audio_format: object) -> None:
with pytest.raises(litellm.UnsupportedParamsError) as raised:
VertexGeminiConfig().map_openai_params(
non_default_params={"modalities": ["audio"], "audio": {"voice": "Kore", "format": audio_format}},
optional_params={},
model="gemini-3.8-flash-tts",
drop_params=False,
)
assert f"audio format {repr(audio_format)[:100]}. Set format" in str(raised.value)
def test_map_audio_params_rejects_an_unsupported_format() -> None:
with pytest.raises(litellm.UnsupportedParamsError, match="audio format 'mp3'"):
VertexGeminiConfig().map_audio_params({"voice": "Kore", "format": "mp3"})
def test_map_audio_params_maps_a_pcm16_voice() -> None:
assert VertexGeminiConfig().map_audio_params({"voice": "Kore", "format": "pcm16"}) == {
"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}
}
_CACHED_PROMPT_USAGE_METADATA = {
"promptTokenCount": 1536,
"candidatesTokenCount": 40,

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

27
tests/unit/proxy/a2a/test_agent_card.py
View file

@ -8,6 +8,7 @@ from litellm.proxy.a2a.agent_card import (
LITELLM_SECURITY_SCHEMES,
merge_agent_card,
normalize_protocol_version,
point_agent_card_at_proxy,
resolve_served_protocol_version,
)
@ -15,6 +16,32 @@ PROXY_URL = "https://proxy.example/a2a/agent-xyz"
PROXY_BASE = "https://proxy.example"
def test_point_agent_card_at_proxy_returns_a_deep_copy() -> None:
card: dict[str, object] = {
"url": "http://internal/",
"metadata": {"nested": ["original"]},
"supportedInterfaces": [
{"url": "http://internal/", "nested": {"value": "original"}},
{"url": "http://secondary/"},
],
}
pointed: dict[str, object] = point_agent_card_at_proxy(card, PROXY_URL)
assert pointed["url"] == PROXY_URL
assert pointed["metadata"] == card["metadata"]
assert pointed["metadata"] is not card["metadata"]
assert pointed["supportedInterfaces"] == [
{"url": PROXY_URL, "nested": {"value": "original"}},
{"url": "http://secondary/"},
]
assert card["url"] == "http://internal/"
assert card["supportedInterfaces"] == [
{"url": "http://internal/", "nested": {"value": "original"}},
{"url": "http://secondary/"},
]
def _full_upstream_card() -> dict:
return {
"protocolVersion": "0.9",

162
tests/unit/proxy/agent_endpoints/test_endpoints.py
View file

@ -1,7 +1,6 @@
import json
from collections.abc import Mapping
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -13,7 +12,12 @@ from fastapi.testclient import TestClient
from prisma.models import LiteLLM_AgentsTable
from litellm.constants import REDACTED_BY_LITELM_STRING
from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy._types import (
LiteLLM_AuditLogs,
LitellmTableNames,
LitellmUserRoles,
UserAPIKeyAuth,
)
from litellm.proxy.agent_endpoints import endpoints as agent_endpoints
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
RestrictedAgentAccess,
@ -43,6 +47,29 @@ def _sample_agent_card_params() -> dict:
}
_AGENT_CARD_BACKEND_URL: Final = "http://agent.internal:9000/plain/"
def _sample_agent_card_params_with_backend_url() -> dict[str, object]:
return {
**_sample_agent_card_params(),
"url": _AGENT_CARD_BACKEND_URL,
"supportedInterfaces": [
{
"url": _AGENT_CARD_BACKEND_URL,
"protocolBinding": "JSONRPC",
"protocolVersion": "1.0",
},
{
"url": _AGENT_CARD_BACKEND_URL,
"protocolBinding": "GRPC",
"protocolVersion": "1.0",
},
],
"additionalInterfaces": [{"url": _AGENT_CARD_BACKEND_URL, "transport": "GRPC"}],
}
def _sample_agent_config() -> dict:
return {
"agent_name": "Test Agent",
@ -641,7 +668,7 @@ class TestAgentRBACProxyAdmin:
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry)
def test_should_allow_admin_to_create_agent(self, monkeypatch):
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
with patch("litellm.proxy.proxy_server.prisma_client"):
self.mock_registry.get_agent_by_name = MagicMock(return_value=None)
self.mock_registry.add_agent_to_db = AsyncMock(
return_value=_sample_agent_response()
@ -1015,6 +1042,129 @@ class TestAgentHealthCheck:
assert len(resp.json()) == 2
def _client_with_agent_card(
monkeypatch: pytest.MonkeyPatch,
role: LitellmUserRoles,
agent: AgentResponse,
) -> TestClient:
from litellm.proxy import proxy_server
from litellm.proxy.agent_endpoints import agent_registry
registry: Final = MagicMock()
registry.get_agent_by_id = MagicMock(return_value=agent)
registry.get_agent_list = MagicMock(return_value=(agent,))
registry.ids_for_agent = MagicMock(return_value=frozenset({agent.agent_id}))
registry.stable_agent_id = MagicMock(side_effect=lambda agent_id: agent_id)
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry)
monkeypatch.setattr(agent_registry, "global_agent_registry", registry)
prisma: Final = MagicMock()
prisma.writer_db = prisma.db
prisma.db.litellm_agentstable.find_many = AsyncMock(return_value=[])
prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
test_app: Final = FastAPI()
test_app.include_router(router)
test_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="test-user",
user_role=role,
object_permission={"object_permission_id": "agent-card-test", "agents": [agent.agent_id]},
)
return TestClient(test_app)
@pytest.mark.parametrize("path", ["/v1/agents", "/v1/agents/agent-123"])
def test_non_admin_agent_card_urls_point_at_proxy(
monkeypatch: pytest.MonkeyPatch,
path: str,
) -> None:
agent: Final = AgentResponse(
agent_id="agent-123",
agent_name="Test Agent",
agent_card_params=_sample_agent_card_params_with_backend_url(),
)
test_client: Final = _client_with_agent_card(monkeypatch, LitellmUserRoles.INTERNAL_USER, agent)
response: Final = test_client.get(path, headers={"Authorization": "Bearer k"})
assert response.status_code == 200, response.text
agents: Final = response.json()
result: Final = agents[0] if path == "/v1/agents" else agents
card: Final = result["agent_card_params"]
proxy_url: Final = "http://testserver/a2a/agent-123"
assert card["url"] == proxy_url
assert "agent.internal" not in response.text
assert "additionalInterfaces" not in card
assert card["supportedInterfaces"] == [
{
"url": proxy_url,
"protocolBinding": "JSONRPC",
"protocolVersion": "1.0",
}
]
def test_proxy_admin_agent_listing_keeps_backend_card_urls(monkeypatch: pytest.MonkeyPatch) -> None:
agent: Final = AgentResponse(
agent_id="agent-123",
agent_name="Test Agent",
agent_card_params=_sample_agent_card_params_with_backend_url(),
)
test_client: Final = _client_with_agent_card(monkeypatch, LitellmUserRoles.PROXY_ADMIN, agent)
response: Final = test_client.get("/v1/agents", headers={"Authorization": "Bearer k"})
assert response.status_code == 200, response.text
card: Final = response.json()[0]["agent_card_params"]
assert card["url"] == _AGENT_CARD_BACKEND_URL
assert card["supportedInterfaces"] == [
{
"url": _AGENT_CARD_BACKEND_URL,
"protocolBinding": "JSONRPC",
"protocolVersion": "1.0",
},
{
"url": _AGENT_CARD_BACKEND_URL,
"protocolBinding": "GRPC",
"protocolVersion": "1.0",
},
]
assert card["additionalInterfaces"] == [{"url": _AGENT_CARD_BACKEND_URL, "transport": "GRPC"}]
def test_non_admin_health_check_probes_backend_and_returns_proxy_card(monkeypatch: pytest.MonkeyPatch) -> None:
agent: Final = AgentResponse(
agent_id="agent-123",
agent_name="Test Agent",
agent_card_params=_sample_agent_card_params_with_backend_url(),
)
test_client: Final = _client_with_agent_card(monkeypatch, LitellmUserRoles.INTERNAL_USER, agent)
http_get: Final = AsyncMock(return_value=httpx.Response(200))
http_client: Final = MagicMock(get=http_get)
monkeypatch.setattr(
agent_endpoints,
"get_async_httpx_client",
MagicMock(return_value=http_client),
)
response: Final = test_client.get("/v1/agents?health_check=true", headers={"Authorization": "Bearer k"})
assert response.status_code == 200, response.text
assert http_get.await_args.args[0] == _AGENT_CARD_BACKEND_URL
assert response.json()[0]["agent_card_params"]["url"] == "http://testserver/a2a/agent-123"
def test_non_admin_agent_card_without_interfaces_does_not_add_them(monkeypatch: pytest.MonkeyPatch) -> None:
agent: Final = _sample_agent_response()
test_client: Final = _client_with_agent_card(monkeypatch, LitellmUserRoles.INTERNAL_USER, agent)
response: Final = test_client.get("/v1/agents/agent-123", headers={"Authorization": "Bearer k"})
assert response.status_code == 200, response.text
assert response.json()["agent_card_params"]["url"] == "http://testserver/a2a/agent-123"
assert "supportedInterfaces" not in response.json()["agent_card_params"]
class TestCheckAgentUrlHealth:
"""Unit tests for the _check_agent_url_health helper."""
@ -1230,7 +1380,11 @@ def test_jwt_authentication_status_does_not_require_virtual_keys(
monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": enabled})
monkeypatch.setattr(proxy_server, "jwt_handler", handler)
agent: Final = _sample_agent_response()
response: Final = agent_endpoints._redact_sensitive_agent_fields((agent,), is_admin=True)[0]
response: Final = agent_endpoints._redact_sensitive_agent_fields(
(agent,),
is_admin=True,
request_base_url="http://testserver",
)[0]
assert response.jwt_auth_configured is expected
assert agent.jwt_auth_configured is False

67
tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -27,8 +27,9 @@ from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockGuardrail,
_redact_pii_matches,
)
from litellm.proxy.guardrails.guardrail_initializers import initialize_bedrock
from litellm.proxy.utils import ProxyLogging
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.guardrails import BedrockGuardrailOutputScope, Guardrail, GuardrailEventHooks, LitellmParams
from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockContentItem,
BedrockTextContent,
@ -1020,6 +1021,64 @@ def _anthropic_tool_result_conversation(
]
@pytest.mark.asyncio
@pytest.mark.parametrize("scope", [None, "INTERVENTIONS", "FULL"])
async def test_configured_output_scope_reaches_bedrock_and_preserves_healthy_assessments(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, scope: str | None
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
params: Final = LitellmParams.model_validate(
{
"guardrail": "bedrock",
"mode": "post_call",
"guardrailIdentifier": "test-guardrail",
"guardrailVersion": "DRAFT",
"contextual_grounding_from_messages": True,
"aws_access_key_id": "test-access-key",
"aws_secret_access_key": "test-secret-key",
"aws_region_name": "us-east-1",
"aws_bedrock_runtime_endpoint": "https://bedrock.example",
**({"outputScope": scope} if scope is not None else {}),
}
)
guardrail: Final = initialize_bedrock(params, Guardrail(guardrail_name="grounding", litellm_params=params))
assessment: Final = {
"contextualGroundingPolicy": {
"filters": [{"type": "GROUNDING", "score": 0.96, "threshold": 0.8, "detected": False, "action": "NONE"}]
}
}
def provider(request: httpx.Request) -> httpx.Response:
body: Final = json.loads(request.content)
return httpx.Response(
200,
json={"action": "NONE", "assessments": [assessment] if body.get("outputScope") == "FULL" else []},
)
route: Final = respx_mock.post("https://bedrock.example/guardrail/test-guardrail/version/DRAFT/apply").mock(
side_effect=provider
)
request_data: Final = {
"messages": [
{"role": "system", "content": "The tower is 330 metres tall."},
{"role": "user", "content": "How tall is the tower?"},
]
}
result: Final = await guardrail.apply_guardrail(
inputs={"texts": ["The tower is 330 metres tall."]}, request_data=request_data, input_type="response"
)
sent: Final = json.loads(route.calls.last.request.content)
if scope is None:
assert "outputScope" not in sent
else:
assert sent["outputScope"] == scope
assert sent["source"] == "OUTPUT"
assert result["texts"] == ["The tower is 330 metres tall."]
recorded: Final = request_data["metadata"]["standard_logging_guardrail_information"][-1]["guardrail_response"]
assert recorded["assessments"] == ([assessment] if scope == "FULL" else [])
@pytest.mark.asyncio
async def test_during_call_hook_skips_bedrock_call_for_tool_result_only_turn():
"""A tool_result-only latest user turn must not post an empty content list to Bedrock.
@ -1805,7 +1864,11 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs():
@pytest.mark.asyncio
async def test_make_bedrock_api_request_filters_dynamic_evaluation_overrides():
guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT")
guardrail = BedrockGuardrail(
guardrailIdentifier="test-guardrail",
guardrailVersion="DRAFT",
outputScope=BedrockGuardrailOutputScope.INTERVENTIONS,
)
mock_credentials = MagicMock()
mock_credentials.access_key = "test-access-key"
mock_credentials.secret_key = "test-secret-key"

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

854
tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix.py
View file

File diff suppressed because it is too large Load diff

443
tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py Normal file
View file

@ -0,0 +1,443 @@
import base64
from typing import Final
from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix_extraction import (
extract_file_parts_from_images,
extract_file_parts_from_messages,
extract_tool_results,
make_tool_data,
tool_call_to_tool_data,
tool_result_text_indices,
)
def _b64(raw: bytes) -> str:
return base64.b64encode(raw).decode()
def test_file_block_data_url_decoded():
msgs = [
{
"role": "user",
"content": [
{"type": "file", "file": {"filename": "a.txt", "file_data": f"data:text/plain;base64,{_b64(b'hi')}"}}
],
}
]
parts = extract_file_parts_from_messages(msgs, size_limit=1000)
assert len(parts) == 1 and parts[0].data == b"hi" and parts[0].name == "a.txt" and parts[0].inline
def test_image_url_reference_tracked_by_name():
msgs = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://x.test/pic.png"}}]}]
parts = extract_file_parts_from_messages(msgs, size_limit=1000)
assert len(parts) == 1 and parts[0].data is None and parts[0].inline is False and parts[0].name == "pic.png"
def test_oversize_file_flagged_no_data():
big = _b64(b"x" * 100)
msgs = [
{
"role": "user",
"content": [
{
"type": "file",
"file": {"filename": "big.bin", "file_data": f"data:application/octet-stream;base64,{big}"},
}
],
}
]
parts = extract_file_parts_from_messages(msgs, size_limit=10)
assert len(parts) == 1 and parts[0].oversize is True and parts[0].data is None
def test_images_field_data_url():
parts = extract_file_parts_from_images([f"data:image/png;base64,{_b64(b'png')}"], size_limit=1000)
assert len(parts) == 1 and parts[0].data == b"png" and parts[0].inline
def test_tool_call_to_tool_data_parses_arguments():
td = tool_call_to_tool_data(
{"id": "c1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "TLV"}'}}
)
assert (
td["content"] == '{"city": "TLV"}'
and td["tool_name"] == "get_weather"
and td["action_name"] == "get_weather"
and td["tool_input"] == {"city": "TLV"}
)
def test_tool_call_malformed_dropped():
assert tool_call_to_tool_data({"id": "c1", "type": "function", "function": {"name": ""}}) is None
assert tool_call_to_tool_data({"id": "c1"}) is None
def test_extract_tool_results_correlates_name():
msgs = [
{"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "get_weather"}}]},
{"role": "tool", "tool_call_id": "c1", "content": "sunny"},
{"role": "tool", "tool_call_id": "unknown", "content": "orphan"},
]
results = extract_tool_results(msgs)
assert ("get_weather", "sunny", "c1") in results
assert ("tool_result", "orphan", "unknown") in results
def test_make_tool_data_truncates_and_defaults_name():
td = make_tool_data(" ", "content")
assert td["tool_name"] == "tool_result" and td["action_name"] == "tool_result"
long = "x" * 200
td2 = make_tool_data(long, "c")
assert len(td2["tool_name"]) == 100 and td2["action_name"] == long
def test_unhashable_block_type_skipped_without_raising():
msgs = [{"role": "user", "content": [{"type": ["file"], "file": {"filename": "a.txt"}}]}]
parts = extract_file_parts_from_messages(msgs, size_limit=1000)
assert parts == ()
def test_unhashable_tool_call_id_skipped_without_raising():
msgs = [
{"role": "assistant", "tool_calls": [{"id": ["c1"], "function": {"name": "get_weather"}}]},
{"role": "tool", "tool_call_id": ["c1"], "content": "sunny"},
]
results = extract_tool_results(msgs)
assert results == (("tool_result", "sunny", ["c1"]),)
def test_extract_tool_results_list_form_content():
msgs = [{"role": "tool", "tool_call_id": "c1", "content": [{"type": "text", "text": "part1"}, "part2"]}]
results = extract_tool_results(msgs)
assert ("tool_result", "part1\npart2", "c1") in results
def test_extract_tool_results_dict_form_content():
msgs = [{"role": "tool", "tool_call_id": "c1", "content": {"city": "TLV"}}]
results = extract_tool_results(msgs)
assert ("tool_result", '{"city": "TLV"}', "c1") in results
def test_images_field_oversize_flagged_no_data():
big = _b64(b"x" * 100)
parts = extract_file_parts_from_images([f"data:image/png;base64,{big}"], size_limit=10)
assert len(parts) == 1 and parts[0].oversize is True and parts[0].data is None
class _StubFunction:
def __init__(self, name, arguments):
self.name = name
self.arguments = arguments
class _StubToolCall:
def __init__(self, function):
self.function = function
def test_tool_call_to_tool_data_accepts_object_style_tool_call():
tool_call = _StubToolCall(_StubFunction("get_weather", '{"city": "TLV"}'))
td = tool_call_to_tool_data(tool_call)
assert td["tool_name"] == "get_weather" and td["tool_input"] == {"city": "TLV"}
def _msgs(block):
return [{"role": "user", "content": [block]}]
def test_image_url_block_data_url_decoded_from_messages():
block = {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_b64(b'png')}"}}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].data == b"png" and parts[0].inline and parts[0].name is None
def test_image_url_block_non_string_url_skipped():
block = {"type": "image_url", "image_url": {"url": 123}}
assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == ()
def test_input_image_block_data_url_decoded():
block = {"type": "input_image", "image_url": f"data:image/png;base64,{_b64(b'img')}"}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].data == b"img" and parts[0].inline
def test_file_block_non_dict_file_skipped():
block = {"type": "file", "file": "not-a-dict"}
assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == ()
def test_file_block_reference_without_bytes_is_name_only():
block = {"type": "file", "file": {"file_id": "file-abc"}}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].name == "file-abc" and parts[0].data is None and parts[0].inline is False
def test_file_block_urlsafe_base64_decoded_via_fallback():
raw = b"\xff\xff\xfe" # encodes with url-unsafe chars '+'/'/' in standard b64
urlsafe = base64.urlsafe_b64encode(raw).decode()
assert "-" in urlsafe or "_" in urlsafe
block = {
"type": "file",
"file": {"filename": "b.bin", "file_data": f"data:application/octet-stream;base64,{urlsafe}"},
}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].data == raw
def test_file_block_data_url_without_base64_marker_is_undecodable():
block = {"type": "file", "file": {"filename": "n.txt", "file_data": "data:text/plain,hello"}}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].undecodable and parts[0].data is None and parts[0].name == "n.txt"
def test_input_file_block_data_url_decoded():
block = {"type": "input_file", "filename": "doc.pdf", "file_data": f"data:application/pdf;base64,{_b64(b'pdf')}"}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].data == b"pdf" and parts[0].name == "doc.pdf" and parts[0].inline
def test_input_file_block_file_url_reference_name_only():
block = {"type": "input_file", "file_url": "https://x.test/report.csv"}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].name == "report.csv" and parts[0].data is None and parts[0].inline is False
def test_input_audio_block_decoded():
block = {"type": "input_audio", "input_audio": {"data": _b64(b"wav"), "format": "wav"}}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].data == b"wav" and parts[0].name == "audio.wav" and parts[0].inline
def test_input_audio_block_undecodable_is_flagged():
block = {"type": "input_audio", "input_audio": {"data": "!!!not-base64!!!"}}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].undecodable and parts[0].data is None and parts[0].name == "audio.bin"
def test_input_audio_block_non_dict_skipped():
block = {"type": "input_audio", "input_audio": "nope"}
assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == ()
def test_tool_call_dict_arguments_serialized_and_parsed():
td = tool_call_to_tool_data({"function": {"name": "f", "arguments": {"a": 1}}})
assert td["content"] == '{"a": 1}' and td["tool_input"] == {"a": 1}
def test_tool_call_none_arguments_yields_empty_content():
td = tool_call_to_tool_data({"function": {"name": "f", "arguments": None}})
assert td["content"] == "" and td["tool_input"] == {}
def test_tool_call_non_string_non_dict_arguments_serialized():
td = tool_call_to_tool_data({"function": {"name": "f", "arguments": [1, 2]}})
assert td["content"] == "[1, 2]" and td["tool_input"] == {}
def test_tool_call_invalid_json_string_arguments_kept_as_content():
td = tool_call_to_tool_data({"function": {"name": "f", "arguments": "{not json"}})
assert td["content"] == "{not json" and td["tool_input"] == {}
def test_tool_result_with_non_string_tool_call_id_uses_default_name():
msgs = [{"role": "tool", "tool_call_id": ["c1"], "content": "orphan"}]
results = extract_tool_results(msgs)
assert results == (("tool_result", "orphan", ["c1"]),)
def test_images_field_http_url_is_name_only_reference():
parts = extract_file_parts_from_images(["https://x.test/pic.png"], size_limit=1000)
assert len(parts) == 1 and parts[0].name == "pic.png" and parts[0].data is None and parts[0].inline is False
def test_messages_skip_non_dict_and_unknown_blocks():
msgs = [
"not-a-message",
{
"role": "user",
"content": [
"bare-string-block",
{"type": "text", "text": "hi"},
{"type": "file", "file": {"filename": "a.txt", "file_data": f"data:text/plain;base64,{_b64(b'x')}"}},
],
},
]
parts = extract_file_parts_from_messages(msgs, size_limit=1000)
assert len(parts) == 1 and parts[0].name == "a.txt" and parts[0].data == b"x"
def test_image_url_block_invalid_data_url_is_flagged():
block = {"type": "image_url", "image_url": {"url": "data:image/png;base64,%%%invalid%%%"}}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].undecodable and parts[0].data is None
def test_tool_message_with_empty_content_is_skipped():
msgs = [{"role": "tool", "tool_call_id": "c1", "content": " "}]
assert extract_tool_results(msgs) == ()
def test_extract_tool_results_skips_non_dict_messages_and_tool_calls():
msgs = [
"junk",
{"role": "assistant", "tool_calls": ["not-a-dict", {"id": "c1", "function": {"name": "f"}}]},
{"role": "tool", "tool_call_id": "c1", "content": "ok"},
]
assert extract_tool_results(msgs) == (("f", "ok", "c1"),)
def test_tool_result_takes_its_name_from_the_nearest_earlier_declaration():
msgs: Final = [
{"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "first"}}]},
{"role": "tool", "tool_call_id": "c1", "content": "r1"},
{
"role": "assistant",
"tool_calls": [{"id": "c1", "function": {"name": "second"}}, {"id": "c1", "function": {"name": "third"}}],
},
{"role": "tool", "tool_call_id": "c1", "content": "r2"},
{"role": "tool", "tool_call_id": "c2", "content": "r3"},
{"role": "assistant", "tool_calls": [{"id": "c2", "function": {"name": "too-late"}}]},
]
assert extract_tool_results(msgs) == (("first", "r1", "c1"), ("third", "r2", "c1"), ("tool_result", "r3", "c2"))
class _ReadCountingMessage(dict):
def __init__(self, **fields):
super().__init__(**fields)
self.reads = 0
def get(self, key, default=None):
self.reads += 1
return super().get(key, default)
def _reads_of_a_declaration_followed_by(tool_results):
declaration: Final = _ReadCountingMessage(role="assistant", tool_calls=[{"id": "c0", "function": {"name": "f"}}])
tool_messages: Final = [
{"role": "tool", "tool_call_id": f"c{index}", "content": "x"} for index in range(tool_results)
]
extract_tool_results([declaration, *tool_messages])
return declaration.reads
def test_tool_result_names_resolve_without_rescanning_earlier_messages():
assert _reads_of_a_declaration_followed_by(1000) == _reads_of_a_declaration_followed_by(1)
def test_malformed_data_url_is_undecodable():
block = {"type": "file", "file": {"filename": "x.bin", "file_data": "data:garbage-no-comma"}}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].undecodable and parts[0].data is None and parts[0].name == "x.bin"
def test_input_audio_block_without_data_skipped():
block = {"type": "input_audio", "input_audio": {"format": "wav"}}
assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == ()
def test_images_field_non_string_entries_skipped():
assert extract_file_parts_from_images([123, None, ""], size_limit=1000) == ()
def test_make_tool_data_whitespace_after_truncation_defaults_name():
name = " " * 100 + "x"
td = make_tool_data(name, "c")
assert td["tool_name"] == "tool_result" and td["action_name"] == name
def test_raw_base64_file_data_without_data_url_prefix_decoded():
block = {"type": "file", "file": {"filename": "a.bin", "file_data": _b64(b"rawbytes")}}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].data == b"rawbytes" and parts[0].mime_hint is None
def test_images_field_raw_base64_without_data_url_prefix_decoded():
parts = extract_file_parts_from_images([_b64(b"rawimg")], size_limit=1000)
assert len(parts) == 1 and parts[0].data == b"rawimg"
def test_whitespace_only_base64_payload_yields_no_bytes():
block = {"type": "file", "file": {"filename": "n.txt", "file_data": "data:text/plain;base64, "}}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].data is None and parts[0].inline is False and parts[0].name == "n.txt"
def test_base64_invalid_even_after_urlsafe_translate_is_undecodable():
block = {"type": "file", "file": {"filename": "b.bin", "file_data": "_@@@"}}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
assert len(parts) == 1 and parts[0].undecodable and parts[0].data is None and parts[0].name == "b.bin"
def test_file_oversize_detected_after_decode_when_estimate_passes():
block = {"type": "file", "file": {"filename": "s.bin", "file_data": _b64(b"abc")}}
parts = extract_file_parts_from_messages(_msgs(block), size_limit=2)
assert len(parts) == 1 and parts[0].oversize is True and parts[0].data is None
def test_image_http_url_that_urlparse_rejects_is_dropped():
parts = extract_file_parts_from_images(["http://["], size_limit=1000)
assert parts == ()
def test_message_content_not_a_list_is_skipped():
parts = extract_file_parts_from_messages(
[{"role": "user", "content": "just a plain string prompt"}], size_limit=1000
)
assert parts == ()
def test_tool_call_dict_arguments_non_serializable_falls_back_to_str():
td = tool_call_to_tool_data({"function": {"name": "f", "arguments": {"a": {1, 2}}}})
assert isinstance(td["content"], str) and td["tool_input"] == {"a": {1, 2}}
def test_tool_call_non_serializable_other_arguments_falls_back_to_str():
td = tool_call_to_tool_data({"function": {"name": "f", "arguments": {1, 2}}})
assert isinstance(td["content"], str) and td["tool_input"] == {}
def test_tool_result_dict_content_non_serializable_falls_back_to_str():
results = extract_tool_results([{"role": "tool", "tool_call_id": "c1", "content": {"x": {1, 2}}}])
assert len(results) == 1 and isinstance(results[0][1], str) and results[0][1].strip()
def test_tool_result_text_indices_marks_only_tool_role_positions():
messages = [
{"role": "user", "content": "ask"},
{"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "f"}}]},
{"role": "tool", "tool_call_id": "c1", "content": "result"},
]
assert tool_result_text_indices(messages, ["ask", "result"]) == frozenset({1})
def test_tool_result_text_indices_covers_every_text_block_of_a_tool_message():
messages = [
{"role": "user", "content": "ask"},
{
"role": "tool",
"tool_call_id": "c1",
"content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}],
},
]
assert tool_result_text_indices(messages, ["ask", "a", "b"]) == frozenset({1, 2})
def test_tool_result_text_indices_empty_when_texts_do_not_replay_messages():
messages = [
{"role": "user", "content": "ask"},
{"role": "tool", "tool_call_id": "c1", "content": "result"},
]
assert tool_result_text_indices(messages, ["result"]) == frozenset()
assert tool_result_text_indices(messages, ["ask", "tampered"]) == frozenset()
def test_tool_result_text_indices_skips_blank_tool_content_never_submitted_as_tool():
messages = [{"role": "tool", "tool_call_id": "c1", "content": " "}]
assert extract_tool_results(messages) == ()
assert tool_result_text_indices(messages, [" "]) == frozenset()
def test_tool_result_text_indices_empty_without_structured_messages():
assert tool_result_text_indices(None, ["ask"]) == frozenset()

220
tests/unit/proxy/guardrails/test_guardrail_endpoints.py
View file

@ -3,17 +3,20 @@ import time
from collections.abc import Callable
from datetime import datetime
from types import MappingProxyType
from typing import Dict, Final, List, Optional
from typing import Dict, Final, List, Literal, Optional
from unittest.mock import AsyncMock
import httpx
import pytest
import yaml
from fastapi import HTTPException
from fastapi import FastAPI, HTTPException
from pydantic import JsonValue
from pytest_mock import MockerFixture
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
from litellm.proxy.guardrails.content_filter_data import DATA_ROOTS
from litellm.proxy.guardrails.guardrail_endpoints import (
CreateGuardrailRequest,
ListGuardrailsResponse,
@ -37,7 +40,6 @@ from litellm.proxy.guardrails.guardrail_endpoints import (
reject_guardrail_submission,
update_guardrail,
)
from litellm.proxy.guardrails.content_filter_data import DATA_ROOTS
from litellm.proxy.guardrails.guardrail_endpoints import (
test_custom_code_guardrail as run_custom_code_test_endpoint,
)
@ -53,8 +55,8 @@ from litellm.proxy.guardrails.guardrail_registry import (
from litellm.types.guardrails import (
ApplyGuardrailRequest,
BaseLitellmParams,
GuardrailEventHooks,
Guardrail,
GuardrailEventHooks,
GuardrailInfoResponse,
LitellmParams,
)
@ -2198,6 +2200,205 @@ def _patch_apply_guardrail_env(mocker, guardrail_result):
return mock_guardrail
class _NativeFindingGuardrail(CustomGuardrail):
def __init__(self, findings: tuple[object, ...]):
super().__init__(guardrail_name="grounding", event_hook=GuardrailEventHooks.post_call)
self.findings = findings
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: Logging | None = None,
) -> GenericGuardrailAPIInputs:
for finding in self.findings:
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=finding, request_data=request_data, guardrail_status="success"
)
return inputs
class _InputDirectionGuardrail(_NativeFindingGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: Logging | None = None,
) -> GenericGuardrailAPIInputs:
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"input_type": input_type},
request_data=request_data,
guardrail_status="success",
)
return {"texts": [input_type]}
async def _apply_finding_guardrail_over_http(
mocker: MockerFixture,
findings: tuple[object, ...],
body: dict[str, object],
hook_result: object = None,
guardrail: CustomGuardrail | None = None,
) -> httpx.Response:
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.guardrails.guardrail_endpoints import router
from litellm.proxy.proxy_server import openai_exception_handler
configured: Final = guardrail if guardrail is not None else _NativeFindingGuardrail(findings)
patched: Final = _patch_apply_guardrail_env(mocker, {})
patched.apply_guardrail = configured.apply_guardrail
patched.event_hook = configured.event_hook
mocker.patch(
"litellm.proxy.proxy_server.proxy_logging_obj.post_call_success_hook", AsyncMock(return_value=hook_result)
)
app: Final = FastAPI()
app.include_router(router)
app.add_exception_handler(ProxyException, openai_exception_handler)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth()
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://guardrails.test") as client:
return await client.post("/guardrails/apply_guardrail", json=body)
@pytest.mark.asyncio
@pytest.mark.parametrize("is_grounded", [True, False])
async def test_apply_guardrail_native_response_preserves_structured_findings(
mocker: MockerFixture, is_grounded: bool
) -> None:
findings: Final = {"is_grounded": is_grounded, "reason": "Compared answer with evidence", "source_ids": ["doc-1"]}
result: Final = await _apply_finding_guardrail_over_http(
mocker,
(findings,),
{"guardrail_name": "grounding", "text": "answer", "response_format": "native"},
)
assert result.status_code == 200
assert result.json() == findings
@pytest.mark.asyncio
async def test_apply_guardrail_standard_response_stays_unchanged_with_native_findings(mocker: MockerFixture) -> None:
result: Final = await _apply_finding_guardrail_over_http(
mocker, ({"is_grounded": False},), {"guardrail_name": "grounding", "text": "answer"}
)
assert result.status_code == 200
assert result.json() == {"response_text": "answer"}
@pytest.mark.asyncio
async def test_apply_guardrail_native_preserves_all_provider_results(mocker: MockerFixture) -> None:
findings: Final = ({"is_grounded": False}, {"response_text": "native text", "additional_finding": [1, 2]})
result: Final = await _apply_finding_guardrail_over_http(
mocker, findings, {"guardrail_name": "grounding", "text": "answer", "response_format": "native"}
)
assert result.status_code == 200
assert result.json() == list(findings)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("finding", "expected"),
[
({"is_grounded": True, "source_ids": ("doc-1",)}, {"is_grounded": True, "source_ids": ["doc-1"]}),
({"response_text": "native", "score": 0.2}, {"response_text": "native", "score": 0.2}),
([{"is_grounded": False}], [{"is_grounded": False}]),
],
)
async def test_apply_guardrail_native_preserves_provider_json_shape(
mocker: MockerFixture, finding: object, expected: JsonValue
) -> None:
result: Final = await _apply_finding_guardrail_over_http(
mocker,
(finding,),
{"guardrail_name": "grounding", "text": "answer", "response_format": "native"},
)
assert result.status_code == 200
assert result.json() == expected
@pytest.mark.asyncio
@pytest.mark.parametrize("findings", [(), ("allow",)])
async def test_apply_guardrail_native_rejects_missing_findings_and_caller_supplied_verdicts(
mocker: MockerFixture, findings: tuple[JsonValue, ...]
) -> None:
result: Final = await _apply_finding_guardrail_over_http(
mocker,
findings,
{
"guardrail_name": "grounding",
"text": "answer",
"response_format": "native",
"metadata": {
"standard_logging_guardrail_information": [
{"guardrail_name": "grounding", "guardrail_response": {"is_grounded": True}}
]
},
},
)
assert result.status_code == 422
assert "did not report structured native findings" in result.text
@pytest.mark.asyncio
@pytest.mark.parametrize("response_format", ["standard", "native"])
async def test_apply_guardrail_omitted_direction_uses_configured_post_call_hook(
mocker: MockerFixture, response_format: str
) -> None:
result: Final = await _apply_finding_guardrail_over_http(
mocker,
(),
{"guardrail_name": "grounding", "text": "answer", "response_format": response_format},
guardrail=_InputDirectionGuardrail(()),
)
assert result.status_code == 200
assert result.json() == (
{"response_text": "response"} if response_format == "standard" else {"input_type": "response"}
)
@pytest.mark.asyncio
async def test_apply_guardrail_native_honors_explicit_input_direction(mocker: MockerFixture) -> None:
result: Final = await _apply_finding_guardrail_over_http(
mocker,
(),
{"guardrail_name": "grounding", "text": "prompt", "input_type": "request", "response_format": "native"},
guardrail=_InputDirectionGuardrail(()),
)
assert result.status_code == 200
assert result.json() == {"input_type": "request"}
@pytest.mark.asyncio
@pytest.mark.parametrize("replacement", [{"is_grounded": False, "reason": "Revised"}, [{"score": 0.5}]])
async def test_apply_guardrail_native_preserves_success_hook_replacements(
mocker: MockerFixture, replacement: JsonValue
) -> None:
result: Final = await _apply_finding_guardrail_over_http(
mocker,
({"is_grounded": True},),
{"guardrail_name": "grounding", "text": "answer", "response_format": "native"},
hook_result=replacement,
)
assert result.status_code == 200
assert result.json() == replacement
@pytest.mark.asyncio
async def test_apply_guardrail_native_ignores_malformed_caller_logging_records(mocker: MockerFixture) -> None:
result: Final = await _apply_finding_guardrail_over_http(
mocker,
({"is_grounded": False},),
{
"guardrail_name": "grounding",
"text": "answer",
"response_format": "native",
"metadata": {"standard_logging_guardrail_information": "caller-controlled"},
},
)
assert result.status_code == 200
assert result.json() == {"is_grounded": False}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("vendor_status", "expected_type"),
@ -2212,9 +2413,7 @@ async def test_a_guardrail_vendors_rejection_is_labelled_the_way_chat_labels_it(
mock_guardrail.apply_guardrail.side_effect = HTTPException(
status_code=vendor_status, detail="vendor refused the request"
)
mocker.patch(
"litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", AsyncMock(return_value=None)
)
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", AsyncMock(return_value=None))
with pytest.raises(ProxyException) as raised:
await apply_guardrail(
@ -3673,6 +3872,7 @@ async def test_guardrail_success_delivers_sync_callback_once_with_overlapping_wo
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils import litellm_logging, thread_pool_executor
from litellm.proxy.guardrails.guardrail_endpoints import _emit_guardrail_success_logs
from litellm.proxy.guardrails.response_handlers import StandardGuardrailResponseHandler
from litellm.types.guardrails import ApplyGuardrailResponse
class OverlappingHook(CustomLogger):
@ -3711,7 +3911,7 @@ async def test_guardrail_success_delivers_sync_callback_once_with_overlapping_wo
response: Final = ApplyGuardrailResponse(response_text="masked")
returned: Final = await _emit_guardrail_success_logs(
proxy_logging, logging_obj, {}, UserAPIKeyAuth(), response, start
StandardGuardrailResponseHandler(), proxy_logging, logging_obj, {}, UserAPIKeyAuth(), response, start
)
submissions: Final = tuple(queued_executor.submit.call_args_list)
assert submissions

216
tests/unit/proxy/test_dev_reload.py Normal file
View file

@ -0,0 +1,216 @@
import asyncio
import importlib
import logging
import logging.handlers
import traceback
from pathlib import Path
from typing import Final
import pytest
from starlette.types import Message, Scope
from uvicorn import config as uvicorn_config
from litellm.proxy.dev_reload import (
RELOAD_OPTIONS,
RELOAD_WORKER_HEALTHCHECK_SECONDS,
ReloadOptions,
load_app_or_exit,
load_worker_app,
multi_worker_reload_supported,
pool_config,
)
from litellm.proxy.proxy_cli import ProxyInitializationHelpers
UNCHANGED_LOGGING: Final = {"version": 1, "incremental": True}
POOL_WORKERS: Final = 2
REQUIRES_UVICORN_STARTUP_FAILURE: Final = pytest.mark.skipif(
not multi_worker_reload_supported(), reason="uvicorn<0.50.0 has no startup failure exit code"
)
def _options(app: str) -> ReloadOptions:
return {
"app": app,
"log_config": UNCHANGED_LOGGING,
"reload": True,
"reload_dirs": (),
"reload_includes": ("*.py",),
}
@REQUIRES_UVICORN_STARTUP_FAILURE
@pytest.mark.parametrize(
("module", "source"),
(
("reload_syntax_error_app", "def app(\n"),
("reload_name_error_app", "app = undefined_name\n"),
("reload_runtime_error_app", "raise RuntimeError('broken edit')\n"),
("reload_missing_attribute_app", "value = 1\n"),
),
)
def test_a_broken_app_exits_the_worker_as_a_startup_failure(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, module: str, source: str
) -> None:
(tmp_path / f"{module}.py").write_text(source)
monkeypatch.syspath_prepend(str(tmp_path))
with pytest.raises(SystemExit) as exit_info:
pool_config(_options(f"{module}:app"), workers=2).load()
assert exit_info.value.code == uvicorn_config.STARTUP_FAILURE
@REQUIRES_UVICORN_STARTUP_FAILURE
@pytest.mark.parametrize(
("module", "source"),
(
(
"reload_factory_missing_module_app",
"def factory():\n import a_module_that_does_not_exist\n",
),
(
"reload_factory_runtime_error_app",
"def factory():\n raise RuntimeError('broken edit')\n",
),
),
)
def test_a_broken_factory_exits_the_worker_as_a_startup_failure(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, module: str, source: str
) -> None:
(tmp_path / f"{module}.py").write_text(source)
monkeypatch.syspath_prepend(str(tmp_path))
with pytest.raises(SystemExit) as exit_info:
pool_config(
_options(f"{module}:app"),
workers=POOL_WORKERS,
worker_overrides={"app": f"{module}:factory", "factory": True},
).load()
assert exit_info.value.code == uvicorn_config.STARTUP_FAILURE
@REQUIRES_UVICORN_STARTUP_FAILURE
@pytest.mark.parametrize(
("module", "source"),
(
("reload_import_error_app", "import os\nfrom os import missing_name\n"),
("reload_value_error_app", "import os\nraise ValueError('broken edit')\n"),
),
)
def test_a_broken_import_logs_the_line_that_raised(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, module: str, source: str
) -> None:
app_file: Final = tmp_path / f"{module}.py"
app_file.write_text(source)
monkeypatch.syspath_prepend(str(tmp_path))
uvicorn_logger: Final = logging.getLogger("uvicorn.error")
captured: Final = logging.handlers.BufferingHandler(capacity=POOL_WORKERS)
uvicorn_logger.addHandler(captured)
try:
with pytest.raises(SystemExit):
load_app_or_exit(f"{module}:app")
finally:
uvicorn_logger.removeHandler(captured)
(record,) = captured.buffer
assert record.exc_info is not None
raised_at: Final = frozenset((frame.filename, frame.lineno) for frame in traceback.extract_tb(record.exc_info[2]))
assert (str(app_file), 2) in raised_at, record.exc_text
def test_the_reload_pool_defaults_the_worker_healthcheck_and_keeps_an_explicit_one() -> None:
explicit: Final = RELOAD_WORKER_HEALTHCHECK_SECONDS + 7
assert pool_config(_options("app:app"), workers=2).timeout_worker_healthcheck == RELOAD_WORKER_HEALTHCHECK_SECONDS
assert (
pool_config({**_options("app:app"), "timeout_worker_healthcheck": explicit}, workers=2).timeout_worker_healthcheck
== explicit
)
def test_the_reload_pool_accepts_and_keeps_every_uvicorn_arg_the_proxy_cli_builds(tmp_path: Path) -> None:
cli_args: Final = ProxyInitializationHelpers.get_default_unvicorn_init_args(
"127.0.0.1", 4123, keepalive_timeout=61, timeout_worker_healthcheck=RELOAD_WORKER_HEALTHCHECK_SECONDS + 7
)
options: Final = RELOAD_OPTIONS.validate_python(
{
**cli_args,
"log_config": UNCHANGED_LOGGING,
"reload": True,
"reload_dirs": [str(tmp_path)],
"reload_includes": ["*.py"],
}
)
config: Final = pool_config(options, workers=POOL_WORKERS)
server_args: Final = {key: value for key, value in cli_args.items() if key not in ("app", "log_config")}
assert {key: getattr(config, key) for key in server_args} == server_args
def test_the_reload_pool_passes_a_uvicorn_arg_it_does_not_list_through_to_the_workers() -> None:
config: Final = pool_config(RELOAD_OPTIONS.validate_python({**_options("app:app"), "backlog": 123}), workers=POOL_WORKERS)
assert (config.backlog, config.reload) == (123, False)
def test_the_reload_pool_serves_the_overridden_worker_app(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
module: Final = "reload_worker_override_app"
(tmp_path / f"{module}.py").write_text(
"async def base_app(scope, receive, send):\n"
" await receive()\n"
" await send({'type': 'http.response.start', 'status': 200, 'headers': []})\n"
" await send({'type': 'http.response.body', 'body': b'base app', 'more_body': False})\n"
"\n"
"async def overridden_app(scope, receive, send):\n"
" await receive()\n"
" await send({'type': 'http.response.start', 'status': 200, 'headers': []})\n"
" await send({'type': 'http.response.body', 'body': b'overridden app', 'more_body': False})\n"
"\n"
"def base_factory():\n"
" return base_app\n"
"\n"
"def factory():\n"
" return overridden_app\n"
)
monkeypatch.syspath_prepend(str(tmp_path))
config: Final = pool_config(
_options(f"{module}:base_factory"),
workers=POOL_WORKERS,
worker_overrides={"app": f"{module}:factory", "factory": True, "log_config": None},
)
assert config.log_config is None
config.load()
loaded_app: Final = config.loaded_app
assert loaded_app is not None
scope: Final[Scope] = {"type": "http", "method": "GET", "path": "/", "headers": []}
async def run_app() -> tuple[Message, ...]:
messages: Final[asyncio.Queue[Message]] = asyncio.Queue()
async def receive() -> Message:
return {"type": "http.request", "body": b"", "more_body": False}
async def send(message: Message) -> None:
await messages.put(message)
await loaded_app(scope, receive, send)
return tuple(messages.get_nowait() for _ in range(messages.qsize()))
messages: Final = asyncio.run(run_app())
assert messages == (
{"type": "http.response.start", "status": 200, "headers": []},
{"type": "http.response.body", "body": b"overridden app", "more_body": False},
)
def test_load_worker_app_calls_the_overridden_factory(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
module: Final = "reload_worker_app_factory"
(tmp_path / f"{module}.py").write_text(
"def app(scope, receive, send):\n return None\n\ndef factory():\n return app\n"
)
monkeypatch.syspath_prepend(str(tmp_path))
app: Final = importlib.import_module(module).app
assert load_worker_app(f"{module}:factory", worker_app_factory=True) is app
def test_windows_keeps_single_worker_reload_whatever_the_uvicorn_version() -> None:
assert not multi_worker_reload_supported("win32")
def test_a_working_app_loads_as_the_module_defines_it(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
(tmp_path / "reload_working_app.py").write_text("async def app(scope, receive, send):\n return None\n")
monkeypatch.syspath_prepend(str(tmp_path))
assert load_app_or_exit("reload_working_app:app") is importlib.import_module("reload_working_app").app

168
tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
View file

@ -689,16 +689,11 @@ class TestProxySettingEndpoints:
# Verify settings were updated
settings = data["settings"]
assert settings["google_client_id"] == new_sso_settings["google_client_id"]
assert (
settings["google_client_secret"] == new_sso_settings["google_client_secret"]
)
assert settings["google_client_secret"] == "new_****************cret"
assert (
settings["microsoft_client_id"] == new_sso_settings["microsoft_client_id"]
)
assert (
settings["microsoft_client_secret"]
== new_sso_settings["microsoft_client_secret"]
)
assert settings["microsoft_client_secret"] == "new_*******************cret"
assert settings["proxy_base_url"] == new_sso_settings["proxy_base_url"]
assert settings["user_email"] == new_sso_settings["user_email"]
@ -722,6 +717,163 @@ class TestProxySettingEndpoints:
create_sso_settings = json.loads(create_data["sso_settings"])
assert create_sso_settings["google_client_id"] == "new_google_client_id"
def _patch_sso_settings(self, monkeypatch, stored, body, environment_variables=None):
"""PATCH /update/sso_settings against a stored row; returns (response, saved row)."""
import json
from unittest.mock import AsyncMock, MagicMock
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
mock_prisma = MagicMock()
stored_record = None
if stored is not None:
stored_record = MagicMock()
stored_record.sso_settings = stored
mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(
return_value=stored_record
)
mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock()
env_var_entry = None
if environment_variables is not None:
env_var_entry = MagicMock()
env_var_entry.param_value = environment_variables
mock_prisma.writer_db.litellm_config.find_unique = AsyncMock(return_value=env_var_entry)
mock_prisma.db.litellm_config.update = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
from litellm.proxy.proxy_server import proxy_config
monkeypatch.setattr(
proxy_config, "encrypt_env_variables", lambda environment_variables: environment_variables
)
monkeypatch.setattr(proxy_config, "decrypt_db_variables", lambda variables: variables)
response = client.patch("/update/sso_settings", json=body)
upsert_kwargs = mock_prisma.db.litellm_ssoconfig.upsert.call_args.kwargs
return response, json.loads(upsert_kwargs["data"]["update"]["sso_settings"])
_STORED_GENERIC_SSO = {
"generic_client_id": "stored-client-id",
"generic_client_secret": "stored-client-secret",
"generic_authorization_endpoint": "https://idp.example.com/authorize",
"user_email": "admin@example.com",
"role_mappings": {
"provider": "generic",
"group_claim": "groups",
"default_role": "internal_user",
"roles": {"proxy_admin": ["admins"]},
},
}
def test_update_sso_settings_omitted_fields_keep_stored_values(
self, mock_proxy_config, mock_auth, monkeypatch
):
monkeypatch.setenv("GENERIC_CLIENT_SECRET", "stored-client-secret")
new_role_mappings = {
"provider": "generic",
"group_claim": "roles",
"default_role": "internal_user",
"roles": {"proxy_admin": ["platform-admins"]},
}
response, saved = self._patch_sso_settings(
monkeypatch, self._STORED_GENERIC_SSO, {"role_mappings": new_role_mappings}
)
assert response.status_code == 200
assert saved["generic_client_secret"] == "stored-client-secret"
assert saved["generic_client_id"] == "stored-client-id"
assert saved["user_email"] == "admin@example.com"
assert saved["role_mappings"]["group_claim"] == "roles"
assert os.environ["GENERIC_CLIENT_SECRET"] == "stored-client-secret"
def test_update_sso_settings_null_clears_only_that_field(
self, mock_proxy_config, mock_auth, monkeypatch
):
monkeypatch.setenv("GENERIC_CLIENT_SECRET", "stored-client-secret")
monkeypatch.setenv("GENERIC_CLIENT_ID", "stored-client-id")
response, saved = self._patch_sso_settings(
monkeypatch,
self._STORED_GENERIC_SSO,
{"generic_client_secret": None, "role_mappings": None},
)
assert response.status_code == 200
assert saved["generic_client_secret"] is None
assert saved["role_mappings"] is None
assert saved["generic_client_id"] == "stored-client-id"
assert "GENERIC_CLIENT_SECRET" not in os.environ
assert os.environ["GENERIC_CLIENT_ID"] == "stored-client-id"
def test_update_sso_settings_new_value_replaces_stored_value(
self, mock_proxy_config, mock_auth, monkeypatch
):
response, saved = self._patch_sso_settings(
monkeypatch, self._STORED_GENERIC_SSO, {"generic_client_secret": "rotated-secret"}
)
assert response.status_code == 200
assert saved["generic_client_secret"] == "rotated-secret"
assert saved["generic_client_id"] == "stored-client-id"
assert os.environ["GENERIC_CLIENT_SECRET"] == "rotated-secret"
def test_update_sso_settings_omitted_field_leaves_env_only_value_alone(
self, mock_proxy_config, mock_auth, monkeypatch
):
monkeypatch.setenv("GENERIC_CLIENT_SECRET", "env-only-secret")
response, saved = self._patch_sso_settings(
monkeypatch, None, {"user_email": "admin@example.com"}
)
assert response.status_code == 200
assert saved["user_email"] == "admin@example.com"
assert saved["generic_client_secret"] is None
assert os.environ["GENERIC_CLIENT_SECRET"] == "env-only-secret"
def test_update_sso_settings_response_masks_stored_secrets(
self, mock_proxy_config, mock_auth, monkeypatch
):
response, saved = self._patch_sso_settings(
monkeypatch, self._STORED_GENERIC_SSO, {}
)
assert response.status_code == 200
assert saved["generic_client_secret"] == "stored-client-secret"
returned = response.json()["settings"]
assert returned["generic_client_id"] == "stored-client-id"
assert returned["generic_client_secret"] != "stored-client-secret"
assert "stored-client-secret" not in response.text
def test_update_sso_settings_moves_omitted_env_row_values_into_sso_row(
self, mock_proxy_config, mock_auth, monkeypatch
):
import json
environment_variables = {
"MICROSOFT_TENANT": "env-row-tenant",
"MICROSOFT_CLIENT_ID": "env-row-client-id",
"GENERIC_CLIENT_ID": "env-row-generic-id",
"LANGFUSE_PUBLIC_KEY": "unrelated",
}
response, saved = self._patch_sso_settings(
monkeypatch,
self._STORED_GENERIC_SSO,
{"user_email": "new@example.com", "microsoft_client_id": None},
environment_variables=environment_variables,
)
from litellm.proxy.proxy_server import prisma_client
assert response.status_code == 200
assert saved["microsoft_tenant"] == "env-row-tenant"
assert saved["microsoft_client_id"] is None
assert saved["generic_client_id"] == "stored-client-id"
assert saved["user_email"] == "new@example.com"
env_row_update = prisma_client.db.litellm_config.update.call_args.kwargs
assert json.loads(env_row_update["data"]["param_value"]) == {"LANGFUSE_PUBLIC_KEY": "unrelated"}
def test_update_sso_settings_maps_saml_fields_to_env_vars(
self, mock_proxy_config, mock_auth, monkeypatch
):
@ -804,7 +956,7 @@ class TestProxySettingEndpoints:
mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None)
mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock()
mock_prisma.db.litellm_config = MagicMock()
mock_prisma.db.litellm_config.find_unique = AsyncMock(
mock_prisma.writer_db.litellm_config.find_unique = AsyncMock(
side_effect=ValueError("cleanup failed")
)
mock_prisma.db.litellm_config.update = AsyncMock()

66
tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
View file

@ -467,7 +467,6 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior():
assert evt is not None
assert evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED
assert evt.output_index == 0
assert hasattr(evt, "__dict__") and "sequence_number" in evt.__dict__
# Collect all remaining delta events from the pending queue by creating empty chunks
delta_events = []
@ -499,17 +498,11 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior():
assert len(evt.delta) <= 10
assert evt.item_id == "fc_call_test"
assert evt.output_index == 0
assert hasattr(evt, "__dict__") and "sequence_number" in evt.__dict__
# Verify all deltas concatenated equal the original arguments
concatenated = "".join(evt.delta for evt in delta_events)
assert concatenated == large_arguments
# Verify sequence numbers are increasing
sequence_numbers = [evt.__dict__["sequence_number"] for evt in delta_events]
assert sequence_numbers == sorted(sequence_numbers)
assert len(set(sequence_numbers)) == len(sequence_numbers) # All unique
def test_tool_call_delta_without_id_uses_index_mapping():
iterator = LiteLLMCompletionStreamingIterator(
@ -1947,6 +1940,65 @@ def test_the_partial_response_billed_on_disconnect_carries_the_logging_objects_b
), "the bill is priced on the partial's own usage"
def _serialized_sequence_numbers(events: Iterable[BaseLiteLLMOpenAIResponseObject]) -> list[int]:
payloads: Final = tuple(json.loads(event.model_dump_json()) for event in events)
assert all("sequence_number" in payload for payload in payloads), [
payload["type"] for payload in payloads if "sequence_number" not in payload
]
return [payload["sequence_number"] for payload in payloads]
def _long_text_chunks() -> tuple[ModelResponseStream, ...]:
return (*(_chunk(f"word{i} ") for i in range(300)), _chunk("end", finish_reason="stop"))
def _text_then_tool_call_chunks() -> tuple[ModelResponseStream, ...]:
return (_text_opening_a_tool_call_chunk("Checking. "), _tool_arguments_chunk('{"x": 1}', finish_reason="tool_calls"))
@pytest.mark.parametrize(
"stream",
[
lambda: (_chunk("Hello"), _chunk("!", finish_reason="stop")),
_long_text_chunks,
_text_then_tool_call_chunks,
_reasoning_then_text_chunks,
_reasoning_and_text_in_one_chunk,
_reasoning_citation_then_tool_call_finishing_chunks,
lambda: (_tool_call_chunk(), _chunk("", finish_reason="tool_calls")),
lambda: (_refusal_chunk("I can't "), _refusal_chunk("help."), _chunk("", finish_reason="stop")),
lambda: (_chunk("Sure. "), _refusal_chunk("I can't help."), _chunk("", finish_reason="stop")),
],
ids=[
"text",
"long_text",
"text_then_tool",
"reasoning_then_text",
"reasoning_and_text",
"citation_tool",
"tool",
"refusal",
"text_then_refusal",
],
)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_every_event_serializes_its_position_in_the_stream_from_zero(stream: StreamChunks, sync_mode: bool):
iterator: Final = _build_iterator(stream())
events: Final = list(iterator) if sync_mode else [event async for event in iterator]
assert _serialized_sequence_numbers(events) == list(range(len(events)))
def test_bridged_events_answer_dict_reads_of_their_number():
events: Final = list(_build_iterator((_chunk("Hello"), _chunk("!", finish_reason="stop"))))
positions: Final = list(range(len(events)))
assert [event["sequence_number"] for event in events] == positions
assert [event.get("sequence_number") for event in events] == positions
assert all("sequence_number" in event for event in events)
def _refusal_chunk(refusal: str, finish_reason: str | None = None) -> ModelResponseStream:
return ModelResponseStream(
id=CHAT_COMPLETION_ID,

51
tests/unit/responses/test_streaming_iterator.py
View file

@ -2867,3 +2867,54 @@ async def test_websocket_turns_logging_survives_garbage_collection(terminal_even
await _run_websocket_turn(logger, terminal_event)
assert await _logging_tasks_left_after_garbage_collection(logger) == 1
@pytest.mark.parametrize(
"output",
[
[
{
"type": "message",
"id": "msg_seq",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "hello world", "annotations": []}],
}
],
[
{
"type": "reasoning",
"id": "rs_seq",
"summary": [{"type": "summary_text", "text": "thinking about it"}],
},
{
"type": "function_call",
"id": "fc_seq",
"call_id": "call_seq",
"name": "lookup",
"arguments": '{"id":1}',
},
],
],
ids=["message", "reasoning_and_tool"],
)
def test_synthetic_stream_serializes_every_event_position_from_zero(output: list[dict[str, object]]):
transformed: Final = ResponsesAPIResponse.model_validate(
{
"id": "resp_seq",
"created_at": 1741476542,
"status": "completed",
"model": "claude-sonnet-5-5",
"object": "response",
"output": output,
}
)
events: Final = streaming_module.build_synthetic_response_events(
transformed=transformed,
logging_obj=None,
chunk_size=5,
)
payloads: Final = tuple(json.loads(event.model_dump_json()) for event in events)
assert [payload.get("sequence_number") for payload in payloads] == list(range(len(payloads)))

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