mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
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:
commit
18a818530f
593 changed files with 83990 additions and 53361 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
14
.github/codeql/codeql-config.yml
vendored
14
.github/codeql/codeql-config.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
9
.github/workflows/codeql.yml
vendored
9
.github/workflows/codeql.yml
vendored
|
|
@ -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 }}"
|
||||
|
|
|
|||
14
.github/workflows/test-linting.yml
vendored
14
.github/workflows/test-linting.yml
vendored
|
|
@ -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'
|
||||
|
|
|
|||
4
Makefile
4
Makefile
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "rpm" INTEGER;
|
||||
|
|
@ -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 $$;
|
||||
|
|
@ -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")
|
||||
);
|
||||
|
|
@ -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])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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!(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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}))
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")]
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(¶ms).unwrap(),
|
||||
optional_object("optional_params", ¶ms).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
|
||||
);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,))?;
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
},
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
33
litellm-rust/crates/traces-clickhouse/query/trace_agents.sql
Normal file
33
litellm-rust/crates/traces-clickhouse/query/trace_agents.sql
Normal 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}
|
||||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
|
|||
|
|
@ -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>(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()),
|
||||
])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
¶meters,
|
||||
)
|
||||
.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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ pub mod named;
|
|||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum ReadQuery {
|
||||
ListTraces,
|
||||
TraceAgents,
|
||||
TraceSpans,
|
||||
TracePageSpans,
|
||||
TraceIdentity,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
|
|||
|
|
@ -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"}}),
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
||||
|
|
|
|||
|
|
@ -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 -
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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'",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
142
litellm/proxy/lens/signal_repository.py
Normal file
142
litellm/proxy/lens/signal_repository.py
Normal 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,
|
||||
)
|
||||
619
litellm/proxy/lens/signals.py
Normal file
619
litellm/proxy/lens/signals.py
Normal 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)
|
||||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
*(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Add table
Reference in a new issue