mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
Merge remote-tracking branch 'origin/main' into litellm_responses_tool_search_anthropic
# Conflicts: # litellm/llms/anthropic/chat/transformation.py
This commit is contained in:
commit
063a87296f
653 changed files with 75005 additions and 3279 deletions
6
.github/workflows/test-linting.yml
vendored
6
.github/workflows/test-linting.yml
vendored
|
|
@ -600,6 +600,9 @@ jobs:
|
|||
- name: check_unsafe_enterprise_import
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_unsafe_enterprise_import.py
|
||||
|
||||
- name: check_telemetry_imports
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_telemetry_imports.py
|
||||
|
||||
- name: ban_copy_deepcopy_kwargs
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/ban_copy_deepcopy_kwargs.py
|
||||
|
||||
|
|
@ -615,6 +618,9 @@ jobs:
|
|||
- name: check_migrations_no_data_rewrites
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py
|
||||
|
||||
- name: check_migrations_lock_timeout
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_lock_timeout.py
|
||||
|
||||
- name: check_no_publicly_known_master_key
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_no_publicly_known_master_key.py
|
||||
|
||||
|
|
|
|||
5
.github/workflows/test-unit.yml
vendored
5
.github/workflows/test-unit.yml
vendored
|
|
@ -233,6 +233,7 @@ jobs:
|
|||
tests/unit/llms/base_llm
|
||||
tests/unit/llms/baseten
|
||||
tests/unit/llms/black_forest_labs
|
||||
tests/unit/llms/brave
|
||||
tests/unit/llms/bytez
|
||||
tests/unit/llms/cerebras
|
||||
tests/unit/llms/chat
|
||||
|
|
@ -268,6 +269,7 @@ jobs:
|
|||
tests/unit/llms/gdc
|
||||
tests/unit/llms/gemini
|
||||
tests/unit/llms/gigachat
|
||||
tests/unit/llms/google_pse
|
||||
tests/unit/llms/github_copilot
|
||||
tests/unit/llms/gradient_ai
|
||||
tests/unit/llms/groq
|
||||
|
|
@ -300,6 +302,7 @@ jobs:
|
|||
tests/unit/llms/nebius
|
||||
tests/unit/llms/neosantara
|
||||
tests/unit/llms/nimble
|
||||
tests/unit/llms/nlp_cloud
|
||||
tests/unit/llms/novita
|
||||
tests/unit/llms/nscale
|
||||
tests/unit/llms/nvidia_nim
|
||||
|
|
@ -310,6 +313,7 @@ jobs:
|
|||
tests/unit/llms/oobabooga
|
||||
tests/unit/llms/openai_like
|
||||
tests/unit/llms/opencode
|
||||
tests/unit/llms/opensandbox
|
||||
tests/unit/llms/openrouter
|
||||
tests/unit/llms/ovhcloud
|
||||
tests/unit/llms/parallel_ai
|
||||
|
|
@ -446,6 +450,7 @@ jobs:
|
|||
tests/unit/repositories
|
||||
tests/unit/secret_managers
|
||||
tests/unit/skills/test_skills_main.py
|
||||
tests/unit/telemetry tests/unit/proxy/telemetry
|
||||
tests/unit/tracing
|
||||
python-version: "3.12"
|
||||
workers: 2
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ When adding new features, add meaningful tests. Don't add tests that don't check
|
|||
|
||||
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
|
||||
|
||||
Never test structure of code only function of it
|
||||
Never test structure or implementation details of code, only function of it. A test should catch a change that breaks behavior by accident. If the only way it can fail is someone deliberately undoing the change, don't write it
|
||||
|
||||
A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent". Assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken
|
||||
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/token",
|
||||
"/onboarding/",
|
||||
"/audit",
|
||||
"/telemetry/",
|
||||
"/oauth/",
|
||||
"/invitation/",
|
||||
"/jwt/",
|
||||
|
|
|
|||
|
|
@ -322,7 +322,12 @@ def create_migration(
|
|||
|
||||
# Write the SQL to migration.sql
|
||||
migration_file = migration_dir / "migration.sql"
|
||||
migration_file.write_text(diff_sql)
|
||||
migration_sql: Final = (
|
||||
diff_sql
|
||||
if "CONCURRENTLY" in diff_sql.upper()
|
||||
else "SET LOCAL lock_timeout = '10s';\n\n" + diff_sql
|
||||
)
|
||||
migration_file.write_text(migration_sql)
|
||||
|
||||
print(f"Created migration in {migration_dir}")
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -55,6 +55,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
|
|||
RENAME TO "LiteLLM_SpendLogs_legacy_session_id_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_api_key_startTime_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_litellm_call_id_idx";
|
||||
|
||||
CREATE TABLE "LiteLLM_SpendLogs" (
|
||||
LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.76"
|
||||
version = "0.1.77"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.76"
|
||||
version = "0.1.77"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -121,6 +121,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
# Realtime / streaming
|
||||
"/v1/realtime",
|
||||
"/realtime",
|
||||
"/v1/live/sessions",
|
||||
"/live/sessions",
|
||||
# Health & ops
|
||||
"/health",
|
||||
"/metrics",
|
||||
|
|
|
|||
|
|
@ -440,10 +440,10 @@ db:
|
|||
# Prefer the native /health/drain preStop hook over a fixed `sleep`: it marks
|
||||
# the pod NotReady and blocks only until in-flight requests actually finish
|
||||
# (bounded by GRACEFUL_SHUTDOWN_TIMEOUT, default 30s), instead of always
|
||||
# waiting the worst-case duration. The drain runs once (the preStop hook and
|
||||
# the SIGTERM handler share it), so set terminationGracePeriodSeconds a few
|
||||
# seconds above GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before
|
||||
# SIGKILL.
|
||||
# waiting the worst-case duration. After SIGTERM the server waits up to
|
||||
# GRACEFUL_SHUTDOWN_TIMEOUT again for connections still open and then cancels
|
||||
# them, so set terminationGracePeriodSeconds a few seconds above twice
|
||||
# GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before SIGKILL.
|
||||
#
|
||||
# /health/drain is off by default; enable it with
|
||||
# general_settings.enable_drain_endpoint: true. The kubelet calls preStop
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import random
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
|
|
@ -10,12 +9,10 @@ from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
|||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
MIGRATION_LOCK_TIMEOUT_ENV_VAR,
|
||||
migration_ddl_lock_timeout,
|
||||
migration_lock_timeout,
|
||||
)
|
||||
|
||||
MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
|
||||
_LOCK_TIMEOUT_PINNED_RE: Final = re.compile(r"(?:-c\s*|--)lock_timeout=")
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import psycopg
|
||||
|
|
@ -30,29 +27,12 @@ def _migration_url(database_url: str, direct_url: "str | None") -> str:
|
|||
return urlunsplit(direct._replace(query=urlencode((*parameters, ("schema", schema)))))
|
||||
|
||||
|
||||
def _with_ddl_lock_timeout(url: str) -> str:
|
||||
split: Final = urlsplit(url)
|
||||
pairs: Final = tuple(parse_qsl(split.query))
|
||||
if any(key == "pgbouncer" and value == "true" for key, value in pairs):
|
||||
return url
|
||||
existing: Final = next((value for key, value in pairs if key == "options"), "")
|
||||
if _LOCK_TIMEOUT_PINNED_RE.search(existing):
|
||||
return url
|
||||
timeout_ms: Final = max(1, int(migration_ddl_lock_timeout() * 1000))
|
||||
options: Final = f"{existing} -c lock_timeout={timeout_ms}".strip()
|
||||
rewritten: Final = tuple((key, value) for key, value in pairs if key != "options") + (("options", options),)
|
||||
return urlunsplit(split._replace(query=urlencode(rewritten)))
|
||||
|
||||
|
||||
def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
|
||||
database_url: Final = environment.get("DATABASE_URL")
|
||||
if not database_url:
|
||||
return environment
|
||||
migration_url: Final = _migration_url(database_url, environment.get("DIRECT_URL"))
|
||||
return {
|
||||
**environment,
|
||||
"DATABASE_URL": _with_ddl_lock_timeout(migration_url),
|
||||
}
|
||||
return {**environment, "DATABASE_URL": migration_url}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,12 @@
|
|||
CREATE TABLE IF NOT EXISTS "LiteLLM_TelemetryReport" (
|
||||
"id" TEXT NOT NULL,
|
||||
"window_start" DOUBLE PRECISION NOT NULL,
|
||||
"window_end" DOUBLE PRECISION NOT NULL,
|
||||
"report" JSONB NOT NULL,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
CONSTRAINT "LiteLLM_TelemetryReport_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_TelemetryReport_window_end_idx" ON "LiteLLM_TelemetryReport"("window_end");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_TelemetryReport_created_at_idx" ON "LiteLLM_TelemetryReport"("created_at");
|
||||
|
|
@ -60,13 +60,11 @@ PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
|
|||
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT"
|
||||
MIGRATION_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_LOCK_TIMEOUT"
|
||||
MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_DDL_LOCK_TIMEOUT"
|
||||
NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"
|
||||
|
||||
DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
|
||||
DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
|
||||
DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0
|
||||
DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT = 10.0
|
||||
|
||||
BOOTSTRAP_ARG = "--version"
|
||||
PRISMA_CONSOLE_SCRIPT = "prisma"
|
||||
|
|
@ -113,11 +111,6 @@ def migration_lock_timeout() -> float:
|
|||
return _timeout_from_env(MIGRATION_LOCK_TIMEOUT_ENV_VAR, 600.0)
|
||||
|
||||
|
||||
def migration_ddl_lock_timeout() -> float:
|
||||
"""Seconds one migration statement may wait on a table lock before Postgres cancels it."""
|
||||
return _timeout_from_env(MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR, DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT)
|
||||
|
||||
|
||||
def prisma_bootstrap_timeout() -> float:
|
||||
"""Seconds the one-time Node toolchain install may run for."""
|
||||
return _timeout_from_env(
|
||||
|
|
|
|||
|
|
@ -2045,3 +2045,14 @@ model LiteLLM_LensTraceSignal {
|
|||
|
||||
@@id([trace_id, trace_ref])
|
||||
}
|
||||
|
||||
model LiteLLM_TelemetryReport {
|
||||
id String @id
|
||||
window_start Float
|
||||
window_end Float
|
||||
report Json
|
||||
created_at DateTime @default(now())
|
||||
|
||||
@@index([window_end])
|
||||
@@index([created_at])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -18,7 +18,6 @@ from litellm_proxy_extras import prisma_toolchain
|
|||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.migration_lock import held_migration_lock
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR,
|
||||
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
|
||||
ensure_prisma_toolchain,
|
||||
|
|
@ -1165,10 +1164,9 @@ class ProxyExtrasDBManager:
|
|||
raise RuntimeError(
|
||||
f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} "
|
||||
"attempts that made no progress (timeouts, lock timeouts or deadlock retries). "
|
||||
"Check database connectivity, load, long-running transactions on the migrated tables, "
|
||||
"and _prisma_migrations ledger state. Raise "
|
||||
f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out, or "
|
||||
f"{MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR} if they timed out waiting for a table lock."
|
||||
"Check database connectivity, load and _prisma_migrations ledger state. "
|
||||
"If a migration hit SET LOCAL lock_timeout, check for long-running transactions on the migrated tables. "
|
||||
f"Raise {PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
|
||||
)
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
|
@ -1284,8 +1282,8 @@ class ProxyExtrasDBManager:
|
|||
if migration_name and _MIGRATION_LOCK_TIMEOUT_MARKER in stderr:
|
||||
ProxyExtrasDBManager._log_migration_lock_holders(migration_name)
|
||||
logger.warning(
|
||||
"Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying. "
|
||||
f"Raise {MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR} if the database needs longer.",
|
||||
"Migration %s timed out waiting for a table lock after hitting SET LOCAL lock_timeout, rolling "
|
||||
"its ledger row back and retrying. Check for long-running transactions on the migrated tables.",
|
||||
migration_name,
|
||||
)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.108"
|
||||
version = "0.4.109"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.108"
|
||||
version = "0.4.109"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
use std::str;
|
||||
|
||||
use bytes::{Buf, BufMut, BytesMut};
|
||||
use bytes::{Buf, BufMut, Bytes, BytesMut};
|
||||
use tokio_util::codec::{Decoder, Encoder};
|
||||
|
||||
use crate::SseError;
|
||||
|
|
@ -29,38 +29,107 @@ impl Decoder for SseCodec {
|
|||
return Ok(None);
|
||||
}
|
||||
while let Some(end) = block_end(src) {
|
||||
let block = src.split_to(end);
|
||||
let pending = lines(&block)
|
||||
.map(|(line, _)| line)
|
||||
.take_while(|line| !line.is_empty())
|
||||
.try_fold(Pending::default(), Pending::apply)?;
|
||||
if let Some(event) = pending.dispatch() {
|
||||
if let Some(event) = event_in(&src.split_to(end))? {
|
||||
return Ok(Some(event));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn decode_eof(&mut self, _pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
|
||||
Ok(None)
|
||||
/// An event the stream ended inside of is dropped, as the SSE spec requires.
|
||||
fn decode_eof(&mut self, pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
|
||||
self.decode(pending)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn decode_all(wire: &[u8]) -> Result<Vec<SseEvent>, SseError> {
|
||||
let mut codec = SseCodec::default();
|
||||
let mut buffer = BytesMut::from(wire);
|
||||
let mut events = Vec::new();
|
||||
while let Some(event) = codec.decode_eof(&mut buffer)? {
|
||||
events.push(event);
|
||||
}
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
impl SseCodec {
|
||||
fn skip_bom(&mut self, src: &mut BytesMut) -> bool {
|
||||
if self.past_bom {
|
||||
return true;
|
||||
}
|
||||
if src.starts_with(BOM) {
|
||||
src.advance(BOM.len());
|
||||
} else if BOM.starts_with(src) {
|
||||
return false;
|
||||
}
|
||||
self.past_bom = true;
|
||||
true
|
||||
skip_bom(&mut self.past_bom, src)
|
||||
}
|
||||
}
|
||||
|
||||
fn event_in(block: &[u8]) -> Result<Option<SseEvent>, SseError> {
|
||||
Ok(lines(block)
|
||||
.map(|(line, _)| line)
|
||||
.take_while(|line| !line.is_empty())
|
||||
.try_fold(Pending::default(), Pending::apply)?
|
||||
.dispatch())
|
||||
}
|
||||
|
||||
/// One event together with the exact wire bytes that carried it, for relays that forward a
|
||||
/// provider's bytes untouched but still need to see each event go by.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct SseBlock {
|
||||
pub bytes: Bytes,
|
||||
pub event: SseEvent,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct SseBlockCodec {
|
||||
past_bom: bool,
|
||||
}
|
||||
|
||||
impl Decoder for SseBlockCodec {
|
||||
type Item = SseBlock;
|
||||
type Error = SseError;
|
||||
|
||||
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<SseBlock>, SseError> {
|
||||
self.decode_block(src, false)
|
||||
}
|
||||
|
||||
fn decode_eof(&mut self, pending: &mut BytesMut) -> Result<Option<SseBlock>, SseError> {
|
||||
self.decode_block(pending, true)
|
||||
}
|
||||
}
|
||||
|
||||
impl SseBlockCodec {
|
||||
fn decode_block(
|
||||
&mut self,
|
||||
src: &mut BytesMut,
|
||||
eof: bool,
|
||||
) -> Result<Option<SseBlock>, SseError> {
|
||||
if !skip_bom(&mut self.past_bom, src) {
|
||||
return Ok(None);
|
||||
}
|
||||
while let Some(end) = block_end(src) {
|
||||
if !eof && end == src.len() && src.ends_with(b"\r") {
|
||||
return Ok(None);
|
||||
}
|
||||
let block = src.split_to(end);
|
||||
if let Some(event) = event_in(&block)? {
|
||||
return Ok(Some(SseBlock {
|
||||
bytes: block.freeze(),
|
||||
event,
|
||||
}));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
fn skip_bom(past_bom: &mut bool, src: &mut BytesMut) -> bool {
|
||||
if *past_bom {
|
||||
return true;
|
||||
}
|
||||
if src.starts_with(BOM) {
|
||||
src.advance(BOM.len());
|
||||
} else if BOM.starts_with(src) {
|
||||
return false;
|
||||
}
|
||||
*past_bom = true;
|
||||
true
|
||||
}
|
||||
|
||||
fn block_end(bytes: &[u8]) -> Option<usize> {
|
||||
lines(bytes)
|
||||
.find(|(line, _)| line.is_empty())
|
||||
|
|
@ -168,3 +237,179 @@ impl Encoder<SseEvent> for SseCodec {
|
|||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::frames;
|
||||
|
||||
fn pieces(wire: &[u8], size: usize) -> Vec<Bytes> {
|
||||
wire.chunks(size).map(Bytes::copy_from_slice).collect()
|
||||
}
|
||||
|
||||
async fn blocks(wire: &[u8], size: usize) -> Result<Vec<SseBlock>, SseError> {
|
||||
frames(
|
||||
stream::iter(pieces(wire, size).into_iter().map(Ok::<_, std::io::Error>)),
|
||||
SseBlockCodec::default(),
|
||||
)
|
||||
.try_collect()
|
||||
.await
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::lf(
|
||||
b": keepalive\n\nevent: delta\ndata: a\ndata: b\n\n",
|
||||
b"event: delta\ndata: a\ndata: b\n\n"
|
||||
)]
|
||||
#[case::crlf(
|
||||
b": keepalive\r\n\r\nevent: delta\r\ndata: a\r\ndata: b\r\n\r\n",
|
||||
b"event: delta\r\ndata: a\r\ndata: b\r\n\r\n"
|
||||
)]
|
||||
#[case::cr(
|
||||
b": keepalive\r\revent: delta\rdata: a\rdata: b\r\r",
|
||||
b"event: delta\rdata: a\rdata: b\r\r"
|
||||
)]
|
||||
#[case::bom(b"\xEF\xBB\xBFdata: \xe2\x82\xac\n\n", b"data: \xe2\x82\xac\n\n")]
|
||||
#[case::empty_data(b"event: ignored\n\ndata:\n\n", b"data:\n\n")]
|
||||
#[case::extensions(
|
||||
b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n",
|
||||
b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn blocks_carry_the_wire_bytes_and_the_event_they_dispatch(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: &[u8],
|
||||
#[values(1, 3, 4096)] size: usize,
|
||||
) {
|
||||
let blocks = blocks(wire, size).await.unwrap();
|
||||
let expected_events = decode_all(expected).unwrap();
|
||||
assert_eq!(expected_events.len(), 1);
|
||||
assert_eq!(
|
||||
blocks,
|
||||
vec![SseBlock {
|
||||
bytes: Bytes::copy_from_slice(expected),
|
||||
event: expected_events[0].clone(),
|
||||
}]
|
||||
);
|
||||
assert_eq!(expected_events, decode_all(wire).unwrap());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty(b"", vec![])]
|
||||
#[case::unterminated(b"data: partial\n", vec![])]
|
||||
#[case::partial_tail(b"data: complete\n\ndata: partial\n", vec![&b"data: complete\n\n"[..]])]
|
||||
#[case::multiple(b"data: first\n\ndata: second\n\n", vec![&b"data: first\n\n"[..], &b"data: second\n\n"[..]])]
|
||||
#[case::trailing_comment(b"data: first\n\n: bye\n\n", vec![&b"data: first\n\n"[..]])]
|
||||
#[tokio::test]
|
||||
async fn blocks_dispatch_complete_events_and_drop_a_cut_tail(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: Vec<&[u8]>,
|
||||
) {
|
||||
let mut blocks = Box::pin(frames(
|
||||
stream::iter(pieces(wire, 1).into_iter().map(Ok::<_, std::io::Error>)),
|
||||
SseBlockCodec::default(),
|
||||
));
|
||||
for expected in expected {
|
||||
assert_eq!(blocks.next().await.unwrap().unwrap().bytes, expected);
|
||||
}
|
||||
assert!(blocks.next().await.is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn bulk_decode_preserves_every_field() {
|
||||
let events = decode_all(b"id: 1\nretry: 20\nevent: delta\ndata: a\ndata: b\n\n").unwrap();
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![SseEvent {
|
||||
event: Some("delta".into()),
|
||||
data: "a\nb".into(),
|
||||
id: Some("1".into()),
|
||||
retry: Some(20),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::only_a_tail(b"data: partial\n", vec![])]
|
||||
#[case::tail_after_an_event(b"data: complete\n\ndata: partial\n", vec!["complete"])]
|
||||
#[case::no_terminator_at_all(b"data: partial", vec![])]
|
||||
fn bulk_decode_drops_the_event_a_stream_ends_inside(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: Vec<&str>,
|
||||
) {
|
||||
let data: Vec<String> = decode_all(wire)
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.map(|e| e.data)
|
||||
.collect();
|
||||
assert_eq!(data, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::trailing_blank_lines(b"data: x\n\n\n\n")]
|
||||
#[case::trailing_comment(b"data: x\n\n: bye\n\n")]
|
||||
#[case::lone_cr_at_eof(b"data: x\r\r")]
|
||||
fn bulk_decode_accepts_a_stream_that_ends_between_events(#[case] wire: &[u8]) {
|
||||
assert_eq!(decode_all(wire).unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn bulk_decode_rejects_invalid_utf8() {
|
||||
assert!(matches!(
|
||||
decode_all(b"data: \xff\n\n"),
|
||||
Err(SseError::InvalidUtf8(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn block_framing_preserves_transport_errors() {
|
||||
let input = stream::iter([
|
||||
Ok(Bytes::from_static(b"data: first\n\n")),
|
||||
Err(std::io::Error::new(
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"reset",
|
||||
)),
|
||||
]);
|
||||
let mut blocks = Box::pin(frames(input, SseBlockCodec::default()));
|
||||
assert_eq!(
|
||||
blocks.next().await.unwrap().unwrap().bytes,
|
||||
Bytes::from_static(b"data: first\n\n")
|
||||
);
|
||||
let Some(Err(SseError::Body(error))) = blocks.next().await else {
|
||||
panic!("expected body error");
|
||||
};
|
||||
let cause = error
|
||||
.get_ref()
|
||||
.unwrap()
|
||||
.downcast_ref::<std::io::Error>()
|
||||
.unwrap();
|
||||
assert_eq!(cause.kind(), std::io::ErrorKind::ConnectionReset);
|
||||
assert!(blocks.next().await.is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn invalid_utf8_stops_block_framing_after_prior_events() {
|
||||
let mut blocks = Box::pin(frames(
|
||||
stream::iter(
|
||||
pieces(b"data: first\n\ndata: \xff\n\n", 1)
|
||||
.into_iter()
|
||||
.map(Ok::<_, std::io::Error>),
|
||||
),
|
||||
SseBlockCodec::default(),
|
||||
));
|
||||
assert_eq!(
|
||||
blocks.next().await.unwrap().unwrap().bytes,
|
||||
Bytes::from_static(b"data: first\n\n")
|
||||
);
|
||||
assert!(matches!(
|
||||
blocks.next().await,
|
||||
Some(Err(SseError::InvalidUtf8(_)))
|
||||
));
|
||||
assert!(blocks.next().await.is_none());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -97,12 +97,11 @@ async fn dispatch_follows_the_data_buffer(#[case] wire: &[u8], #[case] expected:
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unterminated_single(b"data: partial\n", vec![])]
|
||||
#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
|
||||
#[case::lone_cr_terminates_at_eof(b"data: x\r\r", vec![event(None, "x")])]
|
||||
#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
|
||||
#[case::trailing_blank_lines(b"data: x\n\n\n", vec![event(None, "x")])]
|
||||
#[case::trailing_comment(b"data: x\n\n: bye\n\n", vec![event(None, "x")])]
|
||||
#[tokio::test]
|
||||
async fn eof_dispatches_only_terminated_events(
|
||||
async fn eof_between_events_ends_the_stream_cleanly(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: Vec<SseEvent>,
|
||||
) {
|
||||
|
|
@ -112,6 +111,19 @@ async fn eof_dispatches_only_terminated_events(
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unterminated_single(b"data: partial\n", vec![])]
|
||||
#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
|
||||
#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
|
||||
#[case::partial_bom(b"\xEF\xBB", vec![])]
|
||||
#[tokio::test]
|
||||
async fn eof_inside_an_event_drops_it_after_the_terminated_events(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: Vec<SseEvent>,
|
||||
) {
|
||||
assert_eq!(collect(every(wire, 1)).await.unwrap(), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::inside_the_first_line(vec![&b"data: a\r"[..], &b"\ndata: b\r\n\r\n"[..]])]
|
||||
#[case::inside_the_blank_line(vec![&b"data: a\r\ndata: b\r\n\r"[..], &b"\n"[..]])]
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ use litellm_http::transport::Error as TransportError;
|
|||
use litellm_llms::base_llm::{
|
||||
auth::{Authenticated, resolve_auth},
|
||||
messages::{
|
||||
streaming::{ByteStream, StreamDecoder, encode_anthropic_sse},
|
||||
streaming::{self, ByteStream, StreamDecoder},
|
||||
transformation::BaseMessagesConfig,
|
||||
},
|
||||
};
|
||||
|
|
@ -245,7 +245,7 @@ fn decoded_chunks(
|
|||
let Some(event) = events.try_next().await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let chunk = encode_anthropic_sse(&event)?;
|
||||
let chunk = streaming::encode(&event)?;
|
||||
log_chunk(provider, "client_response", &chunk);
|
||||
Ok(Some((chunk, events)))
|
||||
})
|
||||
|
|
@ -271,7 +271,6 @@ fn log_chunk(provider: &str, stage: &str, data: &bytes::Bytes) {
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_llms::base_llm::messages::streaming::anthropic_sse_event_stream;
|
||||
use rstest::rstest;
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
|
||||
|
||||
|
|
@ -300,7 +299,7 @@ mod tests {
|
|||
.await
|
||||
.unwrap();
|
||||
let MessagesCallResponse::Stream { mut chunks, .. } =
|
||||
streaming_response(response, Some(anthropic_sse_event_stream), "test")
|
||||
streaming_response(response, Some(streaming::event_stream), "test")
|
||||
else {
|
||||
panic!("a streaming response returns chunks");
|
||||
};
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
|
|||
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use litellm_http::websocket::{UpstreamWebSocket, connect_upstream};
|
||||
use litellm_llms_types::formats::responses::streaming_websocket::ResponsesWsEventType;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio_tungstenite::tungstenite::{
|
||||
Message,
|
||||
|
|
@ -12,17 +11,6 @@ use tokio_tungstenite::tungstenite::{
|
|||
|
||||
use super::Error;
|
||||
|
||||
pub fn is_terminal_event(event_type: &ResponsesWsEventType) -> bool {
|
||||
matches!(
|
||||
event_type,
|
||||
ResponsesWsEventType::ResponseCreated
|
||||
| ResponsesWsEventType::ResponseCompleted
|
||||
| ResponsesWsEventType::ResponseFailed
|
||||
| ResponsesWsEventType::ResponseIncomplete
|
||||
| ResponsesWsEventType::Error
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ResponsesWebSocketConnection {
|
||||
socket: Arc<Mutex<Option<UpstreamWebSocket>>>,
|
||||
|
|
|
|||
|
|
@ -33,16 +33,82 @@ pub struct ChatCompletionStreamingChoice {
|
|||
pub logprobs: Option<Value>,
|
||||
}
|
||||
|
||||
/// Compatible hosts omit `id`, `created` and `object` on chunks, and Python's
|
||||
/// `ModelResponseStream` fills them in, so none of the envelope fields is required here.
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
#[derive(Default)]
|
||||
pub struct ChatCompletionChunk {
|
||||
pub id: String,
|
||||
pub created: u64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub created: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub model: Option<String>,
|
||||
pub object: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub object: Option<String>,
|
||||
#[serde(default)]
|
||||
pub choices: Vec<ChatCompletionStreamingChoice>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub usage: Option<ChatCompletionsUsage>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_specific_fields: Option<Map<String, Value>>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::openai(json!({
|
||||
"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m",
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi", "future": 1}, "finish_reason": null, "logprobs": null}],
|
||||
"usage": null
|
||||
}))]
|
||||
#[case::usage_only(json!({
|
||||
"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3,
|
||||
"prompt_tokens_details": {"cached_tokens": 0, "cache_creation_tokens": 0, "text_tokens": 1}
|
||||
}
|
||||
}))]
|
||||
#[case::bare_host(json!({"choices": [{"index": 0, "delta": {"content": "hi"}}]}))]
|
||||
#[case::host_extensions(json!({"id": "x", "choices": [], "citations": ["a"], "future": {"k": null}}))]
|
||||
fn chunks_round_trip_without_losing_fields(#[case] wire: Value) {
|
||||
let chunk: ChatCompletionChunk = serde_json::from_value(wire.clone()).unwrap();
|
||||
assert_eq!(serde_json::to_value(chunk).unwrap(), without_nulls(wire));
|
||||
}
|
||||
|
||||
/// Typed optional fields serialize by omission; extension fields keep their nulls.
|
||||
fn without_nulls(value: Value) -> Value {
|
||||
let Value::Object(fields) = value else {
|
||||
return value;
|
||||
};
|
||||
Value::Object(
|
||||
fields
|
||||
.into_iter()
|
||||
.filter(|(_, value)| !value.is_null())
|
||||
.map(|(key, value)| match (key.as_str(), value) {
|
||||
("choices", Value::Array(choices)) => (
|
||||
key,
|
||||
Value::Array(choices.into_iter().map(without_nulls).collect()),
|
||||
),
|
||||
(_, value) => (key, value),
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::choices_not_a_list(json!({"choices": {}}))]
|
||||
#[case::choice_without_delta(json!({"choices": [{"index": 0}]}))]
|
||||
#[case::created_not_a_number(json!({"created": "now", "choices": []}))]
|
||||
#[case::non_object(json!([]))]
|
||||
fn malformed_chunks_are_rejected(#[case] wire: Value) {
|
||||
assert!(serde_json::from_value::<ChatCompletionChunk>(wire).is_err());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -109,7 +109,9 @@ pub struct MessagesStreamError {
|
|||
}
|
||||
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
#[derive(strum::AsRefStr)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum MessagesStreamEvent {
|
||||
MessageStart {
|
||||
message: MessagesStreamMessage,
|
||||
|
|
@ -175,6 +177,23 @@ mod tests {
|
|||
assert_eq!(serde_json::to_value(event).unwrap(), wire);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::start(json!({"type": "message_start", "message": {"id": "m", "type": "message", "role": "assistant", "model": "x", "content": [], "stop_reason": null, "stop_sequence": null, "usage": {}}}))]
|
||||
#[case::block_start(json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}))]
|
||||
#[case::block_delta(json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "a"}}))]
|
||||
#[case::block_stop(json!({"type": "content_block_stop", "index": 0}))]
|
||||
#[case::delta(json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}))]
|
||||
#[case::stop(json!({"type": "message_stop"}))]
|
||||
#[case::ping(json!({"type": "ping"}))]
|
||||
#[case::error(json!({"type": "error", "error": {"type": "overloaded_error", "message": "busy"}}))]
|
||||
fn the_event_name_is_the_serialized_type_tag(#[case] wire: Value) {
|
||||
let event: MessagesStreamEvent = serde_json::from_value(wire.clone()).unwrap();
|
||||
assert_eq!(
|
||||
Some(event.as_ref()),
|
||||
wire.get("type").and_then(Value::as_str)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::invalid_index(json!({"type": "content_block_stop", "index": "zero"}))]
|
||||
#[case::missing_delta(json!({"type": "content_block_delta", "index": 0}))]
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# rules
|
||||
|
||||
- Responses originated at OpenAI and is now offered by other hosts and gateways. Types here describe the format, not OpenAI's deployment of it
|
||||
- Served over HTTP and over the WebSocket mode. `ResponsesWsEvent` is a wire event and lives here
|
||||
- Served over HTTP and over the WebSocket mode. `ResponsesStreamEvent` and `ResponsesEventType` in `streaming.rs` are the wire event for both
|
||||
- `ResponsesWsTransformResult` wraps a provider transformation's output, so it lives in `llms/src/base_llm/responses/`
|
||||
- Hosts that accept only part of the format, or add output item types, are handled in `llms/src/<provider>/responses/`. Keep unknown output items passing through instead of rejecting them
|
||||
- Responses-to-Chat emulation and model-specific parameter rewriting are transformations, not format rules
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
mod output;
|
||||
mod response;
|
||||
pub mod streaming;
|
||||
pub mod streaming_websocket;
|
||||
|
||||
pub use output::*;
|
||||
|
|
|
|||
|
|
@ -0,0 +1,305 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
/// Every event name on https://developers.openai.com/api/reference/resources/responses/streaming-events
|
||||
/// as of 2026-10-11, plus `response.create` from the WebSocket mode. Names OpenAI adds later ride in
|
||||
/// `Other` untouched.
|
||||
#[derive(
|
||||
Clone,
|
||||
Debug,
|
||||
PartialEq,
|
||||
Eq,
|
||||
strum::AsRefStr,
|
||||
strum::EnumIter,
|
||||
strum::EnumString,
|
||||
strum::Display,
|
||||
serde_with::DeserializeFromStr,
|
||||
serde_with::SerializeDisplay,
|
||||
)]
|
||||
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
|
||||
#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))]
|
||||
pub enum ResponsesEventType {
|
||||
#[strum(serialize = "response.create")]
|
||||
ResponseCreate,
|
||||
#[strum(serialize = "response.created")]
|
||||
ResponseCreated,
|
||||
#[strum(serialize = "response.queued")]
|
||||
ResponseQueued,
|
||||
#[strum(serialize = "response.in_progress")]
|
||||
ResponseInProgress,
|
||||
#[strum(serialize = "response.completed")]
|
||||
ResponseCompleted,
|
||||
#[strum(serialize = "response.failed")]
|
||||
ResponseFailed,
|
||||
#[strum(serialize = "response.incomplete")]
|
||||
ResponseIncomplete,
|
||||
#[strum(serialize = "error")]
|
||||
Error,
|
||||
#[strum(serialize = "response.audio.delta")]
|
||||
AudioDelta,
|
||||
#[strum(serialize = "response.audio.done")]
|
||||
AudioDone,
|
||||
#[strum(serialize = "response.audio.transcript.delta")]
|
||||
AudioTranscriptDelta,
|
||||
#[strum(serialize = "response.audio.transcript.done")]
|
||||
AudioTranscriptDone,
|
||||
#[strum(serialize = "response.code_interpreter_call.completed")]
|
||||
CodeInterpreterCallCompleted,
|
||||
#[strum(serialize = "response.code_interpreter_call.in_progress")]
|
||||
CodeInterpreterCallInProgress,
|
||||
#[strum(serialize = "response.code_interpreter_call.interpreting")]
|
||||
CodeInterpreterCallInterpreting,
|
||||
#[strum(serialize = "response.code_interpreter_call_code.delta")]
|
||||
CodeInterpreterCallCodeDelta,
|
||||
#[strum(serialize = "response.code_interpreter_call_code.done")]
|
||||
CodeInterpreterCallCodeDone,
|
||||
#[strum(serialize = "response.compaction.compacting")]
|
||||
CompactionCompacting,
|
||||
#[strum(serialize = "response.content_part.added")]
|
||||
ContentPartAdded,
|
||||
#[strum(serialize = "response.content_part.done")]
|
||||
ContentPartDone,
|
||||
#[strum(serialize = "response.custom_tool_call_input.delta")]
|
||||
CustomToolCallInputDelta,
|
||||
#[strum(serialize = "response.custom_tool_call_input.done")]
|
||||
CustomToolCallInputDone,
|
||||
#[strum(serialize = "response.file_search_call.completed")]
|
||||
FileSearchCallCompleted,
|
||||
#[strum(serialize = "response.file_search_call.in_progress")]
|
||||
FileSearchCallInProgress,
|
||||
#[strum(serialize = "response.file_search_call.searching")]
|
||||
FileSearchCallSearching,
|
||||
#[strum(serialize = "response.function_call_arguments.delta")]
|
||||
FunctionCallArgumentsDelta,
|
||||
#[strum(serialize = "response.function_call_arguments.done")]
|
||||
FunctionCallArgumentsDone,
|
||||
#[strum(serialize = "response.image_generation_call.completed")]
|
||||
ImageGenerationCallCompleted,
|
||||
#[strum(serialize = "response.image_generation_call.generating")]
|
||||
ImageGenerationCallGenerating,
|
||||
#[strum(serialize = "response.image_generation_call.in_progress")]
|
||||
ImageGenerationCallInProgress,
|
||||
#[strum(serialize = "response.image_generation_call.partial_image")]
|
||||
ImageGenerationCallPartialImage,
|
||||
#[strum(serialize = "response.mcp_call.completed")]
|
||||
McpCallCompleted,
|
||||
#[strum(serialize = "response.mcp_call.failed")]
|
||||
McpCallFailed,
|
||||
#[strum(serialize = "response.mcp_call.in_progress")]
|
||||
McpCallInProgress,
|
||||
#[strum(serialize = "response.mcp_call_arguments.delta")]
|
||||
McpCallArgumentsDelta,
|
||||
#[strum(serialize = "response.mcp_call_arguments.done")]
|
||||
McpCallArgumentsDone,
|
||||
#[strum(serialize = "response.mcp_list_tools.completed")]
|
||||
McpListToolsCompleted,
|
||||
#[strum(serialize = "response.mcp_list_tools.failed")]
|
||||
McpListToolsFailed,
|
||||
#[strum(serialize = "response.mcp_list_tools.in_progress")]
|
||||
McpListToolsInProgress,
|
||||
#[strum(serialize = "response.output_item.added")]
|
||||
OutputItemAdded,
|
||||
#[strum(serialize = "response.output_item.done")]
|
||||
OutputItemDone,
|
||||
#[strum(serialize = "response.output_text.annotation.added")]
|
||||
OutputTextAnnotationAdded,
|
||||
#[strum(serialize = "response.output_text.delta")]
|
||||
OutputTextDelta,
|
||||
#[strum(serialize = "response.output_text.done")]
|
||||
OutputTextDone,
|
||||
#[strum(serialize = "response.reasoning_summary_part.added")]
|
||||
ReasoningSummaryPartAdded,
|
||||
#[strum(serialize = "response.reasoning_summary_part.done")]
|
||||
ReasoningSummaryPartDone,
|
||||
#[strum(serialize = "response.reasoning_summary_text.delta")]
|
||||
ReasoningSummaryTextDelta,
|
||||
#[strum(serialize = "response.reasoning_summary_text.done")]
|
||||
ReasoningSummaryTextDone,
|
||||
#[strum(serialize = "response.reasoning_text.delta")]
|
||||
ReasoningTextDelta,
|
||||
#[strum(serialize = "response.reasoning_text.done")]
|
||||
ReasoningTextDone,
|
||||
#[strum(serialize = "response.refusal.delta")]
|
||||
RefusalDelta,
|
||||
#[strum(serialize = "response.refusal.done")]
|
||||
RefusalDone,
|
||||
#[strum(serialize = "response.shell_call_command.added")]
|
||||
ShellCallCommandAdded,
|
||||
#[strum(serialize = "response.shell_call_command.delta")]
|
||||
ShellCallCommandDelta,
|
||||
#[strum(serialize = "response.shell_call_command.done")]
|
||||
ShellCallCommandDone,
|
||||
#[strum(serialize = "response.shell_call_output_content.delta")]
|
||||
ShellCallOutputContentDelta,
|
||||
#[strum(serialize = "response.shell_call_output_content.done")]
|
||||
ShellCallOutputContentDone,
|
||||
#[strum(serialize = "response.web_search_call.completed")]
|
||||
WebSearchCallCompleted,
|
||||
#[strum(serialize = "response.web_search_call.in_progress")]
|
||||
WebSearchCallInProgress,
|
||||
#[strum(serialize = "response.web_search_call.searching")]
|
||||
WebSearchCallSearching,
|
||||
#[strum(default, transparent)]
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl ResponsesEventType {
|
||||
/// The last event of a response. `response.completed`, `response.incomplete` and
|
||||
/// `response.failed` each close the response and carry its final body; `error` aborts it.
|
||||
pub fn is_terminal(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::ResponseCompleted | Self::ResponseFailed | Self::ResponseIncomplete | Self::Error
|
||||
)
|
||||
}
|
||||
|
||||
/// A terminal event after which the response is not a success.
|
||||
pub fn is_failure(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::ResponseFailed | Self::ResponseIncomplete | Self::Error
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
pub struct ResponsesStreamEvent {
|
||||
#[serde(rename = "type")]
|
||||
pub event_type: ResponsesEventType,
|
||||
#[serde(flatten)]
|
||||
pub data: Map<String, Value>,
|
||||
}
|
||||
|
||||
impl ResponsesStreamEvent {
|
||||
pub fn model(&self) -> Option<&str> {
|
||||
let model = self.data.get("model").and_then(Value::as_str);
|
||||
if model.is_some() {
|
||||
return model;
|
||||
}
|
||||
self.data
|
||||
.get("response")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|response| response.get("model"))
|
||||
.and_then(Value::as_str)
|
||||
}
|
||||
|
||||
pub fn is_response_create(&self) -> bool {
|
||||
self.event_type == ResponsesEventType::ResponseCreate
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
use strum::IntoEnumIterator;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
fn every_named_event_type_round_trips() {
|
||||
let named: Vec<_> = ResponsesEventType::iter()
|
||||
.filter(|event_type| !matches!(event_type, ResponsesEventType::Other(_)))
|
||||
.collect();
|
||||
assert!(named.len() > 50);
|
||||
for event_type in named {
|
||||
let wire = event_type.to_string();
|
||||
assert_eq!(wire, event_type.as_ref());
|
||||
assert_eq!(serde_json::to_value(&event_type).unwrap(), json!(wire));
|
||||
assert_eq!(wire.parse::<ResponsesEventType>().unwrap(), event_type);
|
||||
assert_eq!(
|
||||
serde_json::from_value::<ResponsesEventType>(json!(wire)).unwrap(),
|
||||
event_type
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unknown("future.event", ResponsesEventType::Other("future.event".into()))]
|
||||
#[case::empty("", ResponsesEventType::Other(String::new()))]
|
||||
#[case::case_sensitive(
|
||||
"Response.Completed",
|
||||
ResponsesEventType::Other("Response.Completed".into())
|
||||
)]
|
||||
#[case::escaped("future\"\\\n", ResponsesEventType::Other("future\"\\\n".into()))]
|
||||
fn unknown_event_types_round_trip_through_other(
|
||||
#[case] wire: &str,
|
||||
#[case] expected: ResponsesEventType,
|
||||
) {
|
||||
assert_eq!(
|
||||
serde_json::from_value::<ResponsesEventType>(json!(wire)).unwrap(),
|
||||
expected
|
||||
);
|
||||
assert_eq!(serde_json::to_value(&expected).unwrap(), json!(wire));
|
||||
assert_eq!(expected.as_ref(), wire);
|
||||
assert_eq!(expected.to_string(), wire);
|
||||
}
|
||||
|
||||
/// Python's `ResponsesAPIStreamingIterator` stores the response on `response.completed`,
|
||||
/// `response.incomplete` and `response.failed`, and raises on `response.failed` and `error`.
|
||||
#[rstest]
|
||||
#[case::completed(ResponsesEventType::ResponseCompleted, true, false)]
|
||||
#[case::incomplete(ResponsesEventType::ResponseIncomplete, true, true)]
|
||||
#[case::failed(ResponsesEventType::ResponseFailed, true, true)]
|
||||
#[case::error(ResponsesEventType::Error, true, true)]
|
||||
#[case::create(ResponsesEventType::ResponseCreate, false, false)]
|
||||
#[case::created(ResponsesEventType::ResponseCreated, false, false)]
|
||||
#[case::in_progress(ResponsesEventType::ResponseInProgress, false, false)]
|
||||
#[case::delta(ResponsesEventType::OutputTextDelta, false, false)]
|
||||
#[case::item_done(ResponsesEventType::OutputItemDone, false, false)]
|
||||
#[case::unknown(ResponsesEventType::Other("response.future".into()), false, false)]
|
||||
#[case::tool_failed(ResponsesEventType::McpCallFailed, false, false)]
|
||||
fn terminal_and_failure_events_match_the_python_iterator(
|
||||
#[case] event_type: ResponsesEventType,
|
||||
#[case] terminal: bool,
|
||||
#[case] failure: bool,
|
||||
) {
|
||||
assert_eq!(event_type.is_terminal(), terminal);
|
||||
assert_eq!(event_type.is_failure(), failure);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::number("17")]
|
||||
#[case::boolean("true")]
|
||||
#[case::null("null")]
|
||||
#[case::array("[]")]
|
||||
#[case::object("{}")]
|
||||
fn event_types_reject_non_strings(#[case] wire: &str) {
|
||||
assert!(serde_json::from_str::<ResponsesEventType>(wire).is_err());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::text(json!({"type":"response.output_text.delta","delta":"hello","sequence_number":3,"future":null}))]
|
||||
#[case::tool(json!({"type":"response.function_call_arguments.delta","delta":"{\"x\":","item_id":"fc_1","output_index":0}))]
|
||||
#[case::unknown(json!({"type":"future.event","payload":{"items":[1,null]},"future":true}))]
|
||||
#[case::terminal(json!({"type":"response.completed","response":{"id":"r","output":[],"future":null}}))]
|
||||
fn events_preserve_every_payload_field(#[case] wire: Value) {
|
||||
let event: ResponsesStreamEvent = serde_json::from_value(wire.clone()).unwrap();
|
||||
assert_eq!(serde_json::to_value(event).unwrap(), wire);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::missing_type(json!({"delta":"hello"}))]
|
||||
#[case::numeric_type(json!({"type":7}))]
|
||||
#[case::null_type(json!({"type":null}))]
|
||||
#[case::non_object(json!([]))]
|
||||
fn malformed_events_are_rejected(#[case] wire: Value) {
|
||||
assert!(serde_json::from_value::<ResponsesStreamEvent>(wire).is_err());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::flat(json!({"type":"response.create","model":"model-a"}), Some("model-a"))]
|
||||
#[case::nested(json!({"type":"response.create","response":{"model":"model-b"}}), Some("model-b"))]
|
||||
#[case::none(json!({"type":"response.created","response":{}}), None)]
|
||||
fn model_reads_flat_and_nested_shapes(#[case] wire: Value, #[case] expected: Option<&str>) {
|
||||
let event: ResponsesStreamEvent = serde_json::from_value(wire).unwrap();
|
||||
assert_eq!(event.model(), expected);
|
||||
}
|
||||
|
||||
#[cfg(feature = "schema")]
|
||||
#[rstest]
|
||||
fn event_type_schema_accepts_unknown_strings() {
|
||||
let schema = schemars::schema_for!(ResponsesEventType).to_value();
|
||||
assert_eq!(schema.get("type"), Some(&json!("string")));
|
||||
assert!(schema.get("enum").is_none());
|
||||
}
|
||||
}
|
||||
|
|
@ -2,68 +2,6 @@ use super::ResponsesOutputItem;
|
|||
use crate::recognized::Recognized;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
#[derive(
|
||||
Clone,
|
||||
Debug,
|
||||
PartialEq,
|
||||
Eq,
|
||||
strum::AsRefStr,
|
||||
strum::EnumString,
|
||||
strum::Display,
|
||||
serde_with::DeserializeFromStr,
|
||||
serde_with::SerializeDisplay,
|
||||
)]
|
||||
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
|
||||
#[cfg_attr(feature = "schema", schemars(with = "String"))]
|
||||
pub enum ResponsesWsEventType {
|
||||
#[strum(serialize = "response.create")]
|
||||
ResponseCreate,
|
||||
#[strum(serialize = "response.created")]
|
||||
ResponseCreated,
|
||||
#[strum(serialize = "response.completed")]
|
||||
ResponseCompleted,
|
||||
#[strum(serialize = "response.failed")]
|
||||
ResponseFailed,
|
||||
#[strum(serialize = "response.incomplete")]
|
||||
ResponseIncomplete,
|
||||
#[strum(serialize = "error")]
|
||||
Error,
|
||||
#[strum(default, transparent)]
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl ResponsesWsEventType {
|
||||
pub fn as_str(&self) -> &str {
|
||||
self.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
pub struct ResponsesWsEvent {
|
||||
#[serde(rename = "type")]
|
||||
pub event_type: ResponsesWsEventType,
|
||||
#[serde(flatten)]
|
||||
pub data: Map<String, Value>,
|
||||
}
|
||||
|
||||
impl ResponsesWsEvent {
|
||||
pub fn model(&self) -> Option<&str> {
|
||||
let model = self.data.get("model").and_then(Value::as_str);
|
||||
if model.is_some() {
|
||||
return model;
|
||||
}
|
||||
self.data
|
||||
.get("response")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|response| response.get("model"))
|
||||
.and_then(Value::as_str)
|
||||
}
|
||||
|
||||
pub fn is_response_create(&self) -> bool {
|
||||
self.event_type == ResponsesWsEventType::ResponseCreate
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
#[derive(Eq)]
|
||||
pub struct ResponsesErrorFrame {
|
||||
|
|
@ -112,8 +50,7 @@ mod tests {
|
|||
use crate::formats::responses::ResponsesOutputItem;
|
||||
|
||||
use crate::{
|
||||
formats::responses::streaming_websocket::{ResponsesEventResponse, ResponsesWsEventType},
|
||||
recognized::Recognized,
|
||||
formats::responses::streaming_websocket::ResponsesEventResponse, recognized::Recognized,
|
||||
test_support::*,
|
||||
};
|
||||
|
||||
|
|
@ -134,66 +71,6 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::flat(serde_json::json!({"type":"response.create","model":"model-a"}), Some("model-a"))]
|
||||
#[case::nested(serde_json::json!({"type":"response.create","response":{"model":"model-b"}}), Some("model-b"))]
|
||||
fn model_reads_flat_and_nested_create_shapes(
|
||||
#[case] payload: serde_json::Value,
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
let event: ResponsesWsEvent = serde_json::from_value(payload).expect("valid event");
|
||||
assert_eq!(event.model(), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::create("response.create", ResponsesWsEventType::ResponseCreate)]
|
||||
#[case::created("response.created", ResponsesWsEventType::ResponseCreated)]
|
||||
#[case::completed("response.completed", ResponsesWsEventType::ResponseCompleted)]
|
||||
#[case::failed("response.failed", ResponsesWsEventType::ResponseFailed)]
|
||||
#[case::incomplete("response.incomplete", ResponsesWsEventType::ResponseIncomplete)]
|
||||
#[case::error("error", ResponsesWsEventType::Error)]
|
||||
#[case::unknown(
|
||||
"response.output_text.delta",
|
||||
ResponsesWsEventType::Other("response.output_text.delta".to_string())
|
||||
)]
|
||||
#[case::empty("", ResponsesWsEventType::Other(String::new()))]
|
||||
#[case::case_sensitive(
|
||||
"Response.Completed",
|
||||
ResponsesWsEventType::Other("Response.Completed".into())
|
||||
)]
|
||||
#[case::escaped("future\"\\\n", ResponsesWsEventType::Other("future\"\\\n".into()))]
|
||||
fn websocket_event_type_round_trips(
|
||||
#[case] wire: &str,
|
||||
#[case] expected: ResponsesWsEventType,
|
||||
) {
|
||||
let serialized = serde_json::to_string(&expected).unwrap();
|
||||
assert_eq!(serialized, serde_json::to_string(wire).unwrap());
|
||||
assert_eq!(
|
||||
serde_json::from_str::<ResponsesWsEventType>(&serialized).unwrap(),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::number("17")]
|
||||
#[case::boolean("true")]
|
||||
#[case::null("null")]
|
||||
#[case::array("[]")]
|
||||
#[case::object("{}")]
|
||||
fn websocket_event_type_rejects_non_strings(#[case] wire: &str) {
|
||||
assert!(serde_json::from_str::<ResponsesWsEventType>(wire).is_err());
|
||||
}
|
||||
|
||||
#[cfg(feature = "schema")]
|
||||
#[rstest]
|
||||
fn websocket_event_type_schema_is_open_string() {
|
||||
let schema = schemars::schema_for!(ResponsesWsEventType);
|
||||
assert_eq!(
|
||||
schema.to_value().get("type"),
|
||||
Some(&serde_json::json!("string"))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn nested_event_response_exposes_typed_output_and_preserves_extensions() {
|
||||
let wire = json!({
|
||||
|
|
|
|||
|
|
@ -32,6 +32,13 @@
|
|||
- Provider-wide docs (auth, errors, regions): `src/<provider>/`
|
||||
- One host's endpoint docs: `src/<provider>/<format>/`
|
||||
|
||||
## Streaming
|
||||
|
||||
- Each format's `src/base_llm/<format>/streaming.rs` exposes the same three names: `event_stream` (bytes to typed events), `decode` (a whole stream to typed events) and `encode` (one typed event to the SSE frame LiteLLM's proxy sends clients). Callers use the module path, never a vendor prefix
|
||||
- SSE framing, `[DONE]` and the generic helpers live in `src/base_llm/sse.rs`. Messages frames are `event:` plus `data:`; Chat Completions and Responses frames are bare `data:`, as Python's proxy writes them
|
||||
- Decoding is strict: a malformed event, or `[DONE]` on a Messages stream, is `InvalidResponse` and ends the stream. Python never fails a stream this way; the divergence is deliberate and `src/base_llm/sse.rs` lists what Python does per format. A stream cut inside an event drops that event, as the SSE spec requires
|
||||
- What ends or fails a stream belongs on the format's event type (`is_terminal`, `is_failure`). The cache's success check in `inference/src/caching.rs` still matches type strings; move it onto the typed events rather than adding new string matches
|
||||
|
||||
## Python pairs
|
||||
|
||||
- Python paths identify counterparts but do not dictate Rust module names or class hierarchy. Preserve behavior and concepts, not structure
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ use crate::{
|
|||
Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param,
|
||||
},
|
||||
},
|
||||
messages::streaming::anthropic_sse_event_stream,
|
||||
messages::streaming::event_stream as messages_event_stream,
|
||||
},
|
||||
};
|
||||
|
||||
|
|
@ -196,7 +196,7 @@ impl BaseConfig for AnthropicConfig {
|
|||
|
||||
fn model_response_iterator(&self, shape: StreamShape) -> Option<ChatStream> {
|
||||
Some(ChatStream::new(
|
||||
anthropic_sse_event_stream,
|
||||
messages_event_stream,
|
||||
ModelResponseIterator::new(shape),
|
||||
))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,15 +1,32 @@
|
|||
use std::collections::HashMap;
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures_util::{StreamExt, stream::BoxStream};
|
||||
use litellm_llms_types::formats::chat_completions::ChatCompletionChunk;
|
||||
|
||||
use crate::{
|
||||
Error,
|
||||
base_llm::base_model_iterator::{ByteStream, StreamError, StreamTransformer, transform_stream},
|
||||
base_llm::{
|
||||
base_model_iterator::{ByteStream, StreamError, StreamTransformer, transform_stream},
|
||||
sse::{self, is_done},
|
||||
},
|
||||
};
|
||||
|
||||
pub type ChatChunkStream = BoxStream<'static, Result<ChatCompletionChunk, Error>>;
|
||||
|
||||
/// OpenAI-compatible SSE: bare `data:` frames closed by `data: [DONE]`.
|
||||
pub fn event_stream(bytes: ByteStream) -> ChatChunkStream {
|
||||
sse::event_stream(bytes, is_done)
|
||||
}
|
||||
|
||||
pub fn decode(wire: &[u8]) -> Result<Vec<ChatCompletionChunk>, Error> {
|
||||
sse::decode(wire, is_done)
|
||||
}
|
||||
|
||||
pub fn encode(chunk: &ChatCompletionChunk) -> Result<Bytes, Error> {
|
||||
sse::encode_data(chunk)
|
||||
}
|
||||
|
||||
/// What Python's `map_openai_params` decides about the stream and `completion`
|
||||
/// hands to `ModelResponseIterator`: it is settled while the request is built,
|
||||
/// never re-derived from the body.
|
||||
|
|
@ -52,3 +69,109 @@ impl ChatStream {
|
|||
(self.run)(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn in_pieces(wire: &[u8], size: usize) -> ByteStream {
|
||||
let pieces: Vec<Bytes> = wire.chunks(size).map(Bytes::copy_from_slice).collect();
|
||||
stream::iter(pieces.into_iter().map(Ok)).boxed()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::role(json!({"id":"c","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":""}}]}))]
|
||||
#[case::text(json!({"id":"c","choices":[{"index":0,"delta":{"content":"hi\nthere"}}]}))]
|
||||
#[case::tool(json!({"id":"c","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"t","type":"function","function":{"name":"f","arguments":"{\"x\":"}}]}}]}))]
|
||||
#[case::finish(json!({"id":"c","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}))]
|
||||
#[case::usage(json!({"id":"c","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3,"prompt_tokens_details":{"cached_tokens":0,"cache_creation_tokens":0,"text_tokens":1}}}))]
|
||||
#[case::host_extension(json!({"id":"c","choices":[],"citations":[{"url":"u"}]}))]
|
||||
fn encoded_chunks_are_bare_data_frames_that_decode_back(#[case] wire: Value) {
|
||||
let chunk: ChatCompletionChunk = serde_json::from_value(wire.clone()).unwrap();
|
||||
let encoded = encode(&chunk).unwrap();
|
||||
assert!(encoded.starts_with(b"data: "));
|
||||
assert!(encoded.ends_with(b"\n\n"));
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&encoded[b"data: ".len()..encoded.len() - 2]).unwrap(),
|
||||
wire
|
||||
);
|
||||
assert_eq!(decode(&encoded).unwrap(), vec![chunk]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::lf("\n")]
|
||||
#[case::crlf("\r\n")]
|
||||
#[case::cr("\r")]
|
||||
#[tokio::test]
|
||||
async fn streams_split_anywhere_skip_comments_and_done(
|
||||
#[case] newline: &str,
|
||||
#[values(1, 5, 4096)] size: usize,
|
||||
) {
|
||||
let wire = format!(
|
||||
": keepalive{newline}{newline}data: {{\"id\":\"c\",\"choices\":[{{\"index\":0,\"delta\":{{\"content\":\"a\"}}}}]}}{newline}{newline}data:{{\"id\":\"c\",\"choices\":[{{\"index\":0,\"delta\":{{}},\"finish_reason\":\"stop\"}}]}}{newline}{newline}data: [DONE]{newline}{newline}"
|
||||
);
|
||||
let streamed = event_stream(in_pieces(wire.as_bytes(), size))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(streamed, decode(wire.as_bytes()).unwrap());
|
||||
assert_eq!(streamed.len(), 2);
|
||||
assert_eq!(streamed[0].choices[0].delta.content.as_deref(), Some("a"));
|
||||
assert_eq!(
|
||||
streamed[1].choices[0].finish_reason.as_deref(),
|
||||
Some("stop")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::spaced(b"data: [DONE]\n\n")]
|
||||
#[case::unspaced(b"data:[DONE]\n\n")]
|
||||
#[case::trailing_space(b"data: [DONE] \n\n")]
|
||||
fn every_done_spelling_closes_the_stream_without_a_chunk(#[case] wire: &[u8]) {
|
||||
assert_eq!(decode(wire).unwrap(), vec![]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::invalid_json(b"data: {\n\n")]
|
||||
#[case::choices_not_a_list(b"data: {\"choices\":{}}\n\n")]
|
||||
#[case::invalid_utf8(b"data: \xff\n\n")]
|
||||
#[tokio::test]
|
||||
async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) {
|
||||
assert!(matches!(decode(wire), Err(Error::InvalidResponse(_))));
|
||||
let streamed = event_stream(in_pieces(wire, 1)).collect::<Vec<_>>().await;
|
||||
assert!(matches!(
|
||||
streamed.last(),
|
||||
Some(Err(Error::InvalidResponse(_))) | None
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_body_error_ends_the_stream_after_delivered_chunks() {
|
||||
let mut chunks = event_stream(
|
||||
stream::iter([
|
||||
Ok(Bytes::from_static(
|
||||
b"data: {\"id\":\"c\",\"choices\":[]}\n\n",
|
||||
)),
|
||||
Err(std::io::Error::new(
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"reset",
|
||||
)),
|
||||
])
|
||||
.boxed(),
|
||||
);
|
||||
assert_eq!(
|
||||
chunks.next().await.unwrap().unwrap().id.as_deref(),
|
||||
Some("c")
|
||||
);
|
||||
assert!(matches!(
|
||||
chunks.next().await,
|
||||
Some(Err(Error::InvalidResponse(_)))
|
||||
));
|
||||
assert!(chunks.next().await.is_none());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,38 +1,32 @@
|
|||
use bytes::Bytes;
|
||||
use futures_util::{StreamExt, stream::BoxStream};
|
||||
use litellm_framer::{frames, sse::SseCodec};
|
||||
use litellm_llms_types::formats::messages::streaming::MessagesStreamEvent;
|
||||
|
||||
use crate::Error;
|
||||
pub use crate::base_llm::base_model_iterator::ByteStream;
|
||||
use crate::{Error, base_llm::sse};
|
||||
|
||||
pub type EventStream = BoxStream<'static, Result<MessagesStreamEvent, Error>>;
|
||||
pub type EventStream = sse::EventStream<MessagesStreamEvent>;
|
||||
pub type StreamDecoder = fn(ByteStream) -> EventStream;
|
||||
|
||||
pub fn anthropic_sse_event_stream(bytes: ByteStream) -> EventStream {
|
||||
Box::pin(frames(bytes, SseCodec::default()).map(|event| {
|
||||
let event = event.map_err(|error| {
|
||||
Error::InvalidResponse(crate::ErrorDetail::failed("stream framing", error))
|
||||
})?;
|
||||
serde_json::from_str(&event.data).map_err(|error| {
|
||||
Error::InvalidResponse(crate::ErrorDetail::invalid("Anthropic stream event", error))
|
||||
})
|
||||
}))
|
||||
fn no_control_frames(_: &litellm_framer::sse::SseEvent) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
pub fn encode_anthropic_sse(event: &MessagesStreamEvent) -> Result<Bytes, Error> {
|
||||
let data = serde_json::to_value(event).map_err(|error| {
|
||||
Error::InvalidResponse(crate::ErrorDetail::invalid("Anthropic stream event", error))
|
||||
})?;
|
||||
let name = data
|
||||
.get("type")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.ok_or_else(|| {
|
||||
Error::InvalidResponse(
|
||||
"Anthropic stream event is invalid: stream event has no type".into(),
|
||||
)
|
||||
})?;
|
||||
Ok(Bytes::from(format!("event: {name}\ndata: {data}\n\n")))
|
||||
pub fn event_stream(bytes: ByteStream) -> EventStream {
|
||||
sse::event_stream(bytes, no_control_frames)
|
||||
}
|
||||
|
||||
pub fn decode(wire: &[u8]) -> Result<Vec<MessagesStreamEvent>, Error> {
|
||||
sse::decode(wire, no_control_frames)
|
||||
}
|
||||
|
||||
/// Anthropic names every frame after the event's `type`, and LiteLLM's proxy does the same
|
||||
/// when it re-emits translated events.
|
||||
pub fn encode(event: &MessagesStreamEvent) -> Result<Bytes, Error> {
|
||||
let data = serde_json::to_string(event).map_err(sse::invalid)?;
|
||||
Ok(Bytes::from(format!(
|
||||
"event: {}\ndata: {data}\n\n",
|
||||
event.as_ref()
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
@ -41,6 +35,7 @@ mod tests {
|
|||
use litellm_llms_types::formats::messages::streaming::{
|
||||
MessagesContentBlockDelta, MessagesStreamUsage,
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
|
@ -48,51 +43,11 @@ mod tests {
|
|||
const TEXT_DELTA: &str =
|
||||
r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#;
|
||||
|
||||
fn in_pieces(wire: &[u8]) -> ByteStream {
|
||||
let pieces: Vec<Bytes> = wire.chunks(3).map(Bytes::copy_from_slice).collect();
|
||||
fn in_pieces(wire: &[u8], size: usize) -> ByteStream {
|
||||
let pieces: Vec<Bytes> = wire.chunks(size).map(Bytes::copy_from_slice).collect();
|
||||
stream::iter(pieces.into_iter().map(Ok)).boxed()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sse_frames_split_anywhere_decode_into_typed_events() {
|
||||
let wire = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n");
|
||||
let events = anthropic_sse_event_stream(in_pieces(wire.as_bytes()))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![MessagesStreamEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: MessagesContentBlockDelta::TextDelta {
|
||||
text: "hello".into(),
|
||||
},
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decodes_citations_delta_events() {
|
||||
let wire = concat!(
|
||||
"event: content_block_delta\n",
|
||||
r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#,
|
||||
"\n\n",
|
||||
);
|
||||
let events = anthropic_sse_event_stream(in_pieces(wire.as_bytes()))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
events.as_slice(),
|
||||
[MessagesStreamEvent::ContentBlockDelta {
|
||||
delta: MessagesContentBlockDelta::Citations { .. },
|
||||
..
|
||||
}]
|
||||
));
|
||||
}
|
||||
|
||||
fn events() -> Vec<MessagesStreamEvent> {
|
||||
vec![
|
||||
MessagesStreamEvent::Ping,
|
||||
|
|
@ -110,27 +65,80 @@ mod tests {
|
|||
]
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn encoded_events_decode_back_to_themselves() {
|
||||
let wire = events()
|
||||
.iter()
|
||||
.map(encode_anthropic_sse)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.unwrap();
|
||||
fn wire(newline: &str) -> String {
|
||||
format!(
|
||||
"event: content_block_delta{newline}data: {TEXT_DELTA}{newline}{newline}: keepalive{newline}{newline}event: message_stop{newline}data: {{\"type\":\"message_stop\"}}{newline}{newline}"
|
||||
)
|
||||
}
|
||||
|
||||
let decoded = anthropic_sse_event_stream(stream::iter(wire.into_iter().map(Ok)).boxed())
|
||||
#[rstest]
|
||||
#[case::lf("\n")]
|
||||
#[case::crlf("\r\n")]
|
||||
#[case::cr("\r")]
|
||||
#[tokio::test]
|
||||
async fn frames_split_anywhere_decode_into_typed_events(
|
||||
#[case] newline: &str,
|
||||
#[values(1, 3, 4096)] size: usize,
|
||||
) {
|
||||
let wire = wire(newline);
|
||||
let expected = vec![
|
||||
MessagesStreamEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: MessagesContentBlockDelta::TextDelta {
|
||||
text: "hello".into(),
|
||||
},
|
||||
},
|
||||
MessagesStreamEvent::MessageStop { usage: None },
|
||||
];
|
||||
let streamed = event_stream(in_pieces(wire.as_bytes(), size))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(streamed, expected);
|
||||
assert_eq!(decode(wire.as_bytes()).unwrap(), expected);
|
||||
}
|
||||
|
||||
assert_eq!(decoded, events());
|
||||
#[tokio::test]
|
||||
async fn decodes_citations_delta_events() {
|
||||
let wire = concat!(
|
||||
"event: content_block_delta\n",
|
||||
r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#,
|
||||
"\n\n",
|
||||
);
|
||||
let events = event_stream(in_pieces(wire.as_bytes(), 3))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
events.as_slice(),
|
||||
[MessagesStreamEvent::ContentBlockDelta {
|
||||
delta: MessagesContentBlockDelta::Citations { .. },
|
||||
..
|
||||
}]
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn encoded_events_decode_back_to_themselves() {
|
||||
let wire: Vec<u8> = events()
|
||||
.iter()
|
||||
.map(encode)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect();
|
||||
let streamed = event_stream(in_pieces(&wire, 3))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(streamed, events());
|
||||
assert_eq!(decode(&wire).unwrap(), events());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_event_is_named_by_its_type() {
|
||||
let encoded =
|
||||
encode_anthropic_sse(&MessagesStreamEvent::MessageStop { usage: None }).unwrap();
|
||||
|
||||
let encoded = encode(&MessagesStreamEvent::MessageStop { usage: None }).unwrap();
|
||||
assert_eq!(
|
||||
encoded,
|
||||
Bytes::from(format!(
|
||||
|
|
@ -139,4 +147,68 @@ mod tests {
|
|||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::json(b"data: {\n\n")]
|
||||
#[case::unknown_event(b"data: {\"type\":\"future_event\"}\n\n")]
|
||||
#[case::malformed_tool(b"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":3}}\n\n")]
|
||||
#[case::utf8(b"data: \xff\n\n")]
|
||||
#[case::done_is_not_an_anthropic_frame(b"data: [DONE]\n\n")]
|
||||
#[tokio::test]
|
||||
async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) {
|
||||
assert!(matches!(decode(wire), Err(Error::InvalidResponse(_))));
|
||||
let streamed = event_stream(in_pieces(wire, 1)).collect::<Vec<_>>().await;
|
||||
assert!(matches!(
|
||||
streamed.last(),
|
||||
Some(Err(Error::InvalidResponse(_))) | None
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn decoding_keeps_tool_deltas_and_block_extensions() {
|
||||
let wire = concat!(
|
||||
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"t\",\"name\":\"lookup\",\"input\":{},\"future\":null}}\n\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"x\\\":\"}}\n\n",
|
||||
);
|
||||
let decoded = decode(wire.as_bytes()).unwrap();
|
||||
let [
|
||||
MessagesStreamEvent::ContentBlockStart { content_block, .. },
|
||||
MessagesStreamEvent::ContentBlockDelta {
|
||||
delta: MessagesContentBlockDelta::InputJsonDelta { partial_json },
|
||||
..
|
||||
},
|
||||
] = decoded.as_slice()
|
||||
else {
|
||||
panic!("expected tool events");
|
||||
};
|
||||
assert_eq!(
|
||||
content_block.extra.get("future"),
|
||||
Some(&serde_json::Value::Null)
|
||||
);
|
||||
assert_eq!(partial_json, "{\"x\":");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_body_error_ends_the_stream_after_delivered_events() {
|
||||
let mut events = event_stream(
|
||||
stream::iter([
|
||||
Ok(Bytes::from_static(b"data: {\"type\":\"ping\"}\n\n")),
|
||||
Err(std::io::Error::new(
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"reset",
|
||||
)),
|
||||
])
|
||||
.boxed(),
|
||||
);
|
||||
assert_eq!(
|
||||
events.next().await.unwrap().unwrap(),
|
||||
MessagesStreamEvent::Ping
|
||||
);
|
||||
assert!(matches!(
|
||||
events.next().await,
|
||||
Some(Err(Error::InvalidResponse(_)))
|
||||
));
|
||||
assert!(events.next().await.is_none());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,3 +5,4 @@ pub mod chat;
|
|||
pub mod messages;
|
||||
pub mod ocr;
|
||||
pub mod responses;
|
||||
pub mod sse;
|
||||
|
|
|
|||
|
|
@ -1 +1,2 @@
|
|||
pub mod streaming;
|
||||
pub mod transformation;
|
||||
|
|
|
|||
135
litellm-rust/crates/llms/src/base_llm/responses/streaming.rs
Normal file
135
litellm-rust/crates/llms/src/base_llm/responses/streaming.rs
Normal file
|
|
@ -0,0 +1,135 @@
|
|||
use bytes::Bytes;
|
||||
use litellm_llms_types::formats::responses::streaming::ResponsesStreamEvent;
|
||||
|
||||
pub use crate::base_llm::base_model_iterator::ByteStream;
|
||||
use crate::{
|
||||
Error,
|
||||
base_llm::sse::{self, is_done},
|
||||
};
|
||||
|
||||
pub type EventStream = sse::EventStream<ResponsesStreamEvent>;
|
||||
|
||||
pub fn event_stream(bytes: ByteStream) -> EventStream {
|
||||
sse::event_stream(bytes, is_done)
|
||||
}
|
||||
|
||||
pub fn decode(wire: &[u8]) -> Result<Vec<ResponsesStreamEvent>, Error> {
|
||||
sse::decode(wire, is_done)
|
||||
}
|
||||
|
||||
/// A bare `data:` frame. Python's proxy serializes every Responses event this way and never
|
||||
/// writes an `event:` line, so clients see one shape whichever runtime served them.
|
||||
pub fn encode(event: &ResponsesStreamEvent) -> Result<Bytes, Error> {
|
||||
sse::encode_data(event)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||
use litellm_llms_types::formats::responses::streaming::ResponsesEventType;
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn in_pieces(wire: &[u8], size: usize) -> ByteStream {
|
||||
let pieces: Vec<Bytes> = wire.chunks(size).map(Bytes::copy_from_slice).collect();
|
||||
stream::iter(pieces.into_iter().map(Ok)).boxed()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::created(json!({"type":"response.created","response":{"id":"r","output":[]}}))]
|
||||
#[case::text(json!({"type":"response.output_text.delta","delta":"hello\nworld","sequence_number":3}))]
|
||||
#[case::tool(json!({"type":"response.function_call_arguments.delta","delta":"{\"x\":","item_id":"fc_1"}))]
|
||||
#[case::unknown(json!({"type":"future.event","payload":[1,null],"future":{"nested":true}}))]
|
||||
#[case::completed(json!({"type":"response.completed","response":{"id":"r","usage":{"input_tokens":3}}}))]
|
||||
#[case::failed(json!({"type":"response.failed","response":{"error":{"code":"failure"}}}))]
|
||||
#[case::incomplete(json!({"type":"response.incomplete","response":{"incomplete_details":{"reason":"limit"}}}))]
|
||||
#[case::error(json!({"type":"error","error":{"message":"boom"}}))]
|
||||
fn encoded_events_are_bare_data_frames_that_decode_back(#[case] wire: Value) {
|
||||
let event: ResponsesStreamEvent = serde_json::from_value(wire.clone()).unwrap();
|
||||
let encoded = encode(&event).unwrap();
|
||||
assert!(encoded.starts_with(b"data: "));
|
||||
assert!(encoded.ends_with(b"\n\n"));
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&encoded[b"data: ".len()..encoded.len() - 2]).unwrap(),
|
||||
wire
|
||||
);
|
||||
assert_eq!(decode(&encoded).unwrap(), vec![event]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::lf("\n")]
|
||||
#[case::crlf("\r\n")]
|
||||
#[case::cr("\r")]
|
||||
#[tokio::test]
|
||||
async fn streams_split_anywhere_skip_comments_and_done(
|
||||
#[case] newline: &str,
|
||||
#[values(1, 5, 4096)] size: usize,
|
||||
) {
|
||||
let wire = format!(
|
||||
"\u{feff}event: response.created{newline}data: {{\"type\":\"response.created\"}}{newline}{newline}: keepalive{newline}{newline}event: future.event{newline}data: {{\"type\":\"future.event\",\"payload\":null}}{newline}{newline}data: [DONE]{newline}{newline}"
|
||||
);
|
||||
let streamed = event_stream(in_pieces(wire.as_bytes(), size))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(streamed, decode(wire.as_bytes()).unwrap());
|
||||
assert_eq!(streamed.len(), 2);
|
||||
assert_eq!(streamed[0].event_type, ResponsesEventType::ResponseCreated);
|
||||
assert_eq!(
|
||||
streamed[1].event_type,
|
||||
ResponsesEventType::Other("future.event".into())
|
||||
);
|
||||
assert_eq!(streamed[1].data.get("payload"), Some(&Value::Null));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::spaced(b"data: [DONE]\n\n")]
|
||||
#[case::unspaced(b"data:[DONE]\n\n")]
|
||||
#[case::trailing_space(b"data: [DONE] \n\n")]
|
||||
fn every_done_spelling_closes_the_stream_without_an_event(#[case] wire: &[u8]) {
|
||||
assert_eq!(decode(wire).unwrap(), vec![]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::invalid_json(b"data: {\n\n")]
|
||||
#[case::missing_type(b"data: {}\n\n")]
|
||||
#[case::numeric_type(b"data: {\"type\":3}\n\n")]
|
||||
#[case::invalid_utf8(b"data: \xff\n\n")]
|
||||
#[tokio::test]
|
||||
async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) {
|
||||
assert!(matches!(decode(wire), Err(Error::InvalidResponse(_))));
|
||||
let streamed = event_stream(in_pieces(wire, 1)).collect::<Vec<_>>().await;
|
||||
assert!(matches!(
|
||||
streamed.last(),
|
||||
Some(Err(Error::InvalidResponse(_))) | None
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_body_error_ends_the_stream_after_delivered_events() {
|
||||
let mut events = event_stream(
|
||||
stream::iter([
|
||||
Ok(Bytes::from_static(
|
||||
b"data: {\"type\":\"response.created\"}\n\n",
|
||||
)),
|
||||
Err(std::io::Error::new(
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"reset",
|
||||
)),
|
||||
])
|
||||
.boxed(),
|
||||
);
|
||||
assert_eq!(
|
||||
events.next().await.unwrap().unwrap().event_type,
|
||||
ResponsesEventType::ResponseCreated
|
||||
);
|
||||
assert!(matches!(
|
||||
events.next().await,
|
||||
Some(Err(Error::InvalidResponse(_)))
|
||||
));
|
||||
assert!(events.next().await.is_none());
|
||||
}
|
||||
}
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
use litellm_llms_types::formats::responses::{
|
||||
ResponsesApiResponse, streaming_websocket::ResponsesWsEvent,
|
||||
ResponsesApiResponse, streaming::ResponsesStreamEvent,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
|
@ -38,11 +38,11 @@ pub trait BaseResponsesApiConfig: Sync {
|
|||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ResponsesWsTransformResult {
|
||||
pub events: Vec<ResponsesWsEvent>,
|
||||
pub events: Vec<ResponsesStreamEvent>,
|
||||
}
|
||||
|
||||
impl ResponsesWsTransformResult {
|
||||
pub fn passthrough(event: ResponsesWsEvent) -> Self {
|
||||
pub fn passthrough(event: ResponsesStreamEvent) -> Self {
|
||||
Self {
|
||||
events: vec![event],
|
||||
}
|
||||
|
|
@ -58,18 +58,18 @@ pub trait ResponsesWebSocketProviderConfig: Sync {
|
|||
|
||||
fn transform_ws_request(
|
||||
&self,
|
||||
event: &ResponsesWsEvent,
|
||||
event: &ResponsesStreamEvent,
|
||||
model: &str,
|
||||
) -> Result<ResponsesWsTransformResult, Error>;
|
||||
|
||||
fn transform_ws_response(
|
||||
&self,
|
||||
event: &ResponsesWsEvent,
|
||||
event: &ResponsesStreamEvent,
|
||||
model: &str,
|
||||
) -> Result<ResponsesWsTransformResult, Error>;
|
||||
}
|
||||
|
||||
pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent {
|
||||
pub fn enforce_model(event: &ResponsesStreamEvent, model: &str) -> ResponsesStreamEvent {
|
||||
if !event.is_response_create() {
|
||||
return event.clone();
|
||||
}
|
||||
|
|
@ -103,7 +103,7 @@ pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent
|
|||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn event(value: serde_json::Value) -> ResponsesWsEvent {
|
||||
fn event(value: serde_json::Value) -> ResponsesStreamEvent {
|
||||
serde_json::from_value(value).expect("valid event")
|
||||
}
|
||||
|
||||
|
|
|
|||
102
litellm-rust/crates/llms/src/base_llm/sse.rs
Normal file
102
litellm-rust/crates/llms/src/base_llm/sse.rs
Normal file
|
|
@ -0,0 +1,102 @@
|
|||
//! SSE decoding and encoding shared by every format's `streaming` module.
|
||||
//!
|
||||
//! DIVERGENCE: decoding is strict. An event whose `data` is not a valid event of the format is
|
||||
//! `InvalidResponse` and ends the stream, and so is `[DONE]` on a Messages stream. Python
|
||||
//! never fails a stream on a malformed event:
|
||||
//! - Chat Completions (`BaseModelResponseIterator`) and Responses
|
||||
//! (`BaseResponsesAPIStreamingIterator`) skip the event and keep going
|
||||
//! - Messages (`BaseAnthropicMessagesStreamingIterator`) forwards the provider's bytes to the
|
||||
//! client without parsing them
|
||||
//! - Anthropic to Chat (`anthropic/chat/handler.py` `ModelResponseIterator`) buffers the invalid
|
||||
//! lines on the guess that one payload was split across them
|
||||
//!
|
||||
//! Skipping loses content or a tool-call fragment with nothing to show for it, so Rust reports it.
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures_util::{StreamExt, TryStreamExt, future::ready, stream::BoxStream};
|
||||
use litellm_framer::{
|
||||
frames,
|
||||
sse::{SseCodec, SseEvent, decode_all},
|
||||
};
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
|
||||
use crate::{Error, ErrorDetail, base_llm::base_model_iterator::ByteStream};
|
||||
|
||||
pub type EventStream<E> = BoxStream<'static, Result<E, Error>>;
|
||||
|
||||
/// Which frames are control frames rather than events of the format.
|
||||
pub type IsControl = fn(&SseEvent) -> bool;
|
||||
|
||||
/// OpenAI's end-of-stream marker on Chat Completions and Responses streams. It closes the
|
||||
/// stream and is not an event; Anthropic Messages streams never carry it.
|
||||
pub const DONE_FRAME: &[u8] = b"data: [DONE]\n\n";
|
||||
|
||||
pub fn is_done(event: &SseEvent) -> bool {
|
||||
event.data.trim() == "[DONE]"
|
||||
}
|
||||
|
||||
pub fn event_stream<E: DeserializeOwned + Send + 'static>(
|
||||
bytes: ByteStream,
|
||||
is_control: IsControl,
|
||||
) -> EventStream<E> {
|
||||
frames(bytes, SseCodec::default())
|
||||
.map_err(framing)
|
||||
.try_filter(move |event| ready(!is_control(event)))
|
||||
.and_then(|event| ready(parse(&event)))
|
||||
.boxed()
|
||||
}
|
||||
|
||||
/// Decodes a whole stream at once. An event the stream ends inside is dropped.
|
||||
pub fn decode<E: DeserializeOwned>(wire: &[u8], is_control: IsControl) -> Result<Vec<E>, Error> {
|
||||
decode_all(wire)
|
||||
.map_err(framing)?
|
||||
.iter()
|
||||
.filter(|event| !is_control(event))
|
||||
.map(parse)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// A bare `data:` frame, which is what LiteLLM's proxy sends to clients on OpenAI formats.
|
||||
pub fn encode_data<E: Serialize>(event: &E) -> Result<Bytes, Error> {
|
||||
let data = serde_json::to_string(event).map_err(invalid)?;
|
||||
Ok(Bytes::from(format!("data: {data}\n\n")))
|
||||
}
|
||||
|
||||
fn parse<E: DeserializeOwned>(event: &SseEvent) -> Result<E, Error> {
|
||||
serde_json::from_str(&event.data).map_err(invalid)
|
||||
}
|
||||
|
||||
pub fn invalid(error: serde_json::Error) -> Error {
|
||||
Error::InvalidResponse(ErrorDetail::invalid("stream event", error))
|
||||
}
|
||||
|
||||
fn framing(error: litellm_framer::SseError) -> Error {
|
||||
Error::InvalidResponse(ErrorDetail::failed("stream framing", error))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::exact("[DONE]", true)]
|
||||
#[case::padded(" [DONE] ", true)]
|
||||
#[case::json("{\"type\":\"x\"}", false)]
|
||||
#[case::mentions("[DONE] was sent", false)]
|
||||
#[case::empty("", false)]
|
||||
fn done_is_an_exact_marker(#[case] data: &str, #[case] done: bool) {
|
||||
let event = SseEvent {
|
||||
data: data.into(),
|
||||
..SseEvent::default()
|
||||
};
|
||||
assert_eq!(is_done(&event), done);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn the_done_frame_we_send_decodes_as_a_control_frame() {
|
||||
let events: Vec<serde_json::Value> = decode(DONE_FRAME, is_done).unwrap();
|
||||
assert!(events.is_empty());
|
||||
}
|
||||
}
|
||||
|
|
@ -102,7 +102,7 @@ mod tests {
|
|||
use litellm_llms_types::formats::messages::streaming::MessagesContentBlockDelta;
|
||||
|
||||
use super::*;
|
||||
use crate::base_llm::messages::streaming::anthropic_sse_event_stream;
|
||||
use crate::base_llm::messages::streaming;
|
||||
|
||||
fn in_pieces(wire: &[u8]) -> ByteStream {
|
||||
let pieces: Vec<Bytes> = wire.chunks(3).map(Bytes::copy_from_slice).collect();
|
||||
|
|
@ -147,7 +147,7 @@ mod tests {
|
|||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let from_sse = anthropic_sse_event_stream(in_pieces(sse.as_bytes()))
|
||||
let from_sse = streaming::event_stream(in_pieces(sse.as_bytes()))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
|
|
|||
|
|
@ -332,7 +332,7 @@ mod tests {
|
|||
use litellm_auth_aws::constants::DEFAULT_BEDROCK_REGION;
|
||||
|
||||
use super::*;
|
||||
use crate::base_llm::messages::streaming::encode_anthropic_sse;
|
||||
use crate::base_llm::messages::streaming;
|
||||
|
||||
fn event(value: Value) -> MessagesStreamEvent {
|
||||
serde_json::from_value(value).unwrap()
|
||||
|
|
@ -493,7 +493,7 @@ mod tests {
|
|||
.boxed();
|
||||
|
||||
let sse = bedrock_anthropic_messages_event_stream(bytes)
|
||||
.map_ok(|event| encode_anthropic_sse(&event).unwrap())
|
||||
.map_ok(|event| streaming::encode(&event).unwrap())
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap()
|
||||
|
|
@ -508,7 +508,7 @@ mod tests {
|
|||
)),
|
||||
]
|
||||
.iter()
|
||||
.flat_map(|event| encode_anthropic_sse(event).unwrap())
|
||||
.flat_map(|event| streaming::encode(event).unwrap())
|
||||
.collect();
|
||||
assert_eq!(sse, expected);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
use litellm_llms_types::formats::responses::{
|
||||
ResponsesApiResponse, streaming_websocket::ResponsesWsEvent,
|
||||
ResponsesApiResponse, streaming::ResponsesStreamEvent,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
|
|
@ -34,7 +34,7 @@ impl ResponsesWebSocketProviderConfig for OpenAiResponsesApiConfig {
|
|||
|
||||
fn transform_ws_request(
|
||||
&self,
|
||||
event: &ResponsesWsEvent,
|
||||
event: &ResponsesStreamEvent,
|
||||
model: &str,
|
||||
) -> Result<ResponsesWsTransformResult, Error> {
|
||||
Ok(ResponsesWsTransformResult::passthrough(enforce_model(
|
||||
|
|
@ -44,7 +44,7 @@ impl ResponsesWebSocketProviderConfig for OpenAiResponsesApiConfig {
|
|||
|
||||
fn transform_ws_response(
|
||||
&self,
|
||||
event: &ResponsesWsEvent,
|
||||
event: &ResponsesStreamEvent,
|
||||
_model: &str,
|
||||
) -> Result<ResponsesWsTransformResult, Error> {
|
||||
Ok(ResponsesWsTransformResult::passthrough(event.clone()))
|
||||
|
|
@ -230,7 +230,7 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn openai_config_is_native_and_enforces_model() {
|
||||
let event: ResponsesWsEvent =
|
||||
let event: ResponsesStreamEvent =
|
||||
serde_json::from_value(serde_json::json!({"type":"response.create"}))
|
||||
.expect("valid event");
|
||||
let result = OPENAI_RESPONSES_WS_CONFIG
|
||||
|
|
|
|||
|
|
@ -170,6 +170,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
|
|||
"focus",
|
||||
"mavvrik",
|
||||
"vantage",
|
||||
"ternary",
|
||||
"posthog",
|
||||
"levo",
|
||||
"compression_interception",
|
||||
|
|
@ -440,7 +441,6 @@ autorouter_presets_url: str = os.getenv(
|
|||
"LITELLM_AUTOROUTER_PRESETS_URL",
|
||||
"https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/autorouter_presets.json",
|
||||
)
|
||||
whats_new_url: str | None = os.getenv("LITELLM_WHATS_NEW_URL")
|
||||
suppress_debug_info: bool = False
|
||||
dynamodb_table_name: Optional[str] = None
|
||||
s3_callback_params: Optional[Dict] = None
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ AGENT_CARD_WELL_KNOWN_PATH: str = "/.well-known/agent-card.json"
|
|||
PREV_AGENT_CARD_WELL_KNOWN_PATH: str = "/.well-known/agent.json"
|
||||
FOUNDRY_AGENT_CARD_PATH: Final = "/agentCard/v1.0"
|
||||
AGENT_CARD_PATH_PARAM: Final = "agent_card_path"
|
||||
A2A_PROTOCOL_VERSION_PARAM: Final = "a2a_protocol_version"
|
||||
|
||||
try:
|
||||
from a2a.client import A2ACardResolver as _A2ACardResolver
|
||||
|
|
@ -76,22 +77,43 @@ _CANONICAL_PROTOCOL_BINDINGS: Final = MappingProxyType(
|
|||
)
|
||||
|
||||
_LEGACY_PROTOCOL_VERSION: Final = "0.3"
|
||||
_A2A_PROTOCOL_VERSION_ALIASES: Final = MappingProxyType(
|
||||
{
|
||||
"0.3": "0.3",
|
||||
"0.3.0": "0.3",
|
||||
"1.0": "1.0",
|
||||
"1.0.0": "1.0",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def normalize_agent_card_interfaces(agent_card: "AgentCard") -> "AgentCard":
|
||||
def resolve_a2a_protocol_version(value: object) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool):
|
||||
verbose_logger.warning("Ignoring invalid %s value: %r", A2A_PROTOCOL_VERSION_PARAM, value)
|
||||
return None
|
||||
if not isinstance(value, (str, int, float)):
|
||||
verbose_logger.warning("Ignoring invalid %s value: %r", A2A_PROTOCOL_VERSION_PARAM, value)
|
||||
return None
|
||||
|
||||
protocol_version: Final = _A2A_PROTOCOL_VERSION_ALIASES.get(str(value).strip())
|
||||
if protocol_version is None:
|
||||
verbose_logger.warning("Ignoring invalid %s value: %r", A2A_PROTOCOL_VERSION_PARAM, value)
|
||||
return protocol_version
|
||||
|
||||
|
||||
def normalize_agent_card_interfaces(agent_card: "AgentCard", protocol_version: str | None = None) -> "AgentCard":
|
||||
"""
|
||||
Canonicalize the supported interfaces of spec-adjacent agent cards.
|
||||
|
||||
Some A2A servers (e.g. LangGraph Platform) serve agent cards with lowercase
|
||||
bindings like "jsonrpc", but a2a-sdk's ClientFactory matches bindings
|
||||
case-sensitively against its uppercase TransportProtocol constants and fails
|
||||
with "no compatible transports found." for spec-adjacent casings.
|
||||
|
||||
The same servers also speak the A2A 0.3 JSON dialect ("kind"-discriminated
|
||||
payloads) while declaring protocolVersion "1.0", which a2a-sdk's strict v1
|
||||
proto parsing rejects. A mis-cased binding fingerprints such a server, so its
|
||||
declared version is downgraded to 0.3 to route the SDK's ClientFactory onto
|
||||
its v0.3 compat transport, which speaks that dialect.
|
||||
with "no compatible transports found." for spec-adjacent casings. The casing
|
||||
fingerprint stopped working with langgraph-api 0.15, so operators can pin
|
||||
the version per agent via ``litellm_params.a2a_protocol_version``. An explicit
|
||||
"1.0" opts out of the mis-cased downgrade.
|
||||
"""
|
||||
normalized: Final = type(agent_card)()
|
||||
normalized.CopyFrom(agent_card)
|
||||
|
|
@ -101,6 +123,10 @@ def normalize_agent_card_interfaces(agent_card: "AgentCard") -> "AgentCard":
|
|||
continue
|
||||
interface.protocol_binding = canonical
|
||||
interface.protocol_version = _LEGACY_PROTOCOL_VERSION
|
||||
if protocol_version is not None:
|
||||
for interface in normalized.supported_interfaces:
|
||||
if interface.protocol_binding in _CANONICAL_PROTOCOL_BINDINGS.values():
|
||||
interface.protocol_version = protocol_version
|
||||
return normalized
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -23,6 +23,12 @@ if TYPE_CHECKING:
|
|||
from a2a.client import Client as A2AClientType
|
||||
|
||||
|
||||
_A2A_KIND_DIALECT_HINT: Final = (
|
||||
" (the agent replied in the A2A 0.3 dialect while its card declares 1.0; "
|
||||
'set litellm_params.a2a_protocol_version: "0.3" on this agent)'
|
||||
)
|
||||
|
||||
|
||||
try:
|
||||
from a2a.client import Client, ClientConfig, create_client
|
||||
|
||||
|
|
@ -152,7 +158,7 @@ def map_a2a_exception(
|
|||
|
||||
# Default: wrap in generic A2AError
|
||||
raise A2AError(
|
||||
message=error_str,
|
||||
message=error_str + _A2A_KIND_DIALECT_HINT if 'has no field named "kind"' in error_str else error_str,
|
||||
model=model,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ from typing import Any, Final
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.a2a_protocol.card_resolver import AGENT_CARD_PATH_PARAM
|
||||
from litellm.a2a_protocol.card_resolver import A2A_PROTOCOL_VERSION_PARAM, AGENT_CARD_PATH_PARAM
|
||||
from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
|
||||
A2ACompletionBridgeTransformation,
|
||||
A2AStreamingContext,
|
||||
|
|
@ -38,6 +38,7 @@ _AGENT_ONLY_PARAMS: Final = frozenset(
|
|||
"agent_id",
|
||||
"agent_card_params",
|
||||
AGENT_CARD_PATH_PARAM,
|
||||
A2A_PROTOCOL_VERSION_PARAM,
|
||||
A2A_USER_API_KEY_HASH_PARAM,
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -72,10 +72,12 @@ except ImportError:
|
|||
|
||||
# Import our custom card resolver that supports multiple well-known paths
|
||||
from litellm.a2a_protocol.card_resolver import (
|
||||
A2A_PROTOCOL_VERSION_PARAM,
|
||||
AGENT_CARD_PATH_PARAM,
|
||||
LiteLLMA2ACardResolver,
|
||||
get_agent_card_url,
|
||||
normalize_agent_card_interfaces,
|
||||
resolve_a2a_protocol_version,
|
||||
)
|
||||
from litellm.a2a_protocol.exception_mapping_utils import (
|
||||
handle_a2a_localhost_retry,
|
||||
|
|
@ -153,6 +155,10 @@ def _agent_card_path(litellm_params: Mapping[str, object]) -> str | None:
|
|||
return configured_path if isinstance(configured_path, str) and configured_path else None
|
||||
|
||||
|
||||
def _agent_protocol_version(litellm_params: Mapping[str, object]) -> str | None:
|
||||
return resolve_a2a_protocol_version(litellm_params.get(A2A_PROTOCOL_VERSION_PARAM))
|
||||
|
||||
|
||||
def _set_litellm_params_on_logging_obj(
|
||||
kwargs: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
|
|
@ -498,6 +504,7 @@ async def asend_message(
|
|||
base_url=api_base,
|
||||
extra_headers=extra_headers,
|
||||
relative_card_path=_agent_card_path(litellm_params),
|
||||
protocol_version=_agent_protocol_version(litellm_params),
|
||||
)
|
||||
|
||||
# Type assertion: a2a_client is guaranteed to be non-None here
|
||||
|
|
@ -634,7 +641,7 @@ async def asend_message_streaming(
|
|||
proxy_server_request: dict[str, object] | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
**kwargs: object,
|
||||
) -> AsyncIterator[Any]:
|
||||
) -> AsyncIterator["SendStreamingMessageResponse | Mapping[str, object]"]:
|
||||
"""
|
||||
Async: Send a streaming message to an A2A agent.
|
||||
|
||||
|
|
@ -723,6 +730,7 @@ async def asend_message_streaming(
|
|||
extra_headers=extra_headers,
|
||||
streaming=True,
|
||||
relative_card_path=_agent_card_path(litellm_params),
|
||||
protocol_version=_agent_protocol_version(litellm_params),
|
||||
)
|
||||
|
||||
assert a2a_client is not None
|
||||
|
|
@ -770,6 +778,7 @@ async def create_a2a_client(
|
|||
extra_headers: dict[str, str] | None = None,
|
||||
streaming: bool = False,
|
||||
relative_card_path: str | None = None,
|
||||
protocol_version: str | None = None,
|
||||
) -> "A2AClientType":
|
||||
"""
|
||||
Create an A2A client for the given agent URL.
|
||||
|
|
@ -783,6 +792,7 @@ async def create_a2a_client(
|
|||
extra_headers: Optional additional headers to include in requests
|
||||
relative_card_path: Optional card path relative to ``base_url`` (e.g. ``agentCard/v1.0`` for a
|
||||
Microsoft Foundry agent); when None the well-known paths are probed in order
|
||||
protocol_version: Optional per-agent protocol version override for supported interfaces
|
||||
|
||||
Returns:
|
||||
An initialized a2a.client.A2AClient instance
|
||||
|
|
@ -819,7 +829,8 @@ async def create_a2a_client(
|
|||
await resolver.get_agent_card(
|
||||
relative_card_path=relative_card_path,
|
||||
http_kwargs=_card_http_kwargs(extra_headers),
|
||||
)
|
||||
),
|
||||
protocol_version=protocol_version,
|
||||
)
|
||||
|
||||
a2a_client: Final = await create_client( # pyright: ignore[reportOptionalCall]
|
||||
|
|
|
|||
|
|
@ -117,7 +117,7 @@ class BedrockAgentCoreA2AHandler:
|
|||
response.raise_for_status()
|
||||
|
||||
# Check content type — AgentCore may return JSON instead of SSE
|
||||
content_type: Final = response.headers.get("content-type", "").lower()
|
||||
content_type: Final = ", ".join(response.headers.get_list("content-type")).lower()
|
||||
|
||||
if "application/json" in content_type:
|
||||
# Single JSON response fallback (not SSE)
|
||||
|
|
|
|||
|
|
@ -768,7 +768,7 @@ def _get_response_from_batch_job_output_file(
|
|||
|
||||
|
||||
def _batch_response_was_successful(
|
||||
batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
|
||||
batch_job_output_file: Mapping[str, object], custom_llm_provider: str = "openai"
|
||||
) -> bool:
|
||||
"""
|
||||
Check if the batch job response was successful
|
||||
|
|
|
|||
|
|
@ -104,6 +104,7 @@ S3_PARTITION_GRANULARITY_ENV_VAR: Final = "S3_PARTITION_GRANULARITY"
|
|||
MAX_FILE_LIST_LIMIT: Final = 10000
|
||||
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
|
||||
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
|
||||
GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS: Final = 10
|
||||
budget_reservation_disabled_info_emitted = False
|
||||
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
|
||||
DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512))
|
||||
|
|
@ -391,6 +392,11 @@ MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: Final = "litellm.model_group_alias_resolve
|
|||
MCP_PEEKED_BODY_SCOPE_KEY: Final = "litellm_mcp_peeked_body"
|
||||
MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS: Final = 5.0
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
|
||||
OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS: Final = 30
|
||||
OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS: Final = 5
|
||||
OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES: Final = frozenset(
|
||||
{"response.completed", "response.incomplete", "response.failed"}
|
||||
)
|
||||
REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged"
|
||||
|
||||
# SSL/TLS cipher configuration for faster handshakes
|
||||
|
|
|
|||
|
|
@ -1,14 +1,16 @@
|
|||
# What is this?
|
||||
## File for 'response_cost' calculation in Logging
|
||||
import logging
|
||||
import math
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from functools import lru_cache
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, cast
|
||||
|
||||
from httpx import Response
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
||||
import litellm
|
||||
|
|
@ -17,7 +19,9 @@ from litellm import verbose_logger
|
|||
from litellm.constants import (
|
||||
DEFAULT_MAX_LRU_CACHE_SIZE,
|
||||
DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND,
|
||||
OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.cost_discount import resolve_cost_discount
|
||||
from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, is_decisions_call_type
|
||||
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
|
||||
|
|
@ -110,7 +114,6 @@ from litellm.types.llms.openai import (
|
|||
OpenAIModerationResponse,
|
||||
OpenAIRealtimeStreamList,
|
||||
OpenAIRealtimeStreamResponseBaseObject,
|
||||
OpenAIRealtimeStreamSessionEvents,
|
||||
ResponseAPIUsage,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
|
|
@ -836,6 +839,22 @@ def _cost_map_entry_prices_anything(entry: Mapping[str, object]) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def _declares_rates(model_info: ModelInfo | None) -> bool:
|
||||
entry: Final = _raw_cost_map_entry(model_info["key"]) if model_info is not None else None
|
||||
return entry is not None and _cost_map_entry_prices_anything(entry)
|
||||
|
||||
|
||||
def _base_model_pricing_provider(base_model: str, custom_llm_provider: str | None) -> str | None:
|
||||
"""The provider to price an explicit base_model under, when the deployment's own provider has
|
||||
no rates for it (e.g. chatgpt/ subscription rows) but the base model's own row does"""
|
||||
if custom_llm_provider is None or _declares_rates(_cost_map_model_info(base_model, custom_llm_provider)):
|
||||
return None
|
||||
own_model_info: Final = _cost_map_model_info(base_model, None)
|
||||
if own_model_info is None or not _declares_rates(own_model_info):
|
||||
return None
|
||||
return own_model_info["litellm_provider"]
|
||||
|
||||
|
||||
def select_model_name_for_cost_calc(
|
||||
model: str | None,
|
||||
completion_response: object | None,
|
||||
|
|
@ -1497,10 +1516,21 @@ def completion_cost(
|
|||
|
||||
explicit_pricing: Final = custom_pricing is True or base_model is not None
|
||||
custom_priced_request_model: Final = model if custom_pricing is True else None
|
||||
base_model_provider: Final = (
|
||||
_base_model_pricing_provider(
|
||||
base_model,
|
||||
_get_provider_for_cost_calc(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None,
|
||||
),
|
||||
)
|
||||
if base_model is not None and custom_pricing is not True
|
||||
else None
|
||||
)
|
||||
selected_model: Final = select_model_name_for_cost_calc(
|
||||
model=model,
|
||||
completion_response=completion_response,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
custom_llm_provider=base_model_provider or custom_llm_provider,
|
||||
custom_pricing=custom_pricing,
|
||||
base_model=base_model,
|
||||
router_model_id=router_model_id,
|
||||
|
|
@ -1847,6 +1877,13 @@ def completion_cost(
|
|||
if litellm_logging_obj is not None:
|
||||
request_model_for_cost = litellm_logging_obj.model
|
||||
|
||||
pricing_provider: str | None = (
|
||||
base_model_provider
|
||||
if base_model_provider is not None and idx == 0
|
||||
else custom_llm_provider
|
||||
if isinstance(custom_llm_provider, str)
|
||||
else None
|
||||
)
|
||||
(
|
||||
prompt_tokens_cost_usd_dollar,
|
||||
completion_tokens_cost_usd_dollar,
|
||||
|
|
@ -1854,7 +1891,7 @@ def completion_cost(
|
|||
model=model,
|
||||
prompt_tokens=prompt_tokens or 0,
|
||||
completion_tokens=completion_tokens or 0,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
custom_llm_provider=pricing_provider,
|
||||
response_time_ms=response_time_ms,
|
||||
region_name=None if explicit_pricing else region_name,
|
||||
custom_cost_per_second=custom_cost_per_second,
|
||||
|
|
@ -1951,12 +1988,9 @@ def completion_cost(
|
|||
_cache_creation_cost: float | None = None
|
||||
_billed_token_rates: BilledTokenRates | None = None
|
||||
if cost_per_token_usage_object is not None and model:
|
||||
_breakdown_provider: str | None = (
|
||||
custom_llm_provider if isinstance(custom_llm_provider, str) else None
|
||||
)
|
||||
_token_type_breakdown = get_token_type_cost_breakdown(
|
||||
model=model,
|
||||
custom_llm_provider=_breakdown_provider,
|
||||
custom_llm_provider=pricing_provider,
|
||||
usage=cost_per_token_usage_object,
|
||||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
|
|
@ -2979,6 +3013,141 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
|
|||
_TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"
|
||||
|
||||
|
||||
class _LiveResponsePayload(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str | None = None
|
||||
model: str | None = None
|
||||
usage: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
class _LiveResponseEvent(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
type: str = ""
|
||||
response: _LiveResponsePayload | None = None
|
||||
|
||||
|
||||
class _LiveResponseEnvelope(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
type: str = ""
|
||||
event: _LiveResponseEvent | None = None
|
||||
|
||||
|
||||
def _live_response_payload(result: Mapping[str, object]) -> _LiveResponsePayload | None:
|
||||
try:
|
||||
envelope: Final = _LiveResponseEnvelope.model_validate(result)
|
||||
except ValidationError:
|
||||
return None
|
||||
if (
|
||||
envelope.type != "response.event"
|
||||
or envelope.event is None
|
||||
or envelope.event.type not in OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES
|
||||
or envelope.event.response is None
|
||||
or envelope.event.response.usage is None
|
||||
):
|
||||
return None
|
||||
return envelope.event.response
|
||||
|
||||
|
||||
def _unique_live_response_payloads(results: Sequence[Mapping[str, object]]) -> tuple[_LiveResponsePayload, ...]:
|
||||
response_payloads: Final = tuple(
|
||||
response for result in results if (response := _live_response_payload(result)) is not None
|
||||
)
|
||||
first_payload_by_id: Final = {
|
||||
response.id: response for response in reversed(response_payloads) if response.id is not None
|
||||
}
|
||||
return tuple(reversed(first_payload_by_id.values()))
|
||||
|
||||
|
||||
def _live_session_model(result: Mapping[str, object]) -> str | None:
|
||||
session: Final = as_str_mapping(result.get("session"))
|
||||
if session is None:
|
||||
return None
|
||||
model: Final = session.get("model")
|
||||
return model if isinstance(model, str) else None
|
||||
|
||||
|
||||
def _live_usage_seconds(result: Mapping[str, object]) -> float | None:
|
||||
usage: Final = as_str_mapping(result.get("usage"))
|
||||
if usage is None:
|
||||
return None
|
||||
seconds: Final = usage.get("seconds")
|
||||
if isinstance(seconds, (int, float)) and not isinstance(seconds, bool) and math.isfinite(seconds):
|
||||
return float(seconds)
|
||||
return None
|
||||
|
||||
|
||||
def _last_live_usage_seconds(results: Sequence[Mapping[str, object]], event_type: str) -> float | None:
|
||||
return next(
|
||||
(
|
||||
seconds
|
||||
for result in reversed(results)
|
||||
if result.get("type") == event_type and (seconds := _live_usage_seconds(result)) is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _live_session_wall_clock_seconds(litellm_logging_obj: LitellmLoggingObject | None) -> float:
|
||||
if litellm_logging_obj is None:
|
||||
return 0.0
|
||||
model_call_details: Final = litellm_logging_obj.model_call_details
|
||||
start_time: Final = model_call_details.get("start_time")
|
||||
end_time: Final = model_call_details.get("end_time")
|
||||
if isinstance(start_time, datetime) and isinstance(end_time, datetime):
|
||||
return max((end_time - start_time).total_seconds(), 0.0)
|
||||
return 0.0
|
||||
|
||||
|
||||
def _live_input_cost_per_second(model_name: str, custom_llm_provider: str) -> float | None:
|
||||
model_info: Final = _lookup_model_info_or_none(model_name, custom_llm_provider)
|
||||
if model_info is None:
|
||||
return None
|
||||
rate: Final = model_info.get("input_cost_per_second")
|
||||
if isinstance(rate, (int, float)) and not isinstance(rate, bool) and math.isfinite(rate):
|
||||
return float(rate)
|
||||
return None
|
||||
|
||||
|
||||
def _first_live_input_cost_per_second(
|
||||
potential_model_names: Sequence[str | None],
|
||||
custom_llm_provider: str,
|
||||
) -> float:
|
||||
return next(
|
||||
(
|
||||
rate
|
||||
for model_name in potential_model_names
|
||||
if model_name is not None
|
||||
and (rate := _live_input_cost_per_second(model_name, custom_llm_provider)) is not None
|
||||
),
|
||||
0.0,
|
||||
)
|
||||
|
||||
|
||||
def _live_response_token_costs(
|
||||
response: _LiveResponsePayload,
|
||||
custom_llm_provider: str,
|
||||
data_residency: str | None,
|
||||
) -> tuple[float, float]:
|
||||
if response.id is None or response.model is None or response.usage is None:
|
||||
return 0.0, 0.0
|
||||
try:
|
||||
usage: Final = get_usage_object(response.model_dump())
|
||||
except (ValidationError, TypeError, ValueError, KeyError, AttributeError):
|
||||
return 0.0, 0.0
|
||||
if usage is None:
|
||||
return 0.0, 0.0
|
||||
costs: Final = _candidate_realtime_token_costs(
|
||||
model_name=response.model,
|
||||
combined_usage_object=usage,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
data_residency=data_residency,
|
||||
)
|
||||
return (0.0, 0.0) if costs is None else costs
|
||||
|
||||
|
||||
def _candidate_realtime_token_costs(
|
||||
model_name: str,
|
||||
combined_usage_object: Usage,
|
||||
|
|
@ -3065,14 +3234,21 @@ def handle_realtime_stream_cost_calculation(
|
|||
base_pricing_model: the deployment's resolved base_model, tried ahead of the
|
||||
session-reported model but after custom rates
|
||||
"""
|
||||
received_model = None
|
||||
potential_model_names: Final = [custom_pricing_model, base_pricing_model]
|
||||
for result in results:
|
||||
if result["type"] == "session.created":
|
||||
received_model = cast(OpenAIRealtimeStreamSessionEvents, result)["session"].get("model", None)
|
||||
potential_model_names.append(received_model)
|
||||
|
||||
potential_model_names.append(litellm_model_name)
|
||||
received_model: Final = next(
|
||||
(
|
||||
model_name
|
||||
for result in results
|
||||
if result["type"] in ("session.created", "session.started")
|
||||
and (model_name := _live_session_model(result)) is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
potential_model_names: Final = (
|
||||
custom_pricing_model,
|
||||
base_pricing_model,
|
||||
received_model,
|
||||
litellm_model_name,
|
||||
)
|
||||
input_cost_per_token, output_cost_per_token = _first_priced_realtime_token_costs(
|
||||
potential_model_names=potential_model_names,
|
||||
combined_usage_object=combined_usage_object,
|
||||
|
|
@ -3089,15 +3265,56 @@ def handle_realtime_stream_cost_calculation(
|
|||
if any(r.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE for r in results)
|
||||
else 0.0
|
||||
)
|
||||
total_cost: Final = input_cost_per_token + output_cost_per_token + transcription_cost
|
||||
live_lifecycle_event_types: Final = frozenset({"session.started", "session.usage.updated", "session.closed"})
|
||||
live_session_started: Final = any(result["type"] == "session.started" for result in results)
|
||||
live_session_closed: Final = any(result["type"] == "session.closed" for result in results)
|
||||
live_event_seen: Final = any(result["type"] in live_lifecycle_event_types for result in results)
|
||||
closed_seconds: Final = _last_live_usage_seconds(results, "session.closed")
|
||||
reported_seconds: Final = _last_live_usage_seconds(results, "session.usage.updated")
|
||||
latest_reported_seconds: Final = 0.0 if reported_seconds is None else reported_seconds
|
||||
voice_seconds: Final = (
|
||||
closed_seconds
|
||||
if closed_seconds is not None
|
||||
else max(
|
||||
latest_reported_seconds,
|
||||
_live_session_wall_clock_seconds(litellm_logging_obj),
|
||||
)
|
||||
if live_session_started and not live_session_closed
|
||||
else latest_reported_seconds
|
||||
)
|
||||
live_model_names: Final = (
|
||||
(custom_pricing_model, base_pricing_model, received_model, litellm_model_name) if live_event_seen else ()
|
||||
)
|
||||
live_voice_cost: Final = voice_seconds * _first_live_input_cost_per_second(
|
||||
live_model_names,
|
||||
custom_llm_provider,
|
||||
)
|
||||
live_response_costs: Final = tuple(
|
||||
_live_response_token_costs(
|
||||
response=response,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
data_residency=data_residency,
|
||||
)
|
||||
for response in _unique_live_response_payloads(results)
|
||||
)
|
||||
live_response_input_cost: Final = sum(costs[0] for costs in live_response_costs)
|
||||
live_response_output_cost: Final = sum(costs[1] for costs in live_response_costs)
|
||||
input_cost: Final = input_cost_per_token + live_response_input_cost
|
||||
output_cost: Final = output_cost_per_token + live_response_output_cost
|
||||
total_cost: Final = input_cost + output_cost + transcription_cost + live_voice_cost
|
||||
additional_costs: Final = tuple(
|
||||
(cost_name, cost)
|
||||
for cost_name, cost in (("transcription_cost", transcription_cost), ("live_voice_cost", live_voice_cost))
|
||||
if cost > 0
|
||||
)
|
||||
|
||||
_store_cost_breakdown_in_logging_obj(
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
prompt_tokens_cost_usd_dollar=input_cost_per_token,
|
||||
completion_tokens_cost_usd_dollar=output_cost_per_token,
|
||||
prompt_tokens_cost_usd_dollar=input_cost,
|
||||
completion_tokens_cost_usd_dollar=output_cost,
|
||||
cost_for_built_in_tools_cost_usd_dollar=0.0,
|
||||
total_cost_usd_dollar=total_cost,
|
||||
additional_costs={"transcription_cost": transcription_cost} if transcription_cost > 0 else None,
|
||||
additional_costs=dict(additional_costs) if additional_costs else None,
|
||||
data_residency=data_residency,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -158,7 +158,7 @@ class DeepAgentsHandler(BaseHarnessHandler):
|
|||
agent, deps = self._require_agent()
|
||||
snapshot = await agent.aget_state(self._run_config(ctx, None))
|
||||
messages = (snapshot.values or MappingProxyType({})).get("messages") or ()
|
||||
converted: list[dict[str, Any]] = deps.convert_to_openai_messages( # mutable-ok: LangChain returns a list
|
||||
converted: list[dict[str, object]] = deps.convert_to_openai_messages( # mutable-ok: LangChain returns a list
|
||||
messages
|
||||
)
|
||||
return converted
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
from litellm.integrations.prompt_management_base import PromptManagementClient
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
cache_control_as_prompt_cache_breakpoint,
|
||||
with_prompt_cache_breakpoint,
|
||||
)
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
|
|
@ -72,6 +73,7 @@ OPENAI_API_HOST: Final = "api.openai.com"
|
|||
OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE")
|
||||
_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
|
||||
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||
_CONTENT_PART_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues
|
||||
|
||||
|
|
@ -90,6 +92,12 @@ def _validated_object_list(value: object) -> list[object] | None:
|
|||
return None
|
||||
|
||||
|
||||
def _part_with_cache_control_renamed(part: object) -> object:
|
||||
if not _accepts_prompt_cache_breakpoint(part):
|
||||
return part
|
||||
return cache_control_as_prompt_cache_breakpoint(_CONTENT_PART_ADAPTER.validate_python(part))
|
||||
|
||||
|
||||
def configured_injection_points(value: object) -> Sequence[CacheControlInjectionPoint]:
|
||||
if not isinstance(value, (list, tuple)):
|
||||
return ()
|
||||
|
|
@ -242,6 +250,10 @@ def _accepts_prompt_cache_breakpoint(block: object) -> bool:
|
|||
return isinstance(block, dict) and block.get("type") in OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES
|
||||
|
||||
|
||||
def _may_carry_prompt_cache_breakpoint(message: Mapping[str, object]) -> bool:
|
||||
return message.get("role") != "assistant"
|
||||
|
||||
|
||||
# Set by a caller whose message list is not the one that goes upstream -- today the
|
||||
# Responses API layer, whose `instructions` only becomes a system message further down.
|
||||
# Tells this hook to hand role-targeted points to the pass holding the final messages
|
||||
|
|
@ -620,9 +632,32 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
message_content[-1]["cache_control"] = control # pyright: ignore[reportGeneralTypeIssues] # loose runtime dict
|
||||
return message
|
||||
|
||||
@staticmethod
|
||||
def cache_control_as_prompt_cache_breakpoints(message: AllMessageValues) -> AllMessageValues:
|
||||
"""Every Anthropic ``cache_control`` on the message or its parts moved to OpenAI's ``prompt_cache_breakpoint``,
|
||||
placed where the injection hook would put one. Marks with no OpenAI slot are left for the caller to strip."""
|
||||
fields: Final[Mapping[str, object]] = dict(message)
|
||||
if not _may_carry_prompt_cache_breakpoint(fields):
|
||||
return message
|
||||
parts: Final = _validated_object_list(fields.get("content"))
|
||||
renamed: Final = (
|
||||
message
|
||||
if parts is None
|
||||
else cast( # cast-ok: the same message with its part marks renamed
|
||||
AllMessageValues,
|
||||
{
|
||||
**fields,
|
||||
"content": [_part_with_cache_control_renamed(part) for part in parts],
|
||||
},
|
||||
)
|
||||
)
|
||||
if fields.get("cache_control") is None:
|
||||
return renamed
|
||||
return AnthropicCacheControlHook._insert_prompt_cache_breakpoint_in_message(renamed)
|
||||
|
||||
@staticmethod
|
||||
def _insert_prompt_cache_breakpoint_in_message(message: AllMessageValues) -> AllMessageValues:
|
||||
if message.get("role") == "assistant":
|
||||
if not _may_carry_prompt_cache_breakpoint(message):
|
||||
return message
|
||||
message_content: Final = message.get("content", None)
|
||||
if isinstance(message_content, str):
|
||||
|
|
|
|||
|
|
@ -50,7 +50,7 @@ class ArgillaLogger(CustomBatchLogger):
|
|||
if litellm.argilla_transformation_object is None:
|
||||
raise Exception("'litellm.argilla_transformation_object' is required, to log your payload to Argilla.")
|
||||
self.validate_argilla_transformation_object(litellm.argilla_transformation_object)
|
||||
self.argilla_transformation_object = litellm.argilla_transformation_object
|
||||
self.argilla_transformation_object: Mapping[str, str] = litellm.argilla_transformation_object
|
||||
self.default_credentials = self.get_credentials_from_env(
|
||||
argilla_api_key=argilla_api_key,
|
||||
argilla_dataset_name=argilla_dataset_name,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ Arize Phoenix API client for fetching prompt versions from Arize Phoenix.
|
|||
"""
|
||||
|
||||
import urllib.parse
|
||||
from typing import Any, Final
|
||||
from typing import Final
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
|
|
@ -53,7 +53,7 @@ class ArizePhoenixClient:
|
|||
# Initialize HTTPHandler
|
||||
self.http_handler = HTTPHandler(disable_default_headers=True)
|
||||
|
||||
def get_prompt_version(self, prompt_version_id: str) -> dict[str, Any] | None:
|
||||
def get_prompt_version(self, prompt_version_id: str) -> dict[str, object] | None:
|
||||
"""
|
||||
Fetch a prompt version from Arize Phoenix.
|
||||
|
||||
|
|
|
|||
|
|
@ -80,7 +80,7 @@ def _json(value: object) -> str:
|
|||
return value if isinstance(value, str) else json.dumps(value, default=str)
|
||||
|
||||
|
||||
def _json_mapping(value: Mapping[str, Any]) -> str:
|
||||
def _json_mapping(value: Mapping[str, object]) -> str:
|
||||
return _json(dict(value))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -209,7 +209,7 @@ class CloudZeroStreamer:
|
|||
|
||||
return payload
|
||||
|
||||
def _convert_cbf_to_api_format(self, row: dict[str, object]) -> dict[str, Any] | None:
|
||||
def _convert_cbf_to_api_format(self, row: dict[str, object]) -> dict[str, object] | None:
|
||||
"""Convert CBF row to CloudZero API format - keeping CBF field names as CloudZero expects them."""
|
||||
try:
|
||||
# CloudZero expects CBF format field names directly, not converted names
|
||||
|
|
|
|||
|
|
@ -511,7 +511,7 @@ class CustomGuardrail(CustomLogger):
|
|||
sticky_session_routing=self.sticky_session_routing,
|
||||
)
|
||||
|
||||
def _get_session_id_from_request_data(self, request_data: dict[str, Any]) -> str | None:
|
||||
def _get_session_id_from_request_data(self, request_data: dict[str, object]) -> str | None:
|
||||
"""Extract session_id from request data."""
|
||||
return get_session_id_from_request_data(request_data)
|
||||
|
||||
|
|
|
|||
|
|
@ -70,7 +70,7 @@ class PromptManager:
|
|||
self,
|
||||
prompt_id: str | None = None,
|
||||
prompt_directory: str | None = None,
|
||||
prompt_data: dict[str, dict[str, Any]] | None = None,
|
||||
prompt_data: dict[str, dict[str, object]] | None = None,
|
||||
prompt_file: str | None = None,
|
||||
):
|
||||
self.prompt_directory = Path(prompt_directory) if prompt_directory else None
|
||||
|
|
@ -372,6 +372,6 @@ class PromptManager:
|
|||
}
|
||||
return result
|
||||
|
||||
def load_prompts_from_json_data(self, prompt_data: dict[str, dict[str, Any]]) -> None:
|
||||
def load_prompts_from_json_data(self, prompt_data: dict[str, dict[str, object]]) -> None:
|
||||
"""Load additional prompts from JSON data (merges with existing prompts)."""
|
||||
self._load_prompts_from_json(prompt_data)
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from .factory import FocusDestinationFactory
|
|||
from .gcs_destination import FocusGCSDestination
|
||||
from .mavvrik_destination import FocusMavvrikDestination
|
||||
from .s3_destination import FocusS3Destination
|
||||
from .ternary_destination import FocusTernaryDestination
|
||||
from .vantage_destination import FocusVantageDestination
|
||||
|
||||
__all__ = [
|
||||
|
|
@ -13,6 +14,7 @@ __all__ = [
|
|||
"FocusGCSDestination",
|
||||
"FocusMavvrikDestination",
|
||||
"FocusS3Destination",
|
||||
"FocusTernaryDestination",
|
||||
"FocusTimeWindow",
|
||||
"FocusVantageDestination",
|
||||
]
|
||||
|
|
|
|||
152
litellm/integrations/focus/destinations/ternary_destination.py
Normal file
152
litellm/integrations/focus/destinations/ternary_destination.py
Normal file
|
|
@ -0,0 +1,152 @@
|
|||
"""Ternary API destination for FOCUS export: uploads each day's FOCUS CSV via a Ternary-signed GCS URL."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
from urllib.parse import quote, urlparse
|
||||
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
from .base import FocusDestination, FocusTimeWindow
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import httpx
|
||||
|
||||
TERNARY_UPLOAD_TIMEOUT_SECONDS: Final = 120.0
|
||||
|
||||
_LOOPBACK_HOSTS: Final = frozenset({"localhost", "127.0.0.1", "::1"})
|
||||
|
||||
|
||||
class TernaryUploadUrlBody(TypedDict):
|
||||
url: ReadOnly[NotRequired[str]]
|
||||
method: ReadOnly[NotRequired[str]]
|
||||
headers: ReadOnly[NotRequired[Mapping[str, str]]]
|
||||
|
||||
|
||||
class _UploadUrlResponse(Protocol):
|
||||
def json(self) -> TernaryUploadUrlBody: ...
|
||||
|
||||
|
||||
def _upload_url_body(response: _UploadUrlResponse) -> TernaryUploadUrlBody:
|
||||
return response.json()
|
||||
|
||||
|
||||
def _require_secure_base_url(base_url: str) -> None:
|
||||
parsed: Final = urlparse(base_url)
|
||||
if parsed.scheme == "https":
|
||||
return
|
||||
if parsed.scheme == "http" and (parsed.hostname or "").lower() in _LOOPBACK_HOSTS:
|
||||
return
|
||||
raise ValueError(f"base_url must be an HTTPS URL (got {base_url!r}); http is allowed only for loopback")
|
||||
|
||||
|
||||
def _require_gcs_url(url: str) -> None:
|
||||
parsed: Final = urlparse(url)
|
||||
hostname: Final = (parsed.hostname or "").lower()
|
||||
is_gcs: Final = hostname == "storage.googleapis.com" or hostname.endswith(".storage.googleapis.com")
|
||||
if parsed.scheme != "https" or not is_gcs:
|
||||
raise ValueError(f"Ternary destination: upload URL must be an HTTPS GCS URL (got host {hostname!r})")
|
||||
|
||||
|
||||
class _NoRedirectHTTPHandler(AsyncHTTPHandler):
|
||||
"""AsyncHTTPHandler that disables redirects on every client it creates (SSRF guard)."""
|
||||
|
||||
def create_client(
|
||||
self,
|
||||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: passthrough to AsyncHTTPHandler.create_client
|
||||
) -> httpx.AsyncClient:
|
||||
client: Final = super().create_client(*args, **kwargs) # pyright: ignore[reportArgumentType] # passthrough
|
||||
client.follow_redirects = False
|
||||
return client
|
||||
|
||||
|
||||
class FocusTernaryDestination(FocusDestination):
|
||||
"""Upload one UTC day of FOCUS CSV per delivery to Ternary through a signed GCS URL."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
prefix: str,
|
||||
config: dict[str, str] | None = None, # mutable-ok: FocusDestination(config) factory contract
|
||||
) -> None:
|
||||
resolved_config: Final = config or {}
|
||||
api_key: Final = resolved_config.get("api_key")
|
||||
connection_id: Final = resolved_config.get("connection_id")
|
||||
base_url: Final = resolved_config.get("base_url")
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"api_key must be provided for Ternary destination "
|
||||
"(set TERNARY_API_KEY env var or pass in destination_config)"
|
||||
)
|
||||
if not connection_id:
|
||||
raise ValueError(
|
||||
"connection_id must be provided for Ternary destination "
|
||||
"(set TERNARY_CONNECTION_ID env var or pass in destination_config)"
|
||||
)
|
||||
if "/" in connection_id or ".." in connection_id or any(c.isspace() for c in connection_id):
|
||||
raise ValueError(f"connection_id must not contain '/', '..', or whitespace (got {connection_id!r})")
|
||||
if not base_url:
|
||||
raise ValueError(
|
||||
"base_url must be provided for Ternary destination "
|
||||
"(set TERNARY_BASE_URL env var or pass in destination_config)"
|
||||
)
|
||||
_require_secure_base_url(str(base_url))
|
||||
self.api_key = api_key
|
||||
self.connection_id = connection_id
|
||||
self.base_url = str(base_url).rstrip("/")
|
||||
self.prefix = prefix
|
||||
|
||||
async def deliver(
|
||||
self,
|
||||
*,
|
||||
content: bytes,
|
||||
time_window: FocusTimeWindow,
|
||||
filename: str,
|
||||
) -> None:
|
||||
"""Upload one UTC day: request a signed URL for the window's start day, then PUT the CSV to it."""
|
||||
if not content:
|
||||
verbose_logger.debug("Ternary destination: empty content, skipping upload")
|
||||
return
|
||||
|
||||
day: Final = time_window.start_time.strftime("%Y-%m-%d")
|
||||
handler: Final = _NoRedirectHTTPHandler(timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS)
|
||||
try:
|
||||
body: Final = await self._request_upload_url(handler.client, day)
|
||||
url: Final = body.get("url")
|
||||
headers: Final = body.get("headers")
|
||||
if not url or headers is None:
|
||||
raise RuntimeError(f"Ternary destination: upload-url response for {day} is missing 'url' or 'headers'")
|
||||
if body.get("method") != "PUT":
|
||||
raise RuntimeError(f"Ternary destination: unsupported upload method {body.get('method')!r} for {day}")
|
||||
_require_gcs_url(url)
|
||||
uploaded: Final = await handler.client.put(
|
||||
url, content=content, headers=headers, timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS
|
||||
)
|
||||
if uploaded.status_code >= 400:
|
||||
raise RuntimeError(
|
||||
f"Ternary destination: upload for {day} failed ({uploaded.status_code}): {uploaded.text[:200]}"
|
||||
)
|
||||
finally:
|
||||
await handler.close()
|
||||
|
||||
verbose_logger.debug("Ternary destination: uploaded %d bytes for %s (%s)", len(content), day, filename)
|
||||
|
||||
async def _request_upload_url(self, client: httpx.AsyncClient, day: str) -> TernaryUploadUrlBody:
|
||||
endpoint: Final = f"{self.base_url}/external-cost-sources/v1/{quote(self.connection_id, safe='')}/upload-url"
|
||||
response: Final = await client.post(
|
||||
endpoint,
|
||||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||||
json={"day": day},
|
||||
timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS,
|
||||
)
|
||||
if response.status_code >= 400:
|
||||
raise RuntimeError(
|
||||
f"Ternary destination: upload-url request for {day} failed ({response.status_code}): "
|
||||
f"{response.text[:200]}"
|
||||
)
|
||||
return _upload_url_body(response)
|
||||
|
|
@ -115,7 +115,7 @@ class FocusLogger(CustomLogger):
|
|||
|
||||
pod_lock_manager: PodLockManager | None = None
|
||||
if proxy_logging_obj is not None:
|
||||
writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None)
|
||||
writer: Final[object] = getattr(proxy_logging_obj, "db_spend_update_writer", None)
|
||||
if writer is not None:
|
||||
pod_lock_manager = getattr(writer, "pod_lock_manager", None)
|
||||
|
||||
|
|
|
|||
|
|
@ -335,6 +335,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
|
|||
"LiteLLM_PolicyTable",
|
||||
"LiteLLM_PolicyAttachmentTable",
|
||||
"LiteLLM_ToolTable",
|
||||
"LiteLLM_TelemetryReport",
|
||||
"LiteLLM_AccessGroupTable",
|
||||
"LiteLLM_ClaudeCodePluginTable",
|
||||
"LiteLLM_MemoryTable",
|
||||
|
|
|
|||
0
litellm/integrations/ternary/__init__.py
Normal file
0
litellm/integrations/ternary/__init__.py
Normal file
247
litellm/integrations/ternary/export_engine.py
Normal file
247
litellm/integrations/ternary/export_engine.py
Normal file
|
|
@ -0,0 +1,247 @@
|
|||
"""Ternary FOCUS export engine (polars-backed): one UTC day per upload, with token counts merged into Tags."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from functools import partial
|
||||
from typing import Final, Literal, Protocol
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.focus.database import FocusLiteLLMDatabase
|
||||
from litellm.integrations.focus.destinations.base import FocusTimeWindow
|
||||
from litellm.integrations.focus.destinations.ternary_destination import FocusTernaryDestination
|
||||
from litellm.integrations.focus.export_engine import FocusExportEngine
|
||||
from litellm.integrations.focus.transformer import FocusTransformer
|
||||
|
||||
_TOKEN_TAG_KEYS: Final = (
|
||||
"prompt_tokens",
|
||||
"completion_tokens",
|
||||
"cache_read_input_tokens",
|
||||
"cache_creation_input_tokens",
|
||||
)
|
||||
_SESSION_TIMEZONE_MARGIN: Final = timedelta(days=1)
|
||||
_TAGS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
class _Series(Protocol):
|
||||
def to_list(self) -> Sequence[object]: ...
|
||||
|
||||
def alias(self, name: str) -> _Series: ...
|
||||
|
||||
|
||||
class _Frame(Protocol):
|
||||
"""The polars DataFrame surface this module uses."""
|
||||
|
||||
@property
|
||||
def columns(self) -> Sequence[str]: ...
|
||||
|
||||
@property
|
||||
def height(self) -> int: ...
|
||||
|
||||
def is_empty(self) -> bool: ...
|
||||
|
||||
def get_column(self, name: str) -> _Series: ...
|
||||
|
||||
def select(self, columns: Sequence[str], /) -> _Frame: ...
|
||||
|
||||
def hstack(self, columns: _Frame, /) -> _Frame: ...
|
||||
|
||||
def filter(self, mask: Sequence[bool], /) -> _Frame: ...
|
||||
|
||||
def map_rows(self, function: Callable[[tuple[object, ...]], str], /) -> _Frame: ...
|
||||
|
||||
def to_series(self) -> _Series: ...
|
||||
|
||||
def with_columns(self, column: _Series, /) -> _Frame: ...
|
||||
|
||||
def partition_by(
|
||||
self, by: str, /, *, maintain_order: bool, as_dict: Literal[True]
|
||||
) -> Mapping[tuple[object, ...], _Frame]: ...
|
||||
|
||||
|
||||
class _UsageSource(Protocol):
|
||||
async def get_usage_data(self, *, limit: int | None, start_time_utc: datetime | None = None) -> _Frame: ...
|
||||
|
||||
|
||||
class _Transformer(Protocol):
|
||||
def transform(self, frame: _Frame) -> _Frame: ...
|
||||
|
||||
|
||||
class _Serializer(Protocol):
|
||||
def serialize(self, frame: _Frame) -> bytes: ...
|
||||
|
||||
|
||||
def _existing_tags(raw: object) -> Mapping[str, object]:
|
||||
if not isinstance(raw, str) or not raw:
|
||||
return {}
|
||||
try:
|
||||
return _TAGS.validate_json(raw)
|
||||
except ValidationError:
|
||||
return {}
|
||||
|
||||
|
||||
def _tags_with_tokens(token_cols: Sequence[str], row: tuple[object, ...]) -> str:
|
||||
"""Row is (Tags, *token values in `token_cols` order); a missing or non-object Tags starts from empty."""
|
||||
tokens: Final = {key: str(value) for key, value in zip(token_cols, row[1:]) if value is not None}
|
||||
return json.dumps({**_existing_tags(row[0]), **tokens})
|
||||
|
||||
|
||||
def _merge_token_tags(normalized: _Frame, source: _Frame) -> _Frame:
|
||||
"""Zip transformed rows with the row-aligned source and add token keys to each row's Tags; degrades to unchanged on any mismatch."""
|
||||
if "Tags" not in normalized.columns:
|
||||
verbose_logger.warning(
|
||||
"Ternary export: transformed data has no `Tags` column; token counts not merged into %d row(s)",
|
||||
normalized.height,
|
||||
)
|
||||
return normalized
|
||||
if normalized.height != source.height:
|
||||
verbose_logger.warning(
|
||||
"Ternary export: row-count mismatch (normalized=%d, source=%d); token counts not merged",
|
||||
normalized.height,
|
||||
source.height,
|
||||
)
|
||||
return normalized
|
||||
token_cols: Final = tuple(c for c in _TOKEN_TAG_KEYS if c in source.columns)
|
||||
if normalized.is_empty() or not token_cols:
|
||||
return normalized
|
||||
merged: Final = (
|
||||
normalized.select(("Tags",)).hstack(source.select(token_cols)).map_rows(partial(_tags_with_tokens, token_cols))
|
||||
)
|
||||
return normalized.with_columns(merged.to_series().alias("Tags"))
|
||||
|
||||
|
||||
def _canonical_day(raw: object) -> str | None:
|
||||
"""Return the zero-padded YYYY-MM-DD for a `date` value, or None when it is missing or unparseable."""
|
||||
if isinstance(raw, date):
|
||||
return raw.isoformat()
|
||||
if not isinstance(raw, str):
|
||||
return None
|
||||
try:
|
||||
parsed: Final = datetime.strptime(raw, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||||
except ValueError:
|
||||
return None
|
||||
return parsed.date().isoformat()
|
||||
|
||||
|
||||
def _keep_window_days(data: _Frame, window: FocusTimeWindow) -> _Frame:
|
||||
"""Keep source rows whose UTC `date` falls on a day the window touches; each kept day is then uploaded whole."""
|
||||
if "date" not in data.columns:
|
||||
return data
|
||||
first_day: Final = window.start_time.astimezone(timezone.utc).date().isoformat()
|
||||
last_day: Final = (window.end_time.astimezone(timezone.utc) - timedelta(microseconds=1)).date().isoformat()
|
||||
kept: Final = data.filter(
|
||||
[
|
||||
day is not None and first_day <= day <= last_day
|
||||
for day in map(_canonical_day, data.get_column("date").to_list())
|
||||
]
|
||||
)
|
||||
if data.height > 0 and kept.height == 0:
|
||||
verbose_logger.warning(
|
||||
"Ternary export: day window %s..%s dropped all %d rows (unparseable or out-of-window `date`?)",
|
||||
first_day,
|
||||
last_day,
|
||||
data.height,
|
||||
)
|
||||
return kept
|
||||
|
||||
|
||||
def _partition_by_day(data: _Frame) -> Sequence[tuple[str, _Frame]]:
|
||||
"""Split source rows by their UTC `date` (one upload per day); rows with no parseable date are dropped."""
|
||||
if "date" not in data.columns:
|
||||
verbose_logger.warning("Ternary export: source data has no `date` column; %d row(s) not exported", data.height)
|
||||
return ()
|
||||
parts: Final = tuple(
|
||||
(_canonical_day(key[0]), part)
|
||||
for key, part in data.partition_by("date", maintain_order=True, as_dict=True).items()
|
||||
)
|
||||
undated: Final = sum(part.height for day, part in parts if day is None)
|
||||
if undated:
|
||||
verbose_logger.warning("Ternary export: dropped %d row(s) with a missing or unparseable `date`", undated)
|
||||
return tuple(sorted(((day, part) for day, part in parts if day is not None), key=lambda entry: entry[0]))
|
||||
|
||||
|
||||
def _day_window(day: str, frequency: str) -> FocusTimeWindow:
|
||||
"""The single UTC day a per-day upload covers."""
|
||||
start: Final = datetime.strptime(day, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||||
return FocusTimeWindow(start_time=start, end_time=start + timedelta(days=1), frequency=frequency)
|
||||
|
||||
|
||||
async def _read_usage(database: _UsageSource, *, limit: int | None, window: FocusTimeWindow | None) -> _Frame:
|
||||
"""Read every row, or the rows dated on the days `window` touches.
|
||||
|
||||
`updated_at` is naive UTC filtered as `timestamptz` in the database session timezone, so the read starts one
|
||||
day before the window with no upper bound, and each row's own UTC `date` decides what is kept.
|
||||
"""
|
||||
if window is None:
|
||||
return await database.get_usage_data(limit=limit)
|
||||
data: Final = await database.get_usage_data(
|
||||
limit=limit, start_time_utc=window.start_time - _SESSION_TIMEZONE_MARGIN
|
||||
)
|
||||
return _keep_window_days(data, window)
|
||||
|
||||
|
||||
def _day_payload(data: _Frame, *, transformer: _Transformer, serializer: _Serializer) -> bytes:
|
||||
transformed: Final = transformer.transform(data)
|
||||
if transformed.is_empty():
|
||||
return b""
|
||||
return serializer.serialize(_merge_token_tags(transformed, data))
|
||||
|
||||
|
||||
class TernaryExportEngine(FocusExportEngine):
|
||||
"""FocusExportEngine that uploads one whole UTC day per delivery, with token counts merged into Tags."""
|
||||
|
||||
def __init__(self, *, prefix: str, destination: FocusTernaryDestination) -> None:
|
||||
"""Set the base engine's fields directly, since the generic destination factory has no Ternary destination."""
|
||||
self.provider = "ternary"
|
||||
self.export_format = "csv"
|
||||
self.prefix = prefix
|
||||
self.destination = destination
|
||||
self.serializer = self._init_serializer()
|
||||
self.transformer = FocusTransformer()
|
||||
self.database = FocusLiteLLMDatabase()
|
||||
|
||||
async def export_window(self, *, window: FocusTimeWindow, limit: int | None) -> None:
|
||||
data: Final = await _read_usage(self._database, limit=limit, window=window)
|
||||
await self._deliver_enriched(data, window)
|
||||
|
||||
async def export_all(self, *, limit: int | None) -> None:
|
||||
data: Final = await _read_usage(self._database, limit=limit, window=None)
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
window: Final = FocusTimeWindow(
|
||||
start_time=now.replace(hour=0, minute=0, second=0, microsecond=0),
|
||||
end_time=now,
|
||||
frequency="all",
|
||||
)
|
||||
await self._deliver_enriched(data, window)
|
||||
|
||||
async def _deliver_enriched(self, data: _Frame, window: FocusTimeWindow) -> None:
|
||||
if data.is_empty():
|
||||
verbose_logger.debug("Ternary export: no usage data for window %s", window)
|
||||
return
|
||||
failed: Final[list[str]] = [] # mutable-ok: days whose upload failed, reported once after all days run
|
||||
for day, day_data in _partition_by_day(data):
|
||||
try:
|
||||
await self._deliver_day(day_data, _day_window(day, window.frequency))
|
||||
except Exception as e: # noqa: BLE001 # one day's failure must not block the rest; logged and re-raised
|
||||
verbose_logger.warning("Ternary export: upload for %s failed: %s", day, e)
|
||||
failed.append(day)
|
||||
if failed:
|
||||
raise RuntimeError(f"Ternary export: {len(failed)} day(s) failed to upload: {', '.join(failed)}")
|
||||
|
||||
async def _deliver_day(self, data: _Frame, window: FocusTimeWindow) -> None:
|
||||
payload: Final = _day_payload(data, transformer=self._transformer, serializer=self._serializer)
|
||||
if not payload:
|
||||
verbose_logger.debug("Ternary export: nothing to upload for %s", window.start_time.date())
|
||||
return
|
||||
await self._destination.deliver(
|
||||
content=payload,
|
||||
time_window=window,
|
||||
filename=self._build_filename(window),
|
||||
)
|
||||
|
||||
|
||||
__all__ = ("TernaryExportEngine",)
|
||||
167
litellm/integrations/ternary/ternary_logger.py
Normal file
167
litellm/integrations/ternary/ternary_logger.py
Normal file
|
|
@ -0,0 +1,167 @@
|
|||
"""Ternary FocusLogger preset (CSV, Ternary destination) that adds per-row token counts to the FOCUS Tags."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Awaitable, Callable
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.focus.destinations.base import FocusTimeWindow
|
||||
from litellm.integrations.focus.destinations.ternary_destination import FocusTernaryDestination
|
||||
from litellm.integrations.focus.focus_logger import FocusLogger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.focus.export_engine import FocusExportEngine
|
||||
|
||||
TERNARY_USAGE_DATA_JOB_NAME: Final = "ternary_export_usage_data"
|
||||
|
||||
_SUPPORTED_FREQUENCIES: Final = frozenset({"daily", "interval"})
|
||||
|
||||
|
||||
class _JobScheduler(Protocol):
|
||||
"""The slice of ``AsyncIOScheduler`` this module uses, which ships no type information"""
|
||||
|
||||
def add_job(
|
||||
self,
|
||||
func: Callable[[], Awaitable[None]],
|
||||
/,
|
||||
**trigger: str | int, # kwargs-ok: APScheduler add_job trigger options
|
||||
) -> object: ...
|
||||
|
||||
|
||||
def _parse_interval(raw: str | int | None) -> int | None:
|
||||
"""Parse the export interval-seconds override; a non-numeric value is ignored."""
|
||||
if raw is None:
|
||||
return None
|
||||
try:
|
||||
return int(raw)
|
||||
except (ValueError, TypeError):
|
||||
verbose_logger.warning("Invalid TERNARY_EXPORT_INTERVAL_SECONDS value: %s, ignoring", raw)
|
||||
return None
|
||||
|
||||
|
||||
class TernaryLogger(FocusLogger):
|
||||
"""FocusLogger preset for Ternary (CSV, Ternary cost-import API).
|
||||
|
||||
Env: TERNARY_API_KEY, TERNARY_CONNECTION_ID, TERNARY_BASE_URL (required),
|
||||
TERNARY_EXPORT_FREQUENCY ("daily" default / "interval"), TERNARY_EXPORT_INTERVAL_SECONDS.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
api_key: str | None = None,
|
||||
connection_id: str | None = None,
|
||||
base_url: str | None = None,
|
||||
frequency: str | None = None,
|
||||
interval_seconds: int | None = None,
|
||||
) -> None:
|
||||
resolved_api_key: Final = api_key or os.getenv("TERNARY_API_KEY")
|
||||
resolved_connection_id: Final = connection_id or os.getenv("TERNARY_CONNECTION_ID")
|
||||
resolved_base_url: Final = base_url or os.getenv("TERNARY_BASE_URL")
|
||||
resolved_frequency: Final = (frequency or os.getenv("TERNARY_EXPORT_FREQUENCY") or "daily").lower()
|
||||
if resolved_frequency not in _SUPPORTED_FREQUENCIES:
|
||||
raise ValueError(
|
||||
f"Unsupported TERNARY_EXPORT_FREQUENCY {resolved_frequency!r}; "
|
||||
f"Ternary supports {sorted(_SUPPORTED_FREQUENCIES)}. LiteLLM spend is a daily "
|
||||
"aggregate landed with whole-day replace, so 'hourly' adds no grain -- use "
|
||||
"'interval' only for short setup-validation loops."
|
||||
)
|
||||
|
||||
resolved_interval: Final = _parse_interval(
|
||||
interval_seconds if interval_seconds is not None else os.getenv("TERNARY_EXPORT_INTERVAL_SECONDS")
|
||||
)
|
||||
if resolved_frequency == "interval" and (resolved_interval is None or resolved_interval <= 0):
|
||||
raise ValueError(
|
||||
"TERNARY_EXPORT_INTERVAL_SECONDS must be a positive integer when TERNARY_EXPORT_FREQUENCY is 'interval'"
|
||||
)
|
||||
|
||||
destination_config: Final[dict[str, str]] = {} # mutable-ok: built from the config values present below
|
||||
if resolved_api_key:
|
||||
destination_config["api_key"] = resolved_api_key
|
||||
if resolved_connection_id:
|
||||
destination_config["connection_id"] = resolved_connection_id
|
||||
if resolved_base_url:
|
||||
destination_config["base_url"] = resolved_base_url
|
||||
|
||||
super().__init__(
|
||||
provider="ternary",
|
||||
export_format="csv",
|
||||
frequency=resolved_frequency,
|
||||
interval_seconds=resolved_interval,
|
||||
prefix="ternary_exports",
|
||||
destination_config=destination_config,
|
||||
)
|
||||
|
||||
verbose_logger.debug(
|
||||
"TernaryLogger initialized (connection_id=%s)",
|
||||
(
|
||||
resolved_connection_id[:4] + "***"
|
||||
if resolved_connection_id and len(resolved_connection_id) > 4
|
||||
else "***"
|
||||
),
|
||||
)
|
||||
|
||||
def _compute_time_window(self, now: datetime) -> FocusTimeWindow:
|
||||
"""Start at UTC midnight two days back so each push re-sends whole days, including spend flushed late."""
|
||||
now_utc: Final = now.astimezone(timezone.utc)
|
||||
start_time: Final = (now_utc - timedelta(days=2)).replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
return FocusTimeWindow(start_time=start_time, end_time=now_utc, frequency=self.frequency)
|
||||
|
||||
def _ensure_engine(self) -> FocusExportEngine:
|
||||
"""Instantiate the polars-backed Ternary export engine lazily."""
|
||||
if self._engine is None:
|
||||
from .export_engine import TernaryExportEngine # noqa: PLC0415 # polars is a proxy extra
|
||||
|
||||
self._engine = TernaryExportEngine(
|
||||
prefix=self.prefix,
|
||||
destination=FocusTernaryDestination(prefix=self.prefix, config=self._destination_config),
|
||||
)
|
||||
return self._engine
|
||||
|
||||
async def initialize_focus_export_job(self) -> None:
|
||||
"""Run the scheduled export under a Ternary-specific pod lock (distinct from FocusLogger's)."""
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None) if proxy_logging_obj else None
|
||||
pod_lock_manager: Final = getattr(writer, "pod_lock_manager", None) if writer is not None else None
|
||||
|
||||
if pod_lock_manager and pod_lock_manager.redis_cache:
|
||||
acquired: Final = await pod_lock_manager.acquire_lock(cronjob_id=TERNARY_USAGE_DATA_JOB_NAME)
|
||||
if not acquired:
|
||||
verbose_logger.debug("Ternary export: unable to acquire pod lock")
|
||||
return
|
||||
try:
|
||||
await self._run_scheduled_export()
|
||||
finally:
|
||||
await pod_lock_manager.release_lock(cronjob_id=TERNARY_USAGE_DATA_JOB_NAME)
|
||||
else:
|
||||
await self._run_scheduled_export()
|
||||
|
||||
@staticmethod
|
||||
async def init_ternary_background_job(
|
||||
scheduler: _JobScheduler,
|
||||
) -> None:
|
||||
"""Register the Ternary export job with the provided scheduler."""
|
||||
ternary_loggers: Final[list[CustomLogger]] = ( # mutable-ok: list returned by the shared callback manager
|
||||
litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=TernaryLogger)
|
||||
)
|
||||
if not ternary_loggers:
|
||||
verbose_logger.debug("No Ternary logger registered; skipping scheduler")
|
||||
return
|
||||
|
||||
ternary_logger: Final = ternary_loggers[0]
|
||||
if not isinstance(ternary_logger, TernaryLogger):
|
||||
return
|
||||
trigger_kwargs: Final = ternary_logger._build_scheduler_trigger()
|
||||
scheduler.add_job(
|
||||
ternary_logger.initialize_focus_export_job,
|
||||
**trigger_kwargs,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ("TernaryLogger",)
|
||||
|
|
@ -16,6 +16,8 @@ from litellm.integrations.focus.focus_logger import FocusLogger
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
else:
|
||||
AsyncIOScheduler = Any
|
||||
|
||||
|
|
@ -91,7 +93,7 @@ class VantageLogger(FocusLogger):
|
|||
"""
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
pod_lock_manager = None
|
||||
pod_lock_manager: PodLockManager | None = None
|
||||
if proxy_logging_obj is not None:
|
||||
writer: Final[object] = getattr(proxy_logging_obj, "db_spend_update_writer", None)
|
||||
if writer is not None:
|
||||
|
|
|
|||
|
|
@ -52,6 +52,10 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non
|
|||
_DROP_PARAMS_BOOL: Final = TypeAdapter(bool)
|
||||
|
||||
|
||||
def as_str_mapping(value: object) -> Mapping[str, object] | None:
|
||||
return value if isinstance(value, Mapping) else None
|
||||
|
||||
|
||||
def normalize_drop_params(value: object) -> bool | None:
|
||||
if value is None or isinstance(value, bool):
|
||||
return value
|
||||
|
|
|
|||
|
|
@ -48,6 +48,7 @@ from litellm.integrations.posthog import PostHogLogger
|
|||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.integrations.s3_v2 import S3Logger
|
||||
from litellm.integrations.sqs import SQSLogger
|
||||
from litellm.integrations.ternary.ternary_logger import TernaryLogger
|
||||
from litellm.integrations.vantage.vantage_logger import VantageLogger
|
||||
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
|
||||
VectorStorePreCallHook,
|
||||
|
|
@ -117,6 +118,7 @@ class CustomLoggerRegistry:
|
|||
"focus": FocusLogger,
|
||||
"mavvrik": MavvrikFocusLogger,
|
||||
"vantage": VantageLogger,
|
||||
"ternary": TernaryLogger,
|
||||
"posthog": PostHogLogger,
|
||||
"newrelic": NewRelicLogger,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ Disable remote fetching entirely:
|
|||
export LITELLM_LOCAL_BLOG_POSTS=True
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import xml.etree.ElementTree as ET
|
||||
|
|
@ -17,23 +16,53 @@ from importlib.resources import files
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict
|
||||
|
||||
from litellm import verbose_logger
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
BLOG_POSTS_TTL_SECONDS: Final[int] = 3600 # 1 hour
|
||||
BLOG_POSTS_MAX: Final[int] = 6
|
||||
BLOG_POSTS_MAX: Final[int] = 50
|
||||
|
||||
|
||||
class BlogPost(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
title: str
|
||||
description: str
|
||||
date: str
|
||||
url: str
|
||||
tags: tuple[str, ...] = ()
|
||||
|
||||
|
||||
class BlogPostsResponse(LiteLLMBaseModel):
|
||||
posts: list[BlogPost]
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
posts: tuple[BlogPost, ...]
|
||||
|
||||
|
||||
def _pub_date(pub_date_el: ET.Element | None) -> str:
|
||||
if pub_date_el is None or not pub_date_el.text:
|
||||
return ""
|
||||
try:
|
||||
return parsedate_to_datetime(pub_date_el.text).strftime("%Y-%m-%d")
|
||||
except Exception:
|
||||
return pub_date_el.text
|
||||
|
||||
|
||||
def _parse_item(item: ET.Element) -> BlogPost | None:
|
||||
title_el: Final = item.find("title")
|
||||
link_el: Final = item.find("link")
|
||||
if title_el is None or link_el is None:
|
||||
return None
|
||||
desc_el: Final = item.find("description")
|
||||
return BlogPost(
|
||||
title=title_el.text or "",
|
||||
description=(desc_el.text or "") if desc_el is not None else "",
|
||||
date=_pub_date(item.find("pubDate")),
|
||||
url=link_el.text or "",
|
||||
tags=tuple(category.text for category in item.findall("category") if category.text),
|
||||
)
|
||||
|
||||
|
||||
class GetBlogPosts:
|
||||
|
|
@ -41,19 +70,20 @@ class GetBlogPosts:
|
|||
Fetches, validates, and caches LiteLLM blog posts.
|
||||
|
||||
- Fetches RSS feed from docs site with a 5-second timeout
|
||||
- Parses the XML and extracts the latest BLOG_POSTS_MAX blog posts
|
||||
- Parses the XML and extracts up to BLOG_POSTS_MAX blog posts, with their RSS categories as tags
|
||||
- Caches the result in-process for BLOG_POSTS_TTL_SECONDS (1 hour)
|
||||
- Falls back to the bundled local backup on any failure
|
||||
"""
|
||||
|
||||
_cached_posts: list[dict[str, str]] | None = None
|
||||
_cached_posts: tuple[BlogPost, ...] | None = None
|
||||
_last_fetch_time: float = 0.0
|
||||
|
||||
@staticmethod
|
||||
def load_local_blog_posts() -> list[dict[str, str]]:
|
||||
def load_local_blog_posts() -> tuple[BlogPost, ...]:
|
||||
"""Load the bundled local backup blog posts."""
|
||||
content: Final = json.loads(files("litellm").joinpath("blog_posts.json").read_text(encoding="utf-8"))
|
||||
return content.get("posts", [])
|
||||
return BlogPostsResponse.model_validate_json(
|
||||
files("litellm").joinpath("blog_posts.json").read_text(encoding="utf-8")
|
||||
).posts
|
||||
|
||||
@staticmethod
|
||||
def fetch_rss_feed(url: str, timeout: int = 5) -> str:
|
||||
|
|
@ -67,54 +97,24 @@ class GetBlogPosts:
|
|||
return response.text
|
||||
|
||||
@staticmethod
|
||||
def parse_rss_to_posts(xml_text: str, max_posts: int = BLOG_POSTS_MAX) -> list[dict[str, str]]:
|
||||
def parse_rss_to_posts(xml_text: str, max_posts: int = BLOG_POSTS_MAX) -> tuple[BlogPost, ...]:
|
||||
"""
|
||||
Parse RSS XML and return a list of blog post dicts.
|
||||
Parse RSS XML into blog posts.
|
||||
|
||||
Extracts title, description, date (YYYY-MM-DD), and url from each <item>.
|
||||
Extracts title, description, date (YYYY-MM-DD), url and the <category> tags from each <item>.
|
||||
"""
|
||||
root: Final = ET.fromstring(xml_text)
|
||||
channel: Final = root.find("channel")
|
||||
if channel is None:
|
||||
raise ValueError("RSS feed missing <channel> element")
|
||||
|
||||
posts: Final[list[dict[str, str]]] = []
|
||||
for item in channel.findall("item"):
|
||||
if len(posts) >= max_posts:
|
||||
break
|
||||
|
||||
title_el = item.find("title")
|
||||
link_el = item.find("link")
|
||||
desc_el = item.find("description")
|
||||
pub_date_el = item.find("pubDate")
|
||||
|
||||
if title_el is None or link_el is None:
|
||||
continue
|
||||
|
||||
# Parse RFC 2822 date to YYYY-MM-DD
|
||||
date_str = ""
|
||||
if pub_date_el is not None and pub_date_el.text:
|
||||
try:
|
||||
dt = parsedate_to_datetime(pub_date_el.text)
|
||||
date_str = dt.strftime("%Y-%m-%d")
|
||||
except Exception:
|
||||
date_str = pub_date_el.text
|
||||
|
||||
posts.append(
|
||||
{
|
||||
"title": title_el.text or "",
|
||||
"description": desc_el.text or "" if desc_el is not None else "",
|
||||
"date": date_str,
|
||||
"url": link_el.text or "",
|
||||
}
|
||||
)
|
||||
|
||||
return posts
|
||||
parsed: Final = (_parse_item(item) for item in channel.findall("item"))
|
||||
return tuple(post for post in parsed if post is not None)[:max_posts]
|
||||
|
||||
@staticmethod
|
||||
def validate_blog_posts(posts: list[dict[str, str]]) -> bool:
|
||||
"""Return True if posts is a non-empty list."""
|
||||
if not isinstance(posts, list) or len(posts) == 0:
|
||||
def validate_blog_posts(posts: tuple[BlogPost, ...]) -> bool:
|
||||
"""Return True if there is at least one post."""
|
||||
if len(posts) == 0:
|
||||
verbose_logger.warning(
|
||||
"LiteLLM: Parsed RSS feed has no valid posts. Falling back to local backup.",
|
||||
)
|
||||
|
|
@ -122,9 +122,9 @@ class GetBlogPosts:
|
|||
return True
|
||||
|
||||
@classmethod
|
||||
def get_blog_posts(cls, url: str) -> list[dict[str, str]]:
|
||||
def get_blog_posts(cls, url: str) -> tuple[BlogPost, ...]:
|
||||
"""
|
||||
Return the blog posts list.
|
||||
Return the blog posts.
|
||||
|
||||
Uses the in-process cache if within BLOG_POSTS_TTL_SECONDS.
|
||||
Fetches from ``url`` otherwise, falling back to local backup on failure.
|
||||
|
|
@ -156,6 +156,6 @@ class GetBlogPosts:
|
|||
return posts
|
||||
|
||||
|
||||
def get_blog_posts(url: str) -> list[dict[str, str]]:
|
||||
"""Public entry point — returns the blog posts list."""
|
||||
def get_blog_posts(url: str) -> tuple[BlogPost, ...]:
|
||||
"""Public entry point, returns the blog posts."""
|
||||
return GetBlogPosts.get_blog_posts(url=url)
|
||||
|
|
|
|||
|
|
@ -737,6 +737,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
# Init Caching related details
|
||||
self.caching_details: CachingDetails | None = None
|
||||
self.vertex_ai_cache_creation_input_tokens: int = 0
|
||||
# Timing for results that cannot carry ``_hidden_params`` (plain-dict /v1/messages
|
||||
# responses and the bridge stream wrappers); see ``update_response_metadata``.
|
||||
self.response_timing_metrics: Mapping[str, float] = {}
|
||||
|
|
@ -4950,6 +4951,15 @@ def _init_custom_logger_compatible_class(
|
|||
vantage_logger: Final = VantageLogger()
|
||||
_in_memory_loggers.append(vantage_logger)
|
||||
return vantage_logger
|
||||
elif logging_integration == "ternary":
|
||||
from litellm.integrations.ternary.ternary_logger import TernaryLogger
|
||||
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, TernaryLogger):
|
||||
return callback
|
||||
ternary_logger: Final = TernaryLogger()
|
||||
_in_memory_loggers.append(ternary_logger)
|
||||
return ternary_logger
|
||||
elif logging_integration == "deepeval":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, DeepEvalLogger):
|
||||
|
|
@ -5455,6 +5465,12 @@ def get_custom_logger_compatible_class(
|
|||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, VantageLogger):
|
||||
return callback
|
||||
elif logging_integration == "ternary":
|
||||
from litellm.integrations.ternary.ternary_logger import TernaryLogger
|
||||
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, TernaryLogger):
|
||||
return callback
|
||||
elif logging_integration == "deepeval":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, DeepEvalLogger):
|
||||
|
|
|
|||
|
|
@ -709,7 +709,7 @@ def convert_to_model_response_object(
|
|||
|
||||
message = Message(
|
||||
content=content,
|
||||
role=choice["message"]["role"] or "assistant",
|
||||
role=choice["message"].get("role") or "assistant",
|
||||
function_call=choice["message"].get("function_call", None),
|
||||
tool_calls=tool_calls,
|
||||
audio=choice["message"].get("audio", None),
|
||||
|
|
@ -720,6 +720,8 @@ def convert_to_model_response_object(
|
|||
images=_normalize_images_for_message(choice["message"].get("images", None)),
|
||||
)
|
||||
finish_reason = choice.get("finish_reason", None)
|
||||
if "refusal" in choice["message"]:
|
||||
message.refusal = choice["message"]["refusal"]
|
||||
if finish_reason is None:
|
||||
# gpt-4 vision can return 'finish_reason' or 'finish_details'
|
||||
finish_reason = choice.get("finish_details") or "stop"
|
||||
|
|
@ -728,9 +730,10 @@ def convert_to_model_response_object(
|
|||
|
||||
## PROVIDER SPECIFIC FIELDS ##
|
||||
provider_specific_fields = {f: choice[f] for f in choice.keys() - _CHOICES_FIELDS}
|
||||
|
||||
logprobs = choice.get("logprobs", None)
|
||||
provider_sent_logprobs = "logprobs" in choice
|
||||
enhancements = choice.get("enhancements", None)
|
||||
|
||||
choice = Choices(
|
||||
finish_reason=finish_reason,
|
||||
index=idx,
|
||||
|
|
@ -739,6 +742,9 @@ def convert_to_model_response_object(
|
|||
enhancements=enhancements,
|
||||
provider_specific_fields=provider_specific_fields,
|
||||
)
|
||||
if provider_sent_logprobs and logprobs is None:
|
||||
choice.logprobs = None
|
||||
|
||||
choice_list.append(choice)
|
||||
model_response_object.choices = choice_list
|
||||
|
||||
|
|
|
|||
|
|
@ -466,7 +466,7 @@ class LoggingWorker:
|
|||
# Process all tasks concurrently for maximum speed
|
||||
await asyncio.gather(*[self._process_single_task(task) for task in tasks])
|
||||
|
||||
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine):
|
||||
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, None, None]) -> None:
|
||||
"""
|
||||
Ensure the logging worker is initialized and enqueue the coroutine.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1788,7 +1788,8 @@ _MarkedT: Final = TypeVar("_MarkedT", bound=Mapping[str, object])
|
|||
|
||||
|
||||
def with_prompt_cache_breakpoint(target: _MarkedT, marker: object) -> _MarkedT:
|
||||
if marker is None:
|
||||
"""A breakpoint the client already wrote on the block wins, TTL included."""
|
||||
if marker is None or target.get("prompt_cache_breakpoint") is not None:
|
||||
return target
|
||||
marked: Final = {**target, "prompt_cache_breakpoint": marker}
|
||||
return cast(_MarkedT, marked) # cast-ok: same block shape as the input plus the marker key
|
||||
|
|
|
|||
|
|
@ -108,6 +108,7 @@ def _redact_choice_content(choice):
|
|||
choice.message.reasoning_content = REDACTED_BY_LITELLM
|
||||
if hasattr(choice.message, "thinking_blocks"):
|
||||
choice.message.thinking_blocks = None
|
||||
_redact_provider_specific_fields(choice.message)
|
||||
_redact_tool_calls(getattr(choice.message, "tool_calls", None))
|
||||
_redact_function_call(getattr(choice.message, "function_call", None))
|
||||
elif isinstance(choice, litellm.utils.StreamingChoices):
|
||||
|
|
@ -117,10 +118,39 @@ def _redact_choice_content(choice):
|
|||
choice.delta.reasoning_content = REDACTED_BY_LITELLM
|
||||
if hasattr(choice.delta, "thinking_blocks"):
|
||||
choice.delta.thinking_blocks = None
|
||||
_redact_provider_specific_fields(choice.delta)
|
||||
_redact_tool_calls(getattr(choice.delta, "tool_calls", None))
|
||||
_redact_function_call(getattr(choice.delta, "function_call", None))
|
||||
|
||||
|
||||
def _as_mapping(value: object) -> Mapping[str, object] | None:
|
||||
return value if isinstance(value, Mapping) else None # pyright: ignore[reportUnknownVariableType] # keys are read one by one
|
||||
|
||||
|
||||
def _redact_provider_specific_fields(message: object) -> None:
|
||||
if (fields := _as_mapping(getattr(message, "provider_specific_fields", None))) is not None:
|
||||
setattr(message, "provider_specific_fields", _redacted_provider_specific_fields(fields)) # noqa: B010 # typed as object
|
||||
|
||||
|
||||
def _redact_provider_specific_fields_dict(choice: Mapping[str, object]) -> None:
|
||||
for message in (choice.get("message"), choice.get("delta")):
|
||||
if (
|
||||
(readable := _as_mapping(message)) is not None
|
||||
and (fields := _as_mapping(readable.get("provider_specific_fields"))) is not None
|
||||
and isinstance(message, dict)
|
||||
):
|
||||
message["provider_specific_fields"] = _redacted_provider_specific_fields(fields)
|
||||
|
||||
|
||||
def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""Copy of provider_specific_fields with its reasoning redacted the same way as the message's own."""
|
||||
return {
|
||||
**fields,
|
||||
**({"reasoning_content": REDACTED_BY_LITELLM} if fields.get("reasoning_content") is not None else {}),
|
||||
**{key: None for key in ("thinking_blocks", "reasoningContentBlocks", "citations") if key in fields},
|
||||
}
|
||||
|
||||
|
||||
def _redact_responses_api_output(output_items):
|
||||
"""Helper to redact ResponsesAPIResponse output items."""
|
||||
for output_item in output_items:
|
||||
|
|
@ -228,6 +258,7 @@ def _redact_tool_calls_dict(message: Mapping[str, object]) -> None:
|
|||
def _redact_model_response_dict_choices(choices, redacted_str: str):
|
||||
for choice in choices:
|
||||
if isinstance(choice, dict):
|
||||
_redact_provider_specific_fields_dict(choice)
|
||||
if "message" in choice and isinstance(choice["message"], dict):
|
||||
if choice["message"].get("content") is not None:
|
||||
choice["message"]["content"] = redacted_str
|
||||
|
|
@ -261,8 +292,37 @@ def _redacted_baseline_metadata(metadata: Mapping[str, object]) -> Mapping[str,
|
|||
return {**metadata, "_autorouter_baseline_route": replace(route, request_parameters=None)}
|
||||
|
||||
|
||||
def redacted_litellm_params(params: Mapping[str, object]) -> dict[str, object]:
|
||||
request: Final = params.get("proxy_server_request")
|
||||
def _redacted_proxy_server_request(request: Mapping[str, object], *, redact_body: bool) -> Mapping[str, object]:
|
||||
"""Copy of proxy_server_request; redact_body redacts its body the same way as the top-level request."""
|
||||
body: Final = _as_mapping(request.get("body"))
|
||||
if not redact_body or body is None:
|
||||
return without_classifier_audit(request)
|
||||
redacted_body: Final = {**body}
|
||||
perform_redaction(redacted_body, result=None)
|
||||
return {**without_classifier_audit(request), "body": redacted_body}
|
||||
|
||||
|
||||
def _redact_hidden_provider_specific_fields(params: Mapping[str, object]) -> None:
|
||||
"""Redact the provider_specific_fields copied into metadata hidden_params, keeping each metadata dict."""
|
||||
for metadata in (params.get("metadata"), params.get("litellm_metadata")):
|
||||
if (
|
||||
(readable := _as_mapping(metadata)) is not None
|
||||
and (hidden_params := _as_mapping(readable.get("hidden_params"))) is not None
|
||||
and (fields := _as_mapping(hidden_params.get("provider_specific_fields"))) is not None
|
||||
and isinstance(metadata, dict)
|
||||
):
|
||||
metadata["hidden_params"] = {
|
||||
**hidden_params,
|
||||
"provider_specific_fields": _redacted_provider_specific_fields(fields),
|
||||
}
|
||||
|
||||
|
||||
def redacted_litellm_params(params: Mapping[str, object], *, redact_request_body: bool = True) -> dict[str, object]:
|
||||
"""
|
||||
Copy of litellm_params without classifier audit fields or baseline route request parameters.
|
||||
redact_request_body also redacts the nested proxy_server_request body.
|
||||
"""
|
||||
request: Final = _as_mapping(params.get("proxy_server_request"))
|
||||
return {
|
||||
**params,
|
||||
**{
|
||||
|
|
@ -270,7 +330,11 @@ def redacted_litellm_params(params: Mapping[str, object]) -> dict[str, object]:
|
|||
for key, value in params.items()
|
||||
if key in ("metadata", "litellm_metadata") and isinstance(value, Mapping)
|
||||
},
|
||||
**({"proxy_server_request": without_classifier_audit(request)} if isinstance(request, Mapping) else {}),
|
||||
**(
|
||||
{"proxy_server_request": _redacted_proxy_server_request(request, redact_body=redact_request_body)}
|
||||
if request is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -280,12 +344,17 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
|
|||
|
||||
redact_streaming_responses=False skips the in-place redaction of the shared streaming
|
||||
response entries; per-callback redaction hands each opted-out callback its own redacted
|
||||
copy via redact_streaming_responses_for_custom_logger instead.
|
||||
copy via redact_streaming_responses_for_custom_logger instead. It also leaves the shared
|
||||
proxy request body and hidden provider fields as they are, since the spend log reads them.
|
||||
"""
|
||||
# Redact model_call_details
|
||||
params: Final = model_call_details.get("litellm_params")
|
||||
if isinstance(params, Mapping):
|
||||
model_call_details["litellm_params"] = redacted_litellm_params(params)
|
||||
params: Final = _as_mapping(model_call_details.get("litellm_params"))
|
||||
if params is not None:
|
||||
if redact_streaming_responses:
|
||||
_redact_hidden_provider_specific_fields(params)
|
||||
model_call_details["litellm_params"] = redacted_litellm_params(
|
||||
params, redact_request_body=redact_streaming_responses
|
||||
)
|
||||
model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}]
|
||||
model_call_details["prompt"] = ""
|
||||
model_call_details["input"] = ""
|
||||
|
|
|
|||
|
|
@ -63,8 +63,10 @@ def get_cost_for_web_search_request(custom_llm_provider: str, usage: "Usage", mo
|
|||
|
||||
return cost_per_web_search_request_vertex_ai(usage=usage, model_info=model_info)
|
||||
elif custom_llm_provider == "azure_ai":
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests_from_usage
|
||||
|
||||
azure_ai_model_key: Final[str] = model_info.get("key", "") if model_info else ""
|
||||
if "claude" not in azure_ai_model_key.lower():
|
||||
if "claude" not in azure_ai_model_key.lower() and get_web_search_requests_from_usage(usage) is None:
|
||||
return None
|
||||
from .anthropic.cost_calculation import get_cost_for_anthropic_web_search
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.llms.azure_ai.common_utils import AZURE_ENTRA_LITELLM_PARAM_KEYS, get_azure_ai_agent_entra_token
|
||||
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
|
||||
|
|
@ -30,6 +31,7 @@ if TYPE_CHECKING:
|
|||
_REGISTRY_PARAMS_KEPT_OUT_OF_OPTIONAL_PARAMS: Final = (
|
||||
frozenset({"api_key", "api_base", "headers", "model"}) | AZURE_ENTRA_LITELLM_PARAM_KEYS
|
||||
)
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
|
||||
def _card_declares_no_streaming(agent_card_params: Mapping[str, object]) -> bool:
|
||||
|
|
@ -325,7 +327,7 @@ class A2AConfig(BaseConfig):
|
|||
|
||||
# Check for JSON-RPC error
|
||||
if "error" in response_json:
|
||||
error: Final = response_json["error"]
|
||||
error: Final = _JSON_OBJECT.validate_python(response_json["error"])
|
||||
raise A2AError(
|
||||
status_code=raw_response.status_code,
|
||||
message=f"A2A error: {error.get('message', 'Unknown error')}",
|
||||
|
|
@ -333,7 +335,7 @@ class A2AConfig(BaseConfig):
|
|||
)
|
||||
|
||||
# Extract text from A2A response
|
||||
text: Final = extract_text_from_a2a_response(response_json)
|
||||
text: Final = extract_text_from_a2a_response(_JSON_OBJECT.validate_python(response_json))
|
||||
|
||||
# Populate model response
|
||||
model_response.choices = [
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ class A2AError(BaseLLMException):
|
|||
self,
|
||||
status_code: int,
|
||||
message: str,
|
||||
headers: dict[str, Any] = {},
|
||||
headers: dict[str, str] = {},
|
||||
):
|
||||
super().__init__(
|
||||
status_code=status_code,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
|
|
@ -21,6 +23,10 @@ if TYPE_CHECKING:
|
|||
else:
|
||||
LiteLLMLoggingObj = Any
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter(
|
||||
Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True)
|
||||
)
|
||||
|
||||
OPENAI_STYLE_IMAGE_MODEL_PREFIXES: Final[tuple[str, ...]] = ("openai/",)
|
||||
|
||||
|
|
@ -198,29 +204,33 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig):
|
|||
|
||||
if "data" in response_data and isinstance(response_data["data"], list):
|
||||
# Handle OpenAI-like format: {"data": [{"url": "...", "width": 1024, "height": 768, "content_type": "image/jpeg"}]}
|
||||
for image in response_data["data"]:
|
||||
for raw_image in _JSON_CONTAINER.validate_python(response_data["data"]):
|
||||
image = raw_image if isinstance(raw_image, str) else _JSON_CONTAINER.validate_python(raw_image)
|
||||
if "url" in image:
|
||||
url_image = _JSON_OBJECT.validate_python(image)
|
||||
model_response.data.append(
|
||||
ImageObject(
|
||||
b64_json=None,
|
||||
url=image["url"],
|
||||
revised_prompt=image.get("revised_prompt"),
|
||||
url=url_image["url"],
|
||||
revised_prompt=url_image.get("revised_prompt"),
|
||||
)
|
||||
)
|
||||
elif "b64_json" in image or "image_base64" in image:
|
||||
b64_image = _JSON_OBJECT.validate_python(image)
|
||||
model_response.data.append(
|
||||
ImageObject(
|
||||
b64_json=image.get("b64_json") or image.get("image_base64"),
|
||||
b64_json=b64_image.get("b64_json") or b64_image.get("image_base64"),
|
||||
url=None,
|
||||
revised_prompt=image.get("revised_prompt"),
|
||||
revised_prompt=b64_image.get("revised_prompt"),
|
||||
)
|
||||
)
|
||||
elif "output" in response_data and "choices" in response_data["output"]:
|
||||
for choice in response_data["output"]["choices"]:
|
||||
for raw_choice in _JSON_CONTAINER.validate_python(response_data["output"]["choices"]):
|
||||
choice = raw_choice if isinstance(raw_choice, str) else _JSON_CONTAINER.validate_python(raw_choice)
|
||||
if "image_base64" in choice:
|
||||
model_response.data.append(
|
||||
ImageObject(
|
||||
b64_json=choice["image_base64"],
|
||||
b64_json=_JSON_OBJECT.validate_python(choice)["image_base64"],
|
||||
url=None,
|
||||
)
|
||||
)
|
||||
|
|
@ -228,23 +238,24 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig):
|
|||
model_response.data.append(
|
||||
ImageObject(
|
||||
b64_json=None,
|
||||
url=choice["url"],
|
||||
url=_JSON_OBJECT.validate_python(choice)["url"],
|
||||
)
|
||||
)
|
||||
elif "images" in response_data:
|
||||
# Handle alternative format: {"images": [{"url": "...", "width": 1024, "height": 768, "content_type": "image/jpeg"}]}
|
||||
for image in response_data["images"]:
|
||||
for raw_image in _JSON_CONTAINER.validate_python(response_data["images"]):
|
||||
image = raw_image if isinstance(raw_image, str) else _JSON_CONTAINER.validate_python(raw_image)
|
||||
if "url" in image:
|
||||
model_response.data.append(
|
||||
ImageObject(
|
||||
b64_json=None,
|
||||
url=image["url"],
|
||||
url=_JSON_OBJECT.validate_python(image)["url"],
|
||||
)
|
||||
)
|
||||
elif "image_base64" in image:
|
||||
model_response.data.append(
|
||||
ImageObject(
|
||||
b64_json=image["image_base64"],
|
||||
b64_json=_JSON_OBJECT.validate_python(image)["image_base64"],
|
||||
url=None,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,9 @@ Calling + translation logic for anthropic's `/v1/messages` endpoint
|
|||
"""
|
||||
|
||||
import copy
|
||||
import io
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Union, cast
|
||||
|
|
@ -20,6 +22,11 @@ from litellm.anthropic_beta_headers_manager import (
|
|||
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.litellm_core_utils.core_helpers import map_finish_reason
|
||||
from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator
|
||||
from litellm.llms.base_llm.base_utils import (
|
||||
RESPONSE_FORMAT_ENVELOPE_NAME_KEYS,
|
||||
ResponseFormatSchema,
|
||||
response_format_envelope_contents,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
|
|
@ -57,7 +64,7 @@ from litellm.types.utils import (
|
|||
|
||||
from ...base import BaseLLM
|
||||
from ..common_utils import AnthropicError, process_anthropic_headers
|
||||
from .transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY, AnthropicConfig
|
||||
from .transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY, AnthropicConfig, JsonModeRequest
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -69,6 +76,67 @@ def _loads_stream_chunk(payload: str) -> dict[str, object]:
|
|||
return json.loads(payload)
|
||||
|
||||
|
||||
_EnvelopeVerdict = Literal["undecided", "pass_through", "hold"]
|
||||
_FIRST_OBJECT_KEY: Final = re.compile(r'^\s*\{\s*"((?:[^"\\]|\\.)*)"\s*:')
|
||||
_OBJECT_KEY_PREFIX: Final = re.compile(r'^\s*(?:\{\s*(?:"(?:[^"\\]|\\.)*\\?(?:"\s*)?)?)?$')
|
||||
_ENVELOPE_NAME_KEY_PATTERN: Final = "|".join(sorted(RESPONSE_FORMAT_ENVELOPE_NAME_KEYS))
|
||||
_FIRST_NAME_VALUE: Final = re.compile(r'^\s*\{\s*"(?:' + _ENVELOPE_NAME_KEY_PATTERN + r')"\s*:\s*"((?:[^"\\]|\\.)*)"')
|
||||
_NAME_VALUE_PREFIX: Final = re.compile(
|
||||
r'^\s*\{\s*"(?:' + _ENVELOPE_NAME_KEY_PATTERN + r')"\s*:\s*(?:"(?:[^"\\]|\\.)*\\?)?$'
|
||||
)
|
||||
_STRING_VALUE_THEN_KEY: Final = re.compile(r'^"(?:[^"\\]|\\.)*"\s*,\s*"((?:[^"\\]|\\.)*)"\s*:')
|
||||
_UNDECIDED_PREFIX_LIMIT: Final = 4096
|
||||
|
||||
|
||||
def _key_after_a_string_verdict(string_value_start: str, schema: ResponseFormatSchema) -> _EnvelopeVerdict:
|
||||
next_key: Final = _STRING_VALUE_THEN_KEY.match(string_value_start)
|
||||
if next_key is None:
|
||||
return "undecided"
|
||||
return "pass_through" if next_key.group(1) in schema.property_names else "hold"
|
||||
|
||||
|
||||
def _first_value_verdict(after_first_key: str, schema: ResponseFormatSchema) -> _EnvelopeVerdict:
|
||||
value_start: Final = after_first_key.lstrip()
|
||||
if not value_start:
|
||||
return "undecided"
|
||||
if value_start.startswith("{"):
|
||||
return "hold"
|
||||
if value_start.startswith('"') and schema.required_names:
|
||||
return _key_after_a_string_verdict(value_start, schema)
|
||||
return "pass_through"
|
||||
|
||||
|
||||
def _prefix_verdict(buffered_json: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
|
||||
first_key: Final = _FIRST_OBJECT_KEY.match(buffered_json)
|
||||
if first_key is None:
|
||||
return "undecided" if _OBJECT_KEY_PREFIX.match(buffered_json) else "pass_through"
|
||||
if first_key.group(1) in RESPONSE_FORMAT_ENVELOPE_NAME_KEYS:
|
||||
first_value: Final = _FIRST_NAME_VALUE.match(buffered_json)
|
||||
if first_value is not None and first_value.group(1) == tool_name:
|
||||
return "hold"
|
||||
if first_value is None and _NAME_VALUE_PREFIX.match(buffered_json):
|
||||
return "undecided"
|
||||
if first_key.group(1) in schema.property_names or not schema.property_names or schema.accepts_a_lone_undeclared_key:
|
||||
return "pass_through"
|
||||
return _first_value_verdict(buffered_json[first_key.end() :], schema)
|
||||
|
||||
|
||||
def _held_answer(held: str, schema: ResponseFormatSchema) -> str:
|
||||
try:
|
||||
payload: Final[object] = json.loads(held)
|
||||
except (json.JSONDecodeError, RecursionError):
|
||||
return held
|
||||
contents: Final = response_format_envelope_contents(payload, schema, RESPONSE_FORMAT_TOOL_NAME)
|
||||
return held if contents is None else json.dumps(contents)
|
||||
|
||||
|
||||
def _envelope_verdict(buffered_json: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
|
||||
verdict: Final = _prefix_verdict(buffered_json, schema, tool_name)
|
||||
if verdict == "undecided" and len(buffered_json) > _UNDECIDED_PREFIX_LIMIT:
|
||||
return "hold"
|
||||
return verdict
|
||||
|
||||
|
||||
async def make_call(
|
||||
client: AsyncHTTPHandler | None,
|
||||
api_base: str,
|
||||
|
|
@ -81,6 +149,7 @@ async def make_call(
|
|||
json_mode: bool,
|
||||
speed: str | None = None,
|
||||
tool_name_reverse_map: dict[str, str] | None = None,
|
||||
response_format_schema: ResponseFormatSchema | None = None,
|
||||
) -> tuple["ModelResponseIterator", httpx.Headers]:
|
||||
if client is None:
|
||||
client = litellm.module_level_aclient
|
||||
|
|
@ -116,6 +185,7 @@ async def make_call(
|
|||
json_mode=json_mode,
|
||||
speed=speed,
|
||||
tool_name_reverse_map=tool_name_reverse_map,
|
||||
response_format_schema=response_format_schema,
|
||||
)
|
||||
|
||||
# LOGGING
|
||||
|
|
@ -141,6 +211,7 @@ def make_sync_call(
|
|||
json_mode: bool,
|
||||
speed: str | None = None,
|
||||
tool_name_reverse_map: dict[str, str] | None = None,
|
||||
response_format_schema: ResponseFormatSchema | None = None,
|
||||
) -> tuple["ModelResponseIterator", httpx.Headers]:
|
||||
if client is None:
|
||||
client = litellm.module_level_client # re-use a module level client
|
||||
|
|
@ -184,6 +255,7 @@ def make_sync_call(
|
|||
json_mode=json_mode,
|
||||
speed=speed,
|
||||
tool_name_reverse_map=tool_name_reverse_map,
|
||||
response_format_schema=response_format_schema,
|
||||
)
|
||||
|
||||
# LOGGING
|
||||
|
|
@ -241,6 +313,9 @@ class AnthropicChatCompletion(BaseLLM):
|
|||
tool_name_reverse_map=(
|
||||
litellm_params.get(ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY) if isinstance(litellm_params, dict) else None
|
||||
),
|
||||
response_format_schema=(
|
||||
JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode and optional_params else None
|
||||
),
|
||||
)
|
||||
streamwrapper: Final = CustomStreamWrapper(
|
||||
completion_stream=completion_stream,
|
||||
|
|
@ -485,6 +560,9 @@ class AnthropicChatCompletion(BaseLLM):
|
|||
if isinstance(litellm_params, dict)
|
||||
else None
|
||||
),
|
||||
response_format_schema=(
|
||||
JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode else None
|
||||
),
|
||||
)
|
||||
return CustomStreamWrapper(
|
||||
completion_stream=completion_stream,
|
||||
|
|
@ -550,6 +628,7 @@ class ModelResponseIterator:
|
|||
json_mode: bool | None = False,
|
||||
speed: str | None = None,
|
||||
tool_name_reverse_map: dict[str, str] | None = None,
|
||||
response_format_schema: ResponseFormatSchema | None = None,
|
||||
):
|
||||
self.streaming_response = streaming_response
|
||||
self.response_iterator = self.streaming_response
|
||||
|
|
@ -573,6 +652,9 @@ class ModelResponseIterator:
|
|||
self.is_response_format_tool: bool = False
|
||||
# Track if we've converted any response_format tools (affects finish_reason)
|
||||
self.converted_response_format_tool: bool = False
|
||||
self.response_format_schema = response_format_schema
|
||||
self._response_format_buffer: io.StringIO = io.StringIO()
|
||||
self._response_format_verdict: _EnvelopeVerdict = "undecided"
|
||||
|
||||
# For handling partial JSON chunks from fragmentation
|
||||
# See: https://github.com/BerriAI/litellm/issues/17473
|
||||
|
|
@ -955,6 +1037,8 @@ class ModelResponseIterator:
|
|||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
self._current_server_tool_id = None
|
||||
if self.is_response_format_tool:
|
||||
text = self._flush_response_format_buffer()
|
||||
# Reset response_format tool tracking when block stops
|
||||
self.is_response_format_tool = False
|
||||
# Reset current content block type
|
||||
|
|
@ -1072,6 +1156,10 @@ class ModelResponseIterator:
|
|||
tool_name: Final = tool_use.get("function", {}).get("name", "")
|
||||
self.is_response_format_tool = tool_name == RESPONSE_FORMAT_TOOL_NAME
|
||||
|
||||
if self.is_response_format_tool and self.response_format_schema is not None:
|
||||
self.converted_response_format_tool = True
|
||||
return self._stream_response_format_fragment(tool_use, self.response_format_schema), None
|
||||
|
||||
# Convert tool to content if we're tracking a response_format tool
|
||||
if self.is_response_format_tool:
|
||||
message: Final = AnthropicConfig.convert_tool_response_to_message(tool_calls=[tool_use])
|
||||
|
|
@ -1083,6 +1171,31 @@ class ModelResponseIterator:
|
|||
|
||||
return text, tool_use
|
||||
|
||||
def _stream_response_format_fragment(
|
||||
self, tool_use: ChatCompletionToolCallChunk, schema: ResponseFormatSchema
|
||||
) -> str:
|
||||
if tool_use.get("id") is not None:
|
||||
self._response_format_buffer = io.StringIO()
|
||||
self._response_format_verdict = "undecided"
|
||||
return ""
|
||||
fragment: Final = tool_use.get("function", {}).get("arguments") or ""
|
||||
if self._response_format_verdict == "pass_through":
|
||||
return fragment
|
||||
self._response_format_buffer.write(fragment)
|
||||
if self._response_format_verdict == "hold":
|
||||
return ""
|
||||
buffered: Final = self._response_format_buffer.getvalue()
|
||||
self._response_format_verdict = _envelope_verdict(buffered, schema, RESPONSE_FORMAT_TOOL_NAME)
|
||||
return buffered if self._response_format_verdict == "pass_through" else ""
|
||||
|
||||
def _flush_response_format_buffer(self) -> str:
|
||||
if self.response_format_schema is None or self._response_format_verdict == "pass_through":
|
||||
return ""
|
||||
held: Final = self._response_format_buffer.getvalue()
|
||||
self._response_format_buffer = io.StringIO()
|
||||
self._response_format_verdict = "undecided"
|
||||
return _held_answer(held, self.response_format_schema) if held else ""
|
||||
|
||||
def _handle_message_delta(self, chunk: dict) -> tuple[str, Usage | None, dict[str, Any] | None]:
|
||||
"""
|
||||
Handle message_delta event for finish_reason, usage, and container.
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
|
|
@ -35,7 +35,13 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import
|
|||
place_mid_conversation_system,
|
||||
split_leading_system_run,
|
||||
)
|
||||
from litellm.llms.base_llm.base_utils import type_to_response_format_param
|
||||
from litellm.llms.base_llm.base_utils import (
|
||||
ResponseFormatSchema,
|
||||
as_json_object,
|
||||
response_format_schema,
|
||||
type_to_response_format_param,
|
||||
unwrap_response_format_envelope,
|
||||
)
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.types.llms.anthropic import (
|
||||
ANTHROPIC_ADVISOR_TOOL_TYPE,
|
||||
|
|
@ -1752,7 +1758,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
else:
|
||||
_input_schema.update(cast(AnthropicInputSchema, json_schema))
|
||||
|
||||
_tool: Final = AnthropicMessagesTool(name=RESPONSE_FORMAT_TOOL_NAME, input_schema=_input_schema)
|
||||
_tool: Final = AnthropicMessagesTool(
|
||||
name=RESPONSE_FORMAT_TOOL_NAME,
|
||||
input_schema=_input_schema,
|
||||
)
|
||||
return _tool
|
||||
|
||||
def should_strip_billing_metadata(self) -> bool:
|
||||
|
|
@ -2205,6 +2214,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
self,
|
||||
json_mode: bool | None,
|
||||
tool_calls: list[ChatCompletionToolCallChunk],
|
||||
response_format_schema: ResponseFormatSchema | None = None,
|
||||
) -> tuple[
|
||||
LitellmMessage | None,
|
||||
list[ChatCompletionToolCallChunk],
|
||||
|
|
@ -2224,11 +2234,15 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
json_tool: Final = tool_calls[json_indices[0]]
|
||||
if json_tool.get("function", {}).get("arguments") is None:
|
||||
return None, tool_calls, None
|
||||
_message: Final = AnthropicConfig.convert_tool_response_to_message(tool_calls=[json_tool])
|
||||
_message: Final = AnthropicConfig.convert_tool_response_to_message(
|
||||
tool_calls=[json_tool], response_format_schema=response_format_schema
|
||||
)
|
||||
return _message, [], None
|
||||
|
||||
first_json: Final = tool_calls[json_indices[0]]
|
||||
json_msg: Final = AnthropicConfig.convert_tool_response_to_message([first_json])
|
||||
json_msg: Final = AnthropicConfig.convert_tool_response_to_message(
|
||||
[first_json], response_format_schema=response_format_schema
|
||||
)
|
||||
extra_content: Final[str | None] = json_msg.content if json_msg is not None else None
|
||||
filtered_tools: Final = [t for i, t in enumerate(tool_calls) if i not in json_indices]
|
||||
return None, filtered_tools, extra_content
|
||||
|
|
@ -2664,6 +2678,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
prefix_prompt: str | None = None,
|
||||
speed: str | None = None,
|
||||
tool_name_reverse_map: dict[str, str] | None = None,
|
||||
response_format_schema: ResponseFormatSchema | None = None,
|
||||
):
|
||||
_hidden_params: Final[dict] = {}
|
||||
_hidden_params["additional_headers"] = process_anthropic_headers(dict(raw_response.headers))
|
||||
|
|
@ -2721,6 +2736,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
json_mode_message, tool_calls_for_message, json_extra_content = self._resolve_json_mode_non_streaming(
|
||||
json_mode=json_mode,
|
||||
tool_calls=tool_calls,
|
||||
response_format_schema=response_format_schema,
|
||||
)
|
||||
merged_text = text_content or ""
|
||||
if json_extra_content:
|
||||
|
|
@ -2829,12 +2845,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
prefix_prompt=prefix_prompt,
|
||||
speed=speed,
|
||||
tool_name_reverse_map=tool_name_reverse_map,
|
||||
response_format_schema=JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode else None,
|
||||
)
|
||||
return model_response
|
||||
|
||||
@staticmethod
|
||||
def _convert_tool_response_to_message(
|
||||
tool_calls: list[ChatCompletionToolCallChunk],
|
||||
response_format_schema: ResponseFormatSchema | None = None,
|
||||
) -> LitellmMessage | None:
|
||||
"""
|
||||
In JSON mode, Anthropic API returns JSON schema as a tool call, we need to convert it to a message to follow the OpenAI format
|
||||
|
|
@ -2842,28 +2860,32 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
"""
|
||||
## HANDLE JSON MODE - anthropic returns single function call
|
||||
json_mode_content_str: Final[str | None] = tool_calls[0]["function"].get("arguments")
|
||||
if json_mode_content_str is None:
|
||||
return None
|
||||
try:
|
||||
if json_mode_content_str is not None:
|
||||
args: Final = json.loads(json_mode_content_str)
|
||||
if isinstance(args, dict) and (values := args.get("values")) is not None:
|
||||
_message = litellm.Message(content=json.dumps(values))
|
||||
return _message
|
||||
else:
|
||||
# a lot of the times the `values` key is not present in the tool response
|
||||
# relevant issue: https://github.com/BerriAI/litellm/issues/6741
|
||||
_message = litellm.Message(content=json.dumps(args))
|
||||
return _message
|
||||
except json.JSONDecodeError:
|
||||
payload: Final[object] = json.loads(json_mode_content_str)
|
||||
except (json.JSONDecodeError, RecursionError):
|
||||
# json decode error does occur, return the original tool response str
|
||||
return litellm.Message(content=json_mode_content_str)
|
||||
return None
|
||||
if response_format_schema is not None:
|
||||
return litellm.Message(
|
||||
content=json.dumps(
|
||||
unwrap_response_format_envelope(payload, response_format_schema, RESPONSE_FORMAT_TOOL_NAME)
|
||||
)
|
||||
)
|
||||
# a lot of the times the `values` key is not present in the tool response
|
||||
# relevant issue: https://github.com/BerriAI/litellm/issues/6741
|
||||
payload_object: Final = as_json_object(payload)
|
||||
values: Final = None if payload_object is None else payload_object.get("values")
|
||||
return litellm.Message(content=json.dumps(payload if values is None else values))
|
||||
|
||||
@classmethod
|
||||
def convert_tool_response_to_message(
|
||||
cls,
|
||||
tool_calls: list[ChatCompletionToolCallChunk], # mutable-ok: mirrors override contract
|
||||
response_format_schema: ResponseFormatSchema | None = None,
|
||||
) -> LitellmMessage | None:
|
||||
return cls._convert_tool_response_to_message(tool_calls)
|
||||
return cls._convert_tool_response_to_message(tool_calls, response_format_schema)
|
||||
|
||||
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
|
||||
return AnthropicError(
|
||||
|
|
@ -2873,6 +2895,23 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
)
|
||||
|
||||
|
||||
_REQUEST_TOOLS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
|
||||
|
||||
|
||||
class JsonModeRequest(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
tools: object = None
|
||||
|
||||
def tool_schema(self) -> ResponseFormatSchema | None:
|
||||
try:
|
||||
tools: Final = _REQUEST_TOOLS.validate_python(self.tools)
|
||||
except ValidationError:
|
||||
return None
|
||||
json_tool: Final = next((tool for tool in tools if tool.get("name") == RESPONSE_FORMAT_TOOL_NAME), None)
|
||||
return None if json_tool is None else response_format_schema(json_tool.get("input_schema"))
|
||||
|
||||
|
||||
def _valid_user_id(user_id: str) -> bool:
|
||||
"""
|
||||
Validate that user_id is not an email or phone number.
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ Anthropic Token Counter implementation using the CountTokens API.
|
|||
|
||||
from typing import Any, Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import AuthenticationError
|
||||
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
|
||||
|
|
@ -26,12 +28,12 @@ class AnthropicTokenCounter(BaseTokenCounter):
|
|||
async def count_tokens(
|
||||
self,
|
||||
model_to_use: str,
|
||||
messages: list[dict[str, Any]] | None,
|
||||
messages: list[dict[str, JsonValue]] | None,
|
||||
contents: list[dict[str, object]] | None,
|
||||
deployment: dict[str, Any] | None = None,
|
||||
request_model: str = "",
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
system: Any | None = None,
|
||||
tools: list[dict[str, JsonValue]] | None = None,
|
||||
system: JsonValue = None,
|
||||
) -> TokenCountResponse | None:
|
||||
"""
|
||||
Count tokens using Anthropic's CountTokens API.
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ Anthropic Files API endpoints:
|
|||
|
||||
import calendar
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import Final, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -49,6 +49,7 @@ ANTHROPIC_FILES_BETA_HEADER: Final = "files-api-2025-04-14"
|
|||
ANTHROPIC_MESSAGE_BATCH_ID_PREFIX: Final = "msgbatch_"
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_JSON_ITEMS: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
|
||||
class AnthropicFilesConfig(BaseFilesConfig):
|
||||
|
|
@ -284,9 +285,9 @@ class AnthropicFilesConfig(BaseFilesConfig):
|
|||
"last_id": "..."
|
||||
}
|
||||
"""
|
||||
response_json: Final = raw_response.json()
|
||||
files_data: Final = response_json.get("data", [])
|
||||
return [self._parse_anthropic_file(f) for f in files_data]
|
||||
response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
|
||||
files_data: Final = _JSON_ITEMS.validate_python(response_json.get("data", []))
|
||||
return [self._parse_anthropic_file(_JSON_OBJECT.validate_python(f)) for f in files_data]
|
||||
|
||||
def transform_file_content_request(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -214,36 +214,24 @@ class _CombinedChunkSplitter:
|
|||
pieces: Final = tuple(copy.deepcopy(chunk) for _ in present_groups)
|
||||
for index, (piece, group) in enumerate(zip(pieces, present_groups)):
|
||||
copied_delta = piece.choices[0].delta
|
||||
fields = {field: value for field in group if (value := getattr(copied_delta, field, None))}
|
||||
fields = _CombinedChunkSplitter._normalize_reasoning_fields(fields)
|
||||
reasoning_content = (
|
||||
getattr(copied_delta, "reasoning_content", None) if "reasoning_content" in group else None
|
||||
)
|
||||
thinking_blocks = getattr(copied_delta, "thinking_blocks", None) if "thinking_blocks" in group else None
|
||||
unsigned_text = _CombinedChunkSplitter._unsigned_thinking_text(thinking_blocks) if thinking_blocks else None
|
||||
if unsigned_text:
|
||||
reasoning_content = unsigned_text
|
||||
thinking_blocks = None
|
||||
role = getattr(copied_delta, "role", None) if index == 0 else None
|
||||
piece.choices[0].delta = Delta(role=role, **fields)
|
||||
piece.choices[0].delta = Delta(
|
||||
role=role,
|
||||
reasoning_content=reasoning_content if reasoning_content else None,
|
||||
thinking_blocks=thinking_blocks if thinking_blocks else None,
|
||||
content=(copied_delta.content if "content" in group and copied_delta.content else None),
|
||||
tool_calls=(copied_delta.tool_calls if "tool_calls" in group and copied_delta.tool_calls else None),
|
||||
)
|
||||
return pieces
|
||||
|
||||
@staticmethod
|
||||
def _normalize_reasoning_fields(fields: "dict[str, Any]") -> "dict[str, Any]":
|
||||
"""Collapse signature-less ``thinking_blocks`` into ``reasoning_content``.
|
||||
|
||||
The block opener seeds a ``thinking_blocks`` start body with the full
|
||||
thinking text while the delta re-emits it, so SSE accumulators would
|
||||
collect it twice; the ``reasoning_content`` branch opens an empty body.
|
||||
Signature-carrying blocks are kept intact so ``signature_delta``
|
||||
suppression of the full-text snapshot still applies.
|
||||
"""
|
||||
thinking_blocks: Final = fields.get("thinking_blocks")
|
||||
if not thinking_blocks:
|
||||
return fields
|
||||
if any(block.get("signature") for block in thinking_blocks if isinstance(block, dict)):
|
||||
return fields
|
||||
thinking_text: Final = "".join(
|
||||
block.get("thinking") or ""
|
||||
for block in thinking_blocks
|
||||
if isinstance(block, dict) and block.get("type") == "thinking"
|
||||
)
|
||||
if not thinking_text:
|
||||
return fields
|
||||
return {"reasoning_content": thinking_text}
|
||||
|
||||
@staticmethod
|
||||
def _split(chunk: "ModelResponseStream") -> "list[ModelResponseStream]":
|
||||
"""Return ``[chunk]``, or ``[content_chunk, finish_chunk]`` if combined."""
|
||||
|
|
@ -270,6 +258,27 @@ class _CombinedChunkSplitter:
|
|||
def __iter__(self) -> "Iterator[ModelResponseStream]":
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def _unsigned_thinking_text(
|
||||
blocks: Sequence[Mapping[str, Any]],
|
||||
) -> str | None:
|
||||
"""Return the combined thinking text when no block carries a signature.
|
||||
|
||||
The block opener seeds a ``thinking_blocks`` start body with the full
|
||||
thinking text while the delta re-emits it, so SSE accumulators would
|
||||
collect it twice; the ``reasoning_content`` branch opens an empty body.
|
||||
Signature-carrying blocks are kept intact so ``signature_delta``
|
||||
suppression of the full-text snapshot still applies.
|
||||
"""
|
||||
if any(block.get("signature") for block in blocks if isinstance(block, dict)):
|
||||
return None
|
||||
thinking_text: Final = "".join(
|
||||
block.get("thinking") or ""
|
||||
for block in blocks
|
||||
if isinstance(block, dict) and block.get("type") == "thinking"
|
||||
)
|
||||
return thinking_text or None
|
||||
|
||||
def __next__(self) -> "ModelResponseStream":
|
||||
if self._buffer:
|
||||
return self._buffer.popleft()
|
||||
|
|
|
|||
|
|
@ -1347,7 +1347,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
self,
|
||||
choices: list[Choices],
|
||||
tool_name_mapping: dict[str, str] | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
) -> list[dict[str, object]]:
|
||||
new_content: Final[list[dict[str, Any]]] = []
|
||||
for choice, compaction_blocks in (
|
||||
(choice, _compaction_blocks(_optional_attr(choice.message, "provider_specific_fields")))
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from openai import (
|
|||
BadRequestError,
|
||||
OpenAI,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
|
||||
import litellm
|
||||
from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES
|
||||
|
|
@ -41,6 +42,7 @@ from ..base import BaseLLM
|
|||
from ..openai.common_utils import (
|
||||
build_output_token_limit_response,
|
||||
is_output_token_limit_error,
|
||||
model_dump_provider_response,
|
||||
)
|
||||
from .common_utils import (
|
||||
AzureOpenAIError,
|
||||
|
|
@ -352,12 +354,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
headers, response = self.make_sync_azure_openai_chat_completion_request(
|
||||
azure_client=azure_client, data=data, timeout=timeout
|
||||
)
|
||||
if isinstance(response, str):
|
||||
if not isinstance(response, BaseModel):
|
||||
raise AzureOpenAIError(
|
||||
status_code=500,
|
||||
message=f"Unexpected string response from Azure: {response[:500]}",
|
||||
message=f"Unexpected response from Azure: {response!r:.500}",
|
||||
)
|
||||
stringified_response: Final = response.model_dump()
|
||||
stringified_response: Final = model_dump_provider_response(response)
|
||||
## LOGGING
|
||||
logging_obj.post_call(
|
||||
input=messages,
|
||||
|
|
@ -447,12 +449,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
)
|
||||
logging_obj.model_call_details["response_headers"] = headers
|
||||
|
||||
if isinstance(response, str):
|
||||
if not isinstance(response, BaseModel):
|
||||
raise AzureOpenAIError(
|
||||
status_code=500,
|
||||
message=f"Unexpected string response from Azure: {response[:500]}",
|
||||
message=f"Unexpected response from Azure: {response!r:.500}",
|
||||
)
|
||||
stringified_response: Final = response.model_dump()
|
||||
stringified_response: Final = model_dump_provider_response(response)
|
||||
logging_obj.post_call(
|
||||
input=data["messages"],
|
||||
api_key=api_key,
|
||||
|
|
|
|||
|
|
@ -4,9 +4,10 @@ Azure AI Anthropic CountTokens API handler.
|
|||
Uses httpx for HTTP requests with Azure authentication.
|
||||
"""
|
||||
|
||||
from typing import Any, Final
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -28,14 +29,14 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
|
|||
async def handle_count_tokens_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
messages: list[dict[str, JsonValue]],
|
||||
api_key: str,
|
||||
api_base: str,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tools: list[dict[str, JsonValue]] | None = None,
|
||||
system: object = None,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Handle a CountTokens request using httpx with Azure authentication.
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ Azure AI Anthropic Token Counter implementation using the CountTokens API.
|
|||
import os
|
||||
from typing import Any, Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.azure_ai.anthropic.count_tokens.handler import (
|
||||
AzureAIAnthropicCountTokensHandler,
|
||||
|
|
@ -28,11 +30,11 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
|
|||
async def count_tokens(
|
||||
self,
|
||||
model_to_use: str,
|
||||
messages: list[dict[str, object]] | None,
|
||||
messages: list[dict[str, JsonValue]] | None,
|
||||
contents: list[dict[str, object]] | None,
|
||||
deployment: dict[str, Any] | None = None,
|
||||
request_model: str = "",
|
||||
tools: list[dict[str, object]] | None = None,
|
||||
tools: list[dict[str, JsonValue]] | None = None,
|
||||
system: object | None = None,
|
||||
) -> TokenCountResponse | None:
|
||||
"""
|
||||
|
|
@ -86,12 +88,14 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
|
|||
)
|
||||
|
||||
if result is not None:
|
||||
return TokenCountResponse(
|
||||
total_tokens=result.get("input_tokens", 0),
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="azure_ai_anthropic_api",
|
||||
original_response=result,
|
||||
return TokenCountResponse.model_validate(
|
||||
{
|
||||
"total_tokens": result.get("input_tokens", 0),
|
||||
"request_model": request_model,
|
||||
"model_used": model_to_use,
|
||||
"tokenizer_type": "azure_ai_anthropic_api",
|
||||
"original_response": result,
|
||||
}
|
||||
)
|
||||
except AnthropicError as e:
|
||||
verbose_logger.warning(
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import httpx
|
|||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
from litellm.llms.anthropic.chat.transformation import JsonModeRequest
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
|
|
@ -162,6 +163,9 @@ class AzureAnthropicChatCompletion(AnthropicChatCompletion):
|
|||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
json_mode=json_mode,
|
||||
response_format_schema=(
|
||||
JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode else None
|
||||
),
|
||||
)
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
process_anthropic_headers,
|
||||
|
|
|
|||
|
|
@ -113,7 +113,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
for index, reference_image in enumerate(images, start=1)
|
||||
}
|
||||
)
|
||||
request_body: Final[dict[str, Any]] = {
|
||||
request_body: Final[dict[str, object]] = {
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
**reference_images,
|
||||
|
|
|
|||
|
|
@ -3,13 +3,15 @@ Utility functions for base LLM classes.
|
|||
"""
|
||||
|
||||
import copy
|
||||
import itertools
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Final
|
||||
|
||||
from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # SDK parser internals
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
|
||||
|
|
@ -157,6 +159,136 @@ def convert_tool_response_to_message(
|
|||
_convert_tool_response_to_message = convert_tool_response_to_message
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ResponseFormatSchema:
|
||||
property_names: frozenset[str]
|
||||
required_names: frozenset[str] = frozenset()
|
||||
allows_undeclared_keys: bool = False
|
||||
|
||||
@property
|
||||
def accepts_a_lone_undeclared_key(self) -> bool:
|
||||
return self.allows_undeclared_keys and not self.required_names
|
||||
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, JsonValue])
|
||||
_NAMES: Final = TypeAdapter(frozenset[str])
|
||||
|
||||
|
||||
def as_json_object(value: object) -> Mapping[str, JsonValue] | None:
|
||||
try:
|
||||
return _JSON_OBJECT.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _names(value: object) -> frozenset[str]:
|
||||
try:
|
||||
return _NAMES.validate_python(value)
|
||||
except ValidationError:
|
||||
return frozenset()
|
||||
|
||||
|
||||
def _allows_undeclared_keys(document: Mapping[str, JsonValue]) -> bool:
|
||||
additional_properties: Final = document.get("additionalProperties")
|
||||
return (
|
||||
additional_properties is True
|
||||
or as_json_object(additional_properties) is not None
|
||||
or bool(as_json_object(document.get("patternProperties")))
|
||||
)
|
||||
|
||||
|
||||
def response_format_schema(value: object) -> ResponseFormatSchema | None:
|
||||
document: Final = as_json_object(value)
|
||||
if document is None:
|
||||
return None
|
||||
properties: Final = as_json_object(document.get("properties", {}))
|
||||
if properties is None:
|
||||
return None
|
||||
return ResponseFormatSchema(
|
||||
property_names=frozenset(properties),
|
||||
required_names=_names(document.get("required", ())),
|
||||
allows_undeclared_keys=_allows_undeclared_keys(document),
|
||||
)
|
||||
|
||||
|
||||
RESPONSE_FORMAT_ENVELOPE_NAME_KEYS: Final = frozenset({"name", "tool_name"})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _EnvelopeLayer:
|
||||
inner: Mapping[str, JsonValue]
|
||||
names_the_tool: bool
|
||||
|
||||
|
||||
def _json_object_or_encoded_object(value: JsonValue) -> Mapping[str, JsonValue] | None:
|
||||
if not isinstance(value, str):
|
||||
return as_json_object(value)
|
||||
try:
|
||||
return as_json_object(json.loads(value))
|
||||
except (json.JSONDecodeError, RecursionError):
|
||||
return None
|
||||
|
||||
|
||||
def _label_keys(payload_object: Mapping[str, JsonValue], schema: ResponseFormatSchema) -> frozenset[str]:
|
||||
if not schema.required_names:
|
||||
return frozenset()
|
||||
declared: Final = schema.property_names | schema.required_names
|
||||
return frozenset(key for key, value in payload_object.items() if isinstance(value, str) and key not in declared)
|
||||
|
||||
|
||||
def _envelope_layer(
|
||||
payload_object: Mapping[str, JsonValue], schema: ResponseFormatSchema, tool_name: str
|
||||
) -> _EnvelopeLayer | None:
|
||||
name_keys: Final = frozenset(
|
||||
key for key in RESPONSE_FORMAT_ENVELOPE_NAME_KEYS if payload_object.get(key) == tool_name
|
||||
)
|
||||
label_keys: Final = name_keys or _label_keys(payload_object, schema)
|
||||
wrapper_keys: Final = tuple(key for key in payload_object if key not in label_keys)
|
||||
if len(wrapper_keys) != 1 or wrapper_keys[0] in schema.property_names:
|
||||
return None
|
||||
wrapped: Final = payload_object[wrapper_keys[0]]
|
||||
inner: Final = _json_object_or_encoded_object(wrapped) if name_keys else as_json_object(wrapped)
|
||||
return None if inner is None else _EnvelopeLayer(inner=inner, names_the_tool=bool(name_keys))
|
||||
|
||||
|
||||
def _envelope_layers(
|
||||
payload_object: Mapping[str, JsonValue], schema: ResponseFormatSchema, tool_name: str
|
||||
) -> Iterator[_EnvelopeLayer]:
|
||||
layer = _envelope_layer(payload_object, schema, tool_name) # rebind-ok: the recursion gate needs a loop
|
||||
while layer is not None:
|
||||
yield layer
|
||||
layer = _envelope_layer(layer.inner, schema, tool_name)
|
||||
|
||||
|
||||
def _schema_object_inside_envelope(
|
||||
payload_object: Mapping[str, JsonValue], schema: ResponseFormatSchema, tool_name: str
|
||||
) -> Mapping[str, JsonValue] | None:
|
||||
layers: Final = _envelope_layers(payload_object, schema, tool_name)
|
||||
candidates: Final = layers if schema.property_names else itertools.islice(layers, 1)
|
||||
return next((layer.inner for layer in candidates if _holds_the_schema_object(layer, schema)), None)
|
||||
|
||||
|
||||
def _holds_the_schema_object(layer: _EnvelopeLayer, schema: ResponseFormatSchema) -> bool:
|
||||
if layer.names_the_tool:
|
||||
return True
|
||||
return not schema.accepts_a_lone_undeclared_key and not schema.property_names.isdisjoint(layer.inner)
|
||||
|
||||
|
||||
def response_format_envelope_contents(payload: object, schema: ResponseFormatSchema, tool_name: str) -> object | None:
|
||||
payload_object: Final = as_json_object(payload)
|
||||
if payload_object is None:
|
||||
return None
|
||||
schema_object: Final = _schema_object_inside_envelope(payload_object, schema, tool_name)
|
||||
if schema_object is not None:
|
||||
return schema_object
|
||||
return None if "values" in schema.property_names else payload_object.get("values")
|
||||
|
||||
|
||||
def unwrap_response_format_envelope(payload: object, schema: ResponseFormatSchema, tool_name: str) -> object:
|
||||
contents: Final = response_format_envelope_contents(payload, schema, tool_name)
|
||||
return payload if contents is None else contents
|
||||
|
||||
|
||||
def _dict_to_response_format_helper(response_format: dict, ref_template: str | None = None) -> dict:
|
||||
if ref_template is not None and response_format.get("type") == "json_schema":
|
||||
# Deep copy to avoid modifying original
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ they created. Callers with no admin role and no identifying ids at all
|
|||
are denied so an empty user_id can never select an unscoped query.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm.types.proxy.auth.user_api_key_auth import (
|
||||
|
|
@ -46,7 +47,7 @@ def resolve_resource_owner_id(
|
|||
return None
|
||||
|
||||
|
||||
def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]:
|
||||
def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, object]:
|
||||
"""Build the OpenAI-style paginated list response shape used by managed
|
||||
file/batch/vector-store listings. ``first_id`` and ``last_id`` are
|
||||
sourced from each item's ``.id`` attribute."""
|
||||
|
|
@ -61,7 +62,7 @@ def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]:
|
|||
|
||||
def build_owner_filter(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> dict[str, Any] | None:
|
||||
) -> dict[str, str | Sequence[Mapping[str, str]]] | None:
|
||||
"""Return a Prisma `where` fragment that scopes a managed-resource listing
|
||||
to records the caller is allowed to see.
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,10 @@
|
|||
import json
|
||||
import types
|
||||
from collections.abc import AsyncIterator, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm import verbose_logger
|
||||
|
|
@ -16,6 +17,8 @@ from litellm.litellm_core_utils.optional_imports import ensure_optional_import
|
|||
from litellm.llms.anthropic.chat.handler import (
|
||||
ModelResponseIterator as AnthropicModelResponseIterator,
|
||||
)
|
||||
from litellm.llms.anthropic.chat.transformation import JsonModeRequest
|
||||
from litellm.llms.base_llm.base_utils import ResponseFormatSchema
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
|
|
@ -174,15 +177,34 @@ class AmazonCohereChatConfig:
|
|||
return optional_params
|
||||
|
||||
|
||||
_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def response_format_schema_from_request(data: str, json_mode: bool | None) -> ResponseFormatSchema | None:
|
||||
if not json_mode:
|
||||
return None
|
||||
try:
|
||||
body: Final = _REQUEST_BODY.validate_python(json.loads(data))
|
||||
except (json.JSONDecodeError, ValidationError):
|
||||
return None
|
||||
return JsonModeRequest.model_validate(body).tool_schema()
|
||||
|
||||
|
||||
def _stream_decoder(
|
||||
bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None,
|
||||
*,
|
||||
model: str,
|
||||
json_mode: bool | None,
|
||||
sync_stream: bool,
|
||||
response_format_schema: ResponseFormatSchema | None = None,
|
||||
) -> "AWSEventStreamDecoder":
|
||||
if bedrock_invoke_provider == "anthropic":
|
||||
return AmazonAnthropicClaudeStreamDecoder(model=model, sync_stream=sync_stream, json_mode=json_mode)
|
||||
return AmazonAnthropicClaudeStreamDecoder(
|
||||
model=model,
|
||||
sync_stream=sync_stream,
|
||||
json_mode=json_mode,
|
||||
response_format_schema=response_format_schema,
|
||||
)
|
||||
if bedrock_invoke_provider == "deepseek_r1":
|
||||
return AmazonDeepSeekR1StreamDecoder(model=model, sync_stream=sync_stream)
|
||||
if bedrock_invoke_provider == "moonshot":
|
||||
|
|
@ -250,7 +272,11 @@ async def make_call(
|
|||
)
|
||||
else:
|
||||
decoder: Final = _stream_decoder(
|
||||
bedrock_invoke_provider, model=model, json_mode=json_mode, sync_stream=False
|
||||
bedrock_invoke_provider,
|
||||
model=model,
|
||||
json_mode=json_mode,
|
||||
sync_stream=False,
|
||||
response_format_schema=response_format_schema_from_request(data, json_mode),
|
||||
)
|
||||
completion_stream = decoder.aiter_bytes(
|
||||
response.aiter_bytes(chunk_size=stream_chunk_size), response_headers=response.headers
|
||||
|
|
@ -341,7 +367,11 @@ def make_sync_call(
|
|||
)
|
||||
else:
|
||||
decoder: Final = _stream_decoder(
|
||||
bedrock_invoke_provider, model=model, json_mode=json_mode, sync_stream=True
|
||||
bedrock_invoke_provider,
|
||||
model=model,
|
||||
json_mode=json_mode,
|
||||
sync_stream=True,
|
||||
response_format_schema=response_format_schema_from_request(data, json_mode),
|
||||
)
|
||||
completion_stream = decoder.iter_bytes(
|
||||
response.iter_bytes(chunk_size=stream_chunk_size), response_headers=response.headers
|
||||
|
|
@ -878,6 +908,7 @@ class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder):
|
|||
model: str,
|
||||
sync_stream: bool,
|
||||
json_mode: bool | None = None,
|
||||
response_format_schema: ResponseFormatSchema | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Child class of AWSEventStreamDecoder that handles the streaming response from the Anthropic family of models
|
||||
|
|
@ -889,6 +920,7 @@ class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder):
|
|||
streaming_response=None,
|
||||
sync_stream=sync_stream,
|
||||
json_mode=json_mode,
|
||||
response_format_schema=response_format_schema,
|
||||
)
|
||||
|
||||
def _chunk_parser(self, chunk_data: dict) -> ModelResponseStream:
|
||||
|
|
|
|||
|
|
@ -7,9 +7,11 @@ The main difference is in the response format: Qwen2 uses "text" field while Qwe
|
|||
Qwen2 + Invoke API Tutorial: https://docs.aws.amazon.com/bedrock/latest/userguide/invoke-imported-model.html
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import (
|
||||
AmazonQwen3Config,
|
||||
|
|
@ -23,6 +25,9 @@ from litellm.types.utils import ModelResponse, Usage
|
|||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_TEXT: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True))
|
||||
|
||||
|
||||
class AmazonQwen2Config(AmazonQwen3Config):
|
||||
"""
|
||||
|
|
@ -54,10 +59,10 @@ class AmazonQwen2Config(AmazonQwen3Config):
|
|||
Qwen2 uses "text" field, but we also support "generation" field for compatibility.
|
||||
"""
|
||||
try:
|
||||
response_data: Final = raw_response.json()
|
||||
response_data: Final = _JSON_OBJECT.validate_python(raw_response.json())
|
||||
|
||||
# Extract the generated text - Qwen2 uses "text" field, but also support "generation" for compatibility
|
||||
generated_text = response_data.get("generation", "") or response_data.get("text", "")
|
||||
generated_text = _TEXT.validate_python(response_data.get("generation", "") or response_data.get("text", ""))
|
||||
|
||||
# Clean up the response (remove assistant start token if present)
|
||||
generated_text = generated_text.removeprefix("<|im_start|>assistant\n")
|
||||
|
|
@ -71,14 +76,16 @@ class AmazonQwen2Config(AmazonQwen3Config):
|
|||
|
||||
# Set usage information if available in response
|
||||
if "usage" in response_data:
|
||||
usage_data: Final = response_data["usage"]
|
||||
usage_data: Final = _JSON_OBJECT.validate_python(response_data["usage"])
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
Usage(
|
||||
prompt_tokens=usage_data.get("prompt_tokens", 0),
|
||||
completion_tokens=usage_data.get("completion_tokens", 0),
|
||||
total_tokens=usage_data.get("total_tokens", 0),
|
||||
Usage.model_validate(
|
||||
{
|
||||
"prompt_tokens": usage_data.get("prompt_tokens", 0),
|
||||
"completion_tokens": usage_data.get("completion_tokens", 0),
|
||||
"total_tokens": usage_data.get("total_tokens", 0),
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from collections.abc import Mapping
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
|
||||
|
|
@ -22,6 +23,9 @@ from litellm.types.utils import ModelResponse, Usage
|
|||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_TEXT: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True))
|
||||
|
||||
|
||||
class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig):
|
||||
"""
|
||||
|
|
@ -179,10 +183,10 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig):
|
|||
Transform Qwen3 Bedrock response to OpenAI format
|
||||
"""
|
||||
try:
|
||||
response_data: Final = raw_response.json()
|
||||
response_data: Final = _JSON_OBJECT.validate_python(raw_response.json())
|
||||
|
||||
# Extract the generated text - Qwen3 uses "generation" field
|
||||
generated_text = response_data.get("generation", "")
|
||||
generated_text = _TEXT.validate_python(response_data.get("generation", ""))
|
||||
|
||||
# Clean up the response (remove assistant start token if present)
|
||||
generated_text = generated_text.removeprefix("<|im_start|>assistant\n")
|
||||
|
|
@ -196,14 +200,16 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig):
|
|||
|
||||
# Set usage information if available in response
|
||||
if "usage" in response_data:
|
||||
usage_data: Final = response_data["usage"]
|
||||
usage_data: Final = _JSON_OBJECT.validate_python(response_data["usage"])
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
Usage(
|
||||
prompt_tokens=usage_data.get("prompt_tokens", 0),
|
||||
completion_tokens=usage_data.get("completion_tokens", 0),
|
||||
total_tokens=usage_data.get("total_tokens", 0),
|
||||
Usage.model_validate(
|
||||
{
|
||||
"prompt_tokens": usage_data.get("prompt_tokens", 0),
|
||||
"completion_tokens": usage_data.get("completion_tokens", 0),
|
||||
"total_tokens": usage_data.get("total_tokens", 0),
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -51,6 +51,10 @@ _GUARDRAIL_CONFIG_EXPECTED_FORMAT: Final = (
|
|||
"{'guardrailIdentifier': str, 'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}"
|
||||
)
|
||||
|
||||
_INVOKE_PROVIDERS: Final = TypeAdapter(tuple[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, ...]).validate_python(
|
||||
get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL)
|
||||
)
|
||||
|
||||
|
||||
def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]":
|
||||
try:
|
||||
|
|
@ -568,7 +572,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
|
|||
if provider is not None:
|
||||
return provider
|
||||
|
||||
for provider in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL):
|
||||
for provider in _INVOKE_PROVIDERS:
|
||||
if provider in model:
|
||||
return provider
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -6,15 +6,20 @@ Documentation: https://api-dashboard.search.brave.com/app/documentation/web-sear
|
|||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Iterable, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final, Literal, TypedDict
|
||||
|
||||
import httpx
|
||||
from dateutil import parser
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
_ISO_YMD: Final = re.compile(r"^\s*\d{4}[-/]\d{1,2}[-/]\d{1,2}\s*$")
|
||||
_UNIX_TIMESTAMP: Final = re.compile(r"^\s*-?\d+(\.\d+)?\s*$")
|
||||
BRAVE_SECTIONS: Final = ["web", "discussions", "faqs", "faq", "news", "videos"]
|
||||
_JSON_VALUE: Final = TypeAdapter(object)
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.search.transformation import (
|
||||
|
|
@ -26,7 +31,7 @@ from litellm.secret_managers.main import get_secret_str
|
|||
|
||||
|
||||
def to_yyyy_mm_dd(
|
||||
s: str | float | None,
|
||||
s: object,
|
||||
*,
|
||||
dayfirst: bool = False,
|
||||
yearfirst: bool = False,
|
||||
|
|
@ -246,7 +251,7 @@ class BraveSearchConfig(BaseSearchConfig):
|
|||
"""
|
||||
Transform Brave Search API response to LiteLLM unified SearchResponse format.
|
||||
"""
|
||||
response_json: Final = raw_response.json()
|
||||
response_body: Final = _JSON_VALUE.validate_python(raw_response.json())
|
||||
|
||||
# Transform results to SearchResult objects
|
||||
results: Final[list[SearchResult]] = []
|
||||
|
|
@ -254,9 +259,11 @@ class BraveSearchConfig(BaseSearchConfig):
|
|||
query_params: Final = raw_response.request.url.params if raw_response.request else {}
|
||||
sections_to_process: Final = self._sections_from_params(dict(query_params))
|
||||
max_results: Final = max(1, min(int(query_params.get("count", 20)), 20))
|
||||
response_json: Final = _JSON_OBJECT.validate_python(response_body)
|
||||
|
||||
for section in sections_to_process:
|
||||
for result in response_json.get(section, {}).get("results", []):
|
||||
section_json = _JSON_OBJECT.validate_python(response_json.get(section, {}))
|
||||
for raw_result in _JSON_VALUES.validate_python(section_json.get("results", [])):
|
||||
# Because the `max_results`/`count` parameters do not affect
|
||||
# the number of "discussion", "faq", "news", or "videos"
|
||||
# results, we need to manually limit the number of results
|
||||
|
|
@ -264,18 +271,21 @@ class BraveSearchConfig(BaseSearchConfig):
|
|||
if len(results) >= max_results:
|
||||
break
|
||||
|
||||
result = _JSON_OBJECT.validate_python(raw_result)
|
||||
title = result.get("title", "")
|
||||
url = result.get("url", "")
|
||||
snippet = result.get("description", "")
|
||||
date = to_yyyy_mm_dd(result.get("page_age") or result.get("age"))
|
||||
last_updated = to_yyyy_mm_dd(result.get("fetched_content_timestamp", ""))
|
||||
|
||||
search_result = SearchResult(
|
||||
title=title,
|
||||
url=url,
|
||||
snippet=snippet,
|
||||
date=date,
|
||||
last_updated=last_updated,
|
||||
search_result = SearchResult.model_validate(
|
||||
{
|
||||
"title": title,
|
||||
"url": url,
|
||||
"snippet": snippet,
|
||||
"date": date,
|
||||
"last_updated": last_updated,
|
||||
}
|
||||
)
|
||||
|
||||
results.append(search_result)
|
||||
|
|
|
|||
|
|
@ -10,10 +10,11 @@ Convers
|
|||
Docs - https://docs.cohere.com/v2/reference/embed
|
||||
"""
|
||||
|
||||
from collections.abc import Sized
|
||||
from collections.abc import Iterable, Mapping, Sized
|
||||
from typing import Final, Protocol, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm import COHERE_DEFAULT_EMBEDDING_INPUT_TYPE
|
||||
|
|
@ -30,6 +31,9 @@ from litellm.utils import is_base64_encoded
|
|||
|
||||
from ..common_utils import CohereError
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_EMBEDDINGS_BY_TYPE: Final = TypeAdapter(Mapping[str, Iterable[object]], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
|
||||
class _SupportsEncode(Protocol):
|
||||
"""Tokenizer handle: the embedding usage path only encodes text to measure its token length."""
|
||||
|
|
@ -198,7 +202,8 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig):
|
|||
'usage'
|
||||
}
|
||||
"""
|
||||
embeddings: Final = response_json["embeddings"]
|
||||
payload: Final = _JSON_OBJECT.validate_python(response_json)
|
||||
embeddings: Final = _EMBEDDINGS_BY_TYPE.validate_python(payload["embeddings"])
|
||||
output_data: Final = []
|
||||
for k, embedding_list in embeddings.items():
|
||||
for idx, embedding in enumerate(embedding_list):
|
||||
|
|
@ -213,7 +218,7 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig):
|
|||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
self._calculate_usage(input, encoding, response_json.get("meta", {})),
|
||||
self._calculate_usage(input, encoding, dict(_JSON_OBJECT.validate_python(payload.get("meta", {})))),
|
||||
)
|
||||
|
||||
return model_response
|
||||
|
|
|
|||
|
|
@ -2,9 +2,11 @@
|
|||
CometAPI Embedding API support - OpenAI compatible
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
|
@ -15,6 +17,9 @@ from litellm.types.utils import EmbeddingResponse, Usage
|
|||
|
||||
from ..common_utils import CometAPIException
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_INT: Final = TypeAdapter(int)
|
||||
|
||||
|
||||
class CometAPIEmbeddingConfig(BaseEmbeddingConfig):
|
||||
"""
|
||||
|
|
@ -136,9 +141,10 @@ class CometAPIEmbeddingConfig(BaseEmbeddingConfig):
|
|||
model_response.data = raw_response_json.get("data")
|
||||
model_response.object = raw_response_json.get("object")
|
||||
|
||||
usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {}))
|
||||
usage: Final = Usage(
|
||||
prompt_tokens=raw_response_json.get("usage", {}).get("prompt_tokens", 0),
|
||||
total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
|
||||
prompt_tokens=_INT.validate_python(usage_payload.get("prompt_tokens", 0) or 0),
|
||||
total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0),
|
||||
)
|
||||
|
||||
model_response.usage = usage
|
||||
|
|
|
|||
|
|
@ -147,7 +147,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
|
|||
)
|
||||
|
||||
if "error" in response_json:
|
||||
error: Final = response_json["error"]
|
||||
error: Final[object] = response_json["error"]
|
||||
message: Final = error.get("message", str(error)) if isinstance(error, dict) else str(error)
|
||||
raise DashScopeError(
|
||||
status_code=raw_response.status_code,
|
||||
|
|
|
|||
|
|
@ -285,6 +285,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
"tool_choice",
|
||||
"reasoning_effort",
|
||||
"thinking",
|
||||
"service_tier",
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -474,7 +475,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
# Move message-level cache_control into a content block when content is a string.
|
||||
if "cache_control" in _message and isinstance(_message.get("content"), str):
|
||||
_message = self._move_cache_control_into_string_content_block(_message)
|
||||
_sanitize_empty_content(cast(dict[str, Any], _message))
|
||||
_sanitize_empty_content(cast(dict[str, object], _message))
|
||||
if _is_bare_assistant_message(_message):
|
||||
continue
|
||||
new_messages.append(_message)
|
||||
|
|
@ -503,7 +504,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
"""
|
||||
content: Final = message.get("content")
|
||||
# Create new message with cache_control moved into content block
|
||||
transformed_message: Final = cast(dict[str, Any], message.copy())
|
||||
transformed_message: Final = cast(dict[str, object], message.copy())
|
||||
cache_control: Final = transformed_message.pop("cache_control")
|
||||
transformed_message["content"] = [
|
||||
{
|
||||
|
|
@ -696,7 +697,14 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
model_response.model = f"{_custom_llm_provider}/{_response_model}"
|
||||
model_response.id = completion_response["id"]
|
||||
model_response.created = completion_response["created"]
|
||||
if "system_fingerprint" in completion_response:
|
||||
model_response.system_fingerprint = ( # rebind-ok: fills the caller's model_response like id above
|
||||
completion_response["system_fingerprint"]
|
||||
)
|
||||
setattr(model_response, "usage", Usage(**completion_response["usage"]))
|
||||
served_tier: Final = completion_response.get("service_tier")
|
||||
if isinstance(served_tier, str) and served_tier:
|
||||
setattr(model_response, "service_tier", served_tier) # noqa: B010 # no declared field
|
||||
|
||||
model_response.choices = self._transform_dbrx_choices(
|
||||
choices=completion_response["choices"],
|
||||
|
|
|
|||
|
|
@ -4,9 +4,11 @@ Calls DataForSEO SERP API to search the web.
|
|||
DataForSEO API Reference: https://docs.dataforseo.com/v3/serp/google/organic/live/advanced/?bash
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.constants import DEFAULT_DATAFORSEO_LOCATION_CODE
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -17,6 +19,9 @@ from litellm.llms.base_llm.search.transformation import (
|
|||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
|
||||
class DataForSEOSearchConfig(BaseSearchConfig):
|
||||
"""
|
||||
|
|
@ -201,18 +206,21 @@ class DataForSEOSearchConfig(BaseSearchConfig):
|
|||
if task.get("status_code") == 20000 and "result" in task:
|
||||
# Result is an array, take first element
|
||||
if len(task["result"]) > 0:
|
||||
result: Final = task["result"][0]
|
||||
result: Final = _JSON_OBJECT.validate_python(task["result"][0])
|
||||
|
||||
# Items contain the actual search results
|
||||
for item in result.get("items", []):
|
||||
for raw_item in _JSON_VALUES.validate_python(result.get("items", [])):
|
||||
item = _JSON_OBJECT.validate_python(raw_item)
|
||||
# Only process organic search results
|
||||
if item.get("type") == "organic":
|
||||
search_result = SearchResult(
|
||||
title=item.get("title", ""),
|
||||
url=item.get("url", ""),
|
||||
snippet=item.get("description", ""),
|
||||
date=None, # DataForSEO doesn't provide date in standard response
|
||||
last_updated=None,
|
||||
search_result = SearchResult.model_validate(
|
||||
{
|
||||
"title": item.get("title", ""),
|
||||
"url": item.get("url", ""),
|
||||
"snippet": item.get("description", ""),
|
||||
"date": None, # DataForSEO doesn't provide date in standard response
|
||||
"last_updated": None,
|
||||
}
|
||||
)
|
||||
results.append(search_result)
|
||||
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ import json
|
|||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol
|
||||
|
||||
from litellm.harness.options import DeepAgentsOptions
|
||||
from litellm.harness.types import (
|
||||
|
|
@ -73,7 +73,7 @@ _APPROVAL_DECISIONS: Final = ("approve", "reject")
|
|||
|
||||
def chat_model_kwargs(
|
||||
ctx: SessionContext,
|
||||
) -> dict[str, Any]: # mutable-ok: ChatLiteLLM constructor kwargs, splatted as **kwargs
|
||||
) -> dict[str, object]: # mutable-ok: ChatLiteLLM constructor kwargs, splatted as **kwargs
|
||||
"""ChatLiteLLM constructor kwargs for gateway or SDK mode."""
|
||||
if not ctx.model:
|
||||
raise ValueError("Harness.DEEPAGENTS needs model=")
|
||||
|
|
@ -107,7 +107,7 @@ def blocked_tools(permissions: str, disable_tools: Sequence[str]) -> frozenset[s
|
|||
|
||||
def interrupt_config(
|
||||
permissions: str, blocked: frozenset[str]
|
||||
) -> dict[str, Any] | None: # mutable-ok: deepagents create_deep_agent(interrupt_on=) takes a dict
|
||||
) -> dict[str, dict[str, list[str]]] | None: # mutable-ok: deepagents create_deep_agent(interrupt_on=) takes a dict
|
||||
"""interrupt_on for permissions='ask': approve/reject every mutating built-in."""
|
||||
if permissions != "ask":
|
||||
return None
|
||||
|
|
@ -221,9 +221,14 @@ def _message_events(message: object, skip_tools: frozenset[str]) -> tuple[Event,
|
|||
return ()
|
||||
|
||||
|
||||
class InterruptLike(Protocol):
|
||||
@property
|
||||
def id(self) -> str: ...
|
||||
|
||||
|
||||
def interrupts_in(
|
||||
update: object,
|
||||
) -> list[Any]: # mutable-ok: returns a list; existing callers/tests compare it to list literals
|
||||
) -> list[InterruptLike]: # mutable-ok: returns a list; existing callers/tests compare it to list literals
|
||||
if not isinstance(update, Mapping):
|
||||
return [] # mutable-ok: returns a list; existing callers/tests compare it to list literals
|
||||
found = update.get("__interrupt__")
|
||||
|
|
@ -260,7 +265,7 @@ def approval_requests(
|
|||
return list(kept) # mutable-ok: list return; callers/tests compare to lists
|
||||
|
||||
|
||||
def decision(allowed: bool, reason: str) -> dict[str, Any]: # mutable-ok: LangGraph resume payload (HITL decision dict)
|
||||
def decision(allowed: bool, reason: str) -> dict[str, str]: # mutable-ok: LangGraph resume payload (HITL decision dict)
|
||||
if allowed:
|
||||
return {"type": "approve"} # mutable-ok: LangGraph resume payload (HITL decision dict)
|
||||
return { # mutable-ok: LangGraph HITL decision
|
||||
|
|
@ -273,7 +278,7 @@ def decision(allowed: bool, reason: str) -> dict[str, Any]: # mutable-ok: LangG
|
|||
class TurnState:
|
||||
"""Mutable state across the stream passes of one turn."""
|
||||
|
||||
interrupts: tuple[Any, ...] = ()
|
||||
interrupts: tuple[InterruptLike, ...] = ()
|
||||
|
||||
|
||||
class DeepAgentsHarnessConfig(BaseHarnessConfig):
|
||||
|
|
|
|||
|
|
@ -4,10 +4,12 @@ Calls DuckDuckGo's Instant Answer API to search the web.
|
|||
DuckDuckGo API Reference: https://duckduckgo.com/api
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import Final, Literal, TypedDict
|
||||
from urllib.parse import urlencode
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.search.transformation import (
|
||||
|
|
@ -17,6 +19,9 @@ from litellm.llms.base_llm.search.transformation import (
|
|||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
|
||||
class _DuckDuckGoSearchRequestRequired(TypedDict):
|
||||
"""Required fields for DuckDuckGo Search API request."""
|
||||
|
|
@ -159,7 +164,7 @@ class DuckDuckGoSearchConfig(BaseSearchConfig):
|
|||
Returns:
|
||||
SearchResponse with standardized format
|
||||
"""
|
||||
response_json: Final = raw_response.json()
|
||||
response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
|
||||
|
||||
# Extract max_results from the request URL params
|
||||
query_params: Final = raw_response.request.url.params if raw_response.request else {}
|
||||
|
|
@ -178,17 +183,19 @@ class DuckDuckGoSearchConfig(BaseSearchConfig):
|
|||
|
||||
# Check if there's an Abstract with URL
|
||||
if response_json.get("AbstractURL") and response_json.get("AbstractText"):
|
||||
abstract_result: Final = SearchResult(
|
||||
title=response_json.get("Heading", ""),
|
||||
url=response_json.get("AbstractURL", ""),
|
||||
snippet=response_json.get("AbstractText", ""),
|
||||
date=None,
|
||||
last_updated=None,
|
||||
abstract_result: Final = SearchResult.model_validate(
|
||||
{
|
||||
"title": response_json.get("Heading", ""),
|
||||
"url": response_json.get("AbstractURL", ""),
|
||||
"snippet": response_json.get("AbstractText", ""),
|
||||
"date": None,
|
||||
"last_updated": None,
|
||||
}
|
||||
)
|
||||
results.append(abstract_result)
|
||||
|
||||
# Process RelatedTopics
|
||||
related_topics: Final = response_json.get("RelatedTopics", [])
|
||||
related_topics: Final = _JSON_VALUES.validate_python(response_json.get("RelatedTopics", []))
|
||||
for topic in related_topics:
|
||||
# Stop if we've reached max_results
|
||||
if max_results is not None and len(results) >= max_results:
|
||||
|
|
|
|||
|
|
@ -4,9 +4,11 @@ Calls Exa AI's /search endpoint to search the web.
|
|||
Exa AI API Reference: https://docs.exa.ai/reference/search
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import Final, TypedDict
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.search.transformation import (
|
||||
|
|
@ -16,6 +18,10 @@ from litellm.llms.base_llm.search.transformation import (
|
|||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_STRS: Final = TypeAdapter(Iterable[str], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
|
||||
class _ExaAISearchRequestRequired(TypedDict):
|
||||
"""Required fields for Exa AI Search API request."""
|
||||
|
|
@ -177,20 +183,23 @@ class ExaAISearchConfig(BaseSearchConfig):
|
|||
Returns:
|
||||
SearchResponse with standardized format
|
||||
"""
|
||||
response_json: Final = raw_response.json()
|
||||
response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
|
||||
|
||||
# Transform results to SearchResult objects
|
||||
results: Final = []
|
||||
for result in response_json.get("results", []):
|
||||
search_result = SearchResult(
|
||||
title=result.get("title", ""),
|
||||
url=result.get("url", ""),
|
||||
snippet=result.get("text")
|
||||
or "\n\n".join(result.get("highlights") or [])
|
||||
or result.get("summary")
|
||||
or "",
|
||||
date=result.get("publishedDate"), # ISO 8601 datetime string
|
||||
last_updated=None, # Exa AI doesn't provide last_updated in response
|
||||
for raw_result in _JSON_VALUES.validate_python(response_json.get("results", [])):
|
||||
result = _JSON_OBJECT.validate_python(raw_result)
|
||||
search_result = SearchResult.model_validate(
|
||||
{
|
||||
"title": result.get("title", ""),
|
||||
"url": result.get("url", ""),
|
||||
"snippet": result.get("text")
|
||||
or "\n\n".join(_STRS.validate_python(result.get("highlights") or []))
|
||||
or result.get("summary")
|
||||
or "",
|
||||
"date": result.get("publishedDate"), # ISO 8601 datetime string
|
||||
"last_updated": None, # Exa AI doesn't provide last_updated in response
|
||||
}
|
||||
)
|
||||
results.append(search_result)
|
||||
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue