mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
Merge remote-tracking branch 'origin/main' into fix41691
# Conflicts: # litellm/proxy/_types.py
This commit is contained in:
commit
586b022744
303 changed files with 25563 additions and 9564 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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`
|
||||
|
|
|
|||
5
litellm-rust/Cargo.lock
generated
5
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
|
|
|
|||
1
litellm-rust/crates/cache/Cargo.toml
vendored
1
litellm-rust/crates/cache/Cargo.toml
vendored
|
|
@ -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]
|
||||
|
|
|
|||
57
litellm-rust/crates/cache/src/cache_type.rs
vendored
57
litellm-rust/crates/cache/src/cache_type.rs
vendored
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
16
litellm-rust/crates/cache/tests/cache_type.rs
vendored
16
litellm-rust/crates/cache/tests/cache_type.rs
vendored
|
|
@ -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());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
39
litellm-rust/crates/config/tests/mcp.rs
Normal file
39
litellm-rust/crates/config/tests/mcp.rs
Normal 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
|
||||
);
|
||||
}
|
||||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
36
litellm-rust/crates/llms/src/anthropic/beta_headers.rs
Normal file
36
litellm-rust/crates/llms/src/anthropic/beta_headers.rs
Normal 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())
|
||||
}
|
||||
}
|
||||
|
|
@ -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<_>>(),
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -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
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
pub mod beta_headers;
|
||||
pub mod common_utils;
|
||||
|
||||
pub mod batches;
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
]
|
||||
}],
|
||||
|
|
|
|||
|
|
@ -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<_>>(),
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
111
litellm-rust/crates/llms/tests/anthropic_beta_headers.rs
Normal file
111
litellm-rust/crates/llms/tests/anthropic_beta_headers.rs
Normal 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);
|
||||
}
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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", ¶ms).unwrap(),
|
||||
|
|
|
|||
|
|
@ -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 })
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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`
|
||||
|
|
|
|||
|
|
@ -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)?;
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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::{
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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!(
|
||||
|
|
|
|||
|
|
@ -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::{
|
||||
|
|
|
|||
102
litellm-rust/crates/python-bridge/src/routes/responses/mod.rs
Normal file
102
litellm-rust/crates/python-bridge/src/routes/responses/mod.rs
Normal 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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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,
|
||||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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() {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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 =
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 (),
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
99
litellm/llms/bedrock/count_tokens/mantle_handler.py
Normal file
99
litellm/llms/bedrock/count_tokens/mantle_handler.py
Normal 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}")
|
||||
|
|
@ -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"})
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Add table
Reference in a new issue