From ea783cc35c91d1d6d989421e8cc045f23156feae Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:41:31 -0700 Subject: [PATCH] refactor(rust): make litellm-core the callable messages() SDK; drop the ai-gateway handler (#35044) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm-rust/ADDING_A_PROVIDER.md | 11 ++-- litellm-rust/AGENTS.md | 26 ++++++-- litellm-rust/CLAUDE.md | 61 +++++++++++++------ litellm-rust/README.md | 41 ++++++++----- .../PROVIDER_CODING_STANDARDS.md | 10 +-- litellm-rust/crates/ai-gateway/AGENTS.md | 9 ++- litellm-rust/crates/ai-gateway/README.md | 10 +-- .../crates/ai-gateway/src/constants.rs | 16 ----- .../crates/ai-gateway/src/io/messages.rs | 1 - litellm-rust/crates/ai-gateway/src/io/mod.rs | 1 - litellm-rust/crates/ai-gateway/src/lib.rs | 5 +- .../crates/ai-gateway/src/messages/mod.rs | 49 --------------- .../crates/ai-gateway/src/messages/types.rs | 24 -------- .../crates/ai-gateway/src/routes/AGENTS.md | 7 ++- .../ai-gateway/src/routes/messages/service.rs | 23 +++---- litellm-rust/crates/core/AGENTS.md | 8 ++- litellm-rust/crates/core/CLAUDE.md | 31 ++++++++-- litellm-rust/crates/core/Cargo.toml | 2 +- litellm-rust/crates/core/src/constants.rs | 16 +++++ .../src/messages/client.rs | 0 .../src/messages/common_utils.rs | 10 +-- .../src/messages/handler.rs | 17 ++---- litellm-rust/crates/core/src/messages/mod.rs | 30 +++++++++ .../src/messages/prepare.rs | 7 +-- .../src/messages/tests.rs | 14 +++-- .../crates/core/src/messages/types.rs | 24 ++++++++ litellm-rust/crates/python-bridge/AGENTS.md | 4 +- litellm-rust/crates/python-bridge/CLAUDE.md | 6 +- litellm-rust/crates/python-bridge/src/lib.rs | 18 ++++-- 29 files changed, 277 insertions(+), 204 deletions(-) delete mode 100644 litellm-rust/crates/ai-gateway/src/io/messages.rs delete mode 100644 litellm-rust/crates/ai-gateway/src/messages/mod.rs delete mode 100644 litellm-rust/crates/ai-gateway/src/messages/types.rs rename litellm-rust/crates/{ai-gateway => core}/src/messages/client.rs (100%) rename litellm-rust/crates/{ai-gateway => core}/src/messages/common_utils.rs (83%) rename litellm-rust/crates/{ai-gateway => core}/src/messages/handler.rs (84%) rename litellm-rust/crates/{ai-gateway => core}/src/messages/prepare.rs (92%) rename litellm-rust/crates/{ai-gateway => core}/src/messages/tests.rs (97%) diff --git a/litellm-rust/ADDING_A_PROVIDER.md b/litellm-rust/ADDING_A_PROVIDER.md index 5f933ec4fa8..857a744e014 100644 --- a/litellm-rust/ADDING_A_PROVIDER.md +++ b/litellm-rust/ADDING_A_PROVIDER.md @@ -1,10 +1,11 @@ # Adding a provider / route to litellm-rust -Three layers, same for every route (see `ocr` and `realtime` as references): +Everything for a route lives in `crates/core/src//`; `crates/core/src/messages` is the reference. A host (the axum gateway, the Python bridge) only calls the route's entrypoint. -1. **Transform contract (pure)** — `crates/core/src//transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) + types in `types.rs`. No network, env, or auth. -2. **Provider config (pure)** — `crates/providers/src///transformation.rs`: implement that trait as a `const __CONFIG`, mirroring the Python provider tree. Add parity unit tests. -3. **HTTP / transport (the host)** — `crates/providers/src/.rs` (e.g. `ocr.rs`, `realtime.rs`): the callable fn (`run_ocr`, `realtime`). It resolves the key, builds the auth header, builds URL + transforms via the config, then does the network call. This is the only layer allowed to do I/O. +1. **Entrypoint** — `mod.rs`: `pub async fn (request) -> CoreResult`, the Rust equivalent of `litellm.()`, plus a `_stream` variant when the route streams. It is the only thing a host touches. +2. **Transform contract** — `transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) with types in `types.rs`. +3. **Provider config** — `crates/core/src/providers///transformation.rs`: implement that trait as a `const __CONFIG`, mirroring the Python provider tree. Add parity unit tests. +4. **Prepare + handler** — `prepare.rs` resolves provider/model, credentials, auth headers, and URL, then transforms the request; `handler.rs` performs the provider call through the shared client in `client.rs` and transforms the response. ## Coding standards @@ -25,4 +26,4 @@ variants of it. The test for a good abstraction is that adding the next provider is a few declarative lines, not a new file of duplicated flow. Only diverge from the base when behavior is genuinely different, and say so explicitly in the PR. -**Calling:** the host invokes the route fn — the Python bridge calls `run_ocr`; the `ai-gateway` server calls `realtime`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`. +**Calling:** hosts invoke the core entrypoint — the Python bridge and the `ai-gateway` route service both call `litellm_core::messages::messages`. Never add a provider handler to `ai-gateway`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`. diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 398eec4685c..36a5ad5a8f4 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -4,14 +4,30 @@ litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes ( ## Crates -| Crate | Role | Pure / I/O | -|-------|------|------------| -| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under providers/), and the router. Builds requests/responses; no network. | Pure | -| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under io/) plus the axum server binary (behind the `server` feature). | I/O | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding | +| Crate | Role | +|-------|------| +| litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. | +| litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. | +| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +## Where a route lives + +A top-level LiteLLM call is a module under `crates/core/src//`, shaped like `messages`: + +``` +core/src/messages/ + mod.rs # pub async fn messages(..) -> CoreResult<..> (+ messages_stream for SSE) + types.rs # request/response types, MessagesRequest + transformation.rs # the provider template trait + prepare.rs # provider resolution, auth headers, URL + handler.rs # the provider call + client.rs # the shared reqwest client +``` + +Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched. + Adding a crate: default to a MODULE. New crate ONLY on a real trigger — separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these. Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional. diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md index 0659e63df39..fe6ceedbb86 100644 --- a/litellm-rust/CLAUDE.md +++ b/litellm-rust/CLAUDE.md @@ -23,21 +23,34 @@ the base when behavior is genuinely different, and say so explicitly in the PR. ## Crates (exactly three — see AGENTS.md) -`litellm-core` describes work; `litellm-ai-gateway` executes it; `litellm-python-bridge` -exposes it to the Python SDK. A crate is a **layer**, not a route — add modules, not crates. +`litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call. +`litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and +`litellm-python-bridge` exposes it to the Python SDK. A crate is a **layer**, not +a route — add modules, not crates. ## Core Boundary -`litellm-core` is the pure translation layer; the `litellm-ai-gateway` host executes work. +`litellm-core` owns the whole call. The Rust equivalent of `litellm.messages()` +is `litellm_core::messages::messages(request).await`: you call it, it does the +provider call, and you get a typed non-streaming response back. Route-level Rust structure mirrors LiteLLM's Python responsibilities: -- `core/src//` owns the route contract, shared types, and provider - template traits. For OCR, this means `core/src/ocr`. +- `core/src//` owns the route end to end: the public entrypoint fn named + after the route in `mod.rs`, the request/response types (`types.rs`), the + provider template trait (`transformation.rs`), the provider/auth/URL + resolution (`prepare.rs`), the HTTP client (`client.rs`), and the handler that + performs the call (`handler.rs`). `core/src/messages` is the reference. - `core/src/providers///transformation.rs` owns the - provider-specific transform. For Mistral OCR, this means - `core/src/providers/mistral/ocr/transformation.rs`. -- Network execution lives in the host crate `ai-gateway` (`ai-gateway/src/io/`), - never inside `core`. + provider-specific transform. For Anthropic Messages, this means + `core/src/providers/anthropic/messages/transformation.rs`. +- Handlers live in `core`, never in a host. `ai-gateway` must not contain a + route handler that talks to a provider; its axum route reads the HTTP request, + picks a deployment, and calls the `core` entrypoint. `python-bridge` marshals + Python objects and calls the same entrypoint. + +Streaming keeps the same shape: the route entrypoint has a `_stream` +variant in `core` that returns the upstream response so a host can splice it to +its own caller; the host still owns no provider logic. Call-hook and lifecycle instrumentation, including phase timing, usage accumulation, and callback payload construction, always lives in `core`. @@ -45,21 +58,31 @@ Hosts feed observed events into core and dispatch the completed payloads through their I/O logger; hosts must not own callback orchestration. Allowed in `core`: -- Pure request transforms -- Pure response transforms -- Pure stream chunk normalization +- The public entrypoint for a top-level LiteLLM call +- Request/response transforms and stream chunk normalization +- Provider resolution, auth header construction, and URL building +- The provider HTTP call itself, through a shared reused client with connect and + request timeouts - Shared data types and validation errors - Deterministic token/cost helper logic Not allowed in `core`: -- Network calls -- Environment variable or secret reads +- Serving HTTP: axum routes, extractors, and transport concerns stay in the host - Filesystem access -- Database or cache access -- Provider SDK signing or auth flows +- Database access +- Config file reading and rollout state - Logging callbacks, spend writes, or custom callbacks - Global mutable runtime state +Env reads in `core` are limited to credential fallback inside a route's +`prepare.rs` (the `env_lookup` closure), mirroring what the Python SDK does when +no key is passed. Everything else config-shaped is resolved by the host and +passed in. + +Routes still hosted in `ai-gateway` (`ocr`, `audio_transcription`, `realtime`) +predate this rule and are being moved into `core` route modules; do not add new +ones there, and prefer moving one when you touch it. + Python owns rollout state and fallback while Rust is being introduced. Rust paths must be off by default until parity tests prove equivalence with Python. A new provider/route may instead be implemented rust-only with no Python @@ -93,10 +116,10 @@ the first PR: - Preserve Python output shape intentionally. If a field is always serialized as `null` for Python parity, leave a short comment explaining that parity choice. -## Host I/O Rules +## Network I/O Rules -These rules apply when adding future crates or modules that execute network I/O, -such as `ai-gateway`, router hosts, or standalone servers: +These rules apply to every module that executes network I/O, whether it is a +`core` route handler or a host such as `ai-gateway`: - Set connect and full-request timeouts. No unbounded waits. - Reuse HTTP clients; do not construct clients per request. diff --git a/litellm-rust/README.md b/litellm-rust/README.md index 1646c90ad76..bcccf93300b 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -2,18 +2,31 @@ This workspace contains the staged Rust implementation for LiteLLM. -Rust starts as a pure transform core used by the existing Python host. Python -continues to own auth, configuration, network I/O, retries, routing, logging, +`litellm-core` is the LiteLLM SDK in Rust: one entrypoint per top-level call +that makes the LLM call and hands back a typed response, the same shape as +`litellm.messages()` in Python. + +```rust +let response = litellm_core::messages::messages(MessagesRequest { + model: "claude-sonnet-4-5", + body, + api_key: Some(key), + .. +}) +.await?; +``` + +Python continues to own configuration, retries, routing policy, logging, callbacks, spend tracking, and customer plugins until each Rust path has parity coverage and production evidence. ## Crates -| Crate | Role | Pure / I/O | -|-------|------|------------| -| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under providers/), and the router. Builds requests/responses; no network. | Pure | -| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under io/) plus the axum server binary (behind the `server` feature). | I/O | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding | +| Crate | Role | +|-------|------| +| litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. | +| litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | +| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. @@ -21,16 +34,16 @@ Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm- ```text crates/ - core/ Route contracts, shared pure types, errors, and templates. - src/ocr/ - providers/ Provider-specific pure transforms. - src/mistral/ocr/transformation.rs + core/ The SDK: route modules + provider transforms. + src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client + src/providers/anthropic/messages/transformation.rs + ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints. python-bridge/ PyO3 bridge for Python LiteLLM. ``` -The folder shape should follow the Python provider tree: -`providers/src///transformation.rs`. The bridge should expose -one function per top-level route, starting with `ocr(payload)`. +The folder shape follows the Python provider tree: +`core/src/providers///transformation.rs`. The bridge exposes one +function per top-level route, mirroring the core entrypoints. ## Checks diff --git a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md index ed44dc4c729..a1860d8a9c9 100644 --- a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md +++ b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md @@ -1,6 +1,6 @@ # Provider coding standards (litellm-rust) -Rules for adding or changing an LLM provider/route in `litellm-rust`. OCR (`MISTRAL_OCR_CONFIG`) is the reference; `messages` (`ANTHROPIC_MESSAGES_CONFIG`) is the next port. +Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` (`core/src/messages`, `ANTHROPIC_MESSAGES_CONFIG`) is the reference: a route is a `core` module with a public entrypoint that makes the call and returns a typed response. ## Provider resolution @@ -16,10 +16,10 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. OCR (`MIST ## Boundaries -7. Layers never cross: `core` = pure transforms/types (no network, env, secrets, auth, logging, global mutable state); `ai-gateway` = all I/O, auth headers, HTTP/SSE, lifecycle hooks; `python-bridge` = thin PyO3 adapter. +7. Layers never cross: `core` = the call itself (entrypoint, types, transforms, provider resolution, auth headers, provider HTTP, lifecycle hooks); `ai-gateway` = serving HTTP/WS (routing, extractors, auth of *our* callers, streaming to the client); `python-bridge` = thin PyO3 adapter. Hosts call the core entrypoint; they never build a provider request. 8. Generic/route files contain zero provider-specific branches. A provider is one module under `core/src/providers///`; a route is a module, never a new crate. -9. Route entry point stays thin: `()` -> `prepare_*` -> `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing. Handlers validate and delegate; no business logic in them. -10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Env reads happen only at the host/config layer, with the `DEFAULT_*` fallback defined in `constants.rs`. +9. Route entry point stays thin: `core::::()` -> `prepare_*` -> handler (or `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing). Axum handlers validate and delegate to a service that calls the entrypoint; no business logic in them. +10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Config-shaped env reads happen at the host/config layer with the `DEFAULT_*` fallback defined in `constants.rs`; the only env read in `core` is the credential fallback in a route's `prepare.rs`. ## Types and errors @@ -33,7 +33,7 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. OCR (`MIST 16. Never log request/response bodies, base64 payloads, document contents, or secrets. Truncate and bound any upstream body before it crosses a host boundary. 17. Treat empty/whitespace credentials, URLs, and config values as absent at the host resolution layer. -18. Host I/O sets connect + request timeouts (no unbounded waits), reuses a shared HTTP client, and prefers rustls TLS. +18. Network I/O sets connect + request timeouts (no unbounded waits), reuses a shared HTTP client, and prefers rustls TLS. ## Tests and rollout diff --git a/litellm-rust/crates/ai-gateway/AGENTS.md b/litellm-rust/crates/ai-gateway/AGENTS.md index d9e6e1adde5..92567091cd3 100644 --- a/litellm-rust/crates/ai-gateway/AGENTS.md +++ b/litellm-rust/crates/ai-gateway/AGENTS.md @@ -1,7 +1,9 @@ # ai-gateway — folder architecture The Axum server that fronts the Rust gateway. It owns transport + config + auth -only; deployment selection lives in `core::router`, transforms in `core`/`providers`. +only; deployment selection lives in `core::router`, and the LLM call itself +(transforms, auth headers, provider HTTP) lives behind a `core` route entrypoint +such as `litellm_core::messages::messages`. No provider handler lives here. ``` src/ @@ -32,6 +34,11 @@ src/ args; it runs during extraction. Never re-implement the check per route. - **Handlers are thin.** A handler validates and delegates to its `service`. No business logic, no provider calls, no transforms in handlers. +- **Services call `core`, they don't reimplement it.** A `service` picks the + deployment and calls the `core` route entrypoint. Provider resolution, auth + headers, URL building, and the HTTP call are `core`'s job; a service that + builds a provider request itself is a bug (`routes/messages/service.rs` is + the reference). - **State is shared and cheap to clone.** Long-lived handles live behind `Arc` in `state.rs`; read env/config only in `main.rs` when building state. diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index f913beff6d5..7a6c620ee84 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -8,11 +8,11 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame. `litellm-rust` is exactly three crates (a crate is a **layer**, not a route): -| Crate | Role | Pure / I/O | -|-------|------|------------| -| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under `providers/`), and the router. Builds requests/responses; no network. | Pure | -| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under `io/`) plus the Axum server binary (behind the `server` feature). | I/O | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding | +| Crate | Role | +|-------|------| +| litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. | +| litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | +| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. diff --git a/litellm-rust/crates/ai-gateway/src/constants.rs b/litellm-rust/crates/ai-gateway/src/constants.rs index 74808cf1ce6..78af374bf70 100644 --- a/litellm-rust/crates/ai-gateway/src/constants.rs +++ b/litellm-rust/crates/ai-gateway/src/constants.rs @@ -29,18 +29,6 @@ pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500; #[cfg(feature = "server")] pub(crate) const DEFAULT_PROVIDER: &str = "openai"; -/// Full-request timeout ceiling for Anthropic Messages provider calls, in -/// seconds. Mirrors the Python Anthropic Messages default. The per-request -/// timeout from `litellm_params` still overrides this on the request builder. -pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600; - -/// Connect timeout for Anthropic Messages provider calls, in seconds. -pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10; - -/// Max characters of an upstream error body echoed across the host boundary -/// before truncation, so provider bodies are bounded and data-minimized. -pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256; - pub(crate) const DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS: u64 = 10; pub(crate) const DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS: u64 = 300; @@ -48,10 +36,6 @@ pub(crate) const DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS: u64 = 300; #[cfg(feature = "server")] pub(crate) const MESSAGES_ROUTE_PATH: &str = "/v1/messages"; -/// Provider name used by the Anthropic Messages route when a deployment's -/// provider model does not carry an explicit provider prefix. -pub(crate) const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic"; - /// Request headers owned by the gateway and never forwarded upstream. #[cfg(feature = "server")] pub(crate) const MESSAGES_HEADERS_NOT_FORWARDED: &[&str] = diff --git a/litellm-rust/crates/ai-gateway/src/io/messages.rs b/litellm-rust/crates/ai-gateway/src/io/messages.rs deleted file mode 100644 index 86170e45678..00000000000 --- a/litellm-rust/crates/ai-gateway/src/io/messages.rs +++ /dev/null @@ -1 +0,0 @@ -pub use crate::messages::{MessagesRequest, messages}; diff --git a/litellm-rust/crates/ai-gateway/src/io/mod.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs index 6129a808965..cce56dd2121 100644 --- a/litellm-rust/crates/ai-gateway/src/io/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs @@ -1,5 +1,4 @@ pub mod audio_transcription; -pub mod messages; pub mod ocr; pub mod realtime; pub mod realtime_pool; diff --git a/litellm-rust/crates/ai-gateway/src/lib.rs b/litellm-rust/crates/ai-gateway/src/lib.rs index c44d661c29e..057db6457c4 100644 --- a/litellm-rust/crates/ai-gateway/src/lib.rs +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -4,7 +4,9 @@ //! without pulling in the HTTP server: //! //! - Call-type modules such as [`ocr`]: provider transforms, lifecycle hooks, -//! and provider I/O. Always available — no feature required. +//! and provider I/O. Always available — no feature required. These predate the +//! rule that a route's entrypoint and handler live in `litellm-core` (see +//! `litellm_core::messages`) and move there as they are touched. //! - [`io`]: compatibility exports and realtime WebSocket splice helpers. //! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling //! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway` @@ -14,7 +16,6 @@ pub mod audio_transcription; mod client; pub mod io; -pub mod messages; pub mod ocr; /// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and diff --git a/litellm-rust/crates/ai-gateway/src/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/messages/mod.rs deleted file mode 100644 index fd2dd546941..00000000000 --- a/litellm-rust/crates/ai-gateway/src/messages/mod.rs +++ /dev/null @@ -1,49 +0,0 @@ -use litellm_core::CoreResult; -use serde_json::Value; - -mod client; -mod common_utils; -mod handler; -mod prepare; -mod types; - -pub use types::MessagesRequest; - -use handler::{execute_messages_provider_call, execute_messages_provider_stream}; -use prepare::prepare_messages_call; - -pub async fn messages(request: MessagesRequest<'_>) -> CoreResult { - match execute_messages(request, false).await? { - MessagesResponse::Json(body) => Ok(body), - MessagesResponse::Stream(response) => { - drop(response); - Err(litellm_core::CoreError::InvalidResponse( - "non-streaming messages execution returned a stream".to_string(), - )) - } - } -} - -pub(crate) enum MessagesResponse { - Json(Value), - Stream(reqwest::Response), -} - -pub(crate) async fn execute_messages( - request: MessagesRequest<'_>, - stream: bool, -) -> CoreResult { - let prepared = prepare_messages_call(request)?; - if stream { - execute_messages_provider_stream(prepared) - .await - .map(MessagesResponse::Stream) - } else { - execute_messages_provider_call(prepared) - .await - .map(MessagesResponse::Json) - } -} - -#[cfg(test)] -mod tests; diff --git a/litellm-rust/crates/ai-gateway/src/messages/types.rs b/litellm-rust/crates/ai-gateway/src/messages/types.rs deleted file mode 100644 index 848fadb4b02..00000000000 --- a/litellm-rust/crates/ai-gateway/src/messages/types.rs +++ /dev/null @@ -1,24 +0,0 @@ -use std::time::Duration; - -use litellm_core::messages::transformation::AnthropicMessagesProviderConfig; -use serde_json::{Map, Value}; - -pub struct MessagesRequest<'a> { - pub model: &'a str, - pub body: Value, - pub api_key: Option<&'a str>, - pub api_base: Option<&'a str>, - pub custom_llm_provider: Option<&'a str>, - pub extra_headers: Option>, - pub timeout: Option, -} - -pub(crate) struct ProviderMessagesRequest { - pub(crate) provider: String, - pub(crate) model: String, - pub(crate) config: &'static dyn AnthropicMessagesProviderConfig, - pub(crate) url: String, - pub(crate) body: Value, - pub(crate) upstream_headers: Vec<(String, String)>, - pub(crate) timeout: Option, -} diff --git a/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md b/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md index 02c5f18c4f3..3eee43e7a2f 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md +++ b/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md @@ -19,7 +19,10 @@ async fn handle(...) -> impl IntoResponse { ... } When a route has business logic worth testing without axum, put it in a sibling `service` (a file, or a folder if the route grows). The route file stays the **axum surface** (router + handler + any socket/SSE adapter); `service` is plain -Rust with **no axum types**. `realtime/` is the example: +Rust with **no axum types**, and its job is to pick the deployment and call the +`core` route entrypoint (see `messages/service.rs` calling +`litellm_core::messages::messages`). Never build a provider request, resolve a +key, or perform the provider call here. `realtime/` is the older example: ``` realtime/ mod.rs # axum surface: router() + handler + the WS<->events adapter @@ -33,6 +36,8 @@ genuinely gets hard to read. `crate::auth::RequireMasterKey` to its arguments; it runs during extraction. Never re-implement the check per route. - **Handlers contain no business logic; `service` contains no axum types.** +- **No provider handlers in this crate.** Transforms, auth headers, and the + provider HTTP call live in `core/src//`. - A route owns its paths in its own `router()`; `mod.rs` only merges. - Cross-cutting concerns (logging, CORS, timeouts) → Tower layers in `mod.rs`, not duplicated in handlers. diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs index 75ed26e5be8..5f4c5fe8de4 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs @@ -1,12 +1,12 @@ use std::sync::Arc; +use litellm_core::constants::ANTHROPIC_MESSAGES_PROVIDER; +use litellm_core::messages::types::MessagesRequest; +use litellm_core::messages::{messages, messages_stream}; use litellm_core::router::Router; use litellm_core::{CoreError, CoreResult}; use serde_json::{Map, Value}; -use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; -use crate::messages::{MessagesRequest, execute_messages}; - pub(crate) enum MessagesResponse { Json(Value), Stream(reqwest::Response), @@ -52,13 +52,14 @@ pub async fn run( extra_headers, timeout: None, }; - let stream = request.body.get("stream").and_then(Value::as_bool) == Some(true); - execute_messages(request, stream) - .await - .map(|response| match response { - crate::messages::MessagesResponse::Json(body) => MessagesResponse::Json(body), - crate::messages::MessagesResponse::Stream(upstream) => { - MessagesResponse::Stream(upstream) - } + if request.body.get("stream").and_then(Value::as_bool) == Some(true) { + return messages_stream(request).await.map(MessagesResponse::Stream); + } + + let response = messages(request).await?; + serde_json::to_value(response) + .map(MessagesResponse::Json) + .map_err(|err| { + CoreError::InvalidResponse(format!("failed to serialize messages response: {err}")) }) } diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md index 8740dccaf01..aee8b4937ef 100644 --- a/litellm-rust/crates/core/AGENTS.md +++ b/litellm-rust/crates/core/AGENTS.md @@ -1,3 +1,7 @@ -litellm-core is the PURE translation layer — types, route contracts (traits), provider transforms (modules under `providers/`), and the router. No network, no I/O, no env reads. +litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-level call is a module under `src//` exposing a public entrypoint named after the route (`messages::messages()`, the Rust equivalent of `litellm.messages()`): you call it and get a typed non-streaming response back. -Routes (ocr, realtime) and providers (mistral, openai) are modules, not crates. +A route module owns everything the call needs: types, the provider template trait, provider transforms (under `providers/`), provider/auth/URL resolution, and the handler that performs the HTTP call. Handlers belong here, never in a host crate. + +Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback dispatch. Env reads are limited to credential fallback in a route's `prepare.rs`. + +Routes (messages, ocr, realtime) and providers (anthropic, mistral, openai) are modules, not crates. diff --git a/litellm-rust/crates/core/CLAUDE.md b/litellm-rust/crates/core/CLAUDE.md index 20873878967..5d36305ded5 100644 --- a/litellm-rust/crates/core/CLAUDE.md +++ b/litellm-rust/crates/core/CLAUDE.md @@ -4,20 +4,28 @@ Rules for `litellm-rust/crates/core`. ## Responsibility -`core` owns shared data types, typed errors, and deterministic helper contracts. -It must stay pure and host-independent. +`core` is the LiteLLM SDK in Rust: it makes the LLM call. Every top-level +LiteLLM call has a public entrypoint here, named after the route +(`messages::messages()` is the Rust equivalent of `litellm.messages()`), and +calling it returns a typed non-streaming response. Allowed: +- The public entrypoint for a route, plus its `_stream` variant when the + route supports streaming. +- Provider resolution, auth header construction, URL building, and the provider + HTTP call (shared reused client, connect + request timeouts). - Shared request/response structs. - Typed errors with stable, non-sensitive messages. - Deterministic validation helpers. - Serialization helpers that intentionally mirror Python output shape. - Route templates that match Python base config responsibilities, such as - `ocr::transformation::OcrProviderConfig`. + `messages::transformation::AnthropicMessagesProviderConfig`. Not allowed: -- Network, filesystem, database, cache, or environment access. -- Secret reads or auth/header construction. +- Serving HTTP: axum routers, extractors, and other transport concerns. +- Filesystem, database, or cache access. +- Config file reading or rollout state; the host resolves those and passes them + in. Env reads are limited to credential fallback in a route's `prepare.rs`. - Logging callbacks, tracing spans, spend writes, or customer callbacks. - Provider-specific branching that belongs in `providers`. - Panics for user/provider-controlled input. @@ -33,10 +41,21 @@ typed field on a struct, not a raw string threaded through the API. ## Structure -Use route names directly under `src/`: `ocr`, future `messages`, +Use route names directly under `src/`: `messages`, `ocr`, future `chat_completions`, `embeddings`, and similar top-level LiteLLM calls. Do not invent broad names like `engine` for route contracts. +`src/messages` is the reference shape for a route module: + +``` +mod.rs pub async fn messages(..) (+ messages_stream) +types.rs request/response types +transformation.rs the provider template trait +prepare.rs provider resolution, auth headers, URL +handler.rs the provider call +client.rs the shared reqwest client +``` + ## Parity Rules - Every shared type used by a provider transform needs unit tests for diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 65c6db7412c..ab8050734f2 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] rand.workspace = true +reqwest.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true @@ -30,5 +31,4 @@ bedrock-auth = [ ] [dev-dependencies] -reqwest.workspace = true tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index 5826a5bc9c1..caada1d98b0 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -1,3 +1,19 @@ pub const OPENAI_DEFAULT_API_BASE: &str = "https://api.openai.com"; pub const OPENAI_RESPONSES_DEFAULT_API_BASE: &str = "https://api.openai.com/v1"; pub const OPENAI_RESPONSES_PATH: &str = "/responses"; + +/// Full-request timeout ceiling for Anthropic Messages provider calls, in +/// seconds. Mirrors the Python Anthropic Messages default. The per-request +/// timeout from the caller still overrides this on the request builder. +pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600; + +/// Connect timeout for Anthropic Messages provider calls, in seconds. +pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10; + +/// Max characters of an upstream error body echoed across the call boundary +/// before truncation, so provider bodies are bounded and data-minimized. +pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256; + +/// Provider name used for Anthropic Messages when a deployment's provider model +/// does not carry an explicit provider prefix. +pub const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic"; diff --git a/litellm-rust/crates/ai-gateway/src/messages/client.rs b/litellm-rust/crates/core/src/messages/client.rs similarity index 100% rename from litellm-rust/crates/ai-gateway/src/messages/client.rs rename to litellm-rust/crates/core/src/messages/client.rs diff --git a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs similarity index 83% rename from litellm-rust/crates/ai-gateway/src/messages/common_utils.rs rename to litellm-rust/crates/core/src/messages/common_utils.rs index 68ecc3f17c1..9dcfcaa71e3 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -1,11 +1,11 @@ -use litellm_core::CoreResult; -use litellm_core::error::{CoreError, json_type_name}; -use litellm_core::messages::transformation::AnthropicMessagesProviderConfig; -use litellm_core::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; -use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; use serde_json::{Map, Value}; use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS; +use crate::error::{CoreError, CoreResult, json_type_name}; +use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; +use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; + +use super::transformation::AnthropicMessagesProviderConfig; pub(super) fn truncate_error_body(body: &str) -> String { if body.chars().count() <= MESSAGES_ERROR_BODY_MAX_CHARS { diff --git a/litellm-rust/crates/ai-gateway/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs similarity index 84% rename from litellm-rust/crates/ai-gateway/src/messages/handler.rs rename to litellm-rust/crates/core/src/messages/handler.rs index 90c12367f50..1c895f66eba 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,15 +1,13 @@ -use litellm_core::CoreResult; -use litellm_core::error::CoreError; -use serde_json::Value; +use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; +use crate::error::{CoreError, CoreResult}; use super::client::http_client; use super::common_utils::truncate_error_body; -use super::types::ProviderMessagesRequest; -use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; +use super::types::{AnthropicMessagesResponse, ProviderMessagesRequest}; pub(super) async fn execute_messages_provider_call( request: ProviderMessagesRequest, -) -> CoreResult { +) -> CoreResult { let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -39,12 +37,7 @@ pub(super) async fn execute_messages_provider_call( let response = serde_json::from_str(&text).map_err(|err| { CoreError::InvalidResponse(format!("invalid messages response JSON: {err}")) })?; - let transformed = request - .config - .transform_response(&request.model, response)?; - serde_json::to_value(transformed).map_err(|err| { - CoreError::InvalidResponse(format!("failed to serialize messages response: {err}")) - }) + request.config.transform_response(&request.model, response) } pub(super) async fn execute_messages_provider_stream( diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index ec2fbb969a6..acb36d89daf 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -1,2 +1,32 @@ +//! The Anthropic Messages call, the Rust equivalent of Python's +//! `litellm.messages()`. +//! +//! [`messages`] is the top-level entrypoint: give it a model, a body, and +//! credentials, and it resolves the provider, transforms the request, calls the +//! provider, and returns a typed non-streaming response. [`messages_stream`] +//! is the streaming variant; it hands the raw upstream response back so a host +//! can splice the event stream to its own caller. + +mod client; +mod common_utils; +mod handler; +mod prepare; pub mod transformation; pub mod types; + +use crate::error::CoreResult; + +use handler::{execute_messages_provider_call, execute_messages_provider_stream}; +use prepare::prepare_messages_call; +use types::{AnthropicMessagesResponse, MessagesRequest}; + +pub async fn messages(request: MessagesRequest<'_>) -> CoreResult { + execute_messages_provider_call(prepare_messages_call(request)?).await +} + +pub async fn messages_stream(request: MessagesRequest<'_>) -> CoreResult { + execute_messages_provider_stream(prepare_messages_call(request)?).await +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs similarity index 92% rename from litellm-rust/crates/ai-gateway/src/messages/prepare.rs rename to litellm-rust/crates/core/src/messages/prepare.rs index 9a027490eb6..94b5b1eaed7 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -1,9 +1,8 @@ -use litellm_core::CoreError; -use litellm_core::CoreResult; -use litellm_core::messages::transformation::MessagesAuthStrategy; -use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; +use crate::error::{CoreError, CoreResult}; +use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; +use super::transformation::MessagesAuthStrategy; use super::types::{MessagesRequest, ProviderMessagesRequest}; pub(super) fn prepare_messages_call( diff --git a/litellm-rust/crates/ai-gateway/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs similarity index 97% rename from litellm-rust/crates/ai-gateway/src/messages/tests.rs rename to litellm-rust/crates/core/src/messages/tests.rs index 23a53e98045..9fc1763683b 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -1,14 +1,16 @@ use std::time::Duration; -use litellm_core::error::CoreError; use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; +use crate::error::CoreError; + use super::common_utils::{ has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body, }; -use super::{MessagesRequest, messages}; +use super::messages; +use super::types::MessagesRequest; async fn read_http_request(socket: &mut TcpStream) -> String { let mut request = Vec::new(); @@ -152,8 +154,8 @@ async fn messages_round_trip_builds_azure_request_and_passes_response_through() .await .expect("messages request succeeds"); - assert_eq!(response["content"][0]["text"], "hi"); - assert_eq!(response["stop_reason"], "end_turn"); + assert_eq!(response.content[0]["text"], "hi"); + assert_eq!(response.stop_reason.as_deref(), Some("end_turn")); let request = server.await.expect("server task completes"); let (head, body) = request.split_once("\r\n\r\n").expect("has body"); @@ -208,8 +210,8 @@ async fn messages_round_trip_builds_native_anthropic_request() { .await .expect("messages request succeeds"); - assert_eq!(response["content"][0]["text"], "hi"); - assert_eq!(response["stop_reason"], "end_turn"); + assert_eq!(response.content[0]["text"], "hi"); + assert_eq!(response.stop_reason.as_deref(), Some("end_turn")); let request = server.await.expect("server task completes"); let (head, _) = request.split_once("\r\n\r\n").expect("has body"); diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs index 11fe17ea40f..b9f807c29fd 100644 --- a/litellm-rust/crates/core/src/messages/types.rs +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -1,6 +1,30 @@ +use std::time::Duration; + use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; +use super::transformation::AnthropicMessagesProviderConfig; + +pub struct MessagesRequest<'a> { + pub model: &'a str, + pub body: Value, + pub api_key: Option<&'a str>, + pub api_base: Option<&'a str>, + pub custom_llm_provider: Option<&'a str>, + pub extra_headers: Option>, + pub timeout: Option, +} + +pub(super) struct ProviderMessagesRequest { + pub(super) provider: String, + pub(super) model: String, + pub(super) config: &'static dyn AnthropicMessagesProviderConfig, + pub(super) url: String, + pub(super) body: Value, + pub(super) upstream_headers: Vec<(String, String)>, + pub(super) timeout: Option, +} + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(untagged)] pub enum SystemPrompt { diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index d6d3d90e6ab..ad3cddfa5fd 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,3 +1,3 @@ -litellm-python-bridge is the PyO3 cdylib that exposes Rust to the litellm Python SDK — a thin adapter (Python objects → Rust calls → Python results) over litellm-ai-gateway. +litellm-python-bridge is the PyO3 cdylib that exposes Rust to the litellm Python SDK — a thin adapter (Python objects → Rust calls → Python results) over the litellm-core route entrypoints (e.g. `litellm_core::messages::messages`). -Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call into litellm-ai-gateway. +Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call the core entrypoint. diff --git a/litellm-rust/crates/python-bridge/CLAUDE.md b/litellm-rust/crates/python-bridge/CLAUDE.md index e5d021ec25b..3ce8b8c639a 100644 --- a/litellm-rust/crates/python-bridge/CLAUDE.md +++ b/litellm-rust/crates/python-bridge/CLAUDE.md @@ -11,11 +11,11 @@ Python-compatible dictionaries. ## Bridge Shape - Prefer one stable method per top-level LiteLLM route, for example - `ocr(payload)`. + `messages(...)`, calling the matching `litellm-core` entrypoint. - Do not add one exported PyO3 function per provider helper unless there is a measured reason. -- Provider dispatch belongs in Rust route modules such as - `litellm_providers::ocr`, not in this PyO3 crate. +- Provider dispatch belongs in the `litellm-core` route module (e.g. + `litellm_core::messages`), not in this PyO3 crate. - Python owns rollout state and fallback. Rust should return errors; Python decides whether to raise or fall back. For a rust-only provider/route (no Python reference), the Python side is a thin dispatch that calls Rust and diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index ee9bdd0b81f..f0cc26a0cca 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -4,10 +4,11 @@ use std::time::Duration; use litellm_ai_gateway::io::audio_transcription::{ AudioTranscriptionRequest, audio_transcription as run_audio_transcription, }; -use litellm_ai_gateway::io::messages::{MessagesRequest, messages as run_messages}; use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use litellm_core::error::CoreError; +use litellm_core::messages::messages as run_messages; +use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyAny, PyDict}; @@ -35,6 +36,15 @@ fn json_to_py(py: Python<'_>, value: Value) -> PyResult> { Ok(json.call_method1("loads", (encoded,))?.unbind()) } +fn messages_response_to_py( + py: Python<'_>, + response: AnthropicMessagesResponse, +) -> PyResult> { + let value = + serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?; + json_to_py(py, value) +} + fn core_error_to_pyerr(err: CoreError) -> PyErr { match err { CoreError::Auth(message) => PyValueError::new_err(message), @@ -382,7 +392,7 @@ fn messages( }); match result { - Ok(value) => json_to_py(py, value), + Ok(response) => messages_response_to_py(py, response), Err(err) => Err(core_error_to_pyerr(err)), } } @@ -404,7 +414,7 @@ fn amessages( marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; pyo3_async_runtimes::tokio::future_into_py(py, async move { - let value = run_messages(MessagesRequest { + let response = run_messages(MessagesRequest { model: &model, body, api_key: api_key.as_deref(), @@ -416,7 +426,7 @@ fn amessages( .await .map_err(core_error_to_pyerr)?; - Python::attach(|py| json_to_py(py, value)) + Python::attach(|py| messages_response_to_py(py, response)) }) }