Merge remote-tracking branch 'origin/main' into litellm_lit_7108_responses_prompt_cache_breakpoint

# Conflicts:
#	litellm/completion_extras/litellm_responses_transformation/transformation.py
#	tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py
This commit is contained in:
mateo-berri 2026-10-07 17:11:36 -07:00
commit 18a818530f
593 changed files with 83990 additions and 53361 deletions

View file

@ -3145,10 +3145,10 @@ jobs:
name: Test provider capture and replay harness
command: |
mkdir -p test-results/provider-replay-harness
uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o pythonpath=tests/e2e -p no:rerunfailures \
uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures \
--junitxml=test-results/provider-replay-harness/junit.xml \
tests/e2e/test_provider_edge.py tests/e2e/test_fixture_bundle.py \
tests/e2e/test_fixture_canonical.py tests/e2e/test_fixture_mode.py \
tests/e2e_harness/test_provider_edge.py tests/e2e_harness/test_fixture_bundle.py \
tests/e2e_harness/test_fixture_canonical.py tests/e2e_harness/test_fixture_mode.py \
tests/code_coverage_tests/test_provider_replay_harness.py \
tests/code_coverage_tests/test_provider_cache.py
- store_test_results:

View file

@ -24,8 +24,8 @@ while IFS= read -r file || [ -n "$file" ]; do
has_mcp_dependencies=true ;;
esac
case "$file" in
tests/e2e/*/*.py) : ;;
tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/unit/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
tests/e2e/*/*.py | tests/e2e_harness/*/*.py) : ;;
tests/e2e/*.py | tests/e2e_harness/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/unit/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
has_provider_harness=true ;;
esac
case "$file" in

View file

@ -193,10 +193,10 @@ fi
if [ "$suite" = providers ]; then
INTEGRATION_RUN_ID="$integration_identity" .venv/bin/python -m pytest --tb=short --noconftest -o addopts= \
--strict-markers --strict-config -p no:pytest-retry -p no:rerunfailures --timeout=30 \
tests/e2e/test_provider_edge.py::TestReplayMode::test_content_drift_returns_the_miss_status_naming_both_keys \
tests/e2e/test_provider_edge.py::TestReplayMode::test_exhausted_key_returns_the_miss_status \
tests/e2e/test_provider_edge.py::TestReplayLeftover::test_partially_consumed_recording_names_the_leftover \
tests/e2e/test_provider_edge.py::TestStreamingFidelity::test_replay_of_a_stream_makes_no_provider_connection \
tests/e2e_harness/test_provider_edge.py::TestReplayMode::test_content_drift_returns_the_miss_status_naming_both_keys \
tests/e2e_harness/test_provider_edge.py::TestReplayMode::test_exhausted_key_returns_the_miss_status \
tests/e2e_harness/test_provider_edge.py::TestReplayLeftover::test_partially_consumed_recording_names_the_leftover \
tests/e2e_harness/test_provider_edge.py::TestStreamingFidelity::test_replay_of_a_stream_makes_no_provider_connection \
--junitxml="$results/replay-controls.xml"
fi

View file

@ -1,8 +1,5 @@
name: "LiteLLM CodeQL config"
queries:
- uses: security-and-quality
# Known OOM queries on large Python codebases:
# CodeQL builds a full data flow graph in memory. These two queries trace
# sensitive data through every log call / regex pattern, causing combinatorial
@ -14,17 +11,6 @@ query-filters:
id: py/clear-text-logging-sensitive-data # CWE-312
- exclude:
id: py/polynomial-redos # CWE-730
# Import resolution confuses stdlib types with management_endpoints/types.py.
# The generic cycle query also reports intentional deferred imports.
- exclude:
id: py/cyclic-import
- exclude:
id: py/unsafe-cyclic-import
# Known false positives on live settings and Protocol placeholders.
- exclude:
id: py/unused-global-variable
- exclude:
id: py/ineffectual-statement
paths-ignore:
- tests

View file

@ -1,8 +1,6 @@
name: "CodeQL"
on:
push:
branches: [main]
pull_request:
branches: [main]
schedule:
@ -43,14 +41,15 @@ jobs:
persist-credentials: false
- name: Initialize CodeQL
uses: github/codeql-action/init@ebcb5b36ded6beda4ceefea6a8bc4cc885255bb3 # v3.34.1
uses: github/codeql-action/init@2892aa5e19bbd11bc0cff5427e3b750a04d9e3c2 # v4.38.2
with:
languages: ${{ matrix.language }}
build-mode: ${{ matrix.build-mode }}
config-file: ./.github/codeql/codeql-config.yml
queries: ${{ github.event_name == 'pull_request' && '+security-extended' || '' }}
- name: Perform CodeQL Analysis
uses: github/codeql-action/analyze@ebcb5b36ded6beda4ceefea6a8bc4cc885255bb3 # v3.34.1
uses: github/codeql-action/analyze@2892aa5e19bbd11bc0cff5427e3b750a04d9e3c2 # v4.38.2
with:
category: "/language:${{ matrix.language }}"
output: sarif-results
@ -83,7 +82,7 @@ jobs:
output: sarif-results/python.sarif
- name: Upload SARIF
uses: github/codeql-action/upload-sarif@ebcb5b36ded6beda4ceefea6a8bc4cc885255bb3 # v3.34.1
uses: github/codeql-action/upload-sarif@2892aa5e19bbd11bc0cff5427e3b750a04d9e3c2 # v4.38.2
with:
sarif_file: sarif-results
category: "/language:${{ matrix.language }}"

View file

@ -174,23 +174,25 @@ jobs:
- name: Check tests/e2e basedpyright (zero errors)
if: steps.changes.outputs.decision != 'skip'
run: |
if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- ':(glob)tests/e2e/**/*.py' | grep -q .; then
uv run --no-sync basedpyright tests/e2e
if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- ':(glob)tests/e2e/**/*.py' ':(glob)tests/e2e_harness/**/*.py' pyrightconfig.json | grep -q .; then
uv run --no-sync basedpyright tests/e2e tests/e2e_harness
else
echo "No changed tests/e2e Python files; skipping."
fi
- name: Run the claude_code harness unit tests
- name: Run the e2e harness tests
if: steps.changes.outputs.decision != 'skip'
env:
LITELLM_MASTER_KEY: sk-e2e-harness-tests-reach-no-proxy
run: |
if ! git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- ':(glob)tests/e2e/claude_code/**/*.py' ':(glob)tests/e2e/*.py' tests/e2e/claude_code/cron_vm/install_claude_code.sh pyproject.toml uv.lock .github/workflows/test-linting.yml | grep -q .; then
echo "No changed claude_code harness files; skipping."
if ! git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- tests/e2e tests/e2e_harness ':(exclude)tests/e2e/ui' pyproject.toml uv.lock .github/workflows/test-linting.yml | grep -q .; then
echo "No changed e2e harness files; skipping."
exit 0
fi
retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; }
CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/claude_code/pr_gate_version_resolver.py)"
tests/e2e/claude_code/cron_vm/install_claude_code.sh "$CLAUDE_VERSION" "$RUNNER_TEMP/claude-cli"
PATH="$RUNNER_TEMP/claude-cli:$PATH" uv run --no-sync pytest -q --noconftest -o addopts= -o pythonpath=tests/e2e -p no:rerunfailures tests/e2e/claude_code/_*_unit_tests
PATH="$RUNNER_TEMP/claude-cli:$PATH" uv run --no-sync pytest -q tests/e2e_harness
- name: Check for circular imports
if: steps.changes.outputs.decision != 'skip'

View file

@ -31,7 +31,7 @@ help:
@echo " make lint - Run all linting (Ruff, basedpyright, format check, circular imports, import safety)"
@echo " make lint-ruff - Run Ruff linting only"
@echo " make lint-basedpyright - Run basedpyright strict, gated by per-rule error counts"
@echo " make lint-e2e-basedpyright - Run basedpyright over tests/e2e (zero errors allowed)"
@echo " make lint-e2e-basedpyright - Run basedpyright over tests/e2e and tests/e2e_harness (zero errors allowed)"
@echo " make lint-basedpyright-budget-update - Ratchet basedpyright limits down by what this branch fixed"
@echo " make lint-format - Check ruff format formatting (matches CI)"
@echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its limit"
@ -211,7 +211,7 @@ lint-basedpyright: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/type_check_gate.py --base "$(BASE_REF)"
lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
$(UV_RUN) basedpyright tests/e2e
$(UV_RUN) basedpyright tests/e2e tests/e2e_harness
# Type-discipline budget (mutable collections / casts / type guards / kwargs /
# unexplained suppressions), the test-linting.yml step `make lint` used to omit.

View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "rpm" INTEGER;

View file

@ -0,0 +1,19 @@
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_attribute
WHERE attrelid = to_regclass('"LiteLLM_AutoRouterDailySpend"')
AND attname = 'total_tokens' AND NOT attisdropped
) THEN
ALTER TABLE "LiteLLM_AutoRouterDailySpend"
ADD COLUMN IF NOT EXISTS "total_tokens" BIGINT NOT NULL DEFAULT 0;
END IF;
IF NOT EXISTS (
SELECT 1 FROM pg_attribute
WHERE attrelid = to_regclass('"LiteLLM_AutoRouterDailySpend"')
AND attname = 'token_recorded_turns' AND NOT attisdropped
) THEN
ALTER TABLE "LiteLLM_AutoRouterDailySpend"
ADD COLUMN IF NOT EXISTS "token_recorded_turns" INTEGER NOT NULL DEFAULT 0;
END IF;
END $$;

View file

@ -0,0 +1,16 @@
CREATE TABLE IF NOT EXISTS "LiteLLM_LensSignalConfig" (
"id" TEXT NOT NULL,
"data" JSONB NOT NULL,
CONSTRAINT "LiteLLM_LensSignalConfig_pkey" PRIMARY KEY ("id")
);
CREATE TABLE IF NOT EXISTS "LiteLLM_LensTraceSignal" (
"trace_id" TEXT NOT NULL,
"trace_ref" TEXT NOT NULL DEFAULT '',
"config_key" TEXT NOT NULL,
"span_count" INTEGER NOT NULL,
"claimed_until" TIMESTAMP(3),
"classified_at" TIMESTAMP(3),
"data" JSONB NOT NULL,
CONSTRAINT "LiteLLM_LensTraceSignal_pkey" PRIMARY KEY ("trace_id", "trace_ref")
);

View file

@ -417,6 +417,7 @@ model LiteLLM_MCPServerTable {
source_url String?
timeout Float?
max_concurrent_requests Int?
rpm Int?
// BYOM submission lifecycle
approval_status String? @default("active")
submitted_by String?
@ -1756,6 +1757,8 @@ model LiteLLM_AutoRouterDailySpend {
router_name String
router_type String
turns Int @default(0)
total_tokens BigInt @default(0)
token_recorded_turns Int @default(0)
spend Float @default(0)
saved_spend Float @default(0)
savings_estimated_turns Int @default(0)
@ -1977,3 +1980,20 @@ model LiteLLM_LensDataset {
@@id([id, revision])
}
model LiteLLM_LensSignalConfig {
id String @id
data Json
}
model LiteLLM_LensTraceSignal {
trace_id String
trace_ref String @default("")
config_key String
span_count Int
claimed_until DateTime?
classified_at DateTime?
data Json
@@id([trace_id, trace_ref])
}

View file

@ -1,8 +1,5 @@
use pyo3::{prelude::*, types::PyDict};
/// The caller's own object for a public argument: the keyword if given, even an explicit
/// `None`, else the bound request's attribute. Every reader of a public Python call uses
/// this rule, so the callbacks and the provider see one object per argument.
pub fn lookup<'py>(
kwargs: &Bound<'py, PyDict>,
request: &Bound<'py, PyAny>,
@ -11,6 +8,9 @@ pub fn lookup<'py>(
if let Some(value) = kwargs.get_item(name)? {
return Ok(Some(value));
}
if let Ok(bound) = request.cast::<PyDict>() {
return bound.get_item(name);
}
request.getattr_opt(name)
}
@ -48,4 +48,32 @@ kwargs = {'api_key': key, 'api_base': None}
assert!(find("model").is_none());
});
}
#[rstest::rstest]
#[case::prepared_value("{'api_key': 'replacement'}", Some("replacement"))]
#[case::explicit_none("{'api_key': None}", None)]
#[case::bound_fallback("{}", Some("original"))]
fn prepared_mapping_overrides_bound_values(
#[case] source: &str,
#[case] expected: Option<&str>,
) {
crate::initialize_python();
Python::attach(|py| {
let bound = PyDict::new(py);
bound.set_item("api_key", "original").unwrap();
let source = std::ffi::CString::new(source).unwrap();
let prepared = py
.eval(&source, None, None)
.unwrap()
.cast_into::<PyDict>()
.unwrap();
let value = lookup(&prepared, bound.as_any(), "api_key")
.unwrap()
.unwrap();
assert_eq!(
value.extract::<Option<String>>().unwrap().as_deref(),
expected
);
});
}
}

View file

@ -14,7 +14,9 @@ use litellm_secrets::source::SecretSource;
use std::sync::Arc;
pub use litellm_inference::RouteError as Error;
pub use types::{MessagesCall, MessagesCallResponse, MessagesShaping, messages_body};
pub use types::{
MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, messages_body,
};
#[derive(Clone)]
pub struct MessagesRoute {

View file

@ -80,12 +80,13 @@ fn prepare_provider_request(
let sanitized = config.shape_request(
MessagesRequest { model, ..body },
shaping.reasoning_auto_summary,
shaping.settings.reasoning_auto_summary,
)?;
let trimmed = without_additional_drop_params(sanitized, &shaping.additional_drop_params)?;
let trimmed =
without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?;
let transformed = config.transform_anthropic_messages_request(
trimmed,
&MessagesTransformContext::new(shaping.capabilities, shaping.drop_params),
&MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params),
)?;
let scoped =
@ -148,7 +149,7 @@ mod tests {
use serde_json::{Map, Value, json};
use super::*;
use crate::MessagesShaping;
use crate::{MessagesSettings, MessagesShaping};
#[fixture]
fn shaping() -> MessagesShaping {
@ -310,10 +311,13 @@ mod tests {
)
};
let shaping = MessagesShaping {
additional_drop_params: additional_drop_params
.iter()
.map(ToString::to_string)
.collect(),
settings: MessagesSettings {
additional_drop_params: additional_drop_params
.iter()
.map(ToString::to_string)
.collect(),
..shaping.settings
},
..shaping
};
assert_eq!(
@ -394,8 +398,11 @@ mod tests {
#[rstest]
fn dropped_thinking_display_is_not_restored_by_auto_summary(shaping: MessagesShaping) {
let shaping = MessagesShaping {
reasoning_auto_summary: true,
additional_drop_params: vec!["thinking.display".to_string()],
settings: MessagesSettings {
reasoning_auto_summary: true,
additional_drop_params: vec!["thinking.display".to_string()],
..shaping.settings
},
..shaping
};
assert_eq!(
@ -420,7 +427,10 @@ mod tests {
#[rstest]
fn dropping_an_invalid_metadata_user_id_does_not_skip_its_validation(shaping: MessagesShaping) {
let shaping = MessagesShaping {
additional_drop_params: vec!["metadata.user_id".to_string()],
settings: MessagesSettings {
additional_drop_params: vec!["metadata.user_id".to_string()],
..shaping.settings
},
..shaping
};
assert!(matches!(

View file

@ -38,6 +38,12 @@ pub type MessagesCallResponse =
pub struct MessagesShaping {
#[serde(default)]
pub capabilities: MessagesModelCapabilities,
#[serde(flatten)]
pub settings: MessagesSettings,
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct MessagesSettings {
#[serde(default)]
pub drop_params: bool,
#[serde(default)]
@ -58,16 +64,25 @@ mod tests {
#[case::nothing_projected(json!({}), MessagesShaping::default())]
#[case::only_drop_params(
json!({"drop_params": true}),
MessagesShaping { drop_params: true, ..MessagesShaping::default() },
MessagesShaping {
settings: MessagesSettings { drop_params: true, ..MessagesSettings::default() },
..MessagesShaping::default()
},
)]
#[case::only_reasoning_auto_summary(
json!({"reasoning_auto_summary": true}),
MessagesShaping { reasoning_auto_summary: true, ..MessagesShaping::default() },
MessagesShaping {
settings: MessagesSettings { reasoning_auto_summary: true, ..MessagesSettings::default() },
..MessagesShaping::default()
},
)]
#[case::only_additional_drop_params(
json!({"additional_drop_params": ["tools[*].input_examples"]}),
MessagesShaping {
additional_drop_params: vec!["tools[*].input_examples".to_string()],
settings: MessagesSettings {
additional_drop_params: vec!["tools[*].input_examples".to_string()],
..MessagesSettings::default()
},
..MessagesShaping::default()
},
)]
@ -98,6 +113,11 @@ mod tests {
"additional_drop_params": ["metadata.user_id", "thinking"]
}),
MessagesShaping {
settings: MessagesSettings {
drop_params: true,
reasoning_auto_summary: true,
additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()],
},
capabilities: MessagesModelCapabilities {
supports_reasoning: true,
supports_adaptive_thinking: true,
@ -115,9 +135,6 @@ mod tests {
max: false,
},
},
drop_params: true,
reasoning_auto_summary: true,
additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()],
},
)]
fn shaping_deserializes_with_defaults_for_absent_fields(
@ -126,5 +143,22 @@ mod tests {
) {
let shaping: MessagesShaping = serde_json::from_value(projected).unwrap();
assert_eq!(shaping, expected);
let serialized = serde_json::to_value(&shaping).unwrap();
assert_eq!(
serialized["drop_params"],
json!(expected.settings.drop_params)
);
assert_eq!(
serialized["reasoning_auto_summary"],
json!(expected.settings.reasoning_auto_summary)
);
assert_eq!(
serialized["additional_drop_params"],
json!(expected.settings.additional_drop_params)
);
assert_eq!(
serialized["capabilities"],
serde_json::to_value(expected.capabilities).unwrap()
);
}
}

View file

@ -380,12 +380,14 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages
let host = RecordingHost::passthrough(authenticated(
MessagesCall {
shaping: MessagesShaping {
settings: MessagesSettings {
drop_params: true,
..MessagesSettings::default()
},
capabilities: AnthropicModelCapabilities {
supports_sampling_params: false,
..AnthropicModelCapabilities::default()
},
drop_params: true,
..MessagesShaping::default()
},
..with_fields(call, json!({"temperature": 0.2}))
},

View file

@ -5,7 +5,7 @@ use std::{
use litellm_http::{HttpSettings, Resolution};
use litellm_inference_messages::{
Error, MessagesCall, MessagesShaping,
Error, MessagesCall, MessagesSettings, MessagesShaping,
route::{Messages, MessagesMachine, MessagesOutput},
};
use litellm_inference_testing::RecordingSecrets;

View file

@ -239,7 +239,10 @@ async fn additional_drop_params_remove_fields_before_sending(call: MessagesCall)
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
additional_drop_params: vec!["temperature".into()],
settings: MessagesSettings {
additional_drop_params: vec!["temperature".into()],
..MessagesSettings::default()
},
..MessagesShaping::default()
},
..with_fields(call, json!({"temperature": 0.5, "top_k": 3}))
@ -390,8 +393,10 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
capabilities,
drop_params,
..MessagesShaping::default()
settings: MessagesSettings {
drop_params,
..MessagesSettings::default()
},
},
body: call.body.clone(),
custom_llm_provider: call.custom_llm_provider.clone(),
@ -436,13 +441,15 @@ async fn reasoning_auto_summary_marks_active_thinking_on_the_wire(
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
settings: MessagesSettings {
reasoning_auto_summary: true,
..MessagesSettings::default()
},
capabilities: MessagesModelCapabilities {
supports_reasoning: true,
supports_adaptive_thinking: true,
..MessagesModelCapabilities::default()
},
reasoning_auto_summary: true,
..MessagesShaping::default()
},
..call
},
@ -634,7 +641,10 @@ async fn system_message_folding_is_selected_by_the_provider(
api_key: Some("sk-azure".into()),
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
additional_drop_params: drop_params.iter().map(ToString::to_string).collect(),
settings: MessagesSettings {
additional_drop_params: drop_params.iter().map(ToString::to_string).collect(),
..call.shaping.settings
},
..call.shaping
},
..call
@ -695,7 +705,10 @@ async fn provider_validation_runs_before_caller_parameter_removal(
api_key: Some("sk-test".into()),
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
additional_drop_params: vec!["metadata".into()],
settings: MessagesSettings {
additional_drop_params: vec!["metadata".into()],
..call.shaping.settings
},
..call.shaping
},
..call

View file

@ -29,10 +29,16 @@ pub struct ModelInfo {
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_32k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_100k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_100k_tokens_batches: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_1hr: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_1hr_above_100k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_1hr_above_200k_tokens: Option<f64>,
@ -82,6 +88,10 @@ pub struct ModelInfo {
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_32k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_100k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_100k_tokens_batches: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(skip_serializing_if = "Option::is_none")]
@ -200,6 +210,10 @@ pub struct ModelInfo {
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_32k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_100k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_100k_tokens_batches: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(skip_serializing_if = "Option::is_none")]
@ -355,6 +369,10 @@ pub struct ModelInfo {
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_32k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_100k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_100k_tokens_batches: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(skip_serializing_if = "Option::is_none")]

View file

@ -74,3 +74,20 @@ fn checked_in_catalog_and_backup_match() {
"invalid registry aliases"
);
}
#[rstest]
#[case::input("input_cost_per_token_above_100k_tokens")]
#[case::input_batches("input_cost_per_token_above_100k_tokens_batches")]
#[case::output("output_cost_per_token_above_100k_tokens")]
#[case::output_batches("output_cost_per_token_above_100k_tokens_batches")]
#[case::cache_creation("cache_creation_input_token_cost_above_100k_tokens")]
#[case::cache_creation_batches("cache_creation_input_token_cost_above_100k_tokens_batches")]
#[case::cache_creation_1hr("cache_creation_input_token_cost_above_1hr_above_100k_tokens")]
#[case::cache_read("cache_read_input_token_cost_above_100k_tokens")]
#[case::cache_read_batches("cache_read_input_token_cost_above_100k_tokens_batches")]
fn registry_validation_keeps_above_100k_tier_rates(#[case] field: &str) {
let mut entry = Map::new();
entry.insert("litellm_provider".into(), "anthropic".into());
entry.insert(field.into(), 5e-7.into());
validate_model_entry("test", &Value::Object(entry)).unwrap();
}

View file

@ -6,6 +6,7 @@ use std::{
use litellm_auth::InputSource;
use litellm_host_python::{from_py, from_py_argument};
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
use serde::de::DeserializeOwned;
use serde_json::{Map, Value};
/// The keyword arguments every value route shares, validated at the Python boundary.
@ -25,18 +26,6 @@ pub(crate) fn messages_argument(value: &Bound<'_, PyAny>) -> PyResult<Vec<Value>
}
}
pub(crate) fn optional_params_argument(
value: &Bound<'_, PyAny>,
) -> PyResult<Option<Map<String, Value>>> {
optional_object("optional_params", value)
}
pub(crate) fn extra_headers_argument(
value: &Bound<'_, PyAny>,
) -> PyResult<Option<Map<String, Value>>> {
optional_object("extra_headers", value)
}
fn required_object(name: &'static str, value: Value) -> PyResult<Map<String, Value>> {
match value {
Value::Object(values) => Ok(values),
@ -71,6 +60,48 @@ pub(crate) fn python_timeout_seconds(py: Python<'_>, timeout: Py<PyAny>) -> PyRe
.extract()
}
pub(crate) fn required_field<'py>(
fields: &Bound<'py, PyDict>,
name: &str,
) -> PyResult<Bound<'py, PyAny>> {
fields
.get_item(name)?
.ok_or_else(|| PyValueError::new_err(format!("{name} is required")))
}
pub(crate) fn optional_field<T: DeserializeOwned>(
fields: &Bound<'_, PyDict>,
name: &str,
) -> PyResult<Option<T>> {
fields
.get_item(name)?
.map(|value| from_py_argument(&value))
.transpose()
.map(Option::flatten)
}
pub(crate) fn optional_object_field(
fields: &Bound<'_, PyDict>,
name: &'static str,
) -> PyResult<Option<Map<String, Value>>> {
fields
.get_item(name)?
.map(|value| optional_object(name, &value))
.transpose()
.map(Option::flatten)
}
pub(crate) fn value_route_options(fields: &Bound<'_, PyDict>) -> PyResult<RouteOptions> {
Ok(RouteOptions {
model: from_py_argument(&required_field(fields, "model")?)?,
api_key: optional_field(fields, "api_key")?,
api_base: optional_field(fields, "api_base")?,
custom_llm_provider: optional_field(fields, "custom_llm_provider")?,
extra_headers: optional_object_field(fields, "extra_headers")?,
timeout: optional_timeout(optional_field(fields, "timeout_seconds")?),
})
}
pub(crate) fn project_optional_fields(
kwargs: &Bound<'_, PyDict>,
names: &[&str],
@ -244,15 +275,15 @@ mod tests {
let params = py.eval(c"{'temperature': 0.2}", None, None).unwrap();
assert_eq!(
optional_params_argument(&params).unwrap(),
optional_object("optional_params", &params).unwrap(),
Some(required_object("optional_params", json!({"temperature": 0.2})).unwrap())
);
assert_eq!(
optional_params_argument(&py.None().into_bound(py)).unwrap(),
optional_object("optional_params", &py.None().into_bound(py)).unwrap(),
None
);
assert_eq!(
extra_headers_argument(&py.None().into_bound(py)).unwrap(),
optional_object("extra_headers", &py.None().into_bound(py)).unwrap(),
None
);
});

View file

@ -1,14 +1,13 @@
use crate::execution::{run_async, run_sync};
use litellm_host_python::from_py_argument;
use litellm_inference_transcription::{
AudioTranscriptionRoute, Error, types::AudioTranscriptionRequest,
};
use pyo3::{prelude::*, types::PyDict};
use pyo3::prelude::*;
use serde_json::{Map, Value};
use crate::{
errors::route_error_to_pyerr,
marshal::{RouteOptions, extra_headers_argument, optional_params_argument, optional_timeout},
marshal::{RouteOptions, optional_object_field, required_field, value_route_options},
};
async fn execute(
@ -41,81 +40,38 @@ async fn execute(
}
#[pyfunction]
#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
#[expect(
clippy::too_many_arguments,
reason = "one parameter per Python keyword"
)]
pub(crate) fn transcription(
py: Python<'_>,
model: String,
#[pyo3(from_py_with = from_py_argument)] audio: Value,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = extra_headers_argument)] extra_headers: Option<Map<String, Value>>,
#[pyo3(from_py_with = optional_params_argument)] optional_params: Option<Map<String, Value>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
let options = RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout: optional_timeout(timeout_seconds),
};
let http = crate::http::provider_client(py, &PyDict::new(py), false)?;
pub(crate) fn transcription(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
let audio: Value =
litellm_host_python::from_py_argument(&required_field(&call.bound, "audio")?)?;
let options = value_route_options(&call.bound)?;
let optional_params =
optional_object_field(&call.bound, "optional_params")?.unwrap_or_default();
let http = crate::http::provider_client(py, &call.kwargs, false)?;
let secrets = crate::secrets::source(py)?;
run_sync(
py,
execute(
http,
secrets,
audio,
optional_params.unwrap_or_default(),
options,
),
execute(http, secrets, audio, optional_params, options),
route_error_to_pyerr,
)
}
#[pyfunction]
#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
#[expect(
clippy::too_many_arguments,
reason = "one parameter per Python keyword"
)]
pub(crate) fn atranscription<'py>(
py: Python<'py>,
model: String,
#[pyo3(from_py_with = from_py_argument)] audio: Value,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = extra_headers_argument)] extra_headers: Option<Map<String, Value>>,
#[pyo3(from_py_with = optional_params_argument)] optional_params: Option<Map<String, Value>>,
timeout_seconds: Option<f64>,
call: Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyAny>> {
let options = RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout: optional_timeout(timeout_seconds),
};
let http = crate::http::provider_client(py, &PyDict::new(py), true)?;
let call = super::NativeCall::extract(&call)?;
let audio: Value =
litellm_host_python::from_py_argument(&required_field(&call.bound, "audio")?)?;
let options = value_route_options(&call.bound)?;
let optional_params =
optional_object_field(&call.bound, "optional_params")?.unwrap_or_default();
let http = crate::http::provider_client(py, &call.kwargs, true)?;
let secrets = crate::secrets::source(py)?;
run_async(
py,
execute(
http,
secrets,
audio,
optional_params.unwrap_or_default(),
options,
),
execute(http, secrets, audio, optional_params, options),
route_error_to_pyerr,
)
}

View file

@ -11,8 +11,7 @@ use serde_json::{Map, Value};
use crate::{
errors::route_error_to_pyerr,
marshal::{
RouteOptions, extra_headers_argument, messages_argument, optional_params_argument,
optional_timeout,
RouteOptions, messages_argument, optional_object_field, required_field, value_route_options,
},
};
@ -50,81 +49,36 @@ async fn execute(
}
#[pyfunction]
#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
#[expect(
clippy::too_many_arguments,
reason = "one parameter per Python keyword"
)]
pub(crate) fn chat_completions(
py: Python<'_>,
model: String,
#[pyo3(from_py_with = messages_argument)] messages: Vec<Value>,
#[pyo3(from_py_with = optional_params_argument)] optional_params: Option<Map<String, Value>>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = extra_headers_argument)] extra_headers: Option<Map<String, Value>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
let options = RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout: optional_timeout(timeout_seconds),
};
let http = crate::http::provider_client(py, &PyDict::new(py), false)?;
pub(crate) fn chat_completions(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
let messages: Vec<Value> = messages_argument(&required_field(&call.bound, "messages")?)?;
let optional_params =
optional_object_field(&call.bound, "optional_params")?.unwrap_or_default();
let options = value_route_options(&call.bound)?;
let http = crate::http::provider_client(py, &call.kwargs, false)?;
let secrets = crate::secrets::source(py)?;
run_sync(
py,
execute(
http,
secrets,
messages,
optional_params.unwrap_or_default(),
options,
),
execute(http, secrets, messages, optional_params, options),
route_error_to_pyerr,
)
}
#[pyfunction]
#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
#[expect(
clippy::too_many_arguments,
reason = "one parameter per Python keyword"
)]
pub(crate) fn achat_completions<'py>(
py: Python<'py>,
model: String,
#[pyo3(from_py_with = messages_argument)] messages: Vec<Value>,
#[pyo3(from_py_with = optional_params_argument)] optional_params: Option<Map<String, Value>>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = extra_headers_argument)] extra_headers: Option<Map<String, Value>>,
timeout_seconds: Option<f64>,
call: Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyAny>> {
let options = RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout: optional_timeout(timeout_seconds),
};
let http = crate::http::provider_client(py, &PyDict::new(py), true)?;
let call = super::NativeCall::extract(&call)?;
let messages: Vec<Value> = messages_argument(&required_field(&call.bound, "messages")?)?;
let optional_params =
optional_object_field(&call.bound, "optional_params")?.unwrap_or_default();
let options = value_route_options(&call.bound)?;
let http = crate::http::provider_client(py, &call.kwargs, true)?;
let secrets = crate::secrets::source(py)?;
run_async(
py,
execute(
http,
secrets,
messages,
optional_params.unwrap_or_default(),
options,
),
execute(http, secrets, messages, optional_params, options),
route_error_to_pyerr,
)
}
@ -184,21 +138,13 @@ fn run_public(
}
#[pyfunction]
pub(crate) fn completion(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_public(py, request, args, kwargs, false)
pub(crate) fn completion(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_public(py, call.bound.into_any(), call.args, call.kwargs, false)
}
#[pyfunction]
pub(crate) fn acompletion(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_public(py, request, args, kwargs, true)
pub(crate) fn acompletion(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_public(py, call.bound.into_any(), call.args, call.kwargs, true)
}

View file

@ -1,58 +1,49 @@
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use pyo3::prelude::*;
use crate::errors::RustBridgeDeclined;
#[pyfunction]
#[pyo3(signature = (request, args, kwargs))]
pub(crate) fn embedding(
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
drop((request, args, kwargs));
pub(crate) fn embedding(call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
drop(super::NativeCall::extract(&call)?);
Err(RustBridgeDeclined::new_err(
"native embeddings route is not implemented",
))
}
#[pyfunction]
#[pyo3(signature = (request, args, kwargs))]
pub(crate) fn aembedding(
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
drop((request, args, kwargs));
Err(RustBridgeDeclined::new_err(
"native embeddings route is not implemented",
))
pub(crate) fn aembedding(call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
embedding(call)
}
#[cfg(test)]
mod tests {
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use pyo3::{prelude::*, types::PyDict};
use rstest::rstest;
use crate::errors::RustBridgeDeclined;
#[test]
fn both_entrypoints_decline_before_provider_execution() {
#[rstest]
#[case::sync(false)]
#[case::asynchronous(true)]
fn both_entrypoints_decline_before_provider_execution(#[case] asynchronous: bool) {
Python::initialize();
Python::attach(|py| {
let request = PyDict::new(py);
let args = PyTuple::empty(py);
let kwargs = PyDict::new(py);
for entrypoint in [super::embedding, super::aembedding] {
let error = entrypoint(request.clone().into_any(), args.clone(), kwargs.clone())
.expect_err("native embeddings must decline until a route machine exists");
assert!(error.is_instance_of::<RustBridgeDeclined>(py));
let locals = PyDict::new(py);
py.run(
c"from types import SimpleNamespace
call = SimpleNamespace(args=(), kwargs={}, bound={'model':'test-model','input':'hello'})",
Some(&locals),
Some(&locals),
)
.unwrap();
let call = locals.get_item("call").unwrap().unwrap();
let error = if asynchronous {
super::aembedding(call)
} else {
super::embedding(call)
}
.expect_err("native embeddings must decline until a route machine exists");
assert!(error.is_instance_of::<RustBridgeDeclined>(py));
});
}
}

View file

@ -86,6 +86,9 @@ impl InferenceHost {
if let Some(value) = lookup(arguments, request, name)? {
return Ok((!value.is_none()).then_some(value));
}
if request.is_instance_of::<PyDict>() {
return Ok(None);
}
let parameter = request
.getattr("parameters")?
.call_method1("get", (name,))?;

View file

@ -5,9 +5,10 @@ use bytes::Bytes;
use litellm_host_python::{InvokeError, PythonBinding, from_py, lookup, to_py};
use litellm_http::transport::Error as TransportError;
use litellm_inference_messages::{
Error, MessagesCall, MessagesShaping, messages_body,
Error, MessagesCall, MessagesSettings, MessagesShaping, messages_body,
route::{Messages, MessagesStreamHead},
};
use litellm_llms::base_llm::messages::context::MessagesModelCapabilities;
use litellm_llms_types::headers::ProviderSpecificHeaders;
use pyo3::{
exceptions::{PyException, PyValueError},
@ -188,18 +189,24 @@ impl MessagesPythonHost {
custom_llm_provider: Option<&str>,
arguments: &Bound<'_, PyDict>,
) -> PyResult<MessagesShaping> {
let projected = py.import(ROUTE_HOST_MODULE)?.getattr("shaping")?.call1((
model,
custom_llm_provider,
arguments,
))?;
from_py(&projected)
let module = py.import(ROUTE_HOST_MODULE)?;
let capabilities: MessagesModelCapabilities = from_py(
&py.import("litellm.rust_bridge.model_capabilities")?
.getattr("anthropic_model_capabilities")?
.call1((model, custom_llm_provider))?,
)?;
let settings: MessagesSettings =
from_py(&module.getattr("settings")?.call1((arguments,))?)?;
Ok(MessagesShaping {
capabilities,
settings,
})
}
fn provider(&self, py: Python<'_>) -> String {
self.request
.bind(py)
.getattr("custom_llm_provider")
.get_item("custom_llm_provider")
.and_then(|value| value.extract::<Option<String>>())
.ok()
.flatten()

View file

@ -64,21 +64,13 @@ fn run_messages(
}
#[pyfunction]
pub(crate) fn messages(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_messages(py, request, args, kwargs, false)
pub(crate) fn messages(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_messages(py, call.bound.into_any(), call.args, call.kwargs, false)
}
#[pyfunction]
pub(crate) fn amessages(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_messages(py, request, args, kwargs, true)
pub(crate) fn amessages(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_messages(py, call.bound.into_any(), call.args, call.kwargs, true)
}

View file

@ -14,9 +14,35 @@ use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol}
use litellm_host_python::{HookChain, PythonBinding, PythonCallHooks, PythonHostCalls};
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
types::{PyDict, PyMapping, PyTuple},
};
struct NativeCall<'py> {
args: Bound<'py, PyTuple>,
kwargs: Bound<'py, PyDict>,
bound: Bound<'py, PyDict>,
}
impl<'py> NativeCall<'py> {
fn extract(call: &Bound<'py, PyAny>) -> PyResult<Self> {
Ok(Self {
args: call.getattr("args")?.cast_into()?,
kwargs: mapping_dict(&call.getattr("kwargs")?)?,
bound: mapping_dict(&call.getattr("bound")?)?,
})
}
}
fn mapping_dict<'py>(value: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyDict>> {
if let Ok(dict) = value.cast::<PyDict>() {
return Ok(dict.clone());
}
let mapping = value.cast::<PyMapping>()?;
let dict = PyDict::new(value.py());
dict.update(mapping)?;
Ok(dict)
}
fn call_hooks(
py: Python<'_>,
operation: LoggingOperation,
@ -74,40 +100,30 @@ mod tests {
types::{PyDict, PyList},
};
#[test]
fn sync_and_async_route_signatures_match_the_python_contract() {
Python::initialize();
Python::attach(|py| {
let module = crate::native_module(py);
let routes = [
(
"transcription",
"atranscription",
"(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)",
),
(
"chat_completions",
"achat_completions",
"(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)",
),
];
for (sync_name, async_name, expected) in routes {
let sync_signature: String = module
.getattr(sync_name)
.and_then(|function| function.getattr("__text_signature__"))
.and_then(|signature| signature.extract())
.expect("sync signature should be available");
let async_signature: String = module
.getattr(async_name)
.and_then(|function| function.getattr("__text_signature__"))
.and_then(|signature| signature.extract())
.expect("async signature should be available");
assert_eq!(sync_signature, expected);
assert_eq!(async_signature, expected);
}
});
fn value_call<'py>(
py: Python<'py>,
payload_name: &str,
payload: &Bound<'py, PyAny>,
kwargs: Option<&Bound<'py, PyDict>>,
) -> Bound<'py, PyAny> {
let fields = PyDict::new(py);
fields.set_item("model", "model").unwrap();
fields.set_item(payload_name, payload).unwrap();
if let Some(kwargs) = kwargs {
fields.update(kwargs.as_mapping()).unwrap();
}
let attributes = PyDict::new(py);
attributes
.set_item("args", pyo3::types::PyTuple::empty(py))
.unwrap();
attributes.set_item("kwargs", &fields).unwrap();
attributes.set_item("bound", &fields).unwrap();
py.import("types")
.unwrap()
.getattr("SimpleNamespace")
.unwrap()
.call((), Some(&attributes))
.unwrap()
}
#[test]
@ -138,7 +154,9 @@ value = Broken()
for name in ["chat_completions", "achat_completions"] {
let error = module
.getattr(name)
.and_then(|function| function.call1(("model", &broken)))
.and_then(|function| {
function.call1((value_call(py, "messages", &broken, None),))
})
.expect_err("route should reject a value it cannot convert");
assert!(
@ -158,11 +176,15 @@ value = Broken()
let invalid_messages = PyDict::new(py);
let sync_chat_error = module
.getattr("chat_completions")
.and_then(|function| function.call1(("model", &invalid_messages)))
.and_then(|function| {
function.call1((value_call(py, "messages", &invalid_messages, None),))
})
.expect_err("sync chat should reject a non-list messages value");
let async_chat_error = module
.getattr("achat_completions")
.and_then(|function| function.call1(("model", &invalid_messages)))
.and_then(|function| {
function.call1((value_call(py, "messages", &invalid_messages, None),))
})
.expect_err("async chat should reject a non-list messages value");
assert_eq!(
@ -180,11 +202,15 @@ value = Broken()
let sync_error = module
.getattr("transcription")
.and_then(|function| function.call(("model", &audio), Some(&kwargs)))
.and_then(|function| {
function.call1((value_call(py, "audio", &audio, Some(&kwargs)),))
})
.expect_err("sync route should reject non-dict extra_headers");
let async_error = module
.getattr("atranscription")
.and_then(|function| function.call(("model", &audio), Some(&kwargs)))
.and_then(|function| {
function.call1((value_call(py, "audio", &audio, Some(&kwargs)),))
})
.expect_err("async route should reject non-dict extra_headers");
assert_eq!(
@ -213,7 +239,12 @@ value = Broken()
let error = module
.getattr("chat_completions")
.and_then(|function| {
function.call(("model", &invalid_messages), Some(&chat_kwargs))
function.call1((value_call(
py,
"messages",
&invalid_messages,
Some(&chat_kwargs),
),))
})
.expect_err("messages should be validated first");
assert_eq!(error.to_string(), "ValueError: messages must be a list");
@ -221,7 +252,14 @@ value = Broken()
let valid_messages = PyList::empty(py);
let error = module
.getattr("chat_completions")
.and_then(|function| function.call(("model", &valid_messages), Some(&chat_kwargs)))
.and_then(|function| {
function.call1((value_call(
py,
"messages",
&valid_messages,
Some(&chat_kwargs),
),))
})
.expect_err("optional_params should be validated before headers");
assert_eq!(
error.to_string(),
@ -237,7 +275,12 @@ value = Broken()
let error = module
.getattr("transcription")
.and_then(|function| {
function.call(("model", &invalid_payload), Some(&headers_kwargs))
function.call1((value_call(
py,
"audio",
&invalid_payload,
Some(&headers_kwargs),
),))
})
.expect_err("payload should be validated before headers");
assert!(!error.to_string().contains("extra_headers"));
@ -265,11 +308,15 @@ value = Broken()
let omitted_error = module
.getattr("chat_completions")
.and_then(|function| function.call(("model", &messages), Some(&omitted)))
.and_then(|function| {
function.call1((value_call(py, "messages", &messages, Some(&omitted)),))
})
.expect_err("omitted optional_params should reach header validation");
let explicit_error = module
.getattr("chat_completions")
.and_then(|function| function.call(("model", &messages), Some(&explicit)))
.and_then(|function| {
function.call1((value_call(py, "messages", &messages, Some(&explicit)),))
})
.expect_err("None optional_params should reach header validation");
assert_eq!(
omitted_error.to_string(),

View file

@ -81,23 +81,15 @@ fn project_provider_defaults(snapshot: &Snapshot<'_>) -> PyResult<OcrSettings> {
}
#[pyfunction]
pub(crate) fn ocr(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_ocr(py, request, args, kwargs, false)
pub(crate) fn ocr(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_ocr(py, call.bound.into_any(), call.args, call.kwargs, false)
}
#[pyfunction]
pub(crate) fn aocr(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_ocr(py, request, args, kwargs, true)
pub(crate) fn aocr(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_ocr(py, call.bound.into_any(), call.args, call.kwargs, true)
}
#[pyfunction]

View file

@ -105,23 +105,15 @@ fn run_public(
}
#[pyfunction]
pub(crate) fn responses(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_public(py, request, args, kwargs, false)
pub(crate) fn responses(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_public(py, call.bound.into_any(), call.args, call.kwargs, false)
}
#[pyfunction]
pub(crate) fn aresponses(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_public(py, request, args, kwargs, true)
pub(crate) fn aresponses(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_public(py, call.bound.into_any(), call.args, call.kwargs, true)
}
#[pyclass]

View file

@ -1,7 +1,7 @@
use std::time::Duration;
use litellm_config::Config;
use litellm_inference_messages::MessagesShaping;
use litellm_inference_messages::{MessagesSettings, MessagesShaping};
use litellm_router::{Deployment, Router};
use rstest::rstest;
@ -74,8 +74,11 @@ fn programmatic_deployments_preserve_overrides_and_last_entry_wins() {
custom_llm_provider: Some("test-provider".into()),
timeout: Some(Duration::from_secs(7)),
shaping: MessagesShaping {
drop_params: true,
additional_drop_params: vec!["metadata.test".into()],
settings: MessagesSettings {
drop_params: true,
additional_drop_params: vec!["metadata.test".into()],
..MessagesSettings::default()
},
..Default::default()
},
};

View file

@ -325,6 +325,9 @@ mod tests {
call_keys: Vec::new(),
call_evidence: None,
tool_call_id: String::new(),
source_type: String::new(),
source_url: String::new(),
source_title: String::new(),
team_id: String::new(),
api_key_hash: String::new(),
user_id: String::new(),

View file

@ -286,6 +286,9 @@ fn span(index: usize) -> TraceSpansRow {
call_keys: Vec::new(),
call_evidence: None,
tool_call_id: String::new(),
source_type: String::new(),
source_url: String::new(),
source_title: String::new(),
team_id: "team".into(),
api_key_hash: "key".into(),
user_id: "user".into(),

View file

@ -39,6 +39,9 @@ fn row(span_id: &str, parent: &str, name: &str, kind: &str, agent: &str) -> Trac
call_keys: Vec::new(),
call_evidence: None,
tool_call_id: String::new(),
source_type: String::new(),
source_url: String::new(),
source_title: String::new(),
team_id: String::new(),
api_key_hash: String::new(),
user_id: String::new(),

View file

@ -0,0 +1,33 @@
WITH runs AS (
SELECT TeamId, ApiKeyHash, TraceId,
toUnixTimestamp64Milli(min(StartTs)) AS start_ms,
min(StartTs) AS trace_start, max(EndTs) AS trace_end,
sum(ErrorCount) > 0 AS failed
FROM agent_traces_by_key
WHERE ({all_teams:UInt8} = 1
OR ({user_id:String} != '' AND UserIds = [{user_id:String}])
OR has({team_ids:Array(String)}, TeamId))
GROUP BY TeamId, ApiKeyHash, TraceId
HAVING min(StartTs) >= fromUnixTimestamp64Milli({start_ms:Int64})
AND min(StartTs) < fromUnixTimestamp64Milli({end_ms:Int64})
),
named AS (
SELECT DISTINCT o.TeamId AS TeamId, o.ApiKeyHash AS ApiKeyHash, o.TraceId AS TraceId,
o.AgentName AS agent_name, toString(o.Framework) AS framework
FROM otel_traces AS o
WHERE o.AgentName != ''
AND o.Timestamp >= (SELECT min(trace_start) FROM runs)
AND o.Timestamp <= (SELECT max(trace_end) FROM runs)
AND (o.TeamId, o.ApiKeyHash, o.TraceId) IN (SELECT TeamId, ApiKeyHash, TraceId FROM runs)
)
SELECT named.agent_name AS agent_name,
uniqExact(named.TeamId, named.ApiKeyHash, named.TraceId) AS runs,
uniqExactIf((named.TeamId, named.ApiKeyHash, named.TraceId), runs.failed) AS failed_runs,
max(runs.start_ms) AS last_seen_ms,
arraySort(groupUniqArrayIf(named.framework, named.framework != '')) AS frameworks
FROM named
INNER JOIN runs ON named.TeamId = runs.TeamId AND named.ApiKeyHash = runs.ApiKeyHash
AND named.TraceId = runs.TraceId
GROUP BY named.agent_name
ORDER BY last_seen_ms DESC, agent_name
LIMIT {limit:UInt32}

View file

@ -13,6 +13,8 @@ SELECT o.TraceId AS trace_id, o.SpanAttributes['lens.original_trace_id'] AS orig
if(o.ToolCallId != '' OR o.ObservationType != 'tool', o.ToolCallId,
coalesce(nullIf(o.SpanAttributes['gen_ai.tool.call.id'], ''), nullIf(o.SpanAttributes['tool.id'], ''), ''))
AS tool_call_id,
o.SpanAttributes['agent.source.type'] AS source_type, o.SpanAttributes['agent.source.url'] AS source_url,
o.SpanAttributes['agent.source.title'] AS source_title,
o.UserId AS user_id, o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
FROM otel_traces AS o
WHERE o.Timestamp >= fromUnixTimestamp64Milli({start_ms:Int64})

View file

@ -12,6 +12,8 @@ SELECT o.TraceId AS trace_id, o.SpanAttributes['lens.original_trace_id'] AS orig
if(o.ToolCallId != '' OR o.ObservationType != 'tool', o.ToolCallId,
coalesce(nullIf(o.SpanAttributes['gen_ai.tool.call.id'], ''), nullIf(o.SpanAttributes['tool.id'], ''), ''))
AS tool_call_id,
o.SpanAttributes['agent.source.type'] AS source_type, o.SpanAttributes['agent.source.url'] AS source_url,
o.SpanAttributes['agent.source.title'] AS source_title,
o.UserId AS user_id, o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
FROM otel_traces AS o
WHERE o.Timestamp >= fromUnixTimestamp64Milli({start_ms:Int64})

View file

@ -13,6 +13,8 @@ SELECT o.TraceId AS trace_id, o.SpanAttributes['lens.original_trace_id'] AS orig
if(o.ToolCallId != '' OR o.ObservationType != 'tool', o.ToolCallId,
coalesce(nullIf(o.SpanAttributes['gen_ai.tool.call.id'], ''), nullIf(o.SpanAttributes['tool.id'], ''), ''))
AS tool_call_id,
o.SpanAttributes['agent.source.type'] AS source_type, o.SpanAttributes['agent.source.url'] AS source_url,
o.SpanAttributes['agent.source.title'] AS source_title,
o.UserId AS user_id, o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
FROM otel_traces AS o
WHERE o.TraceId = {trace_id:String}

View file

@ -12,6 +12,8 @@ SELECT o.TraceId AS trace_id, o.SpanAttributes['lens.original_trace_id'] AS orig
if(o.ToolCallId != '' OR o.ObservationType != 'tool', o.ToolCallId,
coalesce(nullIf(o.SpanAttributes['gen_ai.tool.call.id'], ''), nullIf(o.SpanAttributes['tool.id'], ''), ''))
AS tool_call_id,
o.SpanAttributes['agent.source.type'] AS source_type, o.SpanAttributes['agent.source.url'] AS source_url,
o.SpanAttributes['agent.source.title'] AS source_title,
o.UserId AS user_id, o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
FROM otel_traces AS o
WHERE o.TraceId = {trace_id:String}

View file

@ -6,7 +6,8 @@ const SAMPLE_READ_LIMITS: ReadLimits = ReadLimits {
..litellm_storage_clickhouse::READ_LIMITS
};
pub const LENS_QUERIES: [litellm_traces::ReadQuery; 5] = [
pub const LENS_QUERIES: [litellm_traces::ReadQuery; 6] = [
litellm_traces::ReadQuery::TraceAgents,
litellm_traces::ReadQuery::Availability,
litellm_traces::ReadQuery::Agents,
litellm_traces::ReadQuery::Sample,
@ -107,6 +108,67 @@ impl Query for LensAgents {
const SQL: &'static str = include_str!("../../query/lens_agents.sql");
}
pub struct TraceAgents;
/// Same access shape as `list_traces`: every team, the caller's own traces, or their teams' traces.
#[macro_rules_attribute::apply(wire_type)]
#[derive(Debug)]
#[serde(deny_unknown_fields)]
#[cfg_attr(feature = "schema", schemars(deny_unknown_fields))]
pub struct TraceAgentsParams {
#[serde(
deserialize_with = "super::number::boolean",
serialize_with = "litellm_traces::wire::serialize_flag"
)]
#[cfg_attr(
feature = "schema",
schemars(schema_with = "litellm_traces::schema::flag")
)]
pub all_teams: bool,
pub user_id: String,
pub team_ids: Vec<String>,
#[serde(deserialize_with = "super::number::deserialize")]
pub start_ms: i64,
#[serde(deserialize_with = "super::number::deserialize")]
pub end_ms: i64,
#[serde(deserialize_with = "super::number::deserialize")]
pub limit: u32,
}
#[macro_rules_attribute::apply(wire_type)]
#[derive(Debug)]
#[cfg_attr(feature = "schema", schemars(rename = "TraceAgentRow"))]
pub struct TraceAgentsRow {
pub agent_name: String,
#[serde(deserialize_with = "super::number::deserialize")]
#[cfg_attr(
feature = "schema",
schemars(schema_with = "crate::wire_schema::u64_number")
)]
pub runs: u64,
#[serde(deserialize_with = "super::number::deserialize")]
#[cfg_attr(
feature = "schema",
schemars(schema_with = "crate::wire_schema::u64_number")
)]
pub failed_runs: u64,
#[serde(deserialize_with = "super::number::deserialize")]
#[cfg_attr(
feature = "schema",
schemars(schema_with = "crate::wire_schema::u64_number")
)]
pub last_seen_ms: u64,
#[serde(default)]
pub frameworks: Vec<String>,
}
impl Query for TraceAgents {
type Params = TraceAgentsParams;
type Row = TraceAgentsRow;
const SQL: &'static str = include_str!("../../query/trace_agents.sql");
}
pub struct LensSample;
#[macro_rules_attribute::apply(wire_type)]

View file

@ -128,6 +128,12 @@ struct TraceSpansRowEncoding {
pub call_evidence: Option<litellm_traces::CallEvidenceKind>,
#[serde(default)]
pub tool_call_id: String,
#[serde(default)]
pub source_type: String,
#[serde(default)]
pub source_url: String,
#[serde(default)]
pub source_title: String,
pub team_id: String,
pub api_key_hash: String,
pub user_id: String,
@ -349,7 +355,7 @@ mod tests {
quoted,
);
round_trip::<TraceSpansRow>(
json!({"trace_id": "trace", "original_trace_id": "original", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "claude-agent-sdk", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "team_id": "team", "api_key_hash": "key", "user_id": "user"}),
json!({"trace_id": "trace", "original_trace_id": "original", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "claude-agent-sdk", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "source_type": "slack", "source_url": "https://acme.slack.com/archives/C1/p1", "source_title": "thread", "team_id": "team", "api_key_hash": "key", "user_id": "user"}),
quoted,
);
round_trip::<SpanDetailRow>(

View file

@ -17,6 +17,7 @@ pub async fn execute_named_read(
) -> Result<String, Error> {
match query {
ReadQuery::ListTraces => named_json::<ListTraces>(client, connection, parameters).await,
ReadQuery::TraceAgents => named_json::<TraceAgents>(client, connection, parameters).await,
ReadQuery::TraceIdentity => {
named_json::<TraceIdentity>(client, connection, parameters).await
}

View file

@ -89,6 +89,8 @@ pub fn schemas() -> BTreeMap<&'static str, Schema> {
("PartRow", received::<lens::LensContentRow>()),
("CountRow", received::<lens::LensEvidenceRow>()),
("AgentRow", received::<lens::LensAgentsRow>()),
("TraceAgentsParams", received::<lens::TraceAgentsParams>()),
("TraceAgentRow", received::<lens::TraceAgentsRow>()),
("TraceQueryHelp", crate::query::help_schema()),
])
}

View file

@ -754,6 +754,145 @@ async fn listed_agent_names_preserve_scope_and_cursor(
Ok(())
}
#[rstest]
#[tokio::test]
async fn trace_agents_count_runs_and_failures_within_scope_and_window(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let now = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let old = now - 3 * 86_400_000_000_000_i64;
for (team, trace, span, parent, agent, status, framework, timestamp) in [
(
"alpha",
"run-1",
"root",
"",
"moyai",
"STATUS_CODE_OK",
"pi",
now,
),
(
"alpha",
"run-1",
"tool",
"root",
"moyai",
"STATUS_CODE_ERROR",
"pi",
now,
),
(
"alpha",
"run-2",
"root",
"",
"moyai",
"STATUS_CODE_OK",
"",
now - 1_000_000,
),
(
"alpha",
"run-3",
"root",
"",
"research",
"STATUS_CODE_OK",
"",
now - 2_000_000,
),
(
"alpha",
"old-run",
"root",
"",
"moyai",
"STATUS_CODE_ERROR",
"",
old,
),
(
"beta",
"other-team",
"root",
"",
"moyai",
"STATUS_CODE_ERROR",
"",
now,
),
(
"beta",
"other-agent",
"root",
"",
"hidden_agent",
"STATUS_CODE_OK",
"",
now,
),
] {
insert_rows(
&database,
"otel_traces",
vec![serde_json::from_value(serde_json::json!({
"Timestamp": timestamp, "TraceId": trace, "SpanId": span, "ParentSpanId": parent,
"ServiceName": "app", "SpanName": span, "AgentName": agent, "UserId": "owner",
"StatusCode": status, "Framework": framework, "ObservationType": "agent",
"ResourceAttributes": {"litellm.team_id": team, "litellm.api_key_hash": "key"}
}))?],
)
.await?;
}
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
let parameters = BTreeMap::from([
("all_teams".into(), Parameter::Integer(0)),
("user_id".into(), Parameter::Text(String::new())),
("team_ids".into(), Parameter::Strings(vec!["alpha".into()])),
(
"start_ms".into(),
Parameter::Integer(now / 1_000_000 - 86_400_000),
),
("end_ms".into(), Parameter::Integer(now / 1_000_000 + 1000)),
("limit".into(), Parameter::Integer(10)),
]);
let agents: serde_json::Value = serde_json::from_str(
&execute_named_read(
&database.client,
&connection,
ReadQuery::TraceAgents,
&parameters,
)
.await?,
)?;
let rows = agents["data"].as_array().ok_or("missing agents")?;
let summary = rows
.iter()
.map(|row| {
(
row["agent_name"].as_str().unwrap_or_default(),
(
row["runs"].to_string().trim_matches('"').to_owned(),
row["failed_runs"].to_string().trim_matches('"').to_owned(),
row["frameworks"].clone(),
),
)
})
.collect::<Vec<_>>();
assert_eq!(
summary,
vec![
("moyai", ("2".into(), "1".into(), serde_json::json!(["pi"]))),
("research", ("1".into(), "0".into(), serde_json::json!([]))),
]
);
Ok(())
}
#[rstest]
#[tokio::test]
async fn rollup_merges_spans_across_days_without_losing_root_fields(

View file

@ -44,6 +44,6 @@ pub use tenant::Tenant;
pub use truncate::{truncate_messages, truncate_value};
pub use ui::{ChatRole, UiContent, UiField, UiMessage, UiToolCall, to_ui_content};
pub use view::{
AgentNode, Span, SpanDetail, SpanErrorPage, SpanStatus, SpendMatch, Trace, TracePage,
TraceSummary,
AgentNode, RunSource, RunSourceType, Span, SpanDetail, SpanErrorPage, SpanStatus, SpendMatch,
Trace, TracePage, TraceSummary,
};

View file

@ -5,6 +5,7 @@ pub mod named;
#[strum(serialize_all = "snake_case")]
pub enum ReadQuery {
ListTraces,
TraceAgents,
TraceSpans,
TracePageSpans,
TraceIdentity,

View file

@ -110,6 +110,12 @@ pub struct TraceSpansRow {
pub call_evidence: Option<crate::CallEvidenceKind>,
#[serde(default)]
pub tool_call_id: String,
#[serde(default)]
pub source_type: String,
#[serde(default)]
pub source_url: String,
#[serde(default)]
pub source_title: String,
pub team_id: String,
pub api_key_hash: String,
pub user_id: String,

View file

@ -6,7 +6,9 @@ use time::OffsetDateTime;
use crate::{
normalize::ObservationType,
query::named::{ListTracesRow, SpendByResponseIdsRow as SpendRow, TraceSpansRow},
view::{AgentNode, Span, SpanStatus, SpendMatch, Trace, TraceSummary},
view::{
AgentNode, RunSource, RunSourceType, Span, SpanStatus, SpendMatch, Trace, TraceSummary,
},
};
use super::{
@ -138,6 +140,15 @@ pub fn iso_time(ms: i64) -> String {
)
}
fn source(row: &TraceSpansRow) -> Option<RunSource> {
row.source_url.starts_with("https://").then(|| RunSource {
kind: serde_json::from_value(serde_json::Value::from(row.source_type.as_str()))
.unwrap_or(RunSourceType::Custom),
url: row.source_url.clone(),
title: row.source_title.clone(),
})
}
fn sorted_unique<'a>(values: impl Iterator<Item = &'a str>) -> Vec<String> {
values
.filter(|value| !value.is_empty())
@ -229,6 +240,12 @@ pub fn resolve_trace(
models: sorted_unique(calls.iter().map(|call| rows[*call].model.as_str())),
spend: priced.spend,
priced_calls: priced.priced_calls,
source: source(&rows[root]).or_else(|| {
rows.iter()
.filter_map(|row| Some((row.start_ns, source(row)?)))
.min_by_key(|(start_ns, _)| *start_ns)
.map(|(_, source)| source)
}),
};
Some(Trace {
summary,
@ -266,5 +283,6 @@ pub fn listed_summary(row: &ListTracesRow) -> TraceSummary {
models: row.models.clone(),
spend: None,
priced_calls: 0,
source: None,
}
}

View file

@ -66,6 +66,30 @@ pub struct AgentNode {
pub priced_calls: u64,
}
#[macro_rules_attribute::apply(wire_type)]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum RunSourceType {
Slack,
Teams,
Discord,
Linear,
Github,
Jira,
#[serde(other)]
Custom,
}
/// The conversation that started the run, from the `agent.source.*` span attributes.
#[macro_rules_attribute::apply(response_type)]
#[derive(Clone, Debug, PartialEq)]
pub struct RunSource {
#[serde(rename = "type")]
pub kind: RunSourceType,
pub url: String,
pub title: String,
}
#[macro_rules_attribute::apply(response_type)]
#[derive(Clone, Debug, PartialEq)]
pub struct TraceSummary {
@ -95,6 +119,9 @@ pub struct TraceSummary {
pub models: Vec<String>,
pub spend: Option<f64>,
pub priced_calls: u64,
#[serde(skip_serializing_if = "Option::is_none")]
#[cfg_attr(feature = "schema", schemars(extend("x-python-optional" = true)))]
pub source: Option<RunSource>,
}
#[macro_rules_attribute::apply(response_type)]

View file

@ -207,6 +207,9 @@ fn trace_span(span: DecodedSpan) -> TraceSpansRow {
call_keys,
call_evidence: Some(normalized.calls.kind()),
tool_call_id: normalized.tool_call_id.unwrap_or_default(),
source_type: String::new(),
source_url: String::new(),
source_title: String::new(),
team_id: "fixture-team".into(),
api_key_hash: "fixture-key".into(),
user_id: "fixture-user".into(),
@ -291,6 +294,9 @@ fn unrelated_transport(call: &TraceSpansRow) -> TraceSpansRow {
call_keys: vec![CallKey::Transport],
call_evidence: Some(CallEvidenceKind::Complete),
tool_call_id: String::new(),
source_type: String::new(),
source_url: String::new(),
source_title: String::new(),
team_id: call.team_id.clone(),
api_key_hash: call.api_key_hash.clone(),
user_id: call.user_id.clone(),

View file

@ -3,6 +3,7 @@ use rstest::rstest;
#[rstest]
#[case::list_traces("list_traces", ReadQuery::ListTraces)]
#[case::trace_agents("trace_agents", ReadQuery::TraceAgents)]
#[case::trace_spans("trace_spans", ReadQuery::TraceSpans)]
#[case::span_detail("span_detail", ReadQuery::SpanDetail)]
#[case::span_error("span_error", ReadQuery::SpanError)]

View file

@ -53,7 +53,7 @@ fn result_contracts_preserve_public_field_names() {
json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "STATUS_CODE_OK", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["framework"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}),
);
round_trip::<TraceSpansRow>(
json!({"trace_id": "trace", "original_trace_id": "original", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "framework", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "team_id": "team", "api_key_hash": "key", "user_id": "user"}),
json!({"trace_id": "trace", "original_trace_id": "original", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "framework", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "source_type": "slack", "source_url": "https://acme.slack.com/archives/C1/p1", "source_title": "thread", "team_id": "team", "api_key_hash": "key", "user_id": "user"}),
);
round_trip::<SpanDetailRow>(
json!({"span_id": "span", "input": "input", "output": "output", "attributes": {"count": "42"}}),

View file

@ -1,5 +1,5 @@
use litellm_traces::{
AgentNode, SpanStatus, SpendMatch, iso_time, listed_summary,
AgentNode, RunSourceType, SpanStatus, SpendMatch, iso_time, listed_summary,
query::named::{ListTracesRow, SpendByResponseIdsRow, TraceSpansRow},
resolve_trace,
};
@ -33,6 +33,9 @@ fn row(span_id: &str, parent: &str, name: &str, kind: &str, agent: &str) -> Trac
call_keys: Vec::new(),
call_evidence: None,
tool_call_id: String::new(),
source_type: String::new(),
source_url: String::new(),
source_title: String::new(),
team_id: "team".into(),
api_key_hash: "key".into(),
user_id: String::new(),
@ -173,6 +176,57 @@ fn summary_counts_model_calls_tools_and_agents() {
assert_eq!(summary.spend, None);
}
fn sourced(mut span: TraceSpansRow, url: &str, title: &str) -> TraceSpansRow {
span.source_url = url.into();
span.source_title = title.into();
span
}
const THREAD: &str = "https://acme.slack.com/archives/C1/p1";
#[rstest]
#[case::root_wins(
vec![sourced(at(row("root", "", "agent", "agent", "agent"), 5, 10), THREAD, "root thread"),
sourced(at(row("tool", "root", "tool", "tool", "agent"), 0, 1), "https://other.example/", "child")],
Some((THREAD, "root thread")),
)]
#[case::earliest_child_when_root_has_none(
vec![at(row("root", "", "agent", "agent", "agent"), 0, 10),
sourced(at(row("late", "root", "tool", "tool", "agent"), 5, 1), "https://late.example/", "late"),
sourced(at(row("early", "root", "tool", "tool", "agent"), 2, 1), THREAD, "early")],
Some((THREAD, "early")),
)]
#[case::non_https_is_dropped(
vec![sourced(row("root", "", "agent", "agent", "agent"), "javascript:alert(1)", "x")],
None,
)]
#[case::absent(vec![row("root", "", "agent", "agent", "agent")], None)]
fn summary_source_links_where_the_run_started(
#[case] rows: Vec<TraceSpansRow>,
#[case] expected: Option<(&str, &str)>,
) {
let source = resolve_trace("t", "", &rows, &[]).unwrap().summary.source;
assert_eq!(
source
.as_ref()
.map(|source| (source.url.as_str(), source.title.as_str())),
expected
);
}
#[rstest]
#[case::slack("slack", RunSourceType::Slack)]
#[case::teams("teams", RunSourceType::Teams)]
#[case::custom("custom", RunSourceType::Custom)]
#[case::unknown_is_custom("my-bot", RunSourceType::Custom)]
#[case::missing_is_custom("", RunSourceType::Custom)]
fn summary_source_type_picks_the_app(#[case] source_type: &str, #[case] expected: RunSourceType) {
let mut root = sourced(row("root", "", "agent", "agent", "agent"), THREAD, "t");
root.source_type = source_type.into();
let source = resolve_trace("t", "", &[root], &[]).unwrap().summary.source;
assert_eq!(source.map(|source| source.kind), Some(expected));
}
#[rstest]
fn spans_are_offset_from_the_trace_start() {
let trace = resolve_trace("t1", "", &deep_agent(1), &[]).unwrap();

View file

@ -1,6 +1,5 @@
import inspect
from collections.abc import Awaitable, Callable, Coroutine, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
from litellm import main
@ -8,13 +7,13 @@ from litellm.rust_bridge.catalog import Route, RouteContext
from litellm.rust_bridge.chat_completions.entrypoints import (
NATIVE_ACOMPLETION,
NATIVE_COMPLETION,
LiteLLMChatCompletionsRequest,
)
from litellm.rust_bridge.dispatch import PublicDispatch, call_hook
from litellm.rust_bridge.dispatch import PublicDispatch
from litellm.rust_bridge.public_call import (
NativeCall,
bind,
optional_bool,
optional_mapping,
native_call,
native_call_hook,
optional_sequence,
optional_str,
signature,
@ -51,33 +50,22 @@ _ACOMPLETION: Final = signature(_PYTHON_ACOMPLETION)
def _public_request(
legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> LiteLLMChatCompletionsRequest | None:
) -> NativeCall | None:
fields: Final = bind(legacy, args, kwargs)
if fields is None:
return None
model: Final = fields.get("model")
messages: Final = optional_sequence(fields.get("messages"))
extra: Final = optional_mapping(fields.get("kwargs")) or MappingProxyType({})
if not isinstance(model, str) or messages is None:
return None
return LiteLLMChatCompletionsRequest(
model=model,
messages=messages,
stream=optional_bool(fields.get("stream")),
api_key=optional_str(fields.get("api_key")),
api_base=optional_str(extra.get("api_base")) or optional_str(fields.get("base_url")),
custom_llm_provider=optional_str(extra.get("custom_llm_provider")),
extra_headers=optional_mapping(fields.get("extra_headers")),
kwargs=extra,
parameters=MappingProxyType({name: value for name, value in fields.items() if name != "kwargs"}),
)
return native_call(args, kwargs, fields)
def _context(request: LiteLLMChatCompletionsRequest) -> RouteContext:
def _context(request: NativeCall) -> RouteContext:
return RouteContext(
Route.CHAT_COMPLETIONS,
provider=request.custom_llm_provider,
model=request.model,
provider=optional_str(request.bound.get("custom_llm_provider")),
model=str(request.bound["model"]),
)
@ -105,7 +93,7 @@ def completion(
kwargs,
python=python,
binding=NATIVE_COMPLETION,
native=call_hook,
native=native_call_hook,
)
@ -116,7 +104,7 @@ async def acompletion(*args: object, **kwargs: object) -> ChatResult: # kwargs-
kwargs,
python=python,
binding=NATIVE_ACOMPLETION,
native=call_hook,
native=native_call_hook,
)

View file

@ -24,6 +24,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError
import litellm
from litellm import ModelResponse
from litellm._logging import verbose_logger
from litellm.integrations.anthropic_cache_control_hook import supports_openai_prompt_cache_breakpoint
from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params
from litellm.litellm_core_utils.prompt_templates.common_utils import (
responses_reasoning_items_from_thinking_blocks,
@ -81,6 +82,38 @@ _RESPONSES_API_ONLY_FIELDS: Final = frozenset((*Response.model_fields, *Response
)
def _strip_prompt_cache_breakpoint_from_content_block(value: object) -> object:
if not isinstance(value, dict):
return value
content_block: Final = cast(dict[str, object], value) # cast-ok: isinstance confirms the content block is a mapping
return {key: item for key, item in content_block.items() if key != "prompt_cache_breakpoint"}
def _strip_prompt_cache_breakpoints_from_content(value: object) -> object:
if isinstance(value, list):
list_content: Final = cast(list[object], value) # cast-ok: isinstance confirms a list of content blocks
return [_strip_prompt_cache_breakpoint_from_content_block(item) for item in list_content]
if isinstance(value, tuple):
tuple_content: Final = cast(tuple[object, ...], value) # cast-ok: isinstance confirms a tuple of content blocks
return tuple(_strip_prompt_cache_breakpoint_from_content_block(item) for item in tuple_content)
return _strip_prompt_cache_breakpoint_from_content_block(value)
def _strip_prompt_cache_breakpoints_from_item(value: object) -> object:
if not isinstance(value, dict):
return value
input_item: Final = cast(dict[str, object], value) # cast-ok: isinstance confirms a Responses input item mapping
return {
key: _strip_prompt_cache_breakpoints_from_content(item) if key in ("content", "output") else item
for key, item in input_item.items()
if key != "prompt_cache_breakpoint"
}
def _strip_prompt_cache_breakpoints(input_items: list[object]) -> list[object]:
return [_strip_prompt_cache_breakpoints_from_item(item) for item in input_items]
def _provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]:
return MappingProxyType(
{
@ -378,7 +411,24 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return None, index
def convert_chat_completion_messages_to_responses_api(
self, messages: list["AllMessageValues"], drop_params: bool = False
self,
messages: list["AllMessageValues"],
*,
drop_params: bool = False,
keep_prompt_cache_breakpoints: bool = False,
) -> tuple[list[object], str | None]:
converted_input_items, instructions = self._convert_chat_completion_messages_to_responses_input(
messages, drop_params=drop_params
)
return (
converted_input_items
if keep_prompt_cache_breakpoints
else _strip_prompt_cache_breakpoints(converted_input_items),
instructions,
)
def _convert_chat_completion_messages_to_responses_input(
self, messages: list["AllMessageValues"], *, drop_params: bool = False
) -> tuple[list[object], str | None]:
input_items: Final[list[object]] = []
instructions: str | None = None
@ -614,26 +664,32 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
litellm_logging_obj: "LiteLLMLoggingObj",
client: object | None = None,
) -> dict:
(
input_items,
instructions,
) = self.convert_chat_completion_messages_to_responses_api(
messages, drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params)
base_model: Final = litellm_params.get("base_model")
supports_prompt_cache_breakpoint: Final = supports_openai_prompt_cache_breakpoint(model) or (
isinstance(base_model, str) and bool(base_model) and supports_openai_prompt_cache_breakpoint(base_model)
)
converted_input_items, converted_instructions = self.convert_chat_completion_messages_to_responses_api(
messages,
drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params),
keep_prompt_cache_breakpoints=supports_prompt_cache_breakpoint,
)
# OpenAI's Responses API rejects an empty input. For a system-only
# request, carry the system message as a system-role input item instead
# of instructions, mirroring how non-string system content is already
# handled in convert_chat_completion_messages_to_responses_api.
if not input_items and instructions is not None:
input_items = [
is_system_only_request: Final = not converted_input_items and converted_instructions is not None
input_items: Final = (
[
{
"type": "message",
"role": "system",
"content": [{"type": "input_text", "text": instructions}],
"content": [{"type": "input_text", "text": converted_instructions}],
}
]
instructions = None
if is_system_only_request
else converted_input_items
)
instructions: Final = None if is_system_only_request else converted_instructions
optional_params = self._extract_extra_body_params(optional_params)

View file

@ -60,6 +60,7 @@ TRACE_READ_RETRY_AFTER_SECONDS: Final = get_env_int("TRACE_READ_RETRY_AFTER_SECO
OTLP_MAX_CONCURRENT_INGESTS: Final = get_env_int("OTLP_MAX_CONCURRENT_INGESTS", 2)
AGENT_TRACING_INPUT_PREVIEW_CHARS: Final = get_env_int("AGENT_TRACING_INPUT_PREVIEW_CHARS", 240)
AGENT_TRACING_LIST_PAGE_SIZE: Final = get_env_int("AGENT_TRACING_LIST_PAGE_SIZE", 50)
AGENT_TRACING_AGENT_LIST_LIMIT: Final = get_env_int("AGENT_TRACING_AGENT_LIST_LIMIT", 500)
LENS_DATASET_MAX_CASES: Final = get_env_int("LENS_DATASET_MAX_CASES", 200)
LENS_DATASET_MAX_CASE_CHARS: Final = get_env_int("LENS_DATASET_MAX_CASE_CHARS", 20_000)
LENS_DATASET_TRACE_PAGE_SIZE: Final = 500

View file

@ -1,17 +1,15 @@
import inspect
from collections.abc import Awaitable, Callable, Coroutine, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
from litellm import main
from litellm.rust_bridge.catalog import Route, RouteContext
from litellm.rust_bridge.dispatch import PublicDispatch, call_hook
from litellm.rust_bridge.dispatch import PublicDispatch
from litellm.rust_bridge.embeddings.entrypoints import (
NATIVE_AEMBEDDING,
NATIVE_EMBEDDING,
LiteLLMEmbeddingRequest,
)
from litellm.rust_bridge.public_call import bind, optional_mapping, optional_str, signature
from litellm.rust_bridge.public_call import NativeCall, bind, native_call, native_call_hook, optional_str, signature
from litellm.types.utils import EmbeddingResponse
__all__ = ("aembedding", "embedding")
@ -30,28 +28,24 @@ _EMBEDDING_SIGNATURE: Final = signature(_PYTHON_EMBEDDING)
def _public_request(
legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> LiteLLMEmbeddingRequest | None:
) -> NativeCall | None:
fields: Final = bind(legacy, args, kwargs)
if fields is None:
return None
model: Final = fields.get("model")
if not isinstance(model, str):
return None
extra: Final = optional_mapping(fields.get("kwargs")) or MappingProxyType({})
return LiteLLMEmbeddingRequest(
model=model,
input=fields.get("input"),
api_key=optional_str(fields.get("api_key")),
api_base=optional_str(fields.get("api_base")),
custom_llm_provider=optional_str(fields.get("custom_llm_provider")),
kwargs=extra,
return native_call(args, kwargs, fields)
def _context(request: NativeCall) -> RouteContext:
return RouteContext(
Route.EMBEDDINGS,
provider=optional_str(request.bound.get("custom_llm_provider")),
model=str(request.bound["model"]),
)
def _context(request: LiteLLMEmbeddingRequest) -> RouteContext:
return RouteContext(Route.EMBEDDINGS, provider=request.custom_llm_provider, model=request.model)
_DISPATCH: Final = PublicDispatch(
route=Route.EMBEDDINGS,
request=lambda args, kwargs: _public_request(_EMBEDDING_SIGNATURE, args, kwargs),
@ -75,7 +69,7 @@ def embedding(
kwargs,
python=_PYTHON_EMBEDDING,
binding=NATIVE_EMBEDDING,
native=call_hook,
native=native_call_hook,
)
@ -85,7 +79,7 @@ async def aembedding(*args: object, **kwargs: object) -> EmbeddingResponse: # k
kwargs,
python=_PYTHON_AEMBEDDING,
binding=NATIVE_AEMBEDDING,
native=call_hook,
native=native_call_hook,
)

View file

@ -24,6 +24,7 @@ from litellm.types.guardrails import (
DynamicGuardrailParams,
GuardrailEventHooks,
LitellmParams,
LoggingOnlyScope,
Mode,
)
from litellm.types.llms.openai import AllMessageValues
@ -180,6 +181,7 @@ class CustomGuardrail(CustomLogger):
use_native_lifecycle_hooks: ClassVar[bool] = False
records_own_guardrail_information: ClassVar[bool] = False
logging_only_scope: LoggingOnlyScope | None
timeout: float | httpx.Timeout | None = None
@ -256,6 +258,7 @@ class CustomGuardrail(CustomLogger):
self.run_in_parallel: bool = run_in_parallel
self.scan_raw_request: bool = scan_raw_request
self.only_scan_new_messages: bool = only_scan_new_messages
self.logging_only_scope = None
if timeout is not None:
self.timeout = timeout
@ -817,6 +820,13 @@ class CustomGuardrail(CustomLogger):
def uses_apply_guardrail_interface(self) -> bool:
return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail
@classmethod
def supports_logging_only_scope(cls) -> bool:
return (
cls.apply_guardrail is not CustomGuardrail.apply_guardrail
and cls.async_logging_hook is CustomGuardrail.async_logging_hook
)
def _deployment_hook_target(self) -> "CustomLogger":
if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
return self
@ -948,7 +958,7 @@ class CustomGuardrail(CustomLogger):
result: object,
call_type: str,
) -> tuple[dict, object]: # mutable-ok: CustomLogger.async_logging_hook contract
"""logging_only: run apply_guardrail on copies of the logged request/response and record the verdict."""
"""logging_only: scan copies of the logged request and/or response according to logging_only_scope."""
from litellm.llms import get_guardrail_translation_mapping
if not self.uses_apply_guardrail_interface():
@ -995,6 +1005,28 @@ class CustomGuardrail(CustomLogger):
"standard_logging_object": {**standard_logging_object, "guardrail_information": [*existing, *entries]},
}, result
def _copy_scratch_request_fields(
self,
kwargs: Mapping[str, object],
) -> tuple[object, object] | None:
optional_params: Final = kwargs.get("optional_params")
try:
return (
copy.deepcopy(kwargs.get("messages") or kwargs.get("input")),
copy.deepcopy(optional_params.get("tools") if isinstance(optional_params, Mapping) else None),
)
except Exception as e:
if self.logging_only_scope == "output":
return None
if self.logging_only_scope == "both":
verbose_logger.warning(
"Guardrail %s: logging_only request copy failed, skipping request scan: %s",
self.guardrail_name,
e,
)
return None
raise
async def _scan_logged_call(
self,
kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract
@ -1003,18 +1035,25 @@ class CustomGuardrail(CustomLogger):
output_translation: "BaseTranslation",
scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata
) -> None:
optional_params: Final = kwargs.get("optional_params") or {}
scratch_input: Final = copy.deepcopy(kwargs.get("messages") or kwargs.get("input"))
scratch_fields: Final = self._copy_scratch_request_fields(kwargs)
scratch_input, scratch_tools = scratch_fields or (None, None)
scratch_request: Final = {
"model": kwargs.get("model"),
"messages": scratch_input,
"input": scratch_input,
"tools": copy.deepcopy(optional_params.get("tools")),
"tools": scratch_tools,
"litellm_call_id": kwargs.get("litellm_call_id"),
"metadata": scratch_metadata,
}
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
if response is None:
if self.logging_only_scope != "output" and scratch_fields is not None:
if self.logging_only_scope == "both":
try:
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
except Exception as e: # noqa: BLE001 # one direction's scan failure must not drop the other direction's verdict
verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e)
else:
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
if response is None or self.logging_only_scope == "input":
return
await output_translation.process_output_response(
response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request

View file

@ -353,6 +353,8 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
"LiteLLM_LensRun",
"LiteLLM_LensReview",
"LiteLLM_LensWorker",
"LiteLLM_LensSignalConfig",
"LiteLLM_LensTraceSignal",
)
)
PRISMA_RELATIONS: Final[frozenset[str]] = _PRISMA_MODELS | _PRISMA_VIEWS

View file

@ -291,7 +291,7 @@ else:
_GENERIC_API_LOGGER_CLS: Final = GenericAPILogger
_in_memory_loggers: Final[list[CustomLogger]] = []
_STANDARD_LOGGING_METADATA_RESOLVED_KEYS: Final[frozenset[str]] = frozenset(("used_client_oauth_token",))
_STANDARD_LOGGING_METADATA_RESOLVED_KEYS: Final[frozenset[str]] = frozenset(("used_client_oauth_token", "usage_object"))
_STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = (
frozenset(StandardLoggingMetadata.__annotations__.keys()) - _STANDARD_LOGGING_METADATA_RESOLVED_KEYS
)
@ -5993,7 +5993,7 @@ class StandardLoggingPayloadSetup:
Like get_usage_from_response_obj but returns a plain dict, skipping
the Pydantic Usage construction on the hot path.
"""
_empty: Final[dict] = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
_empty: Final[dict[str, object]] = {}
if combined_usage_object is not None:
return combined_usage_object.model_dump()
if not response_obj:

View file

@ -5188,7 +5188,7 @@ def function_call_prompt(
messages: list[dict[str, object]],
functions: list[object],
) -> list[dict[str, object]]:
function_prompt = """Produce JSON OUTPUT ONLY! Adhere to this format {"name": "function_name", "arguments":{"argument_name": "argument_value"}} The following functions are available to you:"""
function_prompt = """To call a function, reply with JSON ONLY in this format {"name": "function_name", "arguments":{"argument_name": "argument_value"}}. Once a function result answers the request, reply to the user in plain text instead of calling a function again. The following functions are available to you:"""
for function in functions:
function_prompt += f"""\n{function}\n"""

View file

@ -2303,9 +2303,29 @@ class CustomStreamWrapper:
429 (rate-limit) is explicitly exempted from the 4xx filter because
it is transient and the Router should switch to another model group.
An error an inner stream already wrapped (the chat-to-Responses bridge
consumes a Responses stream) is rebuilt around the provider exception
with this wrapper's own bookkeeping, so the Router's one-level unwrap
surfaces the provider exception and is_pre_first_chunk says whether
this wrapper's consumer received anything (the inner stream counts a
lifecycle event the bridge never forwards as its first chunk).
"""
from litellm.exceptions import MidStreamFallbackError
if isinstance(e, MidStreamFallbackError):
self._restore_consumer_correlation_context()
if e.original_exception is None:
raise e
raise MidStreamFallbackError(
message=str(e.original_exception),
model=self.model,
llm_provider=self.custom_llm_provider or "anthropic",
original_exception=e.original_exception,
generated_content=self.response_uptil_now,
is_pre_first_chunk=not self.sent_first_chunk,
)
# Map to OpenAI exception format. Some providers' mappers (e.g.
# _map_anthropic_exception, _map_aleph_alpha_exception) synchronously
# log a debug diagnostic (the raw status code) as part of mapping -

View file

@ -12,6 +12,7 @@ from pydantic import JsonValue, TypeAdapter
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.llms.anthropic.common_utils import AnthropicError
from litellm.llms.anthropic.count_tokens.transformation import (
AnthropicCountTokensConfig,
@ -62,7 +63,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
verbose_logger.debug("Processing Anthropic CountTokens request for model: %s", model)
# Transform request to Anthropic format
request_body: Final = self.transform_request_to_count_tokens(
request_body: Final = await asyncify(self.transform_request_to_count_tokens)(
model=model,
messages=messages,
tools=tools,

View file

@ -13,11 +13,41 @@ from pydantic import JsonValue, TypeAdapter
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers
from litellm.llms.anthropic.wif import resolve_anthropic_base
from litellm.types.llms.openai import ChatCompletionImageObject
_COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue])
_IMAGE_BLOCK: Final = TypeAdapter(ChatCompletionImageObject)
COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config")
def _count_image(block: JsonValue) -> JsonValue:
if not isinstance(block, dict) or block.get("type") != "image_url":
return block
from litellm.litellm_core_utils.prompt_templates.factory import convert_to_anthropic_image_obj
image_block: Final = _IMAGE_BLOCK.validate_python(block)
image_url: Final = image_block["image_url"]
source: Final = convert_to_anthropic_image_obj(
openai_image_url=image_url if isinstance(image_url, str) else image_url["url"],
format=image_url.get("format") if isinstance(image_url, dict) else None,
)
image: Final = _COUNT_REQUEST.validate_python({"type": "image", "source": source})
return {**{key: value for key, value in block.items() if key not in {"type", "image_url"}}, **image}
def _count_block(block: JsonValue) -> JsonValue:
if not isinstance(block, dict) or block.get("type") != "tool_result":
return _count_image(block)
content: Final = block.get("content")
if not isinstance(content, list):
return block
return {**block, "content": [_count_image(part) for part in content]}
def _count_content(content: JsonValue) -> JsonValue:
return [_count_block(block) for block in content] if isinstance(content, list) else content
class AnthropicCountTokensConfig:
"""
Configuration and transformation logic for Anthropic CountTokens API.
@ -62,7 +92,7 @@ class AnthropicCountTokensConfig:
MappingProxyType(
{
"model": model,
"messages": messages,
"messages": [{**message, "content": _count_content(message["content"])} for message in messages],
**MappingProxyType(
{key: value for key, value in (("system", system), ("tools", tools)) if value is not None}
),

View file

@ -10,6 +10,7 @@ import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.llms.anthropic.common_utils import AnthropicError
from litellm.llms.azure_ai.anthropic.count_tokens.transformation import (
AzureAIAnthropicCountTokensConfig,
@ -59,7 +60,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
verbose_logger.debug("Processing Azure AI Anthropic CountTokens request for model: %s", model)
# Transform request to Anthropic format
request_body: Final = self.transform_request_to_count_tokens(
request_body: Final = await asyncify(self.transform_request_to_count_tokens)(
model=model,
messages=messages,
tools=tools,

View file

@ -65,6 +65,26 @@ class BaseLLMException(Exception):
super().__init__(self.message) # Call the base class constructor with the parameters it needs
_NO_ATTRIBUTION_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType({})
def with_attribution_headers(
attribution_headers: Mapping[str, str],
headers: dict[str, str] | None, # mutable-ok: returned as-is when there is nothing to add
) -> dict[str, str] | None: # mutable-ok: becomes the request's outbound headers
"""
`headers` plus any attribution header the caller didn't already set (names
compared case-insensitively). Builds a new dict; `headers` is never mutated.
"""
if not attribution_headers:
return headers
caller_names: Final = {name.lower() for name in headers or {}}
return {
**{name: value for name, value in attribution_headers.items() if name.lower() not in caller_names},
**(headers or {}),
}
class BaseConfig(ABC):
def __init__(self):
pass
@ -89,6 +109,15 @@ class BaseConfig(ABC):
and not callable(v) # Filter out any callable objects including mocks
}
def get_attribution_headers(self) -> Mapping[str, str]:
"""
Headers that tell the provider a request came through LiteLLM.
Sent by default on every request; a caller header with the same name
(any casing) wins. Override in a provider config to opt in.
"""
return _NO_ATTRIBUTION_HEADERS
def get_json_schema_from_pydantic_object(self, response_format: type[BaseModel] | dict | None) -> dict | None:
return type_to_response_format_param(response_format=response_format)

View file

@ -6,6 +6,7 @@ import httpx
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
from litellm.rust_bridge import runtime
from litellm.rust_bridge.catalog import Route, RouteContext
from litellm.rust_bridge.public_call import NativeCall
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.rust_bridge.transcription.native import (
NATIVE_ATRANSCRIPTION,
@ -52,18 +53,18 @@ class BedrockAudioTranscriptionRustDispatch:
timeout: float | httpx.Timeout | None,
) -> TranscriptionResponse:
def native(rust: RustTranscription) -> TranscriptionResponse:
return TranscriptionResponse(
**rust(
model=model,
audio=self._audio_payload(audio_file),
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout_seconds=timeout_to_seconds(timeout),
)
)
fields: Final = {
"model": model,
"audio": self._audio_payload(audio_file),
"api_key": api_key,
"api_base": api_base,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"optional_params": optional_params,
"timeout_seconds": timeout_to_seconds(timeout),
}
call: Final = NativeCall(args=(), kwargs=fields, bound=fields)
return TranscriptionResponse(**rust(call))
return runtime.run(
RouteContext(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model),
@ -85,18 +86,18 @@ class BedrockAudioTranscriptionRustDispatch:
timeout: float | httpx.Timeout | None,
) -> TranscriptionResponse:
async def native(rust: RustAtranscription) -> TranscriptionResponse:
return TranscriptionResponse(
**await rust(
model=model,
audio=self._audio_payload(audio_file),
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout_seconds=timeout_to_seconds(timeout),
)
)
fields: Final = {
"model": model,
"audio": self._audio_payload(audio_file),
"api_key": api_key,
"api_base": api_base,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"optional_params": optional_params,
"timeout_seconds": timeout_to_seconds(timeout),
}
call: Final = NativeCall(args=(), kwargs=fields, bound=fields)
return TranscriptionResponse(**await rust(call))
return await runtime.arun(
RouteContext(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model),

View file

@ -6,11 +6,20 @@ Calls done in OpenAI/openai.py as Novita AI is openai-compatible.
Docs: https://novita.ai/docs/guides/llm-api
"""
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from ....types.llms.openai import AllMessageValues
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
_NOVITA_ATTRIBUTION_HEADERS: Final[Mapping[str, str]] = MappingProxyType({"X-Novita-Source": "litellm"})
class NovitaConfig(OpenAIGPTConfig):
def get_attribution_headers(self) -> Mapping[str, str]:
return _NOVITA_ATTRIBUTION_HEADERS
def validate_environment(
self,
headers: dict,
@ -27,5 +36,6 @@ class NovitaConfig(OpenAIGPTConfig):
)
headers["Authorization"] = f"Bearer {api_key}"
headers["Content-Type"] = "application/json"
headers["X-Novita-Source"] = "litellm"
if not any(name.lower() == "x-novita-source" for name in headers):
headers["X-Novita-Source"] = "litellm"
return headers

View file

@ -2,6 +2,8 @@
Translate from OpenAI's `/v1/chat/completions` to Perplexity's `/v1/chat/completions`
"""
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import httpx
@ -17,12 +19,17 @@ from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage
if TYPE_CHECKING:
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
_PERPLEXITY_ATTRIBUTION_HEADERS: Final[Mapping[str, str]] = MappingProxyType({"X-Pplx-Integration": "litellm"})
class PerplexityChatConfig(OpenAIGPTConfig):
@property
def custom_llm_provider(self) -> str | None:
return "perplexity"
def get_attribution_headers(self) -> Mapping[str, str]:
return _PERPLEXITY_ATTRIBUTION_HEADERS
def _get_openai_compatible_provider_info(
self, api_base: str | None, api_key: str | None
) -> tuple[str | None, str | None]:

View file

@ -118,6 +118,7 @@ from litellm.llms.base_llm import BaseConfig, BaseImageGenerationConfig
from litellm.llms.base_llm.base_model_iterator import (
convert_model_response_to_streaming,
)
from litellm.llms.base_llm.chat.transformation import with_attribution_headers
from litellm.llms.bedrock.common_utils import (
BedrockModelInfo,
bedrock_route_for_request,
@ -2634,6 +2635,11 @@ def _complete_custom_openai(
)
headers = headers or litellm.headers
outbound_headers: Final = (
headers
if provider_config is None
else with_attribution_headers(provider_config.get_attribution_headers(), headers)
)
# Add GitHub Copilot headers (same as /responses endpoint does)
if custom_llm_provider == "github_copilot":
@ -2685,7 +2691,7 @@ def _complete_custom_openai(
acompletion=acompletion,
stream=stream,
api_key=api_key,
headers=headers,
headers=outbound_headers,
client=client,
provider_config=provider_config,
)
@ -2693,7 +2699,7 @@ def _complete_custom_openai(
response = openai_chat_completions.completion(
model=model,
messages=messages,
headers=headers,
headers=outbound_headers,
model_response=model_response,
print_verbose=print_verbose,
api_key=api_key,
@ -2716,7 +2722,7 @@ def _complete_custom_openai(
input=messages,
api_key=api_key,
original_response=str(e),
additional_args={"headers": headers},
additional_args={"headers": outbound_headers},
)
raise e
@ -2726,7 +2732,7 @@ def _complete_custom_openai(
input=messages,
api_key=api_key,
original_response=response,
additional_args={"headers": headers},
additional_args={"headers": outbound_headers},
)
return response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract

View file

@ -1,22 +1,21 @@
import inspect
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.llms.anthropic.pass_through.messages import handler as main
from litellm.rust_bridge.catalog import Route, RouteContext
from litellm.rust_bridge.dispatch import PublicDispatch, call_hook
from litellm.rust_bridge.dispatch import PublicDispatch
from litellm.rust_bridge.messages.entrypoints import (
NATIVE_AMESSAGES,
NATIVE_MESSAGES,
LiteLLMMessagesRequest,
)
from litellm.rust_bridge.public_call import (
NativeCall,
bind,
optional_bool,
optional_mapping,
native_call,
native_call_hook,
optional_sequence,
optional_str,
signature,
@ -52,7 +51,7 @@ _AMESSAGES: Final = signature(_PYTHON_AMESSAGES)
def _public_request(
legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> LiteLLMMessagesRequest | None:
) -> NativeCall | None:
fields: Final = bind(legacy, args, kwargs)
if fields is None:
return None
@ -61,30 +60,21 @@ def _public_request(
max_tokens: Final = fields.get("max_tokens")
if not isinstance(model, str) or messages is None or not isinstance(max_tokens, int):
return None
return LiteLLMMessagesRequest(
model=model,
messages=messages,
max_tokens=max_tokens,
stream=optional_bool(fields.get("stream")),
api_key=optional_str(fields.get("api_key")),
api_base=optional_str(fields.get("api_base")),
custom_llm_provider=optional_str(fields.get("custom_llm_provider")),
kwargs=optional_mapping(fields.get("kwargs")) or MappingProxyType({}),
)
return native_call(args, kwargs, fields)
def _resolved_provider(request: LiteLLMMessagesRequest) -> str | None:
def _resolved_provider(request: NativeCall) -> str | None:
try:
return get_llm_provider(request.model, request.custom_llm_provider)[1]
return get_llm_provider(str(request.bound["model"]), optional_str(request.bound.get("custom_llm_provider")))[1]
except BadRequestError:
return request.custom_llm_provider
return optional_str(request.bound.get("custom_llm_provider"))
def _context(request: LiteLLMMessagesRequest) -> RouteContext:
def _context(request: NativeCall) -> RouteContext:
return RouteContext(
Route.MESSAGES,
provider=_resolved_provider(request),
model=request.model,
model=str(request.bound["model"]),
)
@ -112,7 +102,7 @@ def anthropic_messages_handler(
kwargs,
python=python,
binding=NATIVE_MESSAGES,
native=call_hook,
native=native_call_hook,
)
@ -123,7 +113,7 @@ async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult:
kwargs,
python=python,
binding=NATIVE_AMESSAGES,
native=call_hook,
native=native_call_hook,
)

View file

@ -42380,6 +42380,40 @@
"supports_response_schema": true,
"supports_web_search": true
},
"openrouter/anthropic/claude-haiku-5.5": {
"input_cost_per_token": 1e-07,
"output_cost_per_token": 5e-07,
"cache_read_input_token_cost": 1e-08,
"cache_creation_input_token_cost": 1.25e-07,
"cache_creation_input_token_cost_above_1hr": 2e-07,
"input_cost_per_token_above_100k_tokens": 5e-07,
"output_cost_per_token_above_100k_tokens": 2.5e-06,
"cache_read_input_token_cost_above_100k_tokens": 5e-08,
"cache_creation_input_token_cost_above_100k_tokens": 6.25e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"source": "https://openrouter.ai/api/v1/models",
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_web_search": true,
"supports_adaptive_thinking": true,
"prompt_cache_min_tokens": 512,
"supports_sampling_params": false
},
"openrouter/anthropic/claude-haiku-4.5": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_1hr": 2e-06,
@ -80770,6 +80804,10 @@
"supports_vision": true
},
"claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens_batches": 2.5e-07,
"output_cost_per_token_above_100k_tokens_batches": 1.25e-06,
"cache_creation_input_token_cost_above_100k_tokens_batches": 3.125e-07,
"cache_read_input_token_cost_above_100k_tokens_batches": 2.5e-08,
"supports_anthropic_compaction": true,
"cache_creation_input_token_cost": 1.25e-07,
"cache_creation_input_token_cost_above_1hr": 2e-07,
@ -80809,8 +80847,8 @@
"us": 1.1
},
"supports_output_config": true,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://platform.claude.com/docs/en/models/haiku-5-5/overview",
"supports_web_search": true,
@ -80821,6 +80859,11 @@
"cache_read_input_token_cost_above_100k_tokens": 5e-08
},
"bedrock_mantle/anthropic.claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 5.5e-07,
"output_cost_per_token_above_100k_tokens": 2.75e-06,
"cache_creation_input_token_cost_above_100k_tokens": 6.875e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.1e-06,
"cache_read_input_token_cost_above_100k_tokens": 5.5e-08,
"bedrock_converse_supports_strict_tools": false,
"cache_creation_input_token_cost": 1.375e-07,
"cache_creation_input_token_cost_above_1hr": 2.2e-07,
@ -80856,12 +80899,17 @@
"supports_output_config": true,
"bedrock_output_config_effort_ceiling": "xhigh",
"supports_parallel_tool_use_config": true,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-haiku-5-5.html"
},
"bedrock_mantle/us-gov-west-1/anthropic.claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 6e-07,
"output_cost_per_token_above_100k_tokens": 3e-06,
"cache_creation_input_token_cost_above_100k_tokens": 7.5e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.2e-06,
"cache_read_input_token_cost_above_100k_tokens": 6e-08,
"bedrock_converse_supports_strict_tools": false,
"bedrock_output_config_effort_ceiling": "xhigh",
"cache_creation_input_token_cost": 1.5e-07,
@ -80875,8 +80923,8 @@
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 6e-07,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"supports_adaptive_thinking": true,
"supports_assistant_prefill": false,
@ -80898,6 +80946,11 @@
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-haiku-5-5.html"
},
"anthropic.claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 5e-07,
"output_cost_per_token_above_100k_tokens": 2.5e-06,
"cache_creation_input_token_cost_above_100k_tokens": 6.25e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06,
"cache_read_input_token_cost_above_100k_tokens": 5e-08,
"bedrock_converse_supports_strict_tools": false,
"cache_creation_input_token_cost": 1.25e-07,
"cache_creation_input_token_cost_above_1hr": 2e-07,
@ -80933,12 +80986,17 @@
"supports_output_config": true,
"bedrock_output_config_effort_ceiling": "xhigh",
"supports_parallel_tool_use_config": true,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
},
"apac.anthropic.claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 5.5e-07,
"output_cost_per_token_above_100k_tokens": 2.75e-06,
"cache_creation_input_token_cost_above_100k_tokens": 6.875e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.1e-06,
"cache_read_input_token_cost_above_100k_tokens": 5.5e-08,
"bedrock_converse_supports_strict_tools": false,
"litellm_provider": "bedrock_converse",
"supports_tool_search": true,
@ -80964,8 +81022,8 @@
"supports_output_config": true,
"bedrock_output_config_effort_ceiling": "xhigh",
"supports_parallel_tool_use_config": true,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"cache_creation_input_token_cost_above_1hr": 2.2e-07,
"cache_creation_input_token_cost": 1.375e-07,
@ -80975,6 +81033,11 @@
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
},
"au.anthropic.claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 5.5e-07,
"output_cost_per_token_above_100k_tokens": 2.75e-06,
"cache_creation_input_token_cost_above_100k_tokens": 6.875e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.1e-06,
"cache_read_input_token_cost_above_100k_tokens": 5.5e-08,
"bedrock_converse_supports_strict_tools": false,
"cache_creation_input_token_cost": 1.375e-07,
"cache_creation_input_token_cost_above_1hr": 2.2e-07,
@ -81010,12 +81073,17 @@
"supports_output_config": true,
"bedrock_output_config_effort_ceiling": "xhigh",
"supports_parallel_tool_use_config": true,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
},
"azure_ai/claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 5e-07,
"output_cost_per_token_above_100k_tokens": 2.5e-06,
"cache_creation_input_token_cost_above_100k_tokens": 6.25e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06,
"cache_read_input_token_cost_above_100k_tokens": 5e-08,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 1.25e-07,
"cache_creation_input_token_cost_above_1hr": 2e-07,
@ -81052,6 +81120,11 @@
"source": "https://platform.claude.com/docs/en/models/haiku-5-5/migration-guide"
},
"bedrock/us-gov-east-1/anthropic.claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 6e-07,
"output_cost_per_token_above_100k_tokens": 3e-06,
"cache_creation_input_token_cost_above_100k_tokens": 7.5e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.2e-06,
"cache_read_input_token_cost_above_100k_tokens": 6e-08,
"bedrock_converse_supports_strict_tools": false,
"bedrock_output_config_effort_ceiling": "xhigh",
"cache_creation_input_token_cost": 1.5e-07,
@ -81065,9 +81138,10 @@
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 6e-07,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json",
"supports_adaptive_thinking": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -81087,6 +81161,11 @@
"supports_xhigh_reasoning_effort": true
},
"bedrock/us-gov-west-1/anthropic.claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 6e-07,
"output_cost_per_token_above_100k_tokens": 3e-06,
"cache_creation_input_token_cost_above_100k_tokens": 7.5e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.2e-06,
"cache_read_input_token_cost_above_100k_tokens": 6e-08,
"bedrock_converse_supports_strict_tools": false,
"bedrock_output_config_effort_ceiling": "xhigh",
"cache_creation_input_token_cost": 1.5e-07,
@ -81100,9 +81179,10 @@
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 6e-07,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json",
"supports_adaptive_thinking": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -81122,6 +81202,11 @@
"supports_xhigh_reasoning_effort": true
},
"eu.anthropic.claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 5.5e-07,
"output_cost_per_token_above_100k_tokens": 2.75e-06,
"cache_creation_input_token_cost_above_100k_tokens": 6.875e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.1e-06,
"cache_read_input_token_cost_above_100k_tokens": 5.5e-08,
"bedrock_converse_supports_strict_tools": false,
"cache_creation_input_token_cost": 1.375e-07,
"cache_creation_input_token_cost_above_1hr": 2.2e-07,
@ -81157,12 +81242,17 @@
"supports_output_config": true,
"bedrock_output_config_effort_ceiling": "xhigh",
"supports_parallel_tool_use_config": true,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
},
"global.anthropic.claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 5e-07,
"output_cost_per_token_above_100k_tokens": 2.5e-06,
"cache_creation_input_token_cost_above_100k_tokens": 6.25e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06,
"cache_read_input_token_cost_above_100k_tokens": 5e-08,
"bedrock_converse_supports_strict_tools": false,
"cache_creation_input_token_cost": 1.25e-07,
"cache_creation_input_token_cost_above_1hr": 2e-07,
@ -81198,12 +81288,17 @@
"supports_output_config": true,
"bedrock_output_config_effort_ceiling": "xhigh",
"supports_parallel_tool_use_config": true,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
},
"jp.anthropic.claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 5.5e-07,
"output_cost_per_token_above_100k_tokens": 2.75e-06,
"cache_creation_input_token_cost_above_100k_tokens": 6.875e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.1e-06,
"cache_read_input_token_cost_above_100k_tokens": 5.5e-08,
"bedrock_converse_supports_strict_tools": false,
"cache_creation_input_token_cost": 1.375e-07,
"cache_creation_input_token_cost_above_1hr": 2.2e-07,
@ -81239,10 +81334,10 @@
"supports_output_config": true,
"bedrock_output_config_effort_ceiling": "xhigh",
"supports_parallel_tool_use_config": true,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
},
"perplexity/anthropic/claude-haiku-5-5": {
"litellm_provider": "perplexity",
@ -81256,6 +81351,11 @@
"source": "https://docs.perplexity.ai/docs/agent-api/models"
},
"us-gov.anthropic.claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 6e-07,
"output_cost_per_token_above_100k_tokens": 3e-06,
"cache_creation_input_token_cost_above_100k_tokens": 7.5e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.2e-06,
"cache_read_input_token_cost_above_100k_tokens": 6e-08,
"bedrock_converse_supports_strict_tools": false,
"bedrock_output_config_effort_ceiling": "xhigh",
"cache_creation_input_token_cost": 1.5e-07,
@ -81269,10 +81369,10 @@
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 6e-07,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://aws.amazon.com/bedrock/pricing/",
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json",
"supports_adaptive_thinking": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -81292,6 +81392,11 @@
"supports_xhigh_reasoning_effort": true
},
"us.anthropic.claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 5.5e-07,
"output_cost_per_token_above_100k_tokens": 2.75e-06,
"cache_creation_input_token_cost_above_100k_tokens": 6.875e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.1e-06,
"cache_read_input_token_cost_above_100k_tokens": 5.5e-08,
"bedrock_converse_supports_strict_tools": false,
"cache_creation_input_token_cost": 1.375e-07,
"cache_creation_input_token_cost_above_1hr": 2.2e-07,
@ -81327,12 +81432,17 @@
"supports_output_config": true,
"bedrock_output_config_effort_ceiling": "xhigh",
"supports_parallel_tool_use_config": true,
"supports_forced_tool_use": false,
"thinking_always_on": true,
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
},
"vertex_ai/claude-haiku-5-5": {
"input_cost_per_token_above_100k_tokens": 5e-07,
"output_cost_per_token_above_100k_tokens": 2.5e-06,
"cache_creation_input_token_cost_above_100k_tokens": 6.25e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06,
"cache_read_input_token_cost_above_100k_tokens": 5e-08,
"regional_endpoint_uplift_multiplier": 1.1,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 1.25e-07,
@ -81370,9 +81480,14 @@
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
},
"vertex_ai/claude-haiku-5-5@default": {
"input_cost_per_token_above_100k_tokens": 5e-07,
"output_cost_per_token_above_100k_tokens": 2.5e-06,
"cache_creation_input_token_cost_above_100k_tokens": 6.25e-07,
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06,
"cache_read_input_token_cost_above_100k_tokens": 5e-08,
"regional_endpoint_uplift_multiplier": 1.1,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 1.25e-07,
@ -81410,6 +81525,6 @@
"supports_forced_tool_use": true,
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
}
}

View file

@ -111,6 +111,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
source_url: str | None = None
timeout: float | None = None
max_concurrent_requests: int | None = None
rpm: int | None = None
approval_status: str | None = Field(
default="active",
description="Approval status: 'pending_review', 'active', 'rejected'",

View file

@ -1,4 +1,5 @@
from collections.abc import Coroutine, Mapping
from types import MappingProxyType
from typing import Final
import httpx
@ -6,8 +7,9 @@ import httpx
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge import runtime
from litellm.rust_bridge.catalog import Route, RouteContext
from litellm.rust_bridge.dispatch import PublicDispatch, call_hook
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR, LiteLLMOcrRequest
from litellm.rust_bridge.dispatch import PublicDispatch
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
from litellm.rust_bridge.public_call import NativeCall, native_call, native_call_hook, optional_str
__all__ = ("aocr", "ocr")
@ -21,30 +23,33 @@ def _bind_request(
custom_llm_provider: str | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs: object, # kwargs-ok: public OCR accepts provider-specific options
) -> LiteLLMOcrRequest:
return LiteLLMOcrRequest(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
) -> Mapping[str, object]:
return MappingProxyType(
{
"model": model,
"document": document,
"api_key": api_key,
"api_base": api_base,
"timeout": timeout,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"kwargs": kwargs,
}
)
def _public_request(name: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> LiteLLMOcrRequest:
def _public_request(name: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> NativeCall:
try:
return _bind_request(*args, **kwargs) # pyright: ignore[reportArgumentType] # Python binds the public arguments before native validation
fields: Final = _bind_request(*args, **kwargs) # pyright: ignore[reportArgumentType] # Python binds the public arguments before native validation
return native_call(args, kwargs, fields)
except TypeError as error:
raise TypeError(str(error).replace("_bind_request()", f"{name}()")) from None
def _context(request: LiteLLMOcrRequest) -> RouteContext:
prefix, separator, _ = request.model.partition("/")
provider: Final = request.custom_llm_provider or (prefix if separator else None)
return RouteContext(Route.OCR, provider=provider, model=request.model)
def _context(request: NativeCall) -> RouteContext:
prefix, separator, _ = str(request.bound["model"]).partition("/")
provider: Final = optional_str(request.bound.get("custom_llm_provider")) or (prefix if separator else None)
return RouteContext(Route.OCR, provider=provider, model=str(request.bound["model"]))
_DISPATCH: Final = PublicDispatch(
@ -70,7 +75,7 @@ def ocr(
kwargs,
python=runtime.NO_PYTHON,
binding=NATIVE_OCR,
native=call_hook,
native=native_call_hook,
)
@ -80,5 +85,5 @@ async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: pr
kwargs,
python=runtime.NO_PYTHON,
binding=NATIVE_AOCR,
native=call_hook,
native=native_call_hook,
)

View file

@ -952,24 +952,48 @@ async def aggregate_gateway_tools(
prefetched: Mapping[str, OAuthCredentialPayload],
*,
record_listing: bool = False,
enforce_rate_limits: bool = True,
) -> AggregateToolListing:
import time
from mcp.types import PaginatedRequestParams
from mcp.types import ListToolsResult, PaginatedRequestParams
from pydantic import TypeAdapter
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
SERVER_OUTCOMES_META_KEY,
AggregateToolListing,
ServerOutcome,
classify_list_exception,
)
from litellm.proxy._experimental.mcp_server.operations import _aggregate_server_key, global_mcp_server_manager
from litellm.proxy._experimental.mcp_server.operations import (
_aggregate_server_key,
_mcp_server_rate_limit_rejection,
global_mcp_server_manager,
)
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
async with global_mcp_server_manager.catalog.operation() as snapshot:
servers: Final = {server.server_id: server for server in allowed}
listing_updates: Final = ExitStack()
rejections: Final[list[ProxyRateLimitError]] = [] # mutable-ok: concurrent fetches share first-page errors
async def fetch(server_id: str, cursor: str | None) -> ListToolsResult:
if enforce_rate_limits:
error: Final = await _mcp_server_rate_limit_rejection(servers[server_id], context.user_api_key_auth)
if error is not None:
if cursor is not None:
raise error
rejections.append(error)
return ListToolsResult(
tools=[],
_meta={
SERVER_OUTCOMES_META_KEY: {
_aggregate_server_key(servers[server_id]): classify_list_exception(error).model_dump(
mode="json"
)
}
},
)
result, outcome = await get_filtered_server_tools(
servers[server_id],
context=context,
@ -1003,6 +1027,8 @@ async def aggregate_gateway_tools(
fetch=fetch,
now=int(time.time()),
)
if params.cursor is None and servers and len(rejections) == len(servers):
raise rejections[0]
listing_updates.close()
return AggregateToolListing(
tools=result.tools,
@ -1063,6 +1089,7 @@ async def list_gateway_catalog(
global_mcp_server_manager,
raise_denied_scoped_mcp_access,
)
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
context = replace(context, _caller=await MCPRequestHandler.refresh_catalog_authority(context.user_api_key_auth))
params: Final = request.params or PaginatedRequestParams()
@ -1080,6 +1107,7 @@ async def list_gateway_catalog(
requested_names=list(scope), user_api_key_auth=caller, client_ip=client_ip
)
servers: Final = {server.server_id: server for server in allowed}
rejections: Final[list[ProxyRateLimitError]] = [] # mutable-ok: concurrent fetches share first-page errors
async def fetch(server_id: str, cursor: str | None) -> CatalogListResult:
server: Final = servers[server_id]
@ -1090,7 +1118,26 @@ async def list_gateway_catalog(
SERVER_OUTCOMES_META_KEY,
classify_list_exception,
)
from litellm.proxy._experimental.mcp_server.operations import _aggregate_server_key
from litellm.proxy._experimental.mcp_server.operations import (
_aggregate_server_key,
_mcp_server_rate_limit_rejection,
)
error: Final = await _mcp_server_rate_limit_rejection(server, caller)
if error is not None:
if cursor is not None:
raise error
rejections.append(error)
return combine_optional_catalog(
request,
(),
None,
{
SERVER_OUTCOMES_META_KEY: {
_aggregate_server_key(server): classify_list_exception(error).model_dump(mode="json")
}
},
)
try:
page: Final = await fetch_optional_catalog_page(context, request, server, allowed, cursor)
@ -1126,6 +1173,8 @@ async def list_gateway_catalog(
fetch=fetch,
now=int(time.time()),
)
if params.cursor is None and servers and len(rejections) == len(servers):
raise rejections[0]
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
SERVER_OUTCOMES_META_KEY,
ServerOutcome,

View file

@ -1761,6 +1761,16 @@ def client_supplied_redirect_uris(value: object) -> list[str] | None:
return uris if len(uris) == len(value) else None
_CLIENT_APPLICATION_TYPE: Final = TypeAdapter(Literal["native", "web"] | None)
def client_supplied_application_type(value: object) -> Literal["native", "web"] | None:
try:
return _CLIENT_APPLICATION_TYPE.validate_python(value)
except ValidationError as exc:
raise HTTPException(status_code=400, detail="application_type must be native or web") from exc
async def _post_dcr_registration(
registration_url: str,
register_data: Mapping[str, object],
@ -1925,6 +1935,7 @@ async def register_client_with_server(
fallback_client_id: str | None = None,
persist_credentials: bool = False,
client_redirect_uris: list[str] | None = None,
client_application_type: Literal["native", "web"] | None = None,
):
_raise_if_not_oauth2(mcp_server)
request_base_url: Final = get_request_base_url(request)
@ -1980,6 +1991,11 @@ async def register_client_with_server(
)
register_data: Final = {
**(
{"application_type": client_application_type}
if bridge_relay and client_application_type is not None
else {}
),
"client_name": client_name,
"redirect_uris": client_redirect_uris if bridge_relay else [current_redirect_uri],
"grant_types": grant_types or (["authorization_code", "refresh_token"] if bridge_relay else []),
@ -3094,6 +3110,7 @@ async def register_client(request: Request, mcp_server_name: str | None = None):
return await register_aggregate_client(
request=request, request_body=data, token_exchange_available=token_exchange_available()
)
client_application_type: Final = client_supplied_application_type(data.get("application_type"))
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
async with global_mcp_server_manager.catalog.operation():
@ -3115,6 +3132,7 @@ async def register_client(request: Request, mcp_server_name: str | None = None):
token_endpoint_auth_method=data.get("token_endpoint_auth_method", ""),
fallback_client_id=resolved.server_name or resolved.name,
client_redirect_uris=client_redirect_uris,
client_application_type=client_application_type,
)
return dummy_return
@ -3130,4 +3148,5 @@ async def register_client(request: Request, mcp_server_name: str | None = None):
token_endpoint_auth_method=data.get("token_endpoint_auth_method", ""),
fallback_client_id=mcp_server_name,
client_redirect_uris=client_redirect_uris,
client_application_type=client_application_type,
)

View file

@ -24,11 +24,13 @@ from litellm.proxy._experimental.mcp_server.exceptions import (
MCPUpstreamAuthError,
)
from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.types.llms.base import LiteLLMBaseModel
ListFaultCategory: TypeAlias = Literal[
"auth_required",
"forbidden",
"rate_limited",
"timeout",
"unreachable",
"upstream_error",
@ -125,6 +127,8 @@ def classify_list_exception(exc: BaseException) -> ServerListFault:
if isinstance(exc, MCPUpstreamAuthError):
tag: Final = "forbidden" if exc.status_code == 403 else "auth_required"
return ServerListFault(tag=tag, status_code=exc.status_code)
if isinstance(exc, ProxyRateLimitError):
return ServerListFault(tag="rate_limited", status_code=429)
if isinstance(exc, TimeoutError):
return ServerListFault(tag="timeout")
if isinstance(exc, ConnectionError):
@ -152,7 +156,7 @@ def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]:
match outcome.tag:
case "ok":
return {"status": "ok", "tool_count": outcome.tool_count}
case "auth_required" | "forbidden" | "timeout" | "unreachable" | "upstream_error" | "internal":
case "auth_required" | "forbidden" | "rate_limited" | "timeout" | "unreachable" | "upstream_error" | "internal":
return {
"status": outcome.tag,
**({"http_status": outcome.status_code} if outcome.status_code is not None else {}),
@ -170,6 +174,8 @@ def list_fault_http_status(fault: ServerListFault) -> int:
return fault.status_code or 401
case "forbidden":
return 403
case "rate_limited":
return 429
case "timeout":
return 504
case "unreachable" | "upstream_error":

View file

@ -426,6 +426,7 @@ class MCPServerConfig(TypedDict, total=False):
client_assertion_signing_alg: str
timeout: float
max_concurrent_requests: int
rpm: ReadOnly[int | None]
class _ProtectedResourceMetadataPayload(TypedDict, total=False):
@ -2738,6 +2739,7 @@ class MCPServerManager:
allow_elicitation=bool(server_config.get("allow_elicitation", False)),
timeout=server_config.get("timeout", None),
max_concurrent_requests=server_config.get("max_concurrent_requests", None),
rpm=server_config.get("rpm", None),
token_validation=server_config.get("token_validation", None),
oauth_identity_binding=server_config.get("oauth_identity_binding", None),
)
@ -3327,6 +3329,7 @@ class MCPServerManager:
or "rfc8693",
timeout=getattr(mcp_server, "timeout", None),
max_concurrent_requests=getattr(mcp_server, "max_concurrent_requests", None),
rpm=getattr(mcp_server, "rpm", None),
)
_warn_legacy_delegate_auth_if_applicable(new_server, source="database")
if register_oauth_discovery:
@ -5973,6 +5976,7 @@ class MCPServerManager:
data=synthetic_llm_data,
call_type=CallTypes.call_mcp_tool.value,
)
await proxy_logging_obj.enforce_mcp_server_rate_limits(user_api_key_auth, server)
if modified_data:
# Convert response back to MCP format and apply modifications
modified_kwargs = proxy_logging_obj._convert_mcp_hook_response_to_kwargs(modified_data, pre_hook_kwargs)
@ -7193,6 +7197,7 @@ class MCPServerManager:
instructions=server.instructions,
timeout=server.timeout,
max_concurrent_requests=server.max_concurrent_requests,
rpm=server.rpm,
)
async def get_all_mcp_servers_with_health_and_teams(
@ -7316,6 +7321,7 @@ class MCPServerManager:
instructions=server.instructions,
timeout=server.timeout,
max_concurrent_requests=server.max_concurrent_requests,
rpm=server.rpm,
)
async def get_all_mcp_servers_unfiltered(self) -> list[LiteLLM_MCPServerTable]:

View file

@ -5,6 +5,8 @@ import traceback
import types
import uuid
from collections.abc import Mapping, Sequence
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import datetime
from functools import partial
from typing import Any, Final, NoReturn, TypeAlias, overload
@ -78,6 +80,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
AggregateToolListing,
ServerListOk,
ServerOutcome,
classify_list_exception,
outcome_wire_value,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
@ -128,6 +131,7 @@ from litellm.proxy._types import (
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
publish_auth_cache_invalidation,
)
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
get_chain_id_from_headers,
@ -224,6 +228,65 @@ class ListMCPToolsRestAPIResponseObject(MCPTool):
model_config = ConfigDict(arbitrary_types_allowed=True)
@dataclass(frozen=True, slots=True)
class _MCPServerRateLimitAdmission:
admitted_servers: tuple[MCPServer, ...]
rejected_servers: tuple[tuple[MCPServer, ProxyRateLimitError], ...]
_mcp_server_admission_memo: Final[ContextVar[dict[str, asyncio.Task[ProxyRateLimitError | None]] | None]] = ContextVar(
"mcp_server_admission_memo", default=None
)
async def _enforce_mcp_server_rate_limit(
user_api_key_auth: UserAPIKeyAuth | None,
server: MCPServer,
) -> None:
from litellm.proxy.proxy_server import proxy_logging_obj
if proxy_logging_obj is not None:
await proxy_logging_obj.enforce_mcp_server_rate_limits(user_api_key_auth, server)
async def _admit_mcp_servers(
servers: Sequence[MCPServer],
user_api_key_auth: UserAPIKeyAuth | None,
) -> _MCPServerRateLimitAdmission:
memo: Final = _mcp_server_admission_memo.get()
async def _server_rate_limit_error(server: MCPServer) -> ProxyRateLimitError | None:
try:
await _enforce_mcp_server_rate_limit(user_api_key_auth, server)
except ProxyRateLimitError as error:
return error
return None
async def _admit_server(server: MCPServer) -> tuple[MCPServer, ProxyRateLimitError | None]:
if memo is None:
return server, await _server_rate_limit_error(server)
admission_task: Final = memo.get(server.server_id)
if admission_task is not None:
return server, await admission_task
created_task: Final = asyncio.create_task(_server_rate_limit_error(server))
memo[server.server_id] = created_task
return server, await created_task
results: Final = await asyncio.gather(*(_admit_server(server) for server in servers))
return _MCPServerRateLimitAdmission(
admitted_servers=tuple(server for server, error in results if error is None),
rejected_servers=tuple((server, error) for server, error in results if error is not None),
)
async def _mcp_server_rate_limit_rejection(
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
) -> ProxyRateLimitError | None:
admission: Final = await _admit_mcp_servers((server,), user_api_key_auth)
return admission.rejected_servers[0][1] if admission.rejected_servers else None
async def _build_virtual_call_logging_obj(
name: str,
arguments: dict[str, object],
@ -961,6 +1024,7 @@ async def _get_tools_from_mcp_servers(
protocol_version: str | None = None,
*,
record_listing: bool = False,
enforce_rate_limits: bool = True,
) -> AggregateToolListing:
"""
Helper method to fetch tools from MCP servers based on server filtering criteria.
@ -1093,20 +1157,37 @@ async def _get_tools_from_mcp_servers(
return page.tools, outcome
if params is None:
server_admission: Final = (
await _admit_mcp_servers(allowed_mcp_servers, user_api_key_auth)
if enforce_rate_limits
else _MCPServerRateLimitAdmission(tuple(allowed_mcp_servers), ())
)
if not server_admission.admitted_servers and server_admission.rejected_servers:
raise server_admission.rejected_servers[0][1]
admitted_servers: Final = server_admission.admitted_servers
results: Final = await asyncio.gather(
*(_fetch_and_filter_server_tools(server) for server in allowed_mcp_servers)
*(_fetch_and_filter_server_tools(server) for server in admitted_servers)
)
aggregated = AggregateToolListing(
tools=[tool for tools, _ in results for tool in tools],
outcomes={
_aggregate_server_key(server): outcome for server, (_, outcome) in zip(allowed_mcp_servers, results)
_aggregate_server_key(server): outcome for server, (_, outcome) in zip(admitted_servers, results)
}
| {
_aggregate_server_key(server): classify_list_exception(error)
for server, error in server_admission.rejected_servers
},
)
else:
from litellm.proxy._experimental.mcp_server.catalog import aggregate_gateway_tools
aggregated = await aggregate_gateway_tools(
context, params, allowed_mcp_servers, _prefetched_oauth_creds, record_listing=record_listing
context,
params,
allowed_mcp_servers,
_prefetched_oauth_creds,
record_listing=record_listing,
enforce_rate_limits=enforce_rate_limits,
)
all_tools: Final = aggregated.tools
server_outcomes: Final = aggregated.outcomes
@ -1751,6 +1832,7 @@ async def _list_tools_before_first_call(
raw_headers=raw_headers,
client_ip=client_ip,
record_listing=False,
enforce_rate_limits=False,
)
except Exception as e: # noqa: BLE001 # best effort: resolution below answers as it did before
verbose_logger.debug("MCP tools/call: listing %s before its first call failed: %s", server.name, e)
@ -2464,6 +2546,7 @@ async def mcp_get_prompt(
user_api_key_auth=user_api_key_auth,
)
await _enforce_mcp_server_rate_limit(user_api_key_auth, server)
return await global_mcp_server_manager.get_prompt_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
@ -2517,6 +2600,7 @@ async def mcp_read_resource(
user_api_key_auth=user_api_key_auth,
)
await _enforce_mcp_server_rate_limit(user_api_key_auth, server)
return await global_mcp_server_manager.read_resource_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
@ -3125,6 +3209,7 @@ class GatewayOperations:
if context.mcp_proxy_mode
else (ListPromptsRequest(), ListResourcesRequest(), ListResourceTemplatesRequest())
)
memo_token: Final = _mcp_server_admission_memo.set({})
tasks: Final = (
asyncio.create_task(
_execute_handle_list_tools(
@ -3139,9 +3224,12 @@ class GatewayOperations:
try:
results: Final = await asyncio.gather(*tasks)
finally:
for task in tasks:
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
try:
for task in tasks:
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
finally:
_mcp_server_admission_memo.reset(memo_token)
return build_discovery(
configured=configured_versions(),
revision=context.protocol_version or "2025-11-25",

View file

@ -55,6 +55,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.responses.mcp.request_context import MCPRequestContext
if TYPE_CHECKING:
@ -737,6 +738,8 @@ if MCP_AVAILABLE:
"""
from litellm.proxy.proxy_server import proxy_logging_obj
if apply_tool_filters and proxy_logging_obj is not None:
await proxy_logging_obj.enforce_mcp_server_rate_limits(user_api_key_auth, server)
listed_generation: Final = global_mcp_server_manager.listed_tools_generation(server.server_id)
tools: Final = await _list_server_tools(
server,
@ -900,6 +903,8 @@ if MCP_AVAILABLE:
# matching status code and WWW-Authenticate challenge; that is what
# lets standards-compliant MCP clients run the upstream OAuth flow.
raise
except ProxyRateLimitError:
raise
except MCPServerListError as e:
fault: Final = classify_list_exception(e)
verbose_logger.info("Listing tools from %s failed with a %s fault", server.name, fault.tag)

View file

@ -11572,6 +11572,23 @@
"description": "Google Cloud location/region (e.g., us-central1)",
"title": "Location"
},
"logging_only_scope": {
"anyOf": [
{
"enum": [
"input",
"output",
"both"
],
"type": "string"
},
{
"type": "null"
}
],
"description": "which direction a logging_only scan observes: 'input' (request), 'output' (response), or 'both' (default). Only applies to mode logging_only; pre_call/post_call on the same guardrail keep blocking.",
"title": "Logging Only Scope"
},
"mask_request_content": {
"anyOf": [
{
@ -12727,6 +12744,77 @@
"title": "GuardrailSubmissionSummary",
"type": "object"
},
"GuardrailUIAddGuardrailSettings": {
"properties": {
"content_filter_settings": {
"anyOf": [
{
"additionalProperties": true,
"type": "object"
},
{
"type": "null"
}
],
"title": "Content Filter Settings"
},
"pii_entity_categories": {
"items": {
"$ref": "#/components/schemas/PiiEntityCategoryMap"
},
"title": "Pii Entity Categories",
"type": "array"
},
"providers_without_directional_logging_only_scope": {
"items": {
"type": "string"
},
"title": "Providers Without Directional Logging Only Scope",
"type": "array"
},
"supported_actions": {
"items": {
"type": "string"
},
"title": "Supported Actions",
"type": "array"
},
"supported_entities": {
"items": {
"type": "string"
},
"title": "Supported Entities",
"type": "array"
},
"supported_modes": {
"items": {
"type": "string"
},
"title": "Supported Modes",
"type": "array"
},
"supported_modes_by_provider": {
"additionalProperties": {
"items": {
"type": "string"
},
"type": "array"
},
"title": "Supported Modes By Provider",
"type": "object"
}
},
"required": [
"supported_entities",
"supported_actions",
"supported_modes",
"supported_modes_by_provider",
"providers_without_directional_logging_only_scope",
"pii_entity_categories"
],
"title": "GuardrailUIAddGuardrailSettings",
"type": "object"
},
"HTTPValidationError": {
"properties": {
"detail": {
@ -12869,6 +12957,19 @@
],
"title": "Akto Base Url"
},
"akto_metadata": {
"anyOf": [
{
"additionalProperties": true,
"type": "object"
},
{
"type": "null"
}
],
"description": "JSON object sent to Akto. 'policy_name': comma-separated Akto policies to enforce (empty enforces all). Example: {\"policy_name\": \"PII Strict, Secrets\"}.",
"title": "Akto Metadata"
},
"akto_vxlan_id": {
"anyOf": [
{
@ -13407,6 +13508,22 @@
"description": "Enable content moderation to check for harmful content (harassment, hate speech, etc.).",
"title": "Content Moderation Check"
},
"context_source": {
"anyOf": [
{
"enum": [
"ENDPOINT",
"AGENTIC"
],
"type": "string"
},
{
"type": "null"
}
],
"description": "Akto context the traffic belongs to: 'ENDPOINT' (Atlas) or 'AGENTIC' (Argus). Default: AGENTIC.",
"title": "Context Source"
},
"contextual_grounding_from_messages": {
"default": false,
"description": "ApplyGuardrail: when True, post-call scans of a request with no grounding_source / query content parts send the system and developer messages as the grounding source and the latest user message as the query, so the guardrail's contextual grounding policy can score the response. Bedrock bills contextual grounding units for these scans and rejects queries, sources and responses over its contextual grounding length limits, so leave this off for guardrails without a contextual grounding policy. Default False: plain messages are never sent as grounding context.",
@ -13618,6 +13735,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_guardrail_timeout": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
],
"description": "HTTP timeout in seconds for checking attached files. Default: 10.",
"title": "File Guardrail Timeout"
},
"gateway_name": {
"anyOf": [
{
@ -13822,6 +13951,23 @@
"description": "Google Cloud location/region (e.g., us-central1)",
"title": "Location"
},
"logging_only_scope": {
"anyOf": [
{
"enum": [
"input",
"output",
"both"
],
"type": "string"
},
{
"type": "null"
}
],
"description": "which direction a logging_only scan observes: 'input' (request), 'output' (response), or 'both' (default). Only applies to mode logging_only; pre_call/post_call on the same guardrail keep blocking.",
"title": "Logging Only Scope"
},
"mask": {
"anyOf": [
{
@ -14882,6 +15028,27 @@
"title": "PiiAction",
"type": "string"
},
"PiiEntityCategoryMap": {
"properties": {
"category": {
"title": "Category",
"type": "string"
},
"entities": {
"items": {
"type": "string"
},
"title": "Entities",
"type": "array"
}
},
"required": [
"category",
"entities"
],
"title": "PiiEntityCategoryMap",
"type": "object"
},
"PiiEntityType": {
"enum": [
"CREDIT_CARD",
@ -16323,7 +16490,9 @@
"200": {
"content": {
"application/json": {
"schema": {}
"schema": {
"$ref": "#/components/schemas/GuardrailUIAddGuardrailSettings"
}
}
},
"description": "Successful Response"
@ -31793,6 +31962,18 @@
],
"title": "Registration Url"
},
"rpm": {
"anyOf": [
{
"minimum": 0.0,
"type": "integer"
},
{
"type": "null"
}
],
"title": "Rpm"
},
"server_id": {
"anyOf": [
{
@ -33645,6 +33826,17 @@
],
"title": "Reviewed At"
},
"rpm": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
],
"title": "Rpm"
},
"server_id": {
"title": "Server Id",
"type": "string"
@ -34531,6 +34723,18 @@
],
"title": "Registration Url"
},
"rpm": {
"anyOf": [
{
"minimum": 0.0,
"type": "integer"
},
{
"type": "null"
}
],
"title": "Rpm"
},
"server_id": {
"anyOf": [
{
@ -36929,6 +37133,17 @@
],
"title": "Reviewed At"
},
"rpm": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
],
"title": "Rpm"
},
"server_id": {
"title": "Server Id",
"type": "string"
@ -38659,6 +38874,18 @@
],
"title": "Registration Url"
},
"rpm": {
"anyOf": [
{
"minimum": 0.0,
"type": "integer"
},
{
"type": "null"
}
],
"title": "Rpm"
},
"server_id": {
"anyOf": [
{
@ -39202,6 +39429,18 @@
],
"title": "Registration Url"
},
"rpm": {
"anyOf": [
{
"minimum": 0.0,
"type": "integer"
},
{
"type": "null"
}
],
"title": "Rpm"
},
"server_id": {
"title": "Server Id",
"type": "string"
@ -42048,6 +42287,18 @@
],
"title": "Registration Url"
},
"rpm": {
"anyOf": [
{
"minimum": 0.0,
"type": "integer"
},
{
"type": "null"
}
],
"title": "Rpm"
},
"server_id": {
"anyOf": [
{

View file

@ -1365,6 +1365,7 @@ class KeyRequestBase(GenerateRequestBase):
enforced_params: list[str] | None = None
allowed_routes: list | None = []
allowed_passthrough_routes: list | None = None
denied_passthrough_routes: list[str] | None = None
allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None
rpm_limit_type: Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] | None = (
None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm
@ -1674,6 +1675,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
source_url: str | None = None
timeout: float | None = None
max_concurrent_requests: int | None = None
rpm: int | None = Field(default=None, ge=0)
# BYOM submission fields — set by the endpoint, not by the caller.
# Any caller-provided values are silently overridden before persistence.
approval_status: str | None = Field(
@ -1781,6 +1783,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
source_url: str | None = None
timeout: float | None = None
max_concurrent_requests: int | None = None
rpm: int | None = Field(default=None, ge=0)
@model_validator(mode="after")
def validate_protocol_transport(self) -> "UpdateMCPServerRequest":
@ -2212,6 +2215,7 @@ class NewTeamRequest(TeamBase):
prompts: list[str] | None = None
object_permission: LiteLLM_ObjectPermissionBase | None = None
allowed_passthrough_routes: list | None = None
denied_passthrough_routes: list[str] | None = None
disable_global_guardrails: bool | None = None
secret_manager_settings: dict | None = None
model_rpm_limit: dict[str, int] | None = None
@ -2293,6 +2297,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
team_member_tpm_limit: int | None = None
team_member_key_duration: str | None = None
allowed_passthrough_routes: list | None = None
denied_passthrough_routes: list[str] | None = None
secret_manager_settings: dict | None = None
prompts: list[str] | None = None
model_rpm_limit: dict[str, int] | None = None
@ -3645,6 +3650,8 @@ class UserInfoV2Response(LiteLLMPydanticObjectBase):
user_role: str | None = None
spend: float = 0.0
max_budget: float | None = None
tpm_limit: int | None = None
rpm_limit: int | None = None
models: list[str] = []
budget_duration: str | None = None
budget_reset_at: datetime | None = None
@ -5102,6 +5109,7 @@ LiteLLM_ManagementEndpoint_MetadataFields_Premium: Final = [
"logging",
"secret_manager_settings",
"allowed_passthrough_routes",
"denied_passthrough_routes",
]
# Metadata keys that are immutable once set: preserved when an update omits them,

View file

@ -1648,6 +1648,10 @@ class JWTAuthManager:
):
return True
team_metadata: Final = (team_object.metadata or {}) if team_object else {}
if RouteChecks.matching_denied_passthrough_route(route=route, metadata_sources=(team_metadata,)) is not None:
return False
if RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=team_allowed_routes):
return True
@ -1655,11 +1659,16 @@ class JWTAuthManager:
# so beyond the JWT config grant above, only the selected team's metadata grants access.
return RouteChecks.check_passthrough_route_access(
route=route,
user_api_key_dict=UserAPIKeyAuth(team_metadata=(team_object.metadata or {}) if team_object else {}),
user_api_key_dict=UserAPIKeyAuth(team_metadata=team_metadata),
)
@staticmethod
def _raise_team_passthrough_route_denial(route: str) -> None:
def _raise_team_passthrough_route_denial(route: str, team_object: LiteLLM_TeamTable | None) -> None:
denied_route: Final = RouteChecks.matching_denied_passthrough_route(
route=route, metadata_sources=((team_object.metadata if team_object else None),)
)
if denied_route is not None:
raise RouteChecks.passthrough_route_denied_exception(route=route, denied_route=denied_route)
raise HTTPException(
status_code=403,
detail=(
@ -1683,7 +1692,7 @@ class JWTAuthManager:
"""Find first team with access to the requested model"""
from litellm.proxy.proxy_server import llm_router
denied_auth_enforced_pass_through_route = False
denied_pass_through_team: LiteLLM_TeamTable | None = None
if not team_ids:
if (
@ -1733,7 +1742,7 @@ class JWTAuthManager:
team_allowed_routes=jwt_handler.litellm_jwtauth.team_allowed_routes,
):
is_allowed = False
denied_auth_enforced_pass_through_route = True
denied_pass_through_team = team_object
verbose_proxy_logger.debug(
"JWT team route check: team_id=%s, route=%s, is_allowed=%s", team_id, route, is_allowed
)
@ -1742,8 +1751,8 @@ class JWTAuthManager:
except Exception:
continue
if denied_auth_enforced_pass_through_route:
JWTAuthManager._raise_team_passthrough_route_denial(route=route)
if denied_pass_through_team is not None:
JWTAuthManager._raise_team_passthrough_route_denial(route=route, team_object=denied_pass_through_team)
if requested_model and (any_claim_team_resolved or not jwt_handler.litellm_jwtauth.team_claim_fallback):
# Claim resolved but no model access, or fallback disabled — deny.
@ -2788,7 +2797,7 @@ class JWTAuthManager:
request_method=request_method,
team_allowed_routes=handler.litellm_jwtauth.team_allowed_routes,
):
JWTAuthManager._raise_team_passthrough_route_denial(route=route)
JWTAuthManager._raise_team_passthrough_route_denial(route=route, team_object=selected_team_object)
# Extract alias fields for resolution (if configured)
org_alias: Final = handler.get_org_alias(token=jwt_valid_token, default_value=None)
@ -2858,7 +2867,7 @@ class JWTAuthManager:
request_method=request_method,
team_allowed_routes=handler.litellm_jwtauth.team_allowed_routes,
):
JWTAuthManager._raise_team_passthrough_route_denial(route=route)
JWTAuthManager._raise_team_passthrough_route_denial(route=route, team_object=team_object)
elif selected_team_id is None:
(
team_id,

View file

@ -1,6 +1,7 @@
import itertools
import re
from collections.abc import Collection
from typing import Final
from collections.abc import Collection, Iterable, Mapping
from typing import Final, cast
from fastapi import HTTPException, Request, status
@ -693,6 +694,70 @@ class RouteChecks:
),
)
@staticmethod
def _route_matches_denied_route(route: str, denied_route: str) -> bool:
"""A `/` entry denies every route, since every route sits under the root."""
normalized_denied_route: Final = denied_route.rstrip("/") or "/"
return (
normalized_denied_route == "/"
or RouteChecks._route_matches_allowed_route(route=route, allowed_route=normalized_denied_route)
or RouteChecks.route_matches_wildcard_pattern(route=route, pattern=denied_route)
)
@staticmethod
def matching_denied_passthrough_route(
route: str, metadata_sources: Iterable[Mapping[str, object] | None]
) -> str | None:
"""
First ``denied_passthrough_routes`` entry across ``metadata_sources`` that matches ``route``.
Unlike the allowlist (key list, else team list), every source's deny list applies.
"""
denied_routes: Final = tuple(
itertools.chain.from_iterable(
cast( # cast-ok: management endpoints validate this metadata key as a list of route strings on write
"list[str]", (metadata or {}).get("denied_passthrough_routes") or []
)
for metadata in metadata_sources
)
)
if not denied_routes:
return None
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
InitPassThroughEndpointHelpers,
)
forwarded_routes: Final = InitPassThroughEndpointHelpers.forwarded_routes(route)
return next(
(
denied_route
for denied_route in denied_routes
if any(
RouteChecks._route_matches_denied_route(route=candidate, denied_route=denied_route)
for candidate in forwarded_routes
)
),
None,
)
@staticmethod
def passthrough_route_denied_exception(route: str, denied_route: str) -> HTTPException:
return HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=(
f"Key/team denied access to passthrough route {route}. "
f"Matched `{denied_route}` in `denied_passthrough_routes`."
),
)
@staticmethod
def _raise_if_passthrough_route_denied(route: str, valid_token: UserAPIKeyAuth) -> None:
denied_route: Final = RouteChecks.matching_denied_passthrough_route(
route=route,
metadata_sources=(valid_token.metadata, valid_token.team_metadata),
)
if denied_route is not None:
raise RouteChecks.passthrough_route_denied_exception(route=route, denied_route=denied_route)
@staticmethod
def jwt_team_routes_grant_pass_through(route: str, team_allowed_routes: Collection[str]) -> bool:
"""
@ -724,8 +789,10 @@ class RouteChecks:
) -> None:
"""
Require an explicit grant for auth=true pass-through: ``allowed_passthrough_routes`` on the
key or team, or an explicit JWT ``team_allowed_routes`` entry.
key or team, or an explicit JWT ``team_allowed_routes`` entry. A key or team
``denied_passthrough_routes`` match blocks the route even when one of those grants it.
"""
RouteChecks._raise_if_passthrough_route_denied(route=route, valid_token=valid_token)
if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token):
return
if RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=jwt_team_allowed_routes):

View file

@ -104,6 +104,8 @@ days AS (
router_name,
router_type,
SUM(turns)::int AS turns,
CASE WHEN SUM(token_recorded_turns) = SUM(turns)
THEN SUM(total_tokens)::bigint END AS day_total_tokens,
SUM(spend)::float8 AS spend,
SUM(saved_spend)::float8 AS saved_spend,
SUM(savings_estimated_turns)::int AS savings_estimated_turns,
@ -139,6 +141,7 @@ SELECT
COALESCE(sessions.total_tokens, 0) AS total_tokens,
COALESCE(sessions.session_seconds, 0) AS session_seconds,
COALESCE(days.turns, 0) AS turns,
CASE WHEN days.turns IS NULL THEN 0 ELSE days.day_total_tokens END AS day_total_tokens,
COALESCE(days.spend, 0) AS spend,
COALESCE(days.saved_spend, 0) AS saved_spend,
COALESCE(days.savings_estimated_turns, 0) AS savings_estimated_turns,
@ -175,6 +178,7 @@ class AutoRouterTurnTransaction:
savings_estimated_actual_spend: float = 0.0
savings_estimated_saved_spend: float = 0.0
user_id: str = ""
token_counts_recorded: bool = False
class TurnCacheFacts(NamedTuple):
@ -294,6 +298,11 @@ def build_autorouter_turn_transaction(
)
usage_object_raw: Final = metadata.get("usage_object")
token_counts: Final = (
(usage_object_raw.get("prompt_tokens"), usage_object_raw.get("completion_tokens"))
if isinstance(usage_object_raw, Mapping)
else ()
)
cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None)
tier_raw: Final = routing_decision.get("tier")
baseline_raw: Final = routing_decision.get("savings_baseline_model")
@ -311,6 +320,8 @@ def build_autorouter_turn_transaction(
model=model,
turn_at=turn_at,
total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0),
token_counts_recorded=len(token_counts) == 2
and all(isinstance(value, int) and not isinstance(value, bool) and value >= 0 for value in token_counts),
spend=actual_spend,
saved_spend=saved_spend,
classifier_cost=classifier_cost or 0.0,
@ -435,17 +446,21 @@ ON CONFLICT ({user_column}api_key, session_id, router_name) DO UPDATE SET
_DAY_UPSERT_SQL: Final = f"""
day_rollup AS (
INSERT INTO "LiteLLM_AutoRouterDailySpend" AS d (
date, api_key, user_id, router_name, router_type, turns, spend, saved_spend, savings_estimated_turns,
date, api_key, user_id, router_name, router_type, turns, total_tokens, token_recorded_turns,
spend, saved_spend, savings_estimated_turns,
savings_estimated_actual_spend, savings_estimated_saved_spend, classifier_cost, classifier_cost_recorded_turns
)
VALUES (
({_TURN_AT}::timestamp)::date::text, {_p("api_key")}::text, {_p("user_id")}::text, {_p("router_name")},
{_p("router_type")}, 1, {_p("spend")}::float8, {_p("saved_spend")}::float8, {_p("savings_estimated_turns")}::int,
{_p("router_type")}, 1, {_p("total_tokens")}::bigint, {_p("token_counts_recorded")}::int,
{_p("spend")}::float8, {_p("saved_spend")}::float8, {_p("savings_estimated_turns")}::int,
{_p("savings_estimated_actual_spend")}::float8, {_p("savings_estimated_saved_spend")}::float8,
{_p("classifier_cost")}::float8, 1
)
ON CONFLICT (date, api_key, user_id, router_name, router_type) DO UPDATE SET
turns = d.turns + 1,
total_tokens = d.total_tokens + EXCLUDED.total_tokens,
token_recorded_turns = d.token_recorded_turns + EXCLUDED.token_recorded_turns,
spend = d.spend + EXCLUDED.spend,
saved_spend = d.saved_spend + EXCLUDED.saved_spend,
savings_estimated_turns = d.savings_estimated_turns + EXCLUDED.savings_estimated_turns,

View file

@ -36,9 +36,11 @@ from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import (
)
from litellm.proxy.guardrails.guardrail_registry import (
GuardrailRegistry,
configured_event_hooks,
contains_encrypted_marker,
decrypt_guardrail_litellm_params,
encrypt_guardrail_litellm_params,
parse_tolerant_litellm_params,
)
from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
@ -283,7 +285,13 @@ async def list_guardrails_v2(
number_of_asterisks=4,
)
masked_litellm_params = (
BaseLitellmParams.model_validate(masked_litellm_params_dict) if masked_litellm_params_dict else None
parse_tolerant_litellm_params(
masked_litellm_params_dict,
guardrail.get("guardrail_name") or "Unknown",
params_model=BaseLitellmParams,
)
if masked_litellm_params_dict
else None
)
guardrail_configs.append(
GuardrailInfoResponse(
@ -324,7 +332,11 @@ async def list_guardrails_v2(
number_of_asterisks=4,
)
masked_in_memory_litellm_params_typed = (
BaseLitellmParams.model_validate(masked_in_memory_litellm_params)
parse_tolerant_litellm_params(
masked_in_memory_litellm_params,
guardrail.get("guardrail_name") or "Unknown",
params_model=BaseLitellmParams,
)
if masked_in_memory_litellm_params
else None
)
@ -425,7 +437,11 @@ async def create_guardrail(
guardrail_id: Final = result.get("guardrail_id", "Unknown")
try:
IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail(guardrail=cast(Guardrail, result), source="db")
IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail(
guardrail=cast(Guardrail, result),
source="db",
reject_invalid_logging_only_scope=True,
)
verbose_proxy_logger.info(
"Immediate sync: Successfully initialized guardrail '%s' (ID: %s)", guardrail_name, guardrail_id
)
@ -550,7 +566,10 @@ async def update_guardrail(
guardrail_name: Final = result.get("guardrail_name", "Unknown")
try:
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=cast(Guardrail, result))
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
guardrail=cast(Guardrail, result),
reject_invalid_logging_only_scope=True,
)
verbose_proxy_logger.info(
"Immediate sync: Successfully updated guardrail '%s' (ID: %s)", guardrail_name, guardrail_id
)
@ -1240,19 +1259,35 @@ async def patch_guardrail(
# Update litellm_params if default_on is provided or pii_entities_config is provided
existing_litellm_params: Final = _as_str_object_mapping(dict(existing_guardrail.get("litellm_params", {})))
litellm_params = LitellmParams(**existing_litellm_params)
if request.litellm_params is not None:
requested_litellm_params: Final = request.litellm_params.model_dump(exclude_unset=True)
litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True)
litellm_params_dict.update(requested_litellm_params)
merged_litellm_params: Final = _as_str_object_mapping(litellm_params_dict)
try:
litellm_params = LitellmParams(**merged_litellm_params)
except ValidationError as validation_error:
raise HTTPException(
status_code=422,
detail=f"Invalid guardrail configuration, update rejected: {validation_error}",
) from validation_error
current_litellm_params: Final = parse_tolerant_litellm_params(
existing_litellm_params,
existing_guardrail.get("guardrail_name") or "Unknown",
)
requested_litellm_params: Final[Mapping[str, object]] = (
MappingProxyType(request.litellm_params.model_dump(exclude_unset=True))
if request.litellm_params is not None
else MappingProxyType({})
)
merged_litellm_params: Final = _as_str_object_mapping(
MappingProxyType({**current_litellm_params.model_dump(exclude_unset=True), **requested_litellm_params})
)
try:
parsed_litellm_params: Final = LitellmParams(**merged_litellm_params)
except ValidationError as validation_error:
raise HTTPException(
status_code=422,
detail=f"Invalid guardrail configuration, update rejected: {validation_error}",
) from validation_error
clear_stored_scope: Final = (
"logging_only_scope" not in requested_litellm_params
and parsed_litellm_params.logging_only_scope is not None
and GuardrailEventHooks.logging_only.value not in configured_event_hooks(parsed_litellm_params.mode)
)
litellm_params: Final = (
LitellmParams(**MappingProxyType({**merged_litellm_params, "logging_only_scope": None}))
if clear_stored_scope
else parsed_litellm_params
)
# Update guardrail_info if provided
guardrail_info: Final = (
@ -1281,6 +1316,7 @@ async def patch_guardrail(
try:
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
guardrail=guardrail,
reject_invalid_logging_only_scope="logging_only_scope" in requested_litellm_params,
)
verbose_proxy_logger.info(
"Immediate sync: Successfully updated guardrail '%s' (ID: %s)", guardrail_name, guardrail_id
@ -1294,15 +1330,7 @@ async def patch_guardrail(
# the caller instead of a misleading 200.
await GUARDRAIL_REGISTRY.update_guardrail_in_db(
guardrail_id=guardrail_id,
guardrail=Guardrail(
guardrail_id=guardrail_id,
guardrail_name=existing_guardrail.get("guardrail_name") or "",
litellm_params=LitellmParams(**existing_litellm_params),
guardrail_info=existing_guardrail.get(
"guardrail_info",
{},
),
),
guardrail=existing_guardrail,
prisma_client=prisma_client,
)
raise HTTPException(
@ -1404,7 +1432,13 @@ async def get_guardrail_info(guardrail_id: str):
number_of_asterisks=4,
)
masked_litellm_params = (
BaseLitellmParams.model_validate(masked_litellm_params_dict) if masked_litellm_params_dict else None
parse_tolerant_litellm_params(
masked_litellm_params_dict,
result.get("guardrail_name") or "Unknown",
params_model=BaseLitellmParams,
)
if masked_litellm_params_dict
else None
)
return GuardrailInfoResponse(
@ -1427,7 +1461,7 @@ async def get_guardrail_info(guardrail_id: str):
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def get_guardrail_ui_settings():
async def get_guardrail_ui_settings() -> GuardrailUIAddGuardrailSettings:
"""
Get the UI settings for the guardrails
@ -1461,12 +1495,18 @@ async def get_guardrail_ui_settings():
# above; it only runs on pre_call.
{SupportedGuardrailIntegrations.HIDE_SECRETS.value: [GuardrailEventHooks.pre_call.value]}
)
providers_without_directional_logging_only_scope: Final = tuple(
provider
for provider, guardrail_class in guardrail_class_registry.items()
if not guardrail_class.supports_logging_only_scope()
)
return GuardrailUIAddGuardrailSettings(
supported_entities=[entity.value for entity in PiiEntityType],
supported_actions=[action.value for action in PiiAction],
supported_modes=[mode.value for mode in GuardrailEventHooks],
supported_modes_by_provider=supported_modes_by_provider,
providers_without_directional_logging_only_scope=providers_without_directional_logging_only_scope,
pii_entity_categories=category_maps,
content_filter_settings={
"prebuilt_patterns": get_pattern_metadata(),

View file

@ -2,7 +2,7 @@ from typing import TYPE_CHECKING, Final
from litellm.types.guardrails import SupportedGuardrailIntegrations
from .akto import AktoGuardrail
from .akto import AktoGuardrail, streaming_sampling_rate_from
if TYPE_CHECKING:
from litellm.types.guardrails import Guardrail, LitellmParams
@ -12,12 +12,16 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
import litellm
_akto_callback: Final = AktoGuardrail(
akto_base_url=getattr(litellm_params, "akto_base_url", None),
akto_api_key=getattr(litellm_params, "akto_api_key", None),
akto_account_id=getattr(litellm_params, "akto_account_id", None),
akto_vxlan_id=getattr(litellm_params, "akto_vxlan_id", None),
unreachable_fallback=getattr(litellm_params, "unreachable_fallback", "fail_closed"),
guardrail_timeout=getattr(litellm_params, "guardrail_timeout", None),
akto_base_url=litellm_params.akto_base_url,
akto_api_key=litellm_params.akto_api_key,
akto_account_id=litellm_params.akto_account_id,
akto_vxlan_id=litellm_params.akto_vxlan_id,
context_source=litellm_params.context_source,
akto_metadata=litellm_params.akto_metadata,
streaming_sampling_rate=streaming_sampling_rate_from(litellm_params),
guardrail_timeout=litellm_params.guardrail_timeout,
file_guardrail_timeout=litellm_params.file_guardrail_timeout,
unreachable_fallback=litellm_params.unreachable_fallback,
guardrail_name=guardrail.get("guardrail_name", ""),
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,401 @@
"""Attachment blocks sent to Akto's file guardrail, including those inside ``tool_result`` blocks:
OpenAI chat ``image_url``, ``input_audio``, ``file``, ``video_url``
Anthropic ``image``, ``document`` (except text documents, which stay in the text check)
Responses API ``input_image``, ``input_file``
A block with neither inline bytes nor a URL (an OpenAI ``file_id``) is unsendable.
"""
import base64
import binascii
import mimetypes
import posixpath
from collections.abc import Mapping
from dataclasses import dataclass
from itertools import chain
from types import MappingProxyType
from typing import Annotated, Final, Literal, TypeAlias, TypeVar
from urllib.parse import unquote, unquote_to_bytes, urlparse
from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, TypeAdapter, ValidationError
AttachmentType: TypeAlias = Literal["image", "audio", "file"]
_REMOTE_URI_SCHEMES: Final = ("http://", "https://")
_URL_SAFE_TO_STANDARD: Final = str.maketrans("-_", "+/")
# Per attachment type, the fields dropped from the text check because they hold bytes, URLs or file references
_FILE_CHECKED_FIELDS: Final = MappingProxyType(
{
"image_url": frozenset(("image_url", "url")),
"input_image": frozenset(("image_url", "url", "file_id")),
"input_audio": frozenset(("input_audio",)),
"video_url": frozenset(("video_url",)),
"file": frozenset(("file",)),
"input_file": frozenset(("file_data", "file_url", "file_id")),
"image": frozenset(("source",)),
"document": frozenset(("source",)),
}
)
_TEXT_SOURCE_TYPES: Final = frozenset(("text", "content"))
_FILE_SOURCE_FIELDS: Final = frozenset(("file_data", "file_id"))
_ATTACHMENT_BLOCK_TYPES: Final = frozenset(_FILE_CHECKED_FIELDS)
_OBJECT_MAPPING: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
_T: Final = TypeVar("_T")
@dataclass(frozen=True, slots=True)
class Attachment:
filename: str
type: AttachmentType
content: str | None = None
url: str | None = None
def as_payload(self) -> Mapping[str, str]:
fields: Final = (("filename", self.filename), ("type", self.type), ("content", self.content), ("url", self.url))
return MappingProxyType({key: value for key, value in fields if value is not None})
@dataclass(frozen=True, slots=True)
class RequestAttachments:
attachments: tuple[Attachment, ...]
unsendable_count: int
malformed_count: int = 0
def _text_or_none(value: object) -> object:
return value if isinstance(value, str) else None
# Optional metadata the provider ignores when malformed, so a bad value must not fail the whole block
_Metadata: TypeAlias = Annotated[str | None, BeforeValidator(_text_or_none)]
class _Model(BaseModel):
model_config = ConfigDict(extra="ignore")
class _ImageURL(_Model):
url: str | None = None
class _ImageURLBlock(_Model):
type: Literal["image_url"]
image_url: _ImageURL | str | None = None
url: _ImageURL | str | None = None
class _VideoURLBlock(_Model):
type: Literal["video_url"]
video_url: _ImageURL | str
class _InputImageBlock(_Model):
type: Literal["input_image"]
image_url: _ImageURL | str | None = None
url: _ImageURL | str | None = None
file_id: str | None = None
class _InputAudio(_Model):
data: str | None = None
format: _Metadata = None
class _InputAudioBlock(_Model):
type: Literal["input_audio"]
input_audio: _InputAudio
class _FileData(_Model):
file_data: str | None = None
file_id: str | None = None
filename: _Metadata = None
class _FileBlock(_Model):
type: Literal["file"]
file: _FileData
class _InputFileBlock(_Model):
type: Literal["input_file"]
file_data: str | None = None
file_url: str | None = None
file_id: str | None = None
filename: _Metadata = None
class _Source(_Model):
type: _Metadata = None
data: str | None = None
media_type: _Metadata = None
url: str | None = None
content: object = None
class _ImageBlock(_Model):
type: Literal["image"]
source: _Source
class _DocumentBlock(_Model):
type: Literal["document"]
source: _Source
title: _Metadata = None
class _ToolResultBlock(_Model):
type: Literal["tool_result"]
content: object = None
class _MalformedBlock(_Model):
"""An attachment type that doesn't parse; it can't be checked, so it blocks."""
class _Message(_Model):
content: object = None
output: object = None
_AttachmentBlock: TypeAlias = (
_ImageURLBlock
| _VideoURLBlock
| _InputImageBlock
| _InputAudioBlock
| _FileBlock
| _InputFileBlock
| _ImageBlock
| _DocumentBlock
| _ToolResultBlock
)
_BLOCK_ADAPTER: Final[TypeAdapter[_AttachmentBlock]] = TypeAdapter(
Annotated[_AttachmentBlock, Field(discriminator="type")]
)
_Block: TypeAlias = _AttachmentBlock | _MalformedBlock
_MESSAGE_ADAPTER: Final[TypeAdapter[_Message]] = TypeAdapter(_Message)
_ITEMS_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object])
# (attachment, is_unsendable); (None, False) is a block that isn't an attachment
_Classified: TypeAlias = tuple[Attachment | None, bool]
_NOT_AN_ATTACHMENT: Final[_Classified] = (None, False)
_UNSENDABLE: Final[_Classified] = (None, True)
def request_attachments(request_data: Mapping[str, object]) -> RequestAttachments:
# Both, so a decoy "messages" can't hide attachments in a Responses API "input"
containers: Final = (_parse(_ITEMS_ADAPTER, request_data.get(key)) or () for key in ("messages", "input"))
blocks: Final = tuple(chain.from_iterable(_message_blocks(message) for message in chain.from_iterable(containers)))
classified: Final = tuple(
chain.from_iterable(_block_attachments(block, index) for index, block in enumerate(blocks))
)
return RequestAttachments(
attachments=tuple(attachment for attachment, _ in classified if attachment is not None),
unsendable_count=sum(1 for _, is_unsendable in classified if is_unsendable),
malformed_count=sum(1 for block in blocks if isinstance(block, _MalformedBlock)),
)
def _message_blocks(message: object) -> tuple[_Block, ...]:
parsed: Final = _parse(_MESSAGE_ADAPTER, message)
top: Final = (_blocks(parsed.content) + _blocks(parsed.output)) if parsed else ()
nested: Final = _nested_blocks(top)
# tool_result -> document -> image is the deepest the APIs nest
return top + nested + _nested_blocks(nested)
def _nested_blocks(blocks: tuple[_Block, ...]) -> tuple[_Block, ...]:
return tuple(chain.from_iterable(_blocks(_nested_content(block)) for block in blocks))
def _nested_content(block: _Block) -> object:
match block:
case _ToolResultBlock():
return block.content
case _DocumentBlock(source=_Source(type="content")):
return block.source.content
case _:
return None
def _blocks(content: object) -> tuple[_Block, ...]:
items: Final = _parse(_ITEMS_ADAPTER, content)
parsed: Final = (_block(item) for item in items or ())
return tuple(block for block in parsed if block is not None)
def _block(item: object) -> _Block | None:
parsed: Final = _parse(_BLOCK_ADAPTER, item)
if parsed is not None:
return parsed
block_type: Final = (_parse(_OBJECT_MAPPING, item) or {}).get("type")
return _MalformedBlock() if isinstance(block_type, str) and block_type in _ATTACHMENT_BLOCK_TYPES else None
def _block_attachments(block: _Block, index: int) -> tuple[_Classified, ...]:
"""A file block can name several sources and providers differ on which they send, so all are checked."""
match block:
case _FileBlock():
return _file_sources((block.file.file_data,), block.file.file_id, block.file.filename, index)
case _InputFileBlock():
return _file_sources((block.file_data, block.file_url), block.file_id, block.filename, index)
case _ImageURLBlock():
return _file_sources((_url(block.image_url), _url(block.url)), None, None, index, "image")
case _InputImageBlock():
return _file_sources((_url(block.image_url), _url(block.url)), block.file_id, None, index, "image")
case _DocumentBlock():
return (_from_source(block.source, block.title, index, "file"),)
case _:
return (_classify_block(block, index),)
def _file_sources(
inline: tuple[str | None, ...], file_id: str | None, name: str | None, index: int, kind: AttachmentType = "file"
) -> tuple[_Classified, ...]:
found: Final = (
*(_from_uri(source, name, index, kind) for source in inline if source),
*((_from_file_id(file_id, name, index, kind),) if file_id else ()),
)
return found or (_UNSENDABLE,)
def _from_file_id(file_id: str, name: str | None, index: int, kind: AttachmentType) -> _Classified:
"""A URL is checked; an uploaded file's id has no content to send."""
is_url: Final = file_id.strip().lower().startswith(_REMOTE_URI_SCHEMES)
return _from_uri(file_id, name, index, kind) if is_url else _UNSENDABLE
def _classify_block(block: _Block, index: int) -> _Classified:
match block:
case _VideoURLBlock():
return _from_uri(_url(block.video_url), None, index, "file")
case _InputAudioBlock(input_audio=_InputAudio(data=str(data), format=audio_format)):
name: Final = f"attachment-{index}.{audio_format}" if audio_format else None
return _from_base64(data, name, index, "audio", None)
case _InputAudioBlock():
return _UNSENDABLE
case _ImageBlock(source=source):
return _from_source(source, None, index, "image")
case _:
return _NOT_AN_ATTACHMENT
def _url(value: _ImageURL | str | None) -> str | None:
return value.url if isinstance(value, _ImageURL) else value
def _from_uri(raw_uri: str | None, name: str | None, index: int, kind: AttachmentType) -> _Classified:
uri: Final = (raw_uri or "").strip()
if not uri:
return _UNSENDABLE
if uri.lower().startswith(_REMOTE_URI_SCHEMES):
return Attachment(_filename(name, index, url=uri), kind, url=uri), False
media_type, data = _parse_data_uri(uri)
return _from_base64(data, name, index, kind, media_type)
def _from_source(source: _Source, name: str | None, index: int, kind: AttachmentType) -> _Classified:
"""base64 or a URL; text sources stay in the text check, and a file_id has nothing to send."""
match source:
case _Source(type="base64", data=str(data)):
return _from_base64(data, name, index, kind, source.media_type)
case _Source(type=str(source_type)) if source_type in _TEXT_SOURCE_TYPES and kind == "file":
return _NOT_AN_ATTACHMENT
case _Source(type="url", url=str(url)) if url:
return Attachment(_filename(name, index, url=url), kind, url=url), False
case _:
return _UNSENDABLE
def _from_base64(data: str, name: str | None, index: int, kind: AttachmentType, media_type: str | None) -> _Classified:
content: Final = _standard_base64(data)
if content is None:
return _UNSENDABLE
return Attachment(_filename(name, index, media_type), kind, content=content), False
def _standard_base64(data: str) -> str | None:
"""Padded standard base64, accepting line breaks, missing padding and URL-safe characters."""
compact: Final = "".join(data.split()).translate(_URL_SAFE_TO_STANDARD)
padded: Final = compact + "=" * (-len(compact) % 4)
return padded if compact and _is_base64(padded) else None
def _parse_data_uri(uri: str) -> tuple[str | None, str]:
"""(media type, base64 data); a plain data URI's text is encoded, anything else is taken as raw base64."""
if uri[:5].lower() != "data:" or "," not in uri:
return None, uri
header, data = uri[5:].split(",", 1)
params: Final = header.split(";")
encoded: Final = params[-1].strip().lower() == "base64"
return params[0], data if encoded else base64.b64encode(
unquote_to_bytes(data.encode(errors="surrogatepass"))
).decode()
def _filename(name: str | None, index: int, media_type: str | None = None, url: str | None = None) -> str:
"""The client's name, else the URL's, with an extension from the media type when it has none."""
stem: Final = posixpath.basename((name or "").strip()) or _url_basename(url) or f"attachment-{index}"
extension: Final = mimetypes.guess_extension(media_type.split(";")[0].strip()) if media_type else None
return stem if posixpath.splitext(stem)[1] or not extension else f"{stem}{extension}"
def _url_basename(url: str | None) -> str:
try:
return posixpath.basename(unquote(urlparse(url or "").path))
except ValueError:
return ""
def _is_base64(data: str) -> bool:
try:
base64.b64decode(data, validate=True)
except (binascii.Error, ValueError):
return False
return True
def without_attachment_content(messages: object) -> object:
items: Final = _parse(_ITEMS_ADAPTER, messages)
return (
messages
if items is None
else tuple(_without_content(_without_content(message, "content"), "output") for message in items)
)
def _without_content(value: object, key: str) -> object:
mapping: Final = _parse(_OBJECT_MAPPING, value)
blocks: Final = _parse(_ITEMS_ADAPTER, mapping.get(key)) if mapping else None
if mapping is None or blocks is None:
return value
return {**mapping, key: tuple(_block_without_content(block) for block in blocks)}
def _block_without_content(block: object) -> object:
mapping: Final = _parse(_OBJECT_MAPPING, block) or {}
block_type: Final = mapping.get("type")
if block_type == "tool_result":
return _without_content(block, "content")
dropped: Final = _FILE_CHECKED_FIELDS.get(block_type) if isinstance(block_type, str) else None
if dropped is None:
return block
source: Final = _parse(_OBJECT_MAPPING, mapping.get("source")) or {}
source_type: Final = source.get("type")
if block_type == "document" and isinstance(source_type, str) and source_type in _TEXT_SOURCE_TYPES:
# A text document is prompt text, so it is checked here; only images nested in it go to the file check
return {**mapping, "source": _without_content(source, "content")}
kept: Final = {key: value for key, value in mapping.items() if key not in dropped}
file: Final = _parse(_OBJECT_MAPPING, mapping.get("file")) if block_type == "file" else None
if file is None:
return kept
return {**kept, "file": {key: value for key, value in file.items() if key not in _FILE_SOURCE_FIELDS}}
def _parse(adapter: TypeAdapter[_T], value: object) -> _T | None:
try:
return adapter.validate_python(value)
except ValidationError:
return None

View file

@ -100,7 +100,7 @@ _MCP_EVENT_HOOKS: Final = frozenset(
)
def _configured_event_hooks(mode: str | list[str] | Mode) -> tuple[str, ...]:
def configured_event_hooks(mode: str | list[str] | Mode) -> tuple[str, ...]:
if isinstance(mode, str):
return (mode,)
if isinstance(mode, list):
@ -114,7 +114,7 @@ def _configured_event_hooks(mode: str | list[str] | Mode) -> tuple[str, ...]:
def _is_mcp_only_mode(mode: str | list[str] | Mode) -> bool:
hooks: Final = _configured_event_hooks(mode)
hooks: Final = configured_event_hooks(mode)
return bool(hooks) and all(hook in _MCP_EVENT_HOOKS for hook in hooks)

View file

@ -6,7 +6,8 @@ import os
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
from datetime import datetime, timezone
from itertools import chain, count
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeAlias, cast
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeAlias, TypeVar, cast
from pydantic import BaseModel, TypeAdapter, ValidationError
@ -58,6 +59,7 @@ from .guardrail_hooks.llm_as_a_judge import (
initialize_guardrail as initialize_llm_as_a_judge,
)
from .guardrail_initializers import (
configured_event_hooks,
initialize_bedrock,
initialize_hide_secrets,
initialize_lakera,
@ -572,9 +574,46 @@ def _as_callback_tuple(
return (initialized,)
def _configure_callback_scoping(
def _logging_only_scope_error(
custom_guardrail_callback: CustomGuardrail, guardrail_name: str, litellm_params: LitellmParams
) -> str | None:
logging_only_scope: Final = litellm_params.logging_only_scope
if logging_only_scope is not None and GuardrailEventHooks.logging_only.value not in configured_event_hooks(
litellm_params.mode
):
return (
f"Guardrail {guardrail_name}: logging_only_scope is set, but mode does not include logging_only, "
"so it would never apply. Add logging_only to mode or remove logging_only_scope."
)
if logging_only_scope in ("input", "output") and not custom_guardrail_callback.supports_logging_only_scope():
return (
f"Guardrail {guardrail_name}: logging_only_scope={logging_only_scope!r} is not supported by this "
"guardrail, whose logging_only hook scans on its own. Remove logging_only_scope."
)
return None
def _configure_callback_scoping(
custom_guardrail_callback: CustomGuardrail,
guardrail_name: str,
litellm_params: LitellmParams,
*,
reject_invalid_logging_only_scope: bool = False,
) -> None:
logging_only_scope: Final = litellm_params.logging_only_scope
logging_only_scope_error: Final = _logging_only_scope_error(
custom_guardrail_callback, guardrail_name, litellm_params
)
if logging_only_scope_error is not None:
if reject_invalid_logging_only_scope:
raise ValueError(logging_only_scope_error)
verbose_proxy_logger.error(
"%s Ignoring logging_only_scope; the guardrail keeps its configured mode.",
logging_only_scope_error.replace("\r", "").replace("\n", ""),
)
custom_guardrail_callback.logging_only_scope = None
else:
custom_guardrail_callback.logging_only_scope = logging_only_scope
for scoping_param in (
"skip_system_message_in_guardrail",
"skip_tool_message_in_guardrail",
@ -597,6 +636,28 @@ def _configure_callback_scoping(
_apply_configured_bool_overrides(custom_guardrail_callback, litellm_params)
_ParamsT = TypeVar("_ParamsT", bound=BaseModel)
def parse_tolerant_litellm_params(
litellm_params_data: Mapping[str, object],
guardrail_name: str,
params_model: type[_ParamsT] = LitellmParams,
) -> _ParamsT:
try:
return params_model(**litellm_params_data)
except ValidationError as validation_error:
if any(tuple(error["loc"]) != ("logging_only_scope",) for error in validation_error.errors()):
raise
verbose_proxy_logger.error(
"Guardrail %s: logging_only_scope=%r is not one of 'input', 'output' or 'both'. "
"Ignoring logging_only_scope; the guardrail keeps its configured mode.",
guardrail_name.replace("\r", "").replace("\n", ""),
str(litellm_params_data.get("logging_only_scope")).replace("\r", "").replace("\n", "")[:100],
)
return params_model(**MappingProxyType({**litellm_params_data, "logging_only_scope": None}))
class InMemoryGuardrailHandler:
"""
Class that handles initializing guardrails and adding them to the CallbackManager
@ -633,6 +694,8 @@ class InMemoryGuardrailHandler:
config_file_path: str | None = None,
llm_router: Optional["Router"] = None,
source: Literal["db", "config"] = "config",
*,
reject_invalid_logging_only_scope: bool = False,
) -> Guardrail | None:
"""
Initialize a guardrail from a dictionary and add it to the litellm callback manager
@ -653,7 +716,10 @@ class InMemoryGuardrailHandler:
verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data)
if isinstance(litellm_params_data, dict):
litellm_params = LitellmParams(**litellm_params_data)
if reject_invalid_logging_only_scope:
litellm_params = LitellmParams(**litellm_params_data)
else:
litellm_params = parse_tolerant_litellm_params(litellm_params_data, guardrail["guardrail_name"])
else:
litellm_params = litellm_params_data
@ -679,8 +745,18 @@ class InMemoryGuardrailHandler:
config_file_path=config_file_path,
llm_router=llm_router,
)
for custom_guardrail_callback in created_callbacks:
_configure_callback_scoping(custom_guardrail_callback, guardrail["guardrail_name"], litellm_params)
try:
for custom_guardrail_callback in created_callbacks:
_configure_callback_scoping(
custom_guardrail_callback,
guardrail["guardrail_name"],
litellm_params,
reject_invalid_logging_only_scope=reject_invalid_logging_only_scope,
)
except Exception:
for custom_guardrail_callback in created_callbacks:
litellm.logging_callback_manager.remove_callback_from_all_lists(custom_guardrail_callback)
raise
parsed_guardrail: Final = Guardrail(
guardrail_id=guardrail.get("guardrail_id"),
@ -729,6 +805,28 @@ class InMemoryGuardrailHandler:
siblings: Final = self.guardrail_id_to_sibling_callbacks.get(guardrail_id, ())
return (() if primary is None else (primary,)) + siblings
def _reject_invalid_logging_only_scope(self, guardrail_id: str, guardrail: Guardrail) -> None:
"""
Strictly validate logging_only_scope on a row whose params are otherwise
unchanged, without rebuilding the live callback.
API write paths send the whole object, so an invalid scope must still be
rejected even when the write changed nothing else. But an unchanged row
must not force a teardown + re-append: initialize_guardrail appends the
rebuilt callback at the END of litellm.callbacks, so a no-op PUT would
reorder guardrails and change which one wins between a BLOCK and a MASK
guardrail over the same content.
"""
params: Final = guardrail.get("litellm_params")
if not isinstance(params, (dict, LitellmParams)):
return
litellm_params: Final = LitellmParams(**params) if isinstance(params, dict) else params
guardrail_name: Final = guardrail.get("guardrail_name", "Unknown")
for custom_guardrail_callback in self._tracked_callbacks(guardrail_id):
scope_error = _logging_only_scope_error(custom_guardrail_callback, guardrail_name, litellm_params)
if scope_error is not None:
raise ValueError(scope_error)
def initialize_custom_guardrail(
self,
guardrail: Guardrail,
@ -788,6 +886,8 @@ class InMemoryGuardrailHandler:
guardrail_id: str,
guardrail: Guardrail,
source: Literal["db", "config"] = "db",
*,
reject_invalid_logging_only_scope: bool = False,
) -> None:
"""
Update a guardrail in memory: a changed name or litellm_params rebuilds the
@ -796,8 +896,14 @@ class InMemoryGuardrailHandler:
"""
updated_guardrail: Final = cast(Guardrail, {**guardrail, "guardrail_id": guardrail_id})
if self._has_guardrail_params_changed(guardrail_id, updated_guardrail):
self.reinitialize_guardrail(guardrail=updated_guardrail, source=source)
self.reinitialize_guardrail(
guardrail=updated_guardrail,
source=source,
reject_invalid_logging_only_scope=reject_invalid_logging_only_scope,
)
return
if reject_invalid_logging_only_scope:
self._reject_invalid_logging_only_scope(guardrail_id, updated_guardrail)
self.IN_MEMORY_GUARDRAILS[guardrail_id] = updated_guardrail
self._sources[guardrail_id] = source
@ -883,6 +989,7 @@ class InMemoryGuardrailHandler:
@staticmethod
def _normalize_litellm_params_for_comparison(
params: LitellmParams | Mapping[str, object] | None,
guardrail_name: str,
) -> Mapping[str, object] | None:
"""
Render litellm_params to a canonical dict so an in-memory LitellmParams and
@ -899,7 +1006,7 @@ class InMemoryGuardrailHandler:
return params.model_dump()
if isinstance(params, dict):
try:
return LitellmParams(**params).model_dump()
return parse_tolerant_litellm_params(params, guardrail_name).model_dump()
except ValidationError as e:
verbose_proxy_logger.warning(
"Could not normalize guardrail litellm_params for comparison; treating the guardrail as changed. Error: %s",
@ -922,8 +1029,12 @@ class InMemoryGuardrailHandler:
return True
# Compare litellm_params
existing_dict: Final = self._normalize_litellm_params_for_comparison(existing.get("litellm_params"))
new_dict: Final = self._normalize_litellm_params_for_comparison(new_guardrail.get("litellm_params"))
existing_dict: Final = self._normalize_litellm_params_for_comparison(
existing.get("litellm_params"), existing.get("guardrail_name", "Unknown")
)
new_dict: Final = self._normalize_litellm_params_for_comparison(
new_guardrail.get("litellm_params"), new_guardrail.get("guardrail_name", "Unknown")
)
# Compare and identify specific differences
changed_fields = {}
@ -949,6 +1060,8 @@ class InMemoryGuardrailHandler:
guardrail: Guardrail,
config_file_path: str | None = None,
source: Literal["db", "config"] = "config",
*,
reject_invalid_logging_only_scope: bool = False,
) -> Guardrail | None:
"""
Force re-initialization of a guardrail even if it exists in memory.
@ -978,7 +1091,12 @@ class InMemoryGuardrailHandler:
# instance instead of leaving the guardrail silently removed: a guardrail
# that was enforcing must never fail open because an update was bad.
try:
return self.initialize_guardrail(guardrail=guardrail, config_file_path=config_file_path, source=source)
return self.initialize_guardrail(
guardrail=guardrail,
config_file_path=config_file_path,
source=source,
reject_invalid_logging_only_scope=reject_invalid_logging_only_scope,
)
except Exception as init_error:
if previous_guardrail is not None:
verbose_proxy_logger.exception(
@ -1003,7 +1121,9 @@ class InMemoryGuardrailHandler:
)
if existing is None or db_params is None or not contains_encrypted_marker(db_params):
return guardrail
loaded_params: Final = self._normalize_litellm_params_for_comparison(existing.get("litellm_params"))
loaded_params: Final = self._normalize_litellm_params_for_comparison(
existing.get("litellm_params"), guardrail.get("guardrail_name", "Unknown")
)
verbose_proxy_logger.warning(
"Guardrail %s has litellm_params that do not decrypt with the current key; keeping the loaded values for "
"them. Restart the proxy if the master key was rotated.",
@ -1023,7 +1143,13 @@ class InMemoryGuardrailHandler:
}
)
def sync_guardrail_from_db(self, guardrail: Guardrail, config_file_path: str | None = None) -> Guardrail | None:
def sync_guardrail_from_db(
self,
guardrail: Guardrail,
config_file_path: str | None = None,
*,
reject_invalid_logging_only_scope: bool = False,
) -> Guardrail | None:
"""
Sync a guardrail from DB - initializes if new, re-initializes if changed.
DB values that do not decrypt with the current key keep the loaded guardrail's values.
@ -1044,8 +1170,12 @@ class InMemoryGuardrailHandler:
guardrail=synced,
config_file_path=config_file_path,
source="db",
reject_invalid_logging_only_scope=reject_invalid_logging_only_scope,
)
if reject_invalid_logging_only_scope:
self._reject_invalid_logging_only_scope(guardrail_id, synced)
# Params unchanged but the entry is still DB-backed; make sure the
# source marker reflects that even if it was previously set differently
# (e.g. a config entry whose UUID later collided with a DB row).

View file

@ -76,6 +76,7 @@ from litellm.router_utils.add_retry_fallback_headers import (
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
from litellm.types.utils import (
CallTypes,
@ -2984,6 +2985,48 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
)
async def enforce_mcp_server_rate_limits(
self,
user_api_key_dict: UserAPIKeyAuth | None,
server: MCPServer,
) -> None:
descriptors: Final[list[RateLimitDescriptor]] = [] # mutable-ok: existing descriptor helpers append in place
mcp_server_name: Final = server.alias or server.server_name or server.name
if user_api_key_dict is not None:
self._add_mcp_per_key_rate_limit_descriptor(
user_api_key_dict=user_api_key_dict,
mcp_server_name=mcp_server_name,
descriptors=descriptors,
)
self._add_mcp_per_team_rate_limit_descriptor(
user_api_key_dict=user_api_key_dict,
mcp_server_name=mcp_server_name,
descriptors=descriptors,
)
if server.rpm is not None:
descriptors.append(
RateLimitDescriptor(
key="mcp_server",
value=server.server_id,
rate_limit={
"requests_per_unit": server.rpm,
"tokens_per_unit": None,
"window_size": self.window_size,
},
)
)
if not descriptors:
return
parent_otel_span: Final = user_api_key_dict.parent_otel_span if user_api_key_dict is not None else None
response: Final = await self.atomic_check_and_increment_by_n(
descriptors=descriptors,
increments=[{"requests": 1} for _ in descriptors],
parent_otel_span=parent_otel_span,
)
if response["overall_code"] == "OVER_LIMIT":
self._handle_rate_limit_error(response, descriptors)
def _should_enforce_rate_limit(
self,
limit_type: str | None,
@ -3270,21 +3313,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
descriptors=descriptors,
)
# REST MCP calls pass the raw body through this hook before server
# resolution; only the later synthetic hook payload may carry this key.
if call_type == CallTypes.call_mcp_tool.value and "server_id" not in data:
mcp_server_name: Final = data.get("mcp_server_name", None)
self._add_mcp_per_key_rate_limit_descriptor(
user_api_key_dict=user_api_key_dict,
mcp_server_name=mcp_server_name,
descriptors=descriptors,
)
self._add_mcp_per_team_rate_limit_descriptor(
user_api_key_dict=user_api_key_dict,
mcp_server_name=mcp_server_name,
descriptors=descriptors,
)
self._add_team_model_rate_limit_descriptor_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model if isinstance(requested_model, str) else None,

View file

@ -52,6 +52,8 @@ from litellm.proxy.lens.models import (
from litellm.proxy.lens.release import PROTOCOL_VERSION, release_tag, worker_image
from litellm.proxy.lens.repository import DueLens, LensRepository, WriterDatabase
from litellm.proxy.lens.reviews import criteria_key
from litellm.proxy.lens.signal_repository import SignalRepository
from litellm.proxy.lens.signals import SignalConfig, TraceSignals, trace_signals
from litellm.proxy.lens.sources import ActivityAvailability, SourceReader, Storage, parse_execution
from litellm.proxy.lens.state import (
can_access,
@ -69,6 +71,7 @@ from litellm.proxy.lens.state import (
summarized,
)
from litellm.proxy.tracing_runtime import provide_storage
from litellm.router import Router
from litellm.types.llms.base import LiteLLMBaseModel
router: Final = APIRouter(prefix="/lens", tags=["Lens"])
@ -101,6 +104,14 @@ def repository() -> LensRepository:
return LensRepository(WriterDatabase(writer_wrapper(prisma_client.db)))
def signals_repository() -> SignalRepository:
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(503, "Lens needs a connected Postgres database")
return SignalRepository(WriterDatabase(writer_wrapper(prisma_client.db)))
def source_reader(storage: Storage | None) -> SourceReader:
if storage is None:
raise HTTPException(
@ -118,6 +129,20 @@ def user_scope(auth: UserAPIKeyAuth, write: bool = False) -> Scope:
raise HTTPException(403, "Lens requires proxy administrator access")
def validate_signal_model(config: SignalConfig, llm_router: Router | None) -> None:
if not config.model:
return
message: Final = "Choose a System 1 model (evaluation mode) configured on this proxy"
if llm_router is None:
raise HTTPException(400, message)
try:
model_group: Final = llm_router.get_model_group_info(model_group=config.model)
except Exception as error:
raise HTTPException(400, message) from error
if model_group is None or model_group.mode != "evaluation":
raise HTTPException(400, message)
async def get_lens(lens_id: str, scope: Scope) -> Lens:
lens: Final = await repository().get(lens_id)
if lens is None or not can_access(scope, lens.scope):
@ -261,6 +286,39 @@ async def list_agents(auth: Auth, storage: StorageDep) -> tuple[str, ...]:
return await source_reader(storage).agents(scope) if storage is not None else ()
@router.get("/signals", response_model=SignalConfig)
async def get_signals(auth: Auth) -> SignalConfig:
user_scope(auth)
return await signals_repository().get_config()
@router.put("/signals", response_model=SignalConfig)
async def put_signals(body: SignalConfig, auth: Auth) -> SignalConfig:
user_scope(auth, write=True)
from litellm.proxy.proxy_server import llm_router
validate_signal_model(body, llm_router)
await signals_repository().save_config(body)
return body
@router.post("/traces/signals", response_model=tuple[TraceSignals, ...])
async def trace_signal_statuses(body: TraceFindingsRequest, auth: Auth) -> tuple[TraceSignals, ...]:
user_scope(auth)
repo: Final = signals_repository()
config: Final = await repo.get_config()
existing: Final = await repo.traces(body.traces)
rows: Final = MappingProxyType({(row.trace_id, row.trace_ref): row for row in existing})
return tuple(
trace_signals(
trace,
rows.get((trace.trace_id, trace.trace_ref)),
config,
)
for trace in body.traces
)
@router.post("/traces/findings", response_model=tuple[TraceFindingCount, ...])
async def trace_findings(body: TraceFindingsRequest, auth: Auth) -> tuple[TraceFindingCount, ...]:
user_scope(auth)

View file

@ -0,0 +1,142 @@
import json
from datetime import datetime
from typing import Final
from pydantic import TypeAdapter
from litellm.proxy.lens.models import Execution, TraceIdentity
from litellm.proxy.lens.repository import Database, Row
from litellm.proxy.lens.signals import (
SIGNAL_RECLASSIFY_AFTER,
SIGNAL_RETRY_FAILED_AFTER,
SignalAttempt,
SignalConfig,
StoredTraceSignal,
)
_ROWS: Final[TypeAdapter[tuple[Row, ...]]] = TypeAdapter(tuple[Row, ...])
class SignalRepository:
def __init__(self, db: Database) -> None:
self.db: Final = db
async def get_config(self) -> SignalConfig:
rows: Final = _ROWS.validate_python(
await self.db.query_raw('SELECT data FROM "LiteLLM_LensSignalConfig" WHERE id=$1', "global")
)
return SignalConfig() if not rows else SignalConfig.model_validate(rows[0].data)
async def save_config(self, config: SignalConfig) -> None:
await self.db.execute_raw(
"""INSERT INTO "LiteLLM_LensSignalConfig" (id, data)
VALUES ($1, $2::jsonb)
ON CONFLICT (id) DO UPDATE SET data=EXCLUDED.data""",
"global",
json.dumps(config.model_dump(mode="json")),
)
async def traces(self, identities: tuple[TraceIdentity, ...]) -> tuple[StoredTraceSignal, ...]:
if not identities:
return ()
payload: Final = json.dumps(
tuple({"trace_id": trace.trace_id, "trace_ref": trace.trace_ref} for trace in identities)
)
rows: Final = _ROWS.validate_python(
await self.db.query_raw(
"""SELECT jsonb_build_object(
'trace_id', trace_id,
'trace_ref', trace_ref,
'config_key', config_key,
'span_count', span_count,
'claimed_until', claimed_until,
'classified_at', classified_at,
'data', data
) AS data
FROM "LiteLLM_LensTraceSignal"
WHERE (trace_id, trace_ref) IN (
SELECT trace_id, trace_ref FROM jsonb_to_recordset($1::jsonb) AS requested(
trace_id text, trace_ref text
)
)""",
payload,
)
)
return tuple(StoredTraceSignal.model_validate(row.data) for row in rows)
async def claim(
self,
execution: Execution,
config: SignalConfig,
claimed_until: datetime,
now: datetime,
) -> bool:
data: Final = json.dumps({"status": "pending", "scores": {}, "model": config.model, "error": ""})
rows: Final = _ROWS.validate_python(
await self.db.query_raw(
"""INSERT INTO "LiteLLM_LensTraceSignal" AS stored
(trace_id, trace_ref, config_key, span_count, claimed_until, classified_at, data)
VALUES ($1, $2, $3, $4, $5::timestamp, NULL, $6::jsonb)
ON CONFLICT (trace_id, trace_ref) DO UPDATE SET
config_key=EXCLUDED.config_key,
span_count=EXCLUDED.span_count,
claimed_until=EXCLUDED.claimed_until,
classified_at=NULL,
data=EXCLUDED.data
WHERE (stored.claimed_until IS NULL OR stored.claimed_until < $7::timestamp)
AND (
stored.config_key IS DISTINCT FROM EXCLUDED.config_key
OR (
stored.data->>'status'='pending'
AND stored.claimed_until < $7::timestamp
)
OR (
EXCLUDED.span_count > stored.span_count
AND stored.classified_at < $8::timestamp
)
OR (
stored.data->>'status'='failed'
AND stored.classified_at < $9::timestamp
)
)
RETURNING jsonb_build_object('trace_id', trace_id) AS data""",
execution.trace_id,
execution.trace_ref,
config.key(),
execution.span_count,
claimed_until,
data,
now,
now - SIGNAL_RECLASSIFY_AFTER,
now - SIGNAL_RETRY_FAILED_AFTER,
)
)
return bool(rows)
async def store(
self,
execution: Execution,
config: SignalConfig,
claimed_until: datetime,
classified_at: datetime,
attempt: SignalAttempt,
) -> None:
payload: Final = json.dumps(
{
"status": attempt.status,
"scores": dict(attempt.scores),
"model": attempt.model,
"error": attempt.error,
}
)
await self.db.execute_raw(
"""UPDATE "LiteLLM_LensTraceSignal"
SET classified_at=$1::timestamp, claimed_until=NULL, data=$2::jsonb
WHERE trace_id=$3 AND trace_ref=$4 AND config_key=$5 AND claimed_until=$6::timestamp""",
classified_at,
payload,
execution.trace_id,
execution.trace_ref,
config.key(),
claimed_until,
)

View file

@ -0,0 +1,619 @@
import asyncio
import hashlib
import json
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from itertools import accumulate
from types import MappingProxyType
from typing import Annotated, Final, Literal, Protocol, TypeAlias
from pydantic import ConfigDict, Field, JsonValue, ValidationError, field_validator, model_validator
from litellm.integrations.clickhouse.context import lens_analysis
from litellm.litellm_core_utils.initialize_dynamic_callback_params import inherit_message_logging_privacy
from litellm.litellm_core_utils.secret_redaction import redact_internal_details
from litellm.proxy.lens.models import ActivitySelection, Execution, Record, Scope, TraceIdentity
from litellm.proxy.lens.sources import SourceReader, Storage
SIGNAL_SETTLE: Final = timedelta(seconds=15)
SIGNAL_PAGE_SIZE: Final = 100
SIGNAL_MAX_PER_TICK: Final = 50
SIGNAL_CONCURRENCY: Final = 8
SIGNAL_CLAIM_LEASE: Final = timedelta(minutes=5)
SIGNAL_RECLASSIFY_AFTER: Final = timedelta(minutes=5)
SIGNAL_RETRY_FAILED_AFTER: Final = timedelta(minutes=30)
SIGNAL_MAX_CONTENT_PAGES: Final = 3
SIGNAL_PART_MAX_CHARS: Final = 2000
SIGNAL_PART_HEAD_CHARS: Final = 800
SIGNAL_PART_TAIL_CHARS: Final = 1200
SIGNAL_TRANSCRIPT_MAX_CHARS: Final = 40000
SIGNAL_TRANSCRIPT_HEAD_CHARS: Final = 15000
SIGNAL_TRANSCRIPT_TAIL_CHARS: Final = 25000
SIGNAL_MAX_SCAN_PAGES: Final = 10
@dataclass(frozen=True, slots=True)
class SignalSweep:
lookback: timedelta
interval_seconds: float
max_pages: int
SIGNAL_LIVE_SWEEP: Final = SignalSweep(lookback=timedelta(minutes=15), interval_seconds=2, max_pages=1)
SIGNAL_BACKLOG_SWEEP: Final = SignalSweep(
lookback=timedelta(hours=24), interval_seconds=60, max_pages=SIGNAL_MAX_SCAN_PAGES
)
SIGNAL_TASK: Final = (
"An AI agent run recorded as a trace. Judge only what the user and the agent said and did in these steps."
)
class Signal(Record):
id: str = Field(pattern=r"^[a-z][a-z0-9_]{0,63}$")
name: str = Field(min_length=1, max_length=60)
question: str = Field(min_length=3, max_length=500)
DEFAULT_SIGNALS: Final[tuple[Signal, ...]] = (
Signal(
id="user_frustration",
name="User frustration",
question=(
"Does the user show frustration, annoyance or dissatisfaction with the agent in this run, for example "
"complaints, irritated corrections, all caps, profanity, or giving up on the task?"
),
),
Signal(
id="missing_capability",
name="Missing capability",
question=(
"Does the user ask for something the agent cannot do in this run, so that the agent refuses, says it "
"lacks a tool, permission, integration or data source, or fails because the capability does not exist?"
),
),
Signal(
id="repeated_request",
name="Repeated request",
question=(
"Does the user ask for the same thing more than once in this run, usually because the agent did not "
"deliver it the first time?"
),
),
)
class SignalConfig(Record):
model: str = ""
threshold: float = Field(default=0.5, ge=0.05, le=0.95, allow_inf_nan=False)
signals: tuple[Signal, ...] = DEFAULT_SIGNALS
@model_validator(mode="after")
def validate_signals(self) -> "SignalConfig":
if len(self.signals) > 20:
raise ValueError("A maximum of 20 signals is allowed")
if len(frozenset(signal.id for signal in self.signals)) != len(self.signals):
raise ValueError("Signal IDs must be unique")
return self
@property
def enabled(self) -> bool:
return bool(self.model) and bool(self.signals)
def key(self) -> str:
payload: Final = json.dumps(
{
"model": self.model,
"signals": tuple({"id": signal.id, "question": signal.question} for signal in self.signals),
},
ensure_ascii=False,
sort_keys=True,
separators=(",", ":"),
)
return hashlib.sha256(payload.encode()).hexdigest()
Score: TypeAlias = Annotated[float, Field(ge=0, le=1, allow_inf_nan=False)]
class SignalFlag(Record):
signal_id: str
name: str
score: Score
class TraceSignals(TraceIdentity):
status: Literal["unclassified", "pending", "classified", "failed"]
flags: tuple[SignalFlag, ...] = ()
model: str = ""
classified_at: datetime | None = None
class SignalStep(Record):
kind: str
name: str
content: str
class SignalData(Record):
model_config = ConfigDict(extra="ignore")
status: Literal["pending", "classified", "failed"] = "pending"
scores: Mapping[str, Score] = Field(default_factory=lambda: MappingProxyType({}))
model: str = ""
error: str = ""
class SignalAttempt(Record):
status: Literal["classified", "failed"]
scores: Mapping[str, Score] = Field(default_factory=lambda: MappingProxyType({}))
model: str
error: str = ""
class StoredTraceSignal(Record):
trace_id: str
trace_ref: str = ""
config_key: str
span_count: int
claimed_until: datetime | None = None
classified_at: datetime | None = None
data: JsonValue
@field_validator("claimed_until", "classified_at")
@classmethod
def normalize_database_timestamp(cls, value: datetime | None) -> datetime | None:
if value is not None and value.tzinfo is None:
return value.replace(tzinfo=timezone.utc)
return value
class NoulAnswer(Record):
model_config = ConfigDict(extra="ignore", allow_inf_nan=False, from_attributes=True)
type: Literal["noul"]
noul: float = Field(ge=0, le=1, allow_inf_nan=False)
class DecisionsOutput(Record):
model_config = ConfigDict(extra="ignore", from_attributes=True)
answers: Mapping[str, object]
DecisionState: TypeAlias = Mapping[str, object]
DecisionQuestions: TypeAlias = Mapping[str, Mapping[str, str]]
Clock: TypeAlias = Callable[[], datetime]
RouterReady: TypeAlias = Callable[[], bool]
class DecisionsCall(Protocol):
async def __call__(
self,
*,
model: str,
state: DecisionState,
questions: DecisionQuestions,
timeout: float,
metadata: Mapping[str, object],
) -> object: ...
class SignalRepositoryProtocol(Protocol):
async def get_config(self) -> SignalConfig: ...
async def traces(self, identities: tuple[TraceIdentity, ...]) -> tuple[StoredTraceSignal, ...]: ...
async def claim(
self,
execution: Execution,
config: SignalConfig,
claimed_until: datetime,
now: datetime,
) -> bool: ...
async def store(
self,
execution: Execution,
config: SignalConfig,
claimed_until: datetime,
classified_at: datetime,
attempt: SignalAttempt,
) -> None: ...
def signal_identity(trace: TraceIdentity | StoredTraceSignal | Execution) -> tuple[str, str]:
return trace.trace_id, trace.trace_ref
def candidate(
trace: Execution,
existing: StoredTraceSignal | None,
config_key: str,
now: datetime,
) -> bool:
if existing is None:
return True
if existing.claimed_until is not None and existing.claimed_until > now:
return False
if existing.config_key != config_key:
return True
status: Final = existing.data.get("status") if isinstance(existing.data, dict) else ""
if status == "pending":
return existing.claimed_until is not None and existing.claimed_until <= now
if existing.span_count > trace.span_count:
return False
if existing.span_count < trace.span_count:
return existing.classified_at is not None and existing.classified_at < now - SIGNAL_RECLASSIFY_AFTER
return (
status == "failed"
and existing.classified_at is not None
and existing.classified_at < now - SIGNAL_RETRY_FAILED_AFTER
)
def _take_head(steps: tuple[SignalStep, ...], remaining: int) -> tuple[SignalStep, ...]:
if remaining <= 0:
return ()
cumulative_lengths: Final = tuple(accumulate(len(step.content) for step in steps))
boundary: Final = next((index for index, total in enumerate(cumulative_lengths) if total >= remaining), None)
if boundary is None:
return steps
preceding: Final = steps[:boundary]
last: Final = steps[boundary]
used: Final = cumulative_lengths[boundary - 1] if boundary > 0 else 0
last_length: Final = remaining - used
return (
*preceding,
last
if last_length == len(last.content)
else last.model_copy(update=MappingProxyType({"content": last.content[:last_length]})),
)
def _take_tail(steps: tuple[SignalStep, ...], remaining: int) -> tuple[SignalStep, ...]:
if remaining <= 0:
return ()
reversed_steps: Final = tuple(reversed(steps))
cumulative_lengths: Final = tuple(accumulate(len(step.content) for step in reversed_steps))
boundary: Final = next((index for index, total in enumerate(cumulative_lengths) if total >= remaining), None)
if boundary is None:
return steps
preceding: Final = reversed_steps[:boundary]
last: Final = reversed_steps[boundary]
used: Final = cumulative_lengths[boundary - 1] if boundary > 0 else 0
last_length: Final = remaining - used
selected: Final = (
*preceding,
last
if last_length == len(last.content)
else last.model_copy(update=MappingProxyType({"content": last.content[-last_length:]})),
)
return tuple(reversed(selected))
def _bounded_steps(steps: tuple[SignalStep, ...]) -> tuple[SignalStep, ...]:
if sum(len(step.content) for step in steps) <= SIGNAL_TRANSCRIPT_MAX_CHARS:
return steps
head: Final = _take_head(steps, SIGNAL_TRANSCRIPT_HEAD_CHARS)
tail: Final = _take_tail(steps, SIGNAL_TRANSCRIPT_TAIL_CHARS)
omitted_count: Final = len(steps) - len(head) - len(tail)
marker: Final = SignalStep(kind="omitted", name="", content=f"{omitted_count} steps omitted")
return (*head, marker, *tail)
def _part_excerpt(content: str) -> str:
if len(content) <= SIGNAL_PART_MAX_CHARS:
return content
omitted: Final = len(content) - SIGNAL_PART_MAX_CHARS
marker: Final = f"\n[... {omitted} characters omitted ...]\n"
return f"{content[:SIGNAL_PART_HEAD_CHARS]}{marker}{content[-SIGNAL_PART_TAIL_CHARS:]}"
async def _content_pages(
reader: SourceReader,
scope: Scope,
execution: Execution,
cursor: str,
pages_left: int,
) -> tuple[SignalStep, ...]:
if pages_left == 0:
return ()
content: Final = await reader.content(scope, execution, cursor)
current: Final = tuple(
SignalStep(kind=part.kind, name=part.name, content=_part_excerpt(part.content)) for part in content.parts
)
rest: Final = (
await _content_pages(reader, scope, execution, content.next_cursor, pages_left - 1)
if content.next_cursor is not None
else ()
)
return (*current, *rest)
async def signal_state(reader: SourceReader, scope: Scope, execution: Execution) -> DecisionState:
steps: Final = _bounded_steps(await _content_pages(reader, scope, execution, "", SIGNAL_MAX_CONTENT_PAGES))
return {
"task": SIGNAL_TASK,
"steps": tuple(step.model_dump(mode="json") for step in steps),
}
def _noul_score(value: object) -> float | None:
try:
return NoulAnswer.model_validate(value).noul
except ValidationError:
return None
class SignalClassifier:
def __init__(self, reader: SourceReader, completion: DecisionsCall, clock: Clock) -> None:
self.reader: Final = reader
self.completion: Final = completion
self.clock: Final = clock
async def classify(self, scope: Scope, execution: Execution, config: SignalConfig) -> SignalAttempt:
try:
state: Final = await signal_state(self.reader, scope, execution)
questions: Final = {
signal.id: {"type": "noul", "instructions": signal.question} for signal in config.signals
}
with lens_analysis(), inherit_message_logging_privacy(True):
response: Final = await self.completion(
model=config.model,
state=state,
questions=questions,
timeout=60,
metadata={"tags": ["litellm-lens-signals"]},
)
output: Final = DecisionsOutput.model_validate(response)
scores: Final = MappingProxyType(
{
signal.id: score
for signal in config.signals
if (score := _noul_score(output.answers.get(signal.id))) is not None
}
)
if len(scores) != len(config.signals):
return SignalAttempt(
status="failed",
scores=scores,
model=config.model,
error="Decisions response omitted a configured noul answer",
)
return SignalAttempt(status="classified", scores=scores, model=config.model)
except Exception as error:
detail: Final = redact_internal_details(str(error))[:300]
return SignalAttempt(status="failed", model=config.model, error=detail)
def trace_signals(
trace: TraceIdentity,
existing: StoredTraceSignal | None,
config: SignalConfig,
) -> TraceSignals:
if existing is None or existing.config_key != config.key():
return TraceSignals(trace_id=trace.trace_id, trace_ref=trace.trace_ref, status="unclassified")
data: Final = SignalData.model_validate(existing.data)
if data.status == "pending":
return TraceSignals(
trace_id=trace.trace_id,
trace_ref=trace.trace_ref,
status="pending",
model=data.model,
)
if data.status == "failed" or data.error:
return TraceSignals(
trace_id=trace.trace_id,
trace_ref=trace.trace_ref,
status="failed",
model=data.model,
classified_at=existing.classified_at,
)
flags: Final = tuple(
sorted(
(
SignalFlag(signal_id=signal.id, name=signal.name, score=data.scores[signal.id])
for signal in config.signals
if signal.id in data.scores and data.scores[signal.id] >= config.threshold
),
key=lambda flag: flag.score,
reverse=True,
)
)
return TraceSignals(
trace_id=trace.trace_id,
trace_ref=trace.trace_ref,
status="classified",
flags=flags,
model=data.model,
classified_at=existing.classified_at,
)
async def _process_claimed(
classifier: SignalClassifier,
repository: SignalRepositoryProtocol,
scope: Scope,
execution: Execution,
config: SignalConfig,
claimed_until: datetime,
) -> None:
from litellm._logging import verbose_proxy_logger
attempt: Final = await classifier.classify(scope, execution, config)
try:
await repository.store(execution, config, claimed_until, classifier.clock(), attempt)
except Exception as error:
verbose_proxy_logger.error("Lens signal result could not be stored: %s", redact_internal_details(str(error)))
class _SignalScan:
def __init__(
self,
reader: SourceReader,
repository: SignalRepositoryProtocol,
scope: Scope,
config: SignalConfig,
now: datetime,
cursor: str,
limit: int,
sweep: SignalSweep,
) -> None:
self.reader: Final = reader
self.repository: Final = repository
self.scope: Final = scope
self.config: Final = config
self.now: Final = now
self.cursor: str = cursor
self.limit: Final = limit
self.sweep: Final = sweep
self.executions: tuple[Execution, ...] = ()
self.finished: bool = False
async def _read_page(self, start: int, end: int) -> tuple[tuple[Execution, ...], str | None]:
page_cursor: Final = self.cursor
sample: Final = await self.reader.sample(
self.scope,
ActivitySelection(source="traces"),
start,
end,
page_size=SIGNAL_PAGE_SIZE,
cursor=page_cursor,
)
identities: Final = tuple(
TraceIdentity(trace_id=trace.trace_id, trace_ref=trace.trace_ref) for trace in sample.executions
)
existing_rows: Final = await self.repository.traces(identities)
existing: Final = MappingProxyType({signal_identity(row): row for row in existing_rows})
remaining: Final = self.limit - len(self.executions)
all_eligible: Final = tuple(
execution
for execution in sample.executions
if candidate(execution, existing.get(signal_identity(execution)), self.config.key(), self.now)
)
eligible: Final = all_eligible[:remaining]
next_cursor: Final = page_cursor if len(all_eligible) > remaining else sample.next_cursor
return eligible, next_cursor
async def run(self) -> tuple[tuple[Execution, ...], str]:
start: Final = int((self.now - self.sweep.lookback).timestamp() * 1000)
end: Final = int((self.now - SIGNAL_SETTLE).timestamp() * 1000)
for _ in range(self.sweep.max_pages):
if self.finished or len(self.executions) >= self.limit:
break
eligible, next_cursor = await self._read_page(start, end)
self.executions = (*self.executions, *eligible)
if next_cursor is None:
self.cursor = ""
self.finished = True
else:
self.cursor = next_cursor
return self.executions, self.cursor
async def _scan_pages(
reader: SourceReader,
repository: SignalRepositoryProtocol,
scope: Scope,
config: SignalConfig,
now: datetime,
cursor: str,
remaining: int,
sweep: SignalSweep,
) -> tuple[tuple[Execution, ...], str]:
if remaining <= 0:
return (), cursor
scan: Final = _SignalScan(reader, repository, scope, config, now, cursor, remaining, sweep)
return await scan.run()
@dataclass(frozen=True, slots=True)
class SignalTick:
cursor: str
claimed: int = 0
async def run_signal_tick(
storage: Storage,
repository: SignalRepositoryProtocol | None,
completion: DecisionsCall | None,
clock: Clock,
router_ready: RouterReady = lambda: True,
cursor: str = "",
sweep: SignalSweep = SIGNAL_BACKLOG_SWEEP,
) -> SignalTick:
if repository is None or completion is None or not router_ready():
return SignalTick(cursor)
now: Final = clock()
config: Final = await repository.get_config()
if not config.enabled:
return SignalTick(cursor)
reader: Final = SourceReader(storage)
scope: Final = Scope(all_teams=True)
candidates: Final = await _scan_pages(
reader,
repository,
scope,
config,
now,
cursor,
SIGNAL_MAX_PER_TICK,
sweep,
)
executions, next_cursor = candidates
classifier: Final = SignalClassifier(reader, completion, clock)
semaphore: Final = asyncio.Semaphore(SIGNAL_CONCURRENCY)
async def process(execution: Execution) -> bool:
from litellm._logging import verbose_proxy_logger
async with semaphore:
claimed_at: Final = classifier.clock()
claimed_until: Final = claimed_at + SIGNAL_CLAIM_LEASE
try:
claimed: Final = await repository.claim(execution, config, claimed_until, claimed_at)
except Exception as error:
verbose_proxy_logger.error("Lens signal claim failed: %s", redact_internal_details(str(error)))
return False
if not claimed:
return False
await _process_claimed(classifier, repository, scope, execution, config, claimed_until)
return True
outcomes: Final = await asyncio.gather(*(process(execution) for execution in executions))
return SignalTick(next_cursor, sum(outcomes))
async def _logged_tick(
storage: Storage,
repository: SignalRepositoryProtocol | None,
completion: DecisionsCall | None,
clock: Clock,
router_ready: RouterReady,
cursor: str,
sweep: SignalSweep,
) -> SignalTick:
from litellm._logging import verbose_proxy_logger
try:
return await run_signal_tick(storage, repository, completion, clock, router_ready, cursor, sweep)
except Exception as error:
verbose_proxy_logger.error("Lens signal tick failed: %s", redact_internal_details(str(error)))
return SignalTick(cursor)
class _SignalLoopState:
def __init__(self) -> None:
self.tick: SignalTick = SignalTick("")
async def run_signal_loop(
storage: Storage,
repository: SignalRepositoryProtocol | None,
completion: DecisionsCall | None,
clock: Clock = lambda: datetime.now(timezone.utc),
router_ready: RouterReady = lambda: True,
sweep: SignalSweep = SIGNAL_BACKLOG_SWEEP,
) -> None:
state: Final = _SignalLoopState()
while True:
state.tick = await _logged_tick(storage, repository, completion, clock, router_ready, state.tick.cursor, sweep)
await asyncio.sleep(0 if state.tick.claimed >= SIGNAL_MAX_PER_TICK else sweep.interval_seconds)

View file

@ -650,6 +650,7 @@ class _SessionAggRow(LiteLLMBaseModel):
ttl_5m_turns: int = 0
ttl_1h_turns: int = 0
total_tokens: int = 0
day_total_tokens: int | None = None
session_seconds: float = 0.0
turns: int = 0
spend: float = 0.0
@ -724,6 +725,7 @@ def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals:
return AutoRouterBenchmarkTotals(
sessions=sessions,
turns=row.turns,
total_tokens=row.day_total_tokens,
avg_turns_per_session=_per_session(row, row.session_turns),
avg_session_seconds=_per_session(row, row.session_seconds),
avg_tokens_per_session=_per_session(row, row.total_tokens),
@ -759,6 +761,7 @@ def _benchmark_group(row: _SessionAggRow) -> AutoRouterBenchmarkGroup:
tier_turns=row.tier_turns,
sessions=totals.sessions,
turns=totals.turns,
total_tokens=totals.total_tokens,
avg_turns_per_session=totals.avg_turns_per_session,
avg_session_seconds=totals.avg_session_seconds,
avg_tokens_per_session=totals.avg_tokens_per_session,
@ -796,6 +799,11 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow:
ttl_5m_turns=sum(row.ttl_5m_turns for row in rows),
ttl_1h_turns=sum(row.ttl_1h_turns for row in rows),
total_tokens=sum(row.total_tokens for row in rows),
day_total_tokens=(
sum(row.day_total_tokens or 0 for row in rows)
if all(row.day_total_tokens is not None for row in rows)
else None
),
spend=sum(row.spend for row in rows),
saved_spend=sum(row.saved_spend for row in rows),
savings_estimated_turns=sum(row.savings_estimated_turns for row in rows),

View file

@ -4,7 +4,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Optional, Union
from fastapi import HTTPException, status
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter, ValidationError
# Defined above the `litellm.proxy.*` imports so the name is bound even when
@ -150,31 +150,89 @@ def require_caller_user_id_for_non_admin(
return user_api_key_dict.user_id
_ROUTE_LIST: Final = TypeAdapter(list[str] | None)
def _passthrough_routes_permission_error(field: str, entity: str) -> HTTPException:
return HTTPException(
status_code=403,
detail={"error": f"Only proxy admins can set `{field}` on a {entity}."},
)
def _check_passthrough_routes_caller_permission(
data: BaseModel,
data: BaseModel | None,
user_api_key_dict: UserAPIKeyAuth,
*,
entity: str = "key",
existing_metadata: Mapping[str, object] | None = None,
) -> None:
"""
Only proxy admins may set `allowed_passthrough_routes` or `denied_passthrough_routes`
(top-level or under `metadata`), since the runtime route checker reads both from key and
team metadata.
"""
check_allowed_passthrough_routes_caller_permission(data, user_api_key_dict, entity=entity)
check_denied_passthrough_routes_caller_permission(
data, user_api_key_dict, entity=entity, existing_metadata=existing_metadata
)
def check_allowed_passthrough_routes_caller_permission(
data: BaseModel | None,
user_api_key_dict: UserAPIKeyAuth,
*,
entity: str = "key",
) -> None:
"""
Only proxy admins may set `allowed_passthrough_routes` (top-level or under
`metadata`) — it short-circuits the role-based route gate, so keys and teams
must be gated identically.
"""
if data is None:
return
metadata: Final = getattr(data, "metadata", None)
if isinstance(metadata, dict):
try:
_ROUTE_LIST.validate_python(metadata.get("denied_passthrough_routes"))
except ValidationError as e:
raise HTTPException(
status_code=400,
detail={"error": "`metadata.denied_passthrough_routes` must be a list of route strings."},
) from e
# view-only admins excluded by design; blocked upstream from writes anyway
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
if getattr(data, "allowed_passthrough_routes", None):
raise HTTPException(
status_code=403,
detail={"error": f"Only proxy admins can set `allowed_passthrough_routes` on a {entity}."},
)
metadata: Final = getattr(data, "metadata", None)
raise _passthrough_routes_permission_error("allowed_passthrough_routes", entity)
if isinstance(metadata, dict) and metadata.get("allowed_passthrough_routes"):
raise HTTPException(
status_code=403,
detail={"error": f"Only proxy admins can set `metadata.allowed_passthrough_routes` on a {entity}."},
)
raise _passthrough_routes_permission_error("metadata.allowed_passthrough_routes", entity)
def check_denied_passthrough_routes_caller_permission(
data: BaseModel | None,
user_api_key_dict: UserAPIKeyAuth,
*,
entity: str = "key",
existing_metadata: Mapping[str, object] | None = None,
) -> None:
"""
A non-admin request must leave an existing deny list as it is: clearing it, or replacing
`metadata` without it, would widen access. The outcome depends on the stored deny list, so
run this only after the caller is known to be allowed to edit the object.
"""
if data is None or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
metadata: Final = getattr(data, "metadata", None)
existing_denied: Final = (existing_metadata or {}).get("denied_passthrough_routes") or None
if (
"denied_passthrough_routes" in data.model_fields_set
and (getattr(data, "denied_passthrough_routes", None) or None) != existing_denied
):
raise _passthrough_routes_permission_error("denied_passthrough_routes", entity)
if _metadata_changes_denied_routes(data, metadata, existing_denied):
raise _passthrough_routes_permission_error("metadata.denied_passthrough_routes", entity)
def _metadata_changes_denied_routes(data: BaseModel, metadata: object, existing_denied: object) -> bool:
if isinstance(metadata, dict):
return (metadata.get("denied_passthrough_routes") or None) != existing_denied
return metadata is None and "metadata" in data.model_fields_set and existing_denied is not None
def _check_disable_global_guardrails_caller_permission(

View file

@ -119,7 +119,7 @@ if TYPE_CHECKING:
router: Final = APIRouter()
_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig])
_USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE: Final = 50
_USER_BUDGET_CACHE_FIELDS: Final = frozenset({"max_budget", "model_max_budget"})
_USER_LIMIT_CACHE_FIELDS: Final = frozenset({"max_budget", "model_max_budget", "tpm_limit", "rpm_limit"})
def _user_table(
@ -1158,6 +1158,8 @@ async def user_info_v2(
user_role=user_data.get("user_role"),
spend=user_data.get("spend", 0.0),
max_budget=user_data.get("max_budget"),
tpm_limit=user_data.get("tpm_limit"),
rpm_limit=user_data.get("rpm_limit"),
models=user_data.get("models") or [],
budget_duration=user_data.get("budget_duration"),
budget_reset_at=user_data.get("budget_reset_at"),
@ -1298,7 +1300,7 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda
fields_set: Final = data.fields_set() if hasattr(data, "fields_set") else set()
for k, v in data_json.items():
if k in ("max_budget", "budget_duration"):
if k in ("max_budget", "budget_duration", "tpm_limit", "rpm_limit"):
if k in fields_set:
non_default_values[k] = v
elif k == "model_max_budget":
@ -1627,7 +1629,7 @@ async def _update_single_user_helper(
await _invalidate_user_spend_counter_if_changed(non_default_values)
if not _USER_BUDGET_CACHE_FIELDS.isdisjoint(non_default_values) or "metadata" in data_json:
if not _USER_LIMIT_CACHE_FIELDS.isdisjoint(non_default_values) or "metadata" in data_json:
await evict_and_broadcast(
cache_keys=(non_default_values["user_id"],),
user_api_key_cache=user_api_key_cache,
@ -1985,7 +1987,7 @@ async def bulk_user_update(
),
)
if not _USER_BUDGET_CACHE_FIELDS.isdisjoint(non_default_values):
if not _USER_LIMIT_CACHE_FIELDS.isdisjoint(non_default_values):
for start in range(0, len(all_users_in_db), _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE):
await asyncio.gather(
*(

View file

@ -94,6 +94,8 @@ from litellm.proxy.management_endpoints.common_utils import (
_set_object_metadata_field,
_team_member_has_permission,
_user_has_admin_view,
check_allowed_passthrough_routes_caller_permission,
check_denied_passthrough_routes_caller_permission,
validate_budget_duration,
validate_finite_spend,
)
@ -2012,6 +2014,7 @@ async def generate_key_fn(
- 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.
- denied_passthrough_routes: Optional[list] - List of pass through routes the key may not call, even if allowed by `allowed_passthrough_routes` or `allowed_routes`. Matches exact paths, path prefixes, and trailing `*` wildcards. Applies together with the team's `denied_passthrough_routes`. Example - ["/my-custom-endpoint/admin"].
- 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.
- key_type: Optional[str] - Type of key that determines default allowed routes. Options: "llm_api" (can call LLM API routes), "management" (can call management routes), "read_only" (can only call info/read routes), "default" (uses default allowed routes). Defaults to "default".
- prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts.
@ -2774,6 +2777,11 @@ async def _process_single_key_update(
existing_key_row=existing_key_row,
user_api_key_cache=user_api_key_cache,
)
check_denied_passthrough_routes_caller_permission(
update_key_request,
user_api_key_dict,
existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict
)
# Custom key update hook
if user_custom_key_update is not None:
@ -3091,10 +3099,7 @@ async def _validate_update_key_data(
existing_key_row=existing_key_row,
user_api_key_dict=user_api_key_dict,
)
_check_passthrough_routes_caller_permission(
data=data,
user_api_key_dict=user_api_key_dict,
)
check_allowed_passthrough_routes_caller_permission(data, user_api_key_dict)
_check_permissions_caller_permission(
data=data,
user_api_key_dict=user_api_key_dict,
@ -3240,6 +3245,11 @@ async def _validate_update_key_data(
user_api_key_cache=user_api_key_cache,
route=("/key/update (max_budget/spend)" if _is_budget_change else "/key/update"),
)
check_denied_passthrough_routes_caller_permission(
data,
user_api_key_dict,
existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict
)
# Check team limits if key has a team_id (from request or existing key)
team_obj: LiteLLM_TeamTableCachedObj | None = None
@ -3428,6 +3438,7 @@ async def update_key_fn(
- temp_budget_expiry: Optional[str] - Expiry time for the temporary budget increase (Enterprise only).
- 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 routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through routes the key can access, without specifying the routes. If allowed_routes is specified, allowed_passthrough_routes is ignored.
- denied_passthrough_routes: Optional[list] - List of pass through routes the key may not call, even if allowed by `allowed_passthrough_routes` or `allowed_routes`. Matches exact paths, path prefixes, and trailing `*` wildcards. Applies together with the team's `denied_passthrough_routes`. Example - ["/my-custom-endpoint/admin"].
- prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts.
- 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.
- auto_rotate: Optional[bool] - Whether this key should be automatically rotated
@ -3944,7 +3955,7 @@ async def bulk_update_team_keys(
# Block metadata.allowed_passthrough_routes for non-admins — the runtime
# route checker reads it from key/team metadata to grant passthrough.
_check_passthrough_routes_caller_permission(data=data.update_fields, user_api_key_dict=user_api_key_dict)
check_allowed_passthrough_routes_caller_permission(data.update_fields, user_api_key_dict)
if not requested_tokens:
raise HTTPException(
@ -5824,10 +5835,7 @@ async def regenerate_key_fn(
user_api_key_dict=user_api_key_dict,
allowed_routes_was_provided="allowed_routes" in data.model_fields_set,
)
_check_passthrough_routes_caller_permission(
data=data,
user_api_key_dict=user_api_key_dict,
)
check_allowed_passthrough_routes_caller_permission(data, user_api_key_dict)
_check_permissions_caller_permission(
data=data,
user_api_key_dict=user_api_key_dict,
@ -5945,6 +5953,11 @@ async def regenerate_key_fn(
status_code=status.HTTP_403_FORBIDDEN,
detail={"error": "You are not authorized to regenerate this key"},
)
check_denied_passthrough_routes_caller_permission(
data,
user_api_key_dict,
existing_metadata=_key_in_db.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict
)
if data is not None and (data.access_group_ids or data.object_permission is not None):
regenerate_team_table: LiteLLM_TeamTableCachedObj | None = None

View file

@ -179,6 +179,7 @@ if MCP_AVAILABLE:
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_raise_if_not_oauth2,
authorize_with_server,
client_supplied_application_type,
client_supplied_redirect_uris,
exchange_token_with_server,
get_request_base_url,
@ -1047,6 +1048,7 @@ if MCP_AVAILABLE:
available_on_public_internet=payload.available_on_public_internet,
timeout=payload.timeout,
max_concurrent_requests=payload.max_concurrent_requests,
rpm=payload.rpm,
)
def get_prisma_client_or_throw(message: str):
@ -2426,6 +2428,7 @@ if MCP_AVAILABLE:
request_data: Final = await _read_request_body(request=request)
data: Final[Mapping[str, object]] = {**request_data}
client_redirect_uris: Final = client_supplied_redirect_uris(data.get("redirect_uris"))
client_application_type: Final = client_supplied_application_type(data.get("application_type"))
return await register_client_with_server(
request=request,
@ -2437,6 +2440,7 @@ if MCP_AVAILABLE:
fallback_client_id=server_id,
persist_credentials=_user_is_full_admin(user_api_key_dict),
client_redirect_uris=client_redirect_uris,
client_application_type=client_application_type,
)
@router.delete(

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