Merge remote-tracking branch 'origin/main' into fix41691
Some checks are pending
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run

# Conflicts:
#	litellm/proxy/_types.py
This commit is contained in:
yucheng-berri 2026-10-08 16:41:55 -07:00
commit 586b022744
303 changed files with 25563 additions and 9564 deletions

View file

@ -35,7 +35,7 @@ pub enum EventType {
}
```
Use `#[strum(serialize_all = "snake_case")]` or another supported case style when it exactly matches the contract. Use explicit variant spellings otherwise. With multiple accepted aliases, set `to_string` to the existing canonical output: Strum Display otherwise selects the longest `serialize` spelling
Spell each variant explicitly with `#[strum(serialize = "...")]`; do not use `serialize_all`. With multiple accepted aliases, set `to_string` to the existing canonical output: Strum Display otherwise selects the longest `serialize` spelling
Use only the derives the contract needs. A deserialize-only type should remain deserialize-only. Do not add an unknown variant to a closed enum, or derive Serde for a type that currently has no serialization contract
@ -43,7 +43,7 @@ Plain Serde derives with `rename` or `rename_all` remain appropriate for closed
## Preserve behavior during migration
Read the type, its callers, and existing tests before changing it. Preserve canonical output, accepted aliases, case sensitivity, whitespace handling, unknown values, malformed-input rejection, and existing public conversion APIs. Keep conversions needed by callers or compatibility even when Serde no longer uses them
Read the type, its callers, and existing tests before changing it. Preserve canonical output, accepted aliases, case sensitivity, whitespace handling, unknown values, malformed-input rejection, and existing public conversion contracts. Wrappers that only delegate to a Strum-derived trait are removed rather than kept: callers use `FromStr` and `From<Enum> for &'static str` directly. Keep conversions that add behavior needed by callers or compatibility even when Serde no longer uses them
Use the workspace dependencies and enable `serde_with.workspace = true` in a crate only when needed. Check the versions and enabled features in `Cargo.toml` and `Cargo.lock` rather than upgrading dependencies for this refactor

View file

@ -4,6 +4,10 @@ For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.a
For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md)
For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array; derive Strum string conversions instead of hand-written variant-to-string matches
Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all`
## Test placement
- Never create a `tests.rs` (or `test.rs`) file under `src/`, and never `#[path = "tests.rs"] mod tests;`
@ -17,6 +21,9 @@ For string-valued enums and their Serde conversions, follow [.agents/skills/rust
Use [`#[rstest]`](https://docs.rs/rstest/latest/rstest/attr.rstest.html) for new and updated tests and [`#[fixture]`](https://docs.rs/rstest/latest/rstest/attr.fixture.html) for reusable setup, injected through typed test arguments. Express input variations as named `#[case::name(...)]` cases instead of loops or duplicated tests so each failure identifies its case. Keep behavior assertions in the test body and fixtures focused on setup. Use the workspace `rstest` dependency
- Never loop over inputs (`for`, `.iter().for_each`, `.all`) inside a test body; give each input its own `#[case::name(...)]`, or use `#[values(...)]` for a cross product
- Exception: a test pinning a Rust table against a repo-owned data file (for example `include_str!` of a JSON config) may iterate that file's entries
## Error definitions
- A crate's errors live in `src/error.rs`, defined with `thiserror`, and re-exported from `lib.rs`

View file

@ -3485,6 +3485,7 @@ dependencies = [
"rstest",
"serde",
"serde_json",
"strum",
"thiserror 2.0.19",
"tokio",
]
@ -3692,7 +3693,9 @@ dependencies = [
"litellm-auth-types",
"rstest",
"serde",
"serde_json",
"serde_yaml_ng",
"strum",
"tempfile",
"thiserror 2.0.19",
]
@ -3990,6 +3993,7 @@ dependencies = [
"rustls-native-certs",
"serde",
"serde_json",
"strum",
"tempfile",
"thiserror 2.0.19",
"tokio",
@ -4539,6 +4543,7 @@ dependencies = [
"serde",
"serde_json",
"sha2 0.10.9",
"strum",
"tar",
"target-lexicon",
"tempfile",

View file

@ -8,6 +8,7 @@ repository.workspace = true
[dependencies]
serde.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
strum.workspace = true
thiserror.workspace = true
[dev-dependencies]

View file

@ -1,57 +1,44 @@
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, Hash)]
#[derive(
Clone,
Copy,
Debug,
Deserialize,
Serialize,
PartialEq,
Eq,
Hash,
strum::EnumString,
strum::IntoStaticStr,
strum::VariantArray,
)]
pub enum CacheType {
#[serde(rename = "local")]
#[strum(serialize = "local")]
Local,
#[serde(rename = "redis")]
#[strum(serialize = "redis")]
Redis,
#[serde(rename = "redis-semantic")]
#[strum(serialize = "redis-semantic")]
RedisSemantic,
#[serde(rename = "valkey-semantic")]
#[strum(serialize = "valkey-semantic")]
ValkeySemantic,
#[serde(rename = "s3")]
#[strum(serialize = "s3")]
S3,
#[serde(rename = "disk")]
#[strum(serialize = "disk")]
Disk,
#[serde(rename = "qdrant-semantic")]
#[strum(serialize = "qdrant-semantic")]
QdrantSemantic,
#[serde(rename = "azure-blob")]
#[strum(serialize = "azure-blob")]
AzureBlob,
#[serde(rename = "gcs")]
#[strum(serialize = "gcs")]
Gcs,
}
impl CacheType {
pub const ALL: [Self; 9] = [
Self::Local,
Self::Redis,
Self::RedisSemantic,
Self::ValkeySemantic,
Self::S3,
Self::Disk,
Self::QdrantSemantic,
Self::AzureBlob,
Self::Gcs,
];
pub const fn as_python_name(self) -> &'static str {
match self {
Self::Local => "local",
Self::Redis => "redis",
Self::RedisSemantic => "redis-semantic",
Self::ValkeySemantic => "valkey-semantic",
Self::S3 => "s3",
Self::Disk => "disk",
Self::QdrantSemantic => "qdrant-semantic",
Self::AzureBlob => "azure-blob",
Self::Gcs => "gcs",
}
}
pub fn from_python_name(value: &str) -> Option<Self> {
Self::ALL
.into_iter()
.find(|cache_type| cache_type.as_python_name() == value)
}
}

View file

@ -1,5 +1,6 @@
use litellm_cache::CacheType;
use rstest::rstest;
use strum::VariantArray;
#[rstest]
#[case(CacheType::Local, "local")]
@ -15,16 +16,16 @@ fn every_python_cache_type_has_one_round_trip_identity(
#[case] cache_type: CacheType,
#[case] name: &str,
) {
assert_eq!(cache_type.as_python_name(), name);
assert_eq!(CacheType::from_python_name(name), Some(cache_type));
assert_eq!(<&'static str>::from(cache_type), name);
assert_eq!(name.parse::<CacheType>().ok(), Some(cache_type));
assert_eq!(
serde_json::to_value(cache_type).unwrap(),
serde_json::Value::from(name)
);
assert_eq!(
CacheType::ALL
CacheType::VARIANTS
.iter()
.filter(|candidate| candidate.as_python_name() == name)
.filter(|candidate| <&'static str>::from(**candidate) == name)
.count(),
1
);
@ -33,7 +34,10 @@ fn every_python_cache_type_has_one_round_trip_identity(
#[rstest]
fn python_cache_types_are_listed_in_python_order() {
assert_eq!(
CacheType::ALL.map(CacheType::as_python_name),
CacheType::VARIANTS
.iter()
.map(|cache_type| <&'static str>::from(*cache_type))
.collect::<Vec<_>>(),
[
"local",
"redis",
@ -52,5 +56,5 @@ fn python_cache_types_are_listed_in_python_order() {
#[case::unknown("memcached")]
#[case::case_sensitive("Redis")]
fn unknown_python_names_have_no_cache_type(#[case] name: &str) {
assert_eq!(CacheType::from_python_name(name), None);
assert!(name.parse::<CacheType>().is_err());
}

View file

@ -936,9 +936,6 @@ mod payload_tests {
/// The payload phases of `Logging` on top of `StubLogger`, with `pre_call` handing the
/// payload to the case's `on_pre_call`.
const PAYLOAD_LOGGER: &CStr = c"
class Request:
pass
class PayloadLogger(StubLogger):
def update_from_kwargs(self, **update):
self.update = update
@ -954,7 +951,7 @@ class PayloadLogger(StubLogger):
self.record('post_call', None)
self.post = (original_response, api_key, additional_args)
request = Request()
bound = {}
kwargs = {}
logger = PayloadLogger()
on_pre_call = lambda additional_args: None
@ -1225,10 +1222,10 @@ on_pre_call = lambda args: observed.append(
def check():
assert observed == [(True, True)], observed
")]
#[case::request_attribute_behind_an_omitted_keyword(c"
#[case::bound_value_behind_an_omitted_keyword(c"
document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
pages = [0]
request.document = document
bound['document'] = document
kwargs = {'pages': pages}
observed = []
on_pre_call = lambda args: observed.append(

View file

@ -13,21 +13,21 @@ use pyo3::{
pub struct PublicCall {
args: Py<PyTuple>,
kwargs: Py<PyDict>,
request: Py<PyAny>,
bound: Py<PyDict>,
}
impl PublicCall {
/// Copies the keyword arguments once, so the legacy path's rewrites never reach the
/// caller's own dict while every value keeps its identity.
pub fn capture(
request: &Bound<'_, PyAny>,
bound: &Bound<'_, PyDict>,
args: &Bound<'_, PyTuple>,
kwargs: &Bound<'_, PyDict>,
) -> PyResult<Self> {
Ok(Self {
args: args.clone().unbind(),
kwargs: kwargs.copy()?.unbind(),
request: request.clone().unbind(),
bound: bound.clone().unbind(),
})
}
@ -55,31 +55,30 @@ impl PublicCall {
py: Python<'py>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
lookup(self.kwargs.bind(py), self.request.bind(py), name)
lookup(self.kwargs.bind(py), self.bound.bind(py), name)
}
pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.args)?;
visit.call(&self.kwargs)?;
visit.call(&self.request)
visit.call(&self.bound)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_support::{local, local_dict};
fn capture<'py>(py: Python<'py>, source: &std::ffi::CStr) -> (PublicCall, Bound<'py, PyDict>) {
let locals = PyDict::new(py);
py.run(source, Some(&locals), Some(&locals)).unwrap();
let request = locals.get_item("request").unwrap().unwrap();
let kwargs = locals
.get_item("kwargs")
.unwrap()
.unwrap()
.cast_into::<PyDict>()
.unwrap();
let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap();
let call = PublicCall::capture(
&local_dict(&locals, "bound"),
&PyTuple::empty(py),
&local_dict(&locals, "kwargs"),
)
.unwrap();
(call, locals)
}
@ -91,25 +90,22 @@ mod tests {
py,
c"
pages = [0]
class Request:
pass
request = Request()
bound = {'pages': [1]}
kwargs = {'pages': pages}
",
);
let caller = locals
.get_item("kwargs")
.unwrap()
.unwrap()
.cast_into::<PyDict>()
.unwrap();
let caller = local_dict(&locals, "kwargs");
call.kwargs()
.bind(py)
.set_item("litellm_call_id", "call")
.unwrap();
assert!(!caller.contains("litellm_call_id").unwrap());
let pages = locals.get_item("pages").unwrap().unwrap();
assert!(call.lookup(py, "pages").unwrap().unwrap().is(&pages));
assert!(
call.lookup(py, "pages")
.unwrap()
.unwrap()
.is(local(&locals, "pages"))
);
});
}
}

View file

@ -173,21 +173,23 @@ pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py,
locals.get_item(name).unwrap().unwrap()
}
/// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`).
pub(crate) fn local_dict<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyDict> {
local(locals, name).cast_into().unwrap()
}
/// A legacy call over the namespace's `kwargs` and `bound` dicts, each empty when absent.
pub(crate) fn legacy_call(
py: Python<'_>,
locals: &Bound<'_, PyDict>,
asynchronous: bool,
) -> LegacyLogging {
let request = locals
.get_item("request")
.unwrap()
.unwrap_or_else(|| py.None().into_bound(py));
let kwargs = locals
.get_item("kwargs")
.unwrap()
.map(|kwargs| kwargs.cast_into::<PyDict>().unwrap())
.unwrap_or_else(|| PyDict::new(py));
let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap();
let dict = |name: &str| {
locals
.get_item(name)
.unwrap()
.map(|value| value.cast_into::<PyDict>().unwrap())
.unwrap_or_else(|| PyDict::new(py))
};
let call = PublicCall::capture(&dict("bound"), &PyTuple::empty(py), &dict("kwargs")).unwrap();
LegacyLogging::new(py, crate::LoggingOperation::Ocr, call, asynchronous)
}

View file

@ -9,8 +9,10 @@ repository.workspace = true
litellm-auth-types.workspace = true
serde.workspace = true
serde_yaml_ng = "0.10.0"
strum.workspace = true
thiserror.workspace = true
[dev-dependencies]
rstest.workspace = true
serde_json.workspace = true
tempfile.workspace = true

View file

@ -41,57 +41,43 @@ impl fmt::Debug for McpServer {
}
}
#[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq, strum::IntoStaticStr)]
#[serde(rename_all = "snake_case")]
pub enum McpTransport {
#[default]
#[strum(serialize = "http")]
Http,
#[strum(serialize = "sse")]
Sse,
#[strum(serialize = "stdio")]
Stdio,
}
impl McpTransport {
pub fn as_str(self) -> &'static str {
match self {
Self::Http => "http",
Self::Sse => "sse",
Self::Stdio => "stdio",
}
}
}
#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::IntoStaticStr)]
#[serde(rename_all = "snake_case")]
pub enum McpAuth {
#[strum(serialize = "none")]
None,
#[strum(serialize = "api_key")]
ApiKey,
#[strum(serialize = "bearer_token")]
BearerToken,
#[strum(serialize = "basic")]
Basic,
#[strum(serialize = "authorization")]
Authorization,
#[strum(serialize = "token")]
Token,
#[strum(serialize = "oauth2")]
Oauth2,
#[strum(serialize = "aws_sigv4")]
AwsSigv4,
#[strum(serialize = "oauth2_token_exchange")]
Oauth2TokenExchange,
#[strum(serialize = "oauth2_id_jag")]
Oauth2IdJag,
#[strum(serialize = "true_passthrough")]
TruePassthrough,
#[strum(serialize = "oauth_delegate")]
OauthDelegate,
}
impl McpAuth {
pub fn as_str(self) -> &'static str {
match self {
Self::None => "none",
Self::ApiKey => "api_key",
Self::BearerToken => "bearer_token",
Self::Basic => "basic",
Self::Authorization => "authorization",
Self::Token => "token",
Self::Oauth2 => "oauth2",
Self::AwsSigv4 => "aws_sigv4",
Self::Oauth2TokenExchange => "oauth2_token_exchange",
Self::Oauth2IdJag => "oauth2_id_jag",
Self::TruePassthrough => "true_passthrough",
Self::OauthDelegate => "oauth_delegate",
}
}
}

View file

@ -0,0 +1,39 @@
use litellm_config::{McpAuth, McpTransport};
use rstest::rstest;
use serde_json::json;
#[rstest]
#[case::http(McpTransport::Http, "http")]
#[case::sse(McpTransport::Sse, "sse")]
#[case::stdio(McpTransport::Stdio, "stdio")]
fn mcp_transport_as_str_matches_the_serde_spelling(
#[case] transport: McpTransport,
#[case] name: &str,
) {
assert_eq!(<&'static str>::from(transport), name);
assert_eq!(
serde_json::from_value::<McpTransport>(json!(name)).unwrap(),
transport
);
}
#[rstest]
#[case::none(McpAuth::None, "none")]
#[case::api_key(McpAuth::ApiKey, "api_key")]
#[case::bearer_token(McpAuth::BearerToken, "bearer_token")]
#[case::basic(McpAuth::Basic, "basic")]
#[case::authorization(McpAuth::Authorization, "authorization")]
#[case::token(McpAuth::Token, "token")]
#[case::oauth2(McpAuth::Oauth2, "oauth2")]
#[case::aws_sigv4(McpAuth::AwsSigv4, "aws_sigv4")]
#[case::oauth2_token_exchange(McpAuth::Oauth2TokenExchange, "oauth2_token_exchange")]
#[case::oauth2_id_jag(McpAuth::Oauth2IdJag, "oauth2_id_jag")]
#[case::true_passthrough(McpAuth::TruePassthrough, "true_passthrough")]
#[case::oauth_delegate(McpAuth::OauthDelegate, "oauth_delegate")]
fn mcp_auth_as_str_matches_the_serde_spelling(#[case] auth: McpAuth, #[case] name: &str) {
assert_eq!(<&'static str>::from(auth), name);
assert_eq!(
serde_json::from_value::<McpAuth>(json!(name)).unwrap(),
auth
);
}

View file

@ -7,17 +7,26 @@ pub struct CustomLlmProvider<'a> {
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, EnumString, IntoStaticStr)]
#[strum(serialize_all = "snake_case")]
pub enum LlmProviders {
#[strum(serialize = "anthropic")]
Anthropic,
#[strum(serialize = "aws_textract")]
AwsTextract,
#[strum(serialize = "azure_ai")]
AzureAi,
#[strum(serialize = "bedrock")]
Bedrock,
#[strum(serialize = "cohere")]
Cohere,
#[strum(serialize = "mistral")]
Mistral,
#[strum(serialize = "openai")]
Openai,
#[strum(serialize = "openai_like")]
OpenaiLike,
#[strum(serialize = "reducto")]
Reducto,
#[strum(serialize = "vertex_ai")]
VertexAi,
}

View file

@ -17,18 +17,13 @@ pub const EMPTY_TEXT_PLACEHOLDER: &str =
"[System: Empty message content sanitised to satisfy protocol]";
#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq)]
#[strum(serialize_all = "snake_case")]
pub enum TurnRole {
#[strum(serialize = "user")]
User,
#[strum(serialize = "assistant")]
Assistant,
}
impl TurnRole {
pub fn as_str(self) -> &'static str {
self.into()
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Turn {
pub role: TurnRole,

View file

@ -49,8 +49,8 @@ fn info(name: &str, config: &McpServer) -> ServerInfo {
let identity = format!(
"{name}|{}|{}|{}|{}",
config.url.as_ref().map_or("", SecretValue::expose),
config.transport.as_str(),
config.auth_type.map_or("", McpAuth::as_str),
<&'static str>::from(config.transport),
config.auth_type.map_or("", <&'static str>::from),
config.alias.as_deref().unwrap_or("")
);
ServerInfo {

View file

@ -2,78 +2,97 @@ use pyo3::{prelude::*, types::PyDict};
pub fn lookup<'py>(
kwargs: &Bound<'py, PyDict>,
request: &Bound<'py, PyAny>,
bound: &Bound<'py, PyDict>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
if let Some(value) = kwargs.get_item(name)? {
return Ok(Some(value));
match kwargs.get_item(name)? {
Some(value) => Ok(Some(value)),
None => bound.get_item(name),
}
if let Ok(bound) = request.cast::<PyDict>() {
return bound.get_item(name);
}
request.getattr_opt(name)
}
pub fn present<'py>(
kwargs: &Bound<'py, PyDict>,
bound: &Bound<'py, PyDict>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
Ok(lookup(kwargs, bound, name)?.filter(|value| !value.is_none()))
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[test]
fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() {
fn dicts<'py>(
py: Python<'py>,
kwargs: &str,
bound: &str,
) -> (Bound<'py, PyDict>, Bound<'py, PyDict>) {
let eval = |source: &str| {
py.eval(&std::ffi::CString::new(source).unwrap(), None, None)
.unwrap()
.cast_into::<PyDict>()
.unwrap()
};
(eval(kwargs), eval(bound))
}
#[rstest]
#[case::keyword_wins(
"{'api_key': 'keyword'}",
"{'api_key': 'bound'}",
Some(Some("keyword"))
)]
#[case::explicit_none_wins("{'api_key': None}", "{'api_key': 'bound'}", Some(None))]
#[case::bound_fallback("{}", "{'api_key': 'bound'}", Some(Some("bound")))]
#[case::missing("{}", "{}", None)]
fn lookup_prefers_the_keyword_and_falls_back_to_bound(
#[case] kwargs: &str,
#[case] bound: &str,
#[case] expected: Option<Option<&str>>,
) {
crate::initialize_python();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(
c"
key = object()
document = {'type': 'document_url'}
class Request:
api_key = 'from-request'
api_base = 'from-request'
document = document
request = Request()
kwargs = {'api_key': key, 'api_base': None}
",
Some(&locals),
Some(&locals),
)
.unwrap();
let item = |name: &str| locals.get_item(name).unwrap().unwrap();
let kwargs = item("kwargs").cast_into::<PyDict>().unwrap();
let request = item("request");
let find = |name: &str| lookup(&kwargs, &request, name).unwrap();
assert!(find("api_key").unwrap().is(item("key")));
assert!(find("api_base").unwrap().is_none());
assert!(find("document").unwrap().is(item("document")));
assert!(find("model").is_none());
let (kwargs, bound) = dicts(py, kwargs, bound);
let value = lookup(&kwargs, &bound, "api_key")
.unwrap()
.map(|value| value.extract::<Option<String>>().unwrap());
assert_eq!(value, expected.map(|value| value.map(str::to_owned)));
});
}
#[rstest::rstest]
#[case::prepared_value("{'api_key': 'replacement'}", Some("replacement"))]
#[case::explicit_none("{'api_key': None}", None)]
#[case::bound_fallback("{}", Some("original"))]
fn prepared_mapping_overrides_bound_values(
#[case] source: &str,
#[rstest]
#[case::explicit_none_hides_bound("{'api_key': None}", "{'api_key': 'bound'}", None)]
#[case::bound_none("{}", "{'api_key': None}", None)]
#[case::bound_value("{}", "{'api_key': 'bound'}", Some("bound"))]
fn present_treats_none_as_unset(
#[case] kwargs: &str,
#[case] bound: &str,
#[case] expected: Option<&str>,
) {
crate::initialize_python();
Python::attach(|py| {
let (kwargs, bound) = dicts(py, kwargs, bound);
let value = present(&kwargs, &bound, "api_key")
.unwrap()
.map(|value| value.extract::<String>().unwrap());
assert_eq!(value.as_deref(), expected);
});
}
#[rstest]
fn lookup_returns_the_callers_object() {
crate::initialize_python();
Python::attach(|py| {
let document = PyDict::new(py);
let bound = PyDict::new(py);
bound.set_item("api_key", "original").unwrap();
let source = std::ffi::CString::new(source).unwrap();
let prepared = py
.eval(&source, None, None)
.unwrap()
.cast_into::<PyDict>()
.unwrap();
let value = lookup(&prepared, bound.as_any(), "api_key")
bound.set_item("document", &document).unwrap();
let found = lookup(&PyDict::new(py), &bound, "document")
.unwrap()
.unwrap();
assert_eq!(
value.extract::<Option<String>>().unwrap().as_deref(),
expected
);
assert!(found.is(&document));
});
}
}

View file

@ -19,7 +19,7 @@ mod owned;
mod runtime;
mod services;
pub use argument::lookup;
pub use argument::{lookup, present};
pub use binding::PythonBinding;
pub use conversion_cache::{FromPythonCache, ToPythonCache};
pub use driver::{CallOptions, run_call};

View file

@ -21,6 +21,7 @@ rustls-native-certs.workspace = true
tokio-tungstenite.workspace = true
serde.workspace = true
serde_json.workspace = true
strum.workspace = true
thiserror.workspace = true
tokio.workspace = true
tracing.workspace = true

View file

@ -42,30 +42,28 @@ impl KeyExchangeGroup {
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, strum::EnumString)]
#[strum(
parse_err_ty = Unsupported,
parse_err_fn = unsupported_cipher_token
)]
pub enum Tls12CipherSuite {
#[strum(serialize = "ECDHE-ECDSA-AES128-GCM-SHA256")]
EcdheEcdsaAes128Gcm,
#[strum(serialize = "ECDHE-ECDSA-AES256-GCM-SHA384")]
EcdheEcdsaAes256Gcm,
#[strum(serialize = "ECDHE-ECDSA-CHACHA20-POLY1305")]
EcdheEcdsaChacha20,
#[strum(serialize = "ECDHE-RSA-AES128-GCM-SHA256")]
EcdheRsaAes128Gcm,
#[strum(serialize = "ECDHE-RSA-AES256-GCM-SHA384")]
EcdheRsaAes256Gcm,
#[strum(serialize = "ECDHE-RSA-CHACHA20-POLY1305")]
EcdheRsaChacha20,
}
impl FromStr for Tls12CipherSuite {
type Err = Unsupported;
fn from_str(name: &str) -> Result<Self, Self::Err> {
match name {
"ECDHE-ECDSA-AES128-GCM-SHA256" => Ok(Self::EcdheEcdsaAes128Gcm),
"ECDHE-ECDSA-AES256-GCM-SHA384" => Ok(Self::EcdheEcdsaAes256Gcm),
"ECDHE-ECDSA-CHACHA20-POLY1305" => Ok(Self::EcdheEcdsaChacha20),
"ECDHE-RSA-AES128-GCM-SHA256" => Ok(Self::EcdheRsaAes128Gcm),
"ECDHE-RSA-AES256-GCM-SHA384" => Ok(Self::EcdheRsaAes256Gcm),
"ECDHE-RSA-CHACHA20-POLY1305" => Ok(Self::EcdheRsaChacha20),
_ => Err(Unsupported::CipherToken(name.to_owned())),
}
}
fn unsupported_cipher_token(name: &str) -> Unsupported {
Unsupported::CipherToken(name.to_owned())
}
impl Tls12CipherSuite {
@ -367,6 +365,39 @@ mod tests {
);
}
#[rstest]
#[case::ecdhe_ecdsa_aes128(
"ECDHE-ECDSA-AES128-GCM-SHA256",
Tls12CipherSuite::EcdheEcdsaAes128Gcm
)]
#[case::ecdhe_ecdsa_aes256(
"ECDHE-ECDSA-AES256-GCM-SHA384",
Tls12CipherSuite::EcdheEcdsaAes256Gcm
)]
#[case::ecdhe_ecdsa_chacha20(
"ECDHE-ECDSA-CHACHA20-POLY1305",
Tls12CipherSuite::EcdheEcdsaChacha20
)]
#[case::ecdhe_rsa_aes128("ECDHE-RSA-AES128-GCM-SHA256", Tls12CipherSuite::EcdheRsaAes128Gcm)]
#[case::ecdhe_rsa_aes256("ECDHE-RSA-AES256-GCM-SHA384", Tls12CipherSuite::EcdheRsaAes256Gcm)]
#[case::ecdhe_rsa_chacha20("ECDHE-RSA-CHACHA20-POLY1305", Tls12CipherSuite::EcdheRsaChacha20)]
fn tls12_cipher_suite_parses_each_openssl_name(
#[case] name: &str,
#[case] expected: Tls12CipherSuite,
) {
assert_eq!(name.parse::<Tls12CipherSuite>().unwrap(), expected);
}
#[rstest]
#[case::unknown("AES128-SHA")]
#[case::case_sensitive("ecdhe-rsa-aes256-gcm-sha384")]
fn unsupported_cipher_token_keeps_the_verbatim_name(#[case] name: &str) {
assert_eq!(
name.parse::<Tls12CipherSuite>().unwrap_err(),
Unsupported::CipherToken(name.to_owned())
);
}
#[test]
fn named_suites_are_the_only_tls12_suites_offered_and_tls13_stays() {
let tls = ClientConfig::try_from(&config(HttpSettings {

View file

@ -9,41 +9,25 @@ pub use content::{
/// Reasoning effort level accepted or applied by the model.
#[macro_rules_attribute::apply(crate::wire_type)]
#[derive(Copy, Eq, IntoStaticStr)]
#[derive(Copy, Eq, IntoStaticStr, strum::EnumString, strum::VariantArray)]
#[serde(rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]
pub enum ReasoningEffort {
#[strum(serialize = "none")]
None,
#[strum(serialize = "minimal")]
Minimal,
#[strum(serialize = "low")]
Low,
#[strum(serialize = "medium")]
Medium,
#[strum(serialize = "high")]
High,
#[strum(serialize = "xhigh")]
Xhigh,
#[strum(serialize = "max")]
Max,
}
impl ReasoningEffort {
pub const ALL: [Self; 7] = [
Self::None,
Self::Minimal,
Self::Low,
Self::Medium,
Self::High,
Self::Xhigh,
Self::Max,
];
pub fn as_str(self) -> &'static str {
self.into()
}
pub fn parse(value: &str) -> Option<Self> {
Self::ALL
.into_iter()
.find(|effort| effort.as_str() == value)
}
}
#[macro_rules_attribute::apply(crate::wire_type)]
#[serde(untagged)]
pub enum ChatMessageContent {
@ -195,6 +179,7 @@ pub struct ChatCompletionChunk {
#[cfg(test)]
mod tests {
use rstest::rstest;
use strum::VariantArray;
use super::*;
@ -213,10 +198,10 @@ mod tests {
) {
assert_eq!(
serde_json::to_value(effort).unwrap(),
Value::String(effort.as_str().to_string())
Value::String(<&'static str>::from(effort).to_string())
);
assert_eq!(ReasoningEffort::parse(effort.as_str()), Some(effort));
assert!(ReasoningEffort::ALL.contains(&effort));
assert_eq!(<&'static str>::from(effort).parse(), Ok(effort));
assert!(ReasoningEffort::VARIANTS.contains(&effort));
}
#[rstest]
@ -224,6 +209,6 @@ mod tests {
#[case::uppercase("HIGH")]
#[case::empty("")]
fn reasoning_effort_parse_rejects(#[case] value: &str) {
assert_eq!(ReasoningEffort::parse(value), None);
assert!(value.parse::<ReasoningEffort>().is_err());
}
}

View file

@ -165,10 +165,12 @@ pub struct Safeguard {
)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))]
#[strum(serialize_all = "snake_case")]
pub enum MessageRole {
#[strum(serialize = "user")]
User,
#[strum(serialize = "assistant")]
Assistant,
#[strum(serialize = "system")]
System,
#[strum(default, transparent)]
Other(String),
@ -205,8 +207,8 @@ impl From<MessageRole> for String {
)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))]
#[strum(serialize_all = "snake_case")]
pub enum MessageType {
#[strum(serialize = "message")]
Message,
#[strum(default, transparent)]
Other(String),
@ -243,15 +245,22 @@ impl From<MessageType> for String {
)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))]
#[strum(serialize_all = "snake_case")]
pub enum StopReason {
#[strum(serialize = "end_turn")]
EndTurn,
#[strum(serialize = "max_tokens")]
MaxTokens,
#[strum(serialize = "stop_sequence")]
StopSequence,
#[strum(serialize = "tool_use")]
ToolUse,
#[strum(serialize = "refusal")]
Refusal,
#[strum(serialize = "compaction")]
Compaction,
#[strum(serialize = "pause_turn")]
PauseTurn,
#[strum(serialize = "model_context_window_exceeded")]
ModelContextWindowExceeded,
#[strum(default, transparent)]
Other(String),

View file

@ -30,16 +30,24 @@ pub enum MessageContent {
)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))]
#[strum(serialize_all = "snake_case")]
pub enum ContentBlockType {
#[strum(serialize = "text")]
Text,
#[strum(serialize = "thinking")]
Thinking,
#[strum(serialize = "redacted_thinking")]
RedactedThinking,
#[strum(serialize = "tool_use")]
ToolUse,
#[strum(serialize = "server_tool_use")]
ServerToolUse,
#[strum(serialize = "tool_result")]
ToolResult,
#[strum(serialize = "compaction")]
Compaction,
#[strum(serialize = "advisor_tool_result")]
AdvisorToolResult,
#[strum(serialize = "web_search_tool_result")]
WebSearchToolResult,
#[strum(default, transparent)]
Other(String),
@ -124,23 +132,21 @@ pub struct Message {
}
#[macro_rules_attribute::apply(crate::wire_type)]
#[derive(Copy, Hash, IntoStaticStr, Eq)]
#[derive(Copy, Hash, IntoStaticStr, Eq, strum::VariantArray)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
pub enum EffortLevel {
#[strum(serialize = "low")]
Low,
#[strum(serialize = "medium")]
Medium,
#[strum(serialize = "high")]
High,
#[strum(serialize = "xhigh")]
Xhigh,
#[strum(serialize = "max")]
Max,
}
impl EffortLevel {
pub fn as_str(self) -> &'static str {
self.into()
}
}
impl From<EffortLevel> for ReasoningEffort {
fn from(level: EffortLevel) -> Self {
match level {
@ -156,18 +162,13 @@ impl From<EffortLevel> for ReasoningEffort {
#[macro_rules_attribute::apply(crate::wire_type)]
#[derive(Copy, IntoStaticStr, Eq)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
pub enum Speed {
#[strum(serialize = "fast")]
Fast,
#[strum(serialize = "standard")]
Standard,
}
impl Speed {
pub fn as_str(self) -> &'static str {
self.into()
}
}
/// The tools whose presence changes how the request is sent. Every other tool, custom or
/// server, deserializes as `Recognized::Unrecognized` and passes through verbatim.
#[macro_rules_attribute::apply(crate::wire_type)]
@ -691,7 +692,10 @@ mod tests {
#[rstest]
fn speed_names_match_the_wire(#[values(Speed::Fast, Speed::Standard)] speed: Speed) {
assert_eq!(serde_json::to_value(speed).unwrap(), json!(speed.as_str()));
assert_eq!(
serde_json::to_value(speed).unwrap(),
json!(<&'static str>::from(speed))
);
}
#[rstest]
@ -705,10 +709,13 @@ mod tests {
)]
level: EffortLevel,
) {
assert_eq!(serde_json::to_value(level).unwrap(), json!(level.as_str()));
assert_eq!(
serde_json::to_value(level).unwrap(),
json!(<&'static str>::from(level))
);
assert_eq!(
serde_json::to_value(ReasoningEffort::from(level)).unwrap(),
json!(level.as_str())
json!(<&'static str>::from(level))
);
}
}

View file

@ -4,7 +4,8 @@ These links describe the provider-specific wire types in this directory. Shared
## anthropic.rs
`AnthropicBeta` and `BetaSet` represent the `anthropic-beta` header values and their wire spelling
`AnthropicBeta`, `BetaSet` and `BetaProvider` represent the `anthropic-beta` header values, their
wire spelling and per-host support
- https://platform.claude.com/docs/en/api/beta-headers.md

View file

@ -7,6 +7,39 @@ use std::{
str::FromStr,
};
use strum::VariantArray;
/// A provider column of `litellm/anthropic_beta_headers_config.json`: which betas a host accepts
/// and under which name.
#[derive(
Clone,
Copy,
Debug,
PartialEq,
Eq,
Hash,
strum::AsRefStr,
strum::Display,
strum::EnumString,
strum::VariantArray,
)]
pub enum BetaProvider {
#[strum(serialize = "anthropic")]
Anthropic,
#[strum(serialize = "azure_ai")]
AzureAi,
#[strum(serialize = "bedrock_converse")]
BedrockConverse,
#[strum(serialize = "bedrock")]
Bedrock,
#[strum(serialize = "bedrock_mantle")]
BedrockMantle,
#[strum(serialize = "vertex_ai")]
VertexAi,
#[strum(serialize = "databricks")]
Databricks,
}
/// One value of the `anthropic-beta` header. Equality, ordering and hashing follow the wire
/// string, so a value parsed from a caller's header never disagrees with the matching variant.
#[derive(Clone, Debug, strum::AsRefStr, strum::Display, strum::EnumString)]
@ -35,12 +68,68 @@ pub enum AnthropicBeta {
PerTurnControl20260701,
#[strum(serialize = "dangerous-tool-use-2026-09-03")]
DangerousToolUse20260903,
#[strum(serialize = "bash_20241022")]
Bash20241022,
#[strum(serialize = "bash_20250124")]
Bash20250124,
#[strum(serialize = "claude-code-20250219")]
ClaudeCode20250219,
#[strum(serialize = "code-execution-2025-08-25")]
CodeExecution20250825,
#[strum(serialize = "computer-use-2025-01-24")]
ComputerUse20250124,
#[strum(serialize = "computer-use-2025-11-24")]
ComputerUse20251124,
#[strum(serialize = "context-1m-2025-08-07")]
Context1m20250807,
#[strum(serialize = "effort-2025-11-24")]
Effort20251124,
#[strum(serialize = "files-api-2025-04-14")]
FilesApi20250414,
#[strum(serialize = "fine-grained-tool-streaming-2025-05-14")]
FineGrainedToolStreaming20250514,
#[strum(serialize = "interleaved-thinking-2025-05-14")]
InterleavedThinking20250514,
#[strum(serialize = "mcp-client-2025-04-04")]
McpClient20250404,
#[strum(serialize = "mcp-client-2025-11-20")]
McpClient20251120,
#[strum(serialize = "mcp-servers-2025-12-04")]
McpServers20251204,
#[strum(serialize = "mid-conversation-output-config-2026-07-01")]
MidConversationOutputConfig20260701,
#[strum(serialize = "mid-conversation-tool-changes-2026-07-01")]
MidConversationToolChanges20260701,
#[strum(serialize = "output-128k-2025-02-19")]
Output128k20250219,
#[strum(serialize = "prompt-caching-scope-2026-01-05")]
PromptCachingScope20260105,
#[strum(serialize = "skills-2025-10-02")]
Skills20251002,
#[strum(serialize = "structured-output-2024-03-01")]
StructuredOutput20240301,
#[strum(serialize = "text_editor_20241022")]
TextEditor20241022,
#[strum(serialize = "text_editor_20250124")]
TextEditor20250124,
#[strum(serialize = "thinking-binding-controls-2026-08-01")]
ThinkingBindingControls20260801,
#[strum(serialize = "thinking-display-updates-2026-08-18")]
ThinkingDisplayUpdates20260818,
#[strum(serialize = "token-efficient-tools-2025-02-19")]
TokenEfficientTools20250219,
#[strum(serialize = "tool-examples-2025-10-29")]
ToolExamples20251029,
#[strum(serialize = "tool-search-tool-2025-10-19")]
ToolSearchTool20251019,
#[strum(serialize = "inline-tools-2026-09-15")]
InlineTools20260915,
#[strum(default, transparent)]
Other(String),
}
impl AnthropicBeta {
pub const KNOWN: [Self; 12] = [
pub const KNOWN: [Self; 40] = [
Self::Oauth20250420,
Self::WebFetch20250910,
Self::WebSearch20250305,
@ -53,11 +142,153 @@ impl AnthropicBeta {
Self::AdvisorTool20260301,
Self::PerTurnControl20260701,
Self::DangerousToolUse20260903,
Self::Bash20241022,
Self::Bash20250124,
Self::ClaudeCode20250219,
Self::CodeExecution20250825,
Self::ComputerUse20250124,
Self::ComputerUse20251124,
Self::Context1m20250807,
Self::Effort20251124,
Self::FilesApi20250414,
Self::FineGrainedToolStreaming20250514,
Self::InterleavedThinking20250514,
Self::McpClient20250404,
Self::McpClient20251120,
Self::McpServers20251204,
Self::MidConversationOutputConfig20260701,
Self::MidConversationToolChanges20260701,
Self::Output128k20250219,
Self::PromptCachingScope20260105,
Self::Skills20251002,
Self::StructuredOutput20240301,
Self::TextEditor20241022,
Self::TextEditor20250124,
Self::ThinkingBindingControls20260801,
Self::ThinkingDisplayUpdates20260818,
Self::TokenEfficientTools20250219,
Self::ToolExamples20251029,
Self::ToolSearchTool20251019,
Self::InlineTools20260915,
];
pub fn as_str(&self) -> &str {
self.as_ref()
}
pub fn on(&self, provider: BetaProvider) -> Option<Self> {
let resolved = match self {
Self::Other(raw) => raw.parse().unwrap(),
_ => self.clone(),
};
if resolved.rejected_by().contains(&provider) {
return None;
}
Some(match (resolved, provider) {
(
Self::AdvancedToolUse20251120,
BetaProvider::Bedrock | BetaProvider::BedrockMantle | BetaProvider::VertexAi,
) => Self::ToolSearchTool20251019,
(beta, _) => beta,
})
}
fn rejected_by(&self) -> &'static [BetaProvider] {
match self {
Self::AdvancedToolUse20251120 | Self::ContextManagement20250627 => {
&[BetaProvider::BedrockConverse]
}
Self::AdvisorTool20260301 | Self::Compact20260904 => &[
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::BedrockMantle,
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::Bash20241022
| Self::Bash20250124
| Self::McpServers20251204
| Self::StructuredOutput20240301
| Self::TextEditor20241022
| Self::TextEditor20250124 => BetaProvider::VARIANTS,
Self::ClaudeCode20250219 | Self::ToolExamples20251029 => &[
BetaProvider::Anthropic,
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::CodeExecution20250825
| Self::FilesApi20250414
| Self::McpClient20250404
| Self::McpClient20251120
| Self::PromptCachingScope20260105
| Self::Skills20251002
| Self::WebFetch20250910 => &[
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::BedrockMantle,
BetaProvider::VertexAi,
],
Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse],
Self::ComputerUse20250124 | Self::ComputerUse20251124 | Self::Context1m20250807 => &[],
Self::DangerousToolUse20260903 => {
&[BetaProvider::BedrockConverse, BetaProvider::Databricks]
}
Self::Effort20251124 => &[BetaProvider::VertexAi],
Self::FastMode20260201 | Self::Oauth20250420 => &[
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::BedrockMantle,
BetaProvider::VertexAi,
],
Self::FineGrainedToolStreaming20250514 => {
&[BetaProvider::AzureAi, BetaProvider::VertexAi]
}
Self::InterleavedThinking20250514 | Self::WebSearch20250305 => {
&[BetaProvider::BedrockConverse, BetaProvider::Bedrock]
}
Self::MidConversationOutputConfig20260701
| Self::MidConversationToolChanges20260701 => &[
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::BedrockMantle,
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::Output128k20250219 | Self::TokenEfficientTools20250219 => &[
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::VertexAi,
],
Self::PerTurnControl20260701 => &[
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::Databricks,
],
Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock, BetaProvider::VertexAi],
Self::ThinkingBindingControls20260801 => &[BetaProvider::AzureAi],
Self::ThinkingDisplayUpdates20260818 => &[BetaProvider::Databricks],
Self::ToolSearchTool20251019 => &[
BetaProvider::Anthropic,
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::Databricks,
],
Self::InlineTools20260915 => &[
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
BetaProvider::Bedrock,
BetaProvider::BedrockMantle,
BetaProvider::Databricks,
],
Self::Other(_) => BetaProvider::VARIANTS,
}
}
}
impl PartialEq for AnthropicBeta {
@ -149,10 +380,20 @@ impl fmt::Display for BetaSet {
#[cfg(test)]
mod tests {
use std::collections::BTreeSet;
use indexmap::IndexMap;
use rstest::rstest;
use super::*;
fn beta_headers_config() -> IndexMap<String, serde_json::Value> {
serde_json::from_str(include_str!(
"../../../../../litellm/anthropic_beta_headers_config.json"
))
.unwrap()
}
fn set(header: &str) -> BetaSet {
header.parse().unwrap_or_else(|never| match never {})
}
@ -171,7 +412,35 @@ mod tests {
AnthropicBeta::FastMode20260201,
AnthropicBeta::AdvisorTool20260301,
AnthropicBeta::PerTurnControl20260701,
AnthropicBeta::DangerousToolUse20260903
AnthropicBeta::DangerousToolUse20260903,
AnthropicBeta::Bash20241022,
AnthropicBeta::Bash20250124,
AnthropicBeta::ClaudeCode20250219,
AnthropicBeta::CodeExecution20250825,
AnthropicBeta::ComputerUse20250124,
AnthropicBeta::ComputerUse20251124,
AnthropicBeta::Context1m20250807,
AnthropicBeta::Effort20251124,
AnthropicBeta::FilesApi20250414,
AnthropicBeta::FineGrainedToolStreaming20250514,
AnthropicBeta::InterleavedThinking20250514,
AnthropicBeta::McpClient20250404,
AnthropicBeta::McpClient20251120,
AnthropicBeta::McpServers20251204,
AnthropicBeta::MidConversationOutputConfig20260701,
AnthropicBeta::MidConversationToolChanges20260701,
AnthropicBeta::Output128k20250219,
AnthropicBeta::PromptCachingScope20260105,
AnthropicBeta::Skills20251002,
AnthropicBeta::StructuredOutput20240301,
AnthropicBeta::TextEditor20241022,
AnthropicBeta::TextEditor20250124,
AnthropicBeta::ThinkingBindingControls20260801,
AnthropicBeta::ThinkingDisplayUpdates20260818,
AnthropicBeta::TokenEfficientTools20250219,
AnthropicBeta::ToolExamples20251029,
AnthropicBeta::ToolSearchTool20251019,
AnthropicBeta::InlineTools20260915
)]
beta: AnthropicBeta,
) {
@ -181,6 +450,119 @@ mod tests {
assert!(AnthropicBeta::KNOWN.contains(&beta));
}
#[rstest]
fn known_betas_are_exactly_the_config_keys() {
let config = beta_headers_config();
let config_keys: BTreeSet<String> = config
.iter()
.filter(|(provider, _)| provider.as_str() != "description")
.flat_map(|(_, column)| {
column
.as_object()
.into_iter()
.flat_map(|values| values.keys().cloned())
})
.collect();
let known_keys: BTreeSet<String> = AnthropicBeta::KNOWN
.iter()
.map(|beta| beta.as_str().to_string())
.collect();
assert_eq!(known_keys, config_keys);
}
#[rstest]
#[case::anthropic(BetaProvider::Anthropic)]
#[case::azure_ai(BetaProvider::AzureAi)]
#[case::bedrock_converse(BetaProvider::BedrockConverse)]
#[case::bedrock(BetaProvider::Bedrock)]
#[case::bedrock_mantle(BetaProvider::BedrockMantle)]
#[case::vertex_ai(BetaProvider::VertexAi)]
#[case::databricks(BetaProvider::Databricks)]
fn provider_columns_match_beta_provider_variants(#[case] provider: BetaProvider) {
let config = beta_headers_config();
let config_columns: Vec<String> = config
.keys()
.filter(|column| column.as_str() != "description")
.cloned()
.collect();
let beta_provider_columns: Vec<String> = BetaProvider::VARIANTS
.iter()
.map(ToString::to_string)
.collect();
assert_eq!(config_columns, beta_provider_columns);
assert_eq!(
provider.to_string().parse::<BetaProvider>().unwrap(),
provider
);
}
#[rstest]
fn on_matches_every_config_cell() {
let config = beta_headers_config();
for provider in BetaProvider::VARIANTS {
for beta in &AnthropicBeta::KNOWN {
let expected = config
.get(provider.as_ref())
.and_then(serde_json::Value::as_object)
.and_then(|column| column.get(beta.as_str()))
.and_then(serde_json::Value::as_str)
.map(str::to_string);
let actual = beta.on(*provider).map(|name| name.to_string());
assert_eq!(actual, expected, "provider {provider}, beta {beta}");
if matches!(beta, AnthropicBeta::AdvancedToolUse20251120)
&& matches!(
provider,
BetaProvider::Bedrock
| BetaProvider::BedrockMantle
| BetaProvider::VertexAi
)
{
let parsed: AnthropicBeta = actual.as_deref().unwrap().parse().unwrap();
assert!(AnthropicBeta::KNOWN.contains(&parsed));
assert!(!matches!(parsed, AnthropicBeta::Other(_)));
}
}
}
}
#[rstest]
#[case::renamed(
"advanced-tool-use-2025-11-20",
BetaProvider::Bedrock,
Some(AnthropicBeta::ToolSearchTool20251019)
)]
#[case::kept(
"oauth-2025-04-20",
BetaProvider::Anthropic,
Some(AnthropicBeta::Oauth20250420)
)]
#[case::rejected("effort-2025-11-24", BetaProvider::VertexAi, None)]
#[case::unknown("example-beta-2099-01-01", BetaProvider::Anthropic, None)]
fn on_resolves_known_spellings_held_as_other(
#[case] raw: &str,
#[case] provider: BetaProvider,
#[case] expected: Option<AnthropicBeta>,
) {
let actual = AnthropicBeta::Other(raw.to_string()).on(provider);
assert_eq!(actual, expected);
assert!(!matches!(actual, Some(AnthropicBeta::Other(_))));
}
#[rstest]
#[case::anthropic(BetaProvider::Anthropic)]
#[case::azure_ai(BetaProvider::AzureAi)]
#[case::bedrock_converse(BetaProvider::BedrockConverse)]
#[case::bedrock(BetaProvider::Bedrock)]
#[case::bedrock_mantle(BetaProvider::BedrockMantle)]
#[case::vertex_ai(BetaProvider::VertexAi)]
#[case::databricks(BetaProvider::Databricks)]
fn unknown_beta_is_rejected_by_every_provider(#[case] provider: BetaProvider) {
assert_eq!(
AnthropicBeta::Other("example-beta-2099-01-01".into()).on(provider),
None
);
}
#[test]
fn unknown_values_are_kept_verbatim() {
let parsed: AnthropicBeta = "claude-code-20250219".parse().unwrap();

View file

@ -0,0 +1,36 @@
use litellm_http::request::{with_header, without_headers};
use litellm_llms_types::providers::anthropic::{BetaProvider, BetaSet};
use crate::{anthropic::common_utils::existing_betas, base_llm::auth::Headers};
const BETA_HEADER: &str = "anthropic-beta";
/// How the `anthropic-beta` header is treated before a request leaves for a host.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BetaPolicy {
/// Send the header untouched and let the host decide.
Forward,
/// Keep only the betas the provider accepts, under the names it expects.
Filter(BetaProvider),
/// Remove every `anthropic-beta` header.
Drop,
}
impl BetaPolicy {
pub fn apply(self, headers: Headers) -> Headers {
let provider = match self {
Self::Forward => return headers,
Self::Drop => return without_headers(headers, &[BETA_HEADER]),
Self::Filter(provider) => provider,
};
let accepted: BetaSet = existing_betas(&headers)
.iter()
.filter_map(|beta| beta.on(provider))
.collect();
let headers = without_headers(headers, &[BETA_HEADER]);
if accepted.is_empty() {
return headers;
}
with_header(headers, BETA_HEADER, accepted.to_string())
}
}

View file

@ -241,7 +241,7 @@ fn anthropic_body(
.iter()
.map(|turn| {
json!({
"role": turn.role.as_str(),
"role": <&'static str>::from(turn.role),
"content": turn.texts.iter().map(|text| text_block(text)).collect::<Vec<_>>(),
})
})

View file

@ -596,17 +596,10 @@ mod tests {
use crate::base_llm::messages::context::SupportedEffortTiers;
use rstest::{fixture, rstest};
use serde_json::json;
use strum::VariantArray;
use super::*;
const ALL_LEVELS: [EffortLevel; 5] = [
EffortLevel::Low,
EffortLevel::Medium,
EffortLevel::High,
EffortLevel::Xhigh,
EffortLevel::Max,
];
fn apply(sanitizer: fn(Vec<Message>) -> Vec<Message>, messages: Value) -> Value {
let parsed: Vec<Message> = serde_json::from_value(messages).unwrap();
serde_json::to_value(sanitizer(parsed)).unwrap()
@ -1712,7 +1705,10 @@ mod tests {
..unmapped
};
assert_eq!(
ALL_LEVELS.map(|level| supports_effort_tier(&capabilities, level)),
EffortLevel::VARIANTS
.iter()
.map(|level| supports_effort_tier(&capabilities, *level))
.collect::<Vec<_>>(),
expected
);
}

View file

@ -10,6 +10,7 @@ use litellm_llms_types::{
};
use litellm_python_compat::{json::from_json, repr::repr, truthy::truthy};
use serde_json::Value;
use strum::VariantArray;
use crate::base_llm::messages::context::{
MessagesModelCapabilities, ThinkingBudgets, ThinkingContext,
@ -54,14 +55,17 @@ fn unmapped_effort(effort: &Value) -> Error {
Error::InvalidRequest(crate::ErrorDetail::InvalidChoice {
field: "reasoning effort",
actual: repr(&from_json(effort.clone())),
choices: ReasoningEffort::ALL.map(|effort| effort.as_str()).into(),
choices: ReasoningEffort::VARIANTS
.iter()
.map(|effort| <&'static str>::from(*effort))
.collect(),
})
}
fn unsupported_effort(level: EffortLevel, model: &str) -> Error {
Error::InvalidRequest(crate::ErrorDetail::UnsupportedValue {
field: "effort",
value: level.as_str(),
value: <&'static str>::from(level),
model: model.into(),
})
}
@ -134,7 +138,7 @@ fn legacy_reasoning_effort(
Some(Recognized::Known(level)) => Ok((*level).into()),
Some(Recognized::Unrecognized(value)) if truthy(&from_json(value.clone())) => value
.as_str()
.and_then(ReasoningEffort::parse)
.and_then(|text| text.parse().ok())
.ok_or_else(|| unmapped_effort(value)),
None | Some(Recognized::Unrecognized(_)) => Ok(ReasoningEffort::Medium),
}

View file

@ -270,7 +270,7 @@ fn drop_unsupported_params(
fn speed_text(speed: &Recognized<Speed>) -> String {
match speed {
Recognized::Known(speed) => speed.as_str().to_string(),
Recognized::Known(speed) => <&'static str>::from(*speed).to_string(),
Recognized::Unrecognized(Value::String(text)) => text.clone(),
Recognized::Unrecognized(other) => other.to_string(),
}

View file

@ -1,3 +1,4 @@
pub mod beta_headers;
pub mod common_utils;
pub mod batches;

View file

@ -23,9 +23,11 @@ const HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAA
/// Textract has operations rather than models; the model slot of
/// `aws_textract/<model>` names the one to call.
#[derive(Clone, Copy, Debug, EnumString, IntoStaticStr, VariantNames, PartialEq, Eq)]
#[strum(serialize_all = "kebab-case", ascii_case_insensitive)]
#[strum(ascii_case_insensitive)]
pub enum TextractOperation {
#[strum(serialize = "detect-document-text")]
DetectDocumentText,
#[strum(serialize = "analyze-document")]
AnalyzeDocument,
}

View file

@ -30,20 +30,17 @@ pub struct BedrockAudioTranscriptionConfig;
#[derive(Clone, Copy, Deserialize, IntoStaticStr)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
enum AudioFormat {
#[strum(serialize = "wav")]
Wav,
#[strum(serialize = "mp3")]
Mp3,
#[strum(serialize = "flac")]
Flac,
#[strum(serialize = "ogg")]
Ogg,
}
impl AudioFormat {
fn as_str(self) -> &'static str {
self.into()
}
}
struct AudioInput {
data: String,
format: AudioFormat,
@ -118,7 +115,7 @@ impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig {
"messages": [{
"role": "user",
"content": [
{"audio": {"format": audio.format.as_str(), "source": {"bytes": audio.data}}},
{"audio": {"format": <&'static str>::from(audio.format), "source": {"bytes": audio.data}}},
{"text": instruction}
]
}],

View file

@ -385,7 +385,7 @@ fn converse_body(conversation: &Conversation, optional_params: &Map<String, Valu
.iter()
.map(|turn| {
json!({
"role": turn.role.as_str(),
"role": <&'static str>::from(turn.role),
"content": turn.texts.iter().map(|text| json!({"text": text})).collect::<Vec<_>>(),
})
})

View file

@ -64,27 +64,23 @@ pub fn invoke_anthropic_event_stream(bytes: ByteStream) -> EventStream {
Box::pin(invoke_chunk_stream(bytes).map(|chunk| decode_invoke_anthropic_chunk(chunk?)))
}
#[derive(Clone, Copy)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString)]
enum InvokeProvider {
#[strum(serialize = "anthropic")]
Anthropic,
#[strum(serialize = "deepseek_r1")]
DeepseekR1,
#[strum(serialize = "moonshot")]
Moonshot,
#[strum(disabled)]
Unsupported,
}
impl From<&str> for InvokeProvider {
fn from(value: &str) -> Self {
match value {
"anthropic" => Self::Anthropic,
"deepseek_r1" => Self::DeepseekR1,
"moonshot" => Self::Moonshot,
_ => Self::Unsupported,
}
}
}
pub fn invoke_chat_stream(invoke_provider: &str, shape: StreamShape) -> Result<ChatStream, Error> {
match InvokeProvider::from(invoke_provider) {
match invoke_provider
.parse()
.unwrap_or(InvokeProvider::Unsupported)
{
InvokeProvider::Anthropic => Ok(ChatStream::new(
invoke_anthropic_event_stream,
ModelResponseIterator::new(shape),
@ -116,6 +112,22 @@ mod tests {
const TEXT_DELTA: &str =
r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#;
#[rstest::rstest]
#[case::anthropic("anthropic", InvokeProvider::Anthropic)]
#[case::deepseek_r1("deepseek_r1", InvokeProvider::DeepseekR1)]
#[case::moonshot("moonshot", InvokeProvider::Moonshot)]
#[case::unknown("qwen", InvokeProvider::Unsupported)]
#[case::case_sensitive("Anthropic", InvokeProvider::Unsupported)]
fn invoke_provider_maps_each_model_family_name(
#[case] name: &str,
#[case] expected: InvokeProvider,
) {
assert_eq!(
name.parse().unwrap_or(InvokeProvider::Unsupported),
expected
);
}
fn aws_wire(chunk: &str) -> Vec<u8> {
let payload = serde_json::json!({"bytes": STANDARD.encode(chunk)});
let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header(

View file

@ -25,11 +25,13 @@ const COHERE_API_KEY_ENV: &str = "COHERE_API_KEY";
const COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC";
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, strum::IntoStaticStr)]
#[serde(rename_all = "lowercase")]
pub enum OutputFormat {
#[default]
#[strum(serialize = "markdown")]
Markdown,
#[strum(serialize = "blocks")]
Blocks,
}
@ -267,11 +269,7 @@ fn build_request(model: &str, image_url: String, params: &CohereOptions) -> Cohe
CohereRequest {
model: model.into(),
document: CohereParseDocument::ImageUrl { image_url },
output_format: match params.output_format.unwrap_or_default() {
OutputFormat::Markdown => "markdown",
OutputFormat::Blocks => "blocks",
}
.into(),
output_format: <&'static str>::from(params.output_format.unwrap_or_default()).into(),
}
}

View file

@ -0,0 +1,111 @@
use litellm_llms::anthropic::beta_headers::BetaPolicy;
use litellm_llms_types::providers::anthropic::BetaProvider;
use rstest::rstest;
fn headers(pairs: &[(&str, &str)]) -> Vec<(String, String)> {
pairs
.iter()
.map(|(name, value)| (name.to_string(), value.to_string()))
.collect()
}
#[rstest]
fn forward_keeps_headers_as_is() {
let input = headers(&[
(
"Anthropic-Beta",
"example-beta-2099-01-01,web-search-2025-03-05",
),
("x-api-key", "k"),
]);
assert_eq!(BetaPolicy::Forward.apply(input.clone()), input);
}
#[rstest]
fn drop_removes_every_casing() {
assert_eq!(
BetaPolicy::Drop.apply(headers(&[
("anthropic-beta", "a"),
("x-api-key", "k"),
("ANTHROPIC-BETA", "b"),
])),
headers(&[("x-api-key", "k")])
);
}
#[rstest]
#[case::bedrock(BetaProvider::Bedrock, Some("tool-search-tool-2025-10-19"))]
#[case::bedrock_mantle(BetaProvider::BedrockMantle, Some("tool-search-tool-2025-10-19"))]
#[case::vertex_ai(BetaProvider::VertexAi, Some("tool-search-tool-2025-10-19"))]
#[case::anthropic(BetaProvider::Anthropic, Some("advanced-tool-use-2025-11-20"))]
#[case::azure_ai(BetaProvider::AzureAi, Some("advanced-tool-use-2025-11-20"))]
#[case::databricks(BetaProvider::Databricks, Some("advanced-tool-use-2025-11-20"))]
#[case::bedrock_converse(BetaProvider::BedrockConverse, None)]
fn filter_renames_per_host(#[case] provider: BetaProvider, #[case] expected_beta: Option<&str>) {
let expected = match expected_beta {
Some(beta) => headers(&[("x-api-key", "k"), ("anthropic-beta", beta)]),
None => headers(&[("x-api-key", "k")]),
};
assert_eq!(
BetaPolicy::Filter(provider).apply(headers(&[
("Anthropic-Beta", "advanced-tool-use-2025-11-20"),
("x-api-key", "k"),
])),
expected
);
}
#[rstest]
#[case::azure_ai(BetaProvider::AzureAi, "fast-mode-2026-02-01,example-beta-2099-01-01")]
#[case::anthropic(BetaProvider::Anthropic, "bash_20241022")]
#[case::vertex_ai(BetaProvider::VertexAi, " , ")]
fn filter_removes_header_when_nothing_survives(
#[case] provider: BetaProvider,
#[case] beta_value: &str,
) {
assert_eq!(
BetaPolicy::Filter(provider).apply(headers(&[
("x-api-key", "k"),
("anthropic-beta", beta_value)
])),
headers(&[("x-api-key", "k")])
);
}
#[rstest]
#[case::anthropic(
BetaProvider::Anthropic,
"web-search-2025-03-05",
"oauth-2025-04-20,web-search-2025-03-05,example-beta-2099-01-01",
"oauth-2025-04-20,web-search-2025-03-05"
)]
#[case::vertex_ai_renames_and_deduplicates(
BetaProvider::VertexAi,
"advanced-tool-use-2025-11-20",
"tool-search-tool-2025-10-19",
"tool-search-tool-2025-10-19"
)]
fn filter_deduplicates_and_sorts(
#[case] provider: BetaProvider,
#[case] first_beta: &str,
#[case] second_beta: &str,
#[case] expected_beta: &str,
) {
assert_eq!(
BetaPolicy::Filter(provider).apply(headers(&[
("ANTHROPIC-BETA", first_beta),
("x-api-key", "k"),
("anthropic-beta", second_beta),
])),
headers(&[("x-api-key", "k"), ("anthropic-beta", expected_beta)])
);
}
#[rstest]
#[case::forward(BetaPolicy::Forward)]
#[case::drop(BetaPolicy::Drop)]
#[case::filter(BetaPolicy::Filter(BetaProvider::Bedrock))]
fn headers_without_a_beta_are_untouched(#[case] policy: BetaPolicy) {
let input = headers(&[("x-api-key", "k"), ("anthropic-version", "2023")]);
assert_eq!(policy.apply(input.clone()), input);
}

View file

@ -17,6 +17,10 @@ mod tokenizer;
#[pymodule(gil_used = true)]
mod _native {
#[pymodule_export]
use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
use pyo3::{prelude::*, types::PyModule};
#[cfg(feature = "panic-test")]
#[pymodule_export]
use crate::diagnostics::_panic_for_test;
@ -29,9 +33,7 @@ mod _native {
#[pymodule_export]
use crate::routes::audio_transcription::{atranscription, transcription};
#[pymodule_export]
use crate::routes::chat_completions::{
achat_completions, acompletion, chat_completions, completion,
};
use crate::routes::chat_completions::{acompletion, completion};
#[pymodule_export]
use crate::routes::embeddings::{aembedding, embedding};
#[pymodule_export]
@ -51,9 +53,6 @@ mod _native {
use crate::tokenizer::HuggingFaceEncoding;
#[pymodule_export]
use crate::tokenizer::Tokenizer;
#[pymodule_export]
use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
use pyo3::{prelude::*, types::PyModule};
#[pymodule_init]
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
@ -101,8 +100,6 @@ mod tests {
"atranscription",
"messages",
"amessages",
"chat_completions",
"achat_completions",
"completion",
"acompletion",
"responses",

View file

@ -19,13 +19,6 @@ pub(crate) struct RouteOptions {
pub(crate) timeout: Option<Duration>,
}
pub(crate) fn messages_argument(value: &Bound<'_, PyAny>) -> PyResult<Vec<Value>> {
match from_py_argument(value)? {
Value::Array(values) => Ok(values),
_ => Err(PyValueError::new_err("messages must be a list")),
}
}
fn required_object(name: &'static str, value: Value) -> PyResult<Map<String, Value>> {
match value {
Value::Object(values) => Ok(values),
@ -449,18 +442,6 @@ module.__getattr__ = fail
fn argument_converters_keep_nested_values_and_accept_explicit_none() {
Python::initialize();
Python::attach(|py| {
let messages = py
.eval(
c"[{'role': 'user', 'content': [{'type': 'text', 'text': 'hi'}]}]",
None,
None,
)
.unwrap();
assert_eq!(
Value::Array(messages_argument(&messages).unwrap()),
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}])
);
let params = py.eval(c"{'temperature': 0.2}", None, None).unwrap();
assert_eq!(
optional_object("optional_params", &params).unwrap(),

View file

@ -6,13 +6,16 @@ use crate::coercion::{FieldSpec, ProjectionError};
const MODULE: &str = "litellm.rust_bridge.settings";
#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq)]
#[strum(serialize_all = "snake_case")]
pub(crate) enum PythonSettings {
#[strum(serialize = "http_settings")]
Http,
#[strum(serialize = "url_policy")]
UrlPolicy,
#[strum(serialize = "provider_defaults")]
ProviderDefaults,
#[strum(serialize = "secret_manager")]
SecretManager,
#[strum(serialize = "secret_manager_binding")]
SecretManagerBinding,
}
@ -23,17 +26,16 @@ pub(crate) struct Snapshot<'py> {
impl Snapshot<'_> {
pub(crate) fn read<T>(&self, spec: &FieldSpec<T>) -> Result<T, ProjectionError> {
spec.read(&self.value, self.group.name())
spec.read(&self.value, self.group.into())
}
}
impl PythonSettings {
pub(crate) fn name(self) -> &'static str {
self.into()
}
pub(crate) fn read(self, py: Python<'_>) -> PyResult<Snapshot<'_>> {
let value = py.import(MODULE)?.getattr(self.name())?.call0()?;
let value = py
.import(MODULE)?
.getattr(<&'static str>::from(self))?
.call0()?;
Ok(Snapshot { group: self, value })
}

View file

@ -11,3 +11,7 @@ The host driver owns sequencing and terminal events; the bridge supplies fallibl
Use the shared `run_public_call` boundary with hooks supplied by bridge composition. `callbacks-legacy-python` owns legacy argument sharing and `Logging` dispatch behind `PublicCall` and `LegacyLogging`. Route bindings supply `callbacks-legacy-python::LoggingOperation` when composing legacy logging and may retain the request needed for projection, but must not duplicate the legacy callback contract
Regression tests must observe that an unstarted call does no setup, hook and preflight rewrites affect resource configuration, setup failures reach the selected failure handler once, and provider work is not replayed. Retain existing read-point and object-identity guarantees while changing setup timing
## Layout
`messages/` is the reference shape. Each lifecycle route is a folder: `mod.rs` holds `run_<route>` and the thin sync and async `#[pyfunction]` entrypoints that take `NativeCall` and call it, `host.rs` holds the `PythonBinding` host, and anything else route-specific (projection, error mapping, extra pyclasses) gets its own sibling file. Scaffolded routes without a lifecycle (embeddings, audio transcription) stay a single `<route>.rs`

View file

@ -5,6 +5,8 @@ use litellm_inference_transcription::{
use pyo3::prelude::*;
use serde_json::{Map, Value};
use super::NativeCall;
use crate::{
errors::route_error_to_pyerr,
marshal::{RouteOptions, optional_object_field, required_field, value_route_options},
@ -40,8 +42,7 @@ async fn execute(
}
#[pyfunction]
pub(crate) fn transcription(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
pub(crate) fn transcription(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
let audio: Value =
litellm_host_python::from_py_argument(&required_field(&call.bound, "audio")?)?;
let options = value_route_options(&call.bound)?;
@ -59,9 +60,8 @@ pub(crate) fn transcription(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<
#[pyfunction]
pub(crate) fn atranscription<'py>(
py: Python<'py>,
call: Bound<'py, PyAny>,
call: NativeCall<'py>,
) -> PyResult<Bound<'py, PyAny>> {
let call = super::NativeCall::extract(&call)?;
let audio: Value =
litellm_host_python::from_py_argument(&required_field(&call.bound, "audio")?)?;
let options = value_route_options(&call.bound)?;

View file

@ -1,150 +0,0 @@
mod host;
use pyo3::types::{PyDict, PyTuple};
use crate::execution::{run_async, run_sync};
use litellm_inference_chat::{ChatCompletionsRoute, Error, types::ChatCompletionsRequest};
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
use pyo3::prelude::*;
use serde_json::{Map, Value};
use crate::{
errors::route_error_to_pyerr,
marshal::{
RouteOptions, messages_argument, optional_object_field, required_field, value_route_options,
},
};
async fn execute(
http: Result<litellm_http::Client, litellm_http::Error>,
secrets: std::sync::Arc<dyn litellm_secrets::source::SecretSource>,
messages: Vec<Value>,
optional_params: Map<String, Value>,
options: RouteOptions,
) -> Result<ChatCompletionsResponse, Error> {
let RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout,
} = options;
ChatCompletionsRoute::new(http?, crate::http::resources().auth.clone(), secrets)
.execute(
ChatCompletionsRequest {
model: &model,
messages: Value::Array(messages),
optional_params,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
timeout,
},
&(),
None,
)
.await
}
#[pyfunction]
pub(crate) fn chat_completions(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
let messages: Vec<Value> = messages_argument(&required_field(&call.bound, "messages")?)?;
let optional_params =
optional_object_field(&call.bound, "optional_params")?.unwrap_or_default();
let options = value_route_options(&call.bound)?;
let http = crate::http::provider_client(py, &call.kwargs, false)?;
let secrets = crate::secrets::source(py)?;
run_sync(
py,
execute(http, secrets, messages, optional_params, options),
route_error_to_pyerr,
)
}
#[pyfunction]
pub(crate) fn achat_completions<'py>(
py: Python<'py>,
call: Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyAny>> {
let call = super::NativeCall::extract(&call)?;
let messages: Vec<Value> = messages_argument(&required_field(&call.bound, "messages")?)?;
let optional_params =
optional_object_field(&call.bound, "optional_params")?.unwrap_or_default();
let options = value_route_options(&call.bound)?;
let http = crate::http::provider_client(py, &call.kwargs, true)?;
let secrets = crate::secrets::source(py)?;
run_async(
py,
execute(http, secrets, messages, optional_params, options),
route_error_to_pyerr,
)
}
fn run_public(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
use super::inference::InferenceHost;
use litellm_callbacks_legacy_python::LoggingOperation;
let host = InferenceHost::new(
request.clone().unbind(),
"litellm.rust_bridge.chat_completions.route_host",
);
let cache_call_type = if asynchronous {
"acompletion"
} else {
"completion"
};
crate::cache::admit_native(py, &kwargs, cache_call_type)?;
let (arguments, hooks) = crate::routes::call_hooks(
py,
LoggingOperation::Completion,
&request,
&args,
&kwargs,
asynchronous,
)?;
crate::routes::run_public_call(
py,
arguments,
move |py, arguments, request| {
let route = ChatCompletionsRoute::new(
crate::http::provider_client(py, arguments, asynchronous)?
.map_err(crate::http::client_error)?,
crate::http::resources().auth.clone(),
crate::secrets::source(py)?,
);
let (cache, cache_options) =
crate::cache::configured_native(py, arguments, cache_call_type)?;
let route = match cache {
Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new(
cache,
litellm_cache_response::CacheScope::Shared,
)),
None => route,
};
Ok(route.machine(request, cache_options.policy))
},
host::ChatCompletionsPythonHost(host),
hooks,
asynchronous,
)
}
#[pyfunction]
pub(crate) fn completion(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_public(py, call.bound.into_any(), call.args, call.kwargs, false)
}
#[pyfunction]
pub(crate) fn acompletion(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_public(py, call.bound.into_any(), call.args, call.kwargs, true)
}

View file

@ -1,6 +1,6 @@
use std::convert::Infallible;
use super::super::inference::InferenceHost;
use crate::routes::inference::InferenceHost;
use litellm_host_python::{InvokeError, PythonBinding, PythonHostCalls, PythonOwned};
use litellm_inference_chat::{Error, route::ChatCompletions, types::ChatCompletionsCall};
use pyo3::{

View file

@ -0,0 +1,71 @@
mod host;
use std::sync::Arc;
use host::ChatCompletionsPythonHost;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, ScopedCache};
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_inference_chat::{ChatCompletionsRoute, route::ChatCompletions};
use litellm_secrets::source::SecretSource;
use pyo3::prelude::*;
use super::{
NativeCall,
inference::{InferenceHost, InferenceRoute, run_inference},
};
fn run_chat_completions(
py: Python<'_>,
call: NativeCall<'_>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
let host = InferenceHost::new(
call.bound.clone().unbind(),
"litellm.rust_bridge.chat_completions.route_host",
);
run_inference::<ChatCompletionsRoute, _>(
py,
call,
asynchronous,
ChatCompletionsPythonHost(host),
)
}
#[pyfunction]
pub(crate) fn completion(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
run_chat_completions(py, call, false)
}
#[pyfunction]
pub(crate) fn acompletion(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
run_chat_completions(py, call, true)
}
impl InferenceRoute for ChatCompletionsRoute {
type Protocol = ChatCompletions;
const OPERATION: LoggingOperation = LoggingOperation::Completion;
const SYNC_CALL_TYPE: &'static str = "completion";
const ASYNC_CALL_TYPE: &'static str = "acompletion";
fn new(
http: litellm_http::Client,
auth: Arc<AuthServices>,
secrets: Arc<dyn SecretSource>,
) -> Self {
Self::new(http, auth, secrets)
}
fn with_cache(self, cache: ScopedCache) -> Self {
self.with_cache(cache)
}
fn machine(
self,
call: <ChatCompletions as Protocol>::Request,
policy: CachePolicy,
) -> HostedMachine<ChatCompletions> {
self.machine(call, policy)
}
}

View file

@ -1,17 +1,18 @@
use pyo3::prelude::*;
use super::NativeCall;
use crate::errors::RustBridgeDeclined;
#[pyfunction]
pub(crate) fn embedding(call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
drop(super::NativeCall::extract(&call)?);
pub(crate) fn embedding(call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
drop(call);
Err(RustBridgeDeclined::new_err(
"native embeddings route is not implemented",
))
}
#[pyfunction]
pub(crate) fn aembedding(call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
pub(crate) fn aembedding(call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
embedding(call)
}
@ -36,7 +37,8 @@ call = SimpleNamespace(args=(), kwargs={}, bound={'model':'test-model','input':'
Some(&locals),
)
.unwrap();
let call = locals.get_item("call").unwrap().unwrap();
let call: super::NativeCall<'_> =
locals.get_item("call").unwrap().unwrap().extract().unwrap();
let error = if asynchronous {
super::aembedding(call)
} else {

View file

@ -1,11 +1,19 @@
use std::sync::Arc;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache};
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider;
use litellm_host_python::{from_py, lookup};
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_host_python::{PythonBinding, PythonHostCalls, from_py, present};
use litellm_http::transport::Error as TransportError;
use litellm_inference::RouteError;
use litellm_secrets::source::SecretSource;
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
use serde::Serialize;
use serde_json::{Map, Value};
use super::NativeCall;
use crate::{
errors::{RustUpstreamError, route_error_to_pyerr},
marshal::{
@ -15,7 +23,7 @@ use crate::{
};
pub(super) struct InferenceHost {
pub request: Py<PyAny>,
pub request: Py<PyDict>,
module: &'static str,
}
@ -26,7 +34,7 @@ pub(super) struct ProjectedCall {
}
impl InferenceHost {
pub fn new(request: Py<PyAny>, module: &'static str) -> Self {
pub fn new(request: Py<PyDict>, module: &'static str) -> Self {
Self { request, module }
}
@ -85,21 +93,7 @@ impl InferenceHost {
arguments: &Bound<'py, PyDict>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
let request = self.request.bind(py);
if let Some(value) = lookup(arguments, request, name)? {
return Ok((!value.is_none()).then_some(value));
}
if request.is_instance_of::<PyDict>() {
return Ok(None);
}
let parameter = request
.getattr("parameters")?
.call_method1("get", (name,))?;
if !parameter.is_none() {
return Ok(Some(parameter));
}
let extra = request.getattr("kwargs")?.call_method1("get", (name,))?;
Ok((!extra.is_none()).then_some(extra))
present(arguments, self.request.bind(py), name)
}
pub fn parameters(
@ -140,3 +134,62 @@ impl InferenceHost {
Ok(PyErr::from_value(mapped))
}
}
pub(super) trait InferenceRoute: Sized + 'static {
type Protocol: Protocol<Error = RouteError>;
const OPERATION: LoggingOperation;
const SYNC_CALL_TYPE: &'static str;
const ASYNC_CALL_TYPE: &'static str;
fn new(
http: litellm_http::Client,
auth: Arc<AuthServices>,
secrets: Arc<dyn SecretSource>,
) -> Self;
fn with_cache(self, cache: ScopedCache) -> Self;
fn machine(
self,
call: <Self::Protocol as Protocol>::Request,
policy: CachePolicy,
) -> HostedMachine<Self::Protocol>;
}
pub(super) fn run_inference<R, H>(
py: Python<'_>,
call: NativeCall<'_>,
asynchronous: bool,
host: H,
) -> PyResult<Py<PyAny>>
where
R: InferenceRoute,
H: PythonBinding<Protocol = R::Protocol> + PythonHostCalls<R::Protocol> + 'static,
{
let call_type = if asynchronous {
R::ASYNC_CALL_TYPE
} else {
R::SYNC_CALL_TYPE
};
crate::cache::admit_native(py, &call.kwargs, call_type)?;
let (arguments, hooks) = super::call_hooks(py, R::OPERATION, &call, asynchronous)?;
super::run_public_call(
py,
arguments,
move |py, arguments, request| {
let route = R::new(
crate::http::provider_client(py, arguments, asynchronous)?
.map_err(crate::http::client_error)?,
crate::http::resources().auth.clone(),
crate::secrets::source(py)?,
);
let (cache, cache_options) = crate::cache::configured_native(py, arguments, call_type)?;
let route = match cache {
Some(cache) => route.with_cache(ScopedCache::new(cache, CacheScope::Shared)),
None => route,
};
Ok(route.machine(request, cache_options.policy))
},
host,
hooks,
asynchronous,
)
}

View file

@ -2,7 +2,7 @@ use crate::cache::{CacheCall, Cached, PythonCache, Selection};
use litellm_host_python::{PythonHostCalls, PythonOwned};
use bytes::Bytes;
use litellm_host_python::{InvokeError, PythonBinding, from_py, lookup, to_py};
use litellm_host_python::{InvokeError, PythonBinding, from_py, present, to_py};
use litellm_http::transport::Error as TransportError;
use litellm_inference_messages::{
Error, MessagesCall, MessagesSettings, MessagesShaping, messages_body,
@ -89,12 +89,12 @@ fn native_error(py: Python<'_>, error: Error) -> PyResult<PyErr> {
/// The Python side of the Messages route: projects the prepared arguments and builds the
/// public response, chunks and exceptions.
pub(super) struct MessagesPythonHost {
request: Py<PyAny>,
request: Py<PyDict>,
cache: PythonCache,
}
impl MessagesPythonHost {
pub(super) fn new(request: Py<PyAny>, asynchronous: bool) -> Self {
pub(super) fn new(request: Py<PyDict>, asynchronous: bool) -> Self {
Self {
request,
cache: PythonCache::new(asynchronous),
@ -107,9 +107,7 @@ impl MessagesPythonHost {
arguments: &Bound<'_, PyDict>,
) -> PyResult<Result<MessagesCall, Error>> {
let request = self.request.bind(py);
let argument = |name: &str| -> PyResult<Option<Bound<'_, PyAny>>> {
Ok(lookup(arguments, request, name)?.filter(|value| !value.is_none()))
};
let argument = |name: &str| present(arguments, request, name);
let string = |name: &str| -> PyResult<Option<String>> {
argument(name)?.map(|value| value.extract()).transpose()
};
@ -153,8 +151,7 @@ impl MessagesPythonHost {
) -> PyResult<Option<Map<String, Value>>> {
let request = self.request.bind(py);
let mapping = |name: &str| -> PyResult<Option<Map<String, Value>>> {
lookup(arguments, request, name)?
.filter(|value| !value.is_none())
present(arguments, request, name)?
.map(|value| from_py(&value))
.transpose()
};
@ -169,8 +166,7 @@ impl MessagesPythonHost {
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> PyResult<Option<ProviderSpecificHeaders>> {
lookup(arguments, self.request.bind(py), "provider_specific_header")?
.filter(|value| !value.is_none())
present(arguments, self.request.bind(py), "provider_specific_header")?
.map(|value| from_py(&value))
.transpose()
}
@ -200,9 +196,9 @@ impl MessagesPythonHost {
self.request
.bind(py)
.get_item("custom_llm_provider")
.and_then(|value| value.extract::<Option<String>>())
.ok()
.flatten()
.and_then(|value| value.extract::<Option<String>>().ok().flatten())
.unwrap_or_else(|| "anthropic".into())
}

View file

@ -2,26 +2,13 @@ mod host;
use host::MessagesPythonHost;
use litellm_callbacks_legacy_python::LoggingOperation;
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use pyo3::prelude::*;
fn run_messages(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
let (arguments, hooks) = crate::routes::call_hooks(
py,
LoggingOperation::Messages,
&request,
&args,
&kwargs,
asynchronous,
)?;
use super::NativeCall;
fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
let (arguments, hooks) =
crate::routes::call_hooks(py, LoggingOperation::Messages, &call, asynchronous)?;
crate::routes::run_public_call(
py,
arguments,
@ -57,20 +44,18 @@ fn run_messages(
},
))
},
MessagesPythonHost::new(request.unbind(), asynchronous),
MessagesPythonHost::new(call.bound.unbind(), asynchronous),
hooks,
asynchronous,
)
}
#[pyfunction]
pub(crate) fn messages(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_messages(py, call.bound.into_any(), call.args, call.kwargs, false)
pub(crate) fn messages(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
run_messages(py, call, false)
}
#[pyfunction]
pub(crate) fn amessages(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_messages(py, call.bound.into_any(), call.args, call.kwargs, true)
pub(crate) fn amessages(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
run_messages(py, call, true)
}

View file

@ -8,8 +8,7 @@ pub(crate) mod responses;
pub(crate) mod token_counter;
pub(crate) mod traces;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_callbacks_legacy_python::{LegacyLogging, PublicCall};
use litellm_callbacks_legacy_python::{LegacyLogging, LoggingOperation, PublicCall};
use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol};
use litellm_host_python::{HookChain, PythonBinding, PythonCallHooks, PythonHostCalls};
use pyo3::{
@ -17,14 +16,16 @@ use pyo3::{
types::{PyDict, PyMapping, PyTuple},
};
struct NativeCall<'py> {
pub(crate) struct NativeCall<'py> {
args: Bound<'py, PyTuple>,
kwargs: Bound<'py, PyDict>,
bound: Bound<'py, PyDict>,
}
impl<'py> NativeCall<'py> {
fn extract(call: &Bound<'py, PyAny>) -> PyResult<Self> {
impl<'py> FromPyObject<'_, 'py> for NativeCall<'py> {
type Error = PyErr;
fn extract(call: Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {
Ok(Self {
args: call.getattr("args")?.cast_into()?,
kwargs: mapping_dict(&call.getattr("kwargs")?)?,
@ -46,12 +47,10 @@ fn mapping_dict<'py>(value: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyDict>>
fn call_hooks(
py: Python<'_>,
operation: LoggingOperation,
request: &Bound<'_, PyAny>,
args: &Bound<'_, PyTuple>,
kwargs: &Bound<'_, PyDict>,
call: &NativeCall<'_>,
asynchronous: bool,
) -> PyResult<(Py<PyDict>, impl PythonCallHooks + use<>)> {
let call = PublicCall::capture(request, args, kwargs)?;
let call = PublicCall::capture(&call.bound, &call.args, &call.kwargs)?;
let arguments = call.arguments(py);
Ok((
arguments,
@ -99,6 +98,7 @@ mod tests {
prelude::*,
types::{PyDict, PyList},
};
use rstest::rstest;
fn value_call<'py>(
py: Python<'py>,
@ -126,8 +126,10 @@ mod tests {
.unwrap()
}
#[test]
fn route_arguments_that_fail_to_convert_raise_value_error() {
#[rstest]
#[case::sync("transcription")]
#[case::asynchronous("atranscription")]
fn route_arguments_that_fail_to_convert_raise_value_error(#[case] name: &str) {
Python::initialize();
Python::attach(|py| {
let module = crate::native_module(py);
@ -151,48 +153,24 @@ value = Broken()
.expect("locals should be readable")
.expect("helper value should exist");
for name in ["chat_completions", "achat_completions"] {
let error = module
.getattr(name)
.and_then(|function| {
function.call1((value_call(py, "messages", &broken, None),))
})
.expect_err("route should reject a value it cannot convert");
let error = module
.getattr(name)
.and_then(|function| function.call1((value_call(py, "audio", &broken, None),)))
.expect_err("route should reject a value it cannot convert");
assert!(
error.is_instance_of::<pyo3::exceptions::PyValueError>(py),
"{name} surfaced {error} instead of ValueError"
);
}
assert!(
error.is_instance_of::<pyo3::exceptions::PyValueError>(py),
"{name} surfaced {error} instead of ValueError"
);
});
}
#[test]
#[rstest]
fn sync_and_async_routes_apply_the_same_input_validation() {
Python::initialize();
Python::attach(|py| {
let module = crate::native_module(py);
let invalid_messages = PyDict::new(py);
let sync_chat_error = module
.getattr("chat_completions")
.and_then(|function| {
function.call1((value_call(py, "messages", &invalid_messages, None),))
})
.expect_err("sync chat should reject a non-list messages value");
let async_chat_error = module
.getattr("achat_completions")
.and_then(|function| {
function.call1((value_call(py, "messages", &invalid_messages, None),))
})
.expect_err("async chat should reject a non-list messages value");
assert_eq!(
sync_chat_error.to_string(),
"ValueError: messages must be a list"
);
assert_eq!(async_chat_error.to_string(), sync_chat_error.to_string());
let invalid_headers = PyList::empty(py);
let kwargs = PyDict::new(py);
kwargs
@ -221,51 +199,43 @@ value = Broken()
});
}
#[test]
#[rstest]
fn missing_and_explicit_none_optional_params_share_the_next_error() {
Python::initialize();
Python::attach(|py| {
let module = crate::native_module(py);
let audio = PyDict::new(py);
let transcribe = |optional_params: Option<Bound<'_, PyAny>>| {
let kwargs = PyDict::new(py);
if let Some(optional_params) = optional_params {
kwargs.set_item("optional_params", optional_params).unwrap();
}
module
.getattr("transcription")
.and_then(|function| {
function.call1((value_call(py, "audio", &audio, Some(&kwargs)),))
})
.expect_err("model 'model' has no provider")
.to_string()
};
let omitted = transcribe(None);
assert_eq!(transcribe(Some(py.None().into_bound(py))), omitted);
assert_ne!(omitted, "ValueError: optional_params must be a dict");
assert_eq!(
transcribe(Some(PyList::empty(py).into_any())),
"ValueError: optional_params must be a dict"
);
});
}
#[rstest]
fn route_input_validation_preserves_left_to_right_order() {
Python::initialize();
Python::attach(|py| {
let module = crate::native_module(py);
let invalid = PyList::empty(py);
let chat_kwargs = PyDict::new(py);
chat_kwargs
.set_item("optional_params", &invalid)
.expect("kwargs should accept optional_params");
chat_kwargs
.set_item("extra_headers", &invalid)
.expect("kwargs should accept extra_headers");
let invalid_messages = PyDict::new(py);
let error = module
.getattr("chat_completions")
.and_then(|function| {
function.call1((value_call(
py,
"messages",
&invalid_messages,
Some(&chat_kwargs),
),))
})
.expect_err("messages should be validated first");
assert_eq!(error.to_string(), "ValueError: messages must be a list");
let valid_messages = PyList::empty(py);
let error = module
.getattr("chat_completions")
.and_then(|function| {
function.call1((value_call(
py,
"messages",
&valid_messages,
Some(&chat_kwargs),
),))
})
.expect_err("optional_params should be validated before headers");
assert_eq!(
error.to_string(),
"ValueError: optional_params must be a dict"
);
let headers_kwargs = PyDict::new(py);
headers_kwargs
.set_item("extra_headers", &invalid)
@ -286,43 +256,4 @@ value = Broken()
assert!(!error.to_string().contains("extra_headers"));
});
}
#[test]
fn missing_and_explicit_none_optional_params_share_the_next_error() {
Python::initialize();
Python::attach(|py| {
let module = crate::native_module(py);
let messages = PyList::empty(py);
let headers = PyList::empty(py);
let omitted = PyDict::new(py);
omitted
.set_item("extra_headers", &headers)
.expect("kwargs should accept extra_headers");
let explicit = PyDict::new(py);
explicit
.set_item("optional_params", py.None())
.expect("kwargs should accept optional_params");
explicit
.set_item("extra_headers", &headers)
.expect("kwargs should accept extra_headers");
let omitted_error = module
.getattr("chat_completions")
.and_then(|function| {
function.call1((value_call(py, "messages", &messages, Some(&omitted)),))
})
.expect_err("omitted optional_params should reach header validation");
let explicit_error = module
.getattr("chat_completions")
.and_then(|function| {
function.call1((value_call(py, "messages", &messages, Some(&explicit)),))
})
.expect_err("None optional_params should reach header validation");
assert_eq!(
omitted_error.to_string(),
"ValueError: extra_headers must be a dict"
);
assert_eq!(explicit_error.to_string(), omitted_error.to_string());
});
}
}

View file

@ -27,12 +27,12 @@ enum OcrHostData {
/// document as it goes), acquires Azure AD tokens, and builds the public response and
/// exception.
pub(super) struct OcrPythonHost {
request: Py<PyAny>,
request: Py<PyDict>,
data: OcrHostData,
}
impl OcrPythonHost {
pub(super) fn new(request: Py<PyAny>) -> Self {
pub(super) fn new(request: Py<PyDict>) -> Self {
Self {
request,
data: OcrHostData::Unprojected,
@ -209,7 +209,7 @@ del provider
.unwrap()
.cast_into::<PyDict>()
.unwrap();
let mut host = OcrPythonHost::new(py.None());
let mut host = OcrPythonHost::new(PyDict::new(py).unbind());
assert!(host.decode_request(py, &kwargs).unwrap().caller_token);
locals.del_item("kwargs").unwrap();
drop(kwargs);

View file

@ -9,10 +9,9 @@ use litellm_core_utils::settings::ProcessEnvironment;
use litellm_host_python::to_py;
use litellm_inference_ocr::provider_config;
use litellm_llms::base_llm::ocr::settings::OcrSettings;
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use pyo3::prelude::*;
use super::NativeCall;
use crate::{
coercion::FieldSpec,
@ -29,21 +28,9 @@ const ENABLE_AZURE_AD_TOKEN_REFRESH: FieldSpec<bool> =
Ok(field.exact_true())
});
fn run_ocr(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
let (arguments, hooks) = crate::routes::call_hooks(
py,
LoggingOperation::Ocr,
&request,
&args,
&kwargs,
asynchronous,
)?;
fn run_ocr(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
let (arguments, hooks) =
crate::routes::call_hooks(py, LoggingOperation::Ocr, &call, asynchronous)?;
crate::routes::run_public_call(
py,
arguments,
@ -61,7 +48,7 @@ fn run_ocr(
let route = litellm_inference_ocr::OcrRoute::new(client);
Ok(route.machine(request, None))
},
OcrPythonHost::new(request.unbind()),
OcrPythonHost::new(call.bound.unbind()),
hooks,
asynchronous,
)
@ -81,15 +68,13 @@ fn project_provider_defaults(snapshot: &Snapshot<'_>) -> PyResult<OcrSettings> {
}
#[pyfunction]
pub(crate) fn ocr(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_ocr(py, call.bound.into_any(), call.args, call.kwargs, false)
pub(crate) fn ocr(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
run_ocr(py, call, false)
}
#[pyfunction]
pub(crate) fn aocr(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_ocr(py, call.bound.into_any(), call.args, call.kwargs, true)
pub(crate) fn aocr(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
run_ocr(py, call, true)
}
#[pyfunction]

View file

@ -1,5 +1,5 @@
use litellm_auth::SecretValue;
use litellm_host_python::from_py;
use litellm_host_python::{from_py, present};
use litellm_inference_ocr::{
types::{LiteLLMOcrRequest, OcrDocumentInput},
wire::{OcrWireRequest, consumed_optional_params, decode_document, decode_request_input},
@ -22,13 +22,13 @@ pub(super) struct OcrHostHandles {
}
struct OcrArguments<'a, 'py> {
request: &'a Bound<'py, PyAny>,
bound: &'a Bound<'py, PyDict>,
kwargs: &'a Bound<'py, PyDict>,
}
impl<'py> OcrArguments<'_, 'py> {
fn lookup(&self, name: &str) -> PyResult<Bound<'py, PyAny>> {
litellm_host_python::lookup(self.kwargs, self.request, name)?
litellm_host_python::lookup(self.kwargs, self.bound, name)?
.ok_or_else(|| PyValueError::new_err(format!("missing argument: {name}")))
}
@ -58,7 +58,7 @@ impl<'py> OcrArguments<'_, 'py> {
fn extra_headers(&self) -> PyResult<Option<Map<String, Value>>> {
self.lookup("extra_headers")?
.extract::<Option<Py<PyAny>>>()?
.map(|value| from_py(value.bind(self.request.py())))
.map(|value| from_py(value.bind(self.bound.py())))
.transpose()
}
@ -66,7 +66,7 @@ impl<'py> OcrArguments<'_, 'py> {
Ok(self
.lookup("timeout")?
.extract::<Option<Py<PyAny>>>()?
.map(|value| python_timeout_seconds(self.request.py(), value))
.map(|value| python_timeout_seconds(self.bound.py(), value))
.transpose()?
.flatten())
}
@ -110,10 +110,10 @@ impl ProjectedDocument {
}
pub(super) fn project_request(
request: &Bound<'_, PyAny>,
bound: &Bound<'_, PyDict>,
kwargs: &Bound<'_, PyDict>,
) -> PyResult<(LiteLLMOcrRequest<OcrDocumentInput>, OcrHostHandles)> {
let arguments = OcrArguments { request, kwargs };
let arguments = OcrArguments { bound, kwargs };
let model = arguments.model()?;
let custom_llm_provider = arguments.custom_llm_provider()?;
let document = ProjectedDocument::project(&arguments.document()?)?;
@ -122,7 +122,7 @@ pub(super) fn project_request(
.map_err(ocr_error_to_pyerr)?;
let names = specs.iter().map(|spec| spec.name).collect::<Vec<_>>();
let optional_params =
project_optional_fields(names.iter().copied(), |name| kwargs.get_item(name))?;
project_optional_fields(names.iter().copied(), |name| present(kwargs, bound, name))?;
let input_sources = request_input_sources(
kwargs,
names
@ -136,7 +136,7 @@ pub(super) fn project_request(
let timeout_seconds = arguments.timeout_seconds()?;
let wire = OcrWireRequest {
model,
document: document.resolve(request.py())?,
document: document.resolve(bound.py())?,
api_key,
api_base,
custom_llm_provider,
@ -169,11 +169,20 @@ mod tests {
locals
}
fn dict<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyDict> {
locals
.get_item(name)
.unwrap()
.unwrap()
.cast_into::<PyDict>()
.unwrap()
}
fn arguments<'a, 'py>(
request: &'a Bound<'py, PyAny>,
bound: &'a Bound<'py, PyDict>,
kwargs: &'a Bound<'py, PyDict>,
) -> OcrArguments<'a, 'py> {
OcrArguments { request, kwargs }
OcrArguments { bound, kwargs }
}
fn project_document(document: &Bound<'_, PyAny>) -> PyResult<OcrDocumentInput> {
@ -204,141 +213,26 @@ sys.modules['litellm.rust_bridge.timeouts'] = timeouts
}
#[test]
fn kwargs_override_request_attributes_including_explicit_none() {
fn kwargs_override_bound_values_including_explicit_none() {
Python::initialize();
Python::attach(|py| {
let locals = eval(
py,
c"
class Request:
def __init__(self):
self.accesses = []
def __getattribute__(self, name):
if name != 'accesses':
object.__getattribute__(self, 'accesses').append(name)
return object.__getattribute__(self, name)
request = Request()
request.model = 'from-request'
request.custom_llm_provider = 'mistral'
bound = {'model': 'from-bound', 'custom_llm_provider': 'mistral'}
kwargs = {'model': 'from-kwargs', 'custom_llm_provider': None}
",
);
let request = locals.get_item("request").unwrap().unwrap();
let kwargs = locals
.get_item("kwargs")
.unwrap()
.unwrap()
.cast_into::<PyDict>()
.unwrap();
let arguments = arguments(&request, &kwargs);
let (bound, kwargs) = (dict(&locals, "bound"), dict(&locals, "kwargs"));
let arguments = arguments(&bound, &kwargs);
assert_eq!(arguments.model().unwrap(), "from-kwargs");
assert_eq!(arguments.custom_llm_provider().unwrap(), None);
let accesses: Vec<String> = request.getattr("accesses").unwrap().extract().unwrap();
assert_eq!(accesses, Vec::<String>::new());
});
}
#[test]
fn missing_kwargs_read_the_request_property_once() {
Python::initialize();
Python::attach(|py| {
let locals = eval(
py,
c"
class Request:
def __init__(self):
self.reads = 0
@property
def model(self):
self.reads += 1
return 'mistral-ocr-latest'
request = Request()
kwargs = {}
",
);
let request = locals.get_item("request").unwrap().unwrap();
let kwargs = locals
.get_item("kwargs")
.unwrap()
.unwrap()
.cast_into::<PyDict>()
.unwrap();
assert_eq!(
arguments(&request, &kwargs).model().unwrap(),
"mistral-ocr-latest"
);
assert_eq!(
request.getattr("reads").unwrap().extract::<i32>().unwrap(),
1
);
});
}
#[test]
fn request_property_exceptions_keep_their_identity() {
Python::initialize();
Python::attach(|py| {
let locals = eval(
py,
c"
failure = LookupError('model failed')
class Request:
@property
def model(self):
raise failure
request = Request()
kwargs = {}
",
);
let request = locals.get_item("request").unwrap().unwrap();
let kwargs = locals
.get_item("kwargs")
.unwrap()
.unwrap()
.cast_into::<PyDict>()
.unwrap();
let error = arguments(&request, &kwargs).model().unwrap_err();
assert!(
error
.value(py)
.is(locals.get_item("failure").unwrap().unwrap())
);
});
}
#[test]
fn unused_raising_property_is_never_inspected() {
Python::initialize();
Python::attach(|py| {
let locals = eval(
py,
c"
class Request:
@property
def unused(self):
raise RuntimeError('unused')
model = 'mistral-ocr-latest'
custom_llm_provider = None
request = Request()
kwargs = {}
",
);
let request = locals.get_item("request").unwrap().unwrap();
let kwargs = locals
.get_item("kwargs")
.unwrap()
.unwrap()
.cast_into::<PyDict>()
.unwrap();
let arguments = arguments(&request, &kwargs);
assert_eq!(arguments.model().unwrap(), "mistral-ocr-latest");
assert_eq!(arguments.custom_llm_provider().unwrap(), None);
});
}
/// A reader that rewrites the request while it runs shows which arguments projection
/// read before it and which after: every other argument is read first, and the read
/// happens exactly once.
/// A reader that rewrites the bound arguments while it runs shows which arguments
/// projection read before it and which after: every other argument is read first, and
/// the read happens exactly once.
#[test]
fn document_readers_are_read_once_after_every_other_argument() {
Python::initialize();
@ -347,37 +241,28 @@ kwargs = {}
let locals = eval(
py,
c"
class Request:
model = 'mistral/mistral-ocr-latest'
custom_llm_provider = None
api_key = None
api_base = 'https://original.example.com'
extra_headers = {'x-source': 'original'}
timeout = 1
@property
def document(self):
return document
class Reader:
reads = 0
def read(self):
Reader.reads += 1
Request.api_base = 'https://mutated.example.com'
Request.extra_headers = {'x-source': 'mutated'}
Request.timeout = 9
bound['api_base'] = 'https://mutated.example.com'
bound['extra_headers'] = {'x-source': 'mutated'}
bound['timeout'] = 9
return b'abc'
document = {'type': 'file', 'file': Reader(), 'mime_type': 'application/pdf'}
request = Request()
bound = {
'model': 'mistral/mistral-ocr-latest',
'custom_llm_provider': None,
'api_key': None,
'api_base': 'https://original.example.com',
'extra_headers': {'x-source': 'original'},
'timeout': 1,
'document': {'type': 'file', 'file': Reader(), 'mime_type': 'application/pdf'},
}
kwargs = {}
",
);
let request = locals.get_item("request").unwrap().unwrap();
let kwargs = locals
.get_item("kwargs")
.unwrap()
.unwrap()
.cast_into::<PyDict>()
.unwrap();
let (projected, _) = project_request(&request, &kwargs).unwrap();
let (projected, _) =
project_request(&dict(&locals, "bound"), &dict(&locals, "kwargs")).unwrap();
assert_eq!(
py.eval(c"Reader.reads", Some(&locals), Some(&locals))
.unwrap()
@ -524,34 +409,26 @@ document = Document()
});
}
fn request_and_kwargs<'py>(
fn bound_and_kwargs<'py>(
py: Python<'py>,
kwargs: &std::ffi::CStr,
) -> (Bound<'py, PyAny>, Bound<'py, PyDict>) {
) -> (Bound<'py, PyDict>, Bound<'py, PyDict>) {
let locals = eval(
py,
c"
class Request:
model = 'mistral/mistral-ocr-latest'
custom_llm_provider = 'mistral'
document = {'type': 'document_url', 'document_url': 'https://example.com/request.pdf'}
api_key = None
api_base = 'https://request.example.com'
extra_headers = {'x-source': 'request'}
timeout = 1
request = Request()
bound = {
'model': 'mistral/mistral-ocr-latest',
'custom_llm_provider': 'mistral',
'document': {'type': 'document_url', 'document_url': 'https://example.com/bound.pdf'},
'api_key': None,
'api_base': 'https://bound.example.com',
'extra_headers': {'x-source': 'bound'},
'timeout': 1,
}
",
);
py.run(kwargs, Some(&locals), Some(&locals)).unwrap();
(
locals.get_item("request").unwrap().unwrap(),
locals
.get_item("kwargs")
.unwrap()
.unwrap()
.cast_into::<PyDict>()
.unwrap(),
)
(dict(&locals, "bound"), dict(&locals, "kwargs"))
}
#[test]
@ -559,7 +436,7 @@ request = Request()
Python::initialize();
Python::attach(|py| {
stub_timeout_conversion(py);
let (request, kwargs) = request_and_kwargs(
let (bound, kwargs) = bound_and_kwargs(
py,
c"
kwargs = {
@ -575,7 +452,7 @@ kwargs = {
}
",
);
let (projected, _) = project_request(&request, &kwargs).unwrap();
let (projected, _) = project_request(&bound, &kwargs).unwrap();
assert_eq!(
projected.optional_params.keys().collect::<Vec<_>>(),
["pages"]
@ -584,12 +461,29 @@ kwargs = {
});
}
#[rstest::rstest]
#[case::explicit_none_is_unset(c"bound['pages'] = [1]\nkwargs = {'pages': None}", None)]
#[case::bound_fallback(c"bound['pages'] = [1]\nkwargs = {}", Some(serde_json::json!([1])))]
#[case::keyword_wins(c"bound['pages'] = [1]\nkwargs = {'pages': [0]}", Some(serde_json::json!([0])))]
fn optional_params_read_through_bound_and_drop_none(
#[case] script: &std::ffi::CStr,
#[case] expected: Option<Value>,
) {
Python::initialize();
Python::attach(|py| {
stub_timeout_conversion(py);
let (bound, kwargs) = bound_and_kwargs(py, script);
let (projected, _) = project_request(&bound, &kwargs).unwrap();
assert_eq!(projected.optional_params.get("pages").cloned(), expected);
});
}
#[test]
fn replacement_kwargs_project_provider_connection_and_timeout() {
Python::initialize();
Python::attach(|py| {
stub_timeout_conversion(py);
let (request, kwargs) = request_and_kwargs(
let (bound, kwargs) = bound_and_kwargs(
py,
c"
kwargs = {
@ -602,7 +496,7 @@ kwargs = {
}
",
);
let (projected, handles) = project_request(&request, &kwargs).unwrap();
let (projected, handles) = project_request(&bound, &kwargs).unwrap();
assert_eq!(handles.provider, "azure_ai");
assert_eq!(projected.model, "mistral-ocr-latest");
assert_eq!(

View file

@ -1,6 +1,6 @@
use std::convert::Infallible;
use super::super::inference::InferenceHost;
use crate::routes::inference::InferenceHost;
use litellm_host_python::{InvokeError, PythonBinding, PythonHostCalls, PythonOwned};
use litellm_inference_responses::{Error, route::Responses, types::ResponsesCall};
use pyo3::{

View file

@ -0,0 +1,102 @@
mod host;
mod websocket;
use std::sync::Arc;
use host::ResponsesPythonHost;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, ScopedCache};
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_host_python::present;
use litellm_inference_responses::{ResponsesRoute, route::Responses};
use litellm_secrets::source::SecretSource;
use pyo3::prelude::*;
use serde_json::Value;
pub(crate) use websocket::ResponsesWebSocketConnection;
use super::{
NativeCall,
inference::{InferenceHost, InferenceRoute, run_inference},
};
use crate::errors::RustBridgeDeclined;
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.responses.route_host";
fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
if let Some(reason) = py
.import(ROUTE_HOST_MODULE)?
.getattr("decline_reason")?
.call1((&call.bound,))?
.extract::<Option<String>>()?
{
return Err(RustBridgeDeclined::new_err(reason));
}
let argument = |name: &str| present(&call.kwargs, &call.bound, name);
let model = argument("model")?
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("model is required"))?
.extract::<String>()?;
let provider = argument("custom_llm_provider")?
.map(|value| value.extract::<String>())
.transpose()?;
if provider
.as_deref()
.is_some_and(|provider| provider != "openai")
|| model
.strip_prefix("openai/")
.unwrap_or(&model)
.contains('/')
{
return Err(RustBridgeDeclined::new_err(
"native HTTP responses provider",
));
}
if argument("stream")?
.map(|value| litellm_host_python::from_py::<Value>(&value))
.transpose()?
.is_some_and(|value| value == Value::Bool(true))
{
return Err(RustBridgeDeclined::new_err(
"native Python responses streaming",
));
}
let host = InferenceHost::new(call.bound.clone().unbind(), ROUTE_HOST_MODULE);
run_inference::<ResponsesRoute, _>(py, call, asynchronous, ResponsesPythonHost(host))
}
#[pyfunction]
pub(crate) fn responses(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
run_responses(py, call, false)
}
#[pyfunction]
pub(crate) fn aresponses(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
run_responses(py, call, true)
}
impl InferenceRoute for ResponsesRoute {
type Protocol = Responses;
const OPERATION: LoggingOperation = LoggingOperation::Responses;
const SYNC_CALL_TYPE: &'static str = "responses";
const ASYNC_CALL_TYPE: &'static str = "aresponses";
fn new(
http: litellm_http::Client,
auth: Arc<AuthServices>,
secrets: Arc<dyn SecretSource>,
) -> Self {
Self::new(http, auth, secrets)
}
fn with_cache(self, cache: ScopedCache) -> Self {
self.with_cache(cache)
}
fn machine(
self,
call: <Responses as Protocol>::Request,
policy: CachePolicy,
) -> HostedMachine<Responses> {
self.machine(call, policy)
}
}

View file

@ -1,121 +1,12 @@
mod host;
use litellm_inference_responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use pyo3::prelude::*;
use serde_json::Value;
use crate::{
errors::{RustBridgeDeclined, route_error_to_pyerr},
errors::route_error_to_pyerr,
marshal::{marshal_headers, optional_timeout},
};
fn run_public(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
use super::inference::InferenceHost;
use litellm_callbacks_legacy_python::LoggingOperation;
let host = InferenceHost::new(
request.clone().unbind(),
"litellm.rust_bridge.responses.route_host",
);
if let Some(reason) = py
.import("litellm.rust_bridge.responses.route_host")?
.getattr("decline_reason")?
.call1((&request,))?
.extract::<Option<String>>()?
{
return Err(RustBridgeDeclined::new_err(reason));
}
let model = host
.argument(py, &kwargs, "model")?
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("model is required"))?
.extract::<String>()?;
let provider = host
.argument(py, &kwargs, "custom_llm_provider")?
.map(|value| value.extract::<String>())
.transpose()?;
if provider
.as_deref()
.is_some_and(|provider| provider != "openai")
|| model
.strip_prefix("openai/")
.unwrap_or(&model)
.contains('/')
{
return Err(RustBridgeDeclined::new_err(
"native HTTP responses provider",
));
}
if host
.argument(py, &kwargs, "stream")?
.map(|value| litellm_host_python::from_py::<Value>(&value))
.transpose()?
.is_some_and(|value| value == Value::Bool(true))
{
return Err(RustBridgeDeclined::new_err(
"native Python responses streaming",
));
}
let cache_call_type = if asynchronous {
"aresponses"
} else {
"responses"
};
crate::cache::admit_native(py, &kwargs, cache_call_type)?;
let (arguments, hooks) = crate::routes::call_hooks(
py,
LoggingOperation::Responses,
&request,
&args,
&kwargs,
asynchronous,
)?;
crate::routes::run_public_call(
py,
arguments,
move |py, arguments, request| {
let route = litellm_inference_responses::ResponsesRoute::new(
crate::http::provider_client(py, arguments, asynchronous)?
.map_err(crate::http::client_error)?,
crate::http::resources().auth.clone(),
crate::secrets::source(py)?,
);
let (cache, cache_options) =
crate::cache::configured_native(py, arguments, cache_call_type)?;
let route = match cache {
Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new(
cache,
litellm_cache_response::CacheScope::Shared,
)),
None => route,
};
Ok(route.machine(request, cache_options.policy))
},
host::ResponsesPythonHost(host),
hooks,
asynchronous,
)
}
#[pyfunction]
pub(crate) fn responses(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_public(py, call.bound.into_any(), call.args, call.kwargs, false)
}
#[pyfunction]
pub(crate) fn aresponses(py: Python<'_>, call: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let call = super::NativeCall::extract(&call)?;
run_public(py, call.bound.into_any(), call.args, call.kwargs, true)
}
#[pyclass]
pub(crate) struct ResponsesWebSocketConnection {
inner: RustResponsesWebSocketConnection,

View file

@ -172,7 +172,7 @@ impl NativeTraceStorage {
BTreeMap<String, serde_json::Value>,
>,
) -> PyResult<Bound<'py, PyAny>> {
let table = InsertTable::parse(table).map_err(map_error)?;
let table = table.parse::<InsertTable>().map_err(map_error)?;
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
let connection = self.config.storage().writer().clone();
let database = self.config.storage().database().to_owned();

View file

@ -5,18 +5,37 @@ use strum::IntoStaticStr;
use crate::SecretValue;
#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, IntoStaticStr, PartialEq, Serialize)]
#[derive(
Clone,
Copy,
Debug,
Deserialize,
Eq,
Hash,
IntoStaticStr,
PartialEq,
Serialize,
strum::VariantArray,
)]
#[serde(rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]
pub enum KeyManagementSystem {
#[strum(serialize = "google_kms")]
GoogleKms,
#[strum(serialize = "azure_key_vault")]
AzureKeyVault,
#[strum(serialize = "aws_secret_manager")]
AwsSecretManager,
#[strum(serialize = "google_secret_manager")]
GoogleSecretManager,
#[strum(serialize = "hashicorp_vault")]
HashicorpVault,
#[strum(serialize = "cyberark")]
Cyberark,
#[strum(serialize = "local")]
Local,
#[strum(serialize = "aws_kms")]
AwsKms,
#[strum(serialize = "custom")]
Custom,
}

View file

@ -128,17 +128,21 @@ fn rotation_write_context_preserves_the_operation_context(aws_context: SecretOpe
fn provider_context_accepts_only_its_owner(
#[case] owner: KeyManagementSystem,
#[case] context: SecretOperationContext,
) {
for system in [
KeyManagementSystem::AwsSecretManager,
#[values(
KeyManagementSystem::GoogleKms,
KeyManagementSystem::AzureKeyVault,
KeyManagementSystem::AwsSecretManager,
KeyManagementSystem::GoogleSecretManager,
KeyManagementSystem::HashicorpVault,
KeyManagementSystem::Cyberark,
] {
assert_eq!(context.validate_for(system).is_ok(), system == owner);
assert!(SecretOperationContext::Default.validate_for(system).is_ok());
}
KeyManagementSystem::Local,
KeyManagementSystem::AwsKms,
KeyManagementSystem::Custom
)]
system: KeyManagementSystem,
) {
assert_eq!(context.validate_for(system).is_ok(), system == owner);
assert!(SecretOperationContext::Default.validate_for(system).is_ok());
if matches!(
owner,
KeyManagementSystem::AzureKeyVault | KeyManagementSystem::GoogleSecretManager

View file

@ -23,17 +23,22 @@ const OIDC_ALLOWED_CREDENTIAL_DIRS: &str = "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS
const DEFAULT_CREDENTIAL_DIRS: &str = "/var/run/secrets,/run/secrets";
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, strum::EnumString, strum::AsRefStr)]
#[strum(serialize_all = "snake_case")]
pub enum OidcProvider {
#[strum(serialize = "google")]
Google,
#[strum(serialize = "circleci")]
CircleCi,
#[strum(serialize = "circleci_v2")]
CircleCiV2,
#[strum(serialize = "github")]
Github,
#[strum(serialize = "azure")]
Azure,
#[strum(serialize = "file")]
File,
#[strum(serialize = "env")]
Env,
#[strum(serialize = "env_path")]
EnvPath,
}

View file

@ -13,6 +13,7 @@ serde.workspace = true
semver.workspace = true
serde_json.workspace = true
sha2.workspace = true
strum.workspace = true
tar.workspace = true
target-lexicon.workspace = true
thiserror.workspace = true

View file

@ -5,7 +5,7 @@ use semver::Version;
use crate::Error;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)]
pub enum Wire {
ChatCompletions,
Messages,

View file

@ -2,6 +2,7 @@ use std::path::Path;
use litellm_testkit::{ClaudeCode, Codex, Configure, Error, Opencode, Settings, Version, Wire};
use rstest::rstest;
use strum::VariantArray;
fn settings(wire: Wire) -> Settings {
Settings {
@ -121,7 +122,10 @@ fn opencode_uses_a_different_provider_package_for_every_wire() {
.unwrap()
.to_owned()
};
let packages = [Wire::ChatCompletions, Wire::Responses, Wire::Messages].map(package);
let packages = Wire::VARIANTS
.iter()
.map(|wire| package(*wire))
.collect::<Vec<_>>();
assert_eq!(
packages

View file

@ -30,29 +30,19 @@ fn max_insert_bytes() -> Result<usize, Error> {
pub type InsertRow = BTreeMap<String, Shared<Value>>;
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr)]
#[strum(parse_err_ty = Error, parse_err_fn = invalid_table)]
pub enum InsertTable {
#[strum(serialize = "otel_traces")]
OtelTraces,
#[strum(serialize = "spend_logs")]
SpendLogs,
#[strum(serialize = "lens_feedback")]
LensFeedback,
}
impl InsertTable {
pub fn parse(value: &str) -> Result<Self, Error> {
match value {
"otel_traces" => Ok(Self::OtelTraces),
"spend_logs" => Ok(Self::SpendLogs),
"lens_feedback" => Ok(Self::LensFeedback),
_ => Err(Error::InvalidTable),
}
}
fn name(&self) -> &'static str {
match self {
Self::OtelTraces => "otel_traces",
Self::SpendLogs => "spend_logs",
Self::LensFeedback => "lens_feedback",
}
}
fn invalid_table(_name: &str) -> Error {
Error::InvalidTable
}
pub async fn insert_rows(
@ -81,7 +71,7 @@ pub async fn insert_shared_rows(
client,
connection,
database,
table.name(),
<&'static str>::from(table),
&token,
body,
)
@ -242,7 +232,25 @@ mod tests {
use rstest::rstest;
use serde_json::json;
use super::{Error, shared_rows, write_rows};
use super::{Error, InsertTable, shared_rows, write_rows};
#[rstest]
#[case::otel_traces("otel_traces", InsertTable::OtelTraces)]
#[case::spend_logs("spend_logs", InsertTable::SpendLogs)]
#[case::lens_feedback("lens_feedback", InsertTable::LensFeedback)]
fn insert_table_parses_each_table_name(#[case] name: &str, #[case] expected: InsertTable) {
assert_eq!(name.parse::<InsertTable>().unwrap(), expected);
}
#[rstest]
#[case::unknown("events")]
#[case::case_sensitive("OTEL_TRACES")]
fn insert_table_rejects_unknown_names(#[case] name: &str) {
assert!(matches!(
name.parse::<InsertTable>(),
Err(Error::InvalidTable)
));
}
#[rstest]
fn encoded_limit_counts_utf8_bytes_across_rows() {

View file

@ -56,15 +56,21 @@ enum PathPart {
#[macro_rules_attribute::apply(crate::response_type)]
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, strum::Display)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
#[cfg_attr(feature = "schema", schemars(rename = "MetadataValueType"))]
enum JsonKind {
#[strum(serialize = "array")]
Array,
#[strum(serialize = "boolean")]
Boolean,
#[strum(serialize = "integer")]
Integer,
#[strum(serialize = "null")]
Null,
#[strum(serialize = "number")]
Number,
#[strum(serialize = "object")]
Object,
#[strum(serialize = "string")]
String,
}

View file

@ -4,9 +4,11 @@
Clone, Copy, Debug, strum::Display, strum::AsRefStr, strum::EnumIter, strum::IntoStaticStr,
)]
#[serde(rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]
pub enum TraceTable {
#[strum(serialize = "otel_traces")]
OtelTraces,
#[strum(serialize = "agent_traces_by_key")]
AgentTracesByKey,
#[strum(serialize = "spend_logs")]
SpendLogs,
}

View file

@ -126,10 +126,12 @@ async fn lens_sample_keeps_spans_before_window_start_and_excludes_old_only_trace
}
#[derive(Clone, Copy, strum::AsRefStr)]
#[strum(serialize_all = "snake_case")]
enum ScopeCase {
#[strum(serialize = "admin")]
Admin,
#[strum(serialize = "team")]
Team,
#[strum(serialize = "other_team")]
OtherTeam,
}

View file

@ -15,14 +15,23 @@ use crate::{
/// Claude Code's built-in tracing, identified by its instrumentation scope.
pub(crate) struct ClaudeCode;
#[derive(Debug, PartialEq, Eq, strum::EnumString)]
enum SpanType {
#[strum(serialize = "assistant_response")]
AssistantResponse,
#[strum(serialize = "tool_result")]
ToolResult,
#[strum(serialize = "api_request_body")]
ApiRequestBody,
#[strum(serialize = "compaction")]
Compaction,
#[strum(serialize = "interaction")]
Interaction,
#[strum(serialize = "llm_request")]
LlmRequest,
#[strum(serialize = "tool")]
Tool,
#[strum(disabled)]
Other,
}
@ -33,16 +42,7 @@ fn span_type(name: &str, attributes: &BTreeMap<String, String>) -> SpanType {
} else {
kind
};
match kind {
"assistant_response" => SpanType::AssistantResponse,
"tool_result" => SpanType::ToolResult,
"api_request_body" => SpanType::ApiRequestBody,
"compaction" => SpanType::Compaction,
"interaction" => SpanType::Interaction,
"llm_request" => SpanType::LlmRequest,
"tool" => SpanType::Tool,
_ => SpanType::Other,
}
kind.parse().unwrap_or(SpanType::Other)
}
/// `agent:custom:search_agent` -> `search_agent`: the subagent a request ran for.
@ -318,7 +318,7 @@ mod tests {
use rstest::rstest;
use serde_json::Value;
use super::CLAUDE_CODE_SCOPE;
use super::{CLAUDE_CODE_SCOPE, SpanType, span_type};
use crate::{
Error,
normalize::{Normalization, NormalizedSpan, ObservationType},
@ -355,6 +355,22 @@ mod tests {
.collect()
}
#[rstest]
#[case::assistant_response("assistant_response", SpanType::AssistantResponse)]
#[case::tool_result("tool_result", SpanType::ToolResult)]
#[case::api_request_body("api_request_body", SpanType::ApiRequestBody)]
#[case::compaction("compaction", SpanType::Compaction)]
#[case::interaction("interaction", SpanType::Interaction)]
#[case::llm_request("llm_request", SpanType::LlmRequest)]
#[case::tool("tool", SpanType::Tool)]
#[case::unknown("surprise", SpanType::Other)]
fn span_type_maps_each_recorded_kind(#[case] kind: &str, #[case] expected: SpanType) {
assert_eq!(
span_type("anything", &attributes(&[("span.type", kind)])),
expected
);
}
#[rstest]
fn notification_prompts_keep_user_provenance_and_compaction_is_system() {
let prompt_text =

View file

@ -11,18 +11,24 @@ use crate::{
pub(crate) struct GenAi;
#[derive(strum::EnumString)]
#[strum(serialize_all = "snake_case")]
pub(crate) enum Operation {
#[strum(serialize = "create_agent")]
CreateAgent,
#[strum(serialize = "invoke_agent")]
InvokeAgent,
#[strum(serialize = "invoke_workflow")]
InvokeWorkflow,
#[strum(serialize = "chat")]
Chat,
#[strum(serialize = "text_completion", serialize = "completion")]
TextCompletion,
#[strum(serialize = "generate_content")]
GenerateContent,
#[strum(serialize = "execute_tool")]
ExecuteTool,
#[strum(serialize = "embeddings", serialize = "embedding")]
Embeddings,
#[strum(serialize = "retrieval")]
Retrieval,
}

View file

@ -24,32 +24,56 @@ pub enum AgentType {
serde_with::DeserializeFromStr,
serde_with::SerializeDisplay,
)]
#[strum(serialize_all = "kebab-case")]
pub enum Integration {
#[strum(serialize = "claude-code")]
ClaudeCode,
#[strum(serialize = "claude-agent-sdk")]
ClaudeAgentSdk,
#[strum(serialize = "openai-codex")]
OpenaiCodex,
#[strum(serialize = "deepagents-code")]
DeepagentsCode,
#[strum(serialize = "cursor")]
Cursor,
#[strum(serialize = "pi")]
Pi,
#[strum(serialize = "opencode")]
Opencode,
#[strum(serialize = "copilot")]
Copilot,
#[strum(serialize = "langchain")]
Langchain,
#[strum(serialize = "langgraph")]
Langgraph,
#[strum(serialize = "deepagents")]
Deepagents,
#[strum(serialize = "autogen")]
Autogen,
#[strum(serialize = "crewai")]
Crewai,
#[strum(serialize = "google-adk")]
GoogleAdk,
#[strum(serialize = "llama-index")]
LlamaIndex,
#[strum(serialize = "mastra")]
Mastra,
#[strum(serialize = "microsoft-agent-framework")]
MicrosoftAgentFramework,
#[strum(serialize = "openai-agents")]
OpenaiAgents,
#[strum(serialize = "pydantic-ai")]
PydanticAi,
#[strum(serialize = "semantic-kernel")]
SemanticKernel,
#[strum(serialize = "strands")]
Strands,
#[strum(serialize = "vercel-ai-sdk")]
VercelAiSdk,
#[strum(serialize = "instructor")]
Instructor,
#[strum(serialize = "n8n")]
N8n,
#[strum(serialize = "temporal")]
Temporal,
#[strum(default)]
Other(String),
@ -138,23 +162,36 @@ impl AgentMetadata {
}
#[derive(strum::EnumString, strum::IntoStaticStr)]
#[strum(serialize_all = "snake_case")]
enum MetadataField {
#[strum(serialize = "lc_agent_name")]
LcAgentName,
#[strum(serialize = "ls_integration")]
LsIntegration,
#[strum(serialize = "ls_agent_type")]
LsAgentType,
#[strum(serialize = "ls_agent_purpose")]
LsAgentPurpose,
#[strum(serialize = "ls_agent_runtime")]
LsAgentRuntime,
#[strum(serialize = "ls_agent_runtime_version", to_string = "ls_agent_version")]
LsAgentVersion,
#[strum(serialize = "ls_trace_schema_version")]
LsTraceSchemaVersion,
#[strum(serialize = "thread_id")]
ThreadId,
#[strum(serialize = "ls_subagent_id")]
LsSubagentId,
#[strum(serialize = "ls_subagent_type")]
LsSubagentType,
#[strum(serialize = "ls_tool_name")]
LsToolName,
#[strum(serialize = "ls_model_name")]
LsModelName,
#[strum(serialize = "ls_provider")]
LsProvider,
#[strum(serialize = "git_branch")]
GitBranch,
#[strum(serialize = "git_commit_sha")]
GitCommitSha,
#[strum(serialize = "repository_url", to_string = "git_repo_url")]
GitRepoUrl,

View file

@ -32,20 +32,32 @@ pub use metadata::{AgentMetadata, AgentType, Integration};
#[macro_rules_attribute::apply(crate::wire_type)]
#[derive(Clone, Copy, Debug, Eq, PartialEq, strum::EnumString)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase", ascii_case_insensitive)]
#[strum(ascii_case_insensitive)]
#[cfg_attr(feature = "schema", schemars(rename = "SpanType"))]
pub enum ObservationType {
#[strum(serialize = "agent")]
Agent,
#[strum(serialize = "llm")]
Llm,
#[strum(serialize = "tool")]
Tool,
#[strum(serialize = "chain")]
Chain,
#[strum(serialize = "framework")]
Framework,
#[strum(serialize = "retriever")]
Retriever,
#[strum(serialize = "embedding")]
Embedding,
#[strum(serialize = "reranker")]
Reranker,
#[strum(serialize = "guardrail")]
Guardrail,
#[strum(serialize = "evaluator")]
Evaluator,
#[strum(serialize = "prompt")]
Prompt,
#[strum(serialize = "decision")]
Decision,
}

View file

@ -2,23 +2,38 @@ pub mod guide;
pub mod named;
#[derive(Clone, Copy, Debug, Eq, PartialEq, strum::EnumString, strum::Display, strum::AsRefStr)]
#[strum(serialize_all = "snake_case")]
pub enum ReadQuery {
#[strum(serialize = "list_traces")]
ListTraces,
#[strum(serialize = "trace_agents")]
TraceAgents,
#[strum(serialize = "trace_spans")]
TraceSpans,
#[strum(serialize = "trace_page_spans")]
TracePageSpans,
#[strum(serialize = "trace_identity")]
TraceIdentity,
#[strum(serialize = "span_detail")]
SpanDetail,
#[strum(serialize = "span_error")]
SpanError,
#[strum(serialize = "spend_by_response_ids")]
SpendByResponseIds,
#[strum(serialize = "availability")]
Availability,
#[strum(serialize = "agents")]
Agents,
#[strum(serialize = "sample")]
Sample,
#[strum(serialize = "content")]
Content,
#[strum(serialize = "evidence")]
Evidence,
#[strum(serialize = "feedback_target")]
FeedbackTarget,
#[strum(serialize = "feedback")]
Feedback,
#[strum(serialize = "feedback_summary")]
FeedbackSummary,
}

View file

@ -71,15 +71,25 @@ class CacheMode(str, Enum):
#### LiteLLM.Completion / Embedding Cache ####
def _is_conversation_item(item: object) -> bool:
if isinstance(item, BaseModel):
return True
if not isinstance(item, Mapping):
return False
block: Final = cast(Mapping[str, object], item) # cast-ok: isinstance leaves the key and value types unknown
return block.get("type") != "file"
def _request_message_count(kwargs: Mapping[str, object]) -> int:
"""Chat and Messages API `messages`, else Responses API `input` items; embedding `input` strings count as none"""
"""Chat and Messages API `messages`, else Responses API `input` items; embedding strings and file blocks count as none"""
messages: Final = kwargs.get("messages")
if isinstance(messages, list):
return len(messages)
input_items: Final = kwargs.get("input")
if not isinstance(input_items, list):
return 0
return sum(1 for item in input_items if isinstance(item, (Mapping, BaseModel)))
items: Final = cast(list[object], input_items) # cast-ok: isinstance leaves the element type unknown
return sum(1 for item in items if _is_conversation_item(item))
class Cache:

View file

@ -21,7 +21,7 @@ import time
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping
from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
from pydantic import ConfigDict, ValidationError
from pydantic import ConfigDict, SkipValidation, ValidationError
import litellm
from litellm._internal_context import post_response_phase
@ -39,7 +39,7 @@ from litellm.litellm_core_utils.logging_utils import (
from litellm.types.caching import CACHED_STREAM_EVENTS_KEY, EMBEDDING_CACHE_FORMAT_VERSION, CachedEmbedding
from litellm.types.integrations.custom_logger import converted_stream_requested
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.llms.openai import ChatCompletionFileObject, ResponsesAPIResponse
from litellm.types.rerank import RerankResponse
from litellm.types.utils import (
CachingDetails,
@ -71,6 +71,8 @@ from litellm.litellm_core_utils.core_helpers import (
)
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
EmbeddingCacheInputElement = str | list[int] | ChatCompletionFileObject
class CachingHandlerResponse(LiteLLMBaseModel):
"""
@ -82,7 +84,7 @@ class CachingHandlerResponse(LiteLLMBaseModel):
cached_result: object | None = None
final_embedding_cached_response: EmbeddingResponse | None = None
embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call
embedding_uncached_input: list[str | list[int]] | None = None
embedding_uncached_input: SkipValidation[list[EmbeddingCacheInputElement]] | None = None
in_memory_cache_obj: Final = InMemoryCache()
@ -508,7 +510,7 @@ class LLMCachingHandler:
_sync_get_cache = sync_get_cache
def handle_kwargs_input_list_or_str(self, kwargs: dict[str, object]) -> list[str]:
def handle_kwargs_input_list_or_str(self, kwargs: dict[str, object]) -> list[EmbeddingCacheInputElement]:
"""
Handles the input of kwargs['input'] being a list or a string
"""
@ -517,7 +519,11 @@ class LLMCachingHandler:
elif isinstance(kwargs["input"], list):
return kwargs["input"]
else:
raise ValueError("input must be a string or a list")
raise litellm.BadRequestError(
message="input must be a string or a list of strings and content blocks",
model=str(kwargs.get("model")),
llm_provider=str(kwargs.get("custom_llm_provider")),
)
def _extract_model_from_cached_results(self, non_null_list: list[tuple[int, CachedEmbedding]]) -> str | None:
"""
@ -851,10 +857,7 @@ class LLMCachingHandler:
self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs)
cached_result: object | None = None
if call_type == CallTypes.aembedding.value:
if isinstance(new_kwargs["input"], str):
new_kwargs["input"] = [new_kwargs["input"]]
elif not isinstance(new_kwargs["input"], list):
raise ValueError("input must be a string or a list")
new_kwargs["input"] = self.handle_kwargs_input_list_or_str(new_kwargs)
tasks: Final[list[Awaitable[object]]] = []
for idx, i in enumerate(new_kwargs["input"]):
preset_cache_key = litellm.cache.get_cache_key(**{**new_kwargs, "input": i})

View file

@ -1,10 +1,20 @@
import os
import sys
from enum import Enum
from types import MappingProxyType
from typing import Final, Literal
from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_range, get_env_int_or_none
SERVER_STREAMING_CLASSIFICATION_KEY: Final = "litellm_server_streaming_classification"
class ServerStreamingClassification(str, Enum):
MARKER = "litellm-server-streaming"
SERVER_STREAMING_CLASSIFICATION_MARKER: Final = ServerStreamingClassification.MARKER
DEFER_PYDANTIC_BUILD: Final = os.getenv("DEFER_PYDANTIC_BUILD", "true") in ("true", "1", "on")
DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))

View file

@ -6,7 +6,7 @@ import secrets
from collections.abc import Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
import httpx
@ -21,11 +21,14 @@ from litellm.litellm_core_utils.core_helpers import (
)
from litellm.secret_managers.main import str_to_bool
from litellm.types.guardrails import (
DEFAULT_GUARDRAIL_STREAM_SCOPE,
DynamicGuardrailParams,
GuardrailEventHooks,
GuardrailStreamScope,
LitellmParams,
LoggingOnlyScope,
Mode,
runtime_stream_scope,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
@ -49,6 +52,8 @@ from litellm.constants import (
GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS,
LOGS_GUARDRAIL_INFORMATION_MARKER,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
SERVER_STREAMING_CLASSIFICATION_KEY,
SERVER_STREAMING_CLASSIFICATION_MARKER,
)
from litellm.exceptions import (
BlockedPiiEntityError,
@ -173,6 +178,42 @@ def get_session_id_from_request_data(request_data: dict[str, Any]) -> str | None
return None
_REALTIME_STREAMING_HOOKS: Final = frozenset({GuardrailEventHooks.realtime_input_transcription})
def without_server_streaming_classification(data: Mapping[str, object]) -> dict[str, object]:
return {
key: value
for key, value in data.items()
if key != SERVER_STREAMING_CLASSIFICATION_KEY or value != SERVER_STREAMING_CLASSIFICATION_MARKER
}
def guardrail_request_data_with_streaming(
data: Mapping[str, object],
*,
is_streaming: bool,
) -> dict[str, object]:
data_without_server_classification: Final = without_server_streaming_classification(data)
if not is_streaming:
return data_without_server_classification
return {
**data_without_server_classification,
SERVER_STREAMING_CLASSIFICATION_KEY: SERVER_STREAMING_CLASSIFICATION_MARKER,
}
def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = None) -> bool:
if event_type in _REALTIME_STREAMING_HOOKS:
return True
if not isinstance(data, Mapping):
return False
return (
data.get("stream") is True
or data.get(SERVER_STREAMING_CLASSIFICATION_KEY) is SERVER_STREAMING_CLASSIFICATION_MARKER
)
class CustomGuardrail(CustomLogger):
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
use_native_during_call_hook: ClassVar[bool] = False
@ -183,6 +224,9 @@ class CustomGuardrail(CustomLogger):
records_own_guardrail_information: ClassVar[bool] = False
logging_only_scope: LoggingOnlyScope | None
stream_scope_default: GuardrailStreamScope = DEFAULT_GUARDRAIL_STREAM_SCOPE
stream_scope_by_hook: tuple[tuple[str, GuardrailStreamScope], ...] = ()
timeout: float | httpx.Timeout | None = None
def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: forwarded to cooperative __init_subclass__ hooks
@ -258,6 +302,8 @@ class CustomGuardrail(CustomLogger):
self.run_in_parallel: bool = run_in_parallel
self.scan_raw_request: bool = scan_raw_request
self.only_scan_new_messages: bool = only_scan_new_messages
stream_scope_arg: Final[object] = cast(object, kwargs.pop("stream_scope", None)) # cast-ok: config
self.apply_stream_scope(stream_scope_arg)
self.logging_only_scope = None
if timeout is not None:
self.timeout = timeout
@ -1099,6 +1145,23 @@ class CustomGuardrail(CustomLogger):
return name in suppressed_compression_guardrails()
def apply_stream_scope(self, stream_scope: object) -> None:
default, by_hook = runtime_stream_scope(stream_scope)
self.stream_scope_default = default
self.stream_scope_by_hook = tuple(by_hook.items())
def stream_scope_allows(self, data: object, event_type: GuardrailEventHooks) -> bool:
scope: Final = next(
(scope for hook, scope in self.stream_scope_by_hook if hook == event_type.value),
self.stream_scope_default,
)
if scope == "both":
return True
is_streaming: Final = _request_is_streaming(data, event_type)
if scope == "streaming":
return is_streaming
return not is_streaming
def should_run_guardrail(
self,
data,
@ -1142,8 +1205,10 @@ class CustomGuardrail(CustomLogger):
data, self.event_hook, event_type
)
if result is not None:
return result
return True
tagged_result: Final[bool] = bool(cast(object, result)) # cast-ok: helper return
data_obj: Final[object] = cast(object, data) # cast-ok: data param
return tagged_result and self.stream_scope_allows(data_obj, event_type)
return self.stream_scope_allows(cast(object, data), event_type) # cast-ok: data param
return False
if (
@ -1167,8 +1232,9 @@ class CustomGuardrail(CustomLogger):
)
result = EnterpriseCustomGuardrailHelper._should_run_if_mode_by_tag(data, self.event_hook, event_type)
if result is not None:
return result
return True
mode_tag_result: Final[bool] = bool(cast(object, result)) # cast-ok: helper return
return mode_tag_result and self.stream_scope_allows(cast(object, data), event_type) # cast-ok: data
return self.stream_scope_allows(cast(object, data), event_type) # cast-ok: data param
def _event_hook_is_event_type(self, event_type: GuardrailEventHooks) -> bool:
"""

View file

@ -91,13 +91,32 @@ def span_attribute_limit(span: Span) -> int | None:
return span._limits.max_span_attributes # pyright: ignore[reportPrivateUsage] # SDK has no public getter
def attribute_budget(span: Span, reserved: int) -> int | None:
"""How many mapped attributes fit on ``span`` next to what it already carries and ``reserved`` more."""
def _carried_keys(
attributes: Mapping[str, AttrValue] | tuple[tuple[str, AttrValue], ...],
) -> frozenset[str]:
"""The keys a span already carries: a live span exposes a ``Mapping``, an ended one a tuple of pairs."""
if isinstance(attributes, Mapping):
return frozenset(attributes)
return frozenset(key for key, _value in attributes)
def attribute_budget(span: Span, reserved: int, overwrites: frozenset[str] = frozenset()) -> int | None:
"""How many mapped attributes fit on ``span`` next to what it already carries and ``reserved`` more.
``overwrites`` are the mapped keys already present on ``span``: setting one
replaces the value in place and consumes no slot against the limit, so only
the genuinely new pre-existing keys reduce the budget. Counting the
overwritten ones too reserves slots the fit can never spend and sheds
indexed message attributes for nothing.
"""
limit: Final = span_attribute_limit(span)
if limit is None:
return None
on_span: Final = len(span.attributes or ()) if isinstance(span, ReadableSpan) else 0
return limit - on_span - reserved
if not isinstance(span, ReadableSpan):
return limit - reserved
existing_keys: Final = _carried_keys(span.attributes or ())
fresh: Final = len(existing_keys - overwrites) if overwrites else len(existing_keys)
return limit - fresh - reserved
def stamp_error(
@ -284,7 +303,7 @@ class SpanEmitter:
)
stamped_later: Final = error_attributes(error) if error else _NO_ATTRIBUTES
reserved: Final = len(stamped_later.keys() - mapped.keys())
for key, value in fit_indexed_messages(mapped, attribute_budget(span, reserved)).items():
for key, value in fit_indexed_messages(mapped, attribute_budget(span, reserved, frozenset(mapped))).items():
span.set_attribute(key, value)
if error:
stamped: Final = stamp_error(span, error)

View file

@ -575,6 +575,7 @@ class OpenTelemetryV2(CustomLogger):
request_purpose=call.purpose,
trace=call.trace,
session_id=call.session_id,
metadata_keys=tuple(self.config.baggage_metadata_keys),
)
end_time_ns: Final = to_ns(end_time)
if carrier is not None and carrier.span is not None:

View file

@ -7,7 +7,7 @@ Phoenix + any other OpenInference-aware backend simultaneously.
"""
import json
from collections.abc import Callable, Mapping, Sequence
from collections.abc import Callable, Iterator, Mapping, Sequence
from itertools import accumulate, chain, groupby
from types import MappingProxyType
from typing import Final
@ -15,10 +15,12 @@ from typing import Final
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
from litellm.integrations.otel.mappers.utils import (
MAX_TOOL_DEFINITION_ATTRS_PER_SPAN,
MessageToolCall,
collect,
drop_none,
drop_none_pairs,
json_if,
message_content,
message_tool_calls,
output_messages,
tool_definition_attrs,
)
@ -31,49 +33,146 @@ from litellm.integrations.otel.model.payloads import (
_INPUT_MESSAGES: Final = "llm.input_messages"
_OUTPUT_MESSAGES: Final = "llm.output_messages"
_MESSAGE_FAMILIES: Final = (_INPUT_MESSAGES, _OUTPUT_MESSAGES)
_MESSAGE_BASE: Final = -1
_ParsedMessage = tuple[object, str | None, tuple[MessageToolCall, ...]]
def _message_key_groups(attrs: Mapping[str, AttrValue]) -> Mapping[tuple[str, int], tuple[str, ...]]:
"""Per-index message keys in ``attrs`` grouped by ``(family, index)``."""
tagged: Final = sorted(
(family, int(key.split(".")[2]), key)
for key in attrs
for family in _MESSAGE_FAMILIES
if key.startswith(f"{family}.")
def _parse_message(message: object) -> _ParsedMessage:
role: Final = message.get("role") if isinstance(message, dict) else None
return role, message_content(message), message_tool_calls(message)
def _tool_call_attribute_pairs(
prefix: str, idx: int, tool_calls: tuple[MessageToolCall, ...]
) -> Iterator[tuple[str, str | None]]:
for tool_idx, tool_call in enumerate(tool_calls):
yield f"{prefix}.{idx}.message.tool_calls.{tool_idx}.tool_call.id", tool_call.id
yield f"{prefix}.{idx}.message.tool_calls.{tool_idx}.tool_call.function.name", tool_call.name
yield f"{prefix}.{idx}.message.tool_calls.{tool_idx}.tool_call.function.arguments", tool_call.arguments
def _message_attribute_pairs(
prefix: str,
messages: Sequence[_ParsedMessage],
*,
with_tool_call_attrs: bool,
) -> Iterator[tuple[str, str | None]]:
for idx, (role, content, tool_calls) in enumerate(messages):
yield f"{prefix}.{idx}.message.role", role if isinstance(role, str) else None
yield f"{prefix}.{idx}.message.content", content
if with_tool_call_attrs:
yield from _tool_call_attribute_pairs(prefix, idx, tool_calls)
def _message_value(messages: Sequence[_ParsedMessage]) -> str:
return json.dumps(
[
{
"role": role,
"content": content,
**({"tool_calls": [tool_call.to_openai_dict() for tool_call in tool_calls]} if tool_calls else {}),
}
for role, content, tool_calls in messages
]
)
def _message_key_group(key: str) -> tuple[str, int, int, str] | None:
family: Final = next(
(family for family in _MESSAGE_FAMILIES if key.startswith(f"{family}.")),
None,
)
if family is None:
return None
parts: Final = key.split(".")
message_idx: Final = int(parts[2])
tool_idx: Final = int(parts[5]) if parts[4] == "tool_calls" else _MESSAGE_BASE
return family, message_idx, tool_idx, key
def _message_key_groups(attrs: Mapping[str, AttrValue]) -> Mapping[tuple[str, int, int], tuple[str, ...]]:
"""Message and tool-call keys in ``attrs`` grouped by family, message index, and tool index."""
tagged: Final = tuple(tag for key in attrs if (tag := _message_key_group(key)) is not None)
return MappingProxyType(
{group: tuple(key for _, _, _, key in keys) for group, keys in groupby(sorted(tagged), key=lambda tag: tag[:3])}
)
def _tool_call_groups_by_message(
groups: Mapping[tuple[str, int, int], tuple[str, ...]],
) -> Mapping[tuple[str, int], tuple[tuple[str, int, int], ...]]:
"""Tool-call groups indexed by ``(family, message index)``, each tuple highest tool index first.
Indexing once keeps the shed order linear in the group count: rescanning the
full group map per message made attribute fitting quadratic on long prompts.
"""
ordered: Final = sorted(
(group for group in groups if group[2] != _MESSAGE_BASE),
key=lambda group: (group[0], group[1], group[2]),
)
return MappingProxyType(
{group: tuple(key for _, _, key in keys) for group, keys in groupby(tagged, key=lambda tag: tag[:2])}
{
message: tuple(reversed(tuple(message_tool_groups)))
for message, message_tool_groups in groupby(ordered, key=lambda group: (group[0], group[1]))
}
)
def _shed_order(groups: Mapping[tuple[str, int], tuple[str, ...]]) -> tuple[tuple[str, int], ...]:
"""Message groups least valuable first: middle prompt turns, extra choices, then the opener, the newest turn
and the first choice."""
inputs: Final = sorted(idx for family, idx in groups if family == _INPUT_MESSAGES)
outputs: Final = sorted(idx for family, idx in groups if family == _OUTPUT_MESSAGES)
def _message_shed_groups(
groups: Mapping[tuple[str, int, int], tuple[str, ...]],
tool_call_groups: Mapping[tuple[str, int], tuple[tuple[str, int, int], ...]],
family: str,
message_idx: int,
) -> Iterator[tuple[str, int, int]]:
yield from tool_call_groups.get((family, message_idx), ())
base_group: Final = (family, message_idx, _MESSAGE_BASE)
if base_group in groups:
yield base_group
def _shed_order(groups: Mapping[tuple[str, int, int], tuple[str, ...]]) -> tuple[tuple[str, int, int], ...]:
"""Middle inputs, extra choices, pinned inputs, then the first choice, with tool calls before message keys."""
tool_call_groups: Final = _tool_call_groups_by_message(groups)
inputs: Final = sorted(frozenset(idx for family, idx, _ in groups if family == _INPUT_MESSAGES))
outputs: Final = sorted(frozenset(idx for family, idx, _ in groups if family == _OUTPUT_MESSAGES))
pinned_inputs: Final = tuple(dict.fromkeys((*inputs[:1], *inputs[-1:])))
return (
message_order: Final = (
*((_INPUT_MESSAGES, idx) for idx in inputs[1:-1]),
*((_OUTPUT_MESSAGES, idx) for idx in reversed(outputs[1:])),
*((_INPUT_MESSAGES, idx) for idx in pinned_inputs),
*((_OUTPUT_MESSAGES, idx) for idx in outputs[:1]),
)
return tuple(
chain.from_iterable(
_message_shed_groups(groups, tool_call_groups, family, message_idx) for family, message_idx in message_order
)
)
_METADATA_KEY: Final = "metadata"
def fit_indexed_messages(attrs: Mapping[str, AttrValue], budget: int | None) -> Mapping[str, AttrValue]:
"""``attrs`` with whole per-index messages shed, least valuable first, until at most ``budget`` keys remain.
"""``attrs`` with indexed message attributes shed, least valuable first, until at most ``budget`` keys remain.
``None`` means the span has no attribute count limit. Every message still rides the ``input.value`` and
``output.value`` blobs, so shedding a per-index pair loses no content.
``output.value`` blobs, so shedding a per-index pair loses no content. The ``metadata`` blob sheds only
after every indexed message attribute: message attributes are the indexed, queryable view (the blobs
carry no per-index keys), so a count-squeezed span keeps them and the single metadata key absorbs only
the residual shortfall. Shedding happens here, before ``span.set_attribute``, so the fit is exact and
the SDK's dropped-attributes counter never silently masks the choice.
"""
if budget is None or len(attrs) <= budget:
return attrs
groups: Final = _message_key_groups(attrs)
order: Final = _shed_order(groups)
running: Final = tuple(accumulate(len(groups[group]) for group in order))
sheddable: Final = (*(groups[group] for group in _shed_order(groups)),)
flex: Final = (_METADATA_KEY,) if _METADATA_KEY in attrs else ()
candidates: Final = (*sheddable, flex) if flex else sheddable
running: Final = tuple(accumulate(len(candidate) for candidate in candidates))
excess: Final = len(attrs) - budget
shed_count: Final = next((n + 1 for n, total in enumerate(running) if total >= excess), len(order))
shed: Final = frozenset(chain.from_iterable(groups[group] for group in order[:shed_count]))
shed_count: Final = next((n + 1 for n, total in enumerate(running) if total >= excess), len(candidates))
shed: Final = frozenset(chain.from_iterable(candidates[:shed_count]))
return MappingProxyType({key: value for key, value in attrs.items() if key not in shed})
@ -85,7 +184,9 @@ class OpenInferenceMapper:
- ``llm.model_name`` / ``llm.provider`` / ``llm.invocation_parameters``
- ``llm.input_messages.{i}.message.role`` / ``...content``
- ``llm.output_messages.{i}.message.role`` / ``...content``
- ``llm.output_messages.{i}.message.tool_calls.{j}.tool_call.*``
- ``llm.token_count.prompt`` / ``...completion`` / ``...total``
- ``metadata`` — JSON object of allowlisted promoted request metadata
- ``input.value`` / ``output.value`` — JSON-serialized request / response
"""
@ -121,6 +222,7 @@ class OpenInferenceMapper:
"llm.invocation_parameters": lambda d: json_if(
collect(OpenInferenceMapper._INVOCATION_PARAMS, d.request_params)
),
"metadata": lambda d: json_if(dict(sorted(d.promoted_metadata.items()))),
}
def __init__(self, tool_attr_budget: int = MAX_TOOL_DEFINITION_ATTRS_PER_SPAN) -> None:
@ -137,27 +239,36 @@ class OpenInferenceMapper:
return {
**collect(self._LLM_CALL_ATTRS, data),
**collect(self._BLOB_ATTRS, data),
**self._messages(_INPUT_MESSAGES, "input.value", data.messages_in),
**self._messages(_OUTPUT_MESSAGES, "output.value", output_messages(data)),
**self._messages(
_INPUT_MESSAGES,
"input.value",
data.messages_in,
with_tool_call_attrs=False,
),
**self._messages(
_OUTPUT_MESSAGES,
"output.value",
output_messages(data),
with_tool_call_attrs=True,
),
**self._tools(data),
}
@staticmethod
def _messages(prefix: str, value_key: str, messages: Sequence[object]) -> AttributeMap:
def _messages(
prefix: str,
value_key: str,
messages: Sequence[object],
*,
with_tool_call_attrs: bool,
) -> AttributeMap:
"""``{prefix}.{idx}.message.*`` keys for every message + the ``value_key`` blob of all of them."""
parsed: Final = [(m.get("role") if isinstance(m, dict) else None, message_content(m)) for m in messages]
attrs: Final = drop_none(
{
key: value
for idx, (role, content) in enumerate(parsed)
for key, value in (
(f"{prefix}.{idx}.message.role", role if isinstance(role, str) else None),
(f"{prefix}.{idx}.message.content", content),
)
}
parsed: Final = tuple(_parse_message(message) for message in messages)
attrs: Final = drop_none_pairs(
_message_attribute_pairs(prefix, parsed, with_tool_call_attrs=with_tool_call_attrs)
)
if parsed:
attrs[value_key] = json.dumps([{"role": role, "content": content} for role, content in parsed])
attrs[value_key] = _message_value(parsed)
return attrs
def _tools(self, data: LLMCallSpanData) -> AttributeMap:

View file

@ -7,11 +7,28 @@ they live in one place.
import json
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from typing import Final
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue
from litellm.integrations.otel.model.payloads import LLMCallSpanData, ToolDefinition
@dataclass(frozen=True, slots=True)
class MessageToolCall:
id: str | None
type: str
name: str | None
arguments: str | None
def to_openai_dict(self) -> Mapping[str, object]:
return {
"id": self.id,
"type": self.type,
"function": {"name": self.name, "arguments": self.arguments},
}
DEFAULT_SPAN_ATTRIBUTE_LIMIT: Final = 128
"""The OTel SDK's default per-span attribute count limit."""
@ -121,3 +138,50 @@ def message_content(message: object) -> str | None:
def output_messages(data: LLMCallSpanData) -> list:
"""The ``message`` payload of each response choice."""
return [c.get("message") for c in data.choices_out if isinstance(c, dict)]
def message_tool_calls(message: object) -> tuple[MessageToolCall, ...]:
if not isinstance(message, dict):
return ()
tool_calls: Final = message.get("tool_calls")
if not isinstance(tool_calls, (list, tuple)) or not tool_calls:
return ()
return tuple(tool_call for value in tool_calls if (tool_call := _message_tool_call(value)) is not None)
def _message_tool_call(value: object) -> MessageToolCall | None:
if not isinstance(value, dict):
return None
function: Final = value.get("function")
function_data: Final = function if isinstance(function, dict) else {}
raw_type: Final = value.get("type")
raw_arguments: Final = function_data.get("arguments")
arguments: Final = (
raw_arguments
if isinstance(raw_arguments, str)
else _stringify_tool_arguments(raw_arguments)
if raw_arguments is not None
else None
)
name: Final = function_data.get("name")
identifier: Final = value.get("id")
return MessageToolCall(
id=identifier if isinstance(identifier, str) else None,
type=raw_type if isinstance(raw_type, str) else "function",
name=name if isinstance(name, str) else None,
arguments=arguments,
)
def _stringify_tool_arguments(value: object) -> str:
"""Serialize non-string tool-call arguments, falling back to ``repr``.
Arguments normally arrive as already-JSON strings, but provider adapters and
``model_construct`` responses hand over raw Python objects. ``json.dumps``
raises on those (tuple-keyed dicts, cycles) and the escaping exception would
lose the whole span, so keep a readable ``repr`` instead.
"""
try:
return json.dumps(value, default=str)
except (TypeError, ValueError):
return repr(value)

View file

@ -18,7 +18,7 @@ from collections.abc import Callable, Mapping
from types import MappingProxyType
from typing import Final
from litellm.integrations.otel.model.metadata import REQUESTER_METADATA_PATH, RequestIdentity
from litellm.integrations.otel.model.metadata import RequestIdentity, allowlisted_metadata
from litellm.integrations.otel.model.semconv import GenAI, LiteLLM
# Attribute key -> value extractor over (identity, request_model,
@ -92,9 +92,8 @@ def promoted_metadata(metadata: Mapping[str, str], metadata_keys: tuple[str, ...
"""Allowlisted entries of a flattened metadata mapping under ``litellm.metadata.*``."""
return MappingProxyType(
{
f"{LiteLLM.METADATA_PREFIX}{meta_key.removeprefix(REQUESTER_METADATA_PATH)}": value
for meta_key in metadata_keys
if (value := metadata.get(meta_key))
f"{LiteLLM.METADATA_PREFIX}{meta_key}": value
for meta_key, value in allowlisted_metadata(metadata, metadata_keys).items()
}
)

View file

@ -61,6 +61,16 @@ REQUESTER_METADATA_KEY: Final = "requester_metadata"
REQUESTER_METADATA_PATH: Final = f"{REQUESTER_METADATA_KEY}."
def allowlisted_metadata(metadata: Mapping[str, str], metadata_keys: tuple[str, ...]) -> Mapping[str, str]:
return MappingProxyType(
{
meta_key.removeprefix(REQUESTER_METADATA_PATH): value
for meta_key in metadata_keys
if (value := metadata.get(meta_key))
}
)
@dataclass(frozen=True)
class RequestIdentity:
call_id: str | None = None

View file

@ -12,7 +12,7 @@ from urllib.parse import urlsplit
from typing_extensions import ReadOnly, TypedDict
from litellm.integrations.otel.model.metadata import RequestContext, RequestIdentity
from litellm.integrations.otel.model.metadata import RequestContext, RequestIdentity, allowlisted_metadata
from litellm.integrations.otel.model.semconv import (
GenAIOperation,
GenAIOutputType,
@ -61,6 +61,8 @@ if TYPE_CHECKING:
StandardLoggingPayload,
)
_EMPTY_METADATA: Final[Mapping[str, str]] = MappingProxyType({})
# --- typed sub-structures ---------------------------------------------------- #
@ -437,6 +439,7 @@ class LLMCallSpanData:
trace: TraceControls = field(default_factory=TraceControls)
session_id: str | None = None
embedding_output: EmbeddingOutput | None = None
promoted_metadata: Mapping[str, str] = field(default_factory=lambda: _EMPTY_METADATA)
routing_attributes: Mapping[str, RoutingAttributeValue] = field(default_factory=lambda: MappingProxyType({}))
@classmethod
@ -449,6 +452,8 @@ class LLMCallSpanData:
request_purpose: str | None = None,
trace: TraceControls | None = None,
session_id: str | None = None,
*,
metadata_keys: tuple[str, ...] = (),
) -> LLMCallSpanData:
params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {})
# The single parse of the request's metadata — the request-vs-provider
@ -487,6 +492,7 @@ class LLMCallSpanData:
cost=LLMCost.from_breakdown(cast("Mapping[str, object] | None", payload.get("cost_breakdown"))),
server=ServerInfo.from_api_base(context.api_base),
identity=context.identity,
promoted_metadata=allowlisted_metadata(context.identity.metadata, metadata_keys),
is_streaming=as_bool(payload.get("stream")),
tools=_extract_tools(params),
messages_in=_dicts(payload.get("messages")) if capture_content else (),

View file

@ -782,7 +782,7 @@ class RealTimeStreaming:
isinstance(cb, CustomGuardrail)
and any(
cb.should_run_guardrail(
data=self.request_data,
data={**self.request_data, "stream": True},
event_type=et,
)
for et in event_hooks
@ -847,7 +847,7 @@ class RealTimeStreaming:
if event_hooks is None:
event_hooks = [GuardrailEventHooks.realtime_input_transcription]
_realtime_event_types: Final = event_hooks
_check_data: Final = {**self.request_data, "transcript": transcript}
_check_data: Final = {**self.request_data, "transcript": transcript, "stream": True}
_already_run: Final[set] = set()
for callback in litellm.callbacks:

View file

@ -1,247 +0,0 @@
"""The Jev / System One wire shape and its translation to and from the OpenAI Decisions shape.
System One (TypeSafe, Perplexity, OpenRouter, Cloudflare Clef, Strands Decider) takes
{"model", "state", "questions": {name: question}} and answers with {"model", "answers": {name: answer}, "usage"}.
Predicates are `noul` questions, choice options are a `criteria` map, score levels are a `criteria` list.
"""
import itertools
from collections.abc import Mapping, Sequence
from typing import Final, Literal, TypeAlias
from pydantic import ConfigDict, TypeAdapter
from typing_extensions import assert_never
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.base import LiteLLMPydanticObjectBase
from litellm.types.openai_decisions import (
ChoiceAnswer,
ChoiceProbability,
ChoiceQuestion,
DecisionAnswer,
DecisionChoice,
DecisionInput,
DecisionInputMessage,
DecisionInputPart,
DecisionInputTokensDetails,
DecisionOutputTokensDetails,
DecisionQuestion,
DecisionsRequest,
DecisionsRequestBody,
DecisionsResponse,
DecisionUsage,
PredicateAnswer,
PredicateQuestion,
ScoreAnswer,
ScoreProbability,
ScoreQuestion,
)
class SystemOneObjectBase(LiteLLMPydanticObjectBase):
model_config = ConfigDict(extra="allow", frozen=True)
class SystemOneNoulAnswer(SystemOneObjectBase):
type: Literal["noul"]
noul: float
class SystemOneChoiceAnswer(SystemOneObjectBase):
type: Literal["choice"]
choice: str
confidence: float
probabilities: Mapping[str, float]
class SystemOneScoreAnswer(SystemOneObjectBase):
type: Literal["score"]
score: float
confidence: float
probabilities: Mapping[str, float]
SystemOneAnswer: TypeAlias = SystemOneNoulAnswer | SystemOneChoiceAnswer | SystemOneScoreAnswer
class SystemOneUsage(SystemOneObjectBase):
input_tokens: int = 0
output_tokens: int = 0
class SystemOneResponse(SystemOneObjectBase):
model: str | None = None
answers: Mapping[str, SystemOneAnswer]
usage: SystemOneUsage | None = None
SYSTEM_ONE_RESPONSE_ADAPTER: Final[TypeAdapter[SystemOneResponse]] = TypeAdapter(SystemOneResponse)
def _unsupported(what: str, custom_llm_provider: str) -> BaseLLMException:
return BaseLLMException(
status_code=400,
message=f"Decisions provider '{custom_llm_provider}' does not support {what}",
)
def to_system_one_request(model: str, body: DecisionsRequestBody, custom_llm_provider: str) -> dict[str, object]:
keys: Final = question_keys(body.questions, custom_llm_provider)
return {
"model": model,
"state": _state(body.input, custom_llm_provider),
"questions": {
key: _question(question, custom_llm_provider) for key, question in zip(keys, body.questions, strict=True)
},
}
def question_keys(questions: Sequence[DecisionQuestion], custom_llm_provider: str) -> tuple[str, ...]:
"""System One keys questions and answers by name, so unnamed questions get a positional key."""
names: Final = tuple(question.name for question in questions if question.name is not None)
if len(set(names)) != len(names):
raise BaseLLMException(
status_code=400,
message=f"Decisions provider '{custom_llm_provider}' requires a unique name per question",
)
taken: Final = frozenset(names)
return tuple(
question.name if question.name is not None else _positional_key(index, taken)
for index, question in enumerate(questions)
)
def _positional_key(index: int, taken: frozenset[str]) -> str:
candidates: Final = (f"{'_' * depth}q{index}" for depth in itertools.count())
return next(key for key in candidates if key not in taken)
def _state(input_value: DecisionInput, custom_llm_provider: str) -> str:
if isinstance(input_value, str):
return input_value
return "\n".join(_message_text(message, custom_llm_provider) for message in input_value)
def _message_text(message: DecisionInputMessage, custom_llm_provider: str) -> str:
if isinstance(message.content, str):
return message.content
return "\n".join(_part_text(part, custom_llm_provider) for part in message.content)
def _part_text(part: DecisionInputPart, custom_llm_provider: str) -> str:
if part.type != "input_text":
raise _unsupported("input_image parts", custom_llm_provider)
return part.text
def _question(question: DecisionQuestion, custom_llm_provider: str) -> dict[str, object]:
match question:
case PredicateQuestion():
return {"type": "noul", "instructions": question.instructions}
case ChoiceQuestion():
return {
"type": "choice",
"instructions": question.instructions,
"criteria": _choice_criteria(question, custom_llm_provider),
}
case ScoreQuestion():
return {
"type": "score",
"instructions": question.instructions,
"criteria": [_level_text(level.label, level.description) for level in question.levels],
}
case _:
assert_never(question)
def _choice_criteria(question: ChoiceQuestion, custom_llm_provider: str) -> dict[str, str | None]:
values: Final = tuple(_choice_key(choice, custom_llm_provider) for choice in question.choices)
if len(set(values)) != len(values):
raise _unsupported("repeated choice values", custom_llm_provider)
return {value: choice.description for value, choice in zip(values, question.choices, strict=True)}
def _choice_key(choice: DecisionChoice, custom_llm_provider: str) -> str:
if not isinstance(choice.value, str):
raise _unsupported("boolean choice values", custom_llm_provider)
return choice.value
def _level_text(label: str, description: str | None) -> str:
return description if description is not None else label
def to_decisions_response(
system_one: SystemOneResponse,
request: DecisionsRequest,
custom_llm_provider: str,
) -> DecisionsResponse:
keys: Final = question_keys(request.body.questions, custom_llm_provider)
return DecisionsResponse(
model=system_one.model if system_one.model is not None else request.model,
answers=[
_answer(key, question, system_one.answers, custom_llm_provider)
for key, question in zip(keys, request.body.questions, strict=True)
],
usage=_usage(system_one.usage),
)
def _answer(
key: str,
question: DecisionQuestion,
answers: Mapping[str, SystemOneAnswer],
custom_llm_provider: str,
) -> DecisionAnswer:
match question, answers.get(key):
case PredicateQuestion(), SystemOneNoulAnswer() as answer:
return PredicateAnswer(type="predicate", name=question.name, probability=answer.noul)
case ChoiceQuestion(), SystemOneChoiceAnswer() as answer:
return ChoiceAnswer(
type="choice",
name=question.name,
choice=answer.choice,
probabilities=_choice_probabilities(question, answer),
confidence=answer.confidence,
)
case ScoreQuestion(), SystemOneScoreAnswer() as answer:
return ScoreAnswer(
type="score",
name=question.name,
score=answer.score,
probabilities=_score_probabilities(question, answer),
confidence=answer.confidence,
)
case _:
raise BaseLLMException(
status_code=500,
message=(
f"Decisions provider '{custom_llm_provider}' returned no {question.type} answer "
f"for question '{key}'"
),
)
def _choice_probabilities(question: ChoiceQuestion, answer: SystemOneChoiceAnswer) -> list[ChoiceProbability]:
return [
ChoiceProbability(value=choice.value, probability=answer.probabilities.get(str(choice.value), 0.0))
for choice in question.choices
]
def _score_probabilities(question: ScoreQuestion, answer: SystemOneScoreAnswer) -> list[ScoreProbability]:
return [
ScoreProbability(value=index, label=level.label, probability=answer.probabilities.get(str(index), 0.0))
for index, level in enumerate(question.levels)
]
def _usage(usage: SystemOneUsage | None) -> DecisionUsage:
counted: Final = usage if usage is not None else SystemOneUsage()
return DecisionUsage(
input_tokens=counted.input_tokens,
input_tokens_details=DecisionInputTokensDetails(cached_tokens=0, cache_write_tokens=0),
output_tokens=counted.output_tokens,
output_tokens_details=DecisionOutputTokensDetails(reasoning_tokens=0),
total_tokens=counted.input_tokens + counted.output_tokens,
)

View file

@ -36,6 +36,7 @@ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token
from litellm.llms.bedrock.common_utils import (
BedrockError,
bedrock_model_is_openai_gpt,
bedrock_runtime_chat_completions_serves_reasoning_inline,
split_bedrock_region_path,
)
from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
@ -213,7 +214,12 @@ def split_reasoning_tag(content: str) -> tuple[str | None, str]:
class BedrockRuntimeChatCompletionsStreamingHandler(OpenAIChatCompletionStreamingHandler):
"""OpenAI chunk parsing plus the ``<reasoning>`` split, tracked per choice index."""
"""OpenAI chunk parsing plus the inline ``<reasoning>`` split, tracked per choice index.
Every chunk echoes the model id litellm sent, so the split engages only when that id's price-map
row carries ``supports_bedrock_runtime_chat_completions_inline_reasoning`` (gpt-oss); a GPT 5.6 or Grok
answer that starts with a literal ``<reasoning>`` tag streams as content.
"""
def __init__(
self,
@ -226,6 +232,8 @@ class BedrockRuntimeChatCompletionsStreamingHandler(OpenAIChatCompletionStreamin
def chunk_parser(self, chunk: dict) -> ModelResponseStream: # mutable-ok: BaseModelResponseIterator signature
parsed: Final = super().chunk_parser(chunk)
if not bedrock_runtime_chat_completions_serves_reasoning_inline(parsed.model or ""):
return parsed
for choice in parsed.choices:
next_state, reasoning, content = _split_streamed_content(
self._splitters.get(choice.index, ReasoningTagSplitter()),
@ -475,6 +483,8 @@ class AmazonBedrockRuntimeChatCompletionsConfig(OpenAILikeChatConfig):
json_mode=json_mode,
)
set_provider_response_headers_in_hidden_params(response, raw_response.headers)
if not bedrock_runtime_chat_completions_serves_reasoning_inline(model):
return response
for choice in response.choices:
if not isinstance(choice, Choices) or not isinstance(choice.message.content, str):
continue

View file

@ -909,6 +909,17 @@ def bedrock_model_is_openai_gpt(model: str) -> bool:
return _openai_gpt_version(model) is not None
def bedrock_runtime_chat_completions_serves_reasoning_inline(model: str) -> bool:
"""Whether AWS's native Chat Completions writes this model's reasoning inline in the answer text.
Data-driven from the price-map ``supports_bedrock_runtime_chat_completions_inline_reasoning`` flag (gpt-oss).
A flagged model opens its answer with a ``<reasoning>...</reasoning>`` block instead of a
``reasoning_content`` field, so litellm splits that block out for it and keeps every other model's
text as sent.
"""
return _bedrock_price_map_flag(model, "supports_bedrock_runtime_chat_completions_inline_reasoning")
BEDROCK_CONVERSE_ONLY_REQUEST_KEYS: Final = frozenset(
(
"guardrailConfig",

View file

@ -3,27 +3,100 @@ Bedrock Token Counter implementation using the CountTokens API.
"""
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Final
from pydantic import JsonValue
from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.types.utils import LlmProviders, TokenCountResponse
from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler
from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountResponse
RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api"
MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api"
class _CountTokensReply(LiteLLMPydanticObjectBase):
input_tokens: int
@dataclass(frozen=True, slots=True)
class CountedTokens:
input_tokens: int
original_response: Mapping[str, JsonValue]
tokenizer_type: str
@dataclass(frozen=True, slots=True)
class CountTokensFailure:
status_code: int
message: str
tokenizer_type: str
CountTokensOutcome = CountedTokens | CountTokensFailure
def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool:
return (
isinstance(outcome, CountTokensFailure)
and outcome.status_code == 400
and is_mantle_claude_model(resolved_model)
)
class BedrockTokenCounter(BaseTokenCounter):
"""Token counter implementation for AWS Bedrock provider using the CountTokens API."""
"""Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens
on bedrock-mantle for the Claude models bedrock-runtime answers 400 for"""
def __init__(
self,
runtime_handler: BedrockCountTokensHandler | None = None,
mantle_handler: BedrockMantleCountTokensHandler | None = None,
client: AsyncHTTPHandler | None = None,
) -> None:
self._runtime_handler: Final = runtime_handler or BedrockCountTokensHandler()
self._mantle_handler: Final = mantle_handler or BedrockMantleCountTokensHandler()
self._client: Final = client
def should_use_token_counting_api(
self,
custom_llm_provider: str | None = None,
) -> bool:
"""
Returns True if we should use the Bedrock CountTokens API for token counting.
"""
return custom_llm_provider == LlmProviders.BEDROCK.value
async def _count_with(
self,
handler: BedrockCountTokensHandler | BedrockMantleCountTokensHandler,
tokenizer_type: str,
request_data: dict[str, object],
litellm_params: dict[str, object],
resolved_model: str,
) -> CountTokensOutcome:
try:
result: Final = await handler.handle_count_tokens_request(
request_data=request_data,
litellm_params=litellm_params,
resolved_model=resolved_model,
client=self._client,
)
reply: Final = _CountTokensReply.model_validate(result)
except BedrockError as e:
verbose_logger.debug(
"%s CountTokens API error: status=%s, message=%s", tokenizer_type, e.status_code, e.message
)
return CountTokensFailure(status_code=e.status_code, message=e.message, tokenizer_type=tokenizer_type)
except Exception as e:
verbose_logger.debug("Error calling %s CountTokens API: %s", tokenizer_type, e)
return CountTokensFailure(status_code=500, message=str(e), tokenizer_type=tokenizer_type)
return CountedTokens(input_tokens=reply.input_tokens, original_response=result, tokenizer_type=tokenizer_type)
async def count_tokens(
self,
model_to_use: str,
@ -34,81 +107,52 @@ class BedrockTokenCounter(BaseTokenCounter):
tools: Sequence[Mapping[str, object]] | None = None,
system: object | None = None,
) -> TokenCountResponse | None:
"""
Count tokens using AWS Bedrock's CountTokens API.
This method calls the existing BedrockCountTokensHandler to make an API call
to Bedrock's token counting endpoint, bypassing the local tiktoken-based counting.
Args:
model_to_use: The model identifier
messages: The messages to count tokens for
contents: Alternative content format (not used for Bedrock)
deployment: Deployment configuration containing litellm_params
request_model: The original request model name
Returns:
TokenCountResponse with token count, or None if counting fails
"""
if not messages:
return None
deployment = deployment or {}
litellm_params: Final = deployment.get("litellm_params", {})
# Build request data in the format expected by BedrockCountTokensHandler
litellm_params: Final = (deployment or {}).get("litellm_params", {})
request_data: Final[dict[str, object]] = {
"model": model_to_use,
"messages": messages,
**({"tools": tools} if tools else {}),
**({"system": system} if system else {}),
}
if tools:
request_data["tools"] = tools
if system:
request_data["system"] = system
# Get the resolved model (strip prefixes like bedrock/, converse/, etc.)
resolved_model: Final = get_bedrock_base_model(model_to_use)
try:
handler: Final = BedrockCountTokensHandler()
result: Final = await handler.handle_count_tokens_request(
request_data=request_data,
litellm_params=litellm_params,
resolved_model=resolved_model,
runtime: Final = await self._count_with(
self._runtime_handler, RUNTIME_TOKENIZER_TYPE, request_data, litellm_params, resolved_model
)
outcome: Final = (
await self._count_with(
self._mantle_handler, MANTLE_TOKENIZER_TYPE, request_data, litellm_params, resolved_model
)
# Transform response to TokenCountResponse
if result is not None:
if _runtime_rejected_claude_model(runtime, resolved_model)
else runtime
)
match outcome:
case CountedTokens():
return TokenCountResponse(
total_tokens=result.get("input_tokens", 0),
total_tokens=outcome.input_tokens,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="bedrock_api",
original_response=result,
tokenizer_type=outcome.tokenizer_type,
original_response=dict(outcome.original_response),
)
except BedrockError as e:
verbose_logger.warning("Bedrock CountTokens API error: status=%s, message=%s", e.status_code, e.message)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="bedrock_api",
error=True,
error_message=e.message,
status_code=e.status_code,
)
except Exception as e:
verbose_logger.warning("Error calling Bedrock CountTokens API: %s", e)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="bedrock_api",
error=True,
error_message=str(e),
status_code=500,
)
return None
case CountTokensFailure():
verbose_logger.warning(
"%s CountTokens API error: status=%s, message=%s",
outcome.tokenizer_type,
outcome.status_code,
outcome.message,
)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type=outcome.tokenizer_type,
error=True,
error_message=outcome.message,
status_code=outcome.status_code,
)
case _:
assert_never(outcome)

View file

@ -125,7 +125,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
raise
except httpx.HTTPStatusError as e:
# HTTP errors - preserve the actual status code
verbose_logger.error("HTTP error in CountTokens handler: %s", e)
verbose_logger.debug("HTTP error in CountTokens handler: %s", e)
raise BedrockError(
status_code=e.response.status_code,
message=e.response.text,

View file

@ -0,0 +1,99 @@
from typing import Final
import httpx
from pydantic import JsonValue, TypeAdapter
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.llms.anthropic.count_tokens.transformation import AnthropicCountTokensConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens"
MANTLE_ANTHROPIC_VERSION: Final = "2023-06-01"
class MantleCountTokensRequest(TypedDict):
messages: ReadOnly[list[dict[str, JsonValue]]]
system: ReadOnly[NotRequired[JsonValue]]
tools: ReadOnly[NotRequired[list[dict[str, JsonValue]]]]
_COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest)
_COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue])
class BedrockMantleCountTokensHandler(BaseAWSLLM):
"""Counts tokens through Anthropic's count_tokens on the bedrock-mantle endpoint.
Claude models that Bedrock offers only through cross-region inference answer 400 on
bedrock-runtime's CountTokens; AWS documents Mantle's /anthropic/v1/messages/count_tokens
as the way to count them, with the base model id and the deployment's AWS credentials
"""
def __init__(self, anthropic_config: AnthropicCountTokensConfig | None = None) -> None:
super().__init__()
self._anthropic_config: Final = anthropic_config or AnthropicCountTokensConfig()
def get_mantle_count_tokens_endpoint(self, aws_region_name: str) -> str:
messages_url: Final = build_mantle_messages_url(
api_base=None, aws_bedrock_runtime_endpoint=None, region=aws_region_name
)
return f"{messages_url}{MANTLE_COUNT_TOKENS_SUFFIX}"
async def handle_count_tokens_request(
self,
request_data: dict[str, object],
litellm_params: dict[str, object],
resolved_model: str,
client: AsyncHTTPHandler | None = None,
) -> dict[str, JsonValue]:
try:
request: Final = _COUNT_REQUEST.validate_python(request_data)
aws_region_name: Final = self._get_aws_region_name(
optional_params=litellm_params, model=resolved_model, model_id=None
)
endpoint_url: Final = self.get_mantle_count_tokens_endpoint(aws_region_name)
body: Final = self._anthropic_config.transform_request_to_count_tokens(
model=resolved_model,
messages=request["messages"],
tools=request.get("tools"),
system=request.get("system"),
)
verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url)
api_key: Final = litellm_params.get("api_key")
signed_headers, signed_body = await run_aws_signing(
self._sign_request,
service_name="bedrock",
headers={"Content-Type": "application/json", "anthropic-version": MANTLE_ANTHROPIC_VERSION},
optional_params=litellm_params,
request_data=body,
api_base=endpoint_url,
model=resolved_model,
api_key=api_key if isinstance(api_key, str) else None,
)
async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
response: Final = await async_client.post(
endpoint_url, headers=signed_headers, data=signed_body, timeout=30.0
)
if response.status_code != 200:
raise BedrockError(
status_code=response.status_code,
message=response.text,
headers=response.headers,
response=response,
)
return _COUNT_RESPONSE.validate_json(response.content)
except BedrockError:
raise
except httpx.HTTPStatusError as e:
raise BedrockError(
status_code=e.response.status_code,
message=e.response.text,
headers=e.response.headers,
response=e.response,
)
except Exception as e:
raise BedrockError(status_code=500, message=f"bedrock-mantle count_tokens processing error: {e}")

View file

@ -22,6 +22,13 @@ if TYPE_CHECKING:
from litellm.types.utils import CostResponseTypes
BEDROCK_STREAMING_ACTIONS: Final = frozenset({"invoke-with-response-stream", "converse-stream"})
def is_bedrock_streaming_endpoint(endpoint: str) -> bool:
return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1] in BEDROCK_STREAMING_ACTIONS
_TEXT_ONLY_DELTA_FIELDS: Final = frozenset({"content", "role"})

View file

@ -10,6 +10,7 @@ from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Optional
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import without_server_streaming_classification
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.proxy._types import PassThroughGuardrailSettings
from litellm.types.utils import GenericGuardrailAPIInputs
@ -80,7 +81,9 @@ class PassThroughEndpointHandler(BaseTranslation):
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
payload_to_check: Final = {
k: v for k, v in data.items() if not k.startswith("_") and k not in ("metadata", "litellm_logging_obj")
k: v
for k, v in without_server_streaming_classification(data).items()
if not k.startswith("_") and k not in ("metadata", "litellm_logging_obj")
}
verbose_proxy_logger.debug("PassThroughEndpointHandler: Using full payload for guardrail")
return safe_dumps(payload_to_check)

View file

@ -25,6 +25,22 @@ from litellm.types.llms.vertex_ai import (
from litellm.types.utils import TokenCountResponse
from litellm.utils import supports_response_schema, supports_system_messages
GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType(
{"fps": "fps", "start_offset": "startOffset", "end_offset": "endOffset"}
)
GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/"
def gemini_video_metadata_from_openai(video_metadata: Mapping[str, object]) -> dict[str, object]:
return {
gemini_key: video_metadata[openai_key]
for openai_key, gemini_key in GEMINI_VIDEO_METADATA_KEYS.items()
if openai_key in video_metadata
}
VERTEX_SELF_DEPLOYED_ENDPOINT_UNSUPPORTED_PARAMS: Final = frozenset(
{
"audio",

View file

@ -67,7 +67,7 @@ from litellm.types.llms.openai import (
OpenAIFilesPurpose,
PathLike,
)
from litellm.types.llms.vertex_ai import GcsBucketResponse, GeminiEmbeddingInput
from litellm.types.llms.vertex_ai import GcsBucketResponse, GeminiEmbeddingElement, GeminiEmbeddingInput
from litellm.types.utils import (
Embedding,
EmbeddingResponse,
@ -555,18 +555,27 @@ def _is_responses_batch_entry(openai_entry: Mapping[str, object]) -> bool:
return path == "responses" or path.endswith("/responses")
def _own_embedding_input(
element: GeminiEmbeddingElement | list[str] | list[GeminiEmbeddingElement],
) -> GeminiEmbeddingInput:
if isinstance(element, (str, list)):
return element
file_block_alone: Final[list[GeminiEmbeddingElement]] = [element]
return file_block_alone
def _openai_embedding_input_elements(
embedding_input: GeminiEmbeddingInput,
) -> tuple[str | list[str], ...]:
) -> tuple[GeminiEmbeddingInput, ...]:
"""
Split an OpenAI `input` into the elements that each get their own embedding.
A string is one embedding, a flat array is one embedding per element, and a nested
array is one combined embedding per inner array, matching the online
`batchEmbedContents` path.
A string or a file content block is one embedding, a flat array is one embedding
per element, and a nested array is one combined embedding per inner array,
matching the online `batchEmbedContents` path.
"""
if isinstance(embedding_input, list):
return tuple(embedding_input)
return tuple(_own_embedding_input(element) for element in embedding_input)
return (embedding_input,)

View file

@ -57,7 +57,9 @@ from litellm.types.llms.vertex_ai import (
from litellm.types.utils import GenericImageParsingChunk, LlmProviders
from ..common_utils import (
GEMINI_FILES_API_URI_PREFIX,
check_text_in_content,
gemini_video_metadata_from_openai,
get_supports_response_schema,
get_supports_system_message,
)
@ -69,7 +71,6 @@ _GCS_METADATA_VERTEX_BASE: object | None = None
# Shared sync client for GCS JSON API metadata reads so proxy/SSL settings
# from litellm's HTTP stack apply (see Greptile review on PR #27278).
_GCS_METADATA_HTTP_HANDLER: HTTPHandler | None = None
GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/"
_GEMINI_MIME_TYPE_ALIASES: Final[dict[str, str]] = {
"image/jpg": "image/jpeg",
}
@ -197,13 +198,7 @@ def _apply_gemini_metadata(
part_dict["media_resolution"] = media_resolution_enum
if video_metadata is not None:
gemini_video_metadata: Final = {}
if "fps" in video_metadata:
gemini_video_metadata["fps"] = video_metadata["fps"]
if "start_offset" in video_metadata:
gemini_video_metadata["startOffset"] = video_metadata["start_offset"]
if "end_offset" in video_metadata:
gemini_video_metadata["endOffset"] = video_metadata["end_offset"]
gemini_video_metadata: Final = gemini_video_metadata_from_openai(video_metadata)
if gemini_video_metadata:
part_dict["video_metadata"] = gemini_video_metadata

View file

@ -22,6 +22,8 @@ from litellm.types.utils import EmbeddingResponse
from ..gemini.vertex_and_google_ai_studio_gemini import VertexLLM
from .batch_embed_content_transformation import (
file_reference_name,
flatten_media_sources,
is_file_reference,
process_embed_content_response,
process_response,
@ -38,11 +40,8 @@ class GoogleBatchEmbeddings(VertexLLM):
def _flatten_and_detect_file_refs(
input: GeminiEmbeddingInput,
) -> tuple[list[str], bool]:
"""Flatten nested input lists and detect file references."""
input_list: Final = [input] if isinstance(input, str) else input
flat_elements: Final = [
e for item in input_list for e in (item if isinstance(item, list) else [item]) if isinstance(e, str)
]
"""Flatten nested input lists and file content blocks into their sources and detect file references."""
flat_elements: Final = list(flatten_media_sources(input))
has_file_refs: Final = any(is_file_reference(e) for e in flat_elements)
return flat_elements, has_file_refs
@ -68,7 +67,7 @@ class GoogleBatchEmbeddings(VertexLLM):
for element in input_list:
if isinstance(element, str) and is_file_reference(element):
url = f"https://generativelanguage.googleapis.com/v1beta/{element}"
url = f"https://generativelanguage.googleapis.com/v1beta/{file_reference_name(element)}"
headers = {"x-goog-api-key": api_key}
response = sync_handler.get(url=url, headers=headers)
@ -105,7 +104,7 @@ class GoogleBatchEmbeddings(VertexLLM):
for element in input_list:
if isinstance(element, str) and is_file_reference(element):
url = f"https://generativelanguage.googleapis.com/v1beta/{element}"
url = f"https://generativelanguage.googleapis.com/v1beta/{file_reference_name(element)}"
headers = {"x-goog-api-key": api_key}
response = await async_handler.get(url=url, headers=headers)
@ -139,6 +138,7 @@ class GoogleBatchEmbeddings(VertexLLM):
timeout=300,
client=None,
extra_headers: dict | None = None,
drop_params: bool = False,
) -> EmbeddingResponse:
_auth_header, vertex_project = self._ensure_access_token(
credentials=vertex_credentials,
@ -207,22 +207,26 @@ class GoogleBatchEmbeddings(VertexLLM):
api_key=api_key,
optional_params=optional_params,
logging_obj=logging_obj,
drop_params=drop_params,
)
### TRANSFORMATION (sync path) ###
request_data: VertexAIBatchEmbeddingsRequestBody | dict[str, object]
flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input)
if use_embed_content:
resolved_files = {}
if api_key:
resolved_files = self._resolve_file_references(input=input, api_key=api_key, sync_handler=sync_handler)
resolved_files = self._resolve_file_references(
input=flat_elements, api_key=api_key, sync_handler=sync_handler
)
request_data = transform_openai_input_gemini_embed_content(
input=input,
model=model,
optional_params=optional_params,
resolved_files=resolved_files,
drop_params=drop_params,
)
else:
flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input)
if has_file_refs and not api_key:
raise ValueError(
"An API key is required to resolve Gemini file references (files/...). "
@ -238,6 +242,7 @@ class GoogleBatchEmbeddings(VertexLLM):
model=model,
optional_params=optional_params,
resolved_files=resolved_files,
drop_params=drop_params,
)
## LOGGING
@ -294,6 +299,7 @@ class GoogleBatchEmbeddings(VertexLLM):
api_key: str | None = None,
optional_params: dict | None = None,
logging_obj: "LiteLLMLoggingObj | None" = None,
drop_params: bool = False,
) -> EmbeddingResponse:
if client is None:
_params: Final = {}
@ -312,20 +318,21 @@ class GoogleBatchEmbeddings(VertexLLM):
async_handler = client
### TRANSFORMATION (async path) ###
flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input)
if use_embed_content:
resolved_files = {}
if api_key:
resolved_files = await self._async_resolve_file_references(
input=input, api_key=api_key, async_handler=async_handler
input=flat_elements, api_key=api_key, async_handler=async_handler
)
data = transform_openai_input_gemini_embed_content(
input=input,
model=model,
optional_params=optional_params or {},
resolved_files=resolved_files,
drop_params=drop_params,
)
else:
flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input)
if has_file_refs and not api_key:
raise ValueError(
"An API key is required to resolve Gemini file references (files/...). "
@ -341,6 +348,7 @@ class GoogleBatchEmbeddings(VertexLLM):
model=model,
optional_params=optional_params or {},
resolved_files=resolved_files,
drop_params=drop_params,
)
## LOGGING

View file

@ -4,22 +4,27 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc
Why separate file? Make it easy to see how transformation works
"""
from collections.abc import Mapping, Sequence
from typing import Final
from collections.abc import Iterator, Mapping, Sequence
from typing import Annotated, Final, Literal, cast
from pydantic import TypeAdapter, ValidationError
from pydantic import ConfigDict, Field, TypeAdapter, ValidationError
from litellm.exceptions import BadRequestError
from litellm.llms.vertex_ai.common_utils import GEMINI_FILES_API_URI_PREFIX, gemini_video_metadata_from_openai
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.vertex_ai import (
BlobType,
ContentType,
EmbedContentRequest,
FileDataType,
GeminiEmbeddingElement,
GeminiEmbeddingInput,
PartType,
PromptTokensDetails,
UsageMetadata,
VertexAIBatchEmbeddingsRequestBody,
VertexAIBatchEmbeddingsResponseObject,
VideoMetadataType,
)
from litellm.types.utils import (
Embedding,
@ -40,9 +45,16 @@ SUPPORTED_EMBEDDING_MIME_TYPES: Final = {
}
_GEMINI_API_V1BETA: Final = GEMINI_FILES_API_URI_PREFIX.removesuffix("files/")
def is_file_reference(s: str) -> bool:
"""Check if string is a Gemini file reference (files/...)."""
return isinstance(s, str) and s.startswith("files/")
"""A `files/...` name or the Files API URI that `/v1/files` returns as the file id."""
return isinstance(s, str) and (s.startswith("files/") or s.startswith(GEMINI_FILES_API_URI_PREFIX))
def file_reference_name(reference: str) -> str:
return reference.removeprefix(_GEMINI_API_V1BETA)
_is_file_reference = is_file_reference
@ -88,12 +100,13 @@ def _infer_mime_type_from_gcs_url(gcs_url: str) -> str:
)
def _parse_data_url(data_url: str) -> tuple[str, str]:
def _parse_data_url(data_url: str, mime_type_override: str | None = None) -> tuple[str, str]:
"""
Parse a data URL to extract the media type and base64 data.
Args:
data_url: Data URL in format: data:image/jpeg;base64,/9j/4AAQ...
mime_type_override: An explicit media type that replaces the declared one, skipping the allowlist
Returns:
tuple: (media_type, base64_data)
@ -110,50 +123,225 @@ def _parse_data_url(data_url: str) -> tuple[str, str]:
raise ValueError(f"Invalid data URL format (missing comma): {data_url[:50]}...")
metadata, base64_data = data_url.split(",", 1)
declared_media_type: Final = metadata[5:].split(";")[0]
metadata = metadata[5:]
if mime_type_override is not None:
return mime_type_override, base64_data
if ";" in metadata:
media_type = metadata.split(";")[0]
else:
media_type = metadata
if media_type not in SUPPORTED_EMBEDDING_MIME_TYPES:
if declared_media_type not in SUPPORTED_EMBEDDING_MIME_TYPES:
raise ValueError(
f"Unsupported MIME type for embedding: {media_type}. "
f"Unsupported MIME type for embedding: {declared_media_type}. "
f"Supported types: {', '.join(sorted(SUPPORTED_EMBEDDING_MIME_TYPES))}"
)
return media_type, base64_data
return declared_media_type, base64_data
def _is_data_url(s: str) -> bool:
return s.startswith("data:") and ";base64," in s
class _EmbeddingVideoMetadata(LiteLLMBaseModel):
model_config = ConfigDict(extra="forbid", strict=True)
fps: float | None = None
start_offset: str | None = None
end_offset: str | None = None
class _EmbeddingFile(LiteLLMBaseModel):
model_config = ConfigDict(extra="forbid")
file_id: str | None = None
file_data: str | None = None
filename: str | None = None
format: Annotated[str, Field(min_length=1)] | None = None
video_metadata: _EmbeddingVideoMetadata | None = None
class _EmbeddingFileBlock(LiteLLMBaseModel):
model_config = ConfigDict(extra="forbid")
type: Literal["file"]
file: _EmbeddingFile
_file_block_adapter: Final = TypeAdapter(_EmbeddingFileBlock)
_video_metadata_adapter: Final = TypeAdapter(VideoMetadataType)
_input_shape_adapter: Final[TypeAdapter[str | list[object]]] = TypeAdapter(str | list[object])
_mapping_adapter: Final = TypeAdapter(dict[str, object])
_BLOCK_FIELDS: Final = frozenset(_EmbeddingFileBlock.model_fields)
_FILE_FIELDS: Final = frozenset(_EmbeddingFile.model_fields)
_VIDEO_METADATA_FIELDS: Final = frozenset(_EmbeddingVideoMetadata.model_fields)
_FILE_SOURCE_FORMS: Final = "a data: URI, a gs:// URL, a files/ reference, or a Gemini Files API URI"
def _invalid_input(message: str) -> BadRequestError:
return BadRequestError(message=message, model=None, llm_provider="gemini")
def _validation_error_summary(error: ValidationError) -> str:
return "; ".join(f"{'.'.join(str(loc) for loc in detail['loc'])}: {detail['msg']}" for detail in error.errors())
def _as_mapping(value: object) -> dict[str, object] | None:
try:
return _mapping_adapter.validate_python(value)
except ValidationError:
return None
def _only_fields(mapping: Mapping[str, object], fields: frozenset[str]) -> dict[str, object]:
return {key: value for key, value in mapping.items() if key in fields}
def _dropping_unsupported_keys(block: Mapping[str, object]) -> dict[str, object]:
"""What `drop_params` keeps of a file content block: the keys this surface understands, at every level."""
kept_block: Final = _only_fields(block, _BLOCK_FIELDS)
file: Final = _as_mapping(block.get("file"))
if file is None:
return kept_block
kept_file: Final = _only_fields(file, _FILE_FIELDS)
video_metadata: Final = _as_mapping(file.get("video_metadata"))
if video_metadata is None:
return {**kept_block, "file": kept_file}
return {
**kept_block,
"file": {**kept_file, "video_metadata": _only_fields(video_metadata, _VIDEO_METADATA_FIELDS)},
}
def _parse_file_block(element: object, drop_params: bool) -> _EmbeddingFileBlock:
if not isinstance(element, Mapping):
raise _invalid_input(
f"Embedding input elements must be strings or file content blocks, got {type(element).__name__}"
)
block: Final = cast(Mapping[str, object], element) # cast-ok: isinstance leaves the key and value types unknown
try:
return _file_block_adapter.validate_python(_dropping_unsupported_keys(block) if drop_params else block)
except ValidationError as error:
raise _invalid_input(
f"Invalid file content block in embedding input: {_validation_error_summary(error)}"
) from error
def _file_block_source(block: _EmbeddingFileBlock) -> str:
match (block.file.file_id, block.file.file_data):
case (str() as file_id, None):
return file_id
case (None, str() as file_data):
return file_data
case (None, None):
raise _invalid_input("A file content block in embedding input needs file.file_id or file.file_data")
case _:
raise _invalid_input(
"A file content block in embedding input takes file.file_id or file.file_data, not both"
)
def _gemini_video_metadata(video_metadata: _EmbeddingVideoMetadata) -> VideoMetadataType:
return _video_metadata_adapter.validate_python(
gemini_video_metadata_from_openai(video_metadata.model_dump(exclude_none=True))
)
def _source_mime_type(
source: str,
mime_type_override: str | None,
resolved_files: Mapping[str, Mapping[str, str]],
) -> str | None:
if mime_type_override is not None:
return mime_type_override
if _is_data_url(source):
try:
return _parse_data_url(source)[0]
except ValueError:
return None
if _is_gcs_url(source):
try:
return _infer_mime_type_from_gcs_url(source)
except ValueError:
return None
if is_file_reference(source):
file_info: Final = resolved_files.get(source)
return None if file_info is None else file_info.get("mime_type")
return None
def _media_part(
source: str,
mime_type_override: str | None,
resolved_files: Mapping[str, Mapping[str, str]],
) -> PartType:
if _is_data_url(source):
mime_type, base64_data = _parse_data_url(source, mime_type_override)
return PartType(inline_data=BlobType(mime_type=mime_type, data=base64_data))
if _is_gcs_url(source):
gcs_mime_type: Final = mime_type_override or _infer_mime_type_from_gcs_url(source)
return PartType(file_data=FileDataType(mime_type=gcs_mime_type, file_uri=source))
if is_file_reference(source):
file_info: Final = resolved_files.get(source)
if file_info is None:
raise _invalid_input(
f"File reference {source!r} could not be resolved: "
"Gemini Files API references are only supported through the gemini/ provider"
)
return PartType(
file_data=FileDataType(mime_type=mime_type_override or file_info["mime_type"], file_uri=file_info["uri"])
)
raise _invalid_input(f"A file content block source must be {_FILE_SOURCE_FORMS}, got {source[:50]!r}")
def _top_level_elements(
input: GeminiEmbeddingInput,
) -> Sequence[GeminiEmbeddingElement | list[str] | list[GeminiEmbeddingElement]]:
try:
_input_shape_adapter.validate_python(input)
except ValidationError as error:
raise _invalid_input(
f"Embedding input must be a string or a list of strings and file content blocks, got {type(input).__name__}"
) from error
return [input] if isinstance(input, str) else input
def _elements(input: GeminiEmbeddingInput) -> Iterator[GeminiEmbeddingElement]:
for element in _top_level_elements(input):
if isinstance(element, list):
yield from element
else:
yield element
def _element_source(element: GeminiEmbeddingElement) -> str:
if isinstance(element, str):
return element
return _file_block_source(_parse_file_block(element, drop_params=True))
def flatten_media_sources(input: GeminiEmbeddingInput) -> tuple[str, ...]:
"""Every string element plus every file content block's source, in input order."""
return tuple(_element_source(element) for element in _elements(input))
def _is_multimodal_input(input: GeminiEmbeddingInput) -> bool:
"""
Check if the input contains multimodal data (data URIs, file references,
GCS URLs, or nested lists for combined embeddings).
GCS URLs, file content blocks, or nested lists for combined embeddings).
Args:
input: GeminiEmbeddingInput — str, List[str], or List[List[str]] for combined embeddings
input: GeminiEmbeddingInput — str, List[element], or List[List[element]] for combined embeddings
Returns:
bool: True if any element is multimodal or a nested list
bool: True if any element is multimodal
"""
if isinstance(input, str):
return _is_multimodal_element(input)
for element in input:
if isinstance(element, list):
if any(_is_multimodal_element(sub) for sub in element if isinstance(sub, str)):
return True
elif isinstance(element, str) and _is_multimodal_element(element):
return True
return False
return any(_is_multimodal_element(element) for element in _elements(input))
def _is_multimodal_element(element: str) -> bool:
"""Check if a single string element is multimodal."""
if element.startswith("data:") and ";base64," in element:
def _is_multimodal_element(element: GeminiEmbeddingElement) -> bool:
"""Check if a single element is multimodal."""
if not isinstance(element, str):
return True
if _is_data_url(element):
return True
if is_file_reference(element):
return True
@ -163,37 +351,29 @@ def _is_multimodal_element(element: str) -> bool:
def _build_part_for_input(
element: str,
resolved_files: dict[str, dict[str, str]] | None = None,
element: GeminiEmbeddingElement,
resolved_files: Mapping[str, Mapping[str, str]] | None = None,
drop_params: bool = False,
) -> PartType:
"""
Build a single PartType for an input element, handling text, data URIs,
file references, and GCS URLs.
file references, GCS URLs, and file content blocks carrying a mime type
and video_metadata.
"""
resolved_files = resolved_files or {}
files: Final = resolved_files or {}
if element.startswith("data:") and ";base64," in element:
mime_type, base64_data = _parse_data_url(element)
blob: Final[BlobType] = {"mime_type": mime_type, "data": base64_data}
return PartType(inline_data=blob)
elif _is_gcs_url(element):
mime_type = _infer_mime_type_from_gcs_url(element)
file_data: Final[FileDataType] = {
"mime_type": mime_type,
"file_uri": element,
}
return PartType(file_data=file_data)
elif is_file_reference(element):
if element not in resolved_files:
raise ValueError(f"File reference {element} not resolved")
file_info: Final = resolved_files[element]
file_data_ref: Final[FileDataType] = {
"mime_type": file_info["mime_type"],
"file_uri": file_info["uri"],
}
return PartType(file_data=file_data_ref)
else:
return PartType(text=element)
if isinstance(element, str):
return _media_part(element, None, files) if _is_multimodal_element(element) else PartType(text=element)
block: Final = _parse_file_block(element, drop_params)
part: Final = _media_part(_file_block_source(block), block.file.format, files)
if block.file.video_metadata is None:
return part
video_metadata: Final = _gemini_video_metadata(block.file.video_metadata)
if not video_metadata:
return part
part_with_metadata: Final[PartType] = {**part, "video_metadata": video_metadata}
return part_with_metadata
_SUPPORTED_EMBED_PARAMS: Final = {"outputDimensionality", "taskType", "title"}
@ -214,6 +394,7 @@ def transform_openai_input_gemini_content(
model: str,
optional_params: dict,
resolved_files: dict[str, dict[str, str]] | None = None,
drop_params: bool = False,
) -> VertexAIBatchEmbeddingsRequestBody:
"""
Transform OpenAI embedding input to Gemini batchEmbedContents format.
@ -234,19 +415,18 @@ def transform_openai_input_gemini_content(
gemini_params: Final = _filter_embed_params(optional_params)
input_list: Final = [input] if isinstance(input, str) else input
input_list: Final = _top_level_elements(input)
requests: Final[list[EmbedContentRequest]] = []
for element in input_list:
if isinstance(element, list):
if not element:
raise ValueError("Nested input list must not be empty")
for sub in element:
if not isinstance(sub, str):
raise ValueError(f"Elements inside a nested input list must be strings, got {type(sub)}")
parts = [_build_part_for_input(sub, resolved_files=resolved_files) for sub in element]
parts = [
_build_part_for_input(sub, resolved_files=resolved_files, drop_params=drop_params) for sub in element
]
else:
parts = [_build_part_for_input(element, resolved_files=resolved_files)]
parts = [_build_part_for_input(element, resolved_files=resolved_files, drop_params=drop_params)]
request = EmbedContentRequest(
model=gemini_model_name,
content=ContentType(parts=parts),
@ -262,6 +442,7 @@ def transform_openai_input_gemini_embed_content(
model: str,
optional_params: dict,
resolved_files: dict[str, dict[str, str]] | None = None,
drop_params: bool = False,
) -> dict:
"""
Transform OpenAI embedding input to Gemini embedContent format (multimodal).
@ -279,7 +460,7 @@ def transform_openai_input_gemini_embed_content(
gemini_params: Final = _filter_embed_params(optional_params)
input_list: Final = [input] if isinstance(input, str) else input
input_list: Final = _top_level_elements(input)
parts: Final[list[PartType]] = []
for element in input_list:
@ -288,9 +469,7 @@ def transform_openai_input_gemini_embed_content(
"Nested (combined) embeddings are not supported on the embedContent path. "
"Use the batchEmbedContents path or pass a flat list instead."
)
if not isinstance(element, str):
raise ValueError(f"Unsupported input type: {type(element)}")
parts.append(_build_part_for_input(element, resolved_files=resolved_files))
parts.append(_build_part_for_input(element, resolved_files=resolved_files, drop_params=drop_params))
request_body: Final[dict] = {
"content": ContentType(parts=parts),
@ -313,31 +492,18 @@ def _parse_usage_metadata(raw_usage_metadata: object) -> UsageMetadata | None:
return None
def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]:
if isinstance(input, str):
return (input,)
return tuple(sub for element in input for sub in (element if isinstance(element, list) else [element]))
def _flatten_input(input: GeminiEmbeddingInput) -> tuple[GeminiEmbeddingElement, ...]:
return tuple(_elements(input))
def _is_image_element(
element: str,
element: GeminiEmbeddingElement,
resolved_files: Mapping[str, Mapping[str, str]],
) -> bool:
if element.startswith("data:") and ";base64," in element:
try:
mime_type, _ = _parse_data_url(element)
except ValueError:
return False
return mime_type in _IMAGE_MIME_TYPES
if _is_gcs_url(element):
try:
return _infer_mime_type_from_gcs_url(element) in _IMAGE_MIME_TYPES
except ValueError:
return False
if is_file_reference(element):
file_info: Final = resolved_files.get(element)
return file_info is not None and file_info.get("mime_type") in _IMAGE_MIME_TYPES
return False
if isinstance(element, str):
return _source_mime_type(element, None, resolved_files) in _IMAGE_MIME_TYPES
block: Final = _parse_file_block(element, drop_params=True)
return _source_mime_type(_file_block_source(block), block.file.format, resolved_files) in _IMAGE_MIME_TYPES
def _is_image_only_input(

View file

@ -6379,6 +6379,10 @@ def embedding(
"""
azure: Final = kwargs.get("azure", None)
client: Final = kwargs.pop("client", None)
drop_params_kwarg: Final = (
cast(object, kwargs["drop_params"]) if "drop_params" in kwargs else None # cast-ok: untyped request kwargs
)
drop_unsupported_params: Final = litellm.drop_params is True or normalize_drop_params(drop_params_kwarg) is True
shared_session: Final = kwargs.get("shared_session", None)
max_retries: Final = kwargs.get("max_retries", None)
litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj")
@ -6859,6 +6863,7 @@ def embedding(
api_base=api_base,
client=client,
extra_headers=headers,
drop_params=drop_unsupported_params,
)
elif custom_llm_provider == "vertex_ai":
@ -6911,6 +6916,7 @@ def embedding(
api_base=api_base,
client=client,
extra_headers=headers,
drop_params=drop_unsupported_params,
)
elif (
"image" in optional_params

View file

@ -42320,6 +42320,7 @@
"/v1/chat/completions"
],
"supports_bedrock_runtime_chat_completions_tools_with_reasoning": true,
"supports_bedrock_runtime_chat_completions_inline_reasoning": true,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 128000,
@ -42338,6 +42339,7 @@
"/v1/chat/completions"
],
"supports_bedrock_runtime_chat_completions_tools_with_reasoning": true,
"supports_bedrock_runtime_chat_completions_inline_reasoning": true,
"input_cost_per_token": 7e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 128000,
@ -42363,6 +42365,7 @@
"output_cost_per_token_flex": 3e-07,
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_bedrock_runtime_chat_completions_inline_reasoning": true,
"supports_function_calling": true,
"supports_response_schema": true,
"supports_system_messages": true,
@ -42380,6 +42383,7 @@
"output_cost_per_token_flex": 1e-07,
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_bedrock_runtime_chat_completions_inline_reasoning": true,
"supports_function_calling": true,
"supports_response_schema": true,
"supports_system_messages": true,
@ -48105,6 +48109,7 @@
"/v1/chat/completions"
],
"supports_bedrock_runtime_chat_completions_tools_with_reasoning": true,
"supports_bedrock_runtime_chat_completions_inline_reasoning": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -48123,6 +48128,7 @@
"/v1/chat/completions"
],
"supports_bedrock_runtime_chat_completions_tools_with_reasoning": true,
"supports_bedrock_runtime_chat_completions_inline_reasoning": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -65951,6 +65957,7 @@
"/v1/chat/completions"
],
"supports_bedrock_runtime_chat_completions_tools_with_reasoning": true,
"supports_bedrock_runtime_chat_completions_inline_reasoning": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -65969,6 +65976,7 @@
"/v1/chat/completions"
],
"supports_bedrock_runtime_chat_completions_tools_with_reasoning": true,
"supports_bedrock_runtime_chat_completions_inline_reasoning": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -66219,6 +66227,7 @@
"/v1/chat/completions"
],
"supports_bedrock_runtime_chat_completions_tools_with_reasoning": true,
"supports_bedrock_runtime_chat_completions_inline_reasoning": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -66237,6 +66246,7 @@
"/v1/chat/completions"
],
"supports_bedrock_runtime_chat_completions_tools_with_reasoning": true,
"supports_bedrock_runtime_chat_completions_inline_reasoning": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -80458,29 +80468,6 @@
"supports_tool_choice": true,
"supports_vision": true
},
"xai.grok-4.7": {
"supports_regex_lookaround": false,
"cache_read_input_token_cost": 5e-07,
"cache_read_input_token_cost_flex": 2.5e-07,
"cache_read_input_token_cost_priority": 8.75e-07,
"input_cost_per_token": 2e-06,
"input_cost_per_token_flex": 1e-06,
"input_cost_per_token_priority": 3.5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 500000,
"max_output_tokens": 500000,
"max_tokens": 500000,
"mode": "chat",
"output_cost_per_token": 6e-06,
"output_cost_per_token_flex": 3e-06,
"output_cost_per_token_priority": 1.05e-05,
"source": "https://b0.p.awsstatic.com/pricing/2.0/meteredUnitMaps/bedrock/USD/current/bedrock.json",
"supports_function_calling": true,
"supports_prompt_caching": false,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"openrouter/anthropic/claude-sonnet-5.5:batch": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_1hr": 2e-06,

View file

@ -11906,6 +11906,34 @@
"description": "When True (default), after sensitive data is detected and routed, all subsequent requests in the same session will continue routing to the same model.",
"title": "Sticky Session Routing"
},
"stream_scope": {
"anyOf": [
{
"enum": [
"streaming",
"non_streaming",
"both"
],
"type": "string"
},
{
"additionalProperties": {
"enum": [
"streaming",
"non_streaming",
"both"
],
"type": "string"
},
"type": "object"
},
{
"type": "null"
}
],
"description": "Whether this guardrail runs on streaming requests, non-streaming requests, or both. A string applies to every configured mode. A map overrides named modes (pre_call, during_call, post_call, ...); omitted keys default to both. Unset means both, matching historical behavior.",
"title": "Stream Scope"
},
"template_id": {
"anyOf": [
{
@ -14831,6 +14859,34 @@
"description": "When True (default), after sensitive data is detected and routed, all subsequent requests in the same session will continue routing to the same model.",
"title": "Sticky Session Routing"
},
"stream_scope": {
"anyOf": [
{
"enum": [
"streaming",
"non_streaming",
"both"
],
"type": "string"
},
{
"additionalProperties": {
"enum": [
"streaming",
"non_streaming",
"both"
],
"type": "string"
},
"type": "object"
},
{
"type": "null"
}
],
"description": "Whether this guardrail runs on streaming requests, non-streaming requests, or both. A string applies to every configured mode. A map overrides named modes (pre_call, during_call, post_call, ...); omitted keys default to both. Unset means both, matching historical behavior.",
"title": "Stream Scope"
},
"template_id": {
"anyOf": [
{

View file

@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple, Ty
import httpx
from pydantic import (
AfterValidator,
BaseModel,
BeforeValidator,
ConfigDict,
@ -2743,6 +2744,41 @@ class ScheduledJobStaggerSettings(LiteLLMPydanticObjectBase):
)
SPEND_LOGS_METADATA_ALWAYS_KEPT_FIELDS: Final = frozenset({"status", "cold_storage_object_key"})
def _known_spend_logs_metadata_field(name: str) -> str:
if name not in SpendLogsMetadata.__annotations__:
raise ValueError(f"{name!r} is not a LiteLLM_SpendLogs.metadata field")
return name
SpendLogsMetadataFieldName: TypeAlias = Annotated[str, AfterValidator(_known_spend_logs_metadata_field)]
class SpendLogsMetadataFields(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
include: tuple[SpendLogsMetadataFieldName, ...] | None = None
exclude: tuple[SpendLogsMetadataFieldName, ...] | None = None
@model_validator(mode="after")
def _exactly_one_list(self) -> "SpendLogsMetadataFields":
if (self.include is None) == (self.exclude is None):
raise ValueError("set exactly one of 'include' or 'exclude'")
always_kept_excluded: Final = sorted(SPEND_LOGS_METADATA_ALWAYS_KEPT_FIELDS.intersection(self.exclude or ()))
if always_kept_excluded:
raise ValueError(f"{always_kept_excluded} are always kept and cannot be excluded")
return self
def keeps(self, name: str) -> bool:
if name in SPEND_LOGS_METADATA_ALWAYS_KEPT_FIELDS:
return True
if self.include is not None:
return name in self.include
return name not in (self.exclude or ())
DEFAULT_RESPONSES_WEBSOCKET_SESSION_LIMIT_SECONDS: Final[float] = 3600.0
@ -3095,6 +3131,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="If True, a completed batch logged via aretrieve_batch also emits one callback event per JSONL line item (request paired with its response or error). The aggregate batch callback is unchanged. Default is False.",
)
spend_logs_metadata_fields: SpendLogsMetadataFields | None = Field(
None,
description="Which keys of LiteLLM_SpendLogs.metadata are written to the database. Set exactly one of 'include' (write only these keys) or 'exclude' (drop these keys). 'status' and 'cold_storage_object_key' are always written. Daily spend tables, budgets and logging callbacks still see every key. Unset writes every key",
)
disable_auto_add_proxy_admin_to_teams: bool | None = Field(
None,
description="By default, the user calling /team/new is automatically added to the new team as a team admin. If True, proxy admins are no longer auto-added; members explicitly listed in members_with_roles are unaffected. Default is False.",

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