mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
Merge remote-tracking branch 'origin/main' into litellm_langfuse_sdk_v4
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
This commit is contained in:
commit
879d8b27f0
223 changed files with 22161 additions and 1673 deletions
11
.github/workflows/duplicate_issue_check.yml
vendored
11
.github/workflows/duplicate_issue_check.yml
vendored
|
|
@ -93,12 +93,13 @@ jobs:
|
|||
responses-api-endpoint: ${{ vars.LITELLM_API_BASE }}/v1/responses
|
||||
prompt-file: .github/prompts/duplicate-issue-check.md
|
||||
output-schema-file: .github/prompts/duplicate-issue-check.schema.json
|
||||
sandbox: read-only
|
||||
# read-only denies network, and the whole method is searching the tracker with gh
|
||||
codex-args: '["-c", "sandbox_permissions=[\"network-full-access\"]"]'
|
||||
sandbox: workspace-write
|
||||
# The whole method is searching the tracker with gh, and network is only switchable in workspace-write
|
||||
codex-args: '["-c", "sandbox_workspace_write.network_access=true"]'
|
||||
model: ${{ vars.DUPLICATE_CHECK_MODEL }}
|
||||
# Issue authors have no write access and the action refuses them by default; the
|
||||
# prompt is fixed, the sandbox read-only, and the only token is read-only on a public repo
|
||||
codex-version: "0.154.0"
|
||||
# Issue authors have no write access and the action refuses them by default; the prompt is
|
||||
# fixed, writes stay inside the throwaway checkout, and the only token is read-only on a public repo
|
||||
allow-users: "*"
|
||||
|
||||
- name: Summary
|
||||
|
|
|
|||
3
.github/workflows/osv-scan.yml
vendored
3
.github/workflows/osv-scan.yml
vendored
|
|
@ -41,4 +41,5 @@ jobs:
|
|||
"$RUNNER_TEMP/osv-scanner" scan source \
|
||||
--config osv-scanner.toml \
|
||||
-L uv.lock \
|
||||
-L ui/litellm-dashboard/package-lock.json
|
||||
-L ui/litellm-dashboard/package-lock.json \
|
||||
-L vscode-extension/package-lock.json
|
||||
|
|
|
|||
65
.github/workflows/test-vscode-extension.yml
vendored
Normal file
65
.github/workflows/test-vscode-extension.yml
vendored
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
name: VS Code Extension
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
paths:
|
||||
- "vscode-extension/**"
|
||||
- ".github/workflows/test-vscode-extension.yml"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths:
|
||||
- "vscode-extension/**"
|
||||
- ".github/workflows/test-vscode-extension.yml"
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
vscode-extension:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
defaults:
|
||||
run:
|
||||
working-directory: vscode-extension
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
fetch-depth: 1
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
with:
|
||||
node-version: "24"
|
||||
cache: npm
|
||||
cache-dependency-path: vscode-extension/package-lock.json
|
||||
|
||||
- name: Install dependencies
|
||||
run: npm ci
|
||||
|
||||
- name: Typecheck
|
||||
run: npm run typecheck
|
||||
|
||||
- name: Unit tests
|
||||
run: npm test
|
||||
|
||||
- name: Package extension
|
||||
run: npm run package
|
||||
|
||||
- name: Upload VSIX
|
||||
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
|
||||
with:
|
||||
name: litellm-vscode
|
||||
path: vscode-extension/*.vsix
|
||||
if-no-files-found: error
|
||||
|
|
@ -82,6 +82,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/anthropic/",
|
||||
"/azure/",
|
||||
"/azure_ai/",
|
||||
"/azure_speech/",
|
||||
"/aws/",
|
||||
"/bedrock/",
|
||||
"/comprehendmedical",
|
||||
|
|
@ -94,6 +95,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/vertex-ai/",
|
||||
"/assemblyai/",
|
||||
"/eu.assemblyai/",
|
||||
"/deepgram/",
|
||||
"/langfuse/",
|
||||
"/vllm/",
|
||||
"/mistral/",
|
||||
|
|
|
|||
|
|
@ -66,7 +66,7 @@
|
|||
"/v1/video" "/v1/videos" "/video" "/videos" "/v1/search" "/search"
|
||||
"/v1/containers" "/containers" "/v1/evals" "/v1/memory" "/queue/chat"
|
||||
"/v1beta" "/interactions"
|
||||
"/anthropic" "/azure" "/azure_ai" "/aws" "/bedrock" "/comprehendmedical" "/transcribe" "/cohere" "/gemini" "/google"
|
||||
"/anthropic" "/azure" "/azure_ai" "/azure_speech" "/aws" "/bedrock" "/comprehendmedical" "/transcribe" "/cohere" "/gemini" "/google"
|
||||
"/vertex_ai" "/vertex-ai" "/assemblyai" "/eu.assemblyai" "/langfuse" "/vllm"
|
||||
"/mistral" "/groq" "/voyage" "/cursor" "/milvus" "/openai_passthrough"
|
||||
"/toolset"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,35 @@
|
|||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_DailyGlobalSpend" (
|
||||
"id" TEXT NOT NULL,
|
||||
"date" TEXT NOT NULL,
|
||||
"model" TEXT,
|
||||
"model_group" TEXT,
|
||||
"custom_llm_provider" TEXT,
|
||||
"mcp_namespaced_tool_name" TEXT,
|
||||
"endpoint" TEXT,
|
||||
"prompt_tokens" BIGINT NOT NULL DEFAULT 0,
|
||||
"completion_tokens" BIGINT NOT NULL DEFAULT 0,
|
||||
"cache_read_input_tokens" BIGINT NOT NULL DEFAULT 0,
|
||||
"cache_creation_input_tokens" BIGINT NOT NULL DEFAULT 0,
|
||||
"compression_saved_tokens" BIGINT NOT NULL DEFAULT 0,
|
||||
"compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
||||
"prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
||||
"gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
||||
"autorouter_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
||||
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
||||
"api_requests" BIGINT NOT NULL DEFAULT 0,
|
||||
"successful_requests" BIGINT NOT NULL DEFAULT 0,
|
||||
"failed_requests" BIGINT NOT NULL DEFAULT 0,
|
||||
"total_response_time_ms" BIGINT NOT NULL DEFAULT 0,
|
||||
"timed_requests" BIGINT NOT NULL DEFAULT 0,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL,
|
||||
|
||||
CONSTRAINT "LiteLLM_DailyGlobalSpend_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_DailyGlobalSpend_date_idx" ON "LiteLLM_DailyGlobalSpend"("date");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyGlobalSpend_date_model_model_group_custom_llm__key" ON "LiteLLM_DailyGlobalSpend"("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
|
||||
|
|
@ -820,6 +820,37 @@ model LiteLLM_DailyUserSpend {
|
|||
@@index([endpoint])
|
||||
}
|
||||
|
||||
// Key-free daily rollup of LiteLLM_DailyUserSpend, read by the global usage view
|
||||
model LiteLLM_DailyGlobalSpend {
|
||||
id String @id @default(uuid())
|
||||
date String
|
||||
model String?
|
||||
model_group String?
|
||||
custom_llm_provider String?
|
||||
mcp_namespaced_tool_name String?
|
||||
endpoint String?
|
||||
prompt_tokens BigInt @default(0)
|
||||
completion_tokens BigInt @default(0)
|
||||
cache_read_input_tokens BigInt @default(0)
|
||||
cache_creation_input_tokens BigInt @default(0)
|
||||
compression_saved_tokens BigInt @default(0)
|
||||
compression_savings_spend Float @default(0.0)
|
||||
prompt_caching_savings_spend Float @default(0.0)
|
||||
gateway_injected_caching_savings_spend Float @default(0.0)
|
||||
autorouter_savings_spend Float @default(0.0)
|
||||
spend Float @default(0.0)
|
||||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([date, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
}
|
||||
|
||||
// Track daily organization spend metrics per model and key
|
||||
model LiteLLM_DailyOrganizationSpend {
|
||||
id String @id @default(uuid())
|
||||
|
|
|
|||
28
litellm-rust/Cargo.lock
generated
28
litellm-rust/Cargo.lock
generated
|
|
@ -559,6 +559,21 @@ dependencies = [
|
|||
"vsimd",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bit-set"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3"
|
||||
dependencies = [
|
||||
"bit-vec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bit-vec"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7"
|
||||
|
||||
[[package]]
|
||||
name = "bitflags"
|
||||
version = "2.13.1"
|
||||
|
|
@ -1166,6 +1181,17 @@ dependencies = [
|
|||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fancy-regex"
|
||||
version = "0.19.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d301f5bf187b3c295fce6468d3875037a0bccc5f6b151c63cac2f85babf21912"
|
||||
dependencies = [
|
||||
"bit-set",
|
||||
"regex-automata",
|
||||
"regex-syntax",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fastrand"
|
||||
version = "2.5.0"
|
||||
|
|
@ -2059,7 +2085,9 @@ dependencies = [
|
|||
name = "litellm-core-utils"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"fancy-regex",
|
||||
"litellm-types",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_path_to_error",
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ strum = { version = "0.28.0", features = ["derive"] }
|
|||
url = "2.5.8"
|
||||
time = { version = "0.3.53", features = ["parsing"] }
|
||||
criterion = "0.8.2"
|
||||
fancy-regex = "0.19.2"
|
||||
veil = "0.3.0"
|
||||
|
||||
[profile.release]
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
fancy-regex.workspace = true
|
||||
litellm-types.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
|
@ -13,3 +14,6 @@ serde_path_to_error = "0.1"
|
|||
serde_with.workspace = true
|
||||
thiserror.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
|
|
|
|||
|
|
@ -0,0 +1,115 @@
|
|||
use super::public::PublicError;
|
||||
use super::rules::{Rule, contains_any};
|
||||
|
||||
/// The text branches of `_map_cohere_exception`, in its order.
|
||||
pub(super) const RULES: &[Rule] = &[
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
contains_any(
|
||||
&mapping.error_str,
|
||||
&["invalid api token", "No API key provided."],
|
||||
)
|
||||
},
|
||||
PublicError::Authentication,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| mapping.error_str.contains("invalid type: parameter"),
|
||||
PublicError::BadRequest,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| mapping.error_str.contains("too many tokens"),
|
||||
PublicError::ContextWindowExceeded,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
mapping
|
||||
.error_str
|
||||
.to_lowercase()
|
||||
.contains("internal server error")
|
||||
},
|
||||
PublicError::InternalServer,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| mapping.status.is_none() && mapping.error_str.contains("invalid type:"),
|
||||
PublicError::BadRequest,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| mapping.status.is_none() && mapping.error_str.contains("Unexpected server error"),
|
||||
PublicError::InternalServer,
|
||||
),
|
||||
];
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::rules::first_match;
|
||||
use super::super::testing::mapping;
|
||||
use super::*;
|
||||
|
||||
fn classified(text: &str) -> Option<PublicError> {
|
||||
classified_with(Some(400), text)
|
||||
}
|
||||
|
||||
fn classified_with(status: Option<u16>, text: &str) -> Option<PublicError> {
|
||||
first_match(RULES, &mapping(status, text)).map(|rule| rule.error)
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::invalid_token("invalid api token", PublicError::Authentication)]
|
||||
#[case::no_api_key("No API key provided.", PublicError::Authentication)]
|
||||
#[case::invalid_parameter("invalid type: parameter x", PublicError::BadRequest)]
|
||||
#[case::too_many_tokens("too many tokens", PublicError::ContextWindowExceeded)]
|
||||
#[case::internal_server_text("Internal Server Error", PublicError::InternalServer)]
|
||||
#[case::internal_server_any_case("INTERNAL server ERROR", PublicError::InternalServer)]
|
||||
fn each_text_rule_claims_its_marker(#[case] text: &str, #[case] expected: PublicError) {
|
||||
assert_eq!(classified(text), Some(expected));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::token_before_parameter(
|
||||
"invalid api token invalid type: parameter",
|
||||
PublicError::Authentication
|
||||
)]
|
||||
#[case::parameter_before_tokens(
|
||||
"invalid type: parameter too many tokens",
|
||||
PublicError::BadRequest
|
||||
)]
|
||||
#[case::tokens_before_internal(
|
||||
"too many tokens Internal Server Error",
|
||||
PublicError::ContextWindowExceeded
|
||||
)]
|
||||
fn the_earlier_rule_wins_when_two_apply(#[case] text: &str, #[case] expected: PublicError) {
|
||||
assert_eq!(classified(text), Some(expected));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::invalid_type(None, "invalid type: x", Some(PublicError::BadRequest))]
|
||||
#[case::unexpected_server_error(
|
||||
None,
|
||||
"Unexpected server error",
|
||||
Some(PublicError::InternalServer)
|
||||
)]
|
||||
#[case::invalid_type_before_unexpected(
|
||||
None,
|
||||
"invalid type: x Unexpected server error",
|
||||
Some(PublicError::BadRequest)
|
||||
)]
|
||||
#[case::internal_before_invalid_type(
|
||||
None,
|
||||
"internal server error invalid type: x",
|
||||
Some(PublicError::InternalServer)
|
||||
)]
|
||||
#[case::invalid_type_with_a_status(Some(500), "invalid type: x", None)]
|
||||
#[case::unexpected_with_a_status(Some(400), "Unexpected server error", None)]
|
||||
fn the_trailing_rules_only_claim_failures_without_a_status(
|
||||
#[case] status: Option<u16>,
|
||||
#[case] text: &str,
|
||||
#[case] expected: Option<PublicError>,
|
||||
) {
|
||||
assert_eq!(classified_with(status, text), expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_without_a_marker_is_left_to_the_status_table() {
|
||||
assert_eq!(classified("rejected"), None);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,542 @@
|
|||
//! A port of Python's `exception_type` for the routes that run in Rust. Rust decides the
|
||||
//! public class, the message and the debug text; Python only builds the class.
|
||||
//!
|
||||
//! DIVERGENCES: where the Python mapper is inconsistent, the port follows one rule instead.
|
||||
//! - The message is always `{Provider}Exception - {redacted text}`. Python's per-branch
|
||||
//! labels (`RateLimitError: `, `litellm.RateLimitError: `, `Vertex_aiException BadRequestError`)
|
||||
//! are dropped because every public class already prefixes `litellm.{Class}: `.
|
||||
//! - The upstream response is always the real one. Python swaps in made-up `httpx.Response`
|
||||
//! stubs on some Vertex branches, losing the body and `retry-after`.
|
||||
//! - The debug text is always attached; Python passes it on some branches only.
|
||||
//! - No family rule turns a status into a class; the shared status table owns that. So a
|
||||
//! Vertex 502 is a `BadGatewayError` and an OpenAI-family 403 is a `PermissionDeniedError`.
|
||||
//! Three rules read the status only to gate a text match, as Python does: the standalone
|
||||
//! `429`, Vertex's wrapped 429 behind a 5xx, and Cohere's rules for failures with no status.
|
||||
//! - A timeout text marker on an HTTP failure keeps the upstream response. Python's `Timeout`
|
||||
//! carries none.
|
||||
//! - Every family matches and reports the redacted text. Python's OpenAI mapper builds the
|
||||
//! message from the unredacted text.
|
||||
//! - A refused connection is an `APIConnectionError`, not the 500 Python's HTTP handler
|
||||
//! synthesizes.
|
||||
//! - Dropped Python rules: Vertex's bare `403` substring (it matches `4031 tokens`), Vertex's
|
||||
//! `IndexError` quota marker (a Python client crash), the OpenAI SDK's missing-`api_key`
|
||||
//! text and its `OPENAI` renaming, Cohere's `llm_provider="cohere"` override, and Cohere's
|
||||
//! `CohereConnectionError` check (a Python SDK class name).
|
||||
//!
|
||||
//! KNOWN_GAPS: differences from the Python mapper that no Rust route can reach today. Each
|
||||
//! one stops being acceptable at its trigger.
|
||||
//! - The Vertex partner-model API base for "claude" models is not built into the debug text.
|
||||
//! Trigger: a Vertex route whose models include Anthropic partner models.
|
||||
//! - The debug text's `API Base` line is only the non-streaming Vertex URL. Python prefers an
|
||||
//! explicit or provider-resolved `api_base`, uses `:streamGenerateContent` when streaming,
|
||||
//! and has Gemini and OpenAI defaults. Trigger: the first route wired to this mapper, since
|
||||
//! every route knows its `api_base`.
|
||||
//! - The debug text has no `Messages:` line, which Python adds when
|
||||
//! `redact_messages_in_exceptions` is off. Trigger: a wired route that carries messages.
|
||||
//! - Python reports the provider `get_llm_provider` resolves for a stripped model name when
|
||||
//! that name happens to be in the model cost map. Trigger: a route whose model names
|
||||
//! overlap the cost map; that needs the provider resolution port, not a classifier change.
|
||||
//! - `litellm_proxy` errors are not unwrapped into the proxied exception. Trigger: a Rust
|
||||
//! route that calls a LiteLLM proxy.
|
||||
//! - Only the OpenAI-compatible, Vertex AI and Cohere mappers are ported; every other
|
||||
//! provider goes straight to the status table. Trigger: a Rust route for such a provider.
|
||||
|
||||
use super::secret_redaction::SecretRedactor;
|
||||
|
||||
mod cohere;
|
||||
mod openai;
|
||||
mod original;
|
||||
mod public;
|
||||
mod rules;
|
||||
mod status;
|
||||
mod vertex_ai;
|
||||
|
||||
pub use original::{ExceptionFamily, OriginalException};
|
||||
pub use public::{MappedFailure, PublicError, UpstreamResponse};
|
||||
|
||||
use rules::{Rule, contains_any, first_match};
|
||||
|
||||
const TIMEOUT_MARKERS: &[&str] = &[
|
||||
"Request Timeout Error",
|
||||
"Request timed out",
|
||||
"Timed out generating response",
|
||||
"The read operation timed out",
|
||||
];
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct ExceptionContext {
|
||||
pub model: String,
|
||||
pub custom_llm_provider: String,
|
||||
pub asynchronous: bool,
|
||||
pub vertex_project: Option<String>,
|
||||
pub vertex_location: Option<String>,
|
||||
pub model_group: Option<String>,
|
||||
pub deployment: Option<String>,
|
||||
pub user_api_key_alias: Option<String>,
|
||||
pub user_api_key_team_alias: Option<String>,
|
||||
}
|
||||
|
||||
/// What the rules read: the status of a provider response, if any, and the redacted text.
|
||||
struct Mapping {
|
||||
status: Option<u16>,
|
||||
error_str: String,
|
||||
}
|
||||
|
||||
pub fn exception_type(
|
||||
context: &ExceptionContext,
|
||||
redactor: Option<&SecretRedactor>,
|
||||
original: &OriginalException,
|
||||
) -> MappedFailure {
|
||||
let (status, text, upstream) = match original {
|
||||
OriginalException::Http {
|
||||
status,
|
||||
body,
|
||||
headers,
|
||||
} => (
|
||||
Some(*status),
|
||||
body.clone(),
|
||||
Some(UpstreamResponse {
|
||||
status: *status,
|
||||
body: body.clone(),
|
||||
headers: headers.clone(),
|
||||
}),
|
||||
),
|
||||
OriginalException::Connection { message } | OriginalException::Plain { message } => {
|
||||
(None, message.clone(), None)
|
||||
}
|
||||
OriginalException::Timeout {
|
||||
timeout_seconds,
|
||||
elapsed_seconds,
|
||||
} => (
|
||||
None,
|
||||
timeout_message(context.asynchronous, *timeout_seconds, *elapsed_seconds),
|
||||
None,
|
||||
),
|
||||
};
|
||||
let mapping = Mapping {
|
||||
status,
|
||||
error_str: match redactor {
|
||||
Some(redactor) => redactor.redact(&text),
|
||||
None => text,
|
||||
},
|
||||
};
|
||||
let family = ExceptionFamily::for_provider(&context.custom_llm_provider);
|
||||
let (error, hint) = classify(family, original, &mapping);
|
||||
MappedFailure {
|
||||
error,
|
||||
message: format!(
|
||||
"{} - {}{hint}",
|
||||
exception_provider(&context.custom_llm_provider),
|
||||
mapping.error_str
|
||||
),
|
||||
upstream,
|
||||
debug_info: extra_information(context, api_base(context).as_deref()),
|
||||
}
|
||||
}
|
||||
|
||||
fn classify(
|
||||
family: ExceptionFamily,
|
||||
original: &OriginalException,
|
||||
mapping: &Mapping,
|
||||
) -> (PublicError, &'static str) {
|
||||
const TIMEOUT: PublicError = PublicError::Timeout { status: 408 };
|
||||
if matches!(original, OriginalException::Timeout { .. })
|
||||
|| contains_any(&mapping.error_str, TIMEOUT_MARKERS)
|
||||
{
|
||||
return (TIMEOUT, "");
|
||||
}
|
||||
if let Some(rule) = first_match(family_rules(family), mapping) {
|
||||
return (rule.error, rule.hint);
|
||||
}
|
||||
let by_status = mapping.status.and_then(status::classify);
|
||||
(by_status.unwrap_or(PublicError::ApiConnection), "")
|
||||
}
|
||||
|
||||
fn family_rules(family: ExceptionFamily) -> &'static [Rule] {
|
||||
match family {
|
||||
ExceptionFamily::OpenAiCompatible => openai::RULES,
|
||||
ExceptionFamily::VertexAi => vertex_ai::RULES,
|
||||
ExceptionFamily::Cohere => cohere::RULES,
|
||||
ExceptionFamily::Other => &[],
|
||||
}
|
||||
}
|
||||
|
||||
/// The text the Python HTTP handler's timeout carries: the sync and async handlers word it
|
||||
/// differently.
|
||||
fn timeout_message(
|
||||
asynchronous: bool,
|
||||
timeout_seconds: Option<f64>,
|
||||
elapsed_seconds: Option<f64>,
|
||||
) -> String {
|
||||
let timeout = python_float(timeout_seconds);
|
||||
if asynchronous {
|
||||
let elapsed =
|
||||
python_float(elapsed_seconds.map(|seconds| (seconds * 1000.0).round() / 1000.0));
|
||||
format!("Connection timed out. Timeout passed={timeout}, time taken={elapsed} seconds")
|
||||
} else {
|
||||
format!("Connection timed out after {timeout} seconds.")
|
||||
}
|
||||
}
|
||||
|
||||
fn python_float(value: Option<f64>) -> String {
|
||||
match value {
|
||||
None => "None".to_string(),
|
||||
Some(value) if value.fract() == 0.0 => format!("{value:.1}"),
|
||||
Some(value) => value.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn exception_provider(provider: &str) -> String {
|
||||
if provider == "openai" {
|
||||
return "OpenAIException".to_string();
|
||||
}
|
||||
let mut characters = provider.chars();
|
||||
match characters.next() {
|
||||
Some(first) => format!("{}{}Exception", first.to_uppercase(), characters.as_str()),
|
||||
None => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn api_base(context: &ExceptionContext) -> Option<String> {
|
||||
match (&context.vertex_location, &context.vertex_project) {
|
||||
(Some(location), Some(project)) => Some(format!(
|
||||
"{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/publishers/google/models/{}:generateContent",
|
||||
context.model
|
||||
)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn extra_information(context: &ExceptionContext, api_base: Option<&str>) -> String {
|
||||
let lines = [
|
||||
Some(format!("\nModel: {}", context.model)),
|
||||
api_base.map(|api_base| format!("\nAPI Base: `{api_base}`")),
|
||||
context
|
||||
.model_group
|
||||
.as_ref()
|
||||
.map(|value| format!("\nmodel_group: `{value}`\n")),
|
||||
context
|
||||
.deployment
|
||||
.as_ref()
|
||||
.map(|value| format!("\ndeployment: `{value}`\n")),
|
||||
context
|
||||
.vertex_project
|
||||
.as_ref()
|
||||
.map(|value| format!("\nvertex_project: `{value}`\n")),
|
||||
context
|
||||
.vertex_location
|
||||
.as_ref()
|
||||
.map(|value| format!("\nvertex_location: `{value}`\n")),
|
||||
];
|
||||
let information: String = lines.into_iter().flatten().collect();
|
||||
match &context.user_api_key_alias {
|
||||
Some(alias) => format!(
|
||||
"\n\nKey Name: `{alias}`\nTeam: `{}`{information}",
|
||||
context.user_api_key_team_alias.as_deref().unwrap_or("None")
|
||||
),
|
||||
None => information,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod testing {
|
||||
use super::Mapping;
|
||||
|
||||
pub(super) fn mapping(status: Option<u16>, text: &str) -> Mapping {
|
||||
Mapping {
|
||||
status,
|
||||
error_str: text.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const DEBUG: &str = "\nModel: ocr-model";
|
||||
|
||||
fn context(provider: &str) -> ExceptionContext {
|
||||
ExceptionContext {
|
||||
model: "ocr-model".into(),
|
||||
custom_llm_provider: provider.into(),
|
||||
..ExceptionContext::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn redactor() -> SecretRedactor {
|
||||
SecretRedactor::new(16)
|
||||
}
|
||||
|
||||
fn headers() -> Vec<(String, String)> {
|
||||
vec![("retry-after".into(), "7".into())]
|
||||
}
|
||||
|
||||
fn http(status: u16, body: &str) -> OriginalException {
|
||||
OriginalException::Http {
|
||||
status,
|
||||
body: body.into(),
|
||||
headers: headers(),
|
||||
}
|
||||
}
|
||||
|
||||
fn upstream(status: u16, body: &str) -> Option<UpstreamResponse> {
|
||||
Some(UpstreamResponse {
|
||||
status,
|
||||
body: body.into(),
|
||||
headers: headers(),
|
||||
})
|
||||
}
|
||||
|
||||
fn mapped(provider: &str, original: &OriginalException) -> MappedFailure {
|
||||
exception_type(&context(provider), Some(&redactor()), original)
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::openai_family("mistral", "rate limit reached", PublicError::RateLimit)]
|
||||
#[case::vertex_family("vertex_ai", "Resource exhausted", PublicError::RateLimit)]
|
||||
#[case::cohere_family("cohere", "too many tokens", PublicError::ContextWindowExceeded)]
|
||||
fn a_family_text_rule_beats_the_status_and_keeps_the_real_response(
|
||||
#[case] provider: &str,
|
||||
#[case] body: &str,
|
||||
#[case] expected: PublicError,
|
||||
) {
|
||||
let failure = mapped(provider, &http(401, body));
|
||||
assert_eq!(failure.error, expected);
|
||||
assert_eq!(failure.upstream, upstream(401, body));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_other_family_has_no_text_rules() {
|
||||
assert_eq!(
|
||||
mapped("reducto", &http(401, "rate limit reached")).error,
|
||||
PublicError::Authentication
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::openai_403_is_permission_denied("mistral", 403, PublicError::PermissionDenied)]
|
||||
#[case::openai_409_is_bad_request("mistral", 409, PublicError::BadRequest)]
|
||||
#[case::vertex_502_is_bad_gateway("vertex_ai", 502, PublicError::BadGateway)]
|
||||
#[case::vertex_504_is_a_timeout("vertex_ai", 504, PublicError::Timeout { status: 504 })]
|
||||
#[case::cohere_498_is_bad_request("cohere", 498, PublicError::BadRequest)]
|
||||
#[case::other_503("reducto", 503, PublicError::ServiceUnavailable)]
|
||||
fn without_a_text_rule_every_family_uses_the_status_table(
|
||||
#[case] provider: &str,
|
||||
#[case] status: u16,
|
||||
#[case] expected: PublicError,
|
||||
) {
|
||||
assert_eq!(
|
||||
mapped(provider, &http(status, "rejected")),
|
||||
MappedFailure {
|
||||
error: expected,
|
||||
message: format!("{} - rejected", exception_provider(provider)),
|
||||
upstream: upstream(status, "rejected"),
|
||||
debug_info: DEBUG.into(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::request_timeout_error("Request Timeout Error")]
|
||||
#[case::request_timed_out("Request timed out")]
|
||||
#[case::timed_out_generating("Timed out generating response")]
|
||||
#[case::read_operation("The read operation timed out")]
|
||||
fn timeout_markers_win_over_every_family(#[case] marker: &str) {
|
||||
let body = format!("rate limit invalid api token {marker}");
|
||||
for provider in ["mistral", "vertex_ai", "cohere", "reducto"] {
|
||||
assert_eq!(
|
||||
mapped(provider, &http(429, &body)).error,
|
||||
PublicError::Timeout { status: 408 },
|
||||
"{provider}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_handler_timeout_is_a_408_without_a_response() {
|
||||
let original = OriginalException::Timeout {
|
||||
timeout_seconds: Some(0.5),
|
||||
elapsed_seconds: Some(0.5031),
|
||||
};
|
||||
assert_eq!(
|
||||
mapped("mistral", &original),
|
||||
MappedFailure {
|
||||
error: PublicError::Timeout { status: 408 },
|
||||
message: "MistralException - Connection timed out after 0.5 seconds.".into(),
|
||||
upstream: None,
|
||||
debug_info: DEBUG.into(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::refused_connection(OriginalException::Connection { message: "refused".into() })]
|
||||
#[case::unparseable_response(OriginalException::Plain { message: "refused".into() })]
|
||||
#[case::informational_status(OriginalException::Http { status: 399, body: "refused".into(), headers: Vec::new() })]
|
||||
fn a_failure_no_rule_or_status_claims_is_a_connection_error(
|
||||
#[case] original: OriginalException,
|
||||
) {
|
||||
let failure = mapped("reducto", &original);
|
||||
assert_eq!(failure.error, PublicError::ApiConnection);
|
||||
assert_eq!(failure.message, "ReductoException - refused");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_timeout_marker_on_a_response_keeps_the_response() {
|
||||
let failure = mapped("reducto", &http(429, "Request timed out"));
|
||||
assert_eq!(failure.error, PublicError::Timeout { status: 408 });
|
||||
assert_eq!(failure.upstream, upstream(429, "Request timed out"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn family_text_rules_also_classify_failures_without_a_response() {
|
||||
let original = OriginalException::Plain {
|
||||
message: "Request too large".into(),
|
||||
};
|
||||
assert_eq!(mapped("mistral", &original).error, PublicError::RateLimit);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::openai_family("mistral", "MistralException - rejected REDACTED")]
|
||||
#[case::vertex_family("vertex_ai", "Vertex_aiException - rejected REDACTED")]
|
||||
#[case::other_family("reducto", "ReductoException - rejected REDACTED")]
|
||||
fn every_family_reports_the_redacted_text(#[case] provider: &str, #[case] message: &str) {
|
||||
let failure = mapped(provider, &http(400, "rejected Bearer abcdefghijklmnop"));
|
||||
assert_eq!(failure.message, message);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redaction_runs_before_the_rules_see_the_text() {
|
||||
let body = "db_password=rate_limit";
|
||||
assert_eq!(
|
||||
mapped("mistral", &http(400, body)).error,
|
||||
PublicError::BadRequest
|
||||
);
|
||||
assert_eq!(
|
||||
exception_type(&context("mistral"), None, &http(400, body)).error,
|
||||
PublicError::RateLimit
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn without_a_redactor_the_text_is_kept() {
|
||||
let body = "rejected Bearer abcdefghijklmnop";
|
||||
assert_eq!(
|
||||
exception_type(&context("reducto"), None, &http(400, body)).message,
|
||||
format!("ReductoException - {body}")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_rule_hint_follows_the_message() {
|
||||
let failure = mapped("mistral", &http(400, "invalid_encrypted_content"));
|
||||
assert_eq!(failure.error, PublicError::BadRequest);
|
||||
assert!(
|
||||
failure
|
||||
.message
|
||||
.starts_with("MistralException - invalid_encrypted_content\n\n This error occurs")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::sync(
|
||||
false,
|
||||
Some(0.5),
|
||||
Some(0.5031),
|
||||
"Connection timed out after 0.5 seconds."
|
||||
)]
|
||||
#[case::async_rounds_the_elapsed_time(
|
||||
true,
|
||||
Some(0.5),
|
||||
Some(0.5031),
|
||||
"Connection timed out. Timeout passed=0.5, time taken=0.503 seconds"
|
||||
)]
|
||||
#[case::whole_seconds_keep_a_decimal(
|
||||
true,
|
||||
Some(600.0),
|
||||
Some(2.0),
|
||||
"Connection timed out. Timeout passed=600.0, time taken=2.0 seconds"
|
||||
)]
|
||||
#[case::unknown_values_render_as_none(
|
||||
true,
|
||||
None,
|
||||
None,
|
||||
"Connection timed out. Timeout passed=None, time taken=None seconds"
|
||||
)]
|
||||
fn timeout_text_follows_the_delivery_mode(
|
||||
#[case] asynchronous: bool,
|
||||
#[case] timeout_seconds: Option<f64>,
|
||||
#[case] elapsed_seconds: Option<f64>,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
assert_eq!(
|
||||
timeout_message(asynchronous, timeout_seconds, elapsed_seconds),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn debug_information_follows_the_python_layout() {
|
||||
let context = ExceptionContext {
|
||||
vertex_project: Some("project".into()),
|
||||
vertex_location: Some("region".into()),
|
||||
model_group: Some("ocr".into()),
|
||||
deployment: Some("deployment".into()),
|
||||
user_api_key_alias: Some("key".into()),
|
||||
..context("vertex_ai")
|
||||
};
|
||||
assert_eq!(
|
||||
exception_type(&context, None, &http(400, "rejected")).debug_info,
|
||||
concat!(
|
||||
"\n\nKey Name: `key`\nTeam: `None`",
|
||||
"\nModel: ocr-model",
|
||||
"\nAPI Base: `region-aiplatform.googleapis.com/v1/projects/project/locations/region/publishers/google/models/ocr-model:generateContent`",
|
||||
"\nmodel_group: `ocr`\n",
|
||||
"\ndeployment: `deployment`\n",
|
||||
"\nvertex_project: `project`\n",
|
||||
"\nvertex_location: `region`\n",
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::bare(ExceptionContext::default(), "\nModel: ")]
|
||||
#[case::team_alias(
|
||||
ExceptionContext { model: "m".into(), user_api_key_alias: Some("key".into()), user_api_key_team_alias: Some("team".into()), ..ExceptionContext::default() },
|
||||
"\n\nKey Name: `key`\nTeam: `team`\nModel: m"
|
||||
)]
|
||||
#[case::team_alias_without_key_is_ignored(
|
||||
ExceptionContext { model: "m".into(), user_api_key_team_alias: Some("team".into()), ..ExceptionContext::default() },
|
||||
"\nModel: m"
|
||||
)]
|
||||
#[case::project_without_location_has_no_api_base(
|
||||
ExceptionContext { model: "m".into(), vertex_project: Some("p".into()), ..ExceptionContext::default() },
|
||||
"\nModel: m\nvertex_project: `p`\n"
|
||||
)]
|
||||
#[case::location_without_project_has_no_api_base(
|
||||
ExceptionContext { model: "m".into(), vertex_location: Some("l".into()), ..ExceptionContext::default() },
|
||||
"\nModel: m\nvertex_location: `l`\n"
|
||||
)]
|
||||
fn each_optional_context_field_adds_its_own_line(
|
||||
#[case] context: ExceptionContext,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
assert_eq!(
|
||||
extra_information(&context, api_base(&context).as_deref()),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::openai_keeps_its_brand("openai", "OpenAIException")]
|
||||
#[case::lowercase("mistral", "MistralException")]
|
||||
#[case::keeps_the_rest("azure_ai", "Azure_aiException")]
|
||||
#[case::empty("", "")]
|
||||
fn exception_provider_capitalizes_only_the_first_letter(
|
||||
#[case] provider: &str,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
assert_eq!(exception_provider(provider), expected);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,192 @@
|
|||
use super::public::PublicError;
|
||||
use super::rules::{Rule, contains_any, is_context_window_exceeded, is_rate_limit};
|
||||
|
||||
const ENCRYPTED_CONTENT_HELP: &str = "\n\n This error occurs when load balancing Responses API across deployments with different API keys.\n Encrypted content is tied to the organization that created it and cannot be decrypted by other organizations.\n\n Solution: Enable 'encrypted_content_affinity' to route follow-up requests to the correct deployment:\n\n router_settings:\n enable_pre_call_checks: true\n optional_pre_call_checks:\n - encrypted_content_affinity\n\n Learn more: https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing";
|
||||
|
||||
/// The text branches of `_map_openai_exception`, in its order.
|
||||
pub(super) const RULES: &[Rule] = &[
|
||||
Rule::new(
|
||||
|mapping| is_rate_limit(&mapping.error_str, mapping.status),
|
||||
PublicError::RateLimit,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| is_context_window_exceeded(&mapping.error_str),
|
||||
PublicError::ContextWindowExceeded,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
mapping.error_str.contains("invalid_request_error")
|
||||
&& mapping.error_str.contains("model_not_found")
|
||||
},
|
||||
PublicError::NotFound,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| mapping.error_str.contains("A timeout occurred"),
|
||||
PublicError::Timeout { status: 408 },
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
let error_str = &mapping.error_str;
|
||||
(error_str.contains("invalid_request_error")
|
||||
&& error_str.contains("content_policy_violation"))
|
||||
|| (error_str.contains("Invalid prompt")
|
||||
&& error_str.contains("violating our usage policy"))
|
||||
|| error_str
|
||||
.to_lowercase()
|
||||
.contains("request was rejected as a result of the safety system")
|
||||
},
|
||||
PublicError::ContentPolicyViolation,
|
||||
),
|
||||
Rule {
|
||||
hint: ENCRYPTED_CONTENT_HELP,
|
||||
..Rule::new(
|
||||
|mapping| {
|
||||
contains_any(
|
||||
&mapping.error_str,
|
||||
&["invalid_encrypted_content", "could not be verified"],
|
||||
)
|
||||
},
|
||||
PublicError::BadRequest,
|
||||
)
|
||||
},
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
mapping.error_str.contains("invalid_request_error")
|
||||
&& !mapping.error_str.contains("Incorrect API key provided")
|
||||
},
|
||||
PublicError::BadRequest,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
contains_any(
|
||||
&mapping.error_str,
|
||||
&[
|
||||
"Web server is returning an unknown error",
|
||||
"The server had an error processing your request.",
|
||||
],
|
||||
)
|
||||
},
|
||||
PublicError::InternalServer,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| mapping.error_str.contains("Request too large"),
|
||||
PublicError::RateLimit,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
mapping
|
||||
.error_str
|
||||
.contains("Mistral API raised a streaming error")
|
||||
},
|
||||
PublicError::Api { status: 500 },
|
||||
),
|
||||
];
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::rules::first_match;
|
||||
use super::super::testing::mapping;
|
||||
use super::*;
|
||||
|
||||
fn classified(status: Option<u16>, text: &str) -> Option<PublicError> {
|
||||
first_match(RULES, &mapping(status, text)).map(|rule| rule.error)
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::rate_limit_phrase("rate limit reached", PublicError::RateLimit)]
|
||||
#[case::context_window(
|
||||
"This model's maximum context length is 10",
|
||||
PublicError::ContextWindowExceeded
|
||||
)]
|
||||
#[case::model_not_found("invalid_request_error model_not_found", PublicError::NotFound)]
|
||||
#[case::timeout_occurred("A timeout occurred", PublicError::Timeout { status: 408 })]
|
||||
#[case::content_policy_error_code(
|
||||
"invalid_request_error content_policy_violation",
|
||||
PublicError::ContentPolicyViolation
|
||||
)]
|
||||
#[case::content_policy_usage_policy(
|
||||
"Invalid prompt violating our usage policy",
|
||||
PublicError::ContentPolicyViolation
|
||||
)]
|
||||
#[case::content_policy_safety_system(
|
||||
"Request was rejected as a result of the safety system",
|
||||
PublicError::ContentPolicyViolation
|
||||
)]
|
||||
#[case::encrypted_content("invalid_encrypted_content", PublicError::BadRequest)]
|
||||
#[case::unverifiable_content("could not be verified", PublicError::BadRequest)]
|
||||
#[case::invalid_request("invalid_request_error bad field", PublicError::BadRequest)]
|
||||
#[case::unknown_server_error(
|
||||
"Web server is returning an unknown error",
|
||||
PublicError::InternalServer
|
||||
)]
|
||||
#[case::server_had_an_error(
|
||||
"The server had an error processing your request.",
|
||||
PublicError::InternalServer
|
||||
)]
|
||||
#[case::request_too_large("Request too large", PublicError::RateLimit)]
|
||||
#[case::mistral_streaming_error(
|
||||
"Mistral API raised a streaming error",
|
||||
PublicError::Api { status: 500 }
|
||||
)]
|
||||
fn each_text_rule_claims_its_marker(#[case] text: &str, #[case] expected: PublicError) {
|
||||
assert_eq!(classified(Some(400), text), Some(expected));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::rate_limit_before_context_window(
|
||||
"rate limit and This model's maximum context length is 10",
|
||||
PublicError::RateLimit
|
||||
)]
|
||||
#[case::context_window_before_content_policy(
|
||||
"This model's maximum context length is 10 invalid_request_error content_policy_violation",
|
||||
PublicError::ContextWindowExceeded
|
||||
)]
|
||||
#[case::model_not_found_before_invalid_request(
|
||||
"invalid_request_error model_not_found",
|
||||
PublicError::NotFound
|
||||
)]
|
||||
#[case::timeout_before_invalid_request(
|
||||
"A timeout occurred invalid_request_error",
|
||||
PublicError::Timeout { status: 408 }
|
||||
)]
|
||||
#[case::content_policy_before_invalid_request(
|
||||
"invalid_request_error content_policy_violation",
|
||||
PublicError::ContentPolicyViolation
|
||||
)]
|
||||
#[case::encrypted_content_before_invalid_request(
|
||||
"invalid_request_error invalid_encrypted_content",
|
||||
PublicError::BadRequest
|
||||
)]
|
||||
fn the_earlier_rule_wins_when_two_apply(#[case] text: &str, #[case] expected: PublicError) {
|
||||
assert_eq!(classified(Some(400), text), Some(expected));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::encrypted_content("invalid_encrypted_content", ENCRYPTED_CONTENT_HELP)]
|
||||
#[case::plain_invalid_request("invalid_request_error bad field", "")]
|
||||
fn only_encrypted_content_failures_carry_the_affinity_help(
|
||||
#[case] text: &str,
|
||||
#[case] hint: &str,
|
||||
) {
|
||||
assert_eq!(
|
||||
first_match(RULES, &mapping(Some(400), text)).map(|rule| rule.hint),
|
||||
Some(hint)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::bad_key_is_left_to_the_status("invalid_request_error Incorrect API key provided")]
|
||||
#[case::echoed_429_is_not_a_rate_limit("token 429 in the prompt")]
|
||||
#[case::unmarked("rejected")]
|
||||
fn text_without_a_marker_is_left_to_the_status_table(#[case] text: &str) {
|
||||
assert_eq!(classified(Some(400), text), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_standalone_429_counts_with_a_429_status() {
|
||||
assert_eq!(
|
||||
classified(Some(429), "got 429 back"),
|
||||
Some(PublicError::RateLimit)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,144 @@
|
|||
/// A failure a Rust route produced, before any public class is chosen.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum OriginalException {
|
||||
Http {
|
||||
status: u16,
|
||||
body: String,
|
||||
headers: Vec<(String, String)>,
|
||||
},
|
||||
Connection {
|
||||
message: String,
|
||||
},
|
||||
Timeout {
|
||||
timeout_seconds: Option<f64>,
|
||||
elapsed_seconds: Option<f64>,
|
||||
},
|
||||
/// A failure with no HTTP response behind it, such as an unparseable body or a local
|
||||
/// file error.
|
||||
Plain {
|
||||
message: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// Which provider-specific text rules apply before the shared status table.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum ExceptionFamily {
|
||||
OpenAiCompatible,
|
||||
VertexAi,
|
||||
Cohere,
|
||||
Other,
|
||||
}
|
||||
|
||||
/// `openai_compatible_providers` in `litellm/constants.py`.
|
||||
const OPENAI_COMPATIBLE_PROVIDERS: &[&str] = &[
|
||||
"anyscale",
|
||||
"groq",
|
||||
"nvidia_nim",
|
||||
"cerebras",
|
||||
"baseten",
|
||||
"sambanova",
|
||||
"ai21_chat",
|
||||
"ai21",
|
||||
"volcengine",
|
||||
"codestral",
|
||||
"deepseek",
|
||||
"tencent",
|
||||
"deepinfra",
|
||||
"perplexity",
|
||||
"xinference",
|
||||
"xai",
|
||||
"zai",
|
||||
"together_ai",
|
||||
"fireworks_ai",
|
||||
"empower",
|
||||
"friendliai",
|
||||
"azure_ai",
|
||||
"github",
|
||||
"litellm_proxy",
|
||||
"hosted_vllm",
|
||||
"llamafile",
|
||||
"lm_studio",
|
||||
"galadriel",
|
||||
"github_copilot",
|
||||
"chatgpt",
|
||||
"novita",
|
||||
"meta_llama",
|
||||
"publicai",
|
||||
"synthetic",
|
||||
"tensormesh",
|
||||
"apertis",
|
||||
"nano-gpt",
|
||||
"poe",
|
||||
"chutes",
|
||||
"parasail",
|
||||
"libertai",
|
||||
"featherless_ai",
|
||||
"nscale",
|
||||
"nebius",
|
||||
"dashscope",
|
||||
"qwencloud",
|
||||
"qwen_ai_platform",
|
||||
"modelscope",
|
||||
"moonshot",
|
||||
"v0",
|
||||
"helicone",
|
||||
"morph",
|
||||
"lambda_ai",
|
||||
"inception",
|
||||
"hyperbolic",
|
||||
"vercel_ai_gateway",
|
||||
"aiml",
|
||||
"wandb",
|
||||
"cometapi",
|
||||
"clarifai",
|
||||
"docker_model_runner",
|
||||
"ragflow",
|
||||
"pinstripes",
|
||||
"darkbloom",
|
||||
"meta",
|
||||
"cognition",
|
||||
"scx-ai",
|
||||
];
|
||||
|
||||
impl ExceptionFamily {
|
||||
/// The provider dispatch at the top of Python's `exception_type`, in its order.
|
||||
pub fn for_provider(provider: &str) -> Self {
|
||||
match provider {
|
||||
"openai" | "text-completion-openai" | "custom_openai" | "mistral" | "runwayml" => {
|
||||
Self::OpenAiCompatible
|
||||
}
|
||||
provider if OPENAI_COMPATIBLE_PROVIDERS.contains(&provider) => Self::OpenAiCompatible,
|
||||
"vertex_ai" | "vertex_ai_beta" | "gemini" => Self::VertexAi,
|
||||
"cohere" | "cohere_chat" => Self::Cohere,
|
||||
_ => Self::Other,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::openai("openai", ExceptionFamily::OpenAiCompatible)]
|
||||
#[case::text_completion_openai("text-completion-openai", ExceptionFamily::OpenAiCompatible)]
|
||||
#[case::custom_openai("custom_openai", ExceptionFamily::OpenAiCompatible)]
|
||||
#[case::mistral("mistral", ExceptionFamily::OpenAiCompatible)]
|
||||
#[case::runwayml("runwayml", ExceptionFamily::OpenAiCompatible)]
|
||||
#[case::listed_compatible("azure_ai", ExceptionFamily::OpenAiCompatible)]
|
||||
#[case::compatible_list_wins_over_its_own_mapper(
|
||||
"together_ai",
|
||||
ExceptionFamily::OpenAiCompatible
|
||||
)]
|
||||
#[case::vertex_ai("vertex_ai", ExceptionFamily::VertexAi)]
|
||||
#[case::vertex_ai_beta("vertex_ai_beta", ExceptionFamily::VertexAi)]
|
||||
#[case::gemini("gemini", ExceptionFamily::VertexAi)]
|
||||
#[case::cohere("cohere", ExceptionFamily::Cohere)]
|
||||
#[case::cohere_chat("cohere_chat", ExceptionFamily::Cohere)]
|
||||
#[case::unported_mapper("anthropic", ExceptionFamily::Other)]
|
||||
#[case::unknown("reducto", ExceptionFamily::Other)]
|
||||
#[case::empty("", ExceptionFamily::Other)]
|
||||
fn provider_selects_the_family(#[case] provider: &str, #[case] family: ExceptionFamily) {
|
||||
assert_eq!(ExceptionFamily::for_provider(provider), family);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
/// The public LiteLLM exception classes a Rust route failure can become. Python builds the
|
||||
/// class; Rust decides which one.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum PublicError {
|
||||
BadRequest,
|
||||
ContextWindowExceeded,
|
||||
ContentPolicyViolation,
|
||||
Authentication,
|
||||
PermissionDenied,
|
||||
NotFound,
|
||||
Timeout { status: u16 },
|
||||
RateLimit,
|
||||
InternalServer,
|
||||
BadGateway,
|
||||
ServiceUnavailable,
|
||||
ApiConnection,
|
||||
Api { status: u16 },
|
||||
}
|
||||
|
||||
impl PublicError {
|
||||
/// The `status_code` the Python class carries.
|
||||
pub const fn status_code(self) -> u16 {
|
||||
match self {
|
||||
Self::BadRequest | Self::ContextWindowExceeded | Self::ContentPolicyViolation => 400,
|
||||
Self::Authentication => 401,
|
||||
Self::PermissionDenied => 403,
|
||||
Self::NotFound => 404,
|
||||
Self::RateLimit => 429,
|
||||
Self::InternalServer | Self::ApiConnection => 500,
|
||||
Self::BadGateway => 502,
|
||||
Self::ServiceUnavailable => 503,
|
||||
Self::Timeout { status } | Self::Api { status } => status,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct UpstreamResponse {
|
||||
pub status: u16,
|
||||
pub body: String,
|
||||
pub headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct MappedFailure {
|
||||
pub error: PublicError,
|
||||
pub message: String,
|
||||
pub upstream: Option<UpstreamResponse>,
|
||||
pub debug_info: String,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::bad_request(PublicError::BadRequest, 400)]
|
||||
#[case::context_window(PublicError::ContextWindowExceeded, 400)]
|
||||
#[case::content_policy(PublicError::ContentPolicyViolation, 400)]
|
||||
#[case::authentication(PublicError::Authentication, 401)]
|
||||
#[case::permission_denied(PublicError::PermissionDenied, 403)]
|
||||
#[case::not_found(PublicError::NotFound, 404)]
|
||||
#[case::request_timeout(PublicError::Timeout { status: 408 }, 408)]
|
||||
#[case::gateway_timeout(PublicError::Timeout { status: 504 }, 504)]
|
||||
#[case::rate_limit(PublicError::RateLimit, 429)]
|
||||
#[case::internal_server(PublicError::InternalServer, 500)]
|
||||
#[case::api_connection(PublicError::ApiConnection, 500)]
|
||||
#[case::bad_gateway(PublicError::BadGateway, 502)]
|
||||
#[case::service_unavailable(PublicError::ServiceUnavailable, 503)]
|
||||
#[case::api(PublicError::Api { status: 501 }, 501)]
|
||||
fn status_codes_are_the_ones_the_python_classes_set(
|
||||
#[case] error: PublicError,
|
||||
#[case] status: u16,
|
||||
) {
|
||||
assert_eq!(error.status_code(), status);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,176 @@
|
|||
use std::sync::LazyLock;
|
||||
|
||||
use fancy_regex::Regex;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::Mapping;
|
||||
use super::public::PublicError;
|
||||
|
||||
/// One text branch of a Python `_map_*_exception` function: when it applies, the class it
|
||||
/// raises, and any help text appended to the message.
|
||||
pub(super) struct Rule {
|
||||
pub(super) when: fn(&Mapping) -> bool,
|
||||
pub(super) error: PublicError,
|
||||
pub(super) hint: &'static str,
|
||||
}
|
||||
|
||||
impl Rule {
|
||||
pub(super) const fn new(when: fn(&Mapping) -> bool, error: PublicError) -> Self {
|
||||
Self {
|
||||
when,
|
||||
error,
|
||||
hint: "",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The first rule that applies decides the class, as the `if`/`elif` chain does in Python.
|
||||
pub(super) fn first_match<'r>(rules: &'r [Rule], mapping: &Mapping) -> Option<&'r Rule> {
|
||||
rules.iter().find(|rule| (rule.when)(mapping))
|
||||
}
|
||||
|
||||
pub(super) fn contains_any(text: &str, markers: &[&str]) -> bool {
|
||||
markers.iter().any(|marker| text.contains(marker))
|
||||
}
|
||||
|
||||
static STANDALONE_429: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"\b429\b").expect("valid regex"));
|
||||
static RATE_LIMIT_PHRASE: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"rate[\s_\-]*limit").expect("valid regex"));
|
||||
|
||||
/// `ExceptionCheckers.is_error_str_rate_limit`.
|
||||
pub(super) fn is_rate_limit(error_str: &str, status: Option<u16>) -> bool {
|
||||
if STANDALONE_429.is_match(error_str).unwrap_or(false) && matches!(status, None | Some(429)) {
|
||||
return true;
|
||||
}
|
||||
let lower = error_str.to_lowercase();
|
||||
RATE_LIMIT_PHRASE.is_match(&lower).unwrap_or(false)
|
||||
|| lower.contains("service tier capacity exceeded")
|
||||
}
|
||||
|
||||
/// `ExceptionCheckers.is_error_str_context_window_exceeded`.
|
||||
pub(super) fn is_context_window_exceeded(error_str: &str) -> bool {
|
||||
let lower = error_str.to_lowercase();
|
||||
if lower.contains("string_above_max_length") {
|
||||
return false;
|
||||
}
|
||||
if lower.contains("invalid 'user'") && lower.contains("string too long") {
|
||||
return false;
|
||||
}
|
||||
contains_any(
|
||||
&lower,
|
||||
&[
|
||||
"exceed context limit",
|
||||
"this model's maximum context length is",
|
||||
"string too long. expected a string with maximum length",
|
||||
"model's maximum context limit",
|
||||
"is longer than the model's context length",
|
||||
"input tokens exceed the configured limit",
|
||||
"`inputs` tokens + `max_new_tokens` must be",
|
||||
"exceeds the available context size",
|
||||
"exceeds the maximum number of tokens allowed",
|
||||
],
|
||||
) || (lower.contains("current length is") && lower.contains("while limit is"))
|
||||
|| (lower.contains("maximum input length is") && lower.contains("tokens"))
|
||||
}
|
||||
|
||||
/// The integer `error.code` of a JSON error body, read the way Python's `int()` would.
|
||||
pub(super) fn body_error_code(error_str: &str) -> Option<i64> {
|
||||
let body: Value = serde_json::from_str(error_str).ok()?;
|
||||
let Some(Value::Object(error)) = body.as_object()?.get("error") else {
|
||||
return None;
|
||||
};
|
||||
match error.get("code")? {
|
||||
Value::Number(number) => number
|
||||
.as_i64()
|
||||
.or_else(|| number.as_f64().map(|value| value.trunc() as i64)),
|
||||
Value::String(code) => code.trim().replace('_', "").parse().ok(),
|
||||
Value::Bool(flag) => Some(i64::from(*flag)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::testing::mapping;
|
||||
use super::*;
|
||||
|
||||
const ORDERED: &[Rule] = &[
|
||||
Rule::new(
|
||||
|mapping| mapping.error_str.contains("first"),
|
||||
PublicError::NotFound,
|
||||
),
|
||||
Rule::new(|_| true, PublicError::ApiConnection),
|
||||
];
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::earlier_rule_wins("first and second", PublicError::NotFound)]
|
||||
#[case::later_rule_when_the_earlier_does_not_apply("second", PublicError::ApiConnection)]
|
||||
fn the_first_applicable_rule_decides(#[case] text: &str, #[case] expected: PublicError) {
|
||||
let rule = first_match(ORDERED, &mapping(Some(400), text));
|
||||
assert_eq!(rule.map(|rule| rule.error), Some(expected));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_applicable_rule_leaves_the_failure_to_the_caller() {
|
||||
assert!(first_match(&ORDERED[..1], &mapping(Some(400), "second")).is_none());
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::standalone_429_with_429_status("got 429 back", Some(429), true)]
|
||||
#[case::standalone_429_with_other_status("got 429 back", Some(400), false)]
|
||||
#[case::standalone_429_with_unknown_status("got 429 back", None, true)]
|
||||
#[case::embedded_429("token4290", Some(429), false)]
|
||||
#[case::phrase_spaced("Rate Limit reached", None, true)]
|
||||
#[case::phrase_underscored("rate_limit", None, true)]
|
||||
#[case::phrase_hyphenated("rate-limit", None, true)]
|
||||
#[case::service_tier("Service tier capacity exceeded", None, true)]
|
||||
#[case::unrelated("rejected", Some(429), false)]
|
||||
fn rate_limit_detection(
|
||||
#[case] text: &str,
|
||||
#[case] status: Option<u16>,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
assert_eq!(is_rate_limit(text, status), expected);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::exceed_context_limit("Exceed context limit", true)]
|
||||
#[case::maximum_context_length("This model's maximum context length is 10", true)]
|
||||
#[case::string_too_long("string too long. Expected a string with maximum length 5", true)]
|
||||
#[case::maximum_context_limit("the model's maximum context limit", true)]
|
||||
#[case::longer_than_context("prompt is longer than the model's context length", true)]
|
||||
#[case::configured_limit("input tokens exceed the configured limit", true)]
|
||||
#[case::max_new_tokens("`inputs` tokens + `max_new_tokens` must be <= 10", true)]
|
||||
#[case::available_context("exceeds the available context size", true)]
|
||||
#[case::maximum_tokens("exceeds the maximum number of tokens allowed", true)]
|
||||
#[case::current_and_limit("current length is 9 while limit is 8", true)]
|
||||
#[case::current_without_limit("current length is 9", false)]
|
||||
#[case::maximum_input_tokens("maximum input length is 8 tokens", true)]
|
||||
#[case::maximum_input_without_tokens("maximum input length is 8", false)]
|
||||
#[case::string_above_max_length_wins("string_above_max_length exceed context limit", false)]
|
||||
#[case::user_field_is_not_context(
|
||||
"invalid 'user': string too long. expected a string with maximum length",
|
||||
false
|
||||
)]
|
||||
#[case::unrelated("rejected", false)]
|
||||
fn context_window_detection(#[case] text: &str, #[case] expected: bool) {
|
||||
assert_eq!(is_context_window_exceeded(text), expected);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::integer(r#"{"error": {"code": 429}}"#, Some(429))]
|
||||
#[case::float(r#"{"error": {"code": 429.9}}"#, Some(429))]
|
||||
#[case::string(r#"{"error": {"code": " 4_29 "}}"#, Some(429))]
|
||||
#[case::boolean(r#"{"error": {"code": true}}"#, Some(1))]
|
||||
#[case::unparseable_string(r#"{"error": {"code": "slow"}}"#, None)]
|
||||
#[case::null(r#"{"error": {"code": null}}"#, None)]
|
||||
#[case::no_code(r#"{"error": {}}"#, None)]
|
||||
#[case::error_not_an_object(r#"{"error": "429"}"#, None)]
|
||||
#[case::no_error(r#"{"code": 429}"#, None)]
|
||||
#[case::not_an_object("[429]", None)]
|
||||
#[case::not_json("429", None)]
|
||||
fn body_error_code_reads_the_nested_code(#[case] body: &str, #[case] expected: Option<i64>) {
|
||||
assert_eq!(body_error_code(body), expected);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,49 @@
|
|||
use super::public::PublicError;
|
||||
|
||||
/// `_map_exception_by_status`, the one place a provider status picks a class. Statuses
|
||||
/// below 400 are not failures the table claims.
|
||||
pub(super) fn classify(status: u16) -> Option<PublicError> {
|
||||
let error = match status {
|
||||
..400 => return None,
|
||||
401 => PublicError::Authentication,
|
||||
403 => PublicError::PermissionDenied,
|
||||
404 => PublicError::NotFound,
|
||||
408 | 504 => PublicError::Timeout { status },
|
||||
429 => PublicError::RateLimit,
|
||||
500 => PublicError::InternalServer,
|
||||
502 => PublicError::BadGateway,
|
||||
503 => PublicError::ServiceUnavailable,
|
||||
400..500 => PublicError::BadRequest,
|
||||
_ => PublicError::Api { status },
|
||||
};
|
||||
Some(error)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::below_client_errors(399, None)]
|
||||
#[case::lowest_client_error(400, Some(PublicError::BadRequest))]
|
||||
#[case::authentication(401, Some(PublicError::Authentication))]
|
||||
#[case::permission_denied(403, Some(PublicError::PermissionDenied))]
|
||||
#[case::not_found(404, Some(PublicError::NotFound))]
|
||||
#[case::request_timeout(408, Some(PublicError::Timeout { status: 408 }))]
|
||||
#[case::other_client_error(409, Some(PublicError::BadRequest))]
|
||||
#[case::unprocessable(422, Some(PublicError::BadRequest))]
|
||||
#[case::rate_limited(429, Some(PublicError::RateLimit))]
|
||||
#[case::highest_client_error(499, Some(PublicError::BadRequest))]
|
||||
#[case::internal_server(500, Some(PublicError::InternalServer))]
|
||||
#[case::other_server_error(501, Some(PublicError::Api { status: 501 }))]
|
||||
#[case::bad_gateway(502, Some(PublicError::BadGateway))]
|
||||
#[case::service_unavailable(503, Some(PublicError::ServiceUnavailable))]
|
||||
#[case::gateway_timeout(504, Some(PublicError::Timeout { status: 504 }))]
|
||||
#[case::highest_server_error(599, Some(PublicError::Api { status: 599 }))]
|
||||
fn every_mapped_status_and_the_fallback(
|
||||
#[case] status: u16,
|
||||
#[case] expected: Option<PublicError>,
|
||||
) {
|
||||
assert_eq!(classify(status), expected);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,177 @@
|
|||
use super::public::PublicError;
|
||||
use super::rules::{Rule, body_error_code, contains_any, is_context_window_exceeded};
|
||||
|
||||
const QUOTA_MARKERS: &[&str] = &[
|
||||
"429 Quota exceeded",
|
||||
"Quota exceeded for",
|
||||
"Resource exhausted",
|
||||
"429 Unable to submit request because the service is temporarily out of capacity.",
|
||||
];
|
||||
|
||||
/// The text branches of `_map_vertex_exception`, in its order.
|
||||
pub(super) const RULES: &[Rule] = &[
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
contains_any(
|
||||
&mapping.error_str,
|
||||
&[
|
||||
"Vertex AI API has not been used in project",
|
||||
"Unable to find your project",
|
||||
],
|
||||
)
|
||||
},
|
||||
PublicError::BadRequest,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
mapping
|
||||
.error_str
|
||||
.contains("400 Request payload size exceeds")
|
||||
|| is_context_window_exceeded(&mapping.error_str)
|
||||
},
|
||||
PublicError::ContextWindowExceeded,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
contains_any(
|
||||
&mapping.error_str,
|
||||
&["None Unknown Error.", "Content has no parts."],
|
||||
)
|
||||
},
|
||||
PublicError::InternalServer,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| mapping.error_str.contains("API key not valid."),
|
||||
PublicError::Authentication,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
contains_any(
|
||||
&mapping.error_str,
|
||||
&[
|
||||
"The response was blocked.",
|
||||
"Output blocked by content filtering policy",
|
||||
],
|
||||
)
|
||||
},
|
||||
PublicError::ContentPolicyViolation,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
contains_any(&mapping.error_str, QUOTA_MARKERS)
|
||||
|| (mapping
|
||||
.status
|
||||
.is_some_and(|status| (500..600).contains(&status))
|
||||
&& body_error_code(&mapping.error_str) == Some(429))
|
||||
},
|
||||
PublicError::RateLimit,
|
||||
),
|
||||
Rule::new(
|
||||
|mapping| {
|
||||
contains_any(
|
||||
&mapping.error_str,
|
||||
&["500 Internal Server Error", "The model is overloaded."],
|
||||
)
|
||||
},
|
||||
PublicError::InternalServer,
|
||||
),
|
||||
];
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::rules::first_match;
|
||||
use super::super::testing::mapping;
|
||||
use super::*;
|
||||
|
||||
fn classified(status: Option<u16>, text: &str) -> Option<PublicError> {
|
||||
first_match(RULES, &mapping(status, text)).map(|rule| rule.error)
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::api_not_enabled(
|
||||
"Vertex AI API has not been used in project x",
|
||||
PublicError::BadRequest
|
||||
)]
|
||||
#[case::project_not_found("Unable to find your project", PublicError::BadRequest)]
|
||||
#[case::payload_too_large(
|
||||
"400 Request payload size exceeds the limit",
|
||||
PublicError::ContextWindowExceeded
|
||||
)]
|
||||
#[case::context_window(
|
||||
"This model's maximum context length is 10",
|
||||
PublicError::ContextWindowExceeded
|
||||
)]
|
||||
#[case::unknown_error("None Unknown Error.", PublicError::InternalServer)]
|
||||
#[case::no_parts("Content has no parts.", PublicError::InternalServer)]
|
||||
#[case::api_key_not_valid("API key not valid.", PublicError::Authentication)]
|
||||
#[case::response_blocked("The response was blocked.", PublicError::ContentPolicyViolation)]
|
||||
#[case::output_blocked(
|
||||
"Output blocked by content filtering policy",
|
||||
PublicError::ContentPolicyViolation
|
||||
)]
|
||||
#[case::quota_exceeded_429("429 Quota exceeded", PublicError::RateLimit)]
|
||||
#[case::quota_exceeded_for("Quota exceeded for aiplatform", PublicError::RateLimit)]
|
||||
#[case::resource_exhausted("Resource exhausted", PublicError::RateLimit)]
|
||||
#[case::out_of_capacity(
|
||||
"429 Unable to submit request because the service is temporarily out of capacity.",
|
||||
PublicError::RateLimit
|
||||
)]
|
||||
#[case::internal_server_text("500 Internal Server Error", PublicError::InternalServer)]
|
||||
#[case::overloaded("The model is overloaded.", PublicError::InternalServer)]
|
||||
fn each_text_rule_claims_its_marker(#[case] text: &str, #[case] expected: PublicError) {
|
||||
assert_eq!(classified(Some(400), text), Some(expected));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::server_error_wrapping_a_429(Some(503), Some(PublicError::RateLimit))]
|
||||
#[case::lowest_server_error(Some(500), Some(PublicError::RateLimit))]
|
||||
#[case::highest_server_error(Some(599), Some(PublicError::RateLimit))]
|
||||
#[case::client_error(Some(400), None)]
|
||||
#[case::no_status(None, None)]
|
||||
fn a_wrapped_429_is_a_rate_limit_only_behind_a_server_error(
|
||||
#[case] status: Option<u16>,
|
||||
#[case] expected: Option<PublicError>,
|
||||
) {
|
||||
assert_eq!(
|
||||
classified(status, r#"{"error": {"code": "429"}}"#),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::project_before_payload_size(
|
||||
"Unable to find your project 400 Request payload size exceeds",
|
||||
PublicError::BadRequest
|
||||
)]
|
||||
#[case::context_window_before_unknown_error(
|
||||
"This model's maximum context length is 10 None Unknown Error.",
|
||||
PublicError::ContextWindowExceeded
|
||||
)]
|
||||
#[case::unknown_error_before_api_key(
|
||||
"Content has no parts. API key not valid.",
|
||||
PublicError::InternalServer
|
||||
)]
|
||||
#[case::api_key_before_blocked(
|
||||
"API key not valid. The response was blocked.",
|
||||
PublicError::Authentication
|
||||
)]
|
||||
#[case::blocked_before_quota(
|
||||
"The response was blocked. Resource exhausted",
|
||||
PublicError::ContentPolicyViolation
|
||||
)]
|
||||
#[case::quota_before_overloaded(
|
||||
"Resource exhausted The model is overloaded.",
|
||||
PublicError::RateLimit
|
||||
)]
|
||||
fn the_earlier_rule_wins_when_two_apply(#[case] text: &str, #[case] expected: PublicError) {
|
||||
assert_eq!(classified(Some(400), text), Some(expected));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::a_403_in_the_text("got a 403 from 4031 tokens")]
|
||||
#[case::python_client_crash("IndexError: list index out of range")]
|
||||
#[case::unmarked("rejected")]
|
||||
fn text_without_a_marker_is_left_to_the_status_table(#[case] text: &str) {
|
||||
assert_eq!(classified(Some(400), text), None);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,7 +1,9 @@
|
|||
pub mod call_arguments;
|
||||
pub mod core_helpers;
|
||||
pub mod exception_mapping_utils;
|
||||
pub mod get_llm_provider_logic;
|
||||
pub mod params;
|
||||
pub mod prompt_templates;
|
||||
pub mod secret_redaction;
|
||||
pub mod serde_compat;
|
||||
pub mod url_utils;
|
||||
|
|
|
|||
109
litellm-rust/crates/core-utils/src/secret_redaction.rs
Normal file
109
litellm-rust/crates/core-utils/src/secret_redaction.rs
Normal file
|
|
@ -0,0 +1,109 @@
|
|||
use fancy_regex::Regex;
|
||||
|
||||
pub const REDACTED: &str = "REDACTED";
|
||||
|
||||
const DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH: usize = 16;
|
||||
|
||||
fn minimum_custom_key_length() -> usize {
|
||||
std::env::var("MINIMUM_CUSTOM_KEY_LENGTH")
|
||||
.ok()
|
||||
.and_then(|value| value.trim().parse().ok())
|
||||
.unwrap_or(DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH)
|
||||
}
|
||||
|
||||
fn secret_patterns(minimum_custom_key_length: usize) -> String {
|
||||
let sk_suffix_length = minimum_custom_key_length.saturating_sub("sk-".len());
|
||||
[
|
||||
r"-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----",
|
||||
r"\bya29\.[A-Za-z0-9_.~+/-]+",
|
||||
r#"(?:client_secret|azure_password|azure_username)\s+[^\s,'"})\]{}>]+"#,
|
||||
r"(?:AKIA|ASIA)[0-9A-Z]{16}",
|
||||
r"Bearer\s+[A-Za-z0-9\-._~+/]{10,}=*",
|
||||
r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}",
|
||||
&format!(r"sk-[A-Za-z0-9\-_]{{{sk_suffix_length},}}"),
|
||||
r#"(?<=[?&])(?:api[_-]?key|\w*(?:token|password|passwd|client_secret|secret_key|_secret))=[^\s&'"]+"#,
|
||||
r#"(?:api[_-]?key)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]{8,}"#,
|
||||
r#"(?:x-api-key|api-key)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#,
|
||||
r"x-ak-[A-Za-z0-9\-_]{20,}",
|
||||
r"AIza[0-9A-Za-z\-_]{35}",
|
||||
r#"(?<=[?&])key=[^\s&'"]{8,}"#,
|
||||
r#"(?:^|(?<=\W))\w*(?:password|passwd|client_secret|secret_key|_secret)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#,
|
||||
r#"(?<=://)[^\s'":]{0,4096}:[^\s'"]{1,4096}(?=@)"#,
|
||||
r"dapi[0-9a-f]{32}",
|
||||
r#"litellm\.[A-Za-z0-9_]*_key['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#,
|
||||
r#"private_key['"]?\s*[:=]\s*['"]?(?:-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----|[^\s,'"})\]{}>]+)"#,
|
||||
concat!(
|
||||
r"(?:master_key|xai_key|database_url|db_url|connection_string|",
|
||||
r"aws_secret_access_key|aws_session_token|aws_access_key_id|",
|
||||
r"signing_key|encryption_key|",
|
||||
r"auth_token|access_token|refresh_token|",
|
||||
r"slack_webhook_url|webhook_url|",
|
||||
r"database_connection_string|",
|
||||
r"huggingface_token|jwt_secret)",
|
||||
r#"['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#,
|
||||
),
|
||||
r"\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]*",
|
||||
r"(?<=[?&])sig=[A-Za-z0-9%+/=]+",
|
||||
r#"\{[^{}]*"type"\s*:\s*"service_account"[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"#,
|
||||
]
|
||||
.join("|")
|
||||
}
|
||||
|
||||
/// Python's `_ENABLE_SECRET_REDACTION` pattern set, compiled once per configuration.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SecretRedactor {
|
||||
pattern: Regex,
|
||||
}
|
||||
|
||||
impl SecretRedactor {
|
||||
pub fn new(minimum_custom_key_length: usize) -> Self {
|
||||
let pattern = Regex::new(&format!(
|
||||
"(?i){}",
|
||||
secret_patterns(minimum_custom_key_length)
|
||||
))
|
||||
.expect("secret redaction patterns compile");
|
||||
Self { pattern }
|
||||
}
|
||||
|
||||
/// `None` when `LITELLM_DISABLE_REDACT_SECRETS` turns redaction off.
|
||||
pub fn from_env() -> Option<Self> {
|
||||
let disabled = std::env::var("LITELLM_DISABLE_REDACT_SECRETS")
|
||||
.is_ok_and(|value| value.eq_ignore_ascii_case("true"));
|
||||
(!disabled).then(|| Self::new(minimum_custom_key_length()))
|
||||
}
|
||||
|
||||
pub fn redact(&self, value: &str) -> String {
|
||||
self.pattern.replace_all(value, REDACTED).into_owned()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::bearer("auth failed: Bearer abcdefghijklmnop", "auth failed: REDACTED")]
|
||||
#[case::sk_key("key sk-abcdefghijklmnopqrstuvwxyz rejected", "key REDACTED rejected")]
|
||||
#[case::short_sk_key_is_kept("sk-abc", "sk-abc")]
|
||||
#[case::query_param("GET /v1?api_key=secret123&x=1", "GET /v1?REDACTED&x=1")]
|
||||
#[case::dict_repr("{'api_key': 'abcdefghij'}", "{'REDACTED'}")]
|
||||
#[case::url_credentials("postgres://user:pass@host/db", "postgres://REDACTED@host/db")]
|
||||
#[case::case_insensitive("BEARER ABCDEFGHIJKLMNOP", "REDACTED")]
|
||||
#[case::aws_key("AKIAABCDEFGHIJKLMNOP", "REDACTED")]
|
||||
#[case::sas_signature("https://x.blob/a?sv=1&sig=abc%2B=", "https://x.blob/a?sv=1&REDACTED")]
|
||||
#[case::password_needs_word_boundary("db_password=hunter2", "REDACTED")]
|
||||
#[case::plain_text_is_kept(r#"{"message": "rejected"}"#, r#"{"message": "rejected"}"#)]
|
||||
fn redacts_the_same_spans_as_the_python_patterns(#[case] input: &str, #[case] expected: &str) {
|
||||
assert_eq!(
|
||||
SecretRedactor::new(DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH).redact(input),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sk_threshold_follows_the_minimum_custom_key_length() {
|
||||
let redactor = SecretRedactor::new(8);
|
||||
assert_eq!(redactor.redact("sk-abcde"), REDACTED);
|
||||
assert_eq!(redactor.redact("sk-abcd"), "sk-abcd");
|
||||
}
|
||||
}
|
||||
|
|
@ -1700,6 +1700,9 @@ if TYPE_CHECKING:
|
|||
from .llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import (
|
||||
VertexAIAi21Config as VertexAIAi21Config,
|
||||
)
|
||||
from .llms.vertex_ai.vertex_ai_partner_models.mistral.transformation import (
|
||||
VertexAIMistralConfig as VertexAIMistralConfig,
|
||||
)
|
||||
from .llms.bedrock.chat.invoke_handler import (
|
||||
AmazonCohereChatConfig as AmazonCohereChatConfig,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -184,6 +184,7 @@ LLM_CONFIG_NAMES: Final = (
|
|||
"VertexAIAnthropicConfig",
|
||||
"VertexAILlama3Config",
|
||||
"VertexAIAi21Config",
|
||||
"VertexAIMistralConfig",
|
||||
"AmazonCohereChatConfig",
|
||||
"AmazonBedrockGlobalConfig",
|
||||
"AmazonAI21Config",
|
||||
|
|
@ -771,6 +772,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
|
|||
".llms.vertex_ai.vertex_ai_partner_models.ai21.transformation",
|
||||
"VertexAIAi21Config",
|
||||
),
|
||||
"VertexAIMistralConfig": (
|
||||
".llms.vertex_ai.vertex_ai_partner_models.mistral.transformation",
|
||||
"VertexAIMistralConfig",
|
||||
),
|
||||
"AmazonCohereChatConfig": (
|
||||
".llms.bedrock.chat.invoke_handler",
|
||||
"AmazonCohereChatConfig",
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@
|
|||
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
|
||||
"text_editor_20241022": null,
|
||||
"text_editor_20250124": null,
|
||||
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
|
||||
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
|
||||
"web-fetch-2025-09-10": "web-fetch-2025-09-10",
|
||||
"web-search-2025-03-05": "web-search-2025-03-05"
|
||||
|
|
@ -59,6 +60,7 @@
|
|||
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
|
||||
"text_editor_20241022": null,
|
||||
"text_editor_20250124": null,
|
||||
"thinking-binding-controls-2026-08-01": null,
|
||||
"token-efficient-tools-2025-02-19": null,
|
||||
"web-fetch-2025-09-10": "web-fetch-2025-09-10",
|
||||
"web-search-2025-03-05": "web-search-2025-03-05"
|
||||
|
|
@ -90,6 +92,7 @@
|
|||
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
|
||||
"text_editor_20241022": null,
|
||||
"text_editor_20250124": null,
|
||||
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
|
||||
"token-efficient-tools-2025-02-19": null,
|
||||
"tool-search-tool-2025-10-19": null,
|
||||
"web-fetch-2025-09-10": null,
|
||||
|
|
@ -122,6 +125,7 @@
|
|||
"structured-outputs-2025-11-13": null,
|
||||
"text_editor_20241022": null,
|
||||
"text_editor_20250124": null,
|
||||
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
|
||||
"token-efficient-tools-2025-02-19": null,
|
||||
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
|
||||
"web-fetch-2025-09-10": null,
|
||||
|
|
@ -154,6 +158,7 @@
|
|||
"structured-outputs-2025-11-13": null,
|
||||
"text_editor_20241022": null,
|
||||
"text_editor_20250124": null,
|
||||
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
|
||||
"token-efficient-tools-2025-02-19": null,
|
||||
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
|
||||
"web-fetch-2025-09-10": null,
|
||||
|
|
@ -187,6 +192,7 @@
|
|||
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
|
||||
"text_editor_20241022": null,
|
||||
"text_editor_20250124": null,
|
||||
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
|
||||
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
|
||||
"web-fetch-2025-09-10": "web-fetch-2025-09-10",
|
||||
"web-search-2025-03-05": "web-search-2025-03-05"
|
||||
|
|
|
|||
|
|
@ -533,6 +533,8 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict:
|
|||
"vertex_credentials",
|
||||
"gcs_bucket_name",
|
||||
"bucket_name",
|
||||
"s3_endpoint_url",
|
||||
"s3_region_name",
|
||||
"timeout",
|
||||
"max_retries",
|
||||
"_litellm_internal_model_credentials",
|
||||
|
|
|
|||
|
|
@ -183,6 +183,9 @@ MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIME
|
|||
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
|
||||
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
|
||||
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
|
||||
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH: Final = 8
|
||||
MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60
|
||||
MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE: Final = 4096
|
||||
|
||||
# Allowlist of commands permitted for MCP stdio transport.
|
||||
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.
|
||||
|
|
@ -317,6 +320,9 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = float(
|
|||
# RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
|
||||
|
||||
DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1"
|
||||
DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3"
|
||||
|
||||
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
|
||||
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
|
||||
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
|
||||
|
|
@ -1574,6 +1580,21 @@ ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS: Final = {
|
|||
# Works for all LLM pass-through endpoints (Vertex AI, Anthropic, Bedrock, etc.)
|
||||
PASS_THROUGH_HEADER_PREFIX: Final = "x-pass-"
|
||||
|
||||
AZURE_SPEECH_CUSTOM_LLM_PROVIDER: Final = "azure_speech"
|
||||
AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX: Final = "/azure_speech"
|
||||
AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX: Final = "/speech/"
|
||||
AZURE_SPEECH_BATCH_PATH_PREFIX: Final = "/speechtotext/"
|
||||
AZURE_SPEECH_FAST_TRANSCRIPTION_PATH: Final = "/speechtotext/transcriptions:transcribe"
|
||||
AZURE_SPEECH_STT_DOMAIN: Final = "stt.speech.microsoft.com"
|
||||
AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN: Final = "api.cognitive.microsoft.com"
|
||||
AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER: Final = "Ocp-Apim-Subscription-Key"
|
||||
AZURE_SPEECH_SHORT_AUDIO_MODEL: Final = "short-audio"
|
||||
AZURE_SPEECH_BATCH_MODEL: Final = "batch-transcription"
|
||||
AZURE_SPEECH_FAST_TRANSCRIPTION_MODEL: Final = "fast-transcription"
|
||||
AZURE_SPEECH_PRICING_MODEL: Final = "azure/speech/azure-stt"
|
||||
AZURE_SPEECH_TICKS_PER_SECOND: Final = 10_000_000
|
||||
AZURE_SPEECH_MILLISECONDS_PER_SECOND: Final = 1_000
|
||||
|
||||
BASE_MCP_ROUTE: Final = "/mcp"
|
||||
|
||||
TRANSCRIBE_JOB_POLLING_INTERVAL_SECONDS: Final = 10.0
|
||||
|
|
@ -1658,6 +1679,11 @@ LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS: Final = int(
|
|||
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE: Final = int(
|
||||
os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE", 1000)
|
||||
)
|
||||
LOGIN_THROTTLE_CACHE_KEY_PREFIX: Final = "login_fail"
|
||||
LOGIN_THROTTLE_UNKNOWN_SOURCE: Final = "unknown"
|
||||
LOGIN_THROTTLE_MAX_TRACKED_COUNTERS: Final = 20_000
|
||||
LOGIN_THROTTLE_MAX_TRACKED_BLOCKS: Final = 10_000
|
||||
LOGIN_THROTTLE_NOT_BLOCKED: Final = (0, 0)
|
||||
LITELLM_PROXY_ADMIN_NAME: Final = "default_user_id"
|
||||
LITELLM_PROXY_BUDGET_NAME: Final = "litellm-proxy-budget"
|
||||
GLOBAL_PROXY_SPEND_CACHE_KEY: Final = f"{LITELLM_PROXY_ADMIN_NAME}:spend"
|
||||
|
|
@ -2050,12 +2076,16 @@ MCP_SPEND_LOG_MODEL_PREFIX: Final[str] = "MCP: "
|
|||
PTU_SENTINEL_API_KEY: Final[str] = "__ptu_flat_cost__"
|
||||
PTU_ROLLUP_JOB_ID: Final[str] = "ptu_flat_cost_rollup_job"
|
||||
PTU_ROLLUP_LOCK_TTL_SECONDS: Final[int] = 900
|
||||
USAGE_TOP_API_KEYS_LIMIT: Final[int] = int(os.getenv("USAGE_TOP_API_KEYS_LIMIT", "100"))
|
||||
# Furthest back the catch-up pass looks for unpriced PTU days when a deployment
|
||||
# declares no ptu_effective_from, bounding the scan for an open-ended window.
|
||||
PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90
|
||||
# Deployments named in the lapsed-window alert before it is truncated, so a fleet-wide
|
||||
# expiry cannot produce an alert too large for the channel delivering it.
|
||||
PTU_LAPSED_ALERT_LIMIT: Final[int] = 10
|
||||
DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID: Final[str] = "daily_global_spend_reconcile_job"
|
||||
DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS: Final[int] = 3600
|
||||
DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM: Final[str] = "daily_global_spend_reconciled_through"
|
||||
# Slack allowed when deciding a sentinel row is stale. The row's updated_at and the
|
||||
# run's cutoff are stamped by different hosts, so clock skew between them must not let
|
||||
# one run delete a charge another just wrote. A stale row is hours old and a concurrent
|
||||
|
|
|
|||
|
|
@ -846,7 +846,12 @@ def image_edit(
|
|||
local_vars.update(kwargs)
|
||||
# Get ImageEditOptionalRequestParams with only valid parameters
|
||||
image_edit_optional_params: Final[ImageEditOptionalRequestParams] = (
|
||||
_get_ImageEditRequestUtils().get_requested_image_edit_optional_param(local_vars)
|
||||
_get_ImageEditRequestUtils().get_requested_image_edit_optional_param(
|
||||
local_vars,
|
||||
provider_supported_params=frozenset(
|
||||
image_edit_provider_config.get_supported_openai_params(model)
|
||||
).intersection(non_default_params),
|
||||
)
|
||||
)
|
||||
# Get optional parameters for the responses API
|
||||
image_edit_request_params: Final[dict] = _get_ImageEditRequestUtils().get_optional_params_image_edit(
|
||||
|
|
@ -857,7 +862,7 @@ def image_edit(
|
|||
additional_drop_params=kwargs.get("additional_drop_params"),
|
||||
)
|
||||
|
||||
if (
|
||||
if image_edit_provider_config.use_multipart_form_data() and (
|
||||
custom_llm_provider == "openai"
|
||||
or custom_llm_provider == "azure"
|
||||
or custom_llm_provider in litellm.openai_compatible_providers
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from collections.abc import Collection, Mapping
|
||||
from io import BufferedReader, BytesIO
|
||||
from typing import Any, Final, cast, get_type_hints
|
||||
|
||||
|
|
@ -63,6 +63,7 @@ class ImageEditRequestUtils:
|
|||
@staticmethod
|
||||
def get_requested_image_edit_optional_param(
|
||||
params: Mapping[str, object],
|
||||
provider_supported_params: Collection[str] = (),
|
||||
) -> ImageEditOptionalRequestParams:
|
||||
"""
|
||||
Filter parameters to only include those defined in ImageEditOptionalRequestParams.
|
||||
|
|
@ -73,7 +74,9 @@ class ImageEditRequestUtils:
|
|||
Returns:
|
||||
ImageEditOptionalRequestParams instance with only the valid parameters
|
||||
"""
|
||||
valid_keys: Final = get_type_hints(ImageEditOptionalRequestParams).keys()
|
||||
valid_keys: Final = frozenset(get_type_hints(ImageEditOptionalRequestParams)) | frozenset(
|
||||
provider_supported_params
|
||||
)
|
||||
filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None}
|
||||
return cast(ImageEditOptionalRequestParams, filtered_params)
|
||||
|
||||
|
|
|
|||
|
|
@ -44,6 +44,8 @@ OPTIONAL_KWARGS_KEYS: Final = (
|
|||
"client_side_timeout",
|
||||
"gcs_bucket_name",
|
||||
"bucket_name",
|
||||
"s3_endpoint_url",
|
||||
"s3_region_name",
|
||||
"vertex_credentials",
|
||||
"vertex_project",
|
||||
"vertex_location",
|
||||
|
|
|
|||
|
|
@ -190,7 +190,7 @@ def get_supported_openai_params(
|
|||
elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "vertex_ai_beta":
|
||||
if request_type == "chat_completion":
|
||||
if model.startswith("mistral"):
|
||||
return litellm.MistralConfig().get_supported_openai_params(model=model)
|
||||
return litellm.VertexAIMistralConfig().get_supported_openai_params(model=model)
|
||||
elif model.startswith("codestral"):
|
||||
return litellm.CodestralTextCompletionConfig().get_supported_openai_params(model=model)
|
||||
elif model.startswith("claude"):
|
||||
|
|
|
|||
|
|
@ -1848,6 +1848,9 @@ class CostCalculatorUtils:
|
|||
return azure_ai_image_cost_calculator(
|
||||
model=model,
|
||||
image_response=completion_response,
|
||||
size=resolved_size,
|
||||
n=resolved_n,
|
||||
optional_params=optional_params,
|
||||
)
|
||||
elif custom_llm_provider == litellm.LlmProviders.FAL_AI.value:
|
||||
from litellm.llms.fal_ai.cost_calculator import (
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
import base64
|
||||
from collections.abc import Mapping, Sequence
|
||||
from io import BufferedReader
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
from httpx._types import RequestFiles
|
||||
|
|
@ -24,21 +26,12 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
Azure AI Foundry FLUX 2 image edit config
|
||||
|
||||
Supports FLUX 2 models (e.g., flux.2-pro) for image editing.
|
||||
Uses the same /providers/blackforestlabs/v1/flux-2-pro endpoint as image generation,
|
||||
Uses the model-specific /providers/blackforestlabs/v1/flux-2-* endpoint as image generation,
|
||||
with the image passed as base64 in JSON body.
|
||||
"""
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
"""
|
||||
FLUX 2 supports a subset of OpenAI image edit params
|
||||
"""
|
||||
return [
|
||||
"prompt",
|
||||
"image",
|
||||
"model",
|
||||
"n",
|
||||
"size",
|
||||
]
|
||||
return AzureFoundryFluxImageGenerationConfig().get_supported_openai_params(model)
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
|
|
@ -50,14 +43,14 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
Map OpenAI params to FLUX 2 params.
|
||||
FLUX 2 uses the same param names as OpenAI for supported params.
|
||||
"""
|
||||
mapped_params: Final[dict[str, Any]] = {}
|
||||
supported_params: Final = self.get_supported_openai_params(model)
|
||||
|
||||
for key, value in dict(image_edit_optional_params).items():
|
||||
if key in supported_params and value is not None:
|
||||
mapped_params[key] = value
|
||||
|
||||
return mapped_params
|
||||
return AzureFoundryFluxImageGenerationConfig().map_openai_params(
|
||||
non_default_params=MappingProxyType(
|
||||
{key: value for key, value in image_edit_optional_params.items() if value is not None}
|
||||
),
|
||||
optional_params=MappingProxyType({}),
|
||||
model=model,
|
||||
drop_params=drop_params,
|
||||
)
|
||||
|
||||
def use_multipart_form_data(self) -> bool:
|
||||
"""FLUX 2 uses JSON requests, not multipart/form-data."""
|
||||
|
|
@ -90,7 +83,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
self,
|
||||
model: str,
|
||||
prompt: str | None,
|
||||
image: FileTypes | None,
|
||||
image: FileTypes | Sequence[FileTypes] | None,
|
||||
image_edit_optional_request_params: dict,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
|
|
@ -107,29 +100,29 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
if image is None:
|
||||
raise ValueError("FLUX 2 image edit requires an image.")
|
||||
|
||||
image_b64: Final = self._convert_image_to_base64(image)
|
||||
images: Final = tuple(image) if isinstance(image, list) else (image,)
|
||||
if not images:
|
||||
raise ValueError("FLUX 2 image edit requires at least one image.")
|
||||
max_reference_images: Final = 10 if "flex" in model.lower() else 8
|
||||
if len(images) > max_reference_images:
|
||||
raise ValueError(f"{model} supports at most {max_reference_images} reference images.")
|
||||
|
||||
# Build request body with required params
|
||||
reference_images: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"input_image" if index == 1 else f"input_image_{index}": self._convert_image_to_base64(reference_image)
|
||||
for index, reference_image in enumerate(images, start=1)
|
||||
}
|
||||
)
|
||||
request_body: Final[dict[str, Any]] = {
|
||||
"prompt": prompt,
|
||||
"image": image_b64,
|
||||
"model": model,
|
||||
**reference_images,
|
||||
**image_edit_optional_request_params,
|
||||
}
|
||||
|
||||
# Add mapped optional params (already filtered by map_openai_params)
|
||||
request_body.update(image_edit_optional_request_params)
|
||||
|
||||
# Return JSON body and empty files list (FLUX 2 doesn't use multipart)
|
||||
return request_body, []
|
||||
|
||||
def _convert_image_to_base64(self, image: Any) -> str:
|
||||
"""Convert image file to base64 string"""
|
||||
# Handle list of images (take first one)
|
||||
if isinstance(image, list):
|
||||
if len(image) == 0:
|
||||
raise ValueError("Empty image list provided")
|
||||
image = image[0]
|
||||
|
||||
if isinstance(image, BufferedReader):
|
||||
image_bytes = image.read()
|
||||
image.seek(0) # Reset file pointer for potential reuse
|
||||
|
|
@ -151,7 +144,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
"""
|
||||
Constructs a complete URL for Azure AI Foundry FLUX 2 image edits.
|
||||
|
||||
Uses the same /providers/blackforestlabs/v1/flux-2-pro endpoint as image generation.
|
||||
Uses the same model-specific BFL provider endpoint as image generation.
|
||||
"""
|
||||
api_base = AzureFoundryModelInfo.get_api_base(api_base)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
import litellm
|
||||
|
|
@ -10,6 +11,9 @@ from litellm.types.utils import ImageResponse
|
|||
def cost_calculator(
|
||||
model: str,
|
||||
image_response: Any,
|
||||
size: str | None = None,
|
||||
n: int | None = None,
|
||||
optional_params: Mapping[str, object] | None = None,
|
||||
) -> float:
|
||||
"""
|
||||
Azure AI image generation cost calculator
|
||||
|
|
@ -28,10 +32,29 @@ def cost_calculator(
|
|||
if token_based_cost is not None:
|
||||
return token_based_cost
|
||||
|
||||
num_images: Final = n if n is not None else len(image_response.data or ())
|
||||
output_cost_per_image: Final[float] = _model_info.get("output_cost_per_image") or 0.0
|
||||
num_images: int = 0
|
||||
if image_response.data:
|
||||
num_images = len(image_response.data)
|
||||
return output_cost_per_image * num_images
|
||||
if output_cost_per_image:
|
||||
return output_cost_per_image * num_images
|
||||
|
||||
model_cost: Final = litellm.model_cost[_model_info["key"]]
|
||||
input_cost_per_pixel: Final[float] = model_cost.get("input_cost_per_pixel") or 0.0
|
||||
if input_cost_per_pixel:
|
||||
from litellm.cost_calculator import default_image_cost_calculator
|
||||
|
||||
width: Final = optional_params.get("width") if optional_params else None
|
||||
height: Final = optional_params.get("height") if optional_params else None
|
||||
pixel_size: Final = (
|
||||
f"{width}x{height}"
|
||||
if type(width) is int and type(height) is int and width > 0 and height > 0
|
||||
else size or image_response.size
|
||||
)
|
||||
return default_image_cost_calculator(
|
||||
model=_model_info["key"],
|
||||
custom_llm_provider=litellm.LlmProviders.AZURE_AI.value,
|
||||
size=pixel_size,
|
||||
n=num_images,
|
||||
)
|
||||
return 0.0
|
||||
|
||||
raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}")
|
||||
|
|
|
|||
|
|
@ -1,18 +1,22 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.exceptions import BadRequestError, UnsupportedParamsError
|
||||
from litellm.llms.openai.image_generation import GPTImageGenerationConfig
|
||||
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
|
||||
|
||||
FLUX2_DROPPED_OPENAI_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = (
|
||||
"background",
|
||||
"moderation",
|
||||
"output_compression",
|
||||
"quality",
|
||||
"user",
|
||||
)
|
||||
|
||||
|
||||
class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
|
||||
"""
|
||||
Azure Foundry flux image generation config
|
||||
|
||||
From manual testing it follows the gpt-image-1 image generation config
|
||||
|
||||
(Azure Foundry does not have any docs on supported params at the time of writing)
|
||||
|
||||
From our test suite - following GPTImageGenerationConfig is working for this model
|
||||
"""
|
||||
"""Azure Foundry BFL API configuration for FLUX image generation."""
|
||||
|
||||
@staticmethod
|
||||
def get_flux2_image_generation_url(
|
||||
|
|
@ -25,11 +29,11 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
|
|||
|
||||
FLUX 2 models on Azure AI use a different URL pattern than standard Azure OpenAI:
|
||||
- Standard: /openai/deployments/{model}/images/generations
|
||||
- FLUX 2: /providers/blackforestlabs/v1/flux-2-pro
|
||||
- FLUX 2: /providers/blackforestlabs/v1/{model-path}
|
||||
|
||||
Args:
|
||||
api_base: Base URL (e.g., https://litellm-ci-cd-prod.services.ai.azure.com)
|
||||
model: Model name (e.g., flux.2-pro)
|
||||
model: Model name (e.g., FLUX.2-flex or FLUX.2-pro)
|
||||
api_version: API version (e.g., preview)
|
||||
|
||||
Returns:
|
||||
|
|
@ -47,9 +51,8 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
|
|||
return api_base
|
||||
return f"{api_base}?api-version={api_version}"
|
||||
|
||||
# Construct the FLUX 2 provider path
|
||||
# Model name flux.2-pro maps to endpoint flux-2-pro
|
||||
return f"{api_base}/providers/blackforestlabs/v1/flux-2-pro?api-version={api_version}"
|
||||
provider_model_path: Final = AzureFoundryFluxImageGenerationConfig.get_flux2_provider_model_path(model)
|
||||
return f"{api_base}/providers/blackforestlabs/v1/{provider_model_path}?api-version={api_version}"
|
||||
|
||||
@staticmethod
|
||||
def is_flux2_model(model: str) -> bool:
|
||||
|
|
@ -64,3 +67,90 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
|
|||
"""
|
||||
model_lower: Final = model.lower().replace(".", "-").replace("_", "-")
|
||||
return "flux-2" in model_lower or "flux2" in model_lower
|
||||
|
||||
@staticmethod
|
||||
def get_flux2_provider_model_path(model: str) -> str:
|
||||
normalized_model: Final = model.lower().replace(".", "-").replace("_", "-")
|
||||
return "flux-2-flex" if "flux-2-flex" in normalized_model else "flux-2-pro"
|
||||
|
||||
def get_supported_openai_params( # mutable-ok: inherited config contract returns a list
|
||||
self, model: str
|
||||
) -> list[OpenAIImageGenerationOptionalParams]:
|
||||
if not self.is_flux2_model(model):
|
||||
return super().get_supported_openai_params(model)
|
||||
return [ # mutable-ok: BaseImageGenerationConfig requires a list
|
||||
"n",
|
||||
"size",
|
||||
"output_format",
|
||||
"seed",
|
||||
"safety_tolerance",
|
||||
"aspect_ratio",
|
||||
"width",
|
||||
"height",
|
||||
"num_images",
|
||||
"guidance",
|
||||
"steps",
|
||||
*FLUX2_DROPPED_OPENAI_PARAMS,
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _map_parameter(name: str, value: object, model: str) -> tuple[tuple[str, object], ...]:
|
||||
if name in FLUX2_DROPPED_OPENAI_PARAMS:
|
||||
return ()
|
||||
if isinstance(value, str):
|
||||
if name in ("n", "num_images", "width", "height", "steps", "seed", "safety_tolerance"):
|
||||
return (("num_images" if name == "n" else name, int(value)),)
|
||||
if name == "guidance":
|
||||
return ((name, float(value)),)
|
||||
if name == "n":
|
||||
return (("num_images", value),)
|
||||
if name != "size":
|
||||
return ((name, value),)
|
||||
if str(value).lower() == "auto":
|
||||
return ()
|
||||
|
||||
try:
|
||||
width, height = (int(dimension) for dimension in str(value).lower().split("x"))
|
||||
except (TypeError, ValueError):
|
||||
raise BadRequestError(
|
||||
message=f"Invalid size format '{value}'. Expected 'WxH', for example '1024x1024'.",
|
||||
model=model,
|
||||
llm_provider="azure_ai",
|
||||
)
|
||||
return (("width", width), ("height", height))
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: Mapping[str, object],
|
||||
optional_params: Mapping[str, object],
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict[str, object]: # mutable-ok: inherited config contract returns a dict
|
||||
if not self.is_flux2_model(model):
|
||||
return super().map_openai_params(
|
||||
non_default_params=dict(non_default_params),
|
||||
optional_params=dict(optional_params),
|
||||
model=model,
|
||||
drop_params=drop_params,
|
||||
)
|
||||
supported_params: Final = self.get_supported_openai_params(model)
|
||||
unsupported_params: Final = tuple(name for name in non_default_params if name not in supported_params)
|
||||
if unsupported_params and not drop_params:
|
||||
raise UnsupportedParamsError(
|
||||
message=(
|
||||
f"Parameters {unsupported_params} are not supported for model {model}. "
|
||||
f"Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters."
|
||||
),
|
||||
model=model,
|
||||
llm_provider="azure_ai",
|
||||
)
|
||||
|
||||
mapped_params: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{
|
||||
mapped_name: mapped_value
|
||||
for name, value in non_default_params.items()
|
||||
if name in supported_params
|
||||
for mapped_name, mapped_value in self._map_parameter(name, value, model)
|
||||
}
|
||||
)
|
||||
return {**optional_params, **mapped_params} # mutable-ok: inherited config contract returns a dict
|
||||
|
|
|
|||
|
|
@ -1,5 +1,200 @@
|
|||
import math
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
_WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws", "wss": "wss", "ws": "ws"})
|
||||
DEEPGRAM_LISTEN_CALLBACK_PARAMS: Final = frozenset({"callback", "callback_method"})
|
||||
DEEPGRAM_LISTEN_STREAMING_PRICING_PREFIX: Final = "streaming/"
|
||||
DEEPGRAM_LISTEN_MULTILINGUAL_LANGUAGE: Final = "multi"
|
||||
DEEPGRAM_LISTEN_MULTILINGUAL_PRICING_SUFFIX: Final = "-multilingual"
|
||||
DEEPGRAM_LISTEN_ADDON_PRICING_PARAMS: Final = MappingProxyType(
|
||||
{
|
||||
"redact": "redact",
|
||||
"keyterm": "keyterm",
|
||||
"detect_entities": "detect_entities",
|
||||
"diarize": "diarize",
|
||||
"diarize_model": "diarize",
|
||||
}
|
||||
)
|
||||
_DISABLED_PARAM_VALUES: Final = frozenset({"", "false"})
|
||||
_SINGLE_VALUED_PARAMS: Final = frozenset({"model", "language"})
|
||||
|
||||
|
||||
class DeepgramException(BaseLLMException):
|
||||
pass
|
||||
|
||||
|
||||
def deepgram_listen_requested_model(query_string: str) -> str:
|
||||
return httpx.QueryParams(query_string).get("model") or DEEPGRAM_LISTEN_DEFAULT_MODEL
|
||||
|
||||
|
||||
def _first_occurrences(query_string: str) -> httpx.QueryParams:
|
||||
"""Authorization and pricing read the first ``model`` and ``language`` value; Deepgram must not see a second one."""
|
||||
items: Final = httpx.QueryParams(query_string).multi_items()
|
||||
return httpx.QueryParams(
|
||||
tuple(
|
||||
(key, value)
|
||||
for index, (key, value) in enumerate(items)
|
||||
if key not in _SINGLE_VALUED_PARAMS or all(earlier != key for earlier, _ in items[:index])
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def deepgram_listen_websocket_target(api_base: str | None, query_string: str) -> str:
|
||||
listen_url: Final = httpx.URL(f"{(api_base or DEEPGRAM_DEFAULT_API_BASE).rstrip('/')}/listen")
|
||||
websocket_url: Final = listen_url.copy_with(scheme=_WEBSOCKET_SCHEMES.get(listen_url.scheme, listen_url.scheme))
|
||||
params: Final = _first_occurrences(query_string)
|
||||
query: Final = params if params.get("model") else params.remove("model").add("model", DEEPGRAM_LISTEN_DEFAULT_MODEL)
|
||||
return f"{websocket_url}?{query}"
|
||||
|
||||
|
||||
def deepgram_listen_callback_params(query_string: str) -> tuple[str, ...]:
|
||||
return tuple(sorted(DEEPGRAM_LISTEN_CALLBACK_PARAMS.intersection(httpx.QueryParams(query_string).keys())))
|
||||
|
||||
|
||||
def deepgram_listen_model(upstream_url: str) -> str:
|
||||
models: Final = parse_qs(urlparse(upstream_url).query).get("model")
|
||||
return models[0] if models else DEEPGRAM_LISTEN_DEFAULT_MODEL
|
||||
|
||||
|
||||
def _param_enabled(values: Sequence[str]) -> bool:
|
||||
return any(value.strip().lower() not in _DISABLED_PARAM_VALUES for value in values)
|
||||
|
||||
|
||||
def deepgram_listen_pricing_model(upstream_url: str) -> str:
|
||||
"""Registry key, without the provider prefix, for the per-second base rate Deepgram bills a streaming session at:
|
||||
the multilingual streaming entry when ``language=multi``, otherwise the model's own streaming entry. Pre-recorded
|
||||
entries are never a substitute: Deepgram prices the two products differently."""
|
||||
streaming: Final = f"{DEEPGRAM_LISTEN_STREAMING_PRICING_PREFIX}{deepgram_listen_model(upstream_url)}"
|
||||
language: Final = parse_qs(urlparse(upstream_url).query).get("language", ("",))[0]
|
||||
if language.strip().lower() == DEEPGRAM_LISTEN_MULTILINGUAL_LANGUAGE:
|
||||
return f"{streaming}{DEEPGRAM_LISTEN_MULTILINGUAL_PRICING_SUFFIX}"
|
||||
return streaming
|
||||
|
||||
|
||||
def deepgram_listen_registry_key(upstream_url: str) -> str:
|
||||
return f"{LlmProviders.DEEPGRAM.value}/{deepgram_listen_pricing_model(upstream_url)}"
|
||||
|
||||
|
||||
def deepgram_listen_is_priced(upstream_url: str) -> bool:
|
||||
"""Only an exact registry hit counts: the cost calculator resolves a missing ``streaming/<model>`` row to the
|
||||
pre-recorded ``<model>`` row, which is not the rate Deepgram bills a WebSocket session at."""
|
||||
registry_key: Final = deepgram_listen_registry_key(upstream_url)
|
||||
try:
|
||||
model_info: Final = litellm.get_model_info(model=registry_key, custom_llm_provider=LlmProviders.DEEPGRAM.value)
|
||||
except Exception:
|
||||
return False
|
||||
return model_info["key"] == registry_key
|
||||
|
||||
|
||||
def deepgram_listen_addon_pricing_models(upstream_url: str) -> tuple[str, ...]:
|
||||
params: Final = parse_qs(urlparse(upstream_url).query)
|
||||
return tuple(
|
||||
sorted(
|
||||
frozenset(
|
||||
f"{DEEPGRAM_LISTEN_STREAMING_PRICING_PREFIX}{addon}"
|
||||
for param, addon in DEEPGRAM_LISTEN_ADDON_PRICING_PARAMS.items()
|
||||
if _param_enabled(params.get(param, ()))
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _channel_count(value: object) -> int | None:
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
return None
|
||||
return value if value >= 1 else None
|
||||
|
||||
|
||||
def _results_channel_count(frame: Mapping[str, object]) -> int | None:
|
||||
channel_index: Final = frame.get("channel_index")
|
||||
if not isinstance(channel_index, list) or len(channel_index) != 2:
|
||||
return None
|
||||
return _channel_count(channel_index[1])
|
||||
|
||||
|
||||
def _declared_channel_count(upstream_url: str) -> int | None:
|
||||
declared: Final = parse_qs(urlparse(upstream_url).query).get("channels")
|
||||
if not declared or not declared[0].isdigit():
|
||||
return None
|
||||
return _channel_count(int(declared[0]))
|
||||
|
||||
|
||||
def deepgram_listen_channel_count(websocket_messages: Sequence[Mapping[str, object]], upstream_url: str) -> int:
|
||||
metadata_channels: Final = tuple(
|
||||
channels
|
||||
for frame in websocket_messages
|
||||
if frame.get("type") == "Metadata"
|
||||
if (channels := _channel_count(frame.get("channels"))) is not None
|
||||
)
|
||||
if metadata_channels:
|
||||
return metadata_channels[-1]
|
||||
results_channels: Final = tuple(
|
||||
channels
|
||||
for frame in websocket_messages
|
||||
if frame.get("type") == "Results"
|
||||
if (channels := _results_channel_count(frame)) is not None
|
||||
)
|
||||
if results_channels:
|
||||
return max(results_channels)
|
||||
return _declared_channel_count(upstream_url) or 1
|
||||
|
||||
|
||||
def _seconds(value: object) -> float | None:
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
return None
|
||||
return float(value) if math.isfinite(value) and value >= 0 else None
|
||||
|
||||
|
||||
def _results_frame_end(frame: Mapping[str, object]) -> float | None:
|
||||
start: Final = _seconds(frame.get("start"))
|
||||
duration: Final = _seconds(frame.get("duration"))
|
||||
return None if start is None or duration is None else start + duration
|
||||
|
||||
|
||||
def _final_transcript(frame: Mapping[str, object]) -> str | None:
|
||||
if frame.get("is_final") is not True:
|
||||
return None
|
||||
channel: Final = frame.get("channel")
|
||||
alternatives: Final = channel.get("alternatives") if isinstance(channel, Mapping) else None
|
||||
first: Final = alternatives[0] if isinstance(alternatives, list) and alternatives else None
|
||||
transcript: Final = first.get("transcript") if isinstance(first, Mapping) else None
|
||||
return transcript if isinstance(transcript, str) and transcript else None
|
||||
|
||||
|
||||
def deepgram_listen_audio_seconds(websocket_messages: Sequence[Mapping[str, object]]) -> float:
|
||||
metadata_durations: Final = tuple(
|
||||
duration
|
||||
for frame in websocket_messages
|
||||
if frame.get("type") == "Metadata"
|
||||
if (duration := _seconds(frame.get("duration"))) is not None and duration > 0
|
||||
)
|
||||
if metadata_durations:
|
||||
return metadata_durations[-1]
|
||||
return max(
|
||||
(
|
||||
end
|
||||
for frame in websocket_messages
|
||||
if frame.get("type") == "Results"
|
||||
if (end := _results_frame_end(frame)) is not None
|
||||
),
|
||||
default=0.0,
|
||||
)
|
||||
|
||||
|
||||
def deepgram_listen_transcript(websocket_messages: Sequence[Mapping[str, object]]) -> str:
|
||||
return " ".join(
|
||||
transcript
|
||||
for frame in websocket_messages
|
||||
if frame.get("type") == "Results"
|
||||
if (transcript := _final_transcript(frame)) is not None
|
||||
)
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast, get_type_hints, ove
|
|||
|
||||
import httpx
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
handle_messages_with_content_list_to_str_conversion,
|
||||
|
|
@ -20,16 +21,37 @@ from litellm.llms.openai.chat.gpt_transformation import (
|
|||
OpenAIChatCompletionStreamingHandler,
|
||||
OpenAIGPTConfig,
|
||||
)
|
||||
from litellm.router_utils.reasoning_effort_capability import (
|
||||
declared_reasoning_efforts_for_model,
|
||||
nearest_declared_reasoning_effort,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.mistral import MistralThinkingBlock, MistralToolCallMessage
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import ModelResponse, ModelResponseStream
|
||||
from litellm.utils import convert_to_model_response_object
|
||||
from litellm.utils import convert_to_model_response_object, supports_reasoning
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
|
||||
def _accepted_reasoning_effort(model: str, requested: str, custom_llm_provider: str) -> str:
|
||||
declared: Final = declared_reasoning_efforts_for_model(model, custom_llm_provider)
|
||||
if declared is None:
|
||||
return requested
|
||||
accepted: Final = nearest_declared_reasoning_effort(requested, declared)
|
||||
if accepted != requested:
|
||||
verbose_logger.debug(
|
||||
"%s: %s takes reasoning_effort %s, sending %s in place of %s",
|
||||
custom_llm_provider,
|
||||
model,
|
||||
declared,
|
||||
accepted,
|
||||
requested,
|
||||
)
|
||||
return accepted
|
||||
|
||||
|
||||
class MistralConfig(OpenAIGPTConfig):
|
||||
"""
|
||||
Reference: https://docs.mistral.ai/api/
|
||||
|
|
@ -86,8 +108,16 @@ class MistralConfig(OpenAIGPTConfig):
|
|||
def get_config(cls):
|
||||
return super().get_config()
|
||||
|
||||
@property
|
||||
def custom_llm_provider(self) -> str:
|
||||
return "mistral"
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[str]:
|
||||
supported_params: Final = [
|
||||
is_magistral: Final = "magistral" in model.lower()
|
||||
accepts_reasoning_effort: Final = is_magistral or supports_reasoning(
|
||||
model=model, custom_llm_provider=self.custom_llm_provider
|
||||
)
|
||||
return [
|
||||
"stream",
|
||||
"temperature",
|
||||
"top_p",
|
||||
|
|
@ -99,14 +129,10 @@ class MistralConfig(OpenAIGPTConfig):
|
|||
"stop",
|
||||
"response_format",
|
||||
"parallel_tool_calls",
|
||||
*(("thinking",) if is_magistral else ()),
|
||||
*(("reasoning_effort",) if accepts_reasoning_effort else ()),
|
||||
]
|
||||
|
||||
# Add reasoning support for magistral models
|
||||
if "magistral" in model.lower():
|
||||
supported_params.extend(["thinking", "reasoning_effort"])
|
||||
|
||||
return supported_params
|
||||
|
||||
def _map_tool_choice(self, tool_choice: str) -> str:
|
||||
if tool_choice == "auto" or tool_choice == "none":
|
||||
return tool_choice
|
||||
|
|
@ -171,10 +197,9 @@ class MistralConfig(OpenAIGPTConfig):
|
|||
optional_params["extra_body"] = {"random_seed": value}
|
||||
if param == "response_format":
|
||||
optional_params["response_format"] = value
|
||||
if param == "reasoning_effort" and "magistral" in model.lower():
|
||||
# Flag that we need to add reasoning system prompt
|
||||
optional_params["_add_reasoning_prompt"] = True
|
||||
if param == "thinking" and "magistral" in model.lower():
|
||||
if param == "reasoning_effort" and "magistral" not in model.lower():
|
||||
optional_params["reasoning_effort"] = _accepted_reasoning_effort(model, value, self.custom_llm_provider)
|
||||
if param in ("reasoning_effort", "thinking") and "magistral" in model.lower():
|
||||
# Flag that we need to add reasoning system prompt
|
||||
optional_params["_add_reasoning_prompt"] = True
|
||||
if param == "parallel_tool_calls":
|
||||
|
|
@ -534,11 +559,13 @@ class MistralConfig(OpenAIGPTConfig):
|
|||
if "magistral" in model.lower() and optional_params.get("_add_reasoning_prompt", False):
|
||||
messages = self._add_reasoning_system_prompt_if_needed(messages, optional_params)
|
||||
|
||||
upstream_params: Final = {key: value for key, value in optional_params.items() if key != "client_metadata"}
|
||||
|
||||
# Call parent transform_request which handles _transform_messages
|
||||
return super().transform_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
optional_params=upstream_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -82,7 +82,14 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig):
|
|||
)
|
||||
|
||||
# set optional params
|
||||
image_response.size = image_response.size or optional_params.get("size", "1024x1024")
|
||||
width: Final = optional_params.get("width")
|
||||
height: Final = optional_params.get("height")
|
||||
requested_size: Final = (
|
||||
f"{width}x{height}"
|
||||
if isinstance(width, int) and isinstance(height, int)
|
||||
else optional_params.get("size", "1024x1024")
|
||||
)
|
||||
image_response.size = image_response.size or requested_size
|
||||
image_response.quality = image_response.quality or optional_params.get("quality", "high")
|
||||
image_response.output_format = image_response.output_format or optional_params.get("output_format", "png")
|
||||
|
||||
|
|
|
|||
|
|
@ -101,7 +101,8 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig):
|
|||
extra_body: dict[str, object] | None = None,
|
||||
) -> tuple[str, dict]:
|
||||
encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id")
|
||||
url: Final = f"{api_base}/{encoded_vector_store_id}/search"
|
||||
base_url, query_separator, query_string = api_base.partition("?")
|
||||
url: Final = f"{base_url}/{encoded_vector_store_id}/search{query_separator}{query_string}"
|
||||
typed_request_body: Final = VectorStoreSearchRequest(
|
||||
query=query,
|
||||
filters=vector_store_search_optional_params.get("filters", None),
|
||||
|
|
|
|||
|
|
@ -0,0 +1,7 @@
|
|||
from litellm.llms.mistral.chat.transformation import MistralConfig
|
||||
|
||||
|
||||
class VertexAIMistralConfig(MistralConfig):
|
||||
@property
|
||||
def custom_llm_provider(self) -> str:
|
||||
return "vertex_ai"
|
||||
|
|
@ -10820,6 +10820,25 @@
|
|||
"/v1/images/generations"
|
||||
]
|
||||
},
|
||||
"azure_ai/FLUX.2-flex": {
|
||||
"input_cost_per_pixel": 5e-08,
|
||||
"litellm_provider": "azure_ai",
|
||||
"max_input_tokens": 32000,
|
||||
"max_tokens": 32000,
|
||||
"mode": "image_generation",
|
||||
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"image"
|
||||
]
|
||||
},
|
||||
"azure_ai/FW-DeepSeek-V3.2": {
|
||||
"deprecation_date": "2027-07-01",
|
||||
"cache_read_input_token_cost": 3.1e-07,
|
||||
|
|
@ -16690,6 +16709,46 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"dashscope/qwen3.8-flash": {
|
||||
"cache_creation_input_token_cost": 2e-07,
|
||||
"cache_read_input_token_cost": 1.6e-08,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "dashscope",
|
||||
"max_input_tokens": 991808,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.7e-07,
|
||||
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"dashscope/qwen3.8-omni-flash": {
|
||||
"cache_read_input_token_cost": 1.6e-08,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "dashscope",
|
||||
"max_input_tokens": 991808,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.7e-07,
|
||||
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
|
||||
"supports_audio_input": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"dashscope/qwq-plus": {
|
||||
"input_cost_per_token": 8e-07,
|
||||
"litellm_provider": "dashscope",
|
||||
|
|
@ -18594,6 +18653,46 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"qwen_ai_platform/qwen3.8-flash": {
|
||||
"cache_creation_input_token_cost": 2e-07,
|
||||
"cache_read_input_token_cost": 1.6e-08,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "qwen_ai_platform",
|
||||
"max_input_tokens": 991808,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.7e-07,
|
||||
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"qwen_ai_platform/qwen3.8-omni-flash": {
|
||||
"cache_read_input_token_cost": 1.6e-08,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "qwen_ai_platform",
|
||||
"max_input_tokens": 991808,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.7e-07,
|
||||
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
|
||||
"supports_audio_input": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"qwen_ai_platform/qwq-plus": {
|
||||
"input_cost_per_token": 8e-07,
|
||||
"litellm_provider": "qwen_ai_platform",
|
||||
|
|
@ -20754,6 +20853,96 @@
|
|||
"/v1/audio/transcriptions"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/nova-3": {
|
||||
"input_cost_per_second": 8e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0048/60 seconds = $0.00008000 per second",
|
||||
"note": "Nova-3 monolingual streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0048
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/nova-3-multilingual": {
|
||||
"input_cost_per_second": 9.667e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0058/60 seconds = $0.00009667 per second",
|
||||
"note": "Nova-3 multilingual (language=multi) streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0058
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/redact": {
|
||||
"input_cost_per_second": 3.333e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0020/60 seconds = $0.00003333 per second",
|
||||
"note": "Redaction add-on (redact query param), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.002
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/keyterm": {
|
||||
"input_cost_per_second": 2.167e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0013/60 seconds = $0.00002167 per second",
|
||||
"note": "Keyterm Prompting add-on (keyterm query param), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0013
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/detect_entities": {
|
||||
"input_cost_per_second": 2.833e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0017/60 seconds = $0.00002833 per second",
|
||||
"note": "Entity Detection add-on (detect_entities query param), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0017
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/diarize": {
|
||||
"input_cost_per_second": 3.333e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0020/60 seconds = $0.00003333 per second",
|
||||
"note": "Speaker Diarization add-on (diarize / diarize_model query params), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.002
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/whisper": {
|
||||
"input_cost_per_second": 0.0001,
|
||||
"litellm_provider": "deepgram",
|
||||
|
|
@ -22090,8 +22279,8 @@
|
|||
"embed-english-light-v3.0": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "cohere",
|
||||
"max_input_tokens": 1024,
|
||||
"max_tokens": 1024,
|
||||
"max_input_tokens": 512,
|
||||
"max_tokens": 512,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0
|
||||
},
|
||||
|
|
@ -22108,8 +22297,8 @@
|
|||
"input_cost_per_image": 0.0001,
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "cohere",
|
||||
"max_input_tokens": 1024,
|
||||
"max_tokens": 1024,
|
||||
"max_input_tokens": 512,
|
||||
"max_tokens": 512,
|
||||
"metadata": {
|
||||
"notes": "'supports_image_input' is a deprecated field. Use 'supports_embedding_image_input' instead."
|
||||
},
|
||||
|
|
@ -22130,8 +22319,8 @@
|
|||
"embed-multilingual-v3.0": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "cohere",
|
||||
"max_input_tokens": 1024,
|
||||
"max_tokens": 1024,
|
||||
"max_input_tokens": 512,
|
||||
"max_tokens": 512,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_embedding_image_input": true
|
||||
|
|
@ -22139,8 +22328,8 @@
|
|||
"embed-multilingual-light-v3.0": {
|
||||
"input_cost_per_token": 0.0001,
|
||||
"litellm_provider": "cohere",
|
||||
"max_input_tokens": 1024,
|
||||
"max_tokens": 1024,
|
||||
"max_input_tokens": 512,
|
||||
"max_tokens": 512,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_embedding_image_input": true
|
||||
|
|
@ -36998,6 +37187,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37079,6 +37272,15 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"minimal",
|
||||
"low",
|
||||
"medium",
|
||||
"high",
|
||||
"xhigh",
|
||||
"max"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/zai-glm-5-2",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37096,6 +37298,11 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/zai-glm-5-3",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37113,6 +37320,11 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/zai-glm-5-3",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37130,6 +37342,11 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/zai-glm-5-3",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37147,6 +37364,15 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"minimal",
|
||||
"low",
|
||||
"medium",
|
||||
"high",
|
||||
"xhigh",
|
||||
"max"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/zai-glm-5-2",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37442,6 +37668,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37502,6 +37732,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37519,6 +37753,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37552,6 +37790,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37583,6 +37825,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-07,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -57910,14 +58156,14 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"bedrock_mantle/openai.gpt-5.6-sol": {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 1.1e-05,
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 1.375e-05,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
|
||||
"output_cost_per_token": 3.3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.95e-05,
|
||||
"input_cost_per_token": 4.4e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 8.8e-06,
|
||||
"cache_creation_input_token_cost": 5.5e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 1.1e-05,
|
||||
"cache_read_input_token_cost": 4.4e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 8.8e-07,
|
||||
"output_cost_per_token": 2.2e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 3.3e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.012,
|
||||
"search_context_size_low": 0.012,
|
||||
|
|
@ -59928,6 +60174,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-07,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -63253,6 +63503,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -63270,6 +63524,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -63287,6 +63545,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -63304,6 +63566,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-07,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -65816,9 +66082,9 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/z-ai/glm-5.3": {
|
||||
"input_cost_per_token": 1.4e-06,
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"cache_read_input_token_cost": 2.6e-07,
|
||||
"input_cost_per_token": 9.1e-07,
|
||||
"output_cost_per_token": 2.86e-06,
|
||||
"cache_read_input_token_cost": 1.69e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1310720,
|
||||
"max_output_tokens": 943717,
|
||||
|
|
@ -70539,14 +70805,14 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/~deepseek/deepseek-flash-latest": {
|
||||
"cache_read_input_token_cost": 1.5e-08,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"cache_read_input_token_cost": 4.2e-09,
|
||||
"input_cost_per_token": 1.4e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 393216,
|
||||
"max_tokens": 393216,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-07,
|
||||
"output_cost_per_token": 4.2e-07,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_audio_input": false,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -70831,14 +71097,14 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/~z-ai/glm-latest": {
|
||||
"cache_read_input_token_cost": 1.5e-07,
|
||||
"cache_read_input_token_cost": 1.46625e-07,
|
||||
"input_cost_per_token": 9e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1310720,
|
||||
"max_output_tokens": 235929,
|
||||
"max_tokens": 235929,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-06,
|
||||
"output_cost_per_token": 2.805e-06,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_audio_input": false,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -73982,14 +74248,14 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/tencent/hy3": {
|
||||
"cache_read_input_token_cost": 3.3e-08,
|
||||
"input_cost_per_token": 1.32e-07,
|
||||
"cache_read_input_token_cost": 2.0625e-08,
|
||||
"input_cost_per_token": 8.25e-08,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5.28e-07,
|
||||
"output_cost_per_token": 3.3e-07,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_audio_input": false,
|
||||
"supports_function_calling": true,
|
||||
|
|
|
|||
|
|
@ -1,8 +1,15 @@
|
|||
import sys
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS: Final = 600.0
|
||||
|
||||
_SECONDS: Final = TypeAdapter(float)
|
||||
_NO_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def resolve_pass_through_request_timeout(
|
||||
endpoint_timeout: float | None = None,
|
||||
|
|
@ -31,26 +38,41 @@ def resolve_pass_through_request_timeout(
|
|||
|
||||
|
||||
def resolve_llm_passthrough_timeout(
|
||||
kwargs: dict | None = None,
|
||||
litellm_params: dict | None = None,
|
||||
router_timeout: float | None = None,
|
||||
kwargs: Mapping[str, object] | None = None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
router_timeout: float | str | None = None,
|
||||
router_stream_timeout: float | str | None = None,
|
||||
) -> float:
|
||||
"""
|
||||
Resolve upstream httpx timeout for SDK native passthrough (e.g. Bedrock /converse).
|
||||
Resolve upstream httpx timeout for SDK native passthrough (e.g. Bedrock /converse,
|
||||
Anthropic /v1/messages).
|
||||
|
||||
Precedence: kwargs timeout/request_timeout -> litellm_params timeout/request_timeout
|
||||
-> router_timeout -> general_settings.pass_through_request_timeout -> 600s default.
|
||||
Non-streaming precedence: kwargs timeout/request_timeout -> litellm_params
|
||||
timeout/request_timeout -> router_timeout -> general_settings.pass_through_request_timeout
|
||||
-> 600s default.
|
||||
|
||||
Streaming (``kwargs["stream"]`` truthy) resolves ``stream_timeout`` at every level before
|
||||
any generic timeout, matching ``Router._get_stream_timeout`` on the completion route:
|
||||
kwargs stream_timeout -> litellm_params stream_timeout -> router_stream_timeout, then the
|
||||
non-streaming chain above.
|
||||
|
||||
Only the first set value is validated as seconds, so a value in a lower-precedence
|
||||
field never fails the call.
|
||||
"""
|
||||
kwargs = kwargs or {}
|
||||
litellm_params = litellm_params or {}
|
||||
|
||||
for source in (kwargs, litellm_params):
|
||||
for key in ("timeout", "request_timeout"):
|
||||
val = source.get(key)
|
||||
if val is not None:
|
||||
return float(val)
|
||||
|
||||
if router_timeout is not None:
|
||||
return float(router_timeout)
|
||||
|
||||
return resolve_pass_through_request_timeout()
|
||||
request: Final = kwargs if kwargs is not None else _NO_PARAMS
|
||||
deployment: Final = litellm_params if litellm_params is not None else _NO_PARAMS
|
||||
stream_candidates: Final = (
|
||||
(request.get("stream_timeout"), deployment.get("stream_timeout"), router_stream_timeout)
|
||||
if request.get("stream")
|
||||
else ()
|
||||
)
|
||||
candidates: Final = (
|
||||
*stream_candidates,
|
||||
request.get("timeout"),
|
||||
request.get("request_timeout"),
|
||||
deployment.get("timeout"),
|
||||
deployment.get("request_timeout"),
|
||||
router_timeout,
|
||||
)
|
||||
winner: Final = next((val for val in candidates if val is not None), None)
|
||||
return resolve_pass_through_request_timeout() if winner is None else _SECONDS.validate_python(winner)
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ _PASS_THROUGH_PROTECTED_HEADERS: Final[frozenset] = frozenset(
|
|||
"api-key",
|
||||
"x-api-key",
|
||||
"x-goog-api-key",
|
||||
"ocp-apim-subscription-key",
|
||||
"host",
|
||||
"content-length",
|
||||
"accept-encoding",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,38 @@
|
|||
"""Per-worker cache of stored BYOK credentials, keyed so peer workers can evict it over the auth cache pub/sub."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE, MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS
|
||||
|
||||
_CACHE_KEY_PREFIX: Final = "mcp_byok_credential"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CachedByokCredential:
|
||||
credential: str | None
|
||||
|
||||
|
||||
byok_credential_cache: Final = InMemoryCache(
|
||||
max_size_in_memory=MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE,
|
||||
default_ttl=MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def byok_credential_cache_key(user_id: str, server_id: str) -> str:
|
||||
return f"{_CACHE_KEY_PREFIX}:{user_id}:{server_id}"
|
||||
|
||||
|
||||
def get_cached_byok_credential(user_id: str, server_id: str) -> CachedByokCredential | None:
|
||||
cached: Final = byok_credential_cache.get_cache( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # InMemoryCache is untyped
|
||||
byok_credential_cache_key(user_id, server_id)
|
||||
)
|
||||
return cached if isinstance(cached, CachedByokCredential) else None
|
||||
|
||||
|
||||
def cache_byok_credential(user_id: str, server_id: str, credential: str | None) -> None:
|
||||
byok_credential_cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
|
||||
byok_credential_cache_key(user_id, server_id),
|
||||
CachedByokCredential(credential=credential),
|
||||
)
|
||||
|
|
@ -865,7 +865,7 @@ async def byok_token(
|
|||
_invalidate_byok_cred_cache,
|
||||
)
|
||||
|
||||
_invalidate_byok_cred_cache(user_id, server_id)
|
||||
await _invalidate_byok_cred_cache(user_id, server_id)
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.error(
|
||||
"byok_token: failed to store user credential for user=%s server=%s: %s",
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ from litellm.proxy._types import (
|
|||
MCPApprovalStatus,
|
||||
MCPEnvVar,
|
||||
MCPEnvVarScope,
|
||||
MCPServerUserCredentialListItem,
|
||||
MCPSubmissionsSummary,
|
||||
NewMCPServerRequest,
|
||||
SpecialMCPServerName,
|
||||
|
|
@ -1504,6 +1505,37 @@ async def get_user_oauth_credential(
|
|||
return _parse_oauth_payload(decoded)
|
||||
|
||||
|
||||
def _server_user_credential_item(
|
||||
row: "prisma_db_models.LiteLLM_MCPUserCredentials",
|
||||
) -> MCPServerUserCredentialListItem:
|
||||
oauth_payload: Final = _decode_oauth_payload(row.credential_b64)
|
||||
if oauth_payload is None:
|
||||
return MCPServerUserCredentialListItem(
|
||||
user_id=row.user_id,
|
||||
credential_type="byok",
|
||||
updated_at=row.updated_at.isoformat(),
|
||||
)
|
||||
return MCPServerUserCredentialListItem(
|
||||
user_id=row.user_id,
|
||||
credential_type="oauth2",
|
||||
expires_at=oauth_payload.get("expires_at"),
|
||||
connected_at=oauth_payload.get("connected_at"),
|
||||
updated_at=row.updated_at.isoformat(),
|
||||
)
|
||||
|
||||
|
||||
async def list_server_user_credentials(
|
||||
prisma_client: PrismaClient,
|
||||
server_id: str,
|
||||
) -> tuple[MCPServerUserCredentialListItem, ...]:
|
||||
"""Every user's stored credential for one server, typed but without the secret, for admins."""
|
||||
rows: Final = await _db_find_user_credential_rows(
|
||||
prisma_client,
|
||||
{"server_id": server_id}, # mutable-ok: prisma where-inputs must be plain dicts
|
||||
)
|
||||
return tuple(_server_user_credential_item(row) for row in rows)
|
||||
|
||||
|
||||
async def list_user_oauth_credentials(
|
||||
prisma_client: PrismaClient,
|
||||
user_id: str,
|
||||
|
|
|
|||
|
|
@ -295,12 +295,15 @@ class MCPPerUserTokenCache:
|
|||
)
|
||||
|
||||
async def delete(self, user_id: str, server_id: str) -> None:
|
||||
"""Invalidate the cached token (removes from both in-memory and Redis layers)."""
|
||||
"""Invalidate the cached token in Redis, here, and in every peer worker's in-memory layer."""
|
||||
try:
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( # noqa: PLC0415 # proxy import cycle
|
||||
evict_and_broadcast,
|
||||
)
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
key: Final = self._cache_key(user_id, server_id)
|
||||
await user_api_key_cache.async_delete_cache(key)
|
||||
await evict_and_broadcast((key,), user_api_key_cache)
|
||||
except Exception as exc:
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.delete failed for user=%s server=%s: %s",
|
||||
|
|
|
|||
|
|
@ -28,7 +28,10 @@ from starlette.types import Message, Receive, Scope, Send
|
|||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
|
||||
from litellm.constants import (
|
||||
MAXIMUM_TRACEBACK_LINES_TO_LOG,
|
||||
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
|
|
@ -38,6 +41,12 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
|||
MCPRequestHandler,
|
||||
_is_mcp_admitted_user_subject,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.byok_credential_cache import (
|
||||
byok_credential_cache,
|
||||
byok_credential_cache_key,
|
||||
cache_byok_credential,
|
||||
get_cached_byok_credential,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
get_request_base_url,
|
||||
)
|
||||
|
|
@ -82,6 +91,9 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
||||
publish_auth_cache_invalidation,
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
LiteLLMProxyRequestSetup,
|
||||
get_chain_id_from_headers,
|
||||
|
|
@ -91,6 +103,7 @@ from litellm.types.mcp import (
|
|||
MCPGatewaySession,
|
||||
MCPGatewaySessionGroupCount,
|
||||
MCPGatewaySessionsResponse,
|
||||
MCPGatewaySessionsTerminateResponse,
|
||||
MCPSpecVersion,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
|
||||
|
|
@ -102,13 +115,6 @@ if TYPE_CHECKING:
|
|||
|
||||
from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload
|
||||
|
||||
# Short-lived in-memory cache for BYOK credentials.
|
||||
# Keyed by (user_id, server_id); value is (credential_or_None, monotonic_timestamp).
|
||||
# Storing the credential value (not just a bool) means _get_byok_credential and
|
||||
# _check_byok_credential share a single DB round-trip per TTL window.
|
||||
_byok_cred_cache: Final[dict[tuple[str, str], tuple[str | None, float]]] = {}
|
||||
_BYOK_CRED_CACHE_TTL: Final = 60 # seconds
|
||||
_BYOK_CRED_CACHE_MAX_SIZE: Final = 4096 # cap to prevent unbounded growth
|
||||
_STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS: Final = 30 * 60
|
||||
# Upper bound on concurrent stateful sessions a single caller may hold. Each
|
||||
# `initialize` creates a session that survives until the idle timeout, so
|
||||
|
|
@ -127,20 +133,11 @@ _MCP_TRANSPORT_SPAN_SCOPE_KEY: Final = "litellm_otel_transport_span"
|
|||
_MCP_DESTINATIONS_SCOPE_KEY: Final = "litellm_otel_request_destinations"
|
||||
|
||||
|
||||
def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
|
||||
"""Remove a (user_id, server_id) entry from the BYOK credential cache.
|
||||
|
||||
Call this after storing or deleting a credential so subsequent calls
|
||||
see the fresh value rather than a stale cached result.
|
||||
"""
|
||||
_byok_cred_cache.pop((user_id, server_id), None)
|
||||
|
||||
|
||||
def _write_byok_cred_cache(user_id: str, server_id: str, credential: str | None) -> None:
|
||||
"""Write a credential value to the cache, evicting all entries if at capacity."""
|
||||
if len(_byok_cred_cache) >= _BYOK_CRED_CACHE_MAX_SIZE:
|
||||
_byok_cred_cache.clear()
|
||||
_byok_cred_cache[(user_id, server_id)] = (credential, time.monotonic())
|
||||
async def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
|
||||
"""Drop a stored-or-deleted BYOK credential from this worker's cache and from every peer worker's."""
|
||||
cache_key: Final = byok_credential_cache_key(user_id, server_id)
|
||||
byok_credential_cache.delete_cache(cache_key)
|
||||
await publish_auth_cache_invalidation(cache_key=cache_key)
|
||||
|
||||
|
||||
# Check if MCP is available
|
||||
|
|
@ -618,6 +615,7 @@ if MCP_AVAILABLE:
|
|||
_stateful_session_locks: Final[dict[str, asyncio.Lock]] = {}
|
||||
_stateful_session_active_request_counts: Final[dict[str, int]] = {}
|
||||
_stateful_session_client_info: Final[dict[str, Implementation]] = {} # mutable-ok: cleared on session teardown
|
||||
_admin_terminated_session_ids: Final[dict[str, float]] = {} # mutable-ok: admin-closed id -> last replay
|
||||
|
||||
class _TerminableTransport(Protocol):
|
||||
async def terminate(self) -> None: ...
|
||||
|
|
@ -689,6 +687,7 @@ if MCP_AVAILABLE:
|
|||
for session_id in list(_stateful_session_auth_context_last_seen):
|
||||
if session_id not in _stateful_session_auth_contexts:
|
||||
_remove_stateful_session_tracking(session_id)
|
||||
_forget_expired_admin_terminated_session_ids(now)
|
||||
|
||||
async def _enforce_stateful_session_cap_for_owner(owner: str) -> bool:
|
||||
"""
|
||||
|
|
@ -2811,35 +2810,28 @@ if MCP_AVAILABLE:
|
|||
mcp_server: MCPServer,
|
||||
user_api_key_auth: UserAPIKeyAuth | None,
|
||||
) -> str | None:
|
||||
"""Retrieve the stored BYOK credential for a user+server pair.
|
||||
|
||||
Uses the shared _byok_cred_cache to avoid a DB round-trip on every
|
||||
tool call within the TTL window.
|
||||
"""
|
||||
"""Retrieve the stored BYOK credential for a user+server pair, served from the worker cache within its TTL."""
|
||||
if not mcp_server.is_byok:
|
||||
return None
|
||||
user_id: Final = (user_api_key_auth.user_id if user_api_key_auth else None) or ""
|
||||
if not user_id:
|
||||
return None
|
||||
|
||||
cache_key: Final = (user_id, mcp_server.server_id)
|
||||
cached: Final = _byok_cred_cache.get(cache_key)
|
||||
cached: Final = get_cached_byok_credential(user_id, mcp_server.server_id)
|
||||
if cached is not None:
|
||||
credential, ts = cached
|
||||
if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL:
|
||||
return credential
|
||||
return cached.credential
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import get_user_credential
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
return None
|
||||
credential = await get_user_credential(
|
||||
credential: Final = await get_user_credential(
|
||||
prisma_client=prisma_client,
|
||||
user_id=user_id,
|
||||
server_id=mcp_server.server_id,
|
||||
)
|
||||
_write_byok_cred_cache(user_id, mcp_server.server_id, credential)
|
||||
cache_byok_credential(user_id, mcp_server.server_id, credential)
|
||||
return credential
|
||||
|
||||
async def _check_byok_credential(
|
||||
|
|
@ -2868,27 +2860,23 @@ if MCP_AVAILABLE:
|
|||
headers={"WWW-Authenticate": get_byok_www_authenticate()},
|
||||
)
|
||||
|
||||
# Check shared credential cache before hitting the DB.
|
||||
cache_key: Final = (user_id, mcp_server.server_id)
|
||||
cached: Final = _byok_cred_cache.get(cache_key)
|
||||
cached: Final = get_cached_byok_credential(user_id, mcp_server.server_id)
|
||||
if cached is not None:
|
||||
cached_cred, ts = cached
|
||||
if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL:
|
||||
if cached_cred is None:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": "byok_auth_required",
|
||||
"server_id": mcp_server.server_id,
|
||||
"server_name": mcp_server.server_name or mcp_server.name,
|
||||
"message": (
|
||||
"No stored credential found for this BYOK server. "
|
||||
"Complete the OAuth authorization flow to provide your API key."
|
||||
),
|
||||
},
|
||||
headers={"WWW-Authenticate": get_byok_www_authenticate()},
|
||||
)
|
||||
return
|
||||
if cached.credential is None:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": "byok_auth_required",
|
||||
"server_id": mcp_server.server_id,
|
||||
"server_name": mcp_server.server_name or mcp_server.name,
|
||||
"message": (
|
||||
"No stored credential found for this BYOK server. "
|
||||
"Complete the OAuth authorization flow to provide your API key."
|
||||
),
|
||||
},
|
||||
headers={"WWW-Authenticate": get_byok_www_authenticate()},
|
||||
)
|
||||
return
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import get_user_credential
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
|
@ -2912,7 +2900,7 @@ if MCP_AVAILABLE:
|
|||
user_id=user_id,
|
||||
server_id=mcp_server.server_id,
|
||||
)
|
||||
_write_byok_cred_cache(user_id, mcp_server.server_id, credential)
|
||||
cache_byok_credential(user_id, mcp_server.server_id, credential)
|
||||
if credential is None:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
|
|
@ -3850,7 +3838,7 @@ if MCP_AVAILABLE:
|
|||
client_info: Final = _stateful_session_client_info.get(session_id)
|
||||
key_auth: Final = auth_user.user_api_key_auth
|
||||
return MCPGatewaySession(
|
||||
session_id_prefix=session_id[:8],
|
||||
session_id_prefix=session_id[:MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH],
|
||||
client_name=client_info.name if client_info is not None else None,
|
||||
client_version=client_info.version if client_info is not None else None,
|
||||
user_id=key_auth.user_id if key_auth is not None else None,
|
||||
|
|
@ -3885,6 +3873,72 @@ if MCP_AVAILABLE:
|
|||
sessions=sessions,
|
||||
)
|
||||
|
||||
def _session_matches_admin_selector(
|
||||
session_id: str,
|
||||
auth_user: MCPAuthenticatedUser,
|
||||
session_id_prefix: str | None,
|
||||
user_id: str | None,
|
||||
) -> bool:
|
||||
if session_id_prefix is not None and not session_id.startswith(session_id_prefix):
|
||||
return False
|
||||
if user_id is None:
|
||||
return True
|
||||
key_auth: Final = auth_user.user_api_key_auth
|
||||
return key_auth is not None and key_auth.user_id == user_id
|
||||
|
||||
def _forget_expired_admin_terminated_session_ids(now: float) -> None:
|
||||
for session_id in [
|
||||
session_id
|
||||
for session_id, last_replayed in _admin_terminated_session_ids.items()
|
||||
if now - last_replayed >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS
|
||||
]:
|
||||
del _admin_terminated_session_ids[session_id]
|
||||
|
||||
def _is_admin_terminated_session_id(session_id: str, now: float) -> bool:
|
||||
last_replayed: Final = _admin_terminated_session_ids.get(session_id)
|
||||
if last_replayed is None:
|
||||
return False
|
||||
if now - last_replayed >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS:
|
||||
del _admin_terminated_session_ids[session_id]
|
||||
return False
|
||||
_admin_terminated_session_ids[session_id] = now
|
||||
return True
|
||||
|
||||
async def terminate_mcp_gateway_sessions(
|
||||
*,
|
||||
session_id_prefix: str | None = None,
|
||||
user_id: str | None = None,
|
||||
) -> MCPGatewaySessionsTerminateResponse:
|
||||
"""Force-close every live stateful session on this worker matching the selector.
|
||||
|
||||
The transport is terminated (open streams close), all per-session
|
||||
tracking is dropped, and the id is remembered so a client that keeps
|
||||
sending it receives 404 and has to ``initialize`` again, which re-runs
|
||||
admission. Only sessions held by this worker process are affected.
|
||||
"""
|
||||
now: Final = time.monotonic()
|
||||
_forget_expired_admin_terminated_session_ids(now)
|
||||
server_instances: Final = _stateful_server_instances()
|
||||
targets: Final = tuple(
|
||||
(session_id, auth_user)
|
||||
for session_id, auth_user in tuple(_stateful_session_auth_contexts.items())
|
||||
if session_id in server_instances
|
||||
and _session_matches_admin_selector(session_id, auth_user, session_id_prefix, user_id)
|
||||
)
|
||||
terminated: Final = tuple(_gateway_session_for(session_id, auth_user, now) for session_id, auth_user in targets)
|
||||
for session_id, _ in targets:
|
||||
_admin_terminated_session_ids[session_id] = now
|
||||
transport = server_instances.pop(session_id, None)
|
||||
_remove_stateful_session_tracking(session_id)
|
||||
if transport is not None:
|
||||
await transport.terminate()
|
||||
verbose_logger.warning("MCP session '%s' terminated by an administrator.", session_id)
|
||||
return MCPGatewaySessionsTerminateResponse(
|
||||
worker_pid=os.getpid(),
|
||||
terminated_sessions=len(terminated),
|
||||
sessions=terminated,
|
||||
)
|
||||
|
||||
async def _read_request_body_for_routing(
|
||||
receive: Receive,
|
||||
) -> tuple[list[Message], bytes]:
|
||||
|
|
@ -4009,6 +4063,17 @@ if MCP_AVAILABLE:
|
|||
await success_response(scope, receive, send)
|
||||
return True
|
||||
|
||||
if _is_admin_terminated_session_id(_session_id, time.monotonic()):
|
||||
terminated_response: Final = JSONResponse(
|
||||
status_code=404,
|
||||
content={ # mutable-ok: JSONResponse content must be a plain dict
|
||||
"error": "Not Found",
|
||||
"details": "mcp-session-id was terminated by an administrator. Send initialize to start a new session.",
|
||||
},
|
||||
)
|
||||
await terminated_response(scope, receive, send)
|
||||
return True
|
||||
|
||||
# Non-DELETE: strip stale session ID to allow new session creation
|
||||
verbose_logger.warning(
|
||||
"MCP session ID '%s' not found in this worker's memory. "
|
||||
|
|
|
|||
|
|
@ -196,10 +196,12 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
|
|||
"/assemblyai/",
|
||||
"/azure/",
|
||||
"/azure_ai/",
|
||||
"/azure_speech/",
|
||||
"/bedrock/",
|
||||
"/cohere/",
|
||||
"/comprehendmedical",
|
||||
"/cursor/",
|
||||
"/deepgram/",
|
||||
"/eu.assemblyai/",
|
||||
"/gemini/",
|
||||
"/gigachat/",
|
||||
|
|
|
|||
|
|
@ -3050,6 +3050,18 @@
|
|||
},
|
||||
"DailySpendMetadata": {
|
||||
"properties": {
|
||||
"api_key_limit": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "When set, api_keys and every api_key_breakdown list at most this many keys, ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key.",
|
||||
"title": "Api Key Limit"
|
||||
},
|
||||
"has_more": {
|
||||
"default": false,
|
||||
"title": "Has More",
|
||||
|
|
@ -3060,6 +3072,18 @@
|
|||
"title": "Page",
|
||||
"type": "integer"
|
||||
},
|
||||
"total_api_keys": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key lists are truncated to the highest-spend keys.",
|
||||
"title": "Total Api Keys"
|
||||
},
|
||||
"total_api_requests": {
|
||||
"default": 0,
|
||||
"title": "Total Api Requests",
|
||||
|
|
@ -10030,7 +10054,7 @@
|
|||
},
|
||||
"unreachable_fallback": {
|
||||
"default": "fail_closed",
|
||||
"description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.",
|
||||
"description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.",
|
||||
"enum": [
|
||||
"fail_closed",
|
||||
"fail_open"
|
||||
|
|
@ -17157,6 +17181,228 @@
|
|||
]
|
||||
}
|
||||
},
|
||||
"/azure_speech/{endpoint}": {
|
||||
"delete": {
|
||||
"description": "Pass-through for the Azure AI Speech REST APIs (speech to text), e.g.\n`POST /azure_speech/speech/recognition/conversation/cognitiveservices/v1?language=en-US`\nwith the raw audio as the body, or `POST /azure_speech/speechtotext/v3.2/transcriptions`.\n\nThe body is forwarded byte for byte and the proxy injects its own\n`Ocp-Apim-Subscription-Key`; the caller's `Authorization` header is the LiteLLM key\nand is never forwarded.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/azure_speech)",
|
||||
"operationId": "azure_speech_proxy_route_azure_speech__endpoint__delete",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "endpoint",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Endpoint",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Azure Speech Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
]
|
||||
},
|
||||
"get": {
|
||||
"description": "Pass-through for the Azure AI Speech REST APIs (speech to text), e.g.\n`POST /azure_speech/speech/recognition/conversation/cognitiveservices/v1?language=en-US`\nwith the raw audio as the body, or `POST /azure_speech/speechtotext/v3.2/transcriptions`.\n\nThe body is forwarded byte for byte and the proxy injects its own\n`Ocp-Apim-Subscription-Key`; the caller's `Authorization` header is the LiteLLM key\nand is never forwarded.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/azure_speech)",
|
||||
"operationId": "azure_speech_proxy_route_azure_speech__endpoint__get",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "endpoint",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Endpoint",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Azure Speech Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
]
|
||||
},
|
||||
"patch": {
|
||||
"description": "Pass-through for the Azure AI Speech REST APIs (speech to text), e.g.\n`POST /azure_speech/speech/recognition/conversation/cognitiveservices/v1?language=en-US`\nwith the raw audio as the body, or `POST /azure_speech/speechtotext/v3.2/transcriptions`.\n\nThe body is forwarded byte for byte and the proxy injects its own\n`Ocp-Apim-Subscription-Key`; the caller's `Authorization` header is the LiteLLM key\nand is never forwarded.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/azure_speech)",
|
||||
"operationId": "azure_speech_proxy_route_azure_speech__endpoint__patch",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "endpoint",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Endpoint",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Azure Speech Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
]
|
||||
},
|
||||
"post": {
|
||||
"description": "Pass-through for the Azure AI Speech REST APIs (speech to text), e.g.\n`POST /azure_speech/speech/recognition/conversation/cognitiveservices/v1?language=en-US`\nwith the raw audio as the body, or `POST /azure_speech/speechtotext/v3.2/transcriptions`.\n\nThe body is forwarded byte for byte and the proxy injects its own\n`Ocp-Apim-Subscription-Key`; the caller's `Authorization` header is the LiteLLM key\nand is never forwarded.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/azure_speech)",
|
||||
"operationId": "azure_speech_proxy_route_azure_speech__endpoint__post",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "endpoint",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Endpoint",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Azure Speech Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
]
|
||||
},
|
||||
"put": {
|
||||
"description": "Pass-through for the Azure AI Speech REST APIs (speech to text), e.g.\n`POST /azure_speech/speech/recognition/conversation/cognitiveservices/v1?language=en-US`\nwith the raw audio as the body, or `POST /azure_speech/speechtotext/v3.2/transcriptions`.\n\nThe body is forwarded byte for byte and the proxy injects its own\n`Ocp-Apim-Subscription-Key`; the caller's `Authorization` header is the LiteLLM key\nand is never forwarded.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/azure_speech)",
|
||||
"operationId": "azure_speech_proxy_route_azure_speech__endpoint__put",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "endpoint",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Endpoint",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Azure Speech Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
]
|
||||
}
|
||||
},
|
||||
"/bedrock/{endpoint}": {
|
||||
"delete": {
|
||||
"description": "This is the v1 passthrough for Bedrock.\nV2 is handled by the `/bedrock/v2` endpoint.\n[Docs](https://docs.litellm.ai/docs/pass_through/bedrock)",
|
||||
|
|
@ -27965,6 +28211,32 @@
|
|||
"title": "MCPGatewaySessionsResponse",
|
||||
"type": "object"
|
||||
},
|
||||
"MCPGatewaySessionsTerminateResponse": {
|
||||
"description": "Stateful sessions an administrator force-closed on this proxy worker.",
|
||||
"properties": {
|
||||
"sessions": {
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/MCPGatewaySession"
|
||||
},
|
||||
"title": "Sessions",
|
||||
"type": "array"
|
||||
},
|
||||
"terminated_sessions": {
|
||||
"title": "Terminated Sessions",
|
||||
"type": "integer"
|
||||
},
|
||||
"worker_pid": {
|
||||
"title": "Worker Pid",
|
||||
"type": "integer"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"worker_pid",
|
||||
"terminated_sessions"
|
||||
],
|
||||
"title": "MCPGatewaySessionsTerminateResponse",
|
||||
"type": "object"
|
||||
},
|
||||
"MCPOAuthUserCredentialRequest": {
|
||||
"description": "Stores a user's OAuth2 token for an OpenAPI MCP server.",
|
||||
"properties": {
|
||||
|
|
@ -28061,6 +28333,56 @@
|
|||
"title": "MCPOAuthUserCredentialStatus",
|
||||
"type": "object"
|
||||
},
|
||||
"MCPServerUserCredentialListItem": {
|
||||
"description": "One user's stored credential for an MCP server, as an admin sees it. Never carries the secret.",
|
||||
"properties": {
|
||||
"connected_at": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Connected At"
|
||||
},
|
||||
"credential_type": {
|
||||
"enum": [
|
||||
"oauth2",
|
||||
"byok"
|
||||
],
|
||||
"title": "Credential Type",
|
||||
"type": "string"
|
||||
},
|
||||
"expires_at": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Expires At"
|
||||
},
|
||||
"updated_at": {
|
||||
"title": "Updated At",
|
||||
"type": "string"
|
||||
},
|
||||
"user_id": {
|
||||
"title": "User Id",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"user_id",
|
||||
"credential_type",
|
||||
"updated_at"
|
||||
],
|
||||
"title": "MCPServerUserCredentialListItem",
|
||||
"type": "object"
|
||||
},
|
||||
"MCPSubmissionsSummary": {
|
||||
"properties": {
|
||||
"active": {
|
||||
|
|
@ -30237,7 +30559,7 @@
|
|||
},
|
||||
"/v1/mcp/server/{server_id}/oauth-user-credential": {
|
||||
"delete": {
|
||||
"description": "Revoke the calling user's stored OAuth2 token for an MCP server",
|
||||
"description": "Revoke the calling user's stored OAuth2 token for an MCP server. A proxy admin may pass user_id to revoke another user's stored token.",
|
||||
"operationId": "delete_mcp_oauth_user_credential_v1_mcp_server__server_id__oauth_user_credential_delete",
|
||||
"parameters": [
|
||||
{
|
||||
|
|
@ -30248,6 +30570,23 @@
|
|||
"title": "Server Id",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
{
|
||||
"in": "query",
|
||||
"name": "user_id",
|
||||
"required": false,
|
||||
"schema": {
|
||||
"anyOf": [
|
||||
{
|
||||
"minLength": 1,
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "User Id"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
|
|
@ -30447,7 +30786,7 @@
|
|||
},
|
||||
"/v1/mcp/server/{server_id}/user-credential": {
|
||||
"delete": {
|
||||
"description": "Delete the calling user's stored API key for a BYOK MCP server",
|
||||
"description": "Delete the calling user's stored API key for a BYOK MCP server. A proxy admin may pass user_id to revoke another user's stored key.",
|
||||
"operationId": "delete_mcp_user_credential_v1_mcp_server__server_id__user_credential_delete",
|
||||
"parameters": [
|
||||
{
|
||||
|
|
@ -30458,6 +30797,23 @@
|
|||
"title": "Server Id",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
{
|
||||
"in": "query",
|
||||
"name": "user_id",
|
||||
"required": false,
|
||||
"schema": {
|
||||
"anyOf": [
|
||||
{
|
||||
"minLength": 1,
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "User Id"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
|
|
@ -30549,6 +30905,58 @@
|
|||
]
|
||||
}
|
||||
},
|
||||
"/v1/mcp/server/{server_id}/user-credentials": {
|
||||
"get": {
|
||||
"description": "List every user's stored BYOK or OAuth2 credential for an MCP server (admin only, no secrets)",
|
||||
"operationId": "list_mcp_server_user_credentials_v1_mcp_server__server_id__user_credentials_get",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "server_id",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Server Id",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/MCPServerUserCredentialListItem"
|
||||
},
|
||||
"title": "Response List Mcp Server User Credentials V1 Mcp Server Server Id User Credentials Get",
|
||||
"type": "array"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "List Mcp Server User Credentials",
|
||||
"tags": [
|
||||
"mcp_management"
|
||||
]
|
||||
}
|
||||
},
|
||||
"/v1/mcp/server/{server_id}/user-env-vars": {
|
||||
"delete": {
|
||||
"description": "Clear the calling user's per-user MCP env var values for this server.",
|
||||
|
|
@ -30700,6 +31108,77 @@
|
|||
}
|
||||
},
|
||||
"/v1/mcp/sessions": {
|
||||
"delete": {
|
||||
"description": "Force-close live stateful MCP gateway sessions on this proxy worker, selected by session id prefix and/or by the LiteLLM user that opened them (proxy admin only).",
|
||||
"operationId": "delete_mcp_gateway_sessions_v1_mcp_sessions_delete",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "query",
|
||||
"name": "session_id_prefix",
|
||||
"required": false,
|
||||
"schema": {
|
||||
"anyOf": [
|
||||
{
|
||||
"minLength": 8,
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Session Id Prefix"
|
||||
}
|
||||
},
|
||||
{
|
||||
"in": "query",
|
||||
"name": "user_id",
|
||||
"required": false,
|
||||
"schema": {
|
||||
"anyOf": [
|
||||
{
|
||||
"minLength": 1,
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "User Id"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/MCPGatewaySessionsTerminateResponse"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Delete Mcp Gateway Sessions",
|
||||
"tags": [
|
||||
"mcp_management"
|
||||
]
|
||||
},
|
||||
"get": {
|
||||
"description": "Live stateful MCP gateway sessions on this proxy worker, grouped by AI client and by user.",
|
||||
"operationId": "get_mcp_gateway_sessions_v1_mcp_sessions_get",
|
||||
|
|
@ -38542,6 +39021,7 @@
|
|||
"type": "object"
|
||||
},
|
||||
"SCIMMultiValuedAttribute": {
|
||||
"additionalProperties": true,
|
||||
"properties": {
|
||||
"display": {
|
||||
"anyOf": [
|
||||
|
|
@ -38577,13 +39057,17 @@
|
|||
"title": "Type"
|
||||
},
|
||||
"value": {
|
||||
"title": "Value",
|
||||
"type": "string"
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Value"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"value"
|
||||
],
|
||||
"title": "SCIMMultiValuedAttribute",
|
||||
"type": "object"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -71,6 +71,7 @@ from litellm.types.utils import (
|
|||
StandardLoggingVectorStoreRequest,
|
||||
StandardPassThroughResponseObject,
|
||||
TextCompletionResponse,
|
||||
TranscriptionResponse,
|
||||
)
|
||||
from litellm.types.videos.main import VideoObject
|
||||
|
||||
|
|
@ -469,6 +470,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
mapped_pass_through_routes = [
|
||||
"/bedrock",
|
||||
"/comprehendmedical",
|
||||
"/azure_speech",
|
||||
"/transcribe",
|
||||
"/vertex-ai",
|
||||
"/vertex_ai",
|
||||
|
|
@ -490,6 +492,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/gigachat",
|
||||
"/watsonx",
|
||||
"/nvidia_nim",
|
||||
"/deepgram",
|
||||
]
|
||||
|
||||
#########################################################
|
||||
|
|
@ -662,6 +665,11 @@ class LiteLLMRoutes(enum.Enum):
|
|||
KeyManagementRoutes.AUTO_ROUTER_MANAGE.value,
|
||||
]
|
||||
|
||||
team_service_account_key_routes = (
|
||||
KeyManagementRoutes.KEY_GENERATE.value,
|
||||
KeyManagementRoutes.KEY_UPDATE.value,
|
||||
)
|
||||
|
||||
management_routes = (
|
||||
[
|
||||
# user
|
||||
|
|
@ -1721,6 +1729,16 @@ class MCPUserCredentialListItem(LiteLLMPydanticObjectBase):
|
|||
connected_at: str | None = None # ISO-8601
|
||||
|
||||
|
||||
class MCPServerUserCredentialListItem(LiteLLMPydanticObjectBase):
|
||||
"""One user's stored credential for an MCP server, as an admin sees it. Never carries the secret."""
|
||||
|
||||
user_id: str
|
||||
credential_type: Literal["oauth2", "byok"]
|
||||
expires_at: str | None = None
|
||||
connected_at: str | None = None
|
||||
updated_at: str
|
||||
|
||||
|
||||
class MCPUserEnvVarsRequest(LiteLLMPydanticObjectBase):
|
||||
"""Payload for storing the calling user's per-user env var values."""
|
||||
|
||||
|
|
@ -2763,6 +2781,25 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
description="sends alerts if requests hang for 5min+",
|
||||
)
|
||||
ui_access_mode: Literal["admin_only", "all"] | None = Field("all", description="Control access to the Proxy UI")
|
||||
max_failed_login_attempts_per_source: int | None = Field(
|
||||
None,
|
||||
ge=1,
|
||||
description="Failed Admin UI sign-in attempts allowed from one source address, across every username, within `failed_login_window_seconds`. One more blocks that address for `failed_login_block_seconds`. Half this value, rounded down but at least 1, is the allowance for one username from that address; one more blocks that address for that username only, and its further failures stop counting toward the address limit, so a script stuck on one account does not block everyone behind a shared address. The per-address limit is only enforced when `trusted_proxy_ranges` is set: to the proxies in front of LiteLLM, or to an empty list when clients connect directly. Left unset, the peer address may be a shared ingress and only the per-username half runs. IPv6 addresses are grouped by /64. Set under `general_settings` in config.yaml. Defaults to 10",
|
||||
)
|
||||
max_failed_login_attempts_per_source_overrides: dict[str, int] | None = Field(
|
||||
None,
|
||||
description="Per-address overrides of `max_failed_login_attempts_per_source`, keyed by IP address or CIDR range, e.g. {'1.2.3.4': 200, '5.6.0.0/24': 500}. The most specific matching range wins (between equivalent keys such as '1.2.3.4' and '1.2.3.4/32', an exemption wins, then the higher limit), and the per-username allowance for that address follows as half the override. A value of 0 exempts the address from both limits. Set under `general_settings` in config.yaml",
|
||||
)
|
||||
failed_login_window_seconds: int | None = Field(
|
||||
None,
|
||||
ge=1,
|
||||
description="Fixed window in seconds over which failed Admin UI sign-in attempts are counted. The window starts at the first failure and is not extended by later ones. Set under `general_settings` in config.yaml. Defaults to 60",
|
||||
)
|
||||
failed_login_block_seconds: int | None = Field(
|
||||
None,
|
||||
ge=1,
|
||||
description="How long a blocked source address, or source address and username, stays blocked. Every attempt from a blocked key, right or wrong, is refused with 429 before the password is checked; the block is not extended by refused attempts. Set under `general_settings` in config.yaml. Defaults to 300",
|
||||
)
|
||||
allowed_routes: list | None = Field(None, description="Proxy API Endpoints you want users to be able to access")
|
||||
reject_clientside_metadata_tags: bool | None = Field(
|
||||
None,
|
||||
|
|
@ -2876,7 +2913,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
)
|
||||
trusted_proxy_ranges: list[str] | None = Field(
|
||||
None,
|
||||
description="CIDR ranges of trusted reverse proxies allowed to provide identity headers for header-based auth paths such as enable_oauth2_proxy_auth and custom_ui_sso_sign_in_handler.",
|
||||
description="CIDR ranges of trusted reverse proxies allowed to provide identity headers for header-based auth paths such as enable_oauth2_proxy_auth and custom_ui_sso_sign_in_handler, and whose X-Forwarded-For is used to attribute Admin UI sign-in attempts to a source address. Set it to an empty list when clients connect directly, so the peer address is the source. Left unset, or containing an entry that is not an address or CIDR range, the per-source sign-in limit is off.",
|
||||
)
|
||||
store_model_in_db: bool | None = Field(
|
||||
None,
|
||||
|
|
@ -3279,6 +3316,15 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
|
|||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
@property
|
||||
def is_team_service_account(self) -> bool:
|
||||
return (
|
||||
self.user_id is None
|
||||
and self.team_id is not None
|
||||
and bool(self.metadata)
|
||||
and self.metadata.get("service_account_id") is not None
|
||||
)
|
||||
|
||||
|
||||
def user_api_key_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
"""Return True if the caller's role grants unscoped read access to all
|
||||
|
|
@ -4728,6 +4774,7 @@ PassThroughEndpointLoggingResultValues = (
|
|||
| VideoObject
|
||||
| StandardPassThroughResponseObject
|
||||
| ResponsesAPIResponse
|
||||
| TranscriptionResponse
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
445
litellm/proxy/auth/login_throttle.py
Normal file
445
litellm/proxy/auth/login_throttle.py
Normal file
|
|
@ -0,0 +1,445 @@
|
|||
"""Failed-login accounting for the Admin UI sign-in path.
|
||||
|
||||
Wrong passwords are counted over a short window per source address and per source-and-username
|
||||
pair; too many in one window blocks that key for a fixed time. While a key is blocked every attempt
|
||||
from it, right or wrong, is refused with 429 before the password is checked. A blocked pair stops
|
||||
counting against its source, so one script stuck on one account does not block the whole office.
|
||||
Recovery is the master key over the API, which never passes through here, or waiting out the block.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import ipaddress
|
||||
import math
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from functools import cache
|
||||
from typing import Final, Literal, NamedTuple, Protocol, TypeAlias
|
||||
|
||||
from fastapi import Request, status
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError
|
||||
from litellm.constants import (
|
||||
EMPTY_MAPPING,
|
||||
LOGIN_THROTTLE_CACHE_KEY_PREFIX,
|
||||
LOGIN_THROTTLE_MAX_TRACKED_BLOCKS,
|
||||
LOGIN_THROTTLE_MAX_TRACKED_COUNTERS,
|
||||
LOGIN_THROTTLE_NOT_BLOCKED,
|
||||
LOGIN_THROTTLE_UNKNOWN_SOURCE,
|
||||
)
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||
from litellm.proxy.auth.network import TrustedProxyConfig, resolve_client_ip
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
|
||||
DEFAULT_MAX_FAILED_LOGIN_ATTEMPTS_PER_SOURCE: Final = 10
|
||||
DEFAULT_FAILED_LOGIN_WINDOW_SECONDS: Final = 60
|
||||
DEFAULT_FAILED_LOGIN_BLOCK_SECONDS: Final = 300
|
||||
|
||||
IPV6_SOURCE_PREFIX_LENGTH: Final = 64
|
||||
EXEMPT: Final = 0
|
||||
|
||||
SOURCE_LIMIT_KEY: Final = "max_failed_login_attempts_per_source"
|
||||
SOURCE_LIMIT_OVERRIDES_KEY: Final = "max_failed_login_attempts_per_source_overrides"
|
||||
WINDOW_KEY: Final = "failed_login_window_seconds"
|
||||
BLOCK_KEY: Final = "failed_login_block_seconds"
|
||||
TRUSTED_PROXY_RANGES_KEY: Final = "trusted_proxy_ranges"
|
||||
|
||||
_REDIS_FAILURES: Final = (RedisError, RedisCircuitBreakerOpenError, OSError, asyncio.TimeoutError)
|
||||
_LOCAL_BLOCK_EXPIRY: Final = TypeAdapter[float | None](float | None)
|
||||
_SOURCE_LIMIT_OVERRIDES: Final = TypeAdapter[Mapping[str, object]](Mapping[str, object])
|
||||
_RANGE_ENTRIES: Final = TypeAdapter[tuple[object, ...]](tuple[object, ...])
|
||||
|
||||
Scope: TypeAlias = Literal["user", "source"]
|
||||
|
||||
_BlockTtls: TypeAlias = tuple[int, int]
|
||||
_LUA_BLOCK_TTLS: Final = TypeAdapter[_BlockTtls](_BlockTtls)
|
||||
_Network: TypeAlias = ipaddress.IPv4Network | ipaddress.IPv6Network
|
||||
|
||||
|
||||
class LocalStore(Protocol):
|
||||
"""The per-worker store behind the counters and blocks; ``InMemoryCache`` satisfies it."""
|
||||
|
||||
def get_cache(self, key: str) -> object: ...
|
||||
|
||||
def set_cache(self, key: str, value: float, *, ttl: int) -> None: ...
|
||||
|
||||
def increment_cache(self, key: str, value: float, *, ttl: int) -> float: ...
|
||||
|
||||
def delete_cache(self, key: str) -> None: ...
|
||||
|
||||
|
||||
# KEYS: pair counter, pair block, source counter, source block (one cluster slot via the source hash tag)
|
||||
# ARGV: pair limit, source limit (0 = source scope off), window seconds, block seconds
|
||||
# Both scripts return {pair block TTL, source block TTL}; 0 or below means not blocked
|
||||
_BLOCK_TTLS_LUA: Final = "return {redis.call('TTL', KEYS[2]), redis.call('TTL', KEYS[4])}"
|
||||
_RECORD_FAILURE_LUA: Final = (
|
||||
"local function bump(count_key, block_key, limit) "
|
||||
"local blocked = redis.call('TTL', block_key) "
|
||||
"if blocked > 0 then return blocked end "
|
||||
"local count = redis.call('INCR', count_key) "
|
||||
"if redis.call('TTL', count_key) < 0 then redis.call('EXPIRE', count_key, ARGV[3]) end "
|
||||
"if count > limit then redis.call('SET', block_key, '1', 'EX', ARGV[4]) return tonumber(ARGV[4]) end "
|
||||
"return 0 end "
|
||||
"local user_block = bump(KEYS[1], KEYS[2], tonumber(ARGV[1])) "
|
||||
"local source_block = 0 "
|
||||
"if tonumber(ARGV[2]) > 0 and user_block == 0 then "
|
||||
"source_block = bump(KEYS[3], KEYS[4], tonumber(ARGV[2])) end "
|
||||
"return {user_block, source_block}"
|
||||
)
|
||||
|
||||
_COUNTERS: Final = InMemoryCache(
|
||||
max_size_in_memory=LOGIN_THROTTLE_MAX_TRACKED_COUNTERS, default_ttl=DEFAULT_FAILED_LOGIN_WINDOW_SECONDS
|
||||
)
|
||||
_BLOCKS: Final = InMemoryCache(
|
||||
max_size_in_memory=LOGIN_THROTTLE_MAX_TRACKED_BLOCKS, default_ttl=DEFAULT_FAILED_LOGIN_BLOCK_SECONDS
|
||||
)
|
||||
|
||||
|
||||
@cache
|
||||
def _rate_limit_disabled() -> bool:
|
||||
return get_secret_bool("LITELLM_DISABLE_LOGIN_RATE_LIMIT", default_value=False) is True
|
||||
|
||||
|
||||
@cache
|
||||
def warn_login_counters_are_per_worker(num_workers: str) -> None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Running %s workers but Redis is not configured. Failed Admin UI sign-in attempts are counted "
|
||||
"per worker, so the effective limits are %s times the configured values. Configure Redis "
|
||||
"to share one count across workers.",
|
||||
num_workers,
|
||||
num_workers,
|
||||
)
|
||||
|
||||
|
||||
@cache
|
||||
def warn_source_login_limit_is_off() -> None:
|
||||
verbose_proxy_logger.warning(
|
||||
"%s is not set or not a valid list of ranges, so failed Admin UI sign-in attempts are limited per "
|
||||
"source address and username only. Set it to the address ranges of the proxies in front of LiteLLM, "
|
||||
"or to an empty list when clients connect directly, to also limit each source address across usernames.",
|
||||
TRUSTED_PROXY_RANGES_KEY,
|
||||
)
|
||||
|
||||
|
||||
def declared_proxy_ranges(settings: Mapping[str, object]) -> tuple[str, ...] | None:
|
||||
"""What the operator says fronts LiteLLM: the proxy ranges, an empty tuple for none, None when unsaid.
|
||||
|
||||
Only a declared topology makes the source address trustworthy enough to limit across usernames.
|
||||
An unset key, a value that is not a list of ranges, or a list with an entry that is not an address
|
||||
or range leaves it unknown and the source scope off.
|
||||
"""
|
||||
entries: Final = _configured_range_entries(settings.get(TRUSTED_PROXY_RANGES_KEY))
|
||||
if entries is None or any(_parse_network(entry, TRUSTED_PROXY_RANGES_KEY) is None for entry in entries):
|
||||
return None
|
||||
return entries
|
||||
|
||||
|
||||
def _configured_range_entries(raw_ranges: object) -> tuple[str, ...] | None:
|
||||
"""Every configured entry, blanks included, so a stray empty string fails validation like any other typo."""
|
||||
if raw_ranges is None:
|
||||
return None
|
||||
if isinstance(raw_ranges, str):
|
||||
return tuple(part.strip() for part in raw_ranges.split(","))
|
||||
try:
|
||||
return tuple(str(entry).strip() for entry in _RANGE_ENTRIES.validate_python(raw_ranges))
|
||||
except ValidationError:
|
||||
verbose_proxy_logger.warning(
|
||||
"Invalid %s value: expected a list of address ranges, got %s",
|
||||
TRUSTED_PROXY_RANGES_KEY,
|
||||
type(raw_ranges).__name__,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _positive_int(raw: object, key: str, default: int) -> int:
|
||||
if raw is None:
|
||||
return default
|
||||
try:
|
||||
value: Final = int(str(raw))
|
||||
except (TypeError, ValueError):
|
||||
verbose_proxy_logger.warning("Invalid %s value %r; using %s", key, raw, default)
|
||||
return default
|
||||
if value < 1:
|
||||
verbose_proxy_logger.warning("Invalid %s value %s (must be >= 1); using %s", key, value, default)
|
||||
return default
|
||||
return value
|
||||
|
||||
|
||||
def _int_setting(settings: Mapping[str, object], key: str, default: int) -> int:
|
||||
return _positive_int(settings.get(key), key, default)
|
||||
|
||||
|
||||
def _override_limit(raw: object, default: int) -> int:
|
||||
"""A per-address override: a limit of 1 or more, or ``EXEMPT`` (0) to leave that address unlimited."""
|
||||
if str(raw).strip() == str(EXEMPT):
|
||||
return EXEMPT
|
||||
return _positive_int(raw, SOURCE_LIMIT_OVERRIDES_KEY, default)
|
||||
|
||||
|
||||
def _parse_address(client_ip: str) -> ipaddress.IPv4Address | ipaddress.IPv6Address | None:
|
||||
"""The address as it is limited and counted: an IPv4-mapped IPv6 address is its IPv4 address."""
|
||||
try:
|
||||
address: Final = ipaddress.ip_address(client_ip)
|
||||
except ValueError:
|
||||
return None
|
||||
if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped is not None:
|
||||
return address.ipv4_mapped
|
||||
return address
|
||||
|
||||
|
||||
def _parse_network(raw_range: str, setting_name: str = SOURCE_LIMIT_OVERRIDES_KEY) -> _Network | None:
|
||||
try:
|
||||
return ipaddress.ip_network(raw_range.strip(), strict=False)
|
||||
except ValueError:
|
||||
verbose_proxy_logger.warning("Invalid address or range %r in %s; skipping", raw_range, setting_name)
|
||||
return None
|
||||
|
||||
|
||||
def _precedence(network: _Network, limit: int) -> tuple[int, bool, int]:
|
||||
"""Sort key for competing overrides: the longest prefix wins, then an exemption, then the higher limit."""
|
||||
return (network.prefixlen, limit == EXEMPT, limit)
|
||||
|
||||
|
||||
def _source_limit(settings: Mapping[str, object], client_ip: str) -> int:
|
||||
"""Failure allowance for this address: the most specific configured range containing it, else the default.
|
||||
|
||||
``EXEMPT`` (0) means the operator opted this address out of both limits. Between equivalent keys such as
|
||||
``1.2.3.4`` and ``1.2.3.4/32`` an exemption wins, then the higher limit.
|
||||
"""
|
||||
default: Final = _int_setting(settings, SOURCE_LIMIT_KEY, DEFAULT_MAX_FAILED_LOGIN_ATTEMPTS_PER_SOURCE)
|
||||
raw_overrides: Final = settings.get(SOURCE_LIMIT_OVERRIDES_KEY)
|
||||
if raw_overrides is None:
|
||||
return default
|
||||
try:
|
||||
overrides: Final = _SOURCE_LIMIT_OVERRIDES.validate_python(raw_overrides)
|
||||
except ValidationError:
|
||||
verbose_proxy_logger.warning(
|
||||
"Invalid %s value; expected a mapping of address or range to limit", SOURCE_LIMIT_OVERRIDES_KEY
|
||||
)
|
||||
return default
|
||||
address: Final = _parse_address(client_ip)
|
||||
if address is None:
|
||||
return default
|
||||
matches: Final = sorted(
|
||||
_precedence(network, _override_limit(raw_limit, default))
|
||||
for raw_range, raw_limit in overrides.items()
|
||||
if (network := _parse_network(raw_range)) is not None and address in network
|
||||
)
|
||||
return matches[-1][-1] if matches else default
|
||||
|
||||
|
||||
def user_limit_for(source_limit: int) -> int:
|
||||
"""Failures allowed for one username from one address: half the address allowance, rounded down, at least 1."""
|
||||
return max(source_limit // 2, 1)
|
||||
|
||||
|
||||
def source_group(client_ip: str) -> str:
|
||||
"""The bucket an address is counted in: IPv4 as is, IPv6 by its /64, so one prefix holder cannot rotate."""
|
||||
address: Final = _parse_address(client_ip)
|
||||
if address is None:
|
||||
return client_ip
|
||||
if isinstance(address, ipaddress.IPv6Address):
|
||||
return str(ipaddress.ip_network((address, IPV6_SOURCE_PREFIX_LENGTH), strict=False))
|
||||
return str(address)
|
||||
|
||||
|
||||
class _Keys(NamedTuple):
|
||||
pair_counter: str
|
||||
pair_block: str
|
||||
source_counter: str
|
||||
source_block: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Block:
|
||||
scope: Scope
|
||||
retry_after: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class LoginThrottle:
|
||||
"""Failed-login limits for one request's source address.
|
||||
|
||||
``source_limit`` is None when the source scope is off: ``trusted_proxy_ranges`` is unset, so the peer
|
||||
address may be a shared ingress. An empty list means clients connect directly and the peer is the source.
|
||||
``user_limit`` is derived from the address allowance either way, see ``user_limit_for``. An address whose
|
||||
override is ``EXEMPT`` gets a disabled throttle: nothing is counted or blocked for it.
|
||||
"""
|
||||
|
||||
client_ip: str
|
||||
source_limit: int | None
|
||||
user_limit: int
|
||||
window_seconds: int
|
||||
block_seconds: int
|
||||
counters: LocalStore
|
||||
blocks: LocalStore
|
||||
redis_cache: RedisCache | None = None
|
||||
enabled: bool = True
|
||||
|
||||
@classmethod
|
||||
def from_request(
|
||||
cls,
|
||||
request: Request,
|
||||
general_settings: Mapping[str, object] | None,
|
||||
redis_cache: RedisCache | None,
|
||||
) -> LoginThrottle:
|
||||
settings: Final[Mapping[str, object]] = general_settings if general_settings is not None else EMPTY_MAPPING
|
||||
proxies: Final = declared_proxy_ranges(settings)
|
||||
resolved, _ = resolve_client_ip(
|
||||
request, TrustedProxyConfig(use_forwarded_for=bool(proxies), trusted_proxy_cidrs=proxies or ())
|
||||
)
|
||||
source_limit: Final = _source_limit(settings, resolved or LOGIN_THROTTLE_UNKNOWN_SOURCE)
|
||||
exempt: Final = source_limit == EXEMPT
|
||||
return cls(
|
||||
client_ip=resolved or LOGIN_THROTTLE_UNKNOWN_SOURCE,
|
||||
source_limit=source_limit if proxies is not None and resolved is not None and not exempt else None,
|
||||
user_limit=user_limit_for(source_limit),
|
||||
window_seconds=_int_setting(settings, WINDOW_KEY, DEFAULT_FAILED_LOGIN_WINDOW_SECONDS),
|
||||
block_seconds=_int_setting(settings, BLOCK_KEY, DEFAULT_FAILED_LOGIN_BLOCK_SECONDS),
|
||||
counters=_COUNTERS,
|
||||
blocks=_BLOCKS,
|
||||
redis_cache=redis_cache,
|
||||
enabled=not exempt and not _rate_limit_disabled(),
|
||||
)
|
||||
|
||||
def _keys(self, username: str) -> _Keys:
|
||||
group: Final = source_group(self.client_ip)
|
||||
user: Final = hashlib.sha256(username.casefold().encode("utf-8")).hexdigest()
|
||||
return _Keys(
|
||||
pair_counter=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:user:{user}",
|
||||
pair_block=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:block:user:{user}",
|
||||
source_counter=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:source",
|
||||
source_block=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:block:source",
|
||||
)
|
||||
|
||||
async def attempt(self, username: str) -> LoginAttempt:
|
||||
"""Refuses a blocked key before any credential is looked at; otherwise hands back the attempt to settle."""
|
||||
if not self.enabled:
|
||||
return LoginAttempt(throttle=self, username=username)
|
||||
block: Final = await self._active_block(self._keys(username))
|
||||
if block is None:
|
||||
return LoginAttempt(throttle=self, username=username)
|
||||
verbose_proxy_logger.warning(
|
||||
"Admin UI sign-in refused: the %s is blocked for %s more seconds; username=%r source=%s",
|
||||
block.scope,
|
||||
block.retry_after,
|
||||
username,
|
||||
self.client_ip,
|
||||
)
|
||||
raise self.refused(block.retry_after)
|
||||
|
||||
async def _active_block(self, keys: _Keys) -> Block | None:
|
||||
local: Final = self._local_block_ttls(keys)
|
||||
shared: Final = await self._shared_block_ttls(keys)
|
||||
user_ttl: Final = max(local[0], shared[0])
|
||||
source_ttl: Final = max(local[1], shared[1])
|
||||
if self.source_limit is not None and source_ttl > 0:
|
||||
return Block(scope="source", retry_after=source_ttl)
|
||||
if user_ttl > 0:
|
||||
return Block(scope="user", retry_after=user_ttl)
|
||||
return None
|
||||
|
||||
async def _shared_block_ttls(self, keys: _Keys) -> _BlockTtls:
|
||||
if self.redis_cache is None:
|
||||
return LOGIN_THROTTLE_NOT_BLOCKED
|
||||
try:
|
||||
return _LUA_BLOCK_TTLS.validate_python(
|
||||
await self.redis_cache.async_register_script(_BLOCK_TTLS_LUA)(keys, ())
|
||||
)
|
||||
except _REDIS_FAILURES as err:
|
||||
self._warn_redis(err)
|
||||
return LOGIN_THROTTLE_NOT_BLOCKED
|
||||
|
||||
def _local_block_ttls(self, keys: _Keys) -> _BlockTtls:
|
||||
return self._local_block_ttl(keys.pair_block), self._local_block_ttl(keys.source_block)
|
||||
|
||||
def _local_block_ttl(self, block_key: str) -> int:
|
||||
expires_at: Final = _LOCAL_BLOCK_EXPIRY.validate_python(self.blocks.get_cache(block_key))
|
||||
if expires_at is None:
|
||||
return 0
|
||||
return max(math.ceil(expires_at - time.time()), 0)
|
||||
|
||||
async def record_failure(self, username: str) -> _BlockTtls:
|
||||
keys: Final = self._keys(username)
|
||||
source_limit: Final = self.source_limit or 0
|
||||
if self.redis_cache is not None:
|
||||
try:
|
||||
return _LUA_BLOCK_TTLS.validate_python(
|
||||
await self.redis_cache.async_register_script(_RECORD_FAILURE_LUA)(
|
||||
keys, (self.user_limit, source_limit, self.window_seconds, self.block_seconds)
|
||||
)
|
||||
)
|
||||
except _REDIS_FAILURES as err:
|
||||
self._warn_redis(err)
|
||||
user_block: Final = self._local_bump(keys.pair_counter, keys.pair_block, self.user_limit)
|
||||
if source_limit == 0 or user_block > 0:
|
||||
return user_block, 0
|
||||
return user_block, self._local_bump(keys.source_counter, keys.source_block, source_limit)
|
||||
|
||||
def _local_bump(self, count_key: str, block_key: str, limit: int) -> int:
|
||||
blocked: Final = self._local_block_ttl(block_key)
|
||||
if blocked > 0:
|
||||
return blocked
|
||||
count: Final = int(self.counters.increment_cache(count_key, 1, ttl=self.window_seconds))
|
||||
if count <= limit:
|
||||
return 0
|
||||
self.blocks.set_cache(block_key, time.time() + self.block_seconds, ttl=self.block_seconds)
|
||||
return self.block_seconds
|
||||
|
||||
async def clear_pair(self, username: str) -> None:
|
||||
pair_counter: Final = self._keys(username).pair_counter
|
||||
if self.redis_cache is not None:
|
||||
try:
|
||||
await self.redis_cache.async_delete_cache(pair_counter)
|
||||
except _REDIS_FAILURES as err:
|
||||
self._warn_redis(err)
|
||||
self.counters.delete_cache(pair_counter)
|
||||
|
||||
def _warn_redis(self, err: Exception) -> None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Redis failed while counting Admin UI sign-in attempts; using this worker's own counters "
|
||||
"until it recovers: %s",
|
||||
err,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def refused(retry_after: int) -> ProxyException:
|
||||
return ProxyException(
|
||||
message="Too many failed sign-in attempts. Try again later.",
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="username",
|
||||
code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||||
headers={"Retry-After": str(retry_after)}, # mutable-ok: ProxyException writes into its headers dict
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class LoginAttempt:
|
||||
throttle: LoginThrottle
|
||||
username: str
|
||||
|
||||
async def succeeded(self) -> None:
|
||||
if not self.throttle.enabled:
|
||||
return
|
||||
await self.throttle.clear_pair(self.username)
|
||||
|
||||
async def failed(self) -> None:
|
||||
if not self.throttle.enabled:
|
||||
return
|
||||
user_block, source_block = await self.throttle.record_failure(self.username)
|
||||
if user_block == 0 and source_block == 0:
|
||||
return
|
||||
verbose_proxy_logger.warning(
|
||||
"Admin UI sign-in blocked for %s seconds after too many failures; scope=%s username=%r source=%s",
|
||||
user_block or source_block,
|
||||
"user" if user_block else "source",
|
||||
self.username,
|
||||
self.throttle.client_ip,
|
||||
)
|
||||
|
|
@ -27,6 +27,7 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured
|
||||
from litellm.proxy.auth.login_throttle import LoginAttempt, LoginThrottle
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
generate_key_helper_fn,
|
||||
|
|
@ -44,6 +45,11 @@ from litellm.repositories.user_repository import UserRepository
|
|||
from litellm.secret_managers.main import get_secret_bool
|
||||
from litellm.types.proxy.ui_sso import ReturnedUITokenObject
|
||||
|
||||
INVALID_UI_CREDENTIALS_MESSAGE: Final = (
|
||||
"Invalid credentials used to access UI. Check 'UI_USERNAME' and 'UI_PASSWORD', or the password set for your user"
|
||||
)
|
||||
INVALID_USER_PASSWORD_MESSAGE: Final = "Invalid credentials used to access UI. Check the password set for your user"
|
||||
|
||||
|
||||
async def _rehash_password_if_needed(user_id: str, password: str, stored: str) -> None:
|
||||
"""Rehash legacy password (SHA256) to scrypt on successful login."""
|
||||
|
|
@ -92,6 +98,21 @@ def _matches_env_credentials(username: str, password: str, master_key: str | Non
|
|||
)
|
||||
|
||||
|
||||
def _admin_credentials_match(
|
||||
username: str, password: str, master_key: str, general_settings: Mapping[str, object]
|
||||
) -> bool:
|
||||
return general_settings.get("disable_env_credential_login") is not True and _matches_env_credentials(
|
||||
username, password, master_key
|
||||
)
|
||||
|
||||
|
||||
def _invalid_credentials_message(general_settings: Mapping[str, object]) -> str:
|
||||
"""One rejection message for unknown usernames and wrong passwords alike, so neither can be enumerated."""
|
||||
if is_env_credential_login_enabled(general_settings):
|
||||
return INVALID_UI_CREDENTIALS_MESSAGE
|
||||
return INVALID_USER_PASSWORD_MESSAGE
|
||||
|
||||
|
||||
def is_env_credential_login_enabled(general_settings: Mapping[str, object]) -> bool:
|
||||
"""Whether a login with UI_USERNAME/UI_PASSWORD (or the master-key fallback) can succeed.
|
||||
|
||||
|
|
@ -137,6 +158,7 @@ async def authenticate_user(
|
|||
password: str,
|
||||
master_key: str | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
throttle: LoginThrottle,
|
||||
general_settings: Mapping[str, object] = MappingProxyType({}),
|
||||
) -> LoginResult:
|
||||
"""
|
||||
|
|
@ -151,6 +173,7 @@ async def authenticate_user(
|
|||
password: Password from the login form
|
||||
master_key: Master key for the proxy (required)
|
||||
prisma_client: Prisma database client (optional)
|
||||
throttle: Failed sign-in accounting for this request's source address
|
||||
general_settings: Proxy general_settings, checked for
|
||||
`disable_password_login_when_sso_enabled` and
|
||||
`disable_env_credential_login`
|
||||
|
|
@ -163,9 +186,11 @@ async def authenticate_user(
|
|||
or if username/password login is disabled while SSO is configured
|
||||
|
||||
Recovery: an admin locked out of the UI by
|
||||
`disable_password_login_when_sso_enabled` can still administer the proxy over
|
||||
the API with the master key (Authorization: Bearer <master_key>), which never
|
||||
goes through this function. To restore UI username/password login, unset the
|
||||
`disable_password_login_when_sso_enabled`, or by the failed sign-in block in
|
||||
`throttle`, can still administer the proxy over the API with the master key
|
||||
(Authorization: Bearer <master_key>), which never goes through this function.
|
||||
No credential, the env admin credentials and the master key included, is
|
||||
exempt from the block. To restore UI username/password login, unset the
|
||||
setting in config.yaml (or the DB-persisted general_settings) and restart the
|
||||
proxy; this is a deliberate, auditable config change rather than a hidden
|
||||
bypass.
|
||||
|
|
@ -194,6 +219,19 @@ async def authenticate_user(
|
|||
code=500,
|
||||
)
|
||||
|
||||
attempt: Final = await throttle.attempt(username)
|
||||
return await _sign_in(username, password, master_key, prisma_client, attempt, general_settings)
|
||||
|
||||
|
||||
async def _sign_in(
|
||||
username: str,
|
||||
password: str,
|
||||
master_key: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
attempt: LoginAttempt,
|
||||
general_settings: Mapping[str, object],
|
||||
) -> LoginResult:
|
||||
admin_credentials_match: Final = _admin_credentials_match(username, password, master_key, general_settings)
|
||||
# Check if we can find the `username` in the db. On the UI, users can enter username=their email
|
||||
_user_row: LiteLLM_UserTable | None = None
|
||||
user_role: (
|
||||
|
|
@ -219,20 +257,13 @@ async def authenticate_user(
|
|||
- Login with UI_USERNAME and UI_PASSWORD
|
||||
- Login with Invite Link `user_email` and `password` combination
|
||||
"""
|
||||
if general_settings.get("disable_env_credential_login") is not True and _matches_env_credentials(
|
||||
username, password, master_key
|
||||
):
|
||||
if admin_credentials_match:
|
||||
# Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin
|
||||
user_role = LitellmUserRoles.PROXY_ADMIN
|
||||
user_id = LITELLM_PROXY_ADMIN_NAME
|
||||
|
||||
# we want the key created to have PROXY_ADMIN_PERMISSIONS
|
||||
key_user_id = LITELLM_PROXY_ADMIN_NAME
|
||||
if (
|
||||
os.getenv("PROXY_ADMIN_ID", None) is not None and os.environ["PROXY_ADMIN_ID"] == user_id
|
||||
) or user_id == LITELLM_PROXY_ADMIN_NAME:
|
||||
# checks if user is admin
|
||||
key_user_id = os.getenv("PROXY_ADMIN_ID", LITELLM_PROXY_ADMIN_NAME)
|
||||
key_user_id: Final = os.getenv("PROXY_ADMIN_ID", LITELLM_PROXY_ADMIN_NAME)
|
||||
|
||||
# Admin is Authe'd in - generate key for the UI to access Proxy
|
||||
|
||||
|
|
@ -294,6 +325,8 @@ async def authenticate_user(
|
|||
|
||||
key = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(user_info)
|
||||
|
||||
await attempt.succeeded()
|
||||
|
||||
return LoginResult(
|
||||
user_id=user_id,
|
||||
key=key,
|
||||
|
|
@ -349,6 +382,8 @@ async def authenticate_user(
|
|||
|
||||
key = response["token"]
|
||||
|
||||
await attempt.succeeded()
|
||||
|
||||
return LoginResult(
|
||||
user_id=user_id,
|
||||
key=key,
|
||||
|
|
@ -357,20 +392,17 @@ async def authenticate_user(
|
|||
login_method="username_password",
|
||||
)
|
||||
else:
|
||||
await attempt.failed()
|
||||
raise ProxyException(
|
||||
message=f"Invalid credentials used to access UI.\nNot valid credentials for {username}",
|
||||
message=_invalid_credentials_message(general_settings),
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="invalid_credentials",
|
||||
code=401,
|
||||
)
|
||||
else:
|
||||
env_credentials_hint: Final = (
|
||||
"\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file"
|
||||
if is_env_credential_login_enabled(general_settings)
|
||||
else ""
|
||||
)
|
||||
await attempt.failed()
|
||||
raise ProxyException(
|
||||
message=f"Invalid credentials used to access UI.{env_credentials_hint}",
|
||||
message=_invalid_credentials_message(general_settings),
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="invalid_credentials",
|
||||
code=401,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import ipaddress
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Final
|
||||
|
||||
from fastapi import Request
|
||||
|
|
@ -19,7 +20,7 @@ class NetworkContext(BaseModel):
|
|||
|
||||
class TrustedProxyConfig(BaseModel):
|
||||
use_forwarded_for: bool = False
|
||||
trusted_proxy_cidrs: list[str] = Field(default_factory=list)
|
||||
trusted_proxy_cidrs: Sequence[str] = Field(default_factory=tuple)
|
||||
|
||||
|
||||
def normalize_cidr_ranges(configured_ranges: Any, *, setting_name: str = "trusted_proxy_cidrs") -> list[str]:
|
||||
|
|
@ -49,6 +50,12 @@ def parse_trusted_proxy_ranges(
|
|||
return networks
|
||||
|
||||
|
||||
def _unmapped(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> ipaddress.IPv4Address | ipaddress.IPv6Address:
|
||||
if isinstance(addr, ipaddress.IPv6Address) and addr.ipv4_mapped is not None:
|
||||
return addr.ipv4_mapped
|
||||
return addr
|
||||
|
||||
|
||||
def ip_in_networks(client_ip: str | None, networks: list[TrustedProxyNetwork]) -> bool:
|
||||
if not client_ip or not networks:
|
||||
return False
|
||||
|
|
@ -56,7 +63,8 @@ def ip_in_networks(client_ip: str | None, networks: list[TrustedProxyNetwork]) -
|
|||
addr: Final = ipaddress.ip_address(client_ip.strip())
|
||||
except ValueError:
|
||||
return False
|
||||
return any(addr in network for network in networks)
|
||||
candidates: Final = (addr, _unmapped(addr))
|
||||
return any(candidate in network for candidate in candidates for network in networks)
|
||||
|
||||
|
||||
def _is_valid_ip(value: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -326,7 +326,12 @@ class RouteChecks:
|
|||
pass
|
||||
elif route.startswith("/v1/mcp/") or route.startswith("/mcp-rest/"):
|
||||
pass # authN/authZ handled by api itself
|
||||
elif RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token):
|
||||
elif RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token) or (
|
||||
valid_token.is_team_service_account
|
||||
and RouteChecks.check_route_access(
|
||||
route=route, allowed_routes=LiteLLMRoutes.team_service_account_key_routes.value
|
||||
)
|
||||
):
|
||||
pass
|
||||
elif valid_token.allowed_routes is not None:
|
||||
# check if route is in allowed_routes (exact match or prefix match)
|
||||
|
|
|
|||
|
|
@ -107,6 +107,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
|
|||
_safe_get_request_headers,
|
||||
_safe_get_request_query_params,
|
||||
_safe_set_request_parsed_body,
|
||||
is_opaque_audio_pass_through_request,
|
||||
populate_request_with_path_params,
|
||||
read_raw_json_body,
|
||||
rewrite_request_model,
|
||||
|
|
@ -631,9 +632,11 @@ def _apply_budget_limits_to_end_user_params(
|
|||
verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id)
|
||||
|
||||
|
||||
async def user_api_key_auth_websocket(websocket: WebSocket):
|
||||
# Accept the WebSocket connection
|
||||
async def user_api_key_auth_websocket(websocket: WebSocket) -> UserAPIKeyAuth:
|
||||
return await user_api_key_auth_websocket_for_model(websocket, model=websocket.query_params.get("model"))
|
||||
|
||||
|
||||
async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str | None) -> UserAPIKeyAuth:
|
||||
ws_scope: Final = websocket.scope or {}
|
||||
scope_headers: Final = list(ws_scope.get("headers") or [])
|
||||
# ``get_request_route`` falls back to ``request.url.path`` when
|
||||
|
|
@ -653,10 +656,6 @@ async def user_api_key_auth_websocket(websocket: WebSocket):
|
|||
|
||||
request._url = websocket.url
|
||||
|
||||
query_params: Final = websocket.query_params
|
||||
|
||||
model: Final = query_params.get("model")
|
||||
|
||||
async def return_body():
|
||||
return _realtime_request_body(model)
|
||||
|
||||
|
|
@ -1356,6 +1355,12 @@ async def _read_request_body_deferring_parse_failure(
|
|||
must run (resolving identity onto the request's trace) before the 400 goes
|
||||
out; the caller re-raises the returned exception once identity is seeded.
|
||||
"""
|
||||
if is_opaque_audio_pass_through_request(
|
||||
route=get_request_route(request=request),
|
||||
content_type=_safe_get_request_headers(request=request).get("content-type", ""),
|
||||
):
|
||||
_safe_set_request_parsed_body(request=request, parsed_body={}) # mutable-ok: the body cache stores a plain dict
|
||||
return {}, None # mutable-ok: request_data is a plain dict across the whole auth path
|
||||
try:
|
||||
parsed_body: Final = await _read_request_body(request=request)
|
||||
except ProxyException as parse_exception:
|
||||
|
|
|
|||
|
|
@ -1934,6 +1934,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
user_api_base: str | None = None,
|
||||
model: str | None = None,
|
||||
llm_router: Router | None = None,
|
||||
rate_limited_model: str | None = None,
|
||||
) -> tuple[dict, LiteLLMLoggingObj]:
|
||||
start_time: Final = datetime.now() # start before calling guardrail hooks
|
||||
|
||||
|
|
@ -2097,8 +2098,15 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# model_info when allow_client_pricing_override is set, so a caller
|
||||
# could otherwise spoof an unguarded model_info.id while requesting
|
||||
# a guarded alias and bypass guardrails (veria-ai HIGH on #29654).
|
||||
merged_for_requested: Final = (
|
||||
self.data
|
||||
if rate_limited_model is None
|
||||
else _check_and_merge_model_level_guardrails(
|
||||
data=self.data, llm_router=llm_router, trust_client_model_info=False, model_alias=rate_limited_model
|
||||
)
|
||||
)
|
||||
self.data = _check_and_merge_model_level_guardrails(
|
||||
data=self.data,
|
||||
data=merged_for_requested,
|
||||
llm_router=llm_router,
|
||||
trust_client_model_info=False,
|
||||
)
|
||||
|
|
@ -2163,7 +2171,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
|
||||
configured_fallbacks: Final = (
|
||||
self._configured_fallbacks(llm_router=llm_router, user_api_key_dict=user_api_key_dict)
|
||||
if llm_router is not None and not self.data.get("disable_fallbacks")
|
||||
if llm_router is not None
|
||||
else None
|
||||
)
|
||||
pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None
|
||||
|
|
@ -2208,7 +2216,6 @@ class ProxyBaseLLMRequestProcessing:
|
|||
original_model,
|
||||
fallback_models,
|
||||
)
|
||||
|
||||
try:
|
||||
for fallback_model in fallback_models:
|
||||
if fallback_model == original_model:
|
||||
|
|
@ -2231,6 +2238,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
model=fallback_model,
|
||||
route_type=route_type,
|
||||
llm_router=llm_router,
|
||||
rate_limited_model=original_model,
|
||||
)
|
||||
except ProxyRateLimitError:
|
||||
continue
|
||||
|
|
|
|||
|
|
@ -9,7 +9,11 @@ from fastapi import Request, UploadFile, status
|
|||
from typing_extensions import NotRequired, ReadOnly, Required
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB
|
||||
from litellm.constants import (
|
||||
AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX,
|
||||
CLIENT_REQUESTED_MODEL_SCOPE_KEY,
|
||||
MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB,
|
||||
)
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
|
|
@ -214,6 +218,14 @@ async def _read_request_body(request: Request | None) -> dict:
|
|||
return {}
|
||||
|
||||
|
||||
def is_opaque_audio_pass_through_request(route: str, content_type: str) -> bool:
|
||||
"""Azure Speech bodies (raw audio, multipart uploads) are forwarded byte for byte, so auth must not consume them."""
|
||||
media_type: Final = _normalize_media_type(content_type)
|
||||
return route.startswith(f"{AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX}/") and (
|
||||
media_type.startswith("audio/") or media_type == "multipart/form-data"
|
||||
)
|
||||
|
||||
|
||||
async def read_raw_json_body(request: Request | None) -> bytes | None:
|
||||
if request is None or _safe_get_request_parsed_body(request=request) is None:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -87,6 +87,12 @@ class SettingsStore(MutableMapping[str, JsonValue]):
|
|||
)
|
||||
self._deleted_runtime_keys = self._deleted_runtime_keys | frozenset((key,))
|
||||
|
||||
def clear(self) -> None:
|
||||
self._deleted_runtime_keys = frozenset(key for key in self._keys() if not self.owned_by_config(key))
|
||||
self._runtime_values = MappingProxyType(
|
||||
{key: value for key, value in self._runtime_values.items() if self.owned_by_config(key)}
|
||||
)
|
||||
|
||||
def __iter__(self) -> Iterator[str]:
|
||||
return iter(
|
||||
key
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload
|
||||
from urllib.parse import quote, unquote
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
from typing_extensions import LiteralString, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -136,6 +136,8 @@ class _SpendBatchManager(Protocol):
|
|||
class _SpendTransaction(Protocol):
|
||||
def batch_(self) -> _SpendBatchManager: ...
|
||||
|
||||
async def execute_raw(self, query: LiteralString, *args: object) -> int: ...
|
||||
|
||||
|
||||
class _SpendTransactionManager(Protocol):
|
||||
async def __aenter__(self) -> _SpendTransaction: ...
|
||||
|
|
@ -161,6 +163,44 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager:
|
|||
return tx
|
||||
|
||||
|
||||
# The per-team advisory lock the team endpoints hold while changing a roster (TEAM_ADVISORY_LOCK_SQL),
|
||||
# so the roster check below cannot interleave with their writes. A row lock would deadlock with the
|
||||
# access-group endpoints, which lock a team row after an access-group lock.
|
||||
_TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
|
||||
|
||||
# One statement adds every member's cost to their membership row. A missing row is created only
|
||||
# while the user is still on the team's roster, so a spend flush landing after a removal never
|
||||
# recreates the member.
|
||||
_TEAM_MEMBER_SPEND_SQL: Final = """
|
||||
INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend)
|
||||
SELECT p.user_id, p.team_id, p.cost, p.cost
|
||||
FROM unnest($1::text[], $2::text[], $3::float8[]) AS p(user_id, team_id, cost)
|
||||
WHERE EXISTS (
|
||||
SELECT 1 FROM "LiteLLM_TeamTable" t
|
||||
WHERE t.team_id = p.team_id
|
||||
AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', p.user_id))
|
||||
)
|
||||
OR EXISTS (SELECT 1 FROM "LiteLLM_TeamMembership" m WHERE m.user_id = p.user_id AND m.team_id = p.team_id)
|
||||
ON CONFLICT (user_id, team_id) DO UPDATE
|
||||
SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend,
|
||||
total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend
|
||||
"""
|
||||
|
||||
|
||||
async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_member_key: Mapping[str, float]) -> None:
|
||||
# key is "team_id::<value>::user_id::<value>"; locks are taken in sorted team_id order like the team endpoints
|
||||
rows: Final = sorted((key.split("::")[1], key.split("::")[3], cost) for key, cost in spend_by_member_key.items())
|
||||
team_ids: Final = tuple(team_id for team_id, _user_id, _cost in rows)
|
||||
for team_id in dict.fromkeys(team_ids):
|
||||
_ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id)
|
||||
_ = await transaction.execute_raw(
|
||||
_TEAM_MEMBER_SPEND_SQL,
|
||||
tuple(user_id for _team_id, user_id, _cost in rows),
|
||||
team_ids,
|
||||
tuple(cost for _team_id, _user_id, cost in rows),
|
||||
)
|
||||
|
||||
|
||||
def get_llm_router():
|
||||
"""The proxy's router, or None outside a running proxy.
|
||||
|
||||
|
|
@ -1685,21 +1725,7 @@ class DBSpendUpdateWriter:
|
|||
start_time = time.time()
|
||||
try:
|
||||
async with _spend_update_tx(prisma_client) as transaction:
|
||||
async with transaction.batch_() as batcher:
|
||||
# Sort by composite key for consistent lock ordering across pods to prevent deadlocks.
|
||||
# Key format "team_id::<v>::user_id::<v>" makes the string sort equivalent to sorting by (team_id, user_id).
|
||||
for key, response_cost in sorted(team_member_list_transactions.items()):
|
||||
# key is "team_id::<value>::user_id::<value>"
|
||||
team_id = key.split("::")[1]
|
||||
user_id = key.split("::")[3]
|
||||
|
||||
batcher.litellm_teammembership.update_many( # 'update_many' prevents error from being raised if no row exists
|
||||
where={"team_id": team_id, "user_id": user_id},
|
||||
data={
|
||||
"spend": {"increment": response_cost},
|
||||
"total_spend": {"increment": response_cost},
|
||||
},
|
||||
)
|
||||
await _write_team_member_spend(transaction, team_member_list_transactions)
|
||||
# Transaction succeeded, break out of retry loop
|
||||
break
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ if TYPE_CHECKING:
|
|||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.router import Router
|
||||
|
||||
COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"})
|
||||
COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr", "typesafe"})
|
||||
_NO_COMPRESSION: Final = "none"
|
||||
|
||||
# A ContextVar, not metadata: metadata reaches spend logs the caller can read, and a
|
||||
|
|
|
|||
|
|
@ -0,0 +1,74 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm.types.guardrails import (
|
||||
GuardrailEventHooks,
|
||||
Mode,
|
||||
SupportedGuardrailIntegrations,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
|
||||
TypeSafeGuardrailOptionalParams,
|
||||
)
|
||||
|
||||
from .typesafe import TypeSafeGuardrail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def _coerce_event_hook(
|
||||
mode: str | list[str] | Mode,
|
||||
) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode:
|
||||
if isinstance(mode, Mode):
|
||||
return mode
|
||||
if isinstance(mode, list):
|
||||
return [ # mutable-ok: CustomGuardrail event_hook contract wants a list
|
||||
GuardrailEventHooks(item) for item in mode
|
||||
]
|
||||
return GuardrailEventHooks(mode)
|
||||
|
||||
|
||||
def _optional_params(litellm_params: LitellmParams) -> TypeSafeGuardrailOptionalParams:
|
||||
value: Final = litellm_params.optional_params
|
||||
if isinstance(value, TypeSafeGuardrailOptionalParams):
|
||||
return value
|
||||
if isinstance(value, BaseModel):
|
||||
return TypeSafeGuardrailOptionalParams.model_validate(value.model_dump())
|
||||
return TypeSafeGuardrailOptionalParams()
|
||||
|
||||
|
||||
def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> TypeSafeGuardrail:
|
||||
import litellm
|
||||
|
||||
optional_params: Final = _optional_params(litellm_params)
|
||||
|
||||
_callback: Final = TypeSafeGuardrail(
|
||||
api_base=litellm_params.api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
model=litellm_params.model,
|
||||
relevance_threshold=optional_params.relevance_threshold,
|
||||
min_chars_to_evaluate=optional_params.min_chars_to_evaluate,
|
||||
max_result_chars_in_state=optional_params.max_result_chars_in_state,
|
||||
guardrail_name=guardrail["guardrail_name"],
|
||||
event_hook=_coerce_event_hook(litellm_params.mode),
|
||||
default_on=litellm_params.default_on or False,
|
||||
unreachable_fallback=(
|
||||
litellm_params.unreachable_fallback if "unreachable_fallback" in litellm_params.model_fields_set else None
|
||||
),
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback manager is untyped
|
||||
_callback
|
||||
)
|
||||
return _callback
|
||||
|
||||
|
||||
guardrail_initializer_registry: Final = { # mutable-ok: guardrail_registry discovery checks isinstance(registry, dict)
|
||||
SupportedGuardrailIntegrations.TYPESAFE.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
guardrail_class_registry: Final = { # mutable-ok: guardrail_registry discovery checks isinstance(registry, dict)
|
||||
SupportedGuardrailIntegrations.TYPESAFE.value: TypeSafeGuardrail,
|
||||
}
|
||||
416
litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py
Normal file
416
litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py
Normal file
|
|
@ -0,0 +1,416 @@
|
|||
"""TypeSafe (Jev) relevance-based compaction guardrail.
|
||||
|
||||
Instead of summarizing tool output, the guardrail asks TypeSafe's Jev model
|
||||
one yes/no question per completed tool exchange ("is this result still needed
|
||||
for the current task?") over ``POST {api_base}/v1/systemone`` and blanks the
|
||||
tool results Jev judges no longer relevant.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Literal
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
from httpx import Response as HttpxResponse
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.compression.compress import get_protected_indices
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # decorator is untyped in custom_guardrail
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # helper is untyped in http_handler
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_hooks.content_text import content_to_text
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
Logging as LiteLLMLoggingObj,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
|
||||
TypeSafeGuardrailConfigModel,
|
||||
)
|
||||
|
||||
DEFAULT_API_BASE: Final = "https://api.typesafe.ai"
|
||||
DEFAULT_MODEL: Final = "jev-latest"
|
||||
DEFAULT_RELEVANCE_THRESHOLD: Final = 0.2
|
||||
DEFAULT_MIN_CHARS_TO_EVALUATE: Final = 200
|
||||
DEFAULT_MAX_RESULT_CHARS_IN_STATE: Final = 4000
|
||||
_MAX_EXCHANGES_EVALUATED: Final = 200
|
||||
_JEV_TIMEOUT_SECONDS: Final = 30.0
|
||||
DROPPED_RESULT_TEXT: Final = (
|
||||
"[Tool result removed by TypeSafe compaction: judged no longer relevant to the current task]"
|
||||
)
|
||||
_ELISION_MARKER: Final = "\n... [middle truncated] ...\n"
|
||||
|
||||
|
||||
_STR_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||
|
||||
|
||||
def _as_str_object_dict(value: object) -> dict[str, object] | None:
|
||||
try:
|
||||
return _STR_OBJECT_DICT_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _as_object_list(value: object) -> list[object] | None:
|
||||
try:
|
||||
return _OBJECT_LIST_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _safe_response_text(response: HttpxResponse | None, limit: int = 500) -> str:
|
||||
if response is None:
|
||||
return ""
|
||||
try:
|
||||
text: Final = response.text
|
||||
except httpx.DecodingError:
|
||||
return "<undecodable response body>"
|
||||
return (text or "")[:limit]
|
||||
|
||||
|
||||
class _JevNoulAnswer(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, allow_inf_nan=False)
|
||||
|
||||
type: Literal["noul"]
|
||||
noul: Annotated[float, Field(ge=0.0, le=1.0)]
|
||||
|
||||
|
||||
class _JevSystemOneResponse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
answers: Mapping[str, _JevNoulAnswer]
|
||||
|
||||
|
||||
_JEV_RESPONSE_ADAPTER: Final = TypeAdapter(_JevSystemOneResponse)
|
||||
|
||||
|
||||
def _truncate_for_state(text: str, max_chars: int) -> str:
|
||||
"""Keeps the head and tail within ``max_chars`` so Jev sees both ends of a long result."""
|
||||
if len(text) <= max_chars:
|
||||
return text
|
||||
if max_chars <= len(_ELISION_MARKER):
|
||||
return text[:max_chars]
|
||||
budget: Final = max_chars - len(_ELISION_MARKER)
|
||||
head: Final = budget // 2
|
||||
return text[:head] + _ELISION_MARKER + text[len(text) - (budget - head) :]
|
||||
|
||||
|
||||
def _question_instructions(question_id: str) -> str:
|
||||
return (
|
||||
f"Is tool exchange `{question_id}` in `tool_exchanges` still needed by the assistant to "
|
||||
"complete `task`? Answer yes if its result contains information the assistant has not yet "
|
||||
"fully used or will need again; answer no if it is off-topic, superseded, or already "
|
||||
"incorporated into later messages."
|
||||
)
|
||||
|
||||
|
||||
def _tool_call_entry(tool_call: object) -> dict[str, object] | None:
|
||||
parsed_call = _as_str_object_dict(tool_call)
|
||||
if parsed_call is None:
|
||||
return None
|
||||
function = _as_str_object_dict(parsed_call.get("function"))
|
||||
fn = function if function is not None else parsed_call
|
||||
return {"name": fn.get("name"), "arguments": fn.get("arguments")} # mutable-ok: serialized to JSON
|
||||
|
||||
|
||||
def _tool_call_entries(assistant_message: Mapping[str, object]) -> tuple[dict[str, object], ...]:
|
||||
tool_calls: Final = _as_object_list(assistant_message.get("tool_calls"))
|
||||
if tool_calls is None:
|
||||
return ()
|
||||
return tuple(entry for tool_call in tool_calls if (entry := _tool_call_entry(tool_call)) is not None)
|
||||
|
||||
|
||||
def _protected_indices(messages: Sequence[Mapping[str, object]]) -> frozenset[int]:
|
||||
"""``get_protected_indices`` expanded over whole tool exchanges, so the most recent exchange is never evaluated."""
|
||||
protected: Final = frozenset(get_protected_indices(messages))
|
||||
return protected | frozenset(
|
||||
index
|
||||
for group in group_tool_exchanges(messages)
|
||||
if any(member in protected for member in group)
|
||||
for index in group
|
||||
)
|
||||
|
||||
|
||||
class TypeSafeGuardrail(CustomGuardrail):
|
||||
def __init__(
|
||||
self,
|
||||
api_base: str | None = None,
|
||||
api_key: str | None = None,
|
||||
model: str | None = None,
|
||||
relevance_threshold: float | None = None,
|
||||
min_chars_to_evaluate: int | None = None,
|
||||
max_result_chars_in_state: int | None = None,
|
||||
unreachable_fallback: str | None = None,
|
||||
guardrail_name: str | None = None,
|
||||
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None,
|
||||
default_on: bool = False,
|
||||
async_handler: AsyncHTTPHandler | None = None,
|
||||
) -> None:
|
||||
raw_api_base: Final = (api_base or get_secret_str("TYPESAFE_API_BASE") or DEFAULT_API_BASE).rstrip("/")
|
||||
self.typesafe_api_base = raw_api_base
|
||||
self.typesafe_api_key = api_key or get_secret_str("TYPESAFE_API_KEY")
|
||||
if not self.typesafe_api_key:
|
||||
raise ValueError(
|
||||
"TypeSafe guardrail requires an API key. Set `api_key` in the "
|
||||
"guardrail config or the TYPESAFE_API_KEY env var."
|
||||
)
|
||||
self.jev_model = model or DEFAULT_MODEL
|
||||
self.relevance_threshold = DEFAULT_RELEVANCE_THRESHOLD if relevance_threshold is None else relevance_threshold
|
||||
self.min_chars_to_evaluate = (
|
||||
DEFAULT_MIN_CHARS_TO_EVALUATE if min_chars_to_evaluate is None else min_chars_to_evaluate
|
||||
)
|
||||
self.max_result_chars_in_state = (
|
||||
DEFAULT_MAX_RESULT_CHARS_IN_STATE if max_result_chars_in_state is None else max_result_chars_in_state
|
||||
)
|
||||
self.unreachable_fallback: Literal["fail_closed", "fail_open"] = (
|
||||
"fail_closed" if unreachable_fallback == "fail_closed" else "fail_open"
|
||||
)
|
||||
self.async_handler: AsyncHTTPHandler = async_handler or get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.GuardrailCallback,
|
||||
)
|
||||
super().__init__( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.__init__ is untyped
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook=event_hook,
|
||||
default_on=default_on,
|
||||
)
|
||||
|
||||
def _handle_failure(self, error: str, log_detail: dict[str, object]) -> None:
|
||||
"""fail_open logs and returns; fail_closed raises a generic 502 (upstream bodies stay in server logs)."""
|
||||
if self.unreachable_fallback == "fail_open":
|
||||
verbose_proxy_logger.warning(
|
||||
"TypeSafe: %s; fail_open configured, forwarding request uncompacted. detail=%s",
|
||||
error,
|
||||
log_detail,
|
||||
)
|
||||
return
|
||||
verbose_proxy_logger.error("TypeSafe: %s. detail=%s", error, log_detail)
|
||||
raise HTTPException(status_code=502, detail={"error": error}) # mutable-ok: FastAPI wants a dict detail
|
||||
|
||||
def _candidate_exchanges(self, messages: Sequence[dict[str, object]]) -> tuple[tuple[int, ...], ...]:
|
||||
"""Completed tool exchanges eligible for evaluation: unprotected, and long enough to be worth a call."""
|
||||
protected: Final = _protected_indices(messages)
|
||||
candidates: Final = tuple(
|
||||
group
|
||||
for group in group_tool_exchanges(messages)
|
||||
if len(group) >= 2
|
||||
and messages[group[0]].get("role") == "assistant"
|
||||
and not any(member in protected for member in group)
|
||||
and len(self._exchange_tool_text(messages, group)) >= self.min_chars_to_evaluate
|
||||
)
|
||||
return candidates[-_MAX_EXCHANGES_EVALUATED:]
|
||||
|
||||
@staticmethod
|
||||
def _exchange_tool_text(messages: Sequence[dict[str, object]], group: tuple[int, ...]) -> str:
|
||||
return "".join(
|
||||
content_to_text(messages[index].get("content"))
|
||||
for index in group[1:]
|
||||
if messages[index].get("role") in ("tool", "function")
|
||||
)
|
||||
|
||||
def _build_state(
|
||||
self, messages: Sequence[dict[str, object]], candidates: tuple[tuple[int, ...], ...]
|
||||
) -> dict[str, object]:
|
||||
task: Final = next(
|
||||
(
|
||||
content_to_text(messages[index].get("content"))
|
||||
for index in range(len(messages) - 1, -1, -1)
|
||||
if messages[index].get("role") == "user"
|
||||
),
|
||||
"",
|
||||
)
|
||||
system: Final = "\n\n".join(
|
||||
content_to_text(message.get("content")) for message in messages if message.get("role") == "system"
|
||||
)
|
||||
tool_exchanges: Final = { # mutable-ok: accumulated once, serialized to JSON
|
||||
f"e{ordinal}": { # mutable-ok: serialized to JSON
|
||||
"tool_calls": _tool_call_entries(messages[group[0]]),
|
||||
"result": _truncate_for_state(
|
||||
self._exchange_tool_text(messages, group), self.max_result_chars_in_state
|
||||
),
|
||||
}
|
||||
for ordinal, group in enumerate(candidates)
|
||||
}
|
||||
return {"task": task, "system": system, "tool_exchanges": tool_exchanges} # mutable-ok: serialized to JSON
|
||||
|
||||
async def _call_systemone(
|
||||
self, state: dict[str, object], question_ids: Sequence[str]
|
||||
) -> _JevSystemOneResponse | None:
|
||||
"""Returns the response, or None when the service failed and fail_open applies."""
|
||||
payload: Final[dict[str, object]] = { # mutable-ok: serialized to JSON by httpx
|
||||
"model": self.jev_model,
|
||||
"state": state,
|
||||
"questions": { # mutable-ok: serialized to JSON
|
||||
question_id: { # mutable-ok: serialized to JSON
|
||||
"type": "noul",
|
||||
"instructions": _question_instructions(question_id),
|
||||
}
|
||||
for question_id in question_ids
|
||||
},
|
||||
}
|
||||
try:
|
||||
raw_response: HttpxResponse = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped
|
||||
url=f"{self.typesafe_api_base}/v1/systemone",
|
||||
json=payload,
|
||||
headers={ # mutable-ok: httpx header contract is a dict
|
||||
"Authorization": f"Bearer {self.typesafe_api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=_JEV_TIMEOUT_SECONDS,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e:
|
||||
detail: Final[dict[str, object]] = (
|
||||
{ # mutable-ok: log detail record
|
||||
"error_type": type(e).__name__,
|
||||
"detail": str(e),
|
||||
"status_code": e.response.status_code,
|
||||
"body": _safe_response_text(e.response),
|
||||
}
|
||||
if isinstance(e, httpx.HTTPStatusError)
|
||||
else {"error_type": type(e).__name__, "detail": str(e)} # mutable-ok: log detail record
|
||||
)
|
||||
self._handle_failure("TypeSafe evaluation service request failed", detail)
|
||||
return None
|
||||
if not 200 <= raw_response.status_code < 300:
|
||||
self._handle_failure(
|
||||
"TypeSafe evaluation service returned an error",
|
||||
{ # mutable-ok: log detail record
|
||||
"status_code": raw_response.status_code,
|
||||
"body": _safe_response_text(raw_response),
|
||||
},
|
||||
)
|
||||
return None
|
||||
try:
|
||||
body: Final[object] = raw_response.json() # pyright: ignore[reportAny] # httpx Response.json() is untyped
|
||||
except (ValueError, httpx.DecodingError, RecursionError):
|
||||
self._handle_failure(
|
||||
"TypeSafe evaluation service returned an unreadable response",
|
||||
{"body": _safe_response_text(raw_response)}, # mutable-ok: log detail record
|
||||
)
|
||||
return None
|
||||
try:
|
||||
return _JEV_RESPONSE_ADAPTER.validate_python(body)
|
||||
except ValidationError:
|
||||
self._handle_failure(
|
||||
"TypeSafe evaluation service returned unexpected response shape",
|
||||
{"body": _safe_response_text(raw_response)}, # mutable-ok: log detail record
|
||||
)
|
||||
return None
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: LiteLLMLoggingObj | None = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
if input_type != "request":
|
||||
return inputs
|
||||
|
||||
structured_messages: Final = _as_object_list(inputs.get("structured_messages"))
|
||||
if not structured_messages:
|
||||
return inputs
|
||||
parsed_messages: Final = tuple(_as_str_object_dict(m) for m in structured_messages)
|
||||
if any(m is None for m in parsed_messages):
|
||||
return inputs
|
||||
messages: Final = tuple(m for m in parsed_messages if m is not None)
|
||||
|
||||
candidates: Final = self._candidate_exchanges(messages)
|
||||
if not candidates:
|
||||
verbose_proxy_logger.debug("TypeSafe: no completed tool exchanges eligible for evaluation")
|
||||
return inputs
|
||||
|
||||
question_ids: Final = tuple(f"e{ordinal}" for ordinal in range(len(candidates)))
|
||||
state: Final = self._build_state(messages, candidates)
|
||||
|
||||
start_time: Final = time.monotonic()
|
||||
response: Final = await self._call_systemone(state, question_ids)
|
||||
end_time: Final = time.monotonic()
|
||||
if response is None:
|
||||
self.add_standard_logging_guardrail_information_to_request_data( # pyright: ignore[reportUnknownMemberType] # untyped base helper
|
||||
guardrail_json_response={ # mutable-ok: must stay JSON-serializable for shared logging
|
||||
"error": "TypeSafe evaluation unavailable; request forwarded uncompacted",
|
||||
"model": self.jev_model,
|
||||
},
|
||||
request_data=request_data,
|
||||
guardrail_status="guardrail_failed_to_respond",
|
||||
guardrail_provider="typesafe",
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
duration=end_time - start_time,
|
||||
)
|
||||
return inputs
|
||||
|
||||
dropped_ordinals: Final = frozenset(
|
||||
ordinal
|
||||
for ordinal in range(len(candidates))
|
||||
if (answer := response.answers.get(f"e{ordinal}")) is not None and answer.noul < self.relevance_threshold
|
||||
)
|
||||
dropped_tool_indices: Final[frozenset[int]] = frozenset(
|
||||
index
|
||||
for ordinal in dropped_ordinals
|
||||
for index in candidates[ordinal][1:]
|
||||
if messages[index].get("role") in ("tool", "function")
|
||||
)
|
||||
if not dropped_tool_indices:
|
||||
verbose_proxy_logger.debug("TypeSafe: all evaluated exchanges still relevant; request unchanged")
|
||||
return inputs
|
||||
|
||||
compacted_messages: Final = [ # mutable-ok: structured_messages contract is a list of dicts
|
||||
{**message, "content": DROPPED_RESULT_TEXT} # mutable-ok: JSON message row
|
||||
if index in dropped_tool_indices
|
||||
else message
|
||||
for index, message in enumerate(messages)
|
||||
]
|
||||
chars_removed: Final = sum(
|
||||
len(content_to_text(messages[index].get("content"))) - len(DROPPED_RESULT_TEXT)
|
||||
for index in dropped_tool_indices
|
||||
)
|
||||
exchanges_dropped: Final = len(dropped_ordinals)
|
||||
verbose_proxy_logger.info(
|
||||
"TypeSafe: evaluated %s tool exchange(s), dropped %s, ~%s chars removed",
|
||||
len(candidates),
|
||||
exchanges_dropped,
|
||||
chars_removed,
|
||||
)
|
||||
self.add_standard_logging_guardrail_information_to_request_data( # pyright: ignore[reportUnknownMemberType] # untyped base helper
|
||||
guardrail_json_response={ # mutable-ok: must stay JSON-serializable for shared logging
|
||||
"exchanges_evaluated": len(candidates),
|
||||
"exchanges_dropped": exchanges_dropped,
|
||||
"chars_removed": chars_removed,
|
||||
"model": self.jev_model,
|
||||
},
|
||||
request_data=request_data,
|
||||
guardrail_status="success",
|
||||
guardrail_provider="typesafe",
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
duration=end_time - start_time,
|
||||
)
|
||||
return {**inputs, "structured_messages": compacted_messages} # pyright: ignore[reportReturnType] # mutable-ok: inputs protocol is a plain dict # plain dicts satisfy AllMessageValues at runtime
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> type[TypeSafeGuardrailConfigModel] | None:
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
|
||||
TypeSafeGuardrailConfigModel,
|
||||
)
|
||||
|
||||
return TypeSafeGuardrailConfigModel
|
||||
|
|
@ -91,6 +91,11 @@ else:
|
|||
_REQUEST_RATE_LIMIT_DATA: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _sibling_counter_keys(window_key: str) -> tuple[str, str]:
|
||||
prefix: Final = window_key.removesuffix(":window")
|
||||
return f"{prefix}:requests", f"{prefix}:tokens"
|
||||
|
||||
|
||||
BATCH_RATE_LIMITER_SCRIPT: Final = """
|
||||
local results = {}
|
||||
local now = tonumber(ARGV[1])
|
||||
|
|
@ -106,6 +111,8 @@ for i = 1, #KEYS, 2 do
|
|||
local window_start = redis.call('GET', window_key)
|
||||
if not window_start or (now - tonumber(window_start)) >= window_size then
|
||||
-- Reset window and counter
|
||||
local prefix = string.sub(window_key, 1, -(#':window') - 1)
|
||||
redis.call('DEL', prefix .. ':requests', prefix .. ':tokens')
|
||||
redis.call('SET', window_key, tostring(now))
|
||||
redis.call('SET', counter_key, increment_value)
|
||||
redis.call('EXPIRE', window_key, window_size)
|
||||
|
|
@ -151,6 +158,7 @@ CHECK_AND_INCREMENT_BY_N_SCRIPT: Final = """
|
|||
local time_reply = redis.call('TIME')
|
||||
local now = tonumber(time_reply[1])
|
||||
local descriptor_count = #KEYS / 2
|
||||
local reset_windows = {}
|
||||
|
||||
-- Pass 1: read state, validate. Abort without writing if any over limit.
|
||||
local descriptor_state = {}
|
||||
|
|
@ -201,6 +209,11 @@ for i = 1, descriptor_count do
|
|||
|
||||
if window_expired then
|
||||
active_window_start = now
|
||||
if not reset_windows[window_key] then
|
||||
local prefix = string.sub(window_key, 1, -(#':window') - 1)
|
||||
redis.call('DEL', prefix .. ':requests', prefix .. ':tokens')
|
||||
reset_windows[window_key] = true
|
||||
end
|
||||
redis.call('SET', window_key, tostring(now))
|
||||
redis.call('SET', counter_key, increment)
|
||||
redis.call('EXPIRE', window_key, window_size)
|
||||
|
|
@ -1018,6 +1031,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
Implement sliding window rate limiting logic using in-memory cache operations.
|
||||
This follows the same logic as the Redis Lua script but uses async cache operations.
|
||||
"""
|
||||
async with self._check_and_increment_lock:
|
||||
return await self._in_memory_cache_sliding_window(keys=keys, now_int=now_int, window_size=window_size)
|
||||
|
||||
async def _in_memory_cache_sliding_window(
|
||||
self,
|
||||
keys: list[str],
|
||||
now_int: int,
|
||||
window_size: int,
|
||||
) -> CacheCounterValues:
|
||||
results: Final[list[CacheCounterValue | None]] = []
|
||||
|
||||
# Process each window/counter pair
|
||||
|
|
@ -1036,6 +1058,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
# Check if window exists and is valid
|
||||
if window_start is None or (now_int - int(window_start)) >= window_size:
|
||||
# Reset window and counter
|
||||
for sibling_counter_key in _sibling_counter_keys(window_key):
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
key=sibling_counter_key,
|
||||
value=0,
|
||||
ttl=window_size,
|
||||
litellm_parent_otel_span=None,
|
||||
local_only=True,
|
||||
)
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
key=window_key,
|
||||
value=str(now_int),
|
||||
|
|
@ -2048,6 +2078,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
|
||||
# Pass 2: apply increments.
|
||||
expired_windows: Final[Mapping[str, int]] = {
|
||||
meta["window_key"]: meta["window_size"]
|
||||
for meta, state in zip(per_counter_meta, descriptor_state)
|
||||
if state["window_expired"]
|
||||
}
|
||||
for window_key, window_size in expired_windows.items():
|
||||
for sibling_counter_key in _sibling_counter_keys(window_key):
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
key=sibling_counter_key,
|
||||
value=0,
|
||||
ttl=window_size,
|
||||
litellm_parent_otel_span=parent_otel_span,
|
||||
local_only=True,
|
||||
)
|
||||
statuses: Final[list[RateLimitStatus]] = []
|
||||
for meta, state in zip(per_counter_meta, descriptor_state):
|
||||
new_counter = meta["increment"] if state["window_expired"] else state["current"] + meta["increment"]
|
||||
|
|
|
|||
|
|
@ -9,8 +9,9 @@ from fastapi import HTTPException, status
|
|||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT
|
||||
from litellm.proxy._types import CommonProxyErrors
|
||||
from litellm.proxy.spend_tracking.daily_global_spend_rollup import GLOBAL_SPEND_TABLE_NAME, reconciled_through
|
||||
from litellm.proxy.spend_tracking.key_metadata_recovery import (
|
||||
attach_user_emails,
|
||||
recover_double_hashed_key_metadata,
|
||||
|
|
@ -146,15 +147,9 @@ class _AggregatedSpendData(TypedDict):
|
|||
totals: SpendMetrics
|
||||
|
||||
|
||||
class _GroupingSetsRow(SimpleNamespace):
|
||||
class _RollupMetricsRow(SimpleNamespace):
|
||||
date: str
|
||||
api_key: str | None
|
||||
model: str | None
|
||||
model_group: str | None
|
||||
custom_llm_provider: str | None
|
||||
mcp_namespaced_tool_name: str | None
|
||||
endpoint: str | None
|
||||
group_level: int
|
||||
spend: float | None
|
||||
prompt_tokens: int | None
|
||||
completion_tokens: int | None
|
||||
|
|
@ -172,12 +167,46 @@ class _GroupingSetsRow(SimpleNamespace):
|
|||
timed_requests: int | None
|
||||
|
||||
|
||||
class _EntityRollupRow(_GroupingSetsRow):
|
||||
class _GroupingSetsRow(_RollupMetricsRow):
|
||||
model: str | None
|
||||
model_group: str | None
|
||||
custom_llm_provider: str | None
|
||||
mcp_namespaced_tool_name: str | None
|
||||
endpoint: str | None
|
||||
group_level: int
|
||||
distinct_api_keys: int | None
|
||||
|
||||
|
||||
class _EntityRollupRow(_RollupMetricsRow):
|
||||
entity_id: str | None
|
||||
api_key_rolled: int
|
||||
|
||||
|
||||
def _reported_flat_cost(record: DailySpendRecord | _GroupingSetsRow) -> float:
|
||||
class _AggregatedQueryKwargs(TypedDict):
|
||||
table_name: ReadOnly[str]
|
||||
entity_id_field: ReadOnly[str]
|
||||
entity_id: ReadOnly[str | list[str] | None]
|
||||
start_date: ReadOnly[str]
|
||||
end_date: ReadOnly[str]
|
||||
model: ReadOnly[str | None]
|
||||
api_key: ReadOnly[str | list[str] | None]
|
||||
exclude_entity_ids: ReadOnly[list[str] | None]
|
||||
timezone_offset_minutes: ReadOnly[int | None]
|
||||
include_current_utc_day: ReadOnly[bool]
|
||||
|
||||
|
||||
_SqlQuery = tuple[str, list[str]]
|
||||
|
||||
|
||||
async def _query_raw_optional(
|
||||
prisma_client: PrismaClient, query: _SqlQuery | None
|
||||
) -> list[dict[str, object]] | None: # mutable-ok: prisma query_raw return shape
|
||||
if query is None:
|
||||
return None
|
||||
return await prisma_client.db.query_raw(query[0], *query[1])
|
||||
|
||||
|
||||
def _reported_flat_cost(record: DailySpendRecord | _RollupMetricsRow) -> float:
|
||||
"""Flat cost a daily row reports, which is zero unless PTU cost attribution is enabled.
|
||||
|
||||
Both read paths funnel through here: the paginated path reads the ``ptu_flat_cost``
|
||||
|
|
@ -699,71 +728,8 @@ def _ptu_flat_cost_select(table_name: str) -> str:
|
|||
return "0::float AS ptu_flat_cost"
|
||||
|
||||
|
||||
def _build_aggregated_sql_query(
|
||||
*,
|
||||
table_name: str,
|
||||
entity_id_field: str,
|
||||
entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
|
||||
start_date: str,
|
||||
end_date: str,
|
||||
model: str | None,
|
||||
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
|
||||
exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path
|
||||
timezone_offset_minutes: int | None = None,
|
||||
include_current_utc_day: bool = False,
|
||||
) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params
|
||||
"""Build a parameterized SQL GROUP BY query for aggregated daily activity.
|
||||
|
||||
Groups by (date, api_key, model, model_group, custom_llm_provider,
|
||||
mcp_namespaced_tool_name, endpoint) with SUMs on all metric columns.
|
||||
The entity_id column is intentionally omitted from GROUP BY to collapse
|
||||
rows across entities — this is where the biggest row reduction comes from.
|
||||
|
||||
Returns:
|
||||
Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw().
|
||||
"""
|
||||
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
|
||||
if pg_table is None:
|
||||
raise ValueError(f"Unknown table name: {table_name}")
|
||||
|
||||
adjusted_start, adjusted_end = _adjust_dates_for_timezone(
|
||||
start_date, end_date, timezone_offset_minutes, include_current_utc_day
|
||||
)
|
||||
|
||||
where_clause, sql_params = _build_aggregated_where_clause(
|
||||
entity_id_field=entity_id_field,
|
||||
entity_id=entity_id,
|
||||
adjusted_start=adjusted_start,
|
||||
adjusted_end=adjusted_end,
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
exclude_entity_ids=exclude_entity_ids,
|
||||
)
|
||||
|
||||
# Postgres computes every rollup level the response needs — per-date
|
||||
# totals, per-(date, model), per-(date, model, api_key), per-provider,
|
||||
# etc. — in a single pass via GROUPING SETS. The GROUPING() bitmask
|
||||
# encodes which level a row belongs to so Python can dispatch rows
|
||||
# straight into their buckets without re-summing. The leaf grouping
|
||||
# is omitted on purpose: nothing in the response shape needs it once
|
||||
# all the rollups are present.
|
||||
#
|
||||
# TODO: drop the successful_requests/failed_requests aggregates (and the
|
||||
# total_successful_requests metadata they feed) once the admin UI reads SGR
|
||||
# only from LiteLLM_DailyGatewayRequests. The remaining spend, token and
|
||||
# api_requests rollups are still served from here.
|
||||
sql_query: Final = f"""
|
||||
SELECT
|
||||
date,
|
||||
api_key,
|
||||
model,
|
||||
COALESCE(NULLIF(model_group, ''), model) AS model_group,
|
||||
custom_llm_provider,
|
||||
mcp_namespaced_tool_name,
|
||||
endpoint,
|
||||
GROUPING(date, api_key, model, COALESCE(NULLIF(model_group, ''), model),
|
||||
custom_llm_provider, mcp_namespaced_tool_name,
|
||||
endpoint) AS group_level,
|
||||
def _rollup_metric_select(table_name: str) -> str:
|
||||
return f"""
|
||||
SUM(spend)::float AS spend,
|
||||
{_ptu_flat_cost_select(table_name)},
|
||||
SUM(prompt_tokens)::bigint AS prompt_tokens,
|
||||
|
|
@ -779,27 +745,175 @@ def _build_aggregated_sql_query(
|
|||
SUM(successful_requests)::bigint AS successful_requests,
|
||||
SUM(failed_requests)::bigint AS failed_requests,
|
||||
SUM(total_response_time_ms)::bigint AS total_response_time_ms,
|
||||
SUM(timed_requests)::bigint AS timed_requests
|
||||
FROM "{pg_table}"
|
||||
WHERE {where_clause}
|
||||
SUM(timed_requests)::bigint AS timed_requests"""
|
||||
|
||||
|
||||
_MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)"
|
||||
|
||||
|
||||
_KEY_FREE_SOURCE_COLUMNS: Final = (
|
||||
"date",
|
||||
"model",
|
||||
"model_group",
|
||||
"custom_llm_provider",
|
||||
"mcp_namespaced_tool_name",
|
||||
"endpoint",
|
||||
"spend",
|
||||
"prompt_tokens",
|
||||
"completion_tokens",
|
||||
"cache_read_input_tokens",
|
||||
"cache_creation_input_tokens",
|
||||
"compression_saved_tokens",
|
||||
"compression_savings_spend",
|
||||
"prompt_caching_savings_spend",
|
||||
"gateway_injected_caching_savings_spend",
|
||||
"autorouter_savings_spend",
|
||||
"api_requests",
|
||||
"successful_requests",
|
||||
"failed_requests",
|
||||
"total_response_time_ms",
|
||||
"timed_requests",
|
||||
)
|
||||
|
||||
|
||||
async def global_rollup_reconciled_through(prisma_client: PrismaClient, query: _AggregatedQueryKwargs) -> str | None:
|
||||
"""The last day ``LiteLLM_DailyGlobalSpend`` can answer the key-free arm for, or None to
|
||||
read it all from the per-key table.
|
||||
|
||||
Only an unfiltered read of the user table sums to the same rows as the global table. The
|
||||
marker read is served from the config cache, so this is not a database round trip per request.
|
||||
"""
|
||||
if query["table_name"] != "litellm_dailyuserspend":
|
||||
return None
|
||||
if query["entity_id"] is not None or query["api_key"] is not None or query["exclude_entity_ids"]:
|
||||
return None
|
||||
try:
|
||||
return await reconciled_through(prisma_client)
|
||||
except Exception as exc: # noqa: BLE001 # the per-key table is always a correct answer, so never fail the read
|
||||
verbose_proxy_logger.warning("Could not read the daily global spend marker, using the per-key table: %s", exc)
|
||||
return None
|
||||
|
||||
|
||||
def _key_free_source(pg_table: str, where_clause: str, marker_param: str | None) -> str:
|
||||
"""The relation the key-free arm aggregates: the per-key table alone, or the global rollup
|
||||
for days through the marker plus the per-key table for the days still open after it."""
|
||||
if marker_param is None:
|
||||
return f'"{pg_table}"\n WHERE {where_clause}'
|
||||
columns: Final = ", ".join(_KEY_FREE_SOURCE_COLUMNS)
|
||||
return f"""(
|
||||
SELECT {columns}
|
||||
FROM "{GLOBAL_SPEND_TABLE_NAME}"
|
||||
WHERE {where_clause} AND date <= {marker_param}
|
||||
UNION ALL
|
||||
SELECT {columns}
|
||||
FROM "{pg_table}"
|
||||
WHERE {where_clause} AND date > {marker_param}
|
||||
) AS key_free_source"""
|
||||
|
||||
|
||||
def _build_aggregated_sql_query(
|
||||
*,
|
||||
table_name: str,
|
||||
entity_id_field: str,
|
||||
entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
|
||||
start_date: str,
|
||||
end_date: str,
|
||||
model: str | None,
|
||||
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
|
||||
exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path
|
||||
timezone_offset_minutes: int | None = None,
|
||||
include_current_utc_day: bool = False,
|
||||
global_rollup_through: str | None = None,
|
||||
) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params
|
||||
"""Build the GROUPING SETS query for aggregated daily activity.
|
||||
|
||||
Returns:
|
||||
Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw().
|
||||
"""
|
||||
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
|
||||
if pg_table is None:
|
||||
raise ValueError(f"Unknown table name: {table_name}")
|
||||
|
||||
adjusted_start, adjusted_end = _adjust_dates_for_timezone(
|
||||
start_date, end_date, timezone_offset_minutes, include_current_utc_day
|
||||
)
|
||||
|
||||
where_clause, where_params = _build_aggregated_where_clause(
|
||||
entity_id_field=entity_id_field,
|
||||
entity_id=entity_id,
|
||||
adjusted_start=adjusted_start,
|
||||
adjusted_end=adjusted_end,
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
exclude_entity_ids=exclude_entity_ids,
|
||||
)
|
||||
sentinel_param: Final = f"${len(where_params) + 1}"
|
||||
marker_param: Final = None if global_rollup_through is None else f"${len(where_params) + 2}"
|
||||
metric_select: Final = _rollup_metric_select(table_name)
|
||||
|
||||
# TODO: drop the successful_requests/failed_requests aggregates (and the
|
||||
# total_successful_requests metadata they feed) once the admin UI reads SGR
|
||||
# only from LiteLLM_DailyGatewayRequests. The remaining spend, token and
|
||||
# api_requests rollups are still served from here.
|
||||
sql_query: Final = f"""
|
||||
(SELECT
|
||||
date,
|
||||
NULL::text AS api_key,
|
||||
model,
|
||||
{_MODEL_GROUP_EXPR} AS model_group,
|
||||
custom_llm_provider,
|
||||
mcp_namespaced_tool_name,
|
||||
endpoint,
|
||||
(GROUPING(date) << 6) | {_API_KEY_ROLLED_UP_BIT}
|
||||
| GROUPING(model, {_MODEL_GROUP_EXPR},
|
||||
custom_llm_provider, mcp_namespaced_tool_name,
|
||||
endpoint) AS group_level,
|
||||
NULL::bigint AS distinct_api_keys,{metric_select}
|
||||
FROM {_key_free_source(pg_table, where_clause, marker_param)}
|
||||
GROUP BY GROUPING SETS (
|
||||
(date),
|
||||
(date, api_key),
|
||||
(date, model),
|
||||
(date, model, api_key),
|
||||
(date, COALESCE(NULLIF(model_group, ''), model)),
|
||||
(date, COALESCE(NULLIF(model_group, ''), model), api_key),
|
||||
(date, {_MODEL_GROUP_EXPR}),
|
||||
(date, custom_llm_provider),
|
||||
(date, custom_llm_provider, api_key),
|
||||
(date, mcp_namespaced_tool_name),
|
||||
(date, mcp_namespaced_tool_name, api_key),
|
||||
(date, endpoint),
|
||||
(date, endpoint, api_key),
|
||||
()
|
||||
))
|
||||
UNION ALL
|
||||
(WITH top_api_keys AS (
|
||||
SELECT api_key, COUNT(*) OVER () AS distinct_api_keys
|
||||
FROM "{pg_table}"
|
||||
WHERE {where_clause} AND api_key <> {sentinel_param}
|
||||
GROUP BY api_key
|
||||
ORDER BY SUM(spend) DESC, api_key
|
||||
LIMIT {USAGE_TOP_API_KEYS_LIMIT}
|
||||
)
|
||||
SELECT
|
||||
date,
|
||||
api_key,
|
||||
model,
|
||||
{_MODEL_GROUP_EXPR} AS model_group,
|
||||
custom_llm_provider,
|
||||
mcp_namespaced_tool_name,
|
||||
endpoint,
|
||||
GROUPING(date, api_key, model, {_MODEL_GROUP_EXPR},
|
||||
custom_llm_provider, mcp_namespaced_tool_name,
|
||||
endpoint) AS group_level,
|
||||
MAX(top_api_keys.distinct_api_keys) AS distinct_api_keys,{metric_select}
|
||||
FROM "{pg_table}" JOIN top_api_keys USING (api_key)
|
||||
WHERE {where_clause}
|
||||
GROUP BY GROUPING SETS (
|
||||
(date, api_key),
|
||||
(date, model, api_key),
|
||||
(date, {_MODEL_GROUP_EXPR}, api_key),
|
||||
(date, custom_llm_provider, api_key),
|
||||
(date, mcp_namespaced_tool_name, api_key),
|
||||
(date, endpoint, api_key)
|
||||
))
|
||||
"""
|
||||
|
||||
return sql_query, sql_params
|
||||
marker_params: Final = () if global_rollup_through is None else (global_rollup_through,)
|
||||
return sql_query, [*where_params, PTU_SENTINEL_API_KEY, *marker_params]
|
||||
|
||||
|
||||
def _build_entity_rollup_sql_query(
|
||||
|
|
@ -844,23 +958,7 @@ def _build_entity_rollup_sql_query(
|
|||
"{entity_id_field}" AS entity_id,
|
||||
date,
|
||||
api_key,
|
||||
GROUPING(api_key) AS api_key_rolled,
|
||||
SUM(spend)::float AS spend,
|
||||
{_ptu_flat_cost_select(table_name)},
|
||||
SUM(prompt_tokens)::bigint AS prompt_tokens,
|
||||
SUM(completion_tokens)::bigint AS completion_tokens,
|
||||
SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens,
|
||||
SUM(cache_creation_input_tokens)::bigint AS cache_creation_input_tokens,
|
||||
SUM(compression_saved_tokens)::bigint AS compression_saved_tokens,
|
||||
SUM(compression_savings_spend)::float AS compression_savings_spend,
|
||||
SUM(prompt_caching_savings_spend)::float AS prompt_caching_savings_spend,
|
||||
SUM(gateway_injected_caching_savings_spend)::float AS gateway_injected_caching_savings_spend,
|
||||
SUM(autorouter_savings_spend)::float AS autorouter_savings_spend,
|
||||
SUM(api_requests)::bigint AS api_requests,
|
||||
SUM(successful_requests)::bigint AS successful_requests,
|
||||
SUM(failed_requests)::bigint AS failed_requests,
|
||||
SUM(total_response_time_ms)::bigint AS total_response_time_ms,
|
||||
SUM(timed_requests)::bigint AS timed_requests
|
||||
GROUPING(api_key) AS api_key_rolled,{_rollup_metric_select(table_name)}
|
||||
FROM "{pg_table}"
|
||||
WHERE {where_clause}
|
||||
GROUP BY GROUPING SETS (
|
||||
|
|
@ -962,6 +1060,7 @@ async def _aggregate_spend_records(
|
|||
# current grouping set's key), 0 when the column is part of the key.
|
||||
_GROUP_GRAND_TOTAL: Final = 127 # 0b1111111 — all rolled up
|
||||
_GROUP_DATE: Final = 63 # 0b0111111 — only date kept
|
||||
_API_KEY_ROLLED_UP_BIT: Final = 32 # 0b0100000
|
||||
_GROUP_DATE_API_KEY: Final = 31 # 0b0011111
|
||||
_GROUP_DATE_MODEL: Final = 47 # 0b0101111
|
||||
_GROUP_DATE_MODEL_API_KEY: Final = 15 # 0b0001111
|
||||
|
|
@ -975,7 +1074,7 @@ _GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110
|
|||
_GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110
|
||||
|
||||
|
||||
def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics:
|
||||
def _record_to_spend_metrics(record: _RollupMetricsRow) -> SpendMetrics:
|
||||
"""Build a SpendMetrics directly from one already-aggregated rollup row.
|
||||
|
||||
SUM() over zero rows is SQL NULL, so rollup rows (notably the grand-total
|
||||
|
|
@ -1329,10 +1428,6 @@ async def get_daily_activity_aggregated(
|
|||
) -> SpendAnalyticsPaginatedResponse:
|
||||
"""Aggregated variant that returns the full result set (no pagination).
|
||||
|
||||
Uses SQL GROUP BY to aggregate rows in the database rather than fetching
|
||||
all individual rows into Python. This collapses rows across entities
|
||||
(users/teams/orgs), reducing ~150k rows to ~2-3k grouped rows.
|
||||
|
||||
include_entity_breakdown runs a small companion rollup query and folds
|
||||
`breakdown.entities` onto the response, as entity-scoped views like Team Usage need.
|
||||
|
||||
|
|
@ -1351,7 +1446,7 @@ async def get_daily_activity_aggregated(
|
|||
)
|
||||
|
||||
try:
|
||||
sql_query, sql_params = _build_aggregated_sql_query(
|
||||
query_kwargs: Final = _AggregatedQueryKwargs(
|
||||
table_name=table_name,
|
||||
entity_id_field=entity_id_field,
|
||||
entity_id=entity_id,
|
||||
|
|
@ -1363,36 +1458,19 @@ async def get_daily_activity_aggregated(
|
|||
timezone_offset_minutes=timezone_offset_minutes,
|
||||
include_current_utc_day=include_current_utc_day,
|
||||
)
|
||||
sql_query, sql_params = _build_aggregated_sql_query(
|
||||
**query_kwargs,
|
||||
global_rollup_through=await global_rollup_reconciled_through(prisma_client, query_kwargs),
|
||||
)
|
||||
entity_query: Final = _build_entity_rollup_sql_query(**query_kwargs) if include_entity_breakdown else None
|
||||
|
||||
entity_query: Final = (
|
||||
_build_entity_rollup_sql_query(
|
||||
table_name=table_name,
|
||||
entity_id_field=entity_id_field,
|
||||
entity_id=entity_id,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
exclude_entity_ids=exclude_entity_ids,
|
||||
timezone_offset_minutes=timezone_offset_minutes,
|
||||
include_current_utc_day=include_current_utc_day,
|
||||
)
|
||||
if include_entity_breakdown
|
||||
else None
|
||||
raw_rows, raw_entity_rows = await asyncio.gather(
|
||||
prisma_client.db.query_raw(sql_query, *sql_params),
|
||||
_query_raw_optional(prisma_client, entity_query),
|
||||
)
|
||||
|
||||
# Execute the GROUPING SETS query (one row per rollup level), alongside
|
||||
# the per-entity companion rollup when the caller wants entities.
|
||||
raw_rows, raw_entity_rows = (
|
||||
await asyncio.gather(
|
||||
prisma_client.db.query_raw(sql_query, *sql_params),
|
||||
prisma_client.db.query_raw(entity_query[0], *entity_query[1]),
|
||||
)
|
||||
if entity_query is not None
|
||||
else (await prisma_client.db.query_raw(sql_query, *sql_params), None)
|
||||
)
|
||||
|
||||
records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or [])]
|
||||
records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or ())]
|
||||
total_api_keys: Final = next((r.distinct_api_keys for r in records if r.distinct_api_keys is not None), 0)
|
||||
|
||||
# The grouping-sets dispatcher places each row directly in its bucket
|
||||
# using the row's GROUPING() bitmask. No Python-side summing needed.
|
||||
|
|
@ -1446,6 +1524,8 @@ async def get_daily_activity_aggregated(
|
|||
page=1,
|
||||
total_pages=1,
|
||||
has_more=False,
|
||||
api_key_limit=USAGE_TOP_API_KEYS_LIMIT,
|
||||
total_api_keys=total_api_keys,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -510,6 +510,16 @@ def _get_user_in_team(team_table: LiteLLM_TeamTableCachedObj, user_id: str | Non
|
|||
return None
|
||||
|
||||
|
||||
def _get_caller_team_role(
|
||||
team_table: LiteLLM_TeamTableCachedObj,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> Literal["admin", "user"] | None:
|
||||
if user_api_key_dict.is_team_service_account and user_api_key_dict.team_id == team_table.team_id:
|
||||
return "user"
|
||||
member: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
|
||||
return None if member is None else member.role
|
||||
|
||||
|
||||
def _calculate_key_rotation_time(rotation_interval: str) -> datetime:
|
||||
"""
|
||||
Helper function to calculate the next rotation time for a key based on the rotation interval.
|
||||
|
|
@ -604,7 +614,7 @@ def _team_key_operation_team_member_check(
|
|||
detail=f"User={assigned_user_id} not assigned to team={team_table.team_id}",
|
||||
)
|
||||
|
||||
team_member_object: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
|
||||
caller_team_role: Final = _get_caller_team_role(team_table=team_table, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
is_admin: Final = (
|
||||
user_api_key_dict.user_role is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
|
|
@ -612,22 +622,22 @@ def _team_key_operation_team_member_check(
|
|||
|
||||
if is_admin:
|
||||
return True
|
||||
elif team_member_object is None:
|
||||
elif caller_team_role is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"User={user_api_key_dict.user_id} not assigned to team={team_table.team_id}",
|
||||
)
|
||||
elif (
|
||||
"allowed_team_member_roles" in team_key_generation
|
||||
and team_member_object.role not in team_key_generation["allowed_team_member_roles"]
|
||||
and caller_team_role not in team_key_generation["allowed_team_member_roles"]
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Team member role {team_member_object.role} not in allowed_team_member_roles={team_key_generation['allowed_team_member_roles']}",
|
||||
detail=f"Team member role {caller_team_role} not in allowed_team_member_roles={team_key_generation['allowed_team_member_roles']}",
|
||||
)
|
||||
|
||||
TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint(
|
||||
team_member_object=team_member_object,
|
||||
team_member_role=caller_team_role,
|
||||
team_table=team_table,
|
||||
route=route,
|
||||
)
|
||||
|
|
@ -748,6 +758,12 @@ def key_generation_check(
|
|||
Check if admin has restricted key creation to certain roles for teams or individuals
|
||||
"""
|
||||
|
||||
if user_api_key_dict.is_team_service_account and data.team_id != user_api_key_dict.team_id:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"Service account keys can only create keys for their own team. team_id={user_api_key_dict.team_id}",
|
||||
)
|
||||
|
||||
## check if key is for team or individual
|
||||
is_team_key: Final = _is_team_key(data=data)
|
||||
_is_admin: Final = (
|
||||
|
|
@ -2233,6 +2249,14 @@ async def generate_service_account_key_fn(
|
|||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
if data.metadata is None or data.metadata.get("service_account_id") is None:
|
||||
service_account_id: Final = data.key_alias or str(uuid.uuid4())
|
||||
stamped_metadata: Final = { # mutable-ok: GenerateKeyRequest.metadata is a plain dict field
|
||||
**(data.metadata or MappingProxyType({})),
|
||||
"service_account_id": service_account_id,
|
||||
}
|
||||
data.metadata = stamped_metadata # rebind-ok: the request carries the stamp so it persists on the key
|
||||
|
||||
verbose_proxy_logger.debug("entered /key/generate")
|
||||
|
||||
custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_generate_hook(
|
||||
|
|
@ -3891,8 +3915,10 @@ async def validate_key_team_change(
|
|||
detail=f"Key={key.token} has a rpm_limit={key.rpm_limit} which is greater than the team's rpm_limit={team.rpm_limit}.",
|
||||
)
|
||||
|
||||
team_table: Final = cast(LiteLLM_TeamTableCachedObj, team)
|
||||
|
||||
# Check if the key's user_id is a member of the team
|
||||
member_object: Final = _get_user_in_team(team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id)
|
||||
member_object: Final = _get_user_in_team(team_table=team_table, user_id=key.user_id)
|
||||
if key.user_id is not None:
|
||||
if not member_object:
|
||||
raise HTTPException(
|
||||
|
|
@ -3908,8 +3934,8 @@ async def validate_key_team_change(
|
|||
team_obj=team,
|
||||
)
|
||||
or TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint(
|
||||
team_member_object=member_object,
|
||||
team_table=cast(LiteLLM_TeamTableCachedObj, team),
|
||||
team_member_role=None if member_object is None else member_object.role,
|
||||
team_table=team_table,
|
||||
route=KeyManagementRoutes.KEY_UPDATE.value,
|
||||
)
|
||||
):
|
||||
|
|
|
|||
|
|
@ -47,7 +47,7 @@ except ImportError:
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
|
||||
from litellm.constants import LITELLM_PROXY_ADMIN_NAME, MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH
|
||||
from litellm.proxy._experimental.mcp_server.utils import (
|
||||
LITELLM_MCP_SERVER_DESCRIPTION,
|
||||
LITELLM_MCP_SERVER_NAME,
|
||||
|
|
@ -145,6 +145,7 @@ if MCP_AVAILABLE:
|
|||
get_user_env_vars,
|
||||
get_user_env_vars_bulk,
|
||||
get_user_oauth_credential,
|
||||
list_server_user_credentials,
|
||||
list_user_oauth_credentials,
|
||||
mcp_oauth_token_identity,
|
||||
merge_user_env_vars,
|
||||
|
|
@ -180,6 +181,7 @@ if MCP_AVAILABLE:
|
|||
MCPApprovalStatus,
|
||||
MCPOAuthUserCredentialRequest,
|
||||
MCPOAuthUserCredentialStatus,
|
||||
MCPServerUserCredentialListItem,
|
||||
MCPSubmissionsSummary,
|
||||
MCPTransport,
|
||||
MCPUserCredentialListItem,
|
||||
|
|
@ -221,6 +223,7 @@ if MCP_AVAILABLE:
|
|||
MCPAuth,
|
||||
MCPCredentials,
|
||||
MCPGatewaySessionsResponse,
|
||||
MCPGatewaySessionsTerminateResponse,
|
||||
normalize_upstream_header_name,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
|
@ -662,6 +665,31 @@ if MCP_AVAILABLE:
|
|||
"""
|
||||
return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
|
||||
def _resolve_credential_target_user_id(user_api_key_dict: UserAPIKeyAuth, requested_user_id: str | None) -> str:
|
||||
"""The user whose stored MCP credential a request acts on.
|
||||
|
||||
Defaults to the caller. Naming another user is a revocation and needs
|
||||
``PROXY_ADMIN``; a read-only admin or a regular user gets 403.
|
||||
"""
|
||||
caller_user_id: Final = user_api_key_dict.user_id or ""
|
||||
if requested_user_id is not None and requested_user_id != caller_user_id:
|
||||
if not _user_is_full_admin(user_api_key_dict):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict
|
||||
"error": "Proxy admin access required to revoke another user's MCP credential.",
|
||||
},
|
||||
)
|
||||
return requested_user_id
|
||||
if not caller_user_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail={
|
||||
"error": "User ID not found in token"
|
||||
}, # mutable-ok: FastAPI HTTPException detail requires a plain dict
|
||||
)
|
||||
return caller_user_id
|
||||
|
||||
def _is_restricted_virtual_key_request(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
"""Best-effort detection for route-restricted virtual keys.
|
||||
|
||||
|
|
@ -1373,6 +1401,41 @@ if MCP_AVAILABLE:
|
|||
|
||||
return get_mcp_gateway_sessions_report()
|
||||
|
||||
@router.delete(
|
||||
"/sessions",
|
||||
description=(
|
||||
"Force-close live stateful MCP gateway sessions on this proxy worker, selected by session id prefix "
|
||||
"and/or by the LiteLLM user that opened them (proxy admin only)."
|
||||
),
|
||||
dependencies=(Depends(user_api_key_auth),),
|
||||
response_model=MCPGatewaySessionsTerminateResponse,
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def delete_mcp_gateway_sessions(
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
session_id_prefix: Annotated[str | None, Query(min_length=MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH)] = None,
|
||||
user_id: Annotated[str | None, Query(min_length=1)] = None,
|
||||
) -> MCPGatewaySessionsTerminateResponse:
|
||||
if not _user_is_full_admin(user_api_key_dict):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict
|
||||
"error": "Proxy admin access required to terminate MCP gateway sessions.",
|
||||
},
|
||||
)
|
||||
if session_id_prefix is None and user_id is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict
|
||||
"error": "Provide session_id_prefix and/or user_id to select the sessions to terminate.",
|
||||
},
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
terminate_mcp_gateway_sessions,
|
||||
)
|
||||
|
||||
return await terminate_mcp_gateway_sessions(session_id_prefix=session_id_prefix, user_id=user_id)
|
||||
|
||||
@router.get(
|
||||
"/server/submissions",
|
||||
description="Returns all MCP servers submitted by non-admin users (admin review queue). Mirrors GET /guardrails/submissions.",
|
||||
|
|
@ -2254,14 +2317,17 @@ if MCP_AVAILABLE:
|
|||
_invalidate_byok_cred_cache,
|
||||
)
|
||||
|
||||
_invalidate_byok_cred_cache(user_id, server_id)
|
||||
await _invalidate_byok_cred_cache(user_id, server_id)
|
||||
return MCPUserCredentialResponse(server_id=server_id, has_credential=True)
|
||||
# save=False: credential not persisted
|
||||
return MCPUserCredentialResponse(server_id=server_id, has_credential=False)
|
||||
|
||||
@router.delete(
|
||||
"/server/{server_id}/user-credential",
|
||||
description="Delete the calling user's stored API key for a BYOK MCP server",
|
||||
description=(
|
||||
"Delete the calling user's stored API key for a BYOK MCP server. "
|
||||
"A proxy admin may pass user_id to revoke another user's stored key."
|
||||
),
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=MCPUserCredentialResponse,
|
||||
)
|
||||
|
|
@ -2269,24 +2335,20 @@ if MCP_AVAILABLE:
|
|||
async def delete_mcp_user_credential(
|
||||
server_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
user_id: Annotated[str | None, Query(min_length=1)] = None,
|
||||
):
|
||||
"""Remove the calling user's BYOK credential."""
|
||||
"""Remove the target user's BYOK credential (the caller unless an admin names another user)."""
|
||||
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
|
||||
user_id: Final = user_api_key_dict.user_id or ""
|
||||
if not user_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail={"error": "User ID not found in token"},
|
||||
)
|
||||
target_user_id: Final = _resolve_credential_target_user_id(user_api_key_dict, user_id)
|
||||
try:
|
||||
await delete_user_credential(prisma_client, user_id, server_id)
|
||||
await delete_user_credential(prisma_client, target_user_id, server_id)
|
||||
except RecordNotFoundError:
|
||||
pass # Already deleted or didn't exist
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_invalidate_byok_cred_cache,
|
||||
)
|
||||
|
||||
_invalidate_byok_cred_cache(user_id, server_id)
|
||||
await _invalidate_byok_cred_cache(target_user_id, server_id)
|
||||
return MCPUserCredentialResponse(server_id=server_id, has_credential=False)
|
||||
|
||||
# ── OAuth2 user-credential endpoints ──────────────────────────────────────
|
||||
|
|
@ -2362,7 +2424,10 @@ if MCP_AVAILABLE:
|
|||
|
||||
@router.delete(
|
||||
"/server/{server_id}/oauth-user-credential",
|
||||
description="Revoke the calling user's stored OAuth2 token for an MCP server",
|
||||
description=(
|
||||
"Revoke the calling user's stored OAuth2 token for an MCP server. "
|
||||
"A proxy admin may pass user_id to revoke another user's stored token."
|
||||
),
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=MCPOAuthUserCredentialStatus,
|
||||
)
|
||||
|
|
@ -2370,29 +2435,25 @@ if MCP_AVAILABLE:
|
|||
async def delete_mcp_oauth_user_credential(
|
||||
server_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
user_id: Annotated[str | None, Query(min_length=1)] = None,
|
||||
):
|
||||
"""Revoke/delete the user's OAuth2 credential."""
|
||||
"""Revoke the target user's OAuth2 credential (the caller unless an admin names another user)."""
|
||||
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
|
||||
user_id: Final = user_api_key_dict.user_id or ""
|
||||
if not user_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail={"error": "User ID not found in token"},
|
||||
)
|
||||
target_user_id: Final = _resolve_credential_target_user_id(user_api_key_dict, user_id)
|
||||
# Only delete if the stored credential is actually an OAuth2 token.
|
||||
# This prevents accidentally deleting a BYOK credential if one exists
|
||||
# for the same (user_id, server_id) pair.
|
||||
cred_to_delete: Final = await get_user_oauth_credential(prisma_client, user_id, server_id)
|
||||
cred_to_delete: Final = await get_user_oauth_credential(prisma_client, target_user_id, server_id)
|
||||
if cred_to_delete is not None:
|
||||
try:
|
||||
await delete_user_credential(prisma_client, user_id, server_id)
|
||||
await delete_user_credential(prisma_client, target_user_id, server_id)
|
||||
except RecordNotFoundError:
|
||||
pass # Already gone — treat as a successful delete
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
|
||||
await global_mcp_server_manager.invalidate_user_oauth_token_cache(user_id, server_id)
|
||||
await global_mcp_server_manager.invalidate_user_oauth_token_cache(target_user_id, server_id)
|
||||
return MCPOAuthUserCredentialStatus(
|
||||
server_id=server_id,
|
||||
has_credential=False,
|
||||
|
|
@ -2481,6 +2542,30 @@ if MCP_AVAILABLE:
|
|||
)
|
||||
return items
|
||||
|
||||
@router.get(
|
||||
"/server/{server_id}/user-credentials",
|
||||
description="List every user's stored BYOK or OAuth2 credential for an MCP server (admin only, no secrets)",
|
||||
dependencies=(Depends(user_api_key_auth),),
|
||||
response_model=list[MCPServerUserCredentialListItem],
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def list_mcp_server_user_credentials(
|
||||
server_id: str,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
) -> tuple[MCPServerUserCredentialListItem, ...]:
|
||||
if user_api_key_dict.user_role not in (
|
||||
LitellmUserRoles.PROXY_ADMIN,
|
||||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict
|
||||
"error": "Admin access required to view MCP server user credentials.",
|
||||
},
|
||||
)
|
||||
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
|
||||
return await list_server_user_credentials(prisma_client, server_id)
|
||||
|
||||
# ── Per-user MCP env var endpoints ────────────────────────────────────────
|
||||
|
||||
async def _authorize_and_fetch_mcp_server(
|
||||
|
|
|
|||
|
|
@ -2107,7 +2107,7 @@ def _handle_multi_valued_attribute_update(path: str, op_type: str, value: object
|
|||
except ValidationError:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": f"Invalid value for {base}: expected a list of objects with a 'value' sub-attribute"},
|
||||
detail={"error": f"Invalid value for {base}: expected a list of objects or strings"},
|
||||
)
|
||||
|
||||
dumped: Final = [attr.model_dump(exclude_none=True) for attr in attrs]
|
||||
|
|
|
|||
|
|
@ -2904,10 +2904,15 @@ async def _process_team_members(
|
|||
if member_allowed_models is None and team_default_member_models:
|
||||
member_allowed_models = team_default_member_models
|
||||
|
||||
if isinstance(data.member, Member):
|
||||
requested_members: Final[Sequence[Member]] = (
|
||||
(data.member,) if isinstance(data.member, Member) else tuple(data.member)
|
||||
)
|
||||
for m in requested_members:
|
||||
if _member_already_in_team(m, complete_team_data):
|
||||
continue
|
||||
try:
|
||||
updated_user, updated_tm = await add_new_member(
|
||||
new_member=data.member,
|
||||
new_member=m,
|
||||
max_budget_in_team=data.max_budget_in_team,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -2921,34 +2926,11 @@ async def _process_team_members(
|
|||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": f"Unable to add user - {data.member}, to team - {data.team_id}, for reason - {e}"},
|
||||
detail={"error": f"Unable to add user - {m}, to team - {data.team_id}, for reason - {e}"},
|
||||
)
|
||||
updated_users.append(updated_user)
|
||||
if updated_tm is not None:
|
||||
updated_team_memberships.append(updated_tm)
|
||||
elif isinstance(data.member, list):
|
||||
for m in data.member:
|
||||
try:
|
||||
updated_user, updated_tm = await add_new_member(
|
||||
new_member=m,
|
||||
max_budget_in_team=data.max_budget_in_team,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
team_id=data.team_id,
|
||||
default_team_budget_id=default_team_budget_id,
|
||||
allowed_models=member_allowed_models,
|
||||
budget_duration=data.budget_duration,
|
||||
tx=tx,
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": f"Unable to add user - {m}, to team - {data.team_id}, for reason - {e}"},
|
||||
)
|
||||
updated_users.append(updated_user)
|
||||
if updated_tm is not None:
|
||||
updated_team_memberships.append(updated_tm)
|
||||
|
||||
return updated_users, updated_team_memberships
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import Final
|
||||
from typing import Final, Literal
|
||||
|
||||
from litellm.proxy._types import (
|
||||
KeyManagementRoutes,
|
||||
|
|
@ -6,7 +6,6 @@ from litellm.proxy._types import (
|
|||
LiteLLM_VerificationToken,
|
||||
LiteLLMRoutes,
|
||||
LitellmUserRoles,
|
||||
Member,
|
||||
ProxyErrorTypes,
|
||||
ProxyException,
|
||||
UserAPIKeyAuth,
|
||||
|
|
@ -27,7 +26,6 @@ DEFAULT_TEAM_MEMBER_PERMISSIONS: Final = BASELINE_TEAM_MEMBER_PERMISSIONS
|
|||
class TeamMemberPermissionChecks:
|
||||
@staticmethod
|
||||
def get_permissions_for_team_member(
|
||||
team_member_object: Member,
|
||||
team_table: LiteLLM_TeamTableCachedObj,
|
||||
) -> list[KeyManagementRoutes]:
|
||||
"""
|
||||
|
|
@ -67,7 +65,7 @@ class TeamMemberPermissionChecks:
|
|||
Main handler for checking if a team member can update a key
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_get_user_in_team,
|
||||
_get_caller_team_role,
|
||||
)
|
||||
|
||||
# 1. Don't execute these checks if the user role is proxy admin
|
||||
|
|
@ -87,12 +85,11 @@ class TeamMemberPermissionChecks:
|
|||
check_db_only=True,
|
||||
)
|
||||
|
||||
# 4. Extract `Member` object from `team_table`
|
||||
key_assigned_user_in_team: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
|
||||
caller_team_role: Final = _get_caller_team_role(team_table=team_table, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
# 5. Check if the team member has permissions for the endpoint
|
||||
# 4. Check if the team member has permissions for the endpoint
|
||||
has_permission: Final = TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint(
|
||||
team_member_object=key_assigned_user_in_team,
|
||||
team_member_role=caller_team_role,
|
||||
team_table=team_table,
|
||||
route=route,
|
||||
)
|
||||
|
|
@ -106,7 +103,7 @@ class TeamMemberPermissionChecks:
|
|||
|
||||
@staticmethod
|
||||
def does_team_member_have_permissions_for_endpoint(
|
||||
team_member_object: Member | None,
|
||||
team_member_role: Literal["admin", "user"] | None,
|
||||
team_table: LiteLLM_TeamTableCachedObj,
|
||||
route: str,
|
||||
) -> bool | None:
|
||||
|
|
@ -116,13 +113,12 @@ class TeamMemberPermissionChecks:
|
|||
|
||||
# permission checks only run for non-admin users
|
||||
# Non-Admin user trying to access information about a team's key
|
||||
if team_member_object is None:
|
||||
if team_member_role is None:
|
||||
return False
|
||||
if team_member_object.role == "admin":
|
||||
if team_member_role == "admin":
|
||||
return True
|
||||
|
||||
_team_member_permissions: Final = TeamMemberPermissionChecks.get_permissions_for_team_member(
|
||||
team_member_object=team_member_object,
|
||||
team_table=team_table,
|
||||
)
|
||||
team_member_permissions = TeamMemberPermissionChecks._get_list_of_route_enum_as_str(_team_member_permissions)
|
||||
|
|
@ -156,7 +152,7 @@ class TeamMemberPermissionChecks:
|
|||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_get_user_in_team,
|
||||
_get_caller_team_role,
|
||||
)
|
||||
|
||||
# No-op when the request does not assign any access groups.
|
||||
|
|
@ -177,20 +173,19 @@ class TeamMemberPermissionChecks:
|
|||
),
|
||||
)
|
||||
|
||||
team_member_object: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
|
||||
caller_team_role: Final = _get_caller_team_role(team_table=team_table, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
# Team admins always bypass (consistent with other member-permission checks).
|
||||
if team_member_object is not None and team_member_object.role == "admin":
|
||||
if caller_team_role == "admin":
|
||||
return
|
||||
|
||||
permissions: Final = (
|
||||
TeamMemberPermissionChecks._get_list_of_route_enum_as_str(
|
||||
TeamMemberPermissionChecks.get_permissions_for_team_member(
|
||||
team_member_object=team_member_object,
|
||||
team_table=team_table,
|
||||
)
|
||||
)
|
||||
if team_member_object is not None
|
||||
if caller_team_role is not None
|
||||
else []
|
||||
)
|
||||
|
||||
|
|
@ -214,7 +209,7 @@ class TeamMemberPermissionChecks:
|
|||
Returns True if the user belongs to the team that the key is assigned to
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_get_user_in_team,
|
||||
_get_caller_team_role,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
||||
|
|
@ -228,9 +223,8 @@ class TeamMemberPermissionChecks:
|
|||
check_db_only=True,
|
||||
)
|
||||
|
||||
# 4. Extract `Member` object from `team_table`
|
||||
team_member_object: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
|
||||
return team_member_object is not None
|
||||
caller_team_role: Final = _get_caller_team_role(team_table=team_table, user_api_key_dict=user_api_key_dict)
|
||||
return caller_team_role is not None
|
||||
|
||||
@staticmethod
|
||||
def get_all_available_team_member_permissions() -> list[str]:
|
||||
|
|
|
|||
|
|
@ -86,7 +86,9 @@ class _PrismaUserTable(Protocol):
|
|||
class _PrismaTeamMembershipTable(Protocol):
|
||||
"""Team membership table actions the management helpers issue."""
|
||||
|
||||
async def create(self, *, data: Mapping[str, object], include: Mapping[str, bool]) -> _PrismaRecord: ...
|
||||
async def upsert(
|
||||
self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]], include: Mapping[str, bool]
|
||||
) -> _PrismaRecord: ...
|
||||
|
||||
|
||||
class MemberWriteTx(Protocol):
|
||||
|
|
@ -348,7 +350,7 @@ async def _resolve_member_budget_id(
|
|||
default member budget is cloned (with ``budget_duration`` overriding its
|
||||
reset window while keeping its other limits). A lone ``budget_duration``
|
||||
with no team default creates a window-only budget. With nothing set the
|
||||
member gets no budget.
|
||||
member gets no budget, though ``add_new_member`` still writes its membership row.
|
||||
"""
|
||||
has_explicit_limit: Final = max_budget_in_team is not None or allowed_models is not None
|
||||
|
||||
|
|
@ -415,9 +417,9 @@ async def add_new_member(
|
|||
Add a new member to a team
|
||||
|
||||
- add team id to user table
|
||||
- add team member w/ budget to team member table
|
||||
- add team member to team member table, linked to a budget when one resolves
|
||||
|
||||
Returns created/existing user + team membership w/ budget id
|
||||
Returns created/existing user + team membership (``budget_id`` is ``None`` when no budget applies)
|
||||
|
||||
Callers already inside a transaction pass it as ``tx`` so every write here runs on that
|
||||
connection instead of borrowing more from the pool while the caller's locks are held.
|
||||
|
|
@ -471,14 +473,15 @@ async def add_new_member(
|
|||
tx=tx,
|
||||
)
|
||||
|
||||
if _budget_id and returned_user is not None and returned_user.user_id is not None:
|
||||
if returned_user is not None and returned_user.user_id is not None:
|
||||
membership_table: Final[_PrismaTeamMembershipTable] = _team_membership_table(prisma_client, tx)
|
||||
_returned_team_membership: Final = await membership_table.create(
|
||||
data={
|
||||
"team_id": team_id,
|
||||
"user_id": returned_user.user_id,
|
||||
"budget_id": _budget_id,
|
||||
},
|
||||
membership_key: Final[Mapping[str, object]] = {"user_id": returned_user.user_id, "team_id": team_id}
|
||||
budget_link: Final[Mapping[str, str]] = (
|
||||
MappingProxyType({"budget_id": _budget_id}) if _budget_id is not None else MappingProxyType({})
|
||||
)
|
||||
_returned_team_membership: Final = await membership_table.upsert(
|
||||
where={"user_id_team_id": membership_key},
|
||||
data={"create": {**membership_key, **budget_link}, "update": {}},
|
||||
include={"litellm_budget_table": True},
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import hmac
|
|||
import inspect
|
||||
import json
|
||||
import os
|
||||
import posixpath
|
||||
import re
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
|
|
@ -30,12 +31,27 @@ from litellm import get_llm_provider
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS,
|
||||
AZURE_SPEECH_BATCH_PATH_PREFIX,
|
||||
AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN,
|
||||
AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
|
||||
AZURE_SPEECH_FAST_TRANSCRIPTION_PATH,
|
||||
AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX,
|
||||
AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX,
|
||||
AZURE_SPEECH_STT_DOMAIN,
|
||||
AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER,
|
||||
BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES,
|
||||
)
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.llms.azure.passthrough.transformation import foreign_azure_deployment
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.llms.deepgram.common_utils import (
|
||||
deepgram_listen_callback_params,
|
||||
deepgram_listen_is_priced,
|
||||
deepgram_listen_registry_key,
|
||||
deepgram_listen_requested_model,
|
||||
deepgram_listen_websocket_target,
|
||||
)
|
||||
from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
||||
|
|
@ -48,6 +64,7 @@ from litellm.proxy.auth.user_api_key_auth import (
|
|||
is_no_auth_dev_mode,
|
||||
user_api_key_auth,
|
||||
user_api_key_auth_websocket,
|
||||
user_api_key_auth_websocket_for_model,
|
||||
)
|
||||
from litellm.proxy.common_request_processing import open_sse_before_first_byte
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
|
|
@ -58,6 +75,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
|
|||
get_request_body,
|
||||
is_json_content_type,
|
||||
)
|
||||
from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
|
||||
from litellm.proxy.common_utils.sse_keepalive import (
|
||||
wrap_passthrough_sse_bytes_with_keepalive_pings,
|
||||
)
|
||||
|
|
@ -1358,6 +1376,145 @@ async def comprehend_medical_sdk_proxy_route(
|
|||
)
|
||||
|
||||
|
||||
AZURE_SPEECH_FORWARDED_REQUEST_HEADERS: Final = ("content-type", "accept")
|
||||
AZURE_SPEECH_ENDPOINT_FAMILY_DOMAINS: Final = MappingProxyType(
|
||||
{
|
||||
AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX: AZURE_SPEECH_STT_DOMAIN,
|
||||
AZURE_SPEECH_BATCH_PATH_PREFIX: AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def resolve_azure_speech_base_url(endpoint_path: str, api_base: str | None, region: str | None) -> httpx.URL | None:
|
||||
"""
|
||||
Azure AI Speech serves the two REST families from different regional hosts: short-audio
|
||||
recognition under ``{region}.stt.speech.microsoft.com`` and batch transcription under
|
||||
``{region}.api.cognitive.microsoft.com``. An operator-configured ``api_base`` (custom
|
||||
domain or private endpoint) serves both and wins over the region. Returns ``None`` when
|
||||
the path is outside both families so the operator key is never sent for an unknown API.
|
||||
"""
|
||||
domain: Final = next(
|
||||
(
|
||||
family_domain
|
||||
for family_prefix, family_domain in AZURE_SPEECH_ENDPOINT_FAMILY_DOMAINS.items()
|
||||
if endpoint_path.startswith(family_prefix)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if domain is None:
|
||||
return None
|
||||
if api_base:
|
||||
return httpx.URL(api_base)
|
||||
if not region:
|
||||
return None
|
||||
return httpx.URL(f"https://{region}.{domain}")
|
||||
|
||||
|
||||
def azure_speech_path_manages_shared_resources(endpoint_path: str) -> bool:
|
||||
return (
|
||||
endpoint_path.startswith(AZURE_SPEECH_BATCH_PATH_PREFIX)
|
||||
and endpoint_path != AZURE_SPEECH_FAST_TRANSCRIPTION_PATH
|
||||
)
|
||||
|
||||
|
||||
def canonical_azure_speech_endpoint_path(endpoint: str) -> str:
|
||||
"""
|
||||
The path Azure will actually serve, with ``.`` and ``..`` segments resolved, so the
|
||||
endpoint family and the admin guard are decided on the same path the upstream request uses.
|
||||
"""
|
||||
raw_path: Final = httpx.URL(endpoint).path
|
||||
resolved_path: Final = posixpath.normpath(f"/{raw_path.lstrip('/')}")
|
||||
if raw_path.endswith("/") and resolved_path != "/":
|
||||
return f"{resolved_path}/"
|
||||
return resolved_path
|
||||
|
||||
|
||||
@router.api_route(
|
||||
f"{AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX}/{{endpoint:path}}",
|
||||
methods=["GET", "POST", "PUT", "DELETE", "PATCH"], # mutable-ok: fastapi route methods must be a list
|
||||
tags=["Azure AI Speech Pass-through", "pass-through"], # mutable-ok: fastapi route tags must be a list
|
||||
)
|
||||
async def azure_speech_proxy_route(
|
||||
endpoint: str,
|
||||
request: Request,
|
||||
fastapi_response: Response,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
):
|
||||
"""
|
||||
Pass-through for the Azure AI Speech REST APIs (speech to text), e.g.
|
||||
`POST /azure_speech/speech/recognition/conversation/cognitiveservices/v1?language=en-US`
|
||||
with the raw audio as the body, or `POST /azure_speech/speechtotext/v3.2/transcriptions`.
|
||||
|
||||
The body is forwarded byte for byte and the proxy injects its own
|
||||
`Ocp-Apim-Subscription-Key`; the caller's `Authorization` header is the LiteLLM key
|
||||
and is never forwarded.
|
||||
|
||||
[Docs](https://docs.litellm.ai/docs/pass_through/azure_speech)
|
||||
"""
|
||||
normalized_endpoint_path: Final = canonical_azure_speech_endpoint_path(endpoint)
|
||||
base_url: Final = resolve_azure_speech_base_url(
|
||||
endpoint_path=normalized_endpoint_path,
|
||||
api_base=get_secret_str(secret_name="AZURE_SPEECH_API_BASE"),
|
||||
region=get_secret_str(secret_name="AZURE_SPEECH_REGION"),
|
||||
)
|
||||
if base_url is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
f"Unsupported Azure Speech path: {normalized_endpoint_path}. Supported prefixes are "
|
||||
f"{AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX} and {AZURE_SPEECH_BATCH_PATH_PREFIX}; set "
|
||||
"AZURE_SPEECH_REGION or AZURE_SPEECH_API_BASE in the proxy environment."
|
||||
),
|
||||
)
|
||||
if azure_speech_path_manages_shared_resources(normalized_endpoint_path) and not is_proxy_admin(user_api_key_dict):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=(
|
||||
f"{request.method} {normalized_endpoint_path} manages batch transcription resources that belong to "
|
||||
"the proxy's Azure Speech subscription and whose cost is unknown at request time, so it is limited "
|
||||
f"to proxy admin keys. Use {AZURE_SPEECH_FAST_TRANSCRIPTION_PATH} for transcription that is priced "
|
||||
"per request."
|
||||
),
|
||||
)
|
||||
azure_speech_api_key: Final = passthrough_endpoint_router.get_credentials(
|
||||
custom_llm_provider=AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
|
||||
region_name=None,
|
||||
)
|
||||
if azure_speech_api_key is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Azure Speech credentials not found. Set AZURE_SPEECH_API_KEY in the proxy environment.",
|
||||
)
|
||||
|
||||
target_url: Final = base_url.copy_with(
|
||||
path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, normalized_endpoint_path)
|
||||
)
|
||||
request_headers: Final = _safe_get_request_headers(request)
|
||||
upstream_headers: Final = MappingProxyType(
|
||||
{
|
||||
header_name: header_value
|
||||
for header_name, header_value in (
|
||||
*(
|
||||
(header_name, request_headers[header_name])
|
||||
for header_name in AZURE_SPEECH_FORWARDED_REQUEST_HEADERS
|
||||
if header_name in request_headers
|
||||
),
|
||||
(AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER, azure_speech_api_key),
|
||||
)
|
||||
}
|
||||
)
|
||||
raw_body: Final = await request.body()
|
||||
|
||||
endpoint_func: Final = create_pass_through_route(
|
||||
endpoint=endpoint,
|
||||
target=str(target_url),
|
||||
custom_headers=upstream_headers,
|
||||
custom_llm_provider=AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
|
||||
)
|
||||
setattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, raw_body)
|
||||
return await endpoint_func(request, fastapi_response, user_api_key_dict)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/transcribe/{operation}",
|
||||
tags=["Amazon Transcribe Pass-through", "pass-through"], # mutable-ok: fastapi route tags must be a list
|
||||
|
|
@ -2776,7 +2933,7 @@ async def _openai_websocket_refusal(
|
|||
return None
|
||||
|
||||
|
||||
class _OpenAIWebsocketRelay(Protocol):
|
||||
class _WebsocketRelay(Protocol):
|
||||
async def __call__(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -2790,7 +2947,7 @@ class _OpenAIWebsocketRelay(Protocol):
|
|||
) -> None: ...
|
||||
|
||||
|
||||
def _openai_websocket_relay() -> _OpenAIWebsocketRelay:
|
||||
def _websocket_relay() -> _WebsocketRelay:
|
||||
return websocket_passthrough_request
|
||||
|
||||
|
||||
|
|
@ -2808,6 +2965,15 @@ def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists:
|
|||
return resolve
|
||||
|
||||
|
||||
def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None:
|
||||
requested_subprotocols: Final = tuple(
|
||||
protocol.strip()
|
||||
for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",")
|
||||
if protocol.strip()
|
||||
)
|
||||
return requested_subprotocols[0] if requested_subprotocols else None
|
||||
|
||||
|
||||
@router.websocket("/openai_passthrough/{endpoint:path}")
|
||||
@router.websocket("/openai/{endpoint:path}")
|
||||
async def openai_websocket_proxy_route(
|
||||
|
|
@ -2815,16 +2981,11 @@ async def openai_websocket_proxy_route(
|
|||
endpoint: str,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)],
|
||||
general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)],
|
||||
relay: Annotated[_OpenAIWebsocketRelay, Depends(_openai_websocket_relay)],
|
||||
relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)],
|
||||
model_allowlists: Annotated[_OpenAIWebsocketModelAllowlists, Depends(_proxy_model_allowlists)],
|
||||
) -> None:
|
||||
"""WebSocket passthrough for OpenAI prefixes (realtime / responses.connect)."""
|
||||
requested_subprotocols: Final = tuple(
|
||||
protocol.strip()
|
||||
for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",")
|
||||
if protocol.strip()
|
||||
)
|
||||
negotiated_subprotocol: Final = requested_subprotocols[0] if requested_subprotocols else None
|
||||
negotiated_subprotocol: Final = _negotiated_websocket_subprotocol(websocket)
|
||||
|
||||
refusal: Final = await _openai_websocket_refusal(user_api_key_dict, general_settings, model_allowlists)
|
||||
if refusal is not None:
|
||||
|
|
@ -2883,6 +3044,69 @@ async def openai_websocket_proxy_route(
|
|||
)
|
||||
|
||||
|
||||
_DEEPGRAM_WS_MISSING_KEY_REASON: Final = (
|
||||
"Required 'DEEPGRAM_API_KEY' in environment to make pass-through calls to Deepgram."
|
||||
)
|
||||
_DEEPGRAM_WS_CALLBACK_REASON: Final = "Deepgram callback delivery is not supported through the proxy: remove {params}"
|
||||
_DEEPGRAM_WS_UNPRICED_REASON: Final = (
|
||||
"No streaming price for '{registry_key}': add it to the model cost map to enable it"
|
||||
)
|
||||
|
||||
|
||||
async def deepgram_listen_user_api_key_auth(websocket: WebSocket) -> UserAPIKeyAuth:
|
||||
return await user_api_key_auth_websocket_for_model(
|
||||
websocket, model=deepgram_listen_requested_model(websocket.url.query)
|
||||
)
|
||||
|
||||
|
||||
@router.websocket("/deepgram/v1/listen")
|
||||
@router.websocket("/deepgram/listen")
|
||||
async def deepgram_listen_websocket_route(
|
||||
websocket: WebSocket,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(deepgram_listen_user_api_key_auth)],
|
||||
relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)],
|
||||
) -> None:
|
||||
deepgram_api_key: Final = passthrough_endpoint_router.get_credentials(
|
||||
custom_llm_provider=litellm.LlmProviders.DEEPGRAM.value,
|
||||
region_name=None,
|
||||
)
|
||||
if deepgram_api_key is None:
|
||||
await websocket.close(code=1011, reason=_DEEPGRAM_WS_MISSING_KEY_REASON)
|
||||
return
|
||||
|
||||
await websocket.accept(subprotocol=_negotiated_websocket_subprotocol(websocket))
|
||||
callback_params: Final = deepgram_listen_callback_params(websocket.url.query)
|
||||
if callback_params:
|
||||
await websocket.close(
|
||||
code=1008,
|
||||
reason=_DEEPGRAM_WS_CALLBACK_REASON.format(params=", ".join(callback_params)),
|
||||
)
|
||||
return
|
||||
|
||||
target: Final = deepgram_listen_websocket_target(
|
||||
api_base=get_secret_str("DEEPGRAM_API_BASE"),
|
||||
query_string=websocket.url.query,
|
||||
)
|
||||
if not deepgram_listen_is_priced(target):
|
||||
await websocket.close(
|
||||
code=1008,
|
||||
reason=_DEEPGRAM_WS_UNPRICED_REASON.format(registry_key=deepgram_listen_registry_key(target)),
|
||||
)
|
||||
return
|
||||
|
||||
await relay(
|
||||
websocket=websocket,
|
||||
target=target,
|
||||
custom_headers={ # mutable-ok: websocket_passthrough_request requires a plain dict of upstream headers
|
||||
"Authorization": f"Token {deepgram_api_key}"
|
||||
},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
forward_headers=False,
|
||||
endpoint=websocket.url.path,
|
||||
accept_websocket=False,
|
||||
)
|
||||
|
||||
|
||||
class BaseOpenAIPassThroughHandler:
|
||||
@staticmethod
|
||||
async def _base_openai_pass_through_handler(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,172 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
AZURE_SPEECH_BATCH_MODEL,
|
||||
AZURE_SPEECH_BATCH_PATH_PREFIX,
|
||||
AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
|
||||
AZURE_SPEECH_FAST_TRANSCRIPTION_MODEL,
|
||||
AZURE_SPEECH_FAST_TRANSCRIPTION_PATH,
|
||||
AZURE_SPEECH_MILLISECONDS_PER_SECOND,
|
||||
AZURE_SPEECH_PRICING_MODEL,
|
||||
AZURE_SPEECH_SHORT_AUDIO_MODEL,
|
||||
AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX,
|
||||
AZURE_SPEECH_TICKS_PER_SECOND,
|
||||
)
|
||||
from litellm.cost_calculator import transcription_cost
|
||||
from litellm.litellm_core_utils.audio_utils.utils import calculate_request_duration
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
get_standard_logging_object_payload,
|
||||
)
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
|
||||
from litellm.types.utils import StandardPassThroughResponseObject
|
||||
|
||||
|
||||
class AzureSpeechPassthroughLoggingHandler:
|
||||
@staticmethod
|
||||
def _is_short_audio_route(url_route: str) -> bool:
|
||||
path: Final = urlparse(url_route).path
|
||||
return path.rfind(AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX) > path.rfind(AZURE_SPEECH_BATCH_PATH_PREFIX)
|
||||
|
||||
@staticmethod
|
||||
def _is_fast_transcription_route(url_route: str) -> bool:
|
||||
return urlparse(url_route).path.endswith(AZURE_SPEECH_FAST_TRANSCRIPTION_PATH)
|
||||
|
||||
@staticmethod
|
||||
def _model_from_url_route(url_route: str) -> str:
|
||||
if AzureSpeechPassthroughLoggingHandler._is_short_audio_route(url_route):
|
||||
return f"{AZURE_SPEECH_CUSTOM_LLM_PROVIDER}/{AZURE_SPEECH_SHORT_AUDIO_MODEL}"
|
||||
if AzureSpeechPassthroughLoggingHandler._is_fast_transcription_route(url_route):
|
||||
return f"{AZURE_SPEECH_CUSTOM_LLM_PROVIDER}/{AZURE_SPEECH_FAST_TRANSCRIPTION_MODEL}"
|
||||
return f"{AZURE_SPEECH_CUSTOM_LLM_PROVIDER}/{AZURE_SPEECH_BATCH_MODEL}"
|
||||
|
||||
@staticmethod
|
||||
def _recognized_audio_seconds(response_body: Mapping[str, object] | Sequence[object] | None) -> float:
|
||||
if not isinstance(response_body, Mapping):
|
||||
return 0.0
|
||||
offset: Final = response_body.get("Offset")
|
||||
duration: Final = response_body.get("Duration")
|
||||
if not isinstance(offset, int) or not isinstance(duration, int):
|
||||
return 0.0
|
||||
return (offset + duration) / AZURE_SPEECH_TICKS_PER_SECOND
|
||||
|
||||
@staticmethod
|
||||
def _uploaded_audio_seconds(httpx_response: httpx.Response) -> float:
|
||||
try:
|
||||
uploaded_audio: Final = httpx_response.request.content
|
||||
except RuntimeError:
|
||||
return 0.0
|
||||
return calculate_request_duration(uploaded_audio) or 0.0
|
||||
|
||||
@staticmethod
|
||||
def _short_audio_seconds(
|
||||
httpx_response: httpx.Response, response_body: Mapping[str, object] | Sequence[object] | None
|
||||
) -> float:
|
||||
return max(
|
||||
AzureSpeechPassthroughLoggingHandler._uploaded_audio_seconds(httpx_response),
|
||||
AzureSpeechPassthroughLoggingHandler._recognized_audio_seconds(response_body),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _fast_transcription_audio_seconds(response_body: Mapping[str, object] | Sequence[object] | None) -> float:
|
||||
if not isinstance(response_body, Mapping):
|
||||
return 0.0
|
||||
duration_milliseconds: Final = response_body.get("durationMilliseconds")
|
||||
if not isinstance(duration_milliseconds, int):
|
||||
return 0.0
|
||||
return duration_milliseconds / AZURE_SPEECH_MILLISECONDS_PER_SECOND
|
||||
|
||||
@staticmethod
|
||||
def _billed_audio_seconds(
|
||||
url_route: str,
|
||||
httpx_response: httpx.Response,
|
||||
response_body: Mapping[str, object] | Sequence[object] | None,
|
||||
) -> float:
|
||||
if AzureSpeechPassthroughLoggingHandler._is_short_audio_route(url_route):
|
||||
return AzureSpeechPassthroughLoggingHandler._short_audio_seconds(httpx_response, response_body)
|
||||
if AzureSpeechPassthroughLoggingHandler._is_fast_transcription_route(url_route):
|
||||
return AzureSpeechPassthroughLoggingHandler._fast_transcription_audio_seconds(response_body)
|
||||
return 0.0
|
||||
|
||||
@staticmethod
|
||||
def _response_cost(
|
||||
url_route: str,
|
||||
httpx_response: httpx.Response,
|
||||
response_body: Mapping[str, object] | Sequence[object] | None,
|
||||
) -> float:
|
||||
audio_seconds: Final = AzureSpeechPassthroughLoggingHandler._billed_audio_seconds(
|
||||
url_route, httpx_response, response_body
|
||||
)
|
||||
if audio_seconds <= 0.0:
|
||||
return 0.0
|
||||
try:
|
||||
prompt_cost, completion_cost = transcription_cost(
|
||||
model=AZURE_SPEECH_PRICING_MODEL,
|
||||
custom_llm_provider="azure",
|
||||
duration=audio_seconds,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # a missing price entry must not drop the spend log row
|
||||
verbose_proxy_logger.warning(
|
||||
"No price for %s, logging Azure Speech call at zero cost: %s", AZURE_SPEECH_PRICING_MODEL, e
|
||||
)
|
||||
return 0.0
|
||||
return prompt_cost + completion_cost
|
||||
|
||||
@staticmethod
|
||||
def azure_speech_passthrough_handler(
|
||||
httpx_response: httpx.Response,
|
||||
response_body: Mapping[str, object] | Sequence[object] | None,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
url_route: str,
|
||||
result: str,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
cache_hit: bool,
|
||||
request_body: Mapping[str, object],
|
||||
**kwargs: object, # kwargs-ok: the passthrough logging dispatch forwards shared logging kwargs to every handler
|
||||
) -> PassThroughEndpointLoggingTypedDict:
|
||||
try:
|
||||
model_name: Final = AzureSpeechPassthroughLoggingHandler._model_from_url_route(url_route)
|
||||
response_cost: Final = AzureSpeechPassthroughLoggingHandler._response_cost(
|
||||
url_route, httpx_response, response_body
|
||||
)
|
||||
|
||||
updated_kwargs: Final = { # mutable-ok: the logging pipeline requires a plain kwargs dict
|
||||
**kwargs,
|
||||
"model": model_name,
|
||||
"custom_llm_provider": AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
|
||||
"response_cost": response_cost,
|
||||
}
|
||||
logging_obj.model_call_details.update(
|
||||
model=model_name,
|
||||
custom_llm_provider=AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
|
||||
response_cost=response_cost,
|
||||
)
|
||||
|
||||
standard_logging_object: Final = get_standard_logging_object_payload(
|
||||
kwargs=updated_kwargs,
|
||||
init_response_obj=StandardPassThroughResponseObject(response=result),
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=logging_obj,
|
||||
status="success",
|
||||
)
|
||||
|
||||
handler_payload: Final[PassThroughEndpointLoggingTypedDict] = {
|
||||
"result": StandardPassThroughResponseObject(response=result),
|
||||
"kwargs": {**updated_kwargs, "standard_logging_object": standard_logging_object},
|
||||
}
|
||||
except Exception as e: # noqa: BLE001 # logging must never fail the forwarded request
|
||||
verbose_proxy_logger.exception("Error in Azure Speech passthrough logging handler: %s", e)
|
||||
fallback_payload: Final[PassThroughEndpointLoggingTypedDict] = {
|
||||
"result": StandardPassThroughResponseObject(response=result),
|
||||
"kwargs": kwargs,
|
||||
}
|
||||
return fallback_payload
|
||||
return handler_payload
|
||||
|
|
@ -0,0 +1,96 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.deepgram.common_utils import (
|
||||
deepgram_listen_addon_pricing_models,
|
||||
deepgram_listen_audio_seconds,
|
||||
deepgram_listen_channel_count,
|
||||
deepgram_listen_is_priced,
|
||||
deepgram_listen_model,
|
||||
deepgram_listen_pricing_model,
|
||||
deepgram_listen_registry_key,
|
||||
deepgram_listen_transcript,
|
||||
)
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
|
||||
from litellm.types.utils import TranscriptionResponse
|
||||
|
||||
DEEPGRAM_LISTEN_ROUTE_SUFFIX: Final = "/listen"
|
||||
|
||||
|
||||
def _registry_cost(response: TranscriptionResponse, pricing_model: str) -> float | None:
|
||||
try:
|
||||
return litellm.completion_cost(
|
||||
completion_response=response,
|
||||
model=pricing_model,
|
||||
custom_llm_provider=litellm.LlmProviders.DEEPGRAM.value,
|
||||
call_type="transcription",
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # an unpriced entry must not lose the spend row, only its cost
|
||||
verbose_proxy_logger.debug("Deepgram listen passthrough: no registry price for '%s': %s", pricing_model, e)
|
||||
return None
|
||||
|
||||
|
||||
def _audio_cost(response: TranscriptionResponse, upstream_url: str) -> float | None:
|
||||
if not deepgram_listen_is_priced(upstream_url):
|
||||
verbose_proxy_logger.warning(
|
||||
"Deepgram listen passthrough: no registry entry '%s'", deepgram_listen_registry_key(upstream_url)
|
||||
)
|
||||
return None
|
||||
base_cost: Final = _registry_cost(response, deepgram_listen_pricing_model(upstream_url))
|
||||
if base_cost is None:
|
||||
return None
|
||||
addon_costs: Final = tuple(
|
||||
_registry_cost(response, pricing_model) for pricing_model in deepgram_listen_addon_pricing_models(upstream_url)
|
||||
)
|
||||
return base_cost + sum(cost for cost in addon_costs if cost is not None)
|
||||
|
||||
|
||||
class DeepgramListenPassthroughLoggingHandler:
|
||||
@staticmethod
|
||||
def is_deepgram_listen_route(url_route: str) -> bool:
|
||||
path: Final = urlparse(url_route).path
|
||||
return "/deepgram/" in path and path.endswith(DEEPGRAM_LISTEN_ROUTE_SUFFIX)
|
||||
|
||||
def deepgram_listen_passthrough_handler(
|
||||
self,
|
||||
websocket_messages: Sequence[Mapping[str, object]],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
upstream_url: str,
|
||||
kwargs: Mapping[str, object] = MappingProxyType({}),
|
||||
) -> PassThroughEndpointLoggingTypedDict:
|
||||
model: Final = deepgram_listen_model(upstream_url)
|
||||
audio_seconds: Final = deepgram_listen_audio_seconds(websocket_messages)
|
||||
channels: Final = deepgram_listen_channel_count(websocket_messages, upstream_url)
|
||||
billed_seconds: Final = audio_seconds * channels
|
||||
response: Final = TranscriptionResponse(text=deepgram_listen_transcript(websocket_messages))
|
||||
response._hidden_params["audio_transcription_duration"] = billed_seconds # pyright: ignore[reportPrivateUsage] # the cost calculator reads the billed duration off the response's hidden params
|
||||
response_cost: Final = _audio_cost(response, upstream_url)
|
||||
response._hidden_params["response_cost"] = response_cost # pyright: ignore[reportPrivateUsage] # the logger reads a precomputed cost off the response's hidden params
|
||||
|
||||
provider: Final = litellm.LlmProviders.DEEPGRAM.value
|
||||
logging_obj.model = model # rebind-ok: the spend logger reads model and cost off the shared logging object
|
||||
logging_obj.model_call_details["model"] = model # rebind-ok: same shared logging object
|
||||
logging_obj.model_call_details["custom_llm_provider"] = provider # rebind-ok: same shared logging object
|
||||
logging_obj.model_call_details["response_cost"] = response_cost # rebind-ok: same shared logging object
|
||||
verbose_proxy_logger.debug(
|
||||
"Deepgram listen passthrough cost tracking: model %s, audio seconds %s, channels %s, cost %s",
|
||||
model,
|
||||
audio_seconds,
|
||||
channels,
|
||||
response_cost,
|
||||
)
|
||||
logging_result: Final[PassThroughEndpointLoggingTypedDict] = {
|
||||
"result": response,
|
||||
"kwargs": {
|
||||
**kwargs,
|
||||
"model": model,
|
||||
"custom_llm_provider": provider,
|
||||
"response_cost": response_cost,
|
||||
},
|
||||
}
|
||||
return logging_result
|
||||
|
|
@ -8,7 +8,7 @@ from base64 import b64encode
|
|||
from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from itertools import groupby
|
||||
from itertools import count, groupby
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
|
||||
from urllib.parse import urlencode, urlparse
|
||||
|
|
@ -53,6 +53,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import validate_no_callback_env_reference
|
||||
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.llms.base_llm.managed_resources.utils import (
|
||||
|
|
@ -1024,7 +1025,7 @@ async def pass_through_request(
|
|||
verbose_proxy_logger.debug(
|
||||
"Pass through endpoint sending request to \nURL %s\nheaders: %s\nbody: %s\n",
|
||||
url,
|
||||
upstream_headers,
|
||||
_get_masked_values(upstream_headers),
|
||||
_parsed_body,
|
||||
)
|
||||
|
||||
|
|
@ -2121,6 +2122,14 @@ def _resolved_vertex_live_setup(
|
|||
return {**setup_data, "model": setup_model_rewriter(setup_model)}
|
||||
|
||||
|
||||
def _json_object_frame(frame: str | bytes) -> dict[str, object] | None:
|
||||
try:
|
||||
decoded: Final = json.loads(frame if isinstance(frame, str) else frame.decode("utf-8"))
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
return None
|
||||
return decoded if isinstance(decoded, dict) else None
|
||||
|
||||
|
||||
def _truncated_close_reason(reason: str) -> str:
|
||||
"""
|
||||
Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character
|
||||
|
|
@ -2402,70 +2411,41 @@ async def websocket_passthrough_request(
|
|||
)
|
||||
await upstream_ws.close()
|
||||
|
||||
def _extract_vertex_live_model_from_setup_response(setup_response: Mapping[str, object]) -> None:
|
||||
extracted_model: Final = _extract_model_from_vertex_ai_setup(setup_response)
|
||||
if not extracted_model:
|
||||
verbose_proxy_logger.warning(
|
||||
"WebSocket passthrough (%s): Failed to extract model from server setup response: %s",
|
||||
endpoint,
|
||||
setup_response,
|
||||
)
|
||||
return
|
||||
kwargs["model"] = extracted_model
|
||||
kwargs["custom_llm_provider"] = "vertex_ai_language_models"
|
||||
logging_obj.model = extracted_model
|
||||
logging_obj.model_call_details["model"] = extracted_model
|
||||
logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai_language_models"
|
||||
|
||||
is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint)
|
||||
json_frame_ordinal: Final = count()
|
||||
|
||||
async def relay_upstream_frame(upstream_message: str | bytes) -> None:
|
||||
if isinstance(upstream_message, bytes):
|
||||
await websocket.send_bytes(upstream_message)
|
||||
else:
|
||||
await websocket.send_text(upstream_message)
|
||||
message_data: Final = _json_object_frame(upstream_message)
|
||||
if message_data is None:
|
||||
return
|
||||
if is_vertex_live and next(json_frame_ordinal) == 0:
|
||||
_extract_vertex_live_model_from_setup_response(message_data)
|
||||
return
|
||||
websocket_messages.append(message_data)
|
||||
|
||||
async def forward_upstream_to_client() -> Close | None:
|
||||
"""Forward messages from upstream to client WebSocket, returning the upstream's close frame"""
|
||||
try:
|
||||
# Wait for the first response from upstream
|
||||
raw_response = await upstream_ws.recv(decode=False)
|
||||
# Ensure raw_response is bytes before decoding
|
||||
if isinstance(raw_response, str):
|
||||
raw_response = raw_response.encode("utf-8")
|
||||
setup_response: Final[Mapping[str, object]] = json.loads(raw_response.decode("utf-8"))
|
||||
verbose_proxy_logger.debug("Setup response: %s", setup_response)
|
||||
|
||||
# Extract model and provider from setup response for Vertex AI Live
|
||||
if endpoint and "/vertex_ai/live" in endpoint:
|
||||
verbose_proxy_logger.debug(
|
||||
"WebSocket passthrough (%s): Processing server setup response for model extraction",
|
||||
endpoint,
|
||||
)
|
||||
extracted_model: Final = _extract_model_from_vertex_ai_setup(setup_response)
|
||||
if extracted_model:
|
||||
kwargs["model"] = extracted_model
|
||||
kwargs["custom_llm_provider"] = "vertex_ai_language_models"
|
||||
# Update logging object with correct model
|
||||
logging_obj.model = extracted_model
|
||||
logging_obj.model_call_details["model"] = extracted_model
|
||||
logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai_language_models"
|
||||
verbose_proxy_logger.debug(
|
||||
"WebSocket passthrough (%s): Successfully extracted model '%s' and set provider to 'vertex_ai' from server setup response",
|
||||
endpoint,
|
||||
extracted_model,
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.warning(
|
||||
"WebSocket passthrough (%s): Failed to extract model from server setup response: %s",
|
||||
endpoint,
|
||||
setup_response,
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.debug(
|
||||
"WebSocket passthrough (%s): Not a Vertex AI Live endpoint, skipping model extraction",
|
||||
endpoint,
|
||||
)
|
||||
|
||||
# Send the setup response to the client
|
||||
await websocket.send_text(json.dumps(setup_response))
|
||||
|
||||
# Now continuously forward messages from upstream to client
|
||||
async for upstream_message in upstream_ws:
|
||||
if isinstance(upstream_message, bytes):
|
||||
await websocket.send_bytes(upstream_message)
|
||||
# Parse and collect for cost tracking
|
||||
try:
|
||||
message_data: dict[str, object] = json.loads(upstream_message.decode())
|
||||
websocket_messages.append(message_data)
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
pass
|
||||
else:
|
||||
await websocket.send_text(upstream_message)
|
||||
# Parse and collect for cost tracking
|
||||
try:
|
||||
message_data = json.loads(upstream_message)
|
||||
websocket_messages.append(message_data)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
while True:
|
||||
await relay_upstream_frame(await upstream_ws.recv())
|
||||
except (ConnectionClosedOK, ConnectionClosedError) as e:
|
||||
verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e)
|
||||
return e.rcvd
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from urllib.parse import urlparse
|
|||
|
||||
import httpx
|
||||
|
||||
from litellm.constants import AZURE_SPEECH_CUSTOM_LLM_PROVIDER
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingResultValues
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
|
|
@ -25,6 +26,9 @@ from .llm_provider_handlers.cohere_passthrough_logging_handler import (
|
|||
from .llm_provider_handlers.cursor_passthrough_logging_handler import (
|
||||
CursorPassthroughLoggingHandler,
|
||||
)
|
||||
from .llm_provider_handlers.deepgram_listen_passthrough_logging_handler import (
|
||||
DeepgramListenPassthroughLoggingHandler,
|
||||
)
|
||||
from .llm_provider_handlers.gemini_passthrough_logging_handler import (
|
||||
GeminiPassthroughLoggingHandler,
|
||||
)
|
||||
|
|
@ -274,6 +278,25 @@ class PassThroughEndpointLogging:
|
|||
)
|
||||
standard_logging_response_object = comprehend_medical_handler_result["result"] # rebind-ok: elif-chain
|
||||
kwargs = comprehend_medical_handler_result["kwargs"] # rebind-ok: elif-chain contract
|
||||
elif self.is_azure_speech_route(custom_llm_provider):
|
||||
from .llm_provider_handlers.azure_speech_passthrough_logging_handler import (
|
||||
AzureSpeechPassthroughLoggingHandler,
|
||||
)
|
||||
|
||||
azure_speech_handler_result: Final = AzureSpeechPassthroughLoggingHandler.azure_speech_passthrough_handler(
|
||||
httpx_response=httpx_response,
|
||||
response_body=response_body,
|
||||
logging_obj=logging_obj,
|
||||
url_route=url_route,
|
||||
result=result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
cache_hit=cache_hit,
|
||||
request_body=request_body,
|
||||
**kwargs,
|
||||
)
|
||||
standard_logging_response_object = azure_speech_handler_result["result"] # rebind-ok: elif-chain
|
||||
kwargs = azure_speech_handler_result["kwargs"] # rebind-ok: elif-chain contract
|
||||
elif self.is_transcribe_route(custom_llm_provider):
|
||||
transcribe_handler_result: Final = TranscribePassthroughLoggingHandler.transcribe_passthrough_handler(
|
||||
httpx_response=httpx_response,
|
||||
|
|
@ -329,6 +352,21 @@ class PassThroughEndpointLogging:
|
|||
|
||||
standard_logging_response_object = vertex_ai_live_handler_result["result"]
|
||||
kwargs = vertex_ai_live_handler_result["kwargs"]
|
||||
elif DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route):
|
||||
deepgram_handler_result: Final = (
|
||||
DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=tuple(
|
||||
message
|
||||
for message in (response_body if isinstance(response_body, list) else ())
|
||||
if isinstance(message, dict)
|
||||
),
|
||||
logging_obj=logging_obj,
|
||||
upstream_url=str(httpx_response.request.url),
|
||||
kwargs=kwargs,
|
||||
)
|
||||
)
|
||||
standard_logging_response_object = deepgram_handler_result["result"] # rebind-ok: elif-chain
|
||||
kwargs = deepgram_handler_result["kwargs"] # rebind-ok: elif-chain contract
|
||||
return_dict["standard_logging_response_object"] = standard_logging_response_object
|
||||
|
||||
return_dict["kwargs"] = kwargs
|
||||
|
|
@ -351,7 +389,7 @@ class PassThroughEndpointLogging:
|
|||
):
|
||||
standard_logging_response_object: PassThroughEndpointLoggingResultValues | None = None
|
||||
logging_obj.model_call_details["passthrough_logging_payload"] = passthrough_logging_payload
|
||||
if self.is_assemblyai_route(url_route):
|
||||
if self.is_assemblyai_route(url_route) and not self.is_azure_speech_route(custom_llm_provider):
|
||||
if AssemblyAIPassthroughLoggingHandler._should_log_request(httpx_response.request.method) is not True:
|
||||
return
|
||||
self.assemblyai_passthrough_logging_handler.assemblyai_passthrough_logging_handler(
|
||||
|
|
@ -458,6 +496,9 @@ class PassThroughEndpointLogging:
|
|||
def is_comprehend_medical_route(self, custom_llm_provider: str | None) -> bool:
|
||||
return custom_llm_provider == "comprehendmedical"
|
||||
|
||||
def is_azure_speech_route(self, custom_llm_provider: str | None) -> bool:
|
||||
return custom_llm_provider == AZURE_SPEECH_CUSTOM_LLM_PROVIDER
|
||||
|
||||
def is_transcribe_route(self, custom_llm_provider: str | None) -> bool:
|
||||
return custom_llm_provider == TRANSCRIBE_CUSTOM_LLM_PROVIDER
|
||||
|
||||
|
|
|
|||
|
|
@ -1411,6 +1411,8 @@ def run_server(
|
|||
# DO NOT DELETE - enables global variables to work across files
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
os.environ["NUM_WORKERS"] = str(num_workers)
|
||||
|
||||
# Auto-create PROMETHEUS_MULTIPROC_DIR for multi-worker setups
|
||||
prometheus_multiproc_dir: Final = ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
|
||||
num_workers=num_workers,
|
||||
|
|
|
|||
|
|
@ -155,7 +155,7 @@ from litellm.types.utils import (
|
|||
TextCompletionResponse,
|
||||
TokenCountResponse,
|
||||
)
|
||||
from litellm.utils import load_credentials_from_list
|
||||
from litellm.utils import cost_map_omits_token_price, load_credentials_from_list
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from aiohttp import ClientSession
|
||||
|
|
@ -263,6 +263,7 @@ from litellm.constants import (
|
|||
APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
APSCHEDULER_REPLACE_EXISTING,
|
||||
CLI_SSO_SESSION_TTL_SECONDS,
|
||||
DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID,
|
||||
DAYS_IN_A_MONTH,
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL,
|
||||
DEFAULT_MODEL_CREATED_AT_TIME,
|
||||
|
|
@ -310,6 +311,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
|
|||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.proxy._experimental.mcp_server.byok_credential_cache import byok_credential_cache
|
||||
from litellm.proxy._lazy_features import attach_lazy_features, reserve_lazy_slot
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.analytics_endpoints.analytics_endpoints import (
|
||||
|
|
@ -331,6 +333,12 @@ from litellm.proxy.auth.fallback_budget import router_fallback_budget_check
|
|||
from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck
|
||||
from litellm.proxy.auth.login_throttle import (
|
||||
LoginThrottle,
|
||||
declared_proxy_ranges,
|
||||
warn_login_counters_are_per_worker,
|
||||
warn_source_login_limit_is_off,
|
||||
)
|
||||
from litellm.proxy.auth.model_checks import (
|
||||
expand_wildcard_deployments_for_model_info,
|
||||
get_all_fallbacks,
|
||||
|
|
@ -682,6 +690,9 @@ from litellm.proxy.route_priority import hot_routes_first
|
|||
from litellm.proxy.search_endpoints.endpoints import router as search_router
|
||||
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
|
||||
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
|
||||
from litellm.proxy.spend_tracking.daily_global_spend_rollup import (
|
||||
run_scheduled_daily_global_spend_reconcile,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import (
|
||||
PendingSpendIncrement,
|
||||
active_spend_counter_batch,
|
||||
|
|
@ -828,6 +839,7 @@ from fastapi.openapi.docs import get_swagger_ui_html
|
|||
from fastapi.openapi.utils import get_openapi
|
||||
from fastapi.responses import (
|
||||
FileResponse,
|
||||
HTMLResponse,
|
||||
JSONResponse,
|
||||
ORJSONResponse,
|
||||
RedirectResponse,
|
||||
|
|
@ -6054,6 +6066,12 @@ class ProxyConfig:
|
|||
general_settings = config.get("general_settings", {})
|
||||
if general_settings is None:
|
||||
general_settings = {}
|
||||
|
||||
if os.getenv("NUM_WORKERS", "1") != "1" and redis_usage_cache is None:
|
||||
warn_login_counters_are_per_worker(os.getenv("NUM_WORKERS", "1"))
|
||||
if declared_proxy_ranges(general_settings) is None:
|
||||
warn_source_login_limit_is_off()
|
||||
|
||||
_bg_hc_model_groups: Final = parse_background_health_check_model_groups(general_settings)
|
||||
_enable_hc_routing = False
|
||||
_hc_staleness = None
|
||||
|
|
@ -7541,7 +7559,7 @@ class ProxyConfig:
|
|||
subscriber: Final = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=redis_cache,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
additional_in_memory_caches=(spend_counter_cache.in_memory_cache,),
|
||||
additional_in_memory_caches=(spend_counter_cache.in_memory_cache, byok_credential_cache),
|
||||
)
|
||||
self.auth_cache_invalidation_subscriber = subscriber
|
||||
subscriber.start()
|
||||
|
|
@ -10007,6 +10025,12 @@ class ProxyStartupEvent:
|
|||
|
||||
await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler)
|
||||
|
||||
cls._initialize_daily_global_spend_reconcile_job(
|
||||
scheduler=scheduler,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
### PTU DAILY ROLLUP ###
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import (
|
||||
is_ptu_cost_attribution_enabled,
|
||||
|
|
@ -10348,6 +10372,39 @@ class ProxyStartupEvent:
|
|||
"LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true to enable)"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _initialize_daily_global_spend_reconcile_job(
|
||||
cls,
|
||||
scheduler: AsyncIOScheduler,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
async def alert(message: str) -> None:
|
||||
await proxy_logging_obj.alerting_handler(
|
||||
message=message,
|
||||
level="High",
|
||||
alert_type=AlertType.failed_tracking_spend,
|
||||
)
|
||||
|
||||
async def reconcile() -> None:
|
||||
await run_scheduled_daily_global_spend_reconcile(
|
||||
prisma_client,
|
||||
pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager,
|
||||
alert=alert,
|
||||
)
|
||||
|
||||
scheduler.add_job(
|
||||
reconcile,
|
||||
"cron",
|
||||
hour=0,
|
||||
minute=30,
|
||||
timezone="UTC",
|
||||
id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID,
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
next_run_time=datetime.now(timezone.utc) + timedelta(minutes=2),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def _initialize_slack_alerting_jobs(
|
||||
cls,
|
||||
|
|
@ -13622,9 +13679,10 @@ def _enrich_model_info_with_litellm_data(
|
|||
discovered_model_info: Final = (
|
||||
llm_router.get_discovered_model_info(model_info.get("id")) if llm_router is not None else MappingProxyType({})
|
||||
)
|
||||
unpriced: Final = cost_map_omits_token_price(model_info.get("id"), litellm_model_info.get("key"))
|
||||
for k, v in MappingProxyType({**litellm_model_info, **discovered_model_info}).items():
|
||||
if k not in model_info or (model_info[k] is None and k in discovered_model_info):
|
||||
model_info[k] = v
|
||||
model_info[k] = None if unpriced and k in ("input_cost_per_token", "output_cost_per_token") else v
|
||||
model["model_info"] = model_info
|
||||
# don't return the api key / vertex credentials
|
||||
# don't return the llm credentials
|
||||
|
|
@ -15060,45 +15118,7 @@ def _translate_model_name_for_response(model: dict) -> dict:
|
|||
|
||||
|
||||
def _get_proxy_model_info(model: dict) -> dict:
|
||||
# provided model_info in config.yaml
|
||||
model_info: Final = model.get("model_info", {})
|
||||
|
||||
# read litellm model_prices_and_context_window.json to get the following:
|
||||
# input_cost_per_token, output_cost_per_token, max_tokens
|
||||
litellm_model_info = get_litellm_model_info(model=model)
|
||||
|
||||
# 2nd pass on the model, try seeing if we can find model in litellm model_cost map
|
||||
if litellm_model_info == {}:
|
||||
# use litellm_param model_name to get model_info
|
||||
litellm_params = model.get("litellm_params", {})
|
||||
litellm_model = litellm_params.get("model", None)
|
||||
try:
|
||||
litellm_model_info = litellm.get_model_info(model=litellm_model)
|
||||
except Exception:
|
||||
litellm_model_info = {}
|
||||
# 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map
|
||||
if litellm_model_info == {}:
|
||||
# use litellm_param model_name to get model_info
|
||||
litellm_params = model.get("litellm_params", {})
|
||||
litellm_model = litellm_params.get("model", None)
|
||||
split_model: Final = litellm_model.split("/")
|
||||
if len(split_model) > 0:
|
||||
litellm_model = split_model[-1]
|
||||
try:
|
||||
litellm_model_info = litellm.get_model_info(model=litellm_model, custom_llm_provider=split_model[0])
|
||||
except Exception:
|
||||
litellm_model_info = {}
|
||||
discovered_model_info: Final = (
|
||||
llm_router.get_discovered_model_info(model_info.get("id")) if llm_router is not None else MappingProxyType({})
|
||||
)
|
||||
for k, v in MappingProxyType({**litellm_model_info, **discovered_model_info}).items():
|
||||
if k not in model_info or (model_info[k] is None and k in discovered_model_info):
|
||||
model_info[k] = v
|
||||
model["model_info"] = model_info
|
||||
# don't return the llm credentials
|
||||
model = remove_sensitive_info_from_deployment(deployment_dict=model, excluded_keys={"litellm_credential_name"})
|
||||
|
||||
return _translate_model_name_for_response(model)
|
||||
return _translate_model_name_for_response(_enrich_model_info_with_litellm_data(model=model, llm_router=llm_router))
|
||||
|
||||
|
||||
def _model_info_json_response(data: Sequence[Mapping[str, object]] | Mapping[str, object]) -> Response:
|
||||
|
|
@ -15905,8 +15925,6 @@ async def fallback_login(request: Request):
|
|||
else:
|
||||
redirect_url += "/sso/callback"
|
||||
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings)
|
||||
return HTMLResponse(
|
||||
content=build_ui_login_form(
|
||||
|
|
@ -15928,13 +15946,27 @@ async def login(request: Request):
|
|||
password: Final = str(form.get("password"))
|
||||
|
||||
# Authenticate user and get login result
|
||||
login_result: Final = await authenticate_user(
|
||||
username=username,
|
||||
password=password,
|
||||
master_key=master_key,
|
||||
prisma_client=prisma_client,
|
||||
general_settings=general_settings,
|
||||
)
|
||||
try:
|
||||
login_result: Final = await authenticate_user(
|
||||
username=username,
|
||||
password=password,
|
||||
master_key=master_key,
|
||||
prisma_client=prisma_client,
|
||||
throttle=LoginThrottle.from_request(request, general_settings, redis_usage_cache),
|
||||
general_settings=general_settings,
|
||||
)
|
||||
except ProxyException as exc:
|
||||
if int(exc.code) != status.HTTP_429_TOO_MANY_REQUESTS:
|
||||
raise
|
||||
retry_after: Final = exc.headers.get("Retry-After", "30")
|
||||
return HTMLResponse(
|
||||
content=(
|
||||
"<html><body><h1>Too many sign-in attempts</h1>"
|
||||
f"<p>Try again in about {retry_after} seconds</p></body></html>"
|
||||
),
|
||||
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||||
headers=exc.headers,
|
||||
)
|
||||
|
||||
# Create UI token object
|
||||
returned_ui_token_object: Final = create_ui_token_object(
|
||||
|
|
@ -16013,6 +16045,7 @@ async def login_v2(request: Request):
|
|||
password=password,
|
||||
master_key=master_key,
|
||||
prisma_client=prisma_client,
|
||||
throttle=LoginThrottle.from_request(request, general_settings, redis_usage_cache),
|
||||
general_settings=general_settings,
|
||||
)
|
||||
|
||||
|
|
@ -16084,6 +16117,7 @@ async def login_v3(request: Request):
|
|||
password=password,
|
||||
master_key=master_key,
|
||||
prisma_client=prisma_client,
|
||||
throttle=LoginThrottle.from_request(request, general_settings, redis_usage_cache),
|
||||
general_settings=general_settings,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -586,6 +586,24 @@
|
|||
],
|
||||
"default_model_placeholder": "azure_ai/command-r-plus"
|
||||
},
|
||||
{
|
||||
"provider": "Azure_Speech",
|
||||
"provider_display_name": "Azure AI Speech",
|
||||
"litellm_provider": "azure_speech",
|
||||
"credential_fields": [
|
||||
{
|
||||
"key": "api_key",
|
||||
"label": "Azure AI Speech Subscription Key",
|
||||
"placeholder": null,
|
||||
"tooltip": "The Ocp-Apim-Subscription-Key for your Azure AI Speech resource. The proxy injects it on every /azure_speech/* pass-through request. Region and API base come from AZURE_SPEECH_REGION / AZURE_SPEECH_API_BASE",
|
||||
"required": true,
|
||||
"field_type": "password",
|
||||
"options": null,
|
||||
"default_value": null
|
||||
}
|
||||
],
|
||||
"default_model_placeholder": "azure_speech/short-audio"
|
||||
},
|
||||
{
|
||||
"provider": "AZURE_TEXT",
|
||||
"provider_display_name": "Azure Text",
|
||||
|
|
|
|||
|
|
@ -820,6 +820,37 @@ model LiteLLM_DailyUserSpend {
|
|||
@@index([endpoint])
|
||||
}
|
||||
|
||||
// Key-free daily rollup of LiteLLM_DailyUserSpend, read by the global usage view
|
||||
model LiteLLM_DailyGlobalSpend {
|
||||
id String @id @default(uuid())
|
||||
date String
|
||||
model String?
|
||||
model_group String?
|
||||
custom_llm_provider String?
|
||||
mcp_namespaced_tool_name String?
|
||||
endpoint String?
|
||||
prompt_tokens BigInt @default(0)
|
||||
completion_tokens BigInt @default(0)
|
||||
cache_read_input_tokens BigInt @default(0)
|
||||
cache_creation_input_tokens BigInt @default(0)
|
||||
compression_saved_tokens BigInt @default(0)
|
||||
compression_savings_spend Float @default(0.0)
|
||||
prompt_caching_savings_spend Float @default(0.0)
|
||||
gateway_injected_caching_savings_spend Float @default(0.0)
|
||||
autorouter_savings_spend Float @default(0.0)
|
||||
spend Float @default(0.0)
|
||||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([date, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
}
|
||||
|
||||
// Track daily organization spend metrics per model and key
|
||||
model LiteLLM_DailyOrganizationSpend {
|
||||
id String @id @default(uuid())
|
||||
|
|
|
|||
293
litellm/proxy/spend_tracking/daily_global_spend_rollup.py
Normal file
293
litellm/proxy/spend_tracking/daily_global_spend_rollup.py
Normal file
|
|
@ -0,0 +1,293 @@
|
|||
"""Roll closed UTC days of ``LiteLLM_DailyUserSpend`` up into ``LiteLLM_DailyGlobalSpend``.
|
||||
|
||||
Only days that are over get rolled up, so a pod still flushing per-key spend for the current
|
||||
day can never leave the global table short; usage reads serve days through the recorded
|
||||
marker from the global table and later days live from the per-key table. Per-key rows are
|
||||
dated by request start, so spend can land on a day that was already rolled up (a flush
|
||||
straddling midnight, a retry after an outage). Each run therefore also rewrites every closed
|
||||
day that has rows touched since the previous run's scan, whatever the date. The marker lives
|
||||
in ``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on
|
||||
a large deployment the first backfill is minutes of work.
|
||||
"""
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import date, timedelta
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID,
|
||||
DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS,
|
||||
DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
GLOBAL_SPEND_TABLE_NAME: Final = "LiteLLM_DailyGlobalSpend"
|
||||
# The unique constraint, in constraint order. NULL never matches itself in a unique index, so
|
||||
# every column is normalized to '' or the same group would be inserted again on every run.
|
||||
_KEY_COLUMNS: Final = ("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint")
|
||||
_METRIC_COLUMNS: Final = (
|
||||
"prompt_tokens",
|
||||
"completion_tokens",
|
||||
"cache_read_input_tokens",
|
||||
"cache_creation_input_tokens",
|
||||
"compression_saved_tokens",
|
||||
"api_requests",
|
||||
"successful_requests",
|
||||
"failed_requests",
|
||||
"total_response_time_ms",
|
||||
"timed_requests",
|
||||
"compression_savings_spend",
|
||||
"prompt_caching_savings_spend",
|
||||
"gateway_injected_caching_savings_spend",
|
||||
"autorouter_savings_spend",
|
||||
"spend",
|
||||
)
|
||||
|
||||
|
||||
def _quoted(columns: tuple[str, ...]) -> str:
|
||||
return ", ".join(f'"{column}"' for column in columns)
|
||||
|
||||
|
||||
def _reconcile_day_sql() -> str:
|
||||
normalized_keys: Final = ", ".join(f"COALESCE(\"{column}\", '')" for column in _KEY_COLUMNS)
|
||||
sums: Final = ", ".join(f'SUM("{column}")' for column in _METRIC_COLUMNS)
|
||||
overwrite: Final = ", ".join(f'"{column}" = EXCLUDED."{column}"' for column in _METRIC_COLUMNS)
|
||||
return (
|
||||
f'INSERT INTO "{GLOBAL_SPEND_TABLE_NAME}" ("id", {_quoted(_KEY_COLUMNS)}, {_quoted(_METRIC_COLUMNS)}, '
|
||||
'"updated_at")\n'
|
||||
f"SELECT gen_random_uuid()::text, {normalized_keys}, {sums}, (NOW() AT TIME ZONE 'UTC')\n"
|
||||
'FROM "LiteLLM_DailyUserSpend" WHERE "date" = $1\n'
|
||||
f"GROUP BY {normalized_keys}\n"
|
||||
f"ON CONFLICT ({_quoted(_KEY_COLUMNS)}) DO UPDATE SET {overwrite}, "
|
||||
"\"updated_at\" = (NOW() AT TIME ZONE 'UTC')"
|
||||
)
|
||||
|
||||
|
||||
RECONCILE_DAY_SQL: Final = _reconcile_day_sql()
|
||||
_DB_NOW_SQL: Final = "SELECT (NOW() AT TIME ZONE 'UTC')::text AS now, (NOW() AT TIME ZONE 'UTC')::date::text AS today"
|
||||
_ALL_CLOSED_DAYS_SQL: Final = 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ORDER BY "date"'
|
||||
# Pod clocks drift from the database clock and from each other, so rows are picked up from a
|
||||
# little before the previous scan; rewriting a day twice is idempotent.
|
||||
_PENDING_DAYS_SQL: Final = (
|
||||
'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 '
|
||||
'AND ("date" > $2 OR "updated_at" >= $3::timestamp - INTERVAL \'1 hour\') '
|
||||
'ORDER BY "date"'
|
||||
)
|
||||
# Runs can overlap (Redis unreachable, lock expired on a long backfill), so the database keeps the
|
||||
# later of the stored and the incoming day and scan time in one statement; GREATEST skips NULL.
|
||||
_ADVANCE_MARKER_SQL: Final = (
|
||||
'INSERT INTO "LiteLLM_Config" ("param_name", "param_value") '
|
||||
"VALUES ($1, jsonb_build_object('reconciled_through', $2::text, 'scanned_at', $3::text)) "
|
||||
'ON CONFLICT ("param_name") DO UPDATE SET "param_value" = jsonb_build_object('
|
||||
"'reconciled_through', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'reconciled_through', "
|
||||
"EXCLUDED.\"param_value\" ->> 'reconciled_through'), "
|
||||
"'scanned_at', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'scanned_at', "
|
||||
"EXCLUDED.\"param_value\" ->> 'scanned_at'))"
|
||||
)
|
||||
|
||||
|
||||
class ReconciledThrough(BaseModel):
|
||||
"""``reconciled_through`` is the last closed UTC day the global table covers. ``scanned_at`` is
|
||||
the database clock when the scan behind the last fully successful run started: every per-key
|
||||
row written before it, on any day through the marker, is in the global table."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
reconciled_through: str
|
||||
scanned_at: str | None = None
|
||||
|
||||
|
||||
class _MarkerRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True)
|
||||
|
||||
param_value: object = None
|
||||
|
||||
|
||||
class _DateRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
date: str
|
||||
|
||||
|
||||
class _NowRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
now: str
|
||||
today: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ReconcileResult:
|
||||
days_reconciled: tuple[str, ...]
|
||||
reconciled_through: str | None
|
||||
failed_day: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PendingScan:
|
||||
marker: ReconciledThrough | None
|
||||
scanned_at: str
|
||||
days: tuple[str, ...]
|
||||
|
||||
|
||||
def _marker_from_param_value(value: object) -> ReconciledThrough | None:
|
||||
try:
|
||||
return (
|
||||
ReconciledThrough.model_validate_json(value)
|
||||
if isinstance(value, str)
|
||||
else ReconciledThrough.model_validate(value)
|
||||
)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
async def read_marker(prisma_client: "PrismaClient") -> ReconciledThrough | None:
|
||||
from litellm.proxy.utils import get_config_param
|
||||
|
||||
row: Final = await get_config_param(prisma_client, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
|
||||
return None if row is None else _marker_from_param_value(_MarkerRow.model_validate(row).param_value)
|
||||
|
||||
|
||||
async def reconciled_through(prisma_client: "PrismaClient") -> str | None:
|
||||
"""The last UTC day ``LiteLLM_DailyGlobalSpend`` is known to cover, or None before the first run."""
|
||||
marker: Final = await read_marker(prisma_client)
|
||||
return None if marker is None else marker.reconciled_through
|
||||
|
||||
|
||||
async def _advance_marker(prisma_client: "PrismaClient", days: tuple[str, ...], *, scanned_at: str | None) -> None:
|
||||
"""Move the stored marker to the last of ``days`` and to ``scanned_at`` where those are later
|
||||
than what is stored, so a slower overlapping run can only add to a faster run's marker."""
|
||||
from litellm.proxy.utils import invalidate_config_param
|
||||
|
||||
await prisma_client.db.execute_raw(
|
||||
_ADVANCE_MARKER_SQL,
|
||||
DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM,
|
||||
max(days) if days else None,
|
||||
scanned_at,
|
||||
)
|
||||
await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
|
||||
|
||||
|
||||
async def _db_now(prisma_client: "PrismaClient") -> _NowRow:
|
||||
rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL)
|
||||
return _NowRow.model_validate(rows[0])
|
||||
|
||||
|
||||
async def _scan_pending(prisma_client: "PrismaClient") -> _PendingScan:
|
||||
"""Every closed UTC day (strictly before the database's today) still to roll up, oldest first:
|
||||
days past the marker, plus any day with per-key rows written since the scan behind the marker.
|
||||
Before a run has fully succeeded there is no such scan, so every closed day is rolled up."""
|
||||
marker: Final = await read_marker(prisma_client)
|
||||
db_now: Final = await _db_now(prisma_client)
|
||||
last_closed_day: Final = (date.fromisoformat(db_now.today) - timedelta(days=1)).isoformat()
|
||||
rows: Final = (
|
||||
await prisma_client.db.query_raw(_ALL_CLOSED_DAYS_SQL, last_closed_day)
|
||||
if marker is None or marker.scanned_at is None
|
||||
else await prisma_client.db.query_raw(
|
||||
_PENDING_DAYS_SQL, last_closed_day, marker.reconciled_through, marker.scanned_at
|
||||
)
|
||||
)
|
||||
return _PendingScan(marker, db_now.now, tuple(_DateRow.model_validate(row).date for row in rows))
|
||||
|
||||
|
||||
async def pending_days(prisma_client: "PrismaClient") -> tuple[str, ...]:
|
||||
return (await _scan_pending(prisma_client)).days
|
||||
|
||||
|
||||
async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None:
|
||||
"""Rewrite one day of the global table from the per-key sums. Idempotent: a rerun
|
||||
overwrites every group with the same totals."""
|
||||
await prisma_client.db.execute_raw(RECONCILE_DAY_SQL, day)
|
||||
|
||||
|
||||
async def run_daily_global_spend_reconcile(prisma_client: "PrismaClient") -> ReconcileResult:
|
||||
"""Roll up every pending day, advancing the marker after each; a failing day stops the run
|
||||
with the marker on the last good day so the next run resumes there. The scan time is only
|
||||
recorded once every pending day is done, so late rows a failed run saw are found again."""
|
||||
scan: Final = await _scan_pending(prisma_client)
|
||||
done: Final = await _reconcile_until_failure(prisma_client, scan)
|
||||
if len(done) < len(scan.days):
|
||||
marker: Final = await reconciled_through(prisma_client)
|
||||
return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=scan.days[len(done)])
|
||||
if scan.marker is not None or done:
|
||||
await _advance_marker(prisma_client, done, scanned_at=scan.scanned_at)
|
||||
return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client))
|
||||
|
||||
|
||||
async def _reconcile_until_failure(prisma_client: "PrismaClient", scan: _PendingScan) -> tuple[str, ...]:
|
||||
for index, day in enumerate(scan.days):
|
||||
if not await _reconcile_and_record(prisma_client, scan.days[: index + 1]):
|
||||
return scan.days[:index]
|
||||
return scan.days
|
||||
|
||||
|
||||
async def _reconcile_and_record(prisma_client: "PrismaClient", done_with_this: tuple[str, ...]) -> bool:
|
||||
day: Final = done_with_this[-1]
|
||||
try:
|
||||
await reconcile_day(prisma_client, day)
|
||||
await _advance_marker(prisma_client, done_with_this, scanned_at=None)
|
||||
except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done
|
||||
verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
async def run_scheduled_daily_global_spend_reconcile(
|
||||
prisma_client: "PrismaClient",
|
||||
pod_lock_manager: "PodLockManager | None" = None,
|
||||
alert: Callable[[str], Awaitable[None]] | None = None,
|
||||
) -> ReconcileResult | None:
|
||||
"""Run the reconcile under a cross-pod lock so one proxy does the work; the lock only saves
|
||||
effort (each day is an idempotent rewrite), so an unreachable Redis runs unguarded rather than skipping."""
|
||||
redis_cache: Final = None if pod_lock_manager is None else pod_lock_manager.redis_cache
|
||||
if pod_lock_manager is None or redis_cache is None:
|
||||
return await _run_and_alert(prisma_client, alert=alert)
|
||||
|
||||
acquired: Final = await pod_lock_manager.acquire_lock(
|
||||
cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, ttl=DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS
|
||||
)
|
||||
if not acquired and await _lock_is_held(pod_lock_manager, redis_cache):
|
||||
verbose_proxy_logger.info("Daily global spend reconcile: another pod holds the lock, skipping this run")
|
||||
return None
|
||||
try:
|
||||
return await _run_and_alert(prisma_client, alert=alert)
|
||||
finally:
|
||||
if acquired:
|
||||
await pod_lock_manager.release_lock(cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID)
|
||||
|
||||
|
||||
async def _lock_is_held(pod_lock_manager: "PodLockManager", redis_cache: "RedisCache") -> bool:
|
||||
try:
|
||||
lock_key: Final = pod_lock_manager.get_redis_lock_key(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID)
|
||||
return bool(await redis_cache.async_get_cache(lock_key))
|
||||
except Exception as exc: # noqa: BLE001 # an unreadable lock must not skip the run
|
||||
verbose_proxy_logger.warning("Daily global spend reconcile: could not read the lock: %s", exc)
|
||||
return False
|
||||
|
||||
|
||||
async def _run_and_alert(
|
||||
prisma_client: "PrismaClient",
|
||||
*,
|
||||
alert: Callable[[str], Awaitable[None]] | None,
|
||||
) -> ReconcileResult:
|
||||
result: Final = await run_daily_global_spend_reconcile(prisma_client)
|
||||
if result.days_reconciled:
|
||||
verbose_proxy_logger.info(
|
||||
"Daily global spend reconcile: rolled up %d day(s), reconciled through %s",
|
||||
len(result.days_reconciled),
|
||||
result.reconciled_through,
|
||||
)
|
||||
if result.failed_day is not None and alert is not None:
|
||||
await alert(
|
||||
f"Daily global spend reconcile stopped at {result.failed_day}; usage totals keep reading the per-key "
|
||||
f"table for ranges past {result.reconciled_through or 'the beginning'} until the next run succeeds."
|
||||
)
|
||||
return result
|
||||
|
|
@ -195,6 +195,7 @@ from litellm.repositories.user_repository import UserRepository
|
|||
from litellm.repositories.verification_token_repository import (
|
||||
VerificationTokenRepository,
|
||||
)
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
|
@ -7497,6 +7498,7 @@ def _check_and_merge_model_level_guardrails(
|
|||
data: dict,
|
||||
llm_router: Router | None,
|
||||
trust_client_model_info: bool = True,
|
||||
model_alias: str | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Check if the model has guardrails defined and merge them with existing guardrails in the request data.
|
||||
|
|
@ -7504,6 +7506,7 @@ def _check_and_merge_model_level_guardrails(
|
|||
Args:
|
||||
data: The request data dict
|
||||
llm_router: The LLM router instance to get deployment info from
|
||||
model_alias: Resolve guardrails for this model group instead of data["model"]
|
||||
trust_client_model_info: If False, ignore metadata.model_info.id and
|
||||
resolve guardrails by alias-union only. Set to False on the
|
||||
pre_call path because add_litellm_data_to_request preserves
|
||||
|
|
@ -7548,13 +7551,13 @@ def _check_and_merge_model_level_guardrails(
|
|||
# set on ANY eligible deployment still fires (#29652; addresses
|
||||
# veria-ai HIGH on the single-deployment fallback that would skip
|
||||
# non-first deployments).
|
||||
model_alias: Final = data.get("model")
|
||||
if not isinstance(model_alias, str) or not model_alias:
|
||||
alias: Final = model_alias if model_alias is not None else data.get("model")
|
||||
if not isinstance(alias, str) or not alias:
|
||||
return data
|
||||
# Pass team_id so team-scoped public model names resolve the same way
|
||||
# route_request resolves them; otherwise team-scoped deployments are
|
||||
# invisible to this lookup and their guardrails are silently dropped.
|
||||
deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or []
|
||||
deployments: Final = llm_router.get_model_list(model_name=alias, team_id=team_id) or []
|
||||
seen: Final[set] = set()
|
||||
union: Final[list] = []
|
||||
for dep in deployments:
|
||||
|
|
@ -8204,18 +8207,23 @@ def create_model_info_response(
|
|||
"owned_by": provider,
|
||||
}
|
||||
|
||||
listing_info: Final = llm_router.get_model_listing_info(model_id) if llm_router is not None else None
|
||||
alias_target: Final = (
|
||||
resolve_model_group_alias(llm_router.model_group_alias, model_id) if llm_router is not None else None
|
||||
)
|
||||
lookup_model: Final = alias_target if alias_target is not None else model_id
|
||||
|
||||
listing_info: Final = llm_router.get_model_listing_info(lookup_model) if llm_router is not None else None
|
||||
|
||||
# One entry per distinct model behind the listed name; (None,) when the router knows
|
||||
# nothing about it, so the listed name is resolved on its own as before.
|
||||
deployment_models: Final[tuple[str | None, ...]] = (
|
||||
listing_info.cost_map_keys if listing_info is not None and listing_info.cost_map_keys else (None,)
|
||||
)
|
||||
listed_info: Final = _safe_get_model_info(model_id, get_model_info)
|
||||
listed_info: Final = _safe_get_model_info(lookup_model, get_model_info)
|
||||
candidate_sets: Final = tuple(
|
||||
_resolve_listing_model_info(
|
||||
deployment_model=deployment_model,
|
||||
listed_model=model_id,
|
||||
listed_model=lookup_model,
|
||||
listed_info=listed_info,
|
||||
get_model_info=get_model_info,
|
||||
)
|
||||
|
|
@ -8246,7 +8254,7 @@ def create_model_info_response(
|
|||
max_output_tokens = listing_info.max_output_tokens
|
||||
|
||||
if llm_router is not None:
|
||||
configured_mode: Final = llm_router.get_configured_mode(model_id)
|
||||
configured_mode: Final = llm_router.get_configured_mode(lookup_model)
|
||||
if isinstance(configured_mode, str):
|
||||
base["mode"] = configured_mode
|
||||
|
||||
|
|
|
|||
|
|
@ -2589,8 +2589,7 @@ class ManagedResponsesWebSocketHandler:
|
|||
if "litellm_metadata" not in call_kwargs:
|
||||
call_kwargs["litellm_metadata"] = {}
|
||||
call_kwargs["litellm_metadata"]["proxy_server_request"] = proxy_server_request
|
||||
call_kwargs.setdefault("litellm_params", {})
|
||||
call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request
|
||||
call_kwargs["proxy_server_request"] = proxy_server_request
|
||||
|
||||
async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> _MutableJsonObject | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -3940,12 +3940,24 @@ class Router:
|
|||
)
|
||||
|
||||
_router_timeout: Final = (
|
||||
float(self._explicit_timeout) if isinstance(self._explicit_timeout, (int, float)) else None
|
||||
self.request_timeout
|
||||
if self.request_timeout is not None
|
||||
else float(self._explicit_timeout)
|
||||
if isinstance(self._explicit_timeout, (int, float))
|
||||
else None
|
||||
)
|
||||
_router_stream_timeout: Final = (
|
||||
self.stream_timeout
|
||||
if self.stream_timeout is not None
|
||||
else self.request_timeout
|
||||
if self.request_timeout is not None
|
||||
else self.default_litellm_params.get("stream_timeout")
|
||||
)
|
||||
kwargs["timeout"] = resolve_llm_passthrough_timeout(
|
||||
kwargs=kwargs,
|
||||
litellm_params=deployment["litellm_params"],
|
||||
router_timeout=_router_timeout,
|
||||
router_stream_timeout=_router_stream_timeout,
|
||||
)
|
||||
else:
|
||||
kwargs["timeout"] = self._get_timeout(kwargs=kwargs, data=deployment["litellm_params"])
|
||||
|
|
@ -10575,11 +10587,12 @@ class Router:
|
|||
2. If not, check if litellm model name is in model info
|
||||
3. If not, return None
|
||||
"""
|
||||
from litellm.utils import _update_dictionary
|
||||
from litellm.utils import _update_dictionary, cost_map_omits_token_price
|
||||
|
||||
model_info: ModelInfo | None = None
|
||||
custom_model_info: dict | None = None
|
||||
litellm_model_name_model_info: ModelInfo | None = None
|
||||
base_model_key: str | None = None
|
||||
|
||||
try:
|
||||
custom_model_info = (
|
||||
|
|
@ -10606,6 +10619,7 @@ class Router:
|
|||
## update litellm model info with base model info
|
||||
base_model_info: Final = copy.deepcopy(litellm.get_model_info(model=base_model))
|
||||
if base_model_info is not None:
|
||||
base_model_key = base_model_info.get("key")
|
||||
# Base model provides defaults, custom model info overrides
|
||||
custom_model_info = _update_dictionary(
|
||||
cast(dict, base_model_info),
|
||||
|
|
@ -10633,6 +10647,15 @@ class Router:
|
|||
# custom_model_info already includes base_model defaults at this point, if applicable
|
||||
model_info = cast(ModelInfo, custom_model_info)
|
||||
|
||||
if model_info is None:
|
||||
return None
|
||||
builtin_key: Final = (
|
||||
litellm_model_name_model_info.get("key") if litellm_model_name_model_info is not None else None
|
||||
)
|
||||
if cost_map_omits_token_price(model_id, builtin_key, base_model_key):
|
||||
return cast( # cast-ok: TypedDict spread with overridden keys loses its type
|
||||
ModelInfo, {**model_info, "input_cost_per_token": None, "output_cost_per_token": None}
|
||||
)
|
||||
return model_info
|
||||
|
||||
def _set_model_group_info(self, model_group: str, user_facing_model_group_name: str) -> ModelGroupInfo | None:
|
||||
|
|
|
|||
|
|
@ -103,6 +103,25 @@ def declared_reasoning_efforts_for_model(model: str, custom_llm_provider: str) -
|
|||
return declared_reasoning_efforts(entry)
|
||||
|
||||
|
||||
REASONING_EFFORT_STRENGTH_ORDER: Final = ("minimal", "low", "medium", "high", "xhigh", "max")
|
||||
_STRENGTH_RANK: Final = MappingProxyType({effort: rank for rank, effort in enumerate(REASONING_EFFORT_STRENGTH_ORDER)})
|
||||
|
||||
|
||||
def nearest_declared_reasoning_effort(requested: str, declared: Sequence[str]) -> str:
|
||||
"""Rounds a request up to the weakest declared level at least as strong as it, and down to the
|
||||
strongest declared level when it asks for more than the model has, so the caller gets no less
|
||||
reasoning than it asked for instead of a rejected call. none is the off switch rather than a
|
||||
strength, so it is never rounded onto the ladder and no level is rounded down to it: a caller
|
||||
who turned reasoning off must not be billed for it, and a model that cannot turn it off says so
|
||||
itself. A level outside the strength order is likewise returned as is for upstream to judge."""
|
||||
ranked: Final = sorted(
|
||||
(effort for effort in declared if effort in _STRENGTH_RANK), key=lambda effort: _STRENGTH_RANK[effort]
|
||||
)
|
||||
if requested in ranked or requested not in _STRENGTH_RANK or not ranked:
|
||||
return requested
|
||||
return next((effort for effort in ranked if _STRENGTH_RANK[effort] >= _STRENGTH_RANK[requested]), ranked[-1])
|
||||
|
||||
|
||||
def _supports_none_reasoning_effort(model_info: Mapping[str, object], flag: object) -> bool:
|
||||
"""Opt-in only where a request path refuses the level. AzureOpenAIGPT5Config raises
|
||||
UnsupportedParamsError on reasoning_effort='none' without an explicit true, and it is selected
|
||||
|
|
|
|||
|
|
@ -27,13 +27,17 @@ class UpstreamFailure(Exception):
|
|||
self.__cause__ = cause
|
||||
|
||||
|
||||
def _upstream_failure(error: Exception) -> Exception:
|
||||
def _upstream_failure(error: Exception, request: LiteLLMOcrRequest) -> Exception:
|
||||
try:
|
||||
status, body = _UPSTREAM_ARGS.validate_python(error.args)
|
||||
headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None))
|
||||
except ValidationError:
|
||||
return error
|
||||
return UpstreamFailure(httpx.Response(status, content=body.encode(), headers=headers), error)
|
||||
http_request: Final = httpx.Request("POST", request.api_base or "https://docs.litellm.ai/docs")
|
||||
return UpstreamFailure(
|
||||
httpx.Response(status, content=body.encode(), headers=headers, request=http_request),
|
||||
error,
|
||||
)
|
||||
|
||||
|
||||
def response(value: Mapping[str, object]) -> OCRResponse:
|
||||
|
|
@ -57,7 +61,7 @@ def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider:
|
|||
model=request.model.removeprefix(f"{request_provider}/"),
|
||||
llm_provider=request_provider,
|
||||
)
|
||||
original: Final = _upstream_failure(error)
|
||||
original: Final = _upstream_failure(error, request)
|
||||
public_error: Final = failures.map_failure(original, request.model, request_provider, arguments(request))
|
||||
if isinstance(original, UpstreamFailure) and public_error.__context__ is original:
|
||||
public_error.__context__ = error
|
||||
|
|
|
|||
|
|
@ -62,6 +62,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.singulr import (
|
|||
from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
|
||||
ToolPermissionGuardrailConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
|
||||
TypeSafeGuardrailConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import (
|
||||
VigilGuardGuardrailConfigModel,
|
||||
)
|
||||
|
|
@ -138,6 +141,7 @@ class SupportedGuardrailIntegrations(Enum):
|
|||
SINGULR = "singulr"
|
||||
HEADROOM = "headroom"
|
||||
COMPRESR = "compresr"
|
||||
TYPESAFE = "typesafe"
|
||||
STRAIKER = "straiker"
|
||||
ALICE = "alice"
|
||||
AGENT_365 = "agent_365"
|
||||
|
|
@ -1055,7 +1059,7 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
|
|||
default="fail_closed",
|
||||
description=(
|
||||
"Behavior when a guardrail endpoint is unreachable due to network errors. "
|
||||
"Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. "
|
||||
"Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. "
|
||||
"'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed."
|
||||
),
|
||||
)
|
||||
|
|
@ -1171,6 +1175,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o
|
|||
LakeraV2GuardrailConfigModel,
|
||||
HeadroomGuardrailConfigModel,
|
||||
CompresrGuardrailConfigModel,
|
||||
TypeSafeGuardrailConfigModel,
|
||||
RepelloAIGuardrailConfigModel,
|
||||
LassoGuardrailConfigModel,
|
||||
DeepKeepGuardrailConfigModel,
|
||||
|
|
|
|||
|
|
@ -1160,6 +1160,10 @@ OpenAIImageGenerationOptionalParams = Literal[
|
|||
"image_url",
|
||||
"image_prompt_strength",
|
||||
"aspect_ratio",
|
||||
"width",
|
||||
"height",
|
||||
"guidance",
|
||||
"steps",
|
||||
"imageConfig",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -464,3 +464,11 @@ class MCPGatewaySessionsResponse(BaseModel):
|
|||
by_client: list[MCPGatewaySessionGroupCount] = Field(default_factory=list)
|
||||
by_user: list[MCPGatewaySessionGroupCount] = Field(default_factory=list)
|
||||
sessions: list[MCPGatewaySession] = Field(default_factory=list)
|
||||
|
||||
|
||||
class MCPGatewaySessionsTerminateResponse(BaseModel):
|
||||
"""Stateful sessions an administrator force-closed on this proxy worker."""
|
||||
|
||||
worker_pid: int
|
||||
terminated_sessions: int
|
||||
sessions: list[MCPGatewaySession] = Field(default_factory=list)
|
||||
|
|
|
|||
63
litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py
Normal file
63
litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
||||
class TypeSafeGuardrailOptionalParams(BaseModel):
|
||||
"""Optional tuning knobs for the TypeSafe (Jev) compaction guardrail."""
|
||||
|
||||
relevance_threshold: float | None = Field(
|
||||
default=None,
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
description=(
|
||||
"Relevance cutoff in [0, 1]. A completed tool exchange is dropped when Jev "
|
||||
"scores the probability that it is still needed below this value. Defaults to 0.2."
|
||||
),
|
||||
)
|
||||
min_chars_to_evaluate: int | None = Field(
|
||||
default=None,
|
||||
ge=0,
|
||||
description=(
|
||||
"Skip tool exchanges whose combined tool-result text is shorter than this many characters. Defaults to 200."
|
||||
),
|
||||
)
|
||||
max_result_chars_in_state: int | None = Field(
|
||||
default=None,
|
||||
ge=1,
|
||||
description=(
|
||||
"Tool result text is truncated to this many characters when sent to the Jev evaluator, "
|
||||
"keeping the head and tail. Defaults to 4000."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TypeSafeGuardrailConfigModel(GuardrailConfigModel[TypeSafeGuardrailOptionalParams]):
|
||||
api_key: str | None = Field(
|
||||
default=None,
|
||||
description="TypeSafe API key, sent as a Bearer token. Falls back to the TYPESAFE_API_KEY env var.",
|
||||
)
|
||||
api_base: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Base URL of the TypeSafe API. Falls back to the TYPESAFE_API_BASE env var, then https://api.typesafe.ai."
|
||||
),
|
||||
)
|
||||
model: str | None = Field(
|
||||
default=None,
|
||||
description="TypeSafe evaluation model (not the LLM). Defaults to 'jev-latest'.",
|
||||
)
|
||||
unreachable_fallback: Literal["fail_closed", "fail_open"] = Field(
|
||||
default="fail_open",
|
||||
description=(
|
||||
"Behavior when the TypeSafe evaluation service is unreachable or errors. "
|
||||
"'fail_open' (default) forwards the request uncompacted. 'fail_closed' "
|
||||
"raises an error instead."
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "TypeSafe (Jev) Compaction"
|
||||
|
|
@ -100,6 +100,16 @@ class DailySpendMetadata(BaseModel):
|
|||
page: int = Field(default=1)
|
||||
total_pages: int = Field(default=1)
|
||||
has_more: bool = Field(default=False)
|
||||
api_key_limit: int | None = Field(
|
||||
default=None,
|
||||
description="When set, api_keys and every api_key_breakdown list at most this many keys, "
|
||||
"ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key.",
|
||||
)
|
||||
total_api_keys: int | None = Field(
|
||||
default=None,
|
||||
description="Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key "
|
||||
"lists are truncated to the highest-spend keys.",
|
||||
)
|
||||
|
||||
|
||||
class SpendAnalyticsPaginatedResponse(BaseModel):
|
||||
|
|
|
|||
|
|
@ -61,7 +61,9 @@ class SCIMUserGroup(BaseModel):
|
|||
|
||||
|
||||
class SCIMMultiValuedAttribute(BaseModel):
|
||||
value: str
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
value: str | None = None
|
||||
display: str | None = None
|
||||
type: str | None = None
|
||||
primary: bool | None = None
|
||||
|
|
|
|||
|
|
@ -302,6 +302,7 @@ class CredentialLiteLLMParams(BaseModel):
|
|||
aws_bedrock_runtime_endpoint: str | None = None
|
||||
aws_bedrock_project_id: str | None = None
|
||||
s3_bucket_name: str | None = None
|
||||
s3_endpoint_url: str | None = None
|
||||
s3_region_name: str | None = None
|
||||
s3_encryption_key_id: str | None = None
|
||||
aws_batch_role_arn: str | None = None
|
||||
|
|
|
|||
|
|
@ -3776,6 +3776,7 @@ bedrock_batch_litellm_params: Final = (
|
|||
"aws_batch_role_arn",
|
||||
"s3_bucket_name",
|
||||
"s3_region_name",
|
||||
"s3_endpoint_url",
|
||||
"s3_output_bucket_name",
|
||||
"bedrock_tags",
|
||||
)
|
||||
|
|
@ -4074,6 +4075,7 @@ class LlmProviders(str, Enum):
|
|||
TOPAZ = "topaz"
|
||||
SAP_GENERATIVE_AI_HUB = "sap"
|
||||
ASSEMBLYAI = "assemblyai"
|
||||
AZURE_SPEECH = "azure_speech"
|
||||
CHARITY_ENGINE = "charity_engine"
|
||||
GITHUB_COPILOT = "github_copilot"
|
||||
SNOWFLAKE = "snowflake"
|
||||
|
|
|
|||
|
|
@ -3156,6 +3156,22 @@ def reapply_runtime_model_cost_registrations() -> None:
|
|||
register_model(model_cost=dict(_runtime_registered_model_cost)) # mutable-ok: snapshot, replay rewrites it
|
||||
|
||||
|
||||
def cost_map_omits_token_price(*keys: object) -> bool:
|
||||
"""Whether the raw ``litellm.model_cost`` entries under ``keys`` exist but none carries a per-token price.
|
||||
|
||||
``get_model_info`` substitutes 0 for a missing price, which reads exactly like a declared
|
||||
zero. Surfaces that report pricing use this to keep an unpriced deployment at ``None``.
|
||||
"""
|
||||
entries: Final = tuple(
|
||||
entry
|
||||
for entry in (litellm.model_cost.get(key) for key in keys if isinstance(key, str))
|
||||
if isinstance(entry, dict)
|
||||
)
|
||||
return len(entries) > 0 and not any(
|
||||
"input_cost_per_token" in entry or "output_cost_per_token" in entry for entry in entries
|
||||
)
|
||||
|
||||
|
||||
def register_model(
|
||||
model_cost: str | dict,
|
||||
*,
|
||||
|
|
@ -4510,7 +4526,7 @@ def get_optional_params(
|
|||
drop_params=bool(drop_params),
|
||||
)
|
||||
else:
|
||||
optional_params = litellm.MistralConfig().map_openai_params(
|
||||
optional_params = litellm.VertexAIMistralConfig().map_openai_params(
|
||||
model=model,
|
||||
non_default_params=non_default_params,
|
||||
optional_params=optional_params,
|
||||
|
|
@ -8380,7 +8396,7 @@ class ProviderConfigManager:
|
|||
elif model in litellm.vertex_mistral_models:
|
||||
if "codestral" in model:
|
||||
return litellm.CodestralTextCompletionConfig()
|
||||
return litellm.MistralConfig()
|
||||
return litellm.VertexAIMistralConfig()
|
||||
elif model in litellm.vertex_ai_ai21_models:
|
||||
return litellm.VertexAIAi21Config()
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -10820,6 +10820,25 @@
|
|||
"/v1/images/generations"
|
||||
]
|
||||
},
|
||||
"azure_ai/FLUX.2-flex": {
|
||||
"input_cost_per_pixel": 5e-08,
|
||||
"litellm_provider": "azure_ai",
|
||||
"max_input_tokens": 32000,
|
||||
"max_tokens": 32000,
|
||||
"mode": "image_generation",
|
||||
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"image"
|
||||
]
|
||||
},
|
||||
"azure_ai/FW-DeepSeek-V3.2": {
|
||||
"deprecation_date": "2027-07-01",
|
||||
"cache_read_input_token_cost": 3.1e-07,
|
||||
|
|
@ -16690,6 +16709,46 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"dashscope/qwen3.8-flash": {
|
||||
"cache_creation_input_token_cost": 2e-07,
|
||||
"cache_read_input_token_cost": 1.6e-08,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "dashscope",
|
||||
"max_input_tokens": 991808,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.7e-07,
|
||||
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"dashscope/qwen3.8-omni-flash": {
|
||||
"cache_read_input_token_cost": 1.6e-08,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "dashscope",
|
||||
"max_input_tokens": 991808,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.7e-07,
|
||||
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
|
||||
"supports_audio_input": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"dashscope/qwq-plus": {
|
||||
"input_cost_per_token": 8e-07,
|
||||
"litellm_provider": "dashscope",
|
||||
|
|
@ -18594,6 +18653,46 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"qwen_ai_platform/qwen3.8-flash": {
|
||||
"cache_creation_input_token_cost": 2e-07,
|
||||
"cache_read_input_token_cost": 1.6e-08,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "qwen_ai_platform",
|
||||
"max_input_tokens": 991808,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.7e-07,
|
||||
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"qwen_ai_platform/qwen3.8-omni-flash": {
|
||||
"cache_read_input_token_cost": 1.6e-08,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "qwen_ai_platform",
|
||||
"max_input_tokens": 991808,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.7e-07,
|
||||
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
|
||||
"supports_audio_input": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"qwen_ai_platform/qwq-plus": {
|
||||
"input_cost_per_token": 8e-07,
|
||||
"litellm_provider": "qwen_ai_platform",
|
||||
|
|
@ -20754,6 +20853,96 @@
|
|||
"/v1/audio/transcriptions"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/nova-3": {
|
||||
"input_cost_per_second": 8e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0048/60 seconds = $0.00008000 per second",
|
||||
"note": "Nova-3 monolingual streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0048
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/nova-3-multilingual": {
|
||||
"input_cost_per_second": 9.667e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0058/60 seconds = $0.00009667 per second",
|
||||
"note": "Nova-3 multilingual (language=multi) streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0058
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/redact": {
|
||||
"input_cost_per_second": 3.333e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0020/60 seconds = $0.00003333 per second",
|
||||
"note": "Redaction add-on (redact query param), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.002
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/keyterm": {
|
||||
"input_cost_per_second": 2.167e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0013/60 seconds = $0.00002167 per second",
|
||||
"note": "Keyterm Prompting add-on (keyterm query param), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0013
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/detect_entities": {
|
||||
"input_cost_per_second": 2.833e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0017/60 seconds = $0.00002833 per second",
|
||||
"note": "Entity Detection add-on (detect_entities query param), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0017
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/diarize": {
|
||||
"input_cost_per_second": 3.333e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0020/60 seconds = $0.00003333 per second",
|
||||
"note": "Speaker Diarization add-on (diarize / diarize_model query params), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.002
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/whisper": {
|
||||
"input_cost_per_second": 0.0001,
|
||||
"litellm_provider": "deepgram",
|
||||
|
|
@ -22090,8 +22279,8 @@
|
|||
"embed-english-light-v3.0": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "cohere",
|
||||
"max_input_tokens": 1024,
|
||||
"max_tokens": 1024,
|
||||
"max_input_tokens": 512,
|
||||
"max_tokens": 512,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0
|
||||
},
|
||||
|
|
@ -22108,8 +22297,8 @@
|
|||
"input_cost_per_image": 0.0001,
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "cohere",
|
||||
"max_input_tokens": 1024,
|
||||
"max_tokens": 1024,
|
||||
"max_input_tokens": 512,
|
||||
"max_tokens": 512,
|
||||
"metadata": {
|
||||
"notes": "'supports_image_input' is a deprecated field. Use 'supports_embedding_image_input' instead."
|
||||
},
|
||||
|
|
@ -22130,8 +22319,8 @@
|
|||
"embed-multilingual-v3.0": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "cohere",
|
||||
"max_input_tokens": 1024,
|
||||
"max_tokens": 1024,
|
||||
"max_input_tokens": 512,
|
||||
"max_tokens": 512,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_embedding_image_input": true
|
||||
|
|
@ -22139,8 +22328,8 @@
|
|||
"embed-multilingual-light-v3.0": {
|
||||
"input_cost_per_token": 0.0001,
|
||||
"litellm_provider": "cohere",
|
||||
"max_input_tokens": 1024,
|
||||
"max_tokens": 1024,
|
||||
"max_input_tokens": 512,
|
||||
"max_tokens": 512,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"supports_embedding_image_input": true
|
||||
|
|
@ -36998,6 +37187,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37079,6 +37272,15 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"minimal",
|
||||
"low",
|
||||
"medium",
|
||||
"high",
|
||||
"xhigh",
|
||||
"max"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/zai-glm-5-2",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37096,6 +37298,11 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/zai-glm-5-3",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37113,6 +37320,11 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/zai-glm-5-3",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37130,6 +37342,11 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/zai-glm-5-3",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37147,6 +37364,15 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"minimal",
|
||||
"low",
|
||||
"medium",
|
||||
"high",
|
||||
"xhigh",
|
||||
"max"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/zai-glm-5-2",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37442,6 +37668,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37502,6 +37732,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37519,6 +37753,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37552,6 +37790,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37583,6 +37825,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-07,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -57910,14 +58156,14 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"bedrock_mantle/openai.gpt-5.6-sol": {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 1.1e-05,
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 1.375e-05,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
|
||||
"output_cost_per_token": 3.3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.95e-05,
|
||||
"input_cost_per_token": 4.4e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 8.8e-06,
|
||||
"cache_creation_input_token_cost": 5.5e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 1.1e-05,
|
||||
"cache_read_input_token_cost": 4.4e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 8.8e-07,
|
||||
"output_cost_per_token": 2.2e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 3.3e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.012,
|
||||
"search_context_size_low": 0.012,
|
||||
|
|
@ -59928,6 +60174,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-07,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -63253,6 +63503,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -63270,6 +63524,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -63287,6 +63545,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -63304,6 +63566,10 @@
|
|||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-07,
|
||||
"reasoning_effort_levels": [
|
||||
"none",
|
||||
"high"
|
||||
],
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -65816,9 +66082,9 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/z-ai/glm-5.3": {
|
||||
"input_cost_per_token": 1.4e-06,
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"cache_read_input_token_cost": 2.6e-07,
|
||||
"input_cost_per_token": 9.1e-07,
|
||||
"output_cost_per_token": 2.86e-06,
|
||||
"cache_read_input_token_cost": 1.69e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1310720,
|
||||
"max_output_tokens": 943717,
|
||||
|
|
@ -70539,14 +70805,14 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/~deepseek/deepseek-flash-latest": {
|
||||
"cache_read_input_token_cost": 1.5e-08,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"cache_read_input_token_cost": 4.2e-09,
|
||||
"input_cost_per_token": 1.4e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 393216,
|
||||
"max_tokens": 393216,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-07,
|
||||
"output_cost_per_token": 4.2e-07,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_audio_input": false,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -70831,14 +71097,14 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/~z-ai/glm-latest": {
|
||||
"cache_read_input_token_cost": 1.5e-07,
|
||||
"cache_read_input_token_cost": 1.46625e-07,
|
||||
"input_cost_per_token": 9e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1310720,
|
||||
"max_output_tokens": 235929,
|
||||
"max_tokens": 235929,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-06,
|
||||
"output_cost_per_token": 2.805e-06,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_audio_input": false,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -73982,14 +74248,14 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/tencent/hy3": {
|
||||
"cache_read_input_token_cost": 3.3e-08,
|
||||
"input_cost_per_token": 1.32e-07,
|
||||
"cache_read_input_token_cost": 2.0625e-08,
|
||||
"input_cost_per_token": 8.25e-08,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5.28e-07,
|
||||
"output_cost_per_token": 3.3e-07,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_audio_input": false,
|
||||
"supports_function_calling": true,
|
||||
|
|
|
|||
|
|
@ -820,6 +820,37 @@ model LiteLLM_DailyUserSpend {
|
|||
@@index([endpoint])
|
||||
}
|
||||
|
||||
// Key-free daily rollup of LiteLLM_DailyUserSpend, read by the global usage view
|
||||
model LiteLLM_DailyGlobalSpend {
|
||||
id String @id @default(uuid())
|
||||
date String
|
||||
model String?
|
||||
model_group String?
|
||||
custom_llm_provider String?
|
||||
mcp_namespaced_tool_name String?
|
||||
endpoint String?
|
||||
prompt_tokens BigInt @default(0)
|
||||
completion_tokens BigInt @default(0)
|
||||
cache_read_input_tokens BigInt @default(0)
|
||||
cache_creation_input_tokens BigInt @default(0)
|
||||
compression_saved_tokens BigInt @default(0)
|
||||
compression_savings_spend Float @default(0.0)
|
||||
prompt_caching_savings_spend Float @default(0.0)
|
||||
gateway_injected_caching_savings_spend Float @default(0.0)
|
||||
autorouter_savings_spend Float @default(0.0)
|
||||
spend Float @default(0.0)
|
||||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([date, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
}
|
||||
|
||||
// Track daily organization spend metrics per model and key
|
||||
model LiteLLM_DailyOrganizationSpend {
|
||||
id String @id @default(uuid())
|
||||
|
|
|
|||
|
|
@ -86,7 +86,7 @@ locals {
|
|||
"/queue/chat/*",
|
||||
"/v1beta/*",
|
||||
"/interactions/*",
|
||||
"/anthropic/*", "/azure/*", "/azure_ai/*", "/aws/*", "/bedrock/*", "/comprehendmedical*", "/transcribe*",
|
||||
"/anthropic/*", "/azure/*", "/azure_ai/*", "/azure_speech/*", "/aws/*", "/bedrock/*", "/comprehendmedical*", "/transcribe*",
|
||||
"/cohere/*", "/gemini/*", "/google/*",
|
||||
"/vertex_ai/*", "/vertex-ai/*",
|
||||
"/assemblyai/*", "/eu.assemblyai/*",
|
||||
|
|
|
|||
|
|
@ -55,7 +55,7 @@ locals {
|
|||
"/queue/chat/*",
|
||||
"/v1beta/*",
|
||||
"/interactions/*",
|
||||
"/anthropic/*", "/azure/*", "/azure_ai/*", "/aws/*", "/bedrock/*", "/comprehendmedical*", "/transcribe*",
|
||||
"/anthropic/*", "/azure/*", "/azure_ai/*", "/azure_speech/*", "/aws/*", "/bedrock/*", "/comprehendmedical*", "/transcribe*",
|
||||
"/cohere/*", "/gemini/*", "/google/*",
|
||||
"/vertex_ai/*", "/vertex-ai/*",
|
||||
"/assemblyai/*", "/eu.assemblyai/*",
|
||||
|
|
|
|||
|
|
@ -51,9 +51,7 @@ try:
|
|||
if general_settings_section:
|
||||
# Extract the table rows, which contain the documented keys
|
||||
table_content = general_settings_section.group(1)
|
||||
doc_key_pattern = re.compile(
|
||||
r"\|\s*([^\|]+?)\s*\|"
|
||||
) # Capture the key from each row of the table
|
||||
doc_key_pattern = re.compile(r"^\|\s*([^\|]+?)\s*\|", re.MULTILINE)
|
||||
documented_keys.update(doc_key_pattern.findall(table_content))
|
||||
except Exception as e:
|
||||
raise Exception(
|
||||
|
|
|
|||
|
|
@ -331,9 +331,7 @@ class TestMCPPerUserTokenCache:
|
|||
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache):
|
||||
await cache.delete("alice", "slack-test")
|
||||
|
||||
mock_dual_cache.async_delete_cache.assert_called_once_with(
|
||||
"mcp:per_user_token:alice:slack-test"
|
||||
)
|
||||
mock_dual_cache.async_delete_cache.assert_called_once_with(key="mcp:per_user_token:alice:slack-test")
|
||||
mock_dual_cache.async_set_cache.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -1,203 +0,0 @@
|
|||
"""
|
||||
Base test class for OCR functionality across different providers.
|
||||
|
||||
This follows the same pattern as BaseLLMChatTest in tests/llm_translation/base_llm_unit_tests.py
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import litellm
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
# Test resources
|
||||
TEST_IMAGE_PATH = "test_image_edit.png"
|
||||
# Tiny in-repo PDF served via jsdelivr (sha-pinned, immutable). The arxiv
|
||||
# PDF previously used here was several MB — once base64-encoded into the
|
||||
# Vertex OCR request it ballooned cassettes past 100 MB per test. Keep
|
||||
# the URL stable across runs so cassettes don't churn.
|
||||
TEST_PDF_URL = (
|
||||
"https://cdn.jsdelivr.net/gh/BerriAI/litellm"
|
||||
"@d769e81c90d453240c61fc572cdb27fae06a89d0"
|
||||
"/tests/llm_translation/fixtures/dummy.pdf"
|
||||
)
|
||||
|
||||
|
||||
class BaseOCRTest(ABC):
|
||||
"""
|
||||
Abstract base test class that enforces common OCR tests across all providers.
|
||||
|
||||
Each provider-specific test class should inherit from this and implement
|
||||
get_base_ocr_call_args() to return provider-specific configuration.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_base_ocr_call_args(self) -> dict:
|
||||
"""Must return the base OCR call args for the specific provider"""
|
||||
pass
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_basic_ocr_with_url(self, sync_mode):
|
||||
"""
|
||||
Test basic OCR with a public URL.
|
||||
"""
|
||||
litellm._turn_on_debug()
|
||||
base_ocr_call_args = self.get_base_ocr_call_args()
|
||||
print("BASE OCR Call args=", base_ocr_call_args)
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
|
||||
try:
|
||||
if sync_mode:
|
||||
response = litellm.ocr(
|
||||
document={"type": "document_url", "document_url": TEST_PDF_URL},
|
||||
**base_ocr_call_args,
|
||||
)
|
||||
else:
|
||||
response = await litellm.aocr(
|
||||
document={"type": "document_url", "document_url": TEST_PDF_URL},
|
||||
**base_ocr_call_args,
|
||||
)
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print(f"Sync Mode: {sync_mode}")
|
||||
print(f"Response type: {type(response)}")
|
||||
print(
|
||||
f"Response object: {response.object if hasattr(response, 'object') else 'N/A'}"
|
||||
)
|
||||
|
||||
# Check if response has expected OCR format
|
||||
assert hasattr(response, "pages"), "Response should have 'pages' attribute"
|
||||
assert hasattr(response, "model"), "Response should have 'model' attribute"
|
||||
assert hasattr(
|
||||
response, "object"
|
||||
), "Response should have 'object' attribute"
|
||||
assert (
|
||||
response.object == "ocr"
|
||||
), f"Expected object='ocr', got '{response.object}'"
|
||||
|
||||
# Validate pages structure
|
||||
assert isinstance(response.pages, list), "pages should be a list"
|
||||
assert len(response.pages) > 0, "Should have at least one page"
|
||||
|
||||
# Check first page structure
|
||||
first_page = response.pages[0]
|
||||
assert hasattr(first_page, "index"), "Page should have 'index' attribute"
|
||||
assert hasattr(
|
||||
first_page, "markdown"
|
||||
), "Page should have 'markdown' attribute"
|
||||
|
||||
# Extract text from all pages for validation
|
||||
total_text = "\n\n".join(
|
||||
page.markdown for page in response.pages if page.markdown
|
||||
)
|
||||
print(f"Total pages: {len(response.pages)}")
|
||||
print(f"Total extracted text length: {len(total_text)} characters")
|
||||
print(f"First 200 chars: {total_text[:200]}")
|
||||
print(f"Model: {response.model}")
|
||||
if response.usage_info:
|
||||
print(f"Pages processed: {response.usage_info.pages_processed}")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
assert len(total_text) > 0, "Should extract some text from the document"
|
||||
|
||||
#########################################################
|
||||
# validate we get a response cost in hidden parameters
|
||||
#########################################################
|
||||
hidden_params = response._hidden_params
|
||||
assert isinstance(
|
||||
hidden_params, dict
|
||||
), "Hidden parameters should be a dictionary"
|
||||
|
||||
print("response usage_info:", response.usage_info)
|
||||
|
||||
response_cost = hidden_params.get("response_cost")
|
||||
assert (
|
||||
response_cost is not None
|
||||
), "Response cost should be in hidden parameters"
|
||||
assert response_cost > 0, "Response cost should be greater than 0"
|
||||
print("response_cost=", response_cost)
|
||||
|
||||
except litellm.RateLimitError as e:
|
||||
error_msg = str(e)
|
||||
if "Quota exceeded" in error_msg or "RESOURCE_EXHAUSTED" in error_msg:
|
||||
pytest.skip(f"Quota exceeded - {error_msg}")
|
||||
else:
|
||||
pytest.skip(f"Rate limit exceeded - {error_msg}")
|
||||
except litellm.InternalServerError:
|
||||
pytest.skip("Model is overloaded")
|
||||
except litellm.BadRequestError as e:
|
||||
error_msg = str(e)
|
||||
if (
|
||||
"URL_REJECTED" in error_msg
|
||||
or "Cannot fetch content from the provided URL" in error_msg
|
||||
):
|
||||
pytest.skip(f"URL rejected by provider - {error_msg}")
|
||||
else:
|
||||
pytest.fail(f"OCR call failed: {str(e)}")
|
||||
except Exception as e:
|
||||
pytest.fail(f"OCR call failed: {str(e)}")
|
||||
|
||||
def test_ocr_response_structure(self):
|
||||
"""
|
||||
Test that the OCR response has the correct structure.
|
||||
"""
|
||||
litellm.set_verbose = True
|
||||
base_ocr_call_args = self.get_base_ocr_call_args()
|
||||
|
||||
try:
|
||||
response = litellm.ocr(
|
||||
document={"type": "document_url", "document_url": TEST_PDF_URL},
|
||||
**base_ocr_call_args,
|
||||
)
|
||||
|
||||
# Validate response structure
|
||||
assert hasattr(response, "pages"), "Response should have 'pages' attribute"
|
||||
assert hasattr(response, "model"), "Response should have 'model' attribute"
|
||||
assert hasattr(
|
||||
response, "object"
|
||||
), "Response should have 'object' attribute"
|
||||
assert hasattr(
|
||||
response, "usage_info"
|
||||
), "Response should have 'usage_info' attribute"
|
||||
|
||||
assert isinstance(response.pages, list), "pages should be a list"
|
||||
assert len(response.pages) > 0, "Should have at least one page"
|
||||
assert response.object == "ocr", "object should be 'ocr'"
|
||||
|
||||
# Validate first page structure
|
||||
first_page = response.pages[0]
|
||||
assert hasattr(first_page, "index"), "Page should have 'index' attribute"
|
||||
assert hasattr(
|
||||
first_page, "markdown"
|
||||
), "Page should have 'markdown' attribute"
|
||||
assert isinstance(first_page.markdown, str), "markdown should be a string"
|
||||
|
||||
print(f"\nResponse structure validated:")
|
||||
print(f" - object: {response.object}")
|
||||
print(f" - model: {response.model}")
|
||||
print(f" - pages: {len(response.pages)}")
|
||||
if response.usage_info:
|
||||
print(f" - pages_processed: {response.usage_info.pages_processed}")
|
||||
print(f" - doc_size_bytes: {response.usage_info.doc_size_bytes}")
|
||||
|
||||
except litellm.RateLimitError as e:
|
||||
error_msg = str(e)
|
||||
if "Quota exceeded" in error_msg or "RESOURCE_EXHAUSTED" in error_msg:
|
||||
pytest.skip(f"Quota exceeded - {error_msg}")
|
||||
else:
|
||||
pytest.skip(f"Rate limit exceeded - {error_msg}")
|
||||
except litellm.InternalServerError:
|
||||
pytest.skip("Model is overloaded")
|
||||
except litellm.BadRequestError as e:
|
||||
error_msg = str(e)
|
||||
if (
|
||||
"URL_REJECTED" in error_msg
|
||||
or "Cannot fetch content from the provided URL" in error_msg
|
||||
):
|
||||
pytest.skip(f"URL rejected by provider - {error_msg}")
|
||||
else:
|
||||
pytest.fail(f"OCR response structure test failed: {str(e)}")
|
||||
except Exception as e:
|
||||
pytest.fail(f"OCR response structure test failed: {str(e)}")
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue