From a4e655c8bb0d101b7571014dc2d02deb1ed1a53e Mon Sep 17 00:00:00 2001 From: whit3rabbit Date: Wed, 25 Mar 2026 20:09:40 -0500 Subject: [PATCH] feat: LiteLLM gap fill - chat completions input, Azure backend, virtual keys, client SDK Phase 1-8 implementation of the LiteLLM gap fill feature set: - POST /v1/chat/completions: Accept OpenAI-format input, translate through Anthropic pipeline, return OpenAI-format responses (streaming + non-streaming) - Reverse translation layer: openai_to_anthropic_request, anthropic_to_openai_response, ReverseStreamingTranslator (Anthropic SSE -> OpenAI ChatCompletionChunk) - Azure OpenAI backend: BACKEND=azure with deployment-scoped URLs, api-key header, api-version query param (default 2024-10-21) - Virtual key management: SQLite-backed CRUD via admin API (POST/GET/DELETE /admin/api/keys), DashMap in-memory cache, immediate revocation - Per-key rate limiting: RPM sliding window enforcement in auth middleware, 429 with retry-after header on limit exceeded - Client library v0.2.0: ClientBuilder, ToolBuilder, ToolChoiceBuilder, typed streaming, rustdoc examples - New dependencies: dashmap, aws-sigv4, aws-credential-types (prod); opentelemetry stack (feature-gated, optional) 534 tests passing, 0 failures, clippy clean. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/client/Cargo.toml | 4 +- crates/client/src/client.rs | 259 +++++-- crates/client/src/lib.rs | 13 +- crates/client/src/streaming.rs | 82 ++ crates/client/src/tools.rs | 192 +++++ crates/proxy/Cargo.toml | 33 +- crates/proxy/src/admin/db.rs | 140 ++++ crates/proxy/src/admin/keys.rs | 221 ++++++ crates/proxy/src/admin/mod.rs | 2 + crates/proxy/src/admin/routes.rs | 205 ++++- crates/proxy/src/admin/state.rs | 6 + crates/proxy/src/backend/mod.rs | 41 + crates/proxy/src/backend/openai_client.rs | 170 +++- crates/proxy/src/config/mod.rs | 159 +++- crates/proxy/src/main.rs | 535 ++++++++----- crates/proxy/src/server/chat_completions.rs | 429 +++++++++++ crates/proxy/src/server/middleware.rs | 80 +- crates/proxy/src/server/mod.rs | 2 + crates/proxy/src/server/routes.rs | 79 +- crates/proxy/src/server/streaming.rs | 3 +- crates/proxy/tests/live_azure.rs | 117 +++ crates/translator/src/error.rs | 4 + crates/translator/src/lib.rs | 7 +- crates/translator/src/mapping/mod.rs | 6 + .../src/mapping/reverse_message_map.rs | 729 ++++++++++++++++++ .../src/mapping/reverse_streaming_map.rs | 370 +++++++++ crates/translator/src/translate.rs | 45 +- crates/translator/src/util/ids.rs | 5 + docs/ENV.md | 102 +++ .../contracts/admin-keys.md | 93 +++ .../contracts/chat-completions.md | 111 +++ .../data-model.md | 169 ++++ .../20260325-120000-litellm-gap-fill/plan.md | 168 ++++ .../quickstart.md | 103 +++ .../research.md | 237 ++++++ .../20260325-120000-litellm-gap-fill/spec.md | 110 +++ .../20260325-120000-litellm-gap-fill/tasks.md | 300 +++++++ 37 files changed, 5011 insertions(+), 320 deletions(-) create mode 100644 crates/client/src/streaming.rs create mode 100644 crates/client/src/tools.rs create mode 100644 crates/proxy/src/admin/keys.rs create mode 100644 crates/proxy/src/server/chat_completions.rs create mode 100644 crates/proxy/tests/live_azure.rs create mode 100644 crates/translator/src/mapping/reverse_message_map.rs create mode 100644 crates/translator/src/mapping/reverse_streaming_map.rs create mode 100644 specs/20260325-120000-litellm-gap-fill/contracts/admin-keys.md create mode 100644 specs/20260325-120000-litellm-gap-fill/contracts/chat-completions.md create mode 100644 specs/20260325-120000-litellm-gap-fill/data-model.md create mode 100644 specs/20260325-120000-litellm-gap-fill/plan.md create mode 100644 specs/20260325-120000-litellm-gap-fill/quickstart.md create mode 100644 specs/20260325-120000-litellm-gap-fill/research.md create mode 100644 specs/20260325-120000-litellm-gap-fill/spec.md create mode 100644 specs/20260325-120000-litellm-gap-fill/tasks.md diff --git a/crates/client/Cargo.toml b/crates/client/Cargo.toml index e548191..af976d7 100644 --- a/crates/client/Cargo.toml +++ b/crates/client/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "anyllm_client" description = "Async HTTP client for Anthropic-to-OpenAI translation with retry, SSRF protection, and SSE streaming" -version.workspace = true +version = "0.2.0" edition.workspace = true license.workspace = true repository.workspace = true @@ -13,7 +13,7 @@ default = ["ssrf-protection"] ssrf-protection = [] [dependencies] -anyllm_translate = { path = "../translator" } +anyllm_translate = { path = "../translator", version = "0.1.0" } reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "native-tls", "http2"] } tokio = { version = "1", features = ["rt", "sync", "time"] } serde = { version = "1", features = ["derive"] } diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs index 58a4b1a..c5d796d 100644 --- a/crates/client/src/client.rs +++ b/crates/client/src/client.rs @@ -5,17 +5,15 @@ use anyllm_translate::anthropic::messages::MessageResponse; use anyllm_translate::anthropic::streaming::StreamEvent; use anyllm_translate::anthropic::MessageCreateRequest; -use anyllm_translate::openai::{ - ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse, -}; -use anyllm_translate::{mapping, translate_request, translate_response, TranslationConfig}; +use anyllm_translate::openai::{ChatCompletionRequest, ChatCompletionResponse}; +use anyllm_translate::{translate_request, translate_response, TranslationConfig}; use futures::Stream; -use pin_project_lite::pin_project; use crate::error::ClientError; use crate::http::{build_http_client, HttpClientConfig}; use crate::rate_limit::RateLimitHeaders; use crate::retry::{self, RetryableError}; +use crate::streaming::SseTranslatingStream; /// Authentication for the backend API. #[derive(Clone, Debug)] @@ -122,10 +120,126 @@ impl From for ClientError { } } +/// Simplified builder for [`Client`] with sensible defaults. +/// +/// Use this when you want a quick client without manually wiring +/// [`ClientConfig`], [`HttpClientConfig`], and [`TranslationConfig`]. +/// +/// # Examples +/// +/// ```rust,no_run +/// use anyllm_client::ClientBuilder; +/// +/// # fn example() -> Result<(), anyllm_client::ClientError> { +/// let client = ClientBuilder::new() +/// .base_url("https://api.openai.com/v1/chat/completions") +/// .api_key("sk-...") +/// .build()?; +/// # Ok(()) +/// # } +/// ``` +pub struct ClientBuilder { + base_url: Option, + api_key: Option, + timeout: Option, + read_timeout: Option, + max_retries: Option, +} + +impl ClientBuilder { + /// Create a new builder with all fields unset. + pub fn new() -> Self { + Self { + base_url: None, + api_key: None, + timeout: None, + read_timeout: None, + max_retries: None, + } + } + + /// Set the backend URL (e.g., `https://api.openai.com/v1/chat/completions`). + pub fn base_url(mut self, url: &str) -> Self { + self.base_url = Some(url.to_string()); + self + } + + /// Set the API key used as a Bearer token. + pub fn api_key(mut self, key: &str) -> Self { + self.api_key = Some(key.to_string()); + self + } + + /// Set the connection timeout (default: 10s). + pub fn timeout(mut self, duration: std::time::Duration) -> Self { + self.timeout = Some(duration); + self + } + + /// Set the read timeout (default: 900s). + pub fn read_timeout(mut self, duration: std::time::Duration) -> Self { + self.read_timeout = Some(duration); + self + } + + /// Set the maximum number of retries on 429/5xx (default: 3). + /// + /// Note: this value is stored for forward compatibility but the current + /// retry implementation uses the crate-level [`MAX_RETRIES`](crate::retry::MAX_RETRIES) constant. + pub fn max_retries(mut self, n: u32) -> Self { + self.max_retries = Some(n); + self + } + + /// Build the [`Client`], returning an error if `base_url` is missing. + pub fn build(self) -> Result { + let base_url = self.base_url.ok_or_else(|| { + ClientError::ApiError { + status: 0, + message: "ClientBuilder: base_url is required".to_string(), + body: String::new(), + } + })?; + + let http_config = HttpClientConfig { + connect_timeout: self.timeout, + read_timeout: self.read_timeout, + ..HttpClientConfig::new() + }; + + let config = ClientConfig { + chat_completions_url: base_url, + auth: Auth::Bearer(self.api_key.unwrap_or_default()), + http: http_config, + translation: TranslationConfig::default(), + }; + + Ok(Client::new(config)) + } +} + +impl Default for ClientBuilder { + fn default() -> Self { + Self::new() + } +} + /// Async HTTP client for Anthropic-to-OpenAI translation. /// /// Accepts Anthropic Messages API requests, translates to OpenAI format, /// sends to the configured backend, and translates the response back. +/// +/// # Examples +/// +/// ```rust,no_run +/// use anyllm_client::{Client, ClientConfig, Auth}; +/// +/// let config = ClientConfig::builder() +/// .backend_url("https://api.openai.com/v1/chat/completions") +/// .auth(Auth::Bearer("sk-...".into())) +/// .build(); +/// let client = Client::new(config); +/// ``` #[derive(Clone)] pub struct Client { http: reqwest::Client, @@ -139,6 +253,25 @@ impl Client { Self { http, config } } + /// Return a [`ClientBuilder`] for simplified construction. + /// + /// # Examples + /// + /// ```rust,no_run + /// use anyllm_client::Client; + /// + /// # fn example() -> Result<(), anyllm_client::ClientError> { + /// let client = Client::builder() + /// .base_url("https://api.openai.com/v1/chat/completions") + /// .api_key("sk-...") + /// .build()?; + /// # Ok(()) + /// # } + /// ``` + pub fn builder() -> ClientBuilder { + ClientBuilder::new() + } + /// Create from an existing reqwest client and configuration. /// Useful when you want to share an HTTP client across multiple instances. pub fn with_http_client(http: reqwest::Client, config: ClientConfig) -> Self { @@ -236,81 +369,6 @@ impl Client { } } -// -- Streaming implementation -- - -pin_project! { - /// A stream that reads SSE frames from a reqwest response, translates - /// OpenAI chunks to Anthropic StreamEvents, and yields them. - struct SseTranslatingStream { - #[pin] - inner: futures::channel::mpsc::Receiver>, - } -} - -impl SseTranslatingStream { - fn new(response: reqwest::Response, model: String) -> Self { - let (mut tx, rx) = futures::channel::mpsc::channel(32); - - // Spawn a task to read SSE frames and translate them. - tokio::spawn(async move { - let mut translator = mapping::streaming_map::StreamingTranslator::new(model); - let mut done = false; - - let result = crate::sse::read_sse_stream( - response, - |json_str| { - if json_str == "[DONE]" { - done = true; - return Some(translator.finish()); - } - match serde_json::from_str::(json_str) { - Ok(chunk) => Some(translator.process_chunk(&chunk)), - Err(e) => { - tracing::debug!("failed to parse streaming chunk: {e}"); - None - } - } - }, - |events| { - for event in events { - // Block on send; if receiver is dropped, stop. - if tx.try_send(Ok(event.clone())).is_err() { - return false; - } - } - true - }, - ) - .await; - - if let Err(e) = result { - let _ = tx.try_send(Err(ClientError::Sse(e))); - } else if !done { - // Stream ended without [DONE]; flush remaining events. - let events = translator.finish(); - for event in events { - if tx.try_send(Ok(event)).is_err() { - break; - } - } - } - }); - - Self { inner: rx } - } -} - -impl Stream for SseTranslatingStream { - type Item = Result; - - fn poll_next( - self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - ) -> std::task::Poll> { - self.project().inner.poll_next(cx) - } -} - #[cfg(test)] mod tests { use super::*; @@ -358,4 +416,45 @@ mod tests { let _client = Client::new(config); } + + #[test] + fn client_builder_success() { + let client = ClientBuilder::new() + .base_url("https://api.openai.com/v1/chat/completions") + .api_key("sk-test") + .timeout(std::time::Duration::from_secs(5)) + .read_timeout(std::time::Duration::from_secs(30)) + .max_retries(2) + .build(); + assert!(client.is_ok()); + } + + #[test] + fn client_builder_missing_url() { + let result = ClientBuilder::new().api_key("sk-test").build(); + assert!(result.is_err()); + } + + #[test] + fn client_builder_default_api_key() { + // No api_key set: should still build (empty bearer token). + let client = ClientBuilder::new() + .base_url("https://example.com") + .build(); + assert!(client.is_ok()); + } + + #[test] + fn client_builder_via_client() { + let client = Client::builder() + .base_url("https://example.com") + .build(); + assert!(client.is_ok()); + } + + #[test] + fn client_builder_default_trait() { + let builder = ClientBuilder::default(); + assert!(builder.base_url.is_none()); + } } diff --git a/crates/client/src/lib.rs b/crates/client/src/lib.rs index eeec0d1..10126e5 100644 --- a/crates/client/src/lib.rs +++ b/crates/client/src/lib.rs @@ -42,7 +42,8 @@ //! //! # Modules //! -//! - [`client`] -- High-level `Client` for Anthropic-in, Anthropic-out API calls +//! - [`client`] -- High-level `Client` and [`ClientBuilder`] for Anthropic-in, Anthropic-out API calls +//! - [`tools`] -- Builder helpers for [`Tool`] definitions and [`ToolChoice`] //! - [`http`] -- HTTP client builder with TLS and SSRF protection //! - [`retry`] -- Generic retry logic with exponential backoff //! - [`rate_limit`] -- Rate limit header extraction and format conversion @@ -55,11 +56,19 @@ pub mod http; pub mod rate_limit; pub mod retry; pub mod sse; +pub(crate) mod streaming; +pub mod tools; // Convenience re-exports -pub use client::{Auth, Client, ClientConfig, ClientConfigBuilder}; +pub use client::{Auth, Client, ClientBuilder, ClientConfig, ClientConfigBuilder}; pub use error::ClientError; pub use http::{build_http_client, HttpClientConfig}; pub use rate_limit::RateLimitHeaders; pub use retry::{backoff_delay, is_retryable, parse_retry_after, send_with_retry, RetryableError}; pub use sse::{find_double_newline, SseError}; +pub use tools::{ToolBuilder, ToolChoiceBuilder}; + +// Re-export key types from the translator crate so downstream users +// do not need a direct dependency on `anyllm_translate`. +pub use anyllm_translate::anthropic::streaming::StreamEvent; +pub use anyllm_translate::anthropic::{Tool, ToolChoice}; diff --git a/crates/client/src/streaming.rs b/crates/client/src/streaming.rs new file mode 100644 index 0000000..8482a94 --- /dev/null +++ b/crates/client/src/streaming.rs @@ -0,0 +1,82 @@ +//! SSE streaming translation: reads OpenAI chunks, yields Anthropic [`StreamEvent`]s. + +use anyllm_translate::anthropic::streaming::StreamEvent; +use anyllm_translate::mapping; +use anyllm_translate::openai::ChatCompletionChunk; +use futures::Stream; +use pin_project_lite::pin_project; + +use crate::error::ClientError; + +pin_project! { + /// A stream that reads SSE frames from a reqwest response, translates + /// OpenAI chunks to Anthropic StreamEvents, and yields them. + pub(crate) struct SseTranslatingStream { + #[pin] + inner: futures::channel::mpsc::Receiver>, + } +} + +impl SseTranslatingStream { + pub(crate) fn new(response: reqwest::Response, model: String) -> Self { + let (mut tx, rx) = futures::channel::mpsc::channel(32); + + // Spawn a task to read SSE frames and translate them. + tokio::spawn(async move { + let mut translator = mapping::streaming_map::StreamingTranslator::new(model); + let mut done = false; + + let result = crate::sse::read_sse_stream( + response, + |json_str| { + if json_str == "[DONE]" { + done = true; + return Some(translator.finish()); + } + match serde_json::from_str::(json_str) { + Ok(chunk) => Some(translator.process_chunk(&chunk)), + Err(e) => { + tracing::debug!("failed to parse streaming chunk: {e}"); + None + } + } + }, + |events| { + for event in events { + // Block on send; if receiver is dropped, stop. + if tx.try_send(Ok(event.clone())).is_err() { + return false; + } + } + true + }, + ) + .await; + + if let Err(e) = result { + let _ = tx.try_send(Err(ClientError::Sse(e))); + } else if !done { + // Stream ended without [DONE]; flush remaining events. + let events = translator.finish(); + for event in events { + if tx.try_send(Ok(event)).is_err() { + break; + } + } + } + }); + + Self { inner: rx } + } +} + +impl Stream for SseTranslatingStream { + type Item = Result; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + self.project().inner.poll_next(cx) + } +} diff --git a/crates/client/src/tools.rs b/crates/client/src/tools.rs new file mode 100644 index 0000000..714ad73 --- /dev/null +++ b/crates/client/src/tools.rs @@ -0,0 +1,192 @@ +//! Builder helpers for Anthropic tool definitions and tool choice. +//! +//! These builders produce [`Tool`] and [`ToolChoice`] values from +//! `anyllm_translate::anthropic` with a fluent API, avoiding raw JSON +//! construction for common cases. + +use anyllm_translate::anthropic::{Tool, ToolChoice}; +use serde_json::Value; + +/// Fluent builder for an Anthropic [`Tool`] definition. +/// +/// # Examples +/// +/// ``` +/// use anyllm_client::ToolBuilder; +/// use serde_json::json; +/// +/// let tool = ToolBuilder::new("get_weather") +/// .description("Get the current weather for a location") +/// .input_schema(json!({ +/// "type": "object", +/// "properties": { +/// "location": { "type": "string" } +/// }, +/// "required": ["location"] +/// })) +/// .build(); +/// +/// assert_eq!(tool.name, "get_weather"); +/// ``` +pub struct ToolBuilder { + name: String, + description: Option, + input_schema: Value, +} + +impl ToolBuilder { + /// Start building a tool with the given name. + pub fn new(name: &str) -> Self { + Self { + name: name.to_string(), + description: None, + input_schema: Value::Object(serde_json::Map::new()), + } + } + + /// Set the human-readable description shown to the model. + pub fn description(mut self, desc: &str) -> Self { + self.description = Some(desc.to_string()); + self + } + + /// Set the JSON Schema describing the tool's expected input. + pub fn input_schema(mut self, schema: Value) -> Self { + self.input_schema = schema; + self + } + + /// Consume the builder and produce a [`Tool`]. + pub fn build(self) -> Tool { + Tool { + name: self.name, + description: self.description, + input_schema: self.input_schema, + } + } +} + +/// Convenience constructors for [`ToolChoice`] variants. +/// +/// # Examples +/// +/// ``` +/// use anyllm_client::ToolChoiceBuilder; +/// +/// let choice = ToolChoiceBuilder::auto(); +/// let specific = ToolChoiceBuilder::specific("get_weather"); +/// ``` +pub struct ToolChoiceBuilder; + +impl ToolChoiceBuilder { + /// Let the model decide whether to use tools. + pub fn auto() -> ToolChoice { + ToolChoice::Auto { + disable_parallel_tool_use: None, + } + } + + /// Force the model to use at least one tool. + pub fn any() -> ToolChoice { + ToolChoice::Any { + disable_parallel_tool_use: None, + } + } + + /// Prevent the model from using any tools. + pub fn none() -> ToolChoice { + ToolChoice::None + } + + /// Force the model to use a specific tool by name. + pub fn specific(name: &str) -> ToolChoice { + ToolChoice::Tool { + name: name.to_string(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn tool_builder_minimal() { + let tool = ToolBuilder::new("test_tool").build(); + assert_eq!(tool.name, "test_tool"); + assert!(tool.description.is_none()); + assert!(tool.input_schema.is_object()); + } + + #[test] + fn tool_builder_full() { + let schema = json!({ + "type": "object", + "properties": { + "query": { "type": "string" } + }, + "required": ["query"] + }); + + let tool = ToolBuilder::new("search") + .description("Search the web") + .input_schema(schema.clone()) + .build(); + + assert_eq!(tool.name, "search"); + assert_eq!(tool.description.as_deref(), Some("Search the web")); + assert_eq!(tool.input_schema, schema); + } + + #[test] + fn tool_choice_auto() { + let choice = ToolChoiceBuilder::auto(); + assert_eq!( + choice, + ToolChoice::Auto { + disable_parallel_tool_use: None + } + ); + } + + #[test] + fn tool_choice_any() { + let choice = ToolChoiceBuilder::any(); + assert_eq!( + choice, + ToolChoice::Any { + disable_parallel_tool_use: None + } + ); + } + + #[test] + fn tool_choice_none() { + let choice = ToolChoiceBuilder::none(); + assert_eq!(choice, ToolChoice::None); + } + + #[test] + fn tool_choice_specific() { + let choice = ToolChoiceBuilder::specific("get_weather"); + assert_eq!( + choice, + ToolChoice::Tool { + name: "get_weather".to_string() + } + ); + } + + #[test] + fn tool_serializes_correctly() { + let tool = ToolBuilder::new("calc") + .description("Calculator") + .input_schema(json!({"type": "object"})) + .build(); + + let json = serde_json::to_value(&tool).unwrap(); + assert_eq!(json["name"], "calc"); + assert_eq!(json["description"], "Calculator"); + } +} diff --git a/crates/proxy/Cargo.toml b/crates/proxy/Cargo.toml index d54755a..6945f75 100644 --- a/crates/proxy/Cargo.toml +++ b/crates/proxy/Cargo.toml @@ -7,8 +7,8 @@ license.workspace = true repository.workspace = true [dependencies] -anyllm_translate = { path = "../translator" } -anyllm_client = { path = "../client" } +anyllm_translate = { path = "../translator", version = "0.1.0" } +anyllm_client = { path = "../client", version = "0.2.0" } axum = { version = "0.8", features = ["ws"] } tokio = { version = "1", features = ["full"] } reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "native-tls", "http2"] } @@ -29,6 +29,35 @@ subtle = "2" sha2 = "0.10" rusqlite = { version = "0.32", features = ["bundled"] } httpdate = "1" +dashmap = "6" +aws-sigv4 = { version = "1.4", features = ["sign-http"] } +aws-credential-types = "1.2" + +[features] +otel = [ + "opentelemetry", + "opentelemetry_sdk", + "opentelemetry-otlp", + "tracing-opentelemetry", +] + +[dependencies.opentelemetry] +version = "0.31" +optional = true + +[dependencies.opentelemetry_sdk] +version = "0.31" +optional = true + +[dependencies.opentelemetry-otlp] +version = "0.31" +default-features = false +features = ["trace", "http-proto", "reqwest-client"] +optional = true + +[dependencies.tracing-opentelemetry] +version = "0.32" +optional = true [dev-dependencies] pretty_assertions = "1" diff --git a/crates/proxy/src/admin/db.rs b/crates/proxy/src/admin/db.rs index fe9c736..f2235fd 100644 --- a/crates/proxy/src/admin/db.rs +++ b/crates/proxy/src/admin/db.rs @@ -36,6 +36,23 @@ pub fn init_db(conn: &Connection) -> rusqlite::Result<()> { value TEXT NOT NULL, updated_at TEXT NOT NULL ); + + CREATE TABLE IF NOT EXISTS virtual_api_key ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + key_hash TEXT NOT NULL UNIQUE, + key_prefix TEXT NOT NULL, + description TEXT, + created_at TEXT NOT NULL, + expires_at TEXT, + revoked_at TEXT, + spend_limit REAL, + rpm_limit INTEGER, + tpm_limit INTEGER, + total_spend REAL NOT NULL DEFAULT 0, + total_requests INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0 + ); + CREATE INDEX IF NOT EXISTS idx_vak_hash ON virtual_api_key(key_hash); ", )?; Ok(()) @@ -326,6 +343,129 @@ pub fn now_iso8601() -> String { chrono_now() } +// --- Virtual API Key CRUD --- + +use super::keys::VirtualKeyRow; + +/// Insert a new virtual API key. +#[allow(clippy::too_many_arguments)] +pub fn insert_virtual_key( + conn: &Connection, + key_hash: &str, + key_prefix: &str, + description: Option<&str>, + expires_at: Option<&str>, + rpm_limit: Option, + tpm_limit: Option, + spend_limit: Option, +) -> rusqlite::Result { + conn.execute( + "INSERT INTO virtual_api_key (key_hash, key_prefix, description, created_at, expires_at, rpm_limit, tpm_limit, spend_limit) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + params![ + key_hash, + key_prefix, + description, + now_iso8601(), + expires_at, + rpm_limit.map(|v| v as i64), + tpm_limit.map(|v| v as i64), + spend_limit, + ], + )?; + Ok(conn.last_insert_rowid()) +} + +/// List all virtual keys (active, expired, revoked). +pub fn list_virtual_keys(conn: &Connection) -> rusqlite::Result> { + let mut stmt = conn.prepare( + "SELECT id, key_hash, key_prefix, description, created_at, expires_at, revoked_at, + rpm_limit, tpm_limit, spend_limit, total_spend, total_requests, total_tokens + FROM virtual_api_key ORDER BY id DESC", + )?; + let rows = stmt.query_map([], |row| { + Ok(VirtualKeyRow { + id: row.get(0)?, + key_hash: row.get(1)?, + key_prefix: row.get(2)?, + description: row.get(3)?, + created_at: row.get(4)?, + expires_at: row.get(5)?, + revoked_at: row.get(6)?, + rpm_limit: row.get::<_, Option>(7)?.map(|v| v as u32), + tpm_limit: row.get::<_, Option>(8)?.map(|v| v as u32), + spend_limit: row.get(9)?, + total_spend: row.get::<_, f64>(10).unwrap_or(0.0), + total_requests: row.get::<_, i64>(11).unwrap_or(0), + total_tokens: row.get::<_, i64>(12).unwrap_or(0), + }) + })?; + rows.collect() +} + +/// Revoke a virtual key by setting revoked_at. Returns the row if found. +pub fn revoke_virtual_key(conn: &Connection, id: i64) -> rusqlite::Result> { + let now = now_iso8601(); + let updated = conn.execute( + "UPDATE virtual_api_key SET revoked_at = ?1 WHERE id = ?2 AND revoked_at IS NULL", + params![now, id], + )?; + if updated == 0 { + return Ok(None); + } + let mut stmt = conn.prepare( + "SELECT id, key_hash, key_prefix, description, created_at, expires_at, revoked_at, + rpm_limit, tpm_limit, spend_limit, total_spend, total_requests, total_tokens + FROM virtual_api_key WHERE id = ?1", + )?; + stmt.query_row(params![id], |row| { + Ok(Some(VirtualKeyRow { + id: row.get(0)?, + key_hash: row.get(1)?, + key_prefix: row.get(2)?, + description: row.get(3)?, + created_at: row.get(4)?, + expires_at: row.get(5)?, + revoked_at: row.get(6)?, + rpm_limit: row.get::<_, Option>(7)?.map(|v| v as u32), + tpm_limit: row.get::<_, Option>(8)?.map(|v| v as u32), + spend_limit: row.get(9)?, + total_spend: row.get::<_, f64>(10).unwrap_or(0.0), + total_requests: row.get::<_, i64>(11).unwrap_or(0), + total_tokens: row.get::<_, i64>(12).unwrap_or(0), + })) + }) +} + +/// Load all active (non-revoked, non-expired) virtual keys from the database. +pub fn load_active_virtual_keys(conn: &Connection) -> rusqlite::Result> { + let now = now_iso8601(); + let mut stmt = conn.prepare( + "SELECT id, key_hash, key_prefix, description, created_at, expires_at, revoked_at, + rpm_limit, tpm_limit, spend_limit, total_spend, total_requests, total_tokens + FROM virtual_api_key + WHERE revoked_at IS NULL AND (expires_at IS NULL OR expires_at > ?1)", + )?; + let rows = stmt.query_map(params![now], |row| { + Ok(VirtualKeyRow { + id: row.get(0)?, + key_hash: row.get(1)?, + key_prefix: row.get(2)?, + description: row.get(3)?, + created_at: row.get(4)?, + expires_at: row.get(5)?, + revoked_at: row.get(6)?, + rpm_limit: row.get::<_, Option>(7)?.map(|v| v as u32), + tpm_limit: row.get::<_, Option>(8)?.map(|v| v as u32), + spend_limit: row.get(9)?, + total_spend: row.get::<_, f64>(10).unwrap_or(0.0), + total_requests: row.get::<_, i64>(11).unwrap_or(0), + total_tokens: row.get::<_, i64>(12).unwrap_or(0), + }) + })?; + rows.collect() +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/proxy/src/admin/keys.rs b/crates/proxy/src/admin/keys.rs new file mode 100644 index 0000000..6bb24a3 --- /dev/null +++ b/crates/proxy/src/admin/keys.rs @@ -0,0 +1,221 @@ +// Virtual API key generation, hashing, and rate limit state. + +use sha2::{Digest, Sha256}; +use std::collections::VecDeque; +use std::sync::{Arc, Mutex}; + +/// Generate a new virtual API key. +/// Returns (raw_key, key_prefix, key_hash_hex). +/// The raw_key is shown once at creation; key_prefix is for display; key_hash_hex is stored. +pub fn generate_virtual_key() -> (String, String, String) { + let a = uuid::Uuid::new_v4().as_simple().to_string(); + let b = uuid::Uuid::new_v4().as_simple().to_string(); + let raw_key = format!("sk-vk{}{}", a, b); + let key_prefix = raw_key[..8].to_string(); + let key_hash_hex = hash_key(&raw_key); + (raw_key, key_prefix, key_hash_hex) +} + +/// SHA-256 hash a key string and return hex-encoded result. +pub fn hash_key(key: &str) -> String { + let hash: [u8; 32] = Sha256::digest(key.as_bytes()).into(); + bytes_to_hex(&hash) +} + +/// Convert a hex-encoded hash to raw bytes. +pub fn hash_from_hex(hex_str: &str) -> Option<[u8; 32]> { + if hex_str.len() != 64 { + return None; + } + let mut arr = [0u8; 32]; + for i in 0..32 { + arr[i] = u8::from_str_radix(&hex_str[i * 2..i * 2 + 2], 16).ok()?; + } + Some(arr) +} + +fn bytes_to_hex(bytes: &[u8]) -> String { + bytes.iter().map(|b| format!("{b:02x}")).collect() +} + +/// In-memory metadata for a virtual key (stored in DashMap). +#[derive(Debug)] +pub struct VirtualKeyMeta { + pub id: i64, + pub description: Option, + /// Epoch seconds; None = no expiry. + pub expires_at: Option, + pub rpm_limit: Option, + pub tpm_limit: Option, + pub rate_state: Arc, +} + +/// Sliding window rate limit state per virtual key. +#[derive(Debug)] +pub struct RateLimitState { + pub rpm_window: Mutex>, + pub tpm_window: Mutex>, +} + +impl Default for RateLimitState { + fn default() -> Self { + Self::new() + } +} + +impl RateLimitState { + pub fn new() -> Self { + Self { + rpm_window: Mutex::new(VecDeque::new()), + tpm_window: Mutex::new(VecDeque::new()), + } + } + + /// Check if a new request is within the RPM limit. + /// Returns Ok(()) if allowed, Err(retry_after_secs) if exceeded. + pub fn check_rpm(&self, limit: u32, now_ms: u64) -> Result<(), u64> { + let mut window = self.rpm_window.lock().unwrap_or_else(|e| e.into_inner()); + let cutoff = now_ms.saturating_sub(60_000); + // Drain expired entries + while window.front().is_some_and(|&ts| ts < cutoff) { + window.pop_front(); + } + if window.len() >= limit as usize { + // Compute retry-after: time until the oldest entry expires + let oldest = window.front().copied().unwrap_or(now_ms); + let retry_after_ms = (oldest + 60_000).saturating_sub(now_ms); + return Err((retry_after_ms / 1000).max(1)); + } + window.push_back(now_ms); + Ok(()) + } + + /// Record a TPM token count for the current request. + pub fn record_tpm(&self, now_ms: u64, tokens: u32) { + let mut window = self.tpm_window.lock().unwrap_or_else(|e| e.into_inner()); + let cutoff = now_ms.saturating_sub(60_000); + while window.front().is_some_and(|&(ts, _)| ts < cutoff) { + window.pop_front(); + } + window.push_back((now_ms, tokens)); + } + + /// Check if adding `tokens` would exceed the TPM limit. + pub fn check_tpm(&self, limit: u32, now_ms: u64) -> Result<(), u64> { + let mut window = self.tpm_window.lock().unwrap_or_else(|e| e.into_inner()); + let cutoff = now_ms.saturating_sub(60_000); + while window.front().is_some_and(|&(ts, _)| ts < cutoff) { + window.pop_front(); + } + let total: u64 = window.iter().map(|&(_, t)| t as u64).sum(); + if total >= limit as u64 { + let oldest = window.front().map(|&(ts, _)| ts).unwrap_or(now_ms); + let retry_after_ms = (oldest + 60_000).saturating_sub(now_ms); + return Err((retry_after_ms / 1000).max(1)); + } + Ok(()) + } +} + +/// Row from the virtual_api_key table. +#[derive(Debug, Clone, serde::Serialize)] +pub struct VirtualKeyRow { + pub id: i64, + pub key_hash: String, + pub key_prefix: String, + pub description: Option, + pub created_at: String, + pub expires_at: Option, + pub revoked_at: Option, + pub rpm_limit: Option, + pub tpm_limit: Option, + pub spend_limit: Option, + pub total_spend: f64, + pub total_requests: i64, + pub total_tokens: i64, +} + +impl VirtualKeyRow { + /// Compute the effective status of a key. + pub fn status(&self) -> &'static str { + if self.revoked_at.is_some() { + return "revoked"; + } + if let Some(ref exp) = self.expires_at { + if *exp <= super::db::now_iso8601() { + return "expired"; + } + } + "active" + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn key_generation_format() { + let (raw, prefix, hash) = generate_virtual_key(); + assert!(raw.starts_with("sk-vk")); + assert_eq!(prefix.len(), 8); + assert!(prefix.starts_with("sk-vk")); + assert_eq!(hash.len(), 64); // hex SHA-256 + } + + #[test] + fn hash_deterministic() { + let h1 = hash_key("test-key-123"); + let h2 = hash_key("test-key-123"); + assert_eq!(h1, h2); + } + + #[test] + fn hash_from_hex_roundtrip() { + let hex = hash_key("test"); + let bytes = hash_from_hex(&hex).unwrap(); + assert_eq!(bytes_to_hex(&bytes), hex); + } + + #[test] + fn rpm_within_limit() { + let state = RateLimitState::new(); + let now = 1000000; + assert!(state.check_rpm(3, now).is_ok()); + assert!(state.check_rpm(3, now + 1).is_ok()); + assert!(state.check_rpm(3, now + 2).is_ok()); + // 4th request should be rejected + assert!(state.check_rpm(3, now + 3).is_err()); + } + + #[test] + fn rpm_window_expiry() { + let state = RateLimitState::new(); + let now = 1000000; + assert!(state.check_rpm(1, now).is_ok()); + assert!(state.check_rpm(1, now + 100).is_err()); + // After 60 seconds, window should clear + assert!(state.check_rpm(1, now + 60_001).is_ok()); + } + + #[test] + fn tpm_within_limit() { + let state = RateLimitState::new(); + let now = 1000000; + state.record_tpm(now, 50); + assert!(state.check_tpm(100, now + 1).is_ok()); + state.record_tpm(now + 1, 50); + // At limit + assert!(state.check_tpm(100, now + 2).is_err()); + } + + #[test] + fn tpm_window_expiry() { + let state = RateLimitState::new(); + let now = 1000000; + state.record_tpm(now, 100); + assert!(state.check_tpm(100, now + 1).is_err()); + // After 60 seconds + assert!(state.check_tpm(100, now + 60_001).is_ok()); + } +} diff --git a/crates/proxy/src/admin/mod.rs b/crates/proxy/src/admin/mod.rs index 005976c..cd759c6 100644 --- a/crates/proxy/src/admin/mod.rs +++ b/crates/proxy/src/admin/mod.rs @@ -2,6 +2,8 @@ pub mod auth; /// SQLite persistence for request logs and config overrides. pub mod db; +/// Virtual API key generation, hashing, and rate limit state. +pub mod keys; /// Admin HTTP router: config management, request log queries, metrics. pub mod routes; /// Shared mutable state between proxy handlers and admin server. diff --git a/crates/proxy/src/admin/routes.rs b/crates/proxy/src/admin/routes.rs index dfb97bd..fb7266e 100644 --- a/crates/proxy/src/admin/routes.rs +++ b/crates/proxy/src/admin/routes.rs @@ -8,7 +8,7 @@ use axum::{ http::StatusCode, middleware, response::IntoResponse, - routing::{delete, get}, + routing::{delete, get, post}, Json, Router, }; use std::sync::Arc; @@ -86,10 +86,13 @@ pub fn admin_router(shared: SharedState, token: Arc) -> Router { "/admin/api/config/overrides/{key}", delete(delete_config_override), ) + .route("/admin/api/env", get(get_env)) .route("/admin/api/metrics", get(get_metrics)) .route("/admin/api/requests", get(get_requests)) .route("/admin/api/requests/{id}", get(get_request_by_id)) .route("/admin/api/backends", get(get_backends)) + .route("/admin/api/keys", post(create_key).get(list_keys)) + .route("/admin/api/keys/{id}", delete(revoke_key)) .with_state(shared.clone()) .layer(middleware::from_fn_with_state( token.clone(), @@ -140,6 +143,58 @@ async fn serve_spa() -> impl IntoResponse { ) } +// -- Env endpoint -- + +/// GET /admin/api/env -- effective environment variable values. +/// Secrets (API keys, tokens) are masked; plain config values are shown as-is. +async fn get_env() -> Json { + fn plain(key: &str) -> serde_json::Value { + match std::env::var(key) { + Ok(v) if !v.is_empty() => serde_json::Value::String(v), + _ => serde_json::Value::Null, + } + } + fn secret(key: &str) -> serde_json::Value { + match std::env::var(key) { + Ok(v) if !v.is_empty() => { + serde_json::Value::String(anyllm_translate::util::redact::redact_secret(&v)) + } + _ => serde_json::Value::Null, + } + } + + Json(serde_json::json!({ + // Core proxy config + "BACKEND": plain("BACKEND"), + "LISTEN_PORT": plain("LISTEN_PORT"), + "BIG_MODEL": plain("BIG_MODEL"), + "SMALL_MODEL": plain("SMALL_MODEL"), + "RUST_LOG": plain("RUST_LOG"), + "LOG_BODIES": plain("LOG_BODIES"), + "PROXY_CONFIG": plain("PROXY_CONFIG"), + // OpenAI / compatible + "OPENAI_BASE_URL": plain("OPENAI_BASE_URL"), + "OPENAI_API_FORMAT": plain("OPENAI_API_FORMAT"), + "OPENAI_API_KEY": secret("OPENAI_API_KEY"), + // Vertex AI + "VERTEX_PROJECT": plain("VERTEX_PROJECT"), + "VERTEX_REGION": plain("VERTEX_REGION"), + "VERTEX_API_KEY": secret("VERTEX_API_KEY"), + // Gemini + "GEMINI_BASE_URL": plain("GEMINI_BASE_URL"), + "GEMINI_API_KEY": secret("GEMINI_API_KEY"), + // Auth + "PROXY_API_KEYS": secret("PROXY_API_KEYS"), + // TLS + "TLS_CLIENT_CERT_P12": plain("TLS_CLIENT_CERT_P12"), + "TLS_CA_CERT": plain("TLS_CA_CERT"), + // Admin + "ADMIN_PORT": plain("ADMIN_PORT"), + "ADMIN_DB_PATH": plain("ADMIN_DB_PATH"), + "ADMIN_LOG_RETENTION_DAYS": plain("ADMIN_LOG_RETENTION_DAYS"), + })) +} + // -- Config endpoints -- /// GET /admin/api/config -- effective config (env defaults + overrides). @@ -561,6 +616,154 @@ async fn get_backends(State(shared): State) -> Json, + expires_at: Option, + rpm_limit: Option, + tpm_limit: Option, + spend_limit: Option, +} + +/// POST /admin/api/keys -- create a new virtual API key. +async fn create_key( + State(shared): State, + Json(body): Json, +) -> axum::response::Response { + let (raw_key, key_prefix, key_hash_hex) = super::keys::generate_virtual_key(); + let result = super::state::with_db(&shared.db, { + let hash = key_hash_hex.clone(); + let prefix = key_prefix.clone(); + let desc = body.description.clone(); + let exp = body.expires_at.clone(); + let rpm = body.rpm_limit; + let tpm = body.tpm_limit; + let spend = body.spend_limit; + move |conn| { + super::db::insert_virtual_key( + conn, + &hash, + &prefix, + desc.as_deref(), + exp.as_deref(), + rpm, + tpm, + spend, + ) + } + }) + .await; + + match result { + Some(Ok(id)) => { + if let Some(hash_bytes) = super::keys::hash_from_hex(&key_hash_hex) { + shared.virtual_keys.insert( + hash_bytes, + super::keys::VirtualKeyMeta { + id, + description: body.description.clone(), + expires_at: None, + rpm_limit: body.rpm_limit, + tpm_limit: body.tpm_limit, + rate_state: std::sync::Arc::new(super::keys::RateLimitState::new()), + }, + ); + } + ( + StatusCode::CREATED, + Json(serde_json::json!({ + "id": id, + "key": raw_key, + "key_prefix": key_prefix, + "description": body.description, + "created_at": super::db::now_iso8601(), + "expires_at": body.expires_at, + "rpm_limit": body.rpm_limit, + "tpm_limit": body.tpm_limit, + "spend_limit": body.spend_limit, + })), + ) + .into_response() + } + _ => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "Failed to create key"})), + ) + .into_response(), + } +} + +/// GET /admin/api/keys -- list all virtual keys. +async fn list_keys(State(shared): State) -> axum::response::Response { + let result = super::state::with_db(&shared.db, super::db::list_virtual_keys).await; + match result { + Some(Ok(keys)) => { + let enriched: Vec = keys + .iter() + .map(|k| { + serde_json::json!({ + "id": k.id, + "key_prefix": k.key_prefix, + "description": k.description, + "created_at": k.created_at, + "expires_at": k.expires_at, + "revoked_at": k.revoked_at, + "rpm_limit": k.rpm_limit, + "tpm_limit": k.tpm_limit, + "spend_limit": k.spend_limit, + "total_spend": k.total_spend, + "total_requests": k.total_requests, + "total_tokens": k.total_tokens, + "status": k.status(), + }) + }) + .collect(); + Json(serde_json::json!({ "keys": enriched })).into_response() + } + _ => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "Failed to list keys"})), + ) + .into_response(), + } +} + +/// DELETE /admin/api/keys/{id} -- revoke a virtual key. +async fn revoke_key( + State(shared): State, + Path(id): Path, +) -> axum::response::Response { + let result = super::state::with_db(&shared.db, move |conn| { + super::db::revoke_virtual_key(conn, id) + }) + .await; + match result { + Some(Ok(Some(row))) => { + if let Some(hash_bytes) = super::keys::hash_from_hex(&row.key_hash) { + shared.virtual_keys.remove(&hash_bytes); + } + Json(serde_json::json!({ + "id": row.id, + "revoked_at": row.revoked_at, + "status": "revoked", + })) + .into_response() + } + Some(Ok(None)) => ( + StatusCode::NOT_FOUND, + Json(serde_json::json!({"error": "Key not found or already revoked"})), + ) + .into_response(), + _ => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "Failed to revoke key"})), + ) + .into_response(), + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/proxy/src/admin/state.rs b/crates/proxy/src/admin/state.rs index 52ce731..cc7bd1e 100644 --- a/crates/proxy/src/admin/state.rs +++ b/crates/proxy/src/admin/state.rs @@ -1,8 +1,10 @@ // Shared state between the proxy and admin server. // RuntimeConfig holds mutable settings; AdminEvent is broadcast to WebSocket clients. +use crate::admin::keys::VirtualKeyMeta; use crate::config::ModelMapping; use crate::metrics::Metrics; +use dashmap::DashMap; use indexmap::IndexMap; use std::collections::HashMap; use std::sync::{Arc, Mutex, RwLock}; @@ -36,6 +38,9 @@ pub struct SharedState { /// Serializes config write operations (Phase 1: SQLite + Phase 2: in-memory) /// so concurrent PUT /admin/api/config requests cannot interleave. pub config_write_lock: Arc>, + /// In-memory cache of active virtual API keys, keyed by SHA-256 hash bytes. + /// Populated from SQLite at startup; updated on create/revoke via admin API. + pub virtual_keys: Arc>, } /// Run a synchronous closure against the SQLite connection on the blocking @@ -125,6 +130,7 @@ impl SharedState { log_tx, log_reload: None, config_write_lock: Arc::new(tokio::sync::Mutex::new(())), + virtual_keys: Arc::new(DashMap::new()), } } } diff --git a/crates/proxy/src/backend/mod.rs b/crates/proxy/src/backend/mod.rs index 4e8e711..d5ceb09 100644 --- a/crates/proxy/src/backend/mod.rs +++ b/crates/proxy/src/backend/mod.rs @@ -42,6 +42,10 @@ pub(crate) async fn send_with_retry( name: "x-goog-api-key", value: key, }, + BackendAuth::AzureApiKey(key) => anyllm_client::retry::RequestAuth::Header { + name: "api-key", + value: key, + }, }; anyllm_client::retry::send_with_retry(client, url, &request_auth, body, label).await } @@ -55,6 +59,8 @@ pub enum BackendClient { /// with a different request/response shape. Separate variant so callers /// can pattern-match on the API format. OpenAIResponses(OpenAIClient), + /// Azure OpenAI: same Chat Completions format, different auth and URL scheme. + AzureOpenAI(OpenAIClient), Vertex(OpenAIClient), /// Gemini via OpenAI-compatible endpoint (reuses OpenAI translation path). GeminiOpenAI(OpenAIClient), @@ -125,6 +131,39 @@ impl From for BackendError { } impl BackendClient { + /// Forward a raw embeddings request to the backend. No translation — model names pass through. + /// Returns `501 Not Implemented` for the Anthropic backend (no embeddings endpoint). + pub async fn embeddings_passthrough( + &self, + body: bytes::Bytes, + content_type: &str, + ) -> Result<(axum::http::StatusCode, axum::http::HeaderMap, bytes::Bytes), BackendError> { + match self { + Self::OpenAI(c) + | Self::AzureOpenAI(c) + | Self::Vertex(c) + | Self::GeminiOpenAI(c) + | Self::OpenAIResponses(c) => c + .embeddings_passthrough(body, content_type) + .await + .map_err(BackendError::OpenAI), + Self::Anthropic(_) => { + // Anthropic has no embeddings API. + let err = anyllm_translate::mapping::errors_map::create_anthropic_error( + anyllm_translate::anthropic::ErrorType::InvalidRequestError, + "Embeddings are not supported by the Anthropic backend.".to_string(), + None, + ); + let body = serde_json::to_vec(&err).unwrap_or_default(); + Ok(( + axum::http::StatusCode::NOT_IMPLEMENTED, + axum::http::HeaderMap::new(), + bytes::Bytes::from(body), + )) + } + } + } + /// Create a backend client from a single-backend [`Config`]. /// /// Dispatches on [`Config::backend`] and [`Config::openai_api_format`] to construct @@ -135,6 +174,7 @@ impl BackendClient { OpenAIApiFormat::Chat => Self::OpenAI(OpenAIClient::new(config)), OpenAIApiFormat::Responses => Self::OpenAIResponses(OpenAIClient::new(config)), }, + BackendKind::AzureOpenAI => Self::AzureOpenAI(OpenAIClient::new(config)), BackendKind::Vertex => Self::Vertex(OpenAIClient::new(config)), BackendKind::Gemini => Self::GeminiOpenAI(OpenAIClient::new(config)), BackendKind::Anthropic => Self::Anthropic(AnthropicClient::new( @@ -166,6 +206,7 @@ impl BackendClient { OpenAIApiFormat::Chat => Self::OpenAI(OpenAIClient::new(&legacy)), OpenAIApiFormat::Responses => Self::OpenAIResponses(OpenAIClient::new(&legacy)), }, + BackendKind::AzureOpenAI => Self::AzureOpenAI(OpenAIClient::new(&legacy)), BackendKind::Vertex => Self::Vertex(OpenAIClient::new(&legacy)), BackendKind::Gemini => Self::GeminiOpenAI(OpenAIClient::new(&legacy)), BackendKind::Anthropic => Self::Anthropic(AnthropicClient::from_backend_config(bc)), diff --git a/crates/proxy/src/backend/openai_client.rs b/crates/proxy/src/backend/openai_client.rs index 471a1f4..1d85790 100644 --- a/crates/proxy/src/backend/openai_client.rs +++ b/crates/proxy/src/backend/openai_client.rs @@ -14,6 +14,7 @@ pub struct OpenAIClient { client: Client, chat_completions_url: String, responses_url: String, + embeddings_url: String, auth: BackendAuth, } @@ -27,22 +28,52 @@ impl OpenAIClient { // - OpenAI: {base}/v1/chat/completions (base has no path) // - Vertex: {base}/chat/completions (base ends at .../openapi) // - Gemini: {base}/chat/completions (config appends /openai to base) - let (chat_completions_url, responses_url) = match config.backend { + let (chat_completions_url, responses_url, embeddings_url) = match config.backend { BackendKind::OpenAI => ( format!("{}/v1/chat/completions", config.openai_base_url), format!("{}/v1/responses", config.openai_base_url), + format!("{}/v1/embeddings", config.openai_base_url), ), BackendKind::Vertex => ( format!("{}/chat/completions", config.openai_base_url), // Vertex does not support Responses API; URL included for completeness format!("{}/responses", config.openai_base_url), + format!("{}/embeddings", config.openai_base_url), ), BackendKind::Gemini => ( // openai_base_url already has /openai appended by config, // producing .../v1beta/openai/chat/completions format!("{}/chat/completions", config.openai_base_url), format!("{}/responses", config.openai_base_url), + // Gemini embeddings: .../v1beta/openai/embeddings + format!("{}/embeddings", config.openai_base_url), ), + BackendKind::AzureOpenAI => { + // Azure URL is pre-constructed in config (includes deployment + api-version). + // Embeddings and Responses URLs are derived by replacing the path component. + let endpoint = config + .openai_base_url + .split("/openai/deployments/") + .next() + .unwrap_or(&config.openai_base_url); + let api_version = config + .openai_base_url + .split("api-version=") + .nth(1) + .unwrap_or("2024-10-21"); + let deployment = config + .openai_base_url + .split("/openai/deployments/") + .nth(1) + .and_then(|s| s.split('/').next()) + .unwrap_or(""); + ( + config.openai_base_url.clone(), + // Azure Responses API is not widely available; provide URL for completeness + format!("{endpoint}/openai/deployments/{deployment}/responses?api-version={api_version}"), + format!("{endpoint}/openai/deployments/{deployment}/embeddings?api-version={api_version}"), + ) + } BackendKind::Anthropic => { unreachable!("OpenAIClient should not be constructed for Anthropic backend") } @@ -52,6 +83,7 @@ impl OpenAIClient { client, chat_completions_url, responses_url, + embeddings_url, auth: config.backend_auth.clone(), } } @@ -158,6 +190,38 @@ impl OpenAIClient { let rate_limits = RateLimitHeaders::from_openai_headers(response.headers()); Ok((response, rate_limits)) } + + /// Forward a raw embeddings request body to the backend embeddings endpoint. + /// No retry: embeddings are idempotent but we keep it simple — callers can retry. + /// + /// OpenAI: + pub async fn embeddings_passthrough( + &self, + body: bytes::Bytes, + content_type: &str, + ) -> Result<(axum::http::StatusCode, axum::http::HeaderMap, bytes::Bytes), OpenAIClientError> + { + let mut req = self + .client + .post(&self.embeddings_url) + .body(body) + .header("content-type", content_type); + req = match &self.auth { + BackendAuth::BearerToken(token) => req.bearer_auth(token), + BackendAuth::GoogleApiKey(key) => req.header("x-goog-api-key", key), + BackendAuth::AzureApiKey(key) => req.header("api-key", key), + }; + + let response = req.send().await.map_err(OpenAIClientError::Request)?; + let status = axum::http::StatusCode::from_u16(response.status().as_u16()) + .unwrap_or(axum::http::StatusCode::INTERNAL_SERVER_ERROR); + let mut resp_headers = axum::http::HeaderMap::new(); + if let Some(ct) = response.headers().get("content-type") { + resp_headers.insert("content-type", ct.clone()); + } + let resp_body = response.bytes().await.map_err(OpenAIClientError::Request)?; + Ok((status, resp_headers, resp_body)) + } } /// Errors from the OpenAI HTTP client. @@ -346,4 +410,108 @@ mod tests { .chat_completions_url .ends_with("/openapi/chat/completions")); } + + #[test] + fn embeddings_url_openai() { + use crate::config::{BackendKind, ModelMapping, OpenAIApiFormat, TlsConfig}; + let config = Config { + backend: BackendKind::OpenAI, + openai_api_key: "test".into(), + openai_base_url: "https://api.openai.com".into(), + listen_port: 3000, + model_mapping: ModelMapping { + big_model: "gpt-4o".into(), + small_model: "gpt-4o-mini".into(), + }, + tls: TlsConfig::default(), + backend_auth: BackendAuth::BearerToken("test".into()), + log_bodies: false, + openai_api_format: OpenAIApiFormat::Chat, + }; + let client = OpenAIClient::new(&config); + assert_eq!( + client.embeddings_url, + "https://api.openai.com/v1/embeddings" + ); + } + + #[test] + fn embeddings_url_vertex() { + use crate::config::{BackendKind, ModelMapping, OpenAIApiFormat, TlsConfig}; + let config = Config { + backend: BackendKind::Vertex, + openai_api_key: String::new(), + openai_base_url: "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/endpoints/openapi".into(), + listen_port: 3000, + model_mapping: ModelMapping { + big_model: "gemini-2.5-pro".into(), + small_model: "gemini-2.5-flash".into(), + }, + tls: TlsConfig::default(), + backend_auth: BackendAuth::GoogleApiKey("test-key".into()), + log_bodies: false, + openai_api_format: OpenAIApiFormat::Chat, + }; + let client = OpenAIClient::new(&config); + assert!( + client.embeddings_url.ends_with("/openapi/embeddings"), + "got: {}", + client.embeddings_url + ); + } + + #[test] + fn embeddings_url_gemini() { + use crate::config::{BackendKind, ModelMapping, OpenAIApiFormat, TlsConfig}; + let config = Config { + backend: BackendKind::Gemini, + openai_api_key: String::new(), + // Config appends /openai to the base, so this is what arrives here + openai_base_url: "https://generativelanguage.googleapis.com/v1beta/openai".into(), + listen_port: 3000, + model_mapping: ModelMapping { + big_model: "gemini-2.5-pro".into(), + small_model: "gemini-2.5-flash".into(), + }, + tls: TlsConfig::default(), + backend_auth: BackendAuth::GoogleApiKey("test-gemini-key".into()), + log_bodies: false, + openai_api_format: OpenAIApiFormat::Chat, + }; + let client = OpenAIClient::new(&config); + assert_eq!( + client.embeddings_url, + "https://generativelanguage.googleapis.com/v1beta/openai/embeddings" + ); + } + + #[test] + fn azure_url_passthrough() { + use crate::config::{BackendKind, ModelMapping, OpenAIApiFormat, TlsConfig}; + let config = Config { + backend: BackendKind::AzureOpenAI, + openai_api_key: String::new(), + openai_base_url: "https://myresource.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21".into(), + listen_port: 3000, + model_mapping: ModelMapping { + big_model: "gpt-4o".into(), + small_model: "gpt-4o-mini".into(), + }, + tls: TlsConfig::default(), + backend_auth: BackendAuth::AzureApiKey("test-azure-key".into()), + log_bodies: false, + openai_api_format: OpenAIApiFormat::Chat, + }; + let client = OpenAIClient::new(&config); + // Chat completions URL is the pre-built URL, unchanged + assert_eq!( + client.chat_completions_url, + "https://myresource.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21" + ); + // Embeddings URL is derived from the endpoint and deployment + assert_eq!( + client.embeddings_url, + "https://myresource.openai.azure.com/openai/deployments/gpt-4o/embeddings?api-version=2024-10-21" + ); + } } diff --git a/crates/proxy/src/config/mod.rs b/crates/proxy/src/config/mod.rs index d8842a2..5ea48ed 100644 --- a/crates/proxy/src/config/mod.rs +++ b/crates/proxy/src/config/mod.rs @@ -15,6 +15,7 @@ const GEMINI_OPENAI_PATH: &str = "/openai"; #[derive(Debug, Clone, PartialEq, Eq)] pub enum BackendKind { OpenAI, + AzureOpenAI, Vertex, Gemini, Anthropic, @@ -36,6 +37,8 @@ pub enum BackendAuth { BearerToken(String), /// `x-goog-api-key: {key}` (Vertex API key) GoogleApiKey(String), + /// `api-key: {key}` (Azure OpenAI) + AzureApiKey(String), } impl fmt::Debug for BackendAuth { @@ -43,6 +46,7 @@ impl fmt::Debug for BackendAuth { match self { Self::BearerToken(_) => write!(f, "BearerToken([REDACTED])"), Self::GoogleApiKey(_) => write!(f, "GoogleApiKey([REDACTED])"), + Self::AzureApiKey(_) => write!(f, "AzureApiKey([REDACTED])"), } } } @@ -84,11 +88,12 @@ impl Config { let backend_str = std::env::var("BACKEND").unwrap_or_else(|_| "openai".into()); let backend = match backend_str.to_ascii_lowercase().as_str() { "openai" => BackendKind::OpenAI, + "azure" => BackendKind::AzureOpenAI, "vertex" => BackendKind::Vertex, "gemini" => BackendKind::Gemini, "anthropic" => BackendKind::Anthropic, other => { - panic!("unknown BACKEND value '{other}', expected 'openai', 'vertex', 'gemini', or 'anthropic'") + panic!("unknown BACKEND value '{other}', expected 'openai', 'azure', 'vertex', 'gemini', or 'anthropic'") } }; @@ -133,6 +138,43 @@ impl Config { openai_api_format, } } + BackendKind::AzureOpenAI => { + let endpoint = std::env::var("AZURE_OPENAI_ENDPOINT").unwrap_or_else(|_| { + panic!("AZURE_OPENAI_ENDPOINT is required when BACKEND=azure") + }); + let deployment = std::env::var("AZURE_OPENAI_DEPLOYMENT").unwrap_or_else(|_| { + panic!("AZURE_OPENAI_DEPLOYMENT is required when BACKEND=azure") + }); + let api_key = std::env::var("AZURE_OPENAI_API_KEY").unwrap_or_else(|_| { + panic!("AZURE_OPENAI_API_KEY is required when BACKEND=azure") + }); + let api_version = std::env::var("AZURE_OPENAI_API_VERSION") + .unwrap_or_else(|_| "2024-10-21".to_string()); + + // Pre-construct the full URL; no suffix is appended by OpenAIClient. + let base_url = format!( + "{}/openai/deployments/{}/chat/completions?api-version={}", + endpoint.trim_end_matches('/'), + deployment, + api_version + ); + // Validate the endpoint (not the full URL, which has query params) + if let Err(e) = validate_base_url(endpoint.trim_end_matches('/')) { + panic!("AZURE_OPENAI_ENDPOINT rejected: {e}"); + } + + Self { + backend, + openai_api_key: String::new(), + openai_base_url: base_url, + listen_port, + model_mapping: ModelMapping::from_env_with_defaults("gpt-4o", "gpt-4o-mini"), + tls, + backend_auth: BackendAuth::AzureApiKey(api_key), + log_bodies, + openai_api_format: OpenAIApiFormat::Chat, + } + } BackendKind::Vertex => { let project = std::env::var("VERTEX_PROJECT") .unwrap_or_else(|_| panic!("VERTEX_PROJECT is required when BACKEND=vertex")); @@ -353,6 +395,10 @@ struct TomlBackendConfig { // Vertex-specific project: Option, region: Option, + // Azure-specific + endpoint: Option, + deployment: Option, + api_version: Option, // Optional env var name for Google access token (Vertex) access_token: Option, // Strip stream_options from streaming requests (local LLM compat) @@ -385,6 +431,7 @@ impl MultiConfig { fn wrap_config(config: &Config) -> Self { let name = match config.backend { BackendKind::OpenAI => "openai", + BackendKind::AzureOpenAI => "azure", BackendKind::Vertex => "vertex", BackendKind::Gemini => "gemini", BackendKind::Anthropic => "anthropic", @@ -470,6 +517,7 @@ impl MultiConfig { ) -> BackendConfig { let kind = match tb.kind.to_ascii_lowercase().as_str() { "openai" => BackendKind::OpenAI, + "azure" => BackendKind::AzureOpenAI, "vertex" => BackendKind::Vertex, "gemini" => BackendKind::Gemini, "anthropic" => BackendKind::Anthropic, @@ -512,6 +560,38 @@ impl MultiConfig { }; (base_url, auth, mm, fmt) } + BackendKind::AzureOpenAI => { + if api_key.is_empty() { + panic!("backend '{name}': api_key is required for azure"); + } + let endpoint = tb.endpoint.as_deref().unwrap_or_else(|| { + panic!("backend '{name}': 'endpoint' is required for azure") + }); + let deployment = tb.deployment.as_deref().unwrap_or_else(|| { + panic!("backend '{name}': 'deployment' is required for azure") + }); + let api_version = tb.api_version.as_deref().unwrap_or("2024-10-21"); + + if let Err(e) = validate_base_url(endpoint.trim_end_matches('/')) { + panic!("backend '{name}' endpoint rejected: {e}"); + } + + let base_url = format!( + "{}/openai/deployments/{}/chat/completions?api-version={}", + endpoint.trim_end_matches('/'), + deployment, + api_version + ); + let auth = BackendAuth::AzureApiKey(api_key.clone()); + let mm = ModelMapping { + big_model: tb.big_model.clone().unwrap_or_else(|| "gpt-4o".to_string()), + small_model: tb + .small_model + .clone() + .unwrap_or_else(|| "gpt-4o-mini".to_string()), + }; + (base_url, auth, mm, OpenAIApiFormat::Chat) + } BackendKind::Vertex => { let project = tb.project.as_deref().unwrap_or_else(|| { panic!("backend '{name}': 'project' is required for vertex") @@ -724,6 +804,11 @@ mod tests { let debug = format!("{:?}", api_key); assert!(debug.contains("REDACTED")); assert!(!debug.contains("secret-key")); + + let azure_key = BackendAuth::AzureApiKey("azure-secret".into()); + let debug = format!("{:?}", azure_key); + assert!(debug.contains("REDACTED")); + assert!(!debug.contains("azure-secret")); } // --- MultiConfig TOML parsing tests --- @@ -964,4 +1049,76 @@ mod tests { "https://generativelanguage.googleapis.com/v1beta/openai" ); } + + // --- Azure OpenAI tests --- + + #[test] + fn multi_config_parses_azure_backend() { + let toml = r#" + [backends.azure] + kind = "azure" + api_key = "az-test-key" + endpoint = "https://my-resource.openai.azure.com" + deployment = "gpt-4o-deploy" + "#; + let mc = MultiConfig::from_toml_str(toml); + let bc = &mc.backends["azure"]; + assert_eq!(bc.kind, BackendKind::AzureOpenAI); + assert_eq!( + bc.base_url, + "https://my-resource.openai.azure.com/openai/deployments/gpt-4o-deploy/chat/completions?api-version=2024-10-21" + ); + assert!(matches!(bc.backend_auth, BackendAuth::AzureApiKey(_))); + } + + #[test] + fn multi_config_azure_custom_api_version() { + let toml = r#" + [backends.azure] + kind = "azure" + api_key = "az-test-key" + endpoint = "https://my-resource.openai.azure.com" + deployment = "gpt-4o-deploy" + api_version = "2025-01-01" + "#; + let mc = MultiConfig::from_toml_str(toml); + let bc = &mc.backends["azure"]; + assert!(bc.base_url.contains("api-version=2025-01-01")); + } + + #[test] + #[should_panic(expected = "api_key is required for azure")] + fn multi_config_panics_azure_no_key() { + let toml = r#" + [backends.azure] + kind = "azure" + endpoint = "https://my-resource.openai.azure.com" + deployment = "gpt-4o-deploy" + "#; + MultiConfig::from_toml_str(toml); + } + + #[test] + #[should_panic(expected = "endpoint' is required for azure")] + fn multi_config_panics_azure_no_endpoint() { + let toml = r#" + [backends.azure] + kind = "azure" + api_key = "az-test-key" + deployment = "gpt-4o-deploy" + "#; + MultiConfig::from_toml_str(toml); + } + + #[test] + #[should_panic(expected = "deployment' is required for azure")] + fn multi_config_panics_azure_no_deployment() { + let toml = r#" + [backends.azure] + kind = "azure" + api_key = "az-test-key" + endpoint = "https://my-resource.openai.azure.com" + "#; + MultiConfig::from_toml_str(toml); + } } diff --git a/crates/proxy/src/main.rs b/crates/proxy/src/main.rs index 7533185..9dd439e 100644 --- a/crates/proxy/src/main.rs +++ b/crates/proxy/src/main.rs @@ -4,6 +4,24 @@ use tracing_subscriber::prelude::*; #[tokio::main] async fn main() { + // Load env file before anything else so RUST_LOG and backend config are visible. + // Explicit --env-file takes priority; otherwise auto-load .anyllm.env if present. + let args: Vec = std::env::args().collect(); + let env_file_path = args + .windows(2) + .find(|w| w[0] == "--env-file") + .map(|w| w[1].as_str()) + .or_else(|| { + if std::path::Path::new(".anyllm.env").exists() { + Some(".anyllm.env") + } else { + None + } + }); + if let Some(path) = env_file_path { + load_env_file(path); + } + // Use a reload layer so the admin API can change log_level at runtime. let env_filter = tracing_subscriber::EnvFilter::from_default_env(); let (filter, reload_handle) = tracing_subscriber::reload::Layer::new(env_filter); @@ -21,218 +39,265 @@ async fn main() { "configured backends" ); - // --- Admin setup --- - let admin_port: u16 = std::env::var("ADMIN_PORT") - .ok() - .and_then(|p| p.parse().ok()) - .unwrap_or(3001); + // Admin web UI is opt-in: pass --webui or --admin to enable. + // DISABLE_ADMIN=1 overrides the flag (useful in container/scripted environments). + let flag_set = args.iter().any(|a| a == "--webui" || a == "--admin"); + let force_disabled = matches!( + std::env::var("DISABLE_ADMIN").as_deref(), + Ok("1") | Ok("true") | Ok("yes") + ); + let enable_admin = flag_set && !force_disabled; - if admin_port == listen_port { - panic!("ADMIN_PORT ({admin_port}) must differ from LISTEN_PORT ({listen_port})"); - } + // --- Admin setup (enabled only when --webui or --admin flag is passed) --- + // Returns Some((SharedState, admin Router, admin TcpListener)) when enabled. + let admin_parts = if enable_admin { + let admin_port: u16 = std::env::var("ADMIN_PORT") + .ok() + .and_then(|p| p.parse().ok()) + .unwrap_or(3001); - // SQLite: open or create the database file in the current directory. - let db_path = std::env::var("ADMIN_DB_PATH").unwrap_or_else(|_| "admin.db".into()); - let conn = - rusqlite::Connection::open(&db_path).expect("failed to open SQLite database for admin"); - admin::db::init_db(&conn).expect("failed to initialize admin database schema"); + if admin_port == listen_port { + panic!("ADMIN_PORT ({admin_port}) must differ from LISTEN_PORT ({listen_port})"); + } - // Build initial RuntimeConfig from the loaded multi_config. - let mut model_mappings = indexmap::IndexMap::new(); - for (name, bc) in &multi_config.backends { - model_mappings.insert(name.clone(), bc.model_mapping.clone()); - } - let log_level = std::env::var("RUST_LOG").unwrap_or_else(|_| "info".into()); - let mut runtime_config = admin::state::RuntimeConfig { - model_mappings, - log_level, - log_bodies: multi_config.log_bodies, - }; + // SQLite: open or create the database file in the current directory. + let db_path = std::env::var("ADMIN_DB_PATH").unwrap_or_else(|_| "admin.db".into()); + let conn = + rusqlite::Connection::open(&db_path).expect("failed to open SQLite database for admin"); + admin::db::init_db(&conn).expect("failed to initialize admin database schema"); - // Apply config overrides from SQLite (survive restarts). - if let Ok(overrides) = admin::db::get_config_overrides(&conn) { - for (key, value, _) in &overrides { - match key.as_str() { - "log_level" => { - // Apply the same allowlist enforced by the admin API to - // prevent a tampered SQLite database from enabling trace-level - // logging, which would expose API keys in HTTP headers. - const ALLOWED_LOG_LEVELS: &[&str] = &["error", "warn", "info", "debug"]; - let normalized = value.trim().to_lowercase(); - if ALLOWED_LOG_LEVELS.contains(&normalized.as_str()) { - runtime_config.log_level = normalized; - } else { - tracing::warn!( - value = %value, - "ignoring invalid log_level override from database" + // Build initial RuntimeConfig from the loaded multi_config. + let mut model_mappings = indexmap::IndexMap::new(); + for (name, bc) in &multi_config.backends { + model_mappings.insert(name.clone(), bc.model_mapping.clone()); + } + let log_level = std::env::var("RUST_LOG").unwrap_or_else(|_| "info".into()); + let mut runtime_config = admin::state::RuntimeConfig { + model_mappings, + log_level, + log_bodies: multi_config.log_bodies, + }; + + // Apply config overrides from SQLite (survive restarts). + if let Ok(overrides) = admin::db::get_config_overrides(&conn) { + for (key, value, _) in &overrides { + match key.as_str() { + "log_level" => { + // Apply the same allowlist enforced by the admin API to + // prevent a tampered SQLite database from enabling trace-level + // logging, which would expose API keys in HTTP headers. + const ALLOWED_LOG_LEVELS: &[&str] = &["error", "warn", "info", "debug"]; + let normalized = value.trim().to_lowercase(); + if ALLOWED_LOG_LEVELS.contains(&normalized.as_str()) { + runtime_config.log_level = normalized; + } else { + tracing::warn!( + value = %value, + "ignoring invalid log_level override from database" + ); + } + } + "log_bodies" => runtime_config.log_bodies = value == "true", + k if k.ends_with(".big_model") => { + let backend = k.strip_suffix(".big_model").unwrap(); + if let Some(m) = runtime_config.model_mappings.get_mut(backend) { + m.big_model = value.clone(); + } + } + k if k.ends_with(".small_model") => { + let backend = k.strip_suffix(".small_model").unwrap(); + if let Some(m) = runtime_config.model_mappings.get_mut(backend) { + m.small_model = value.clone(); + } + } + _ => { + tracing::debug!(key = %key, "unknown config override, skipping"); + } + } + } + if !overrides.is_empty() { + tracing::info!( + count = overrides.len(), + "applied config overrides from database" + ); + } + } + let runtime_config = Arc::new(std::sync::RwLock::new(runtime_config)); + + // Build the log_reload closure that captures the reload handle. + let log_reload: Arc bool + Send + Sync> = { + let handle = reload_handle; + Arc::new(move |new_filter: &str| { + match tracing_subscriber::EnvFilter::try_new(new_filter) { + Ok(f) => handle.reload(f).is_ok(), + Err(e) => { + tracing::error!(filter = new_filter, error = %e, "invalid log filter string"); + false + } + } + }) + }; + + // Now wrap conn in Arc and start the write buffer. + // Uses std::sync::Mutex because rusqlite is synchronous; all access + // goes through spawn_blocking to avoid stalling the tokio executor. + let db = Arc::new(std::sync::Mutex::new(conn)); + let (events_tx, _) = tokio::sync::broadcast::channel(1024); + let log_tx = admin::db::spawn_write_buffer(db.clone()); + + let backend_metrics: std::collections::HashMap = + std::collections::HashMap::new(); + + // Load active virtual keys from SQLite into in-memory DashMap. + let virtual_keys = Arc::new(dashmap::DashMap::new()); + { + let conn_guard = db.lock().unwrap_or_else(|e| e.into_inner()); + if let Ok(active_keys) = admin::db::load_active_virtual_keys(&conn_guard) { + for key_row in &active_keys { + if let Some(hash_bytes) = admin::keys::hash_from_hex(&key_row.key_hash) { + virtual_keys.insert( + hash_bytes, + admin::keys::VirtualKeyMeta { + id: key_row.id, + description: key_row.description.clone(), + expires_at: None, // Expiry checked from the ISO string at auth time + rpm_limit: key_row.rpm_limit, + tpm_limit: key_row.tpm_limit, + rate_state: Arc::new(admin::keys::RateLimitState::new()), + }, ); } } - "log_bodies" => runtime_config.log_bodies = value == "true", - k if k.ends_with(".big_model") => { - let backend = k.strip_suffix(".big_model").unwrap(); - if let Some(m) = runtime_config.model_mappings.get_mut(backend) { - m.big_model = value.clone(); - } - } - k if k.ends_with(".small_model") => { - let backend = k.strip_suffix(".small_model").unwrap(); - if let Some(m) = runtime_config.model_mappings.get_mut(backend) { - m.small_model = value.clone(); - } - } - _ => { - tracing::debug!(key = %key, "unknown config override, skipping"); - } + tracing::info!( + count = active_keys.len(), + "loaded virtual API keys from database" + ); } } - if !overrides.is_empty() { - tracing::info!( - count = overrides.len(), - "applied config overrides from database" - ); - } - } - let runtime_config = Arc::new(std::sync::RwLock::new(runtime_config)); - // Build the log_reload closure that captures the reload handle. - let log_reload: Arc bool + Send + Sync> = { - let handle = reload_handle; - Arc::new( - move |new_filter: &str| match tracing_subscriber::EnvFilter::try_new(new_filter) { - Ok(f) => handle.reload(f).is_ok(), - Err(e) => { - tracing::error!(filter = new_filter, error = %e, "invalid log filter string"); - false + // Make virtual keys available to the auth middleware. + anyllm_proxy::server::middleware::set_virtual_keys(virtual_keys.clone()); + + let shared = admin::state::SharedState { + db: db.clone(), + events_tx: events_tx.clone(), + runtime_config: runtime_config.clone(), + backend_metrics: Arc::new(backend_metrics), + log_tx, + log_reload: Some(log_reload), + config_write_lock: Arc::new(tokio::sync::Mutex::new(())), + virtual_keys, + }; + + // Admin token: use env var or generate random UUID written to a file. + let admin_token = std::env::var("ADMIN_TOKEN").unwrap_or_else(|_| { + let token = uuid::Uuid::new_v4().to_string(); + let token_path = std::env::var("ADMIN_TOKEN_FILE") + .unwrap_or_else(|_| ".admin_token".into()); + // Write token to file with restrictive permissions instead of stderr, + // because stderr is captured by container log drivers in production. + if let Err(e) = write_token_file(&token_path, &token) { + // Do not print the token to stderr: container log drivers capture + // stderr and persist it in centralized logging systems. + panic!( + "Cannot write admin token to {token_path}: {e}. \ + Set ADMIN_TOKEN env var explicitly or ensure the path is writable." + ); + } else { + // Log the path, not the token itself. + tracing::info!(path = %token_path, "generated admin token written to file (set ADMIN_TOKEN env var to avoid this)"); + } + token + }); + let admin_token = Arc::new(admin_token); + + // Spawn periodic tasks: log retention and metrics snapshot broadcast. + let retention_days: u32 = std::env::var("ADMIN_LOG_RETENTION_DAYS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(7); + + let retention_db = shared.db.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(3600)); + loop { + interval.tick().await; + admin::state::with_db(&retention_db, move |conn| { + match admin::db::purge_old_logs(conn, retention_days) { + Ok(n) if n > 0 => { + tracing::info!(purged = n, "purged old request log entries") + } + Err(e) => tracing::error!(error = %e, "failed to purge old logs"), + _ => {} + } + }) + .await; + } + }); + + // Periodic metrics snapshot broadcast (every 5 seconds) for WebSocket dashboard. + let snapshot_shared = shared.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(5)); + loop { + interval.tick().await; + // Skip computation if no WebSocket clients are listening. + if snapshot_shared.events_tx.receiver_count() == 0 { + continue; } - }, - ) + let mut backends = std::collections::HashMap::new(); + let mut aggregate = anyllm_proxy::metrics::MetricsSnapshot::default(); + for (name, m) in snapshot_shared.backend_metrics.iter() { + let snap = m.snapshot(); + aggregate.requests_total += snap.requests_total; + aggregate.requests_error += snap.requests_error; + aggregate.requests_success += snap.requests_success; + backends.insert(name.clone(), snap); + } + let error_rate = aggregate.error_rate(); + let snapshot = admin::state::MetricsSnapshotData { + backends, + latency_p50_ms: None, // Computed on demand by REST endpoint + latency_p95_ms: None, + latency_p99_ms: None, + requests_per_second: 0.0, // TODO: compute from recent request log + error_rate, + }; + let _ = snapshot_shared + .events_tx + .send(admin::state::AdminEvent::MetricsSnapshot(snapshot)); + } + }); + + // Bind admin listener; spawned after the shutdown channel is created below. + let admin_app = admin::routes::admin_router(shared.clone(), admin_token); + let admin_addr = format!("127.0.0.1:{admin_port}"); + let admin_listener = tokio::net::TcpListener::bind(&admin_addr) + .await + .unwrap_or_else(|e| panic!("failed to bind admin to {admin_addr}: {e}")); + tracing::info!("admin listening on {admin_addr}"); + + Some((shared, admin_app, admin_listener)) + } else { + None }; - // Now wrap conn in Arc and start the write buffer. - // Uses std::sync::Mutex because rusqlite is synchronous; all access - // goes through spawn_blocking to avoid stalling the tokio executor. - let db = Arc::new(std::sync::Mutex::new(conn)); - let (events_tx, _) = tokio::sync::broadcast::channel(1024); - let log_tx = admin::db::spawn_write_buffer(db.clone()); + // Build proxy router with optional shared admin state. + let app = routes::app_multi_with_shared( + multi_config, + admin_parts.as_ref().map(|(s, _, _)| s.clone()), + ); - let backend_metrics: std::collections::HashMap = - std::collections::HashMap::new(); - - let shared = admin::state::SharedState { - db: db.clone(), - events_tx: events_tx.clone(), - runtime_config: runtime_config.clone(), - backend_metrics: Arc::new(backend_metrics), - log_tx, - log_reload: Some(log_reload), - config_write_lock: Arc::new(tokio::sync::Mutex::new(())), - }; - - // Admin token: use env var or generate random UUID written to a file. - let admin_token = std::env::var("ADMIN_TOKEN").unwrap_or_else(|_| { - let token = uuid::Uuid::new_v4().to_string(); - let token_path = std::env::var("ADMIN_TOKEN_FILE") - .unwrap_or_else(|_| ".admin_token".into()); - // Write token to file with restrictive permissions instead of stderr, - // because stderr is captured by container log drivers in production. - if let Err(e) = write_token_file(&token_path, &token) { - // Do not print the token to stderr: container log drivers capture - // stderr and persist it in centralized logging systems. - panic!( - "Cannot write admin token to {token_path}: {e}. \ - Set ADMIN_TOKEN env var explicitly or ensure the path is writable." - ); - } else { - // Log the path, not the token itself. - tracing::info!(path = %token_path, "generated admin token written to file (set ADMIN_TOKEN env var to avoid this)"); - } - token - }); - let admin_token = Arc::new(admin_token); - - // Spawn periodic tasks: log retention and metrics snapshot broadcast. - let retention_days: u32 = std::env::var("ADMIN_LOG_RETENTION_DAYS") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(7); - - let retention_db = shared.db.clone(); - tokio::spawn(async move { - let mut interval = tokio::time::interval(std::time::Duration::from_secs(3600)); - loop { - interval.tick().await; - admin::state::with_db(&retention_db, move |conn| { - match admin::db::purge_old_logs(conn, retention_days) { - Ok(n) if n > 0 => { - tracing::info!(purged = n, "purged old request log entries") - } - Err(e) => tracing::error!(error = %e, "failed to purge old logs"), - _ => {} - } - }) - .await; - } - }); - - // Periodic metrics snapshot broadcast (every 5 seconds) for WebSocket dashboard. - let snapshot_shared = shared.clone(); - tokio::spawn(async move { - let mut interval = tokio::time::interval(std::time::Duration::from_secs(5)); - loop { - interval.tick().await; - // Skip computation if no WebSocket clients are listening. - if snapshot_shared.events_tx.receiver_count() == 0 { - continue; - } - let mut backends = std::collections::HashMap::new(); - let mut aggregate = anyllm_proxy::metrics::MetricsSnapshot::default(); - for (name, m) in snapshot_shared.backend_metrics.iter() { - let snap = m.snapshot(); - aggregate.requests_total += snap.requests_total; - aggregate.requests_error += snap.requests_error; - aggregate.requests_success += snap.requests_success; - backends.insert(name.clone(), snap); - } - let error_rate = aggregate.error_rate(); - let snapshot = admin::state::MetricsSnapshotData { - backends, - latency_p50_ms: None, // Computed on demand by REST endpoint - latency_p95_ms: None, - latency_p99_ms: None, - requests_per_second: 0.0, // TODO: compute from recent request log - error_rate, - }; - let _ = snapshot_shared - .events_tx - .send(admin::state::AdminEvent::MetricsSnapshot(snapshot)); - } - }); - - // Build proxy router with shared state. - let app = routes::app_multi_with_shared(multi_config, Some(shared.clone())); - - // Build admin router. - let admin_app = admin::routes::admin_router(shared, admin_token); - - // --- Start both servers --- + // --- Start servers --- let proxy_addr = format!("0.0.0.0:{listen_port}"); let proxy_listener = tokio::net::TcpListener::bind(&proxy_addr) .await .unwrap_or_else(|e| panic!("failed to bind proxy to {proxy_addr}: {e}")); tracing::info!("proxy listening on {proxy_addr}"); - let admin_addr = format!("127.0.0.1:{admin_port}"); - let admin_listener = tokio::net::TcpListener::bind(&admin_addr) - .await - .unwrap_or_else(|e| panic!("failed to bind admin to {admin_addr}: {e}")); - tracing::info!("admin listening on {admin_addr}"); - - // Share the shutdown signal between both servers via a tokio::sync::watch. + // Single shutdown channel shared by proxy and (optionally) admin. let (shutdown_tx, mut shutdown_rx1) = tokio::sync::watch::channel(false); - let mut shutdown_rx2 = shutdown_tx.subscribe(); - // Spawn the proxy server. let proxy_handle = tokio::spawn(async move { axum::serve(proxy_listener, app) .with_graceful_shutdown(async move { @@ -242,25 +307,91 @@ async fn main() { .expect("proxy server error"); }); - // Spawn the admin server. - let admin_handle = tokio::spawn(async move { - axum::serve(admin_listener, admin_app) - .with_graceful_shutdown(async move { - shutdown_rx2.changed().await.ok(); - }) - .await - .expect("admin server error"); - }); + let admin_handle: Option> = + if let Some((_, admin_app, admin_listener)) = admin_parts { + let mut shutdown_rx2 = shutdown_tx.subscribe(); + Some(tokio::spawn(async move { + axum::serve(admin_listener, admin_app) + .with_graceful_shutdown(async move { + shutdown_rx2.changed().await.ok(); + }) + .await + .expect("admin server error"); + })) + } else { + None + }; - // Wait for shutdown signal, then notify both servers. shutdown_signal().await; let _ = shutdown_tx.send(true); - // Wait for both servers to finish. - let _ = tokio::join!(proxy_handle, admin_handle); + let _ = proxy_handle.await; + if let Some(h) = admin_handle { + let _ = h.await; + } tracing::info!("server shut down gracefully"); } +/// Load a `.env`-format file and apply values to the process environment. +/// +/// Rules: +/// - `KEY=VALUE` sets the variable. Surrounding whitespace is trimmed. +/// - Values may be optionally wrapped in `"double"` or `'single'` quotes. +/// - Lines starting with `#` (after trimming) are comments. +/// - Already-set environment variables are never overwritten; the real +/// environment always takes precedence over the file. +/// - `export KEY=VALUE` syntax is supported (the `export` prefix is stripped). +/// +/// Compatible with Docker `--env-file` and standard dotenv tooling. +fn load_env_file(path: &str) { + let content = match std::fs::read_to_string(path) { + Ok(c) => c, + Err(e) => { + // Print directly; tracing isn't initialized yet. + eprintln!("anyllm_proxy: could not read env file '{path}': {e}"); + return; + } + }; + + let mut loaded = 0usize; + for (lineno, raw) in content.lines().enumerate() { + let line = raw.trim(); + if line.is_empty() || line.starts_with('#') { + continue; + } + // Strip optional `export ` prefix. + let line = line.strip_prefix("export ").map(str::trim).unwrap_or(line); + let Some((key, val)) = line.split_once('=') else { + eprintln!( + "anyllm_proxy: {path}:{}: ignoring malformed line (no '=')", + lineno + 1 + ); + continue; + }; + let key = key.trim(); + if key.is_empty() { + continue; + } + // Strip optional surrounding quotes from the value. + let val = val.trim(); + let val = if (val.starts_with('"') && val.ends_with('"')) + || (val.starts_with('\'') && val.ends_with('\'')) + { + &val[1..val.len() - 1] + } else { + val + }; + // Only set if not already present so the real environment wins. + if std::env::var(key).is_err() { + // SAFETY: called before any threads are spawned (before tokio runtime). + #[allow(deprecated)] + std::env::set_var(key, val); + loaded += 1; + } + } + eprintln!("anyllm_proxy: loaded {loaded} variable(s) from '{path}'"); +} + /// Write the admin token to a file with mode 0600 (owner-only read/write). /// On Unix, sets permissions atomically at creation to avoid a TOCTOU race /// where the file is briefly world-readable before chmod. diff --git a/crates/proxy/src/server/chat_completions.rs b/crates/proxy/src/server/chat_completions.rs new file mode 100644 index 0000000..aaf75f7 --- /dev/null +++ b/crates/proxy/src/server/chat_completions.rs @@ -0,0 +1,429 @@ +// OpenAI Chat Completions input handler. +// +// Accepts POST /v1/chat/completions in OpenAI format, translates through +// the Anthropic pipeline, returns OpenAI-format responses. + +use crate::backend::{find_double_newline, BackendClient, BackendError, MAX_SSE_BUFFER_SIZE}; +use anyllm_translate::{ + anthropic, mapping, openai, translate_anthropic_to_openai_response, + translate_openai_to_anthropic_request, ReverseStreamingTranslator, TranslationWarnings, +}; +use axum::{ + extract::State, + http::StatusCode, + response::{IntoResponse, Json, Response}, +}; +use bytes::BytesMut; +use futures::StreamExt; + +use super::routes::{ + inject_degradation_header, log_request, AppState, ConcurrencyPermit, RequestCtx, +}; + +/// OpenAI-shaped error response body. +fn openai_error_response(message: &str, error_type: &str, status: StatusCode) -> Response { + let body = serde_json::json!({ + "error": { + "message": message, + "type": error_type, + "param": null, + "code": null + } + }); + (status, Json(body)).into_response() +} + +/// Convert a BackendError into an OpenAI-shaped error response. +fn backend_error_to_openai_response(error: BackendError) -> Response { + if let Some((message, status)) = error.api_error_details() { + let error_type = if status == 429 { + "rate_limit_error" + } else if status >= 500 { + "server_error" + } else { + "invalid_request_error" + }; + let http_status = StatusCode::from_u16(status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + return openai_error_response(message, error_type, http_status); + } + tracing::error!("backend client error: {error}"); + openai_error_response( + "An internal error occurred while communicating with the upstream service.", + "server_error", + StatusCode::INTERNAL_SERVER_ERROR, + ) +} + +/// Handler for POST /v1/chat/completions (non-streaming and streaming). +pub(crate) async fn chat_completions( + State(state): State, + headers: axum::http::HeaderMap, + permit: Option>, + body: Result, axum::extract::rejection::JsonRejection>, +) -> Response { + let body = match body { + Ok(Json(b)) => b, + Err(e) => { + return openai_error_response( + &e.body_text(), + "invalid_request_error", + StatusCode::BAD_REQUEST, + ); + } + }; + + let permit = permit.map(|axum::Extension(p)| p); + let ctx = RequestCtx { + request_id: headers + .get("x-request-id") + .and_then(|v| v.to_str().ok()) + .unwrap_or("unknown") + .to_string(), + start: std::time::Instant::now(), + model_requested: body.model.clone(), + }; + state.metrics.record_request(); + + // Translate OpenAI request -> Anthropic request + let mut warnings = TranslationWarnings::default(); + let anthropic_req = match translate_openai_to_anthropic_request(&body, &mut warnings) { + Ok(req) => req, + Err(e) => { + return openai_error_response( + &e.to_string(), + "invalid_request_error", + StatusCode::BAD_REQUEST, + ); + } + }; + + if anthropic_req.messages.is_empty() { + return openai_error_response( + "messages array must not be empty", + "invalid_request_error", + StatusCode::BAD_REQUEST, + ); + } + + let is_streaming = body.stream == Some(true); + let original_model = body.model.clone(); + + if is_streaming { + return chat_completions_stream( + state, + anthropic_req, + ctx, + original_model, + warnings, + permit, + ) + .await; + } + + // Non-streaming path + match &state.backend { + BackendClient::OpenAI(client) + | BackendClient::AzureOpenAI(client) + | BackendClient::Vertex(client) + | BackendClient::GeminiOpenAI(client) => { + let mut openai_req = mapping::message_map::anthropic_to_openai_request(&anthropic_req); + super::routes::inject_gemini_thinking(&anthropic_req, &state.backend, &mut openai_req); + if state.omit_stream_options { + openai_req.stream_options = None; + } + openai_req.model = state.map_model(&openai_req.model); + let mapped_model = openai_req.model.clone(); + + match client.chat_completion(&openai_req).await { + Ok((openai_resp, _status, rate_limits)) => { + state.metrics.record_success(); + // Translate Anthropic response back to OpenAI format + let anthropic_resp = mapping::message_map::openai_to_anthropic_response( + &openai_resp, + &original_model, + ); + let oai_response = + translate_anthropic_to_openai_response(&anthropic_resp, &original_model); + log_request( + &state.shared, + ctx.log_entry( + &state.backend_name, + Some(mapped_model), + 200, + Some(( + anthropic_resp.usage.input_tokens as u64, + anthropic_resp.usage.output_tokens as u64, + )), + false, + None, + ), + ); + let mut response = (StatusCode::OK, Json(oai_response)).into_response(); + rate_limits.inject_anthropic_response_headers(response.headers_mut()); + inject_degradation_header(response.headers_mut(), &warnings); + response + } + Err(e) => { + state.metrics.record_error(); + let status = e.status_code(); + log_request( + &state.shared, + ctx.log_entry( + &state.backend_name, + Some(mapped_model), + status, + None, + false, + Some(e.to_string()), + ), + ); + backend_error_to_openai_response(BackendError::from(e)) + } + } + } + BackendClient::OpenAIResponses(client) => { + let mut responses_req = + mapping::responses_message_map::anthropic_to_responses_request(&anthropic_req); + responses_req.model = state.map_model(&responses_req.model); + let mapped_model = responses_req.model.clone(); + + match client.responses(&responses_req).await { + Ok((resp, _status, rate_limits)) => { + state.metrics.record_success(); + let anthropic_resp = + mapping::responses_message_map::responses_to_anthropic_response( + &resp, + &original_model, + ); + let oai_response = + translate_anthropic_to_openai_response(&anthropic_resp, &original_model); + log_request( + &state.shared, + ctx.log_entry( + &state.backend_name, + Some(mapped_model), + 200, + Some(( + anthropic_resp.usage.input_tokens as u64, + anthropic_resp.usage.output_tokens as u64, + )), + false, + None, + ), + ); + let mut response = (StatusCode::OK, Json(oai_response)).into_response(); + rate_limits.inject_anthropic_response_headers(response.headers_mut()); + inject_degradation_header(response.headers_mut(), &warnings); + response + } + Err(e) => { + state.metrics.record_error(); + let status = e.status_code(); + log_request( + &state.shared, + ctx.log_entry( + &state.backend_name, + Some(mapped_model), + status, + None, + false, + Some(e.to_string()), + ), + ); + backend_error_to_openai_response(BackendError::from(e)) + } + } + } + BackendClient::Anthropic(_) => openai_error_response( + "Anthropic passthrough backend does not support /v1/chat/completions", + "invalid_request_error", + StatusCode::BAD_REQUEST, + ), + } +} + +/// Streaming handler for POST /v1/chat/completions with stream: true. +/// +/// Translates the Anthropic request to OpenAI, streams the backend response, +/// then uses ReverseStreamingTranslator to convert Anthropic SSE events back +/// to OpenAI ChatCompletionChunk SSE format. +async fn chat_completions_stream( + state: AppState, + anthropic_req: anthropic::MessageCreateRequest, + ctx: RequestCtx, + original_model: String, + warnings: TranslationWarnings, + concurrency_permit: Option, +) -> Response { + // Translate to OpenAI format for the backend + let mut openai_req = mapping::message_map::anthropic_to_openai_request(&anthropic_req); + super::routes::inject_gemini_thinking(&anthropic_req, &state.backend, &mut openai_req); + if state.omit_stream_options { + openai_req.stream_options = None; + } + openai_req.model = state.map_model(&openai_req.model); + openai_req.stream = Some(true); + openai_req.stream_options = Some(openai::StreamOptions { + include_usage: true, + }); + + let client = match &state.backend { + BackendClient::OpenAI(c) + | BackendClient::AzureOpenAI(c) + | BackendClient::Vertex(c) + | BackendClient::GeminiOpenAI(c) + | BackendClient::OpenAIResponses(c) => c.clone(), + BackendClient::Anthropic(_) => { + return openai_error_response( + "Anthropic passthrough backend does not support /v1/chat/completions streaming", + "invalid_request_error", + StatusCode::BAD_REQUEST, + ); + } + }; + + let mapped_model = openai_req.model.clone(); + + // Start the backend request + let response = match client.chat_completion_stream(&openai_req).await { + Ok((resp, rate_limits)) => { + // Build the SSE response with OpenAI chunk format + let (tx, rx) = + tokio::sync::mpsc::channel::>(32); + let metrics = state.metrics.clone(); + let log_shared = state.shared.clone(); + let log_backend_name = state.backend_name.clone(); + let model_for_translator = original_model.clone(); + let _permit = concurrency_permit; + + tokio::spawn(async move { + let mut translator = ReverseStreamingTranslator::new( + format!("chatcmpl-{}", uuid::Uuid::new_v4().as_simple()), + model_for_translator.clone(), + ); + let mut stream_translator = + mapping::streaming_map::StreamingTranslator::new(model_for_translator.clone()); + + let mut byte_stream = resp.bytes_stream(); + let mut buffer = BytesMut::new(); + let mut search_from: usize = 0; + + while let Some(chunk_result) = byte_stream.next().await { + let bytes = match chunk_result { + Ok(b) => b, + Err(e) => { + tracing::error!("stream read error: {e}"); + metrics.record_error(); + break; + } + }; + buffer.extend_from_slice(&bytes); + + if buffer.len() > MAX_SSE_BUFFER_SIZE { + tracing::error!("SSE buffer exceeded maximum size"); + metrics.record_error(); + break; + } + + while let Some((pos, delim_len)) = find_double_newline(&buffer, search_from) { + if let Ok(frame_str) = std::str::from_utf8(&buffer[..pos]) { + for line in frame_str.lines() { + let line = line.trim(); + if let Some(json_str) = line.strip_prefix("data: ") { + if json_str == "[DONE]" { + // Emit [DONE] for OpenAI clients + let _ = tx.send(Ok("data: [DONE]\n\n".to_string())).await; + continue; + } + // Parse OpenAI chunk, translate to Anthropic events, + // then reverse-translate to OpenAI chunks + if let Ok(chunk) = + serde_json::from_str::( + json_str, + ) + { + let anthropic_events = + stream_translator.process_chunk(&chunk); + for event in &anthropic_events { + let oai_chunks = translator.process_event(event); + for oai_chunk in &oai_chunks { + if let Ok(json) = serde_json::to_string(oai_chunk) { + let sse_line = format!("data: {}\n\n", json); + if tx.send(Ok(sse_line)).await.is_err() { + return; // Client disconnected + } + } + } + } + } + } + } + } + let _ = buffer.split_to(pos + delim_len); + search_from = 0; + } + search_from = buffer.len().saturating_sub(3); + } + + // Emit any remaining finish events + let finish_events = stream_translator.finish(); + for event in &finish_events { + let oai_chunks = translator.process_event(event); + for oai_chunk in &oai_chunks { + if let Ok(json) = serde_json::to_string(oai_chunk) { + let _ = tx.send(Ok(format!("data: {}\n\n", json))).await; + } + } + } + + if !translator.is_done() { + let _ = tx.send(Ok("data: [DONE]\n\n".to_string())).await; + } + + metrics.record_success(); + log_request( + &log_shared, + ctx.log_entry( + &log_backend_name, + Some(mapped_model), + 200, + None, // Token counts come from usage chunk, hard to capture here + true, + None, + ), + ); + }); + + // Build the SSE response using raw text/event-stream + let body_stream = tokio_stream::wrappers::ReceiverStream::new(rx); + let body = axum::body::Body::from_stream(body_stream); + let mut response = Response::builder() + .status(StatusCode::OK) + .header("content-type", "text/event-stream") + .header("cache-control", "no-cache") + .header("connection", "keep-alive") + .body(body) + .unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response()); + rate_limits.inject_anthropic_response_headers(response.headers_mut()); + inject_degradation_header(response.headers_mut(), &warnings); + response + } + Err(e) => { + state.metrics.record_error(); + log_request( + &state.shared, + ctx.log_entry( + &state.backend_name, + Some(mapped_model), + e.status_code(), + None, + true, + Some(e.to_string()), + ), + ); + backend_error_to_openai_response(BackendError::from(e)) + } + }; + + response +} diff --git a/crates/proxy/src/server/middleware.rs b/crates/proxy/src/server/middleware.rs index b984085..01acfb0 100644 --- a/crates/proxy/src/server/middleware.rs +++ b/crates/proxy/src/server/middleware.rs @@ -1,5 +1,6 @@ // Auth, logging, and request size limit middleware +use crate::admin::keys::VirtualKeyMeta; use anyllm_translate::anthropic; use anyllm_translate::mapping::errors_map::create_anthropic_error; use axum::{ @@ -8,10 +9,20 @@ use axum::{ middleware::Next, response::{IntoResponse, Json, Response}, }; +use dashmap::DashMap; use sha2::{Digest, Sha256}; -use std::sync::LazyLock; +use std::sync::{Arc, LazyLock, OnceLock}; use subtle::ConstantTimeEq; +/// Global reference to the virtual keys DashMap, set once during startup. +/// Checked during auth after the static ALLOWED_KEY_HASHES check. +static VIRTUAL_KEYS: OnceLock>> = OnceLock::new(); + +/// Initialize the global virtual keys reference. Called once from main. +pub fn set_virtual_keys(keys: Arc>) { + let _ = VIRTUAL_KEYS.set(keys); +} + /// Pre-hashed allowed API keys for constant-time comparison without /// leaking key length via timing. Each key is SHA-256 hashed at startup. static ALLOWED_KEY_HASHES: LazyLock> = LazyLock::new(|| { @@ -89,28 +100,59 @@ pub async fn validate_auth( // Hashing eliminates the timing side-channel on key length: all comparisons // operate on fixed-size 32-byte digests regardless of original key length. let credential_hash: [u8; 32] = Sha256::digest(credential.as_bytes()).into(); - let is_allowed = ALLOWED_KEY_HASHES + + // Check 1: static env-var keys (constant-time comparison) + let env_key_match = ALLOWED_KEY_HASHES .iter() .any(|h| bool::from(h.ct_eq(&credential_hash))); - if !ALLOWED_KEY_HASHES.is_empty() && !is_allowed { - let err = create_anthropic_error( - anthropic::ErrorType::AuthenticationError, - "Invalid API key.".to_string(), - None, - ); - return Err((StatusCode::UNAUTHORIZED, Json(err)).into_response()); - } - // Reject if no keys configured and open-relay not explicitly enabled. - if ALLOWED_KEY_HASHES.is_empty() && !*OPEN_RELAY { - let err = create_anthropic_error( - anthropic::ErrorType::AuthenticationError, - "Server not configured for access. Contact the administrator.".to_string(), - None, - ); - return Err((StatusCode::UNAUTHORIZED, Json(err)).into_response()); + + if env_key_match { + return Ok(next.run(request).await); } - Ok(next.run(request).await) + // Check 2: virtual keys from DashMap (with per-key rate limiting) + if let Some(map) = VIRTUAL_KEYS.get() { + if let Some(meta) = map.get(&credential_hash) { + // Enforce RPM limit if configured + if let Some(rpm_limit) = meta.rpm_limit { + let now_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64; + if let Err(retry_after) = meta.rate_state.check_rpm(rpm_limit, now_ms) { + let err = create_anthropic_error( + anthropic::ErrorType::RateLimitError, + "Rate limit exceeded for this API key.".to_string(), + None, + ); + let mut resp = (StatusCode::TOO_MANY_REQUESTS, Json(err)).into_response(); + if let Ok(val) = axum::http::HeaderValue::from_str(&retry_after.to_string()) { + resp.headers_mut().insert("retry-after", val); + } + return Err(resp); + } + } + return Ok(next.run(request).await); + } + } + + // Check 3: open-relay mode (any non-empty key accepted) + if *OPEN_RELAY { + return Ok(next.run(request).await); + } + + // No match found: reject + let message = if ALLOWED_KEY_HASHES.is_empty() { + "Server not configured for access. Contact the administrator." + } else { + "Invalid API key." + }; + let err = create_anthropic_error( + anthropic::ErrorType::AuthenticationError, + message.to_string(), + None, + ); + Err((StatusCode::UNAUTHORIZED, Json(err)).into_response()) } /// Attach a request ID to the request and echo it on the response. diff --git a/crates/proxy/src/server/mod.rs b/crates/proxy/src/server/mod.rs index c431279..a176122 100644 --- a/crates/proxy/src/server/mod.rs +++ b/crates/proxy/src/server/mod.rs @@ -1,3 +1,5 @@ +/// OpenAI Chat Completions input handler (POST /v1/chat/completions). +mod chat_completions; /// Auth validation, request ID injection, size limits, concurrency limits, header logging. pub mod middleware; /// Anthropic passthrough handler (no translation, forwards as-is). diff --git a/crates/proxy/src/server/routes.rs b/crates/proxy/src/server/routes.rs index 71ecbb5..2d48ffe 100644 --- a/crates/proxy/src/server/routes.rs +++ b/crates/proxy/src/server/routes.rs @@ -2,7 +2,7 @@ use crate::admin::state::{AdminEvent, RequestLogEntry, RuntimeConfig, SharedStat use crate::backend::{BackendClient, BackendError}; use crate::config::{BackendKind, Config, MultiConfig}; use crate::metrics::Metrics; -use anyllm_translate::{anthropic, mapping, openai}; +use anyllm_translate::{anthropic, compute_request_warnings, mapping, openai}; use axum::{ extract::{rejection::JsonRejection, DefaultBodyLimit, FromRequest, State}, http::StatusCode, @@ -227,9 +227,14 @@ fn backend_router(state: AppState, is_anthropic: bool) -> Router { } else { Router::new() .route("/v1/messages", post(messages)) + .route( + "/v1/chat/completions", + post(super::chat_completions::chat_completions), + ) .route("/v1/models", get(models)) .route("/v1/messages/count_tokens", post(count_tokens)) .route("/v1/messages/batches", post(batches)) + .route("/v1/embeddings", post(embeddings)) }; api_routes @@ -279,13 +284,27 @@ pub(crate) struct ConcurrencyPermit( static MODELS_RESPONSE: std::sync::LazyLock = std::sync::LazyLock::new(|| { serde_json::json!({ "data": [ - {"id": "claude-opus-4-6", "display_name": "Claude Opus 4.6", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, - {"id": "claude-sonnet-4-6", "display_name": "Claude Sonnet 4.6", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, - {"id": "claude-haiku-4-5-20251001", "display_name": "Claude Haiku 4.5", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, + // Claude 4.x + {"id": "claude-opus-4-6", "display_name": "Claude Opus 4.6", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, + {"id": "claude-sonnet-4-6", "display_name": "Claude Sonnet 4.6", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, + {"id": "claude-opus-4-5", "display_name": "Claude Opus 4.5", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, + {"id": "claude-sonnet-4-5", "display_name": "Claude Sonnet 4.5", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, + {"id": "claude-haiku-4-5", "display_name": "Claude Haiku 4.5", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, + {"id": "claude-haiku-4-5-20251001", "display_name": "Claude Haiku 4.5 (Oct 2025)","created_at": "2025-10-01T00:00:00Z", "type": "model"}, + // Claude 3.7 + {"id": "claude-3-7-sonnet-20250219", "display_name": "Claude 3.7 Sonnet", "created_at": "2025-02-19T00:00:00Z", "type": "model"}, + // Claude 3.5 + {"id": "claude-3-5-sonnet-20241022", "display_name": "Claude 3.5 Sonnet (Oct 2024)","created_at": "2024-10-22T00:00:00Z", "type": "model"}, + {"id": "claude-3-5-sonnet-20240620", "display_name": "Claude 3.5 Sonnet (Jun 2024)","created_at": "2024-06-20T00:00:00Z", "type": "model"}, + {"id": "claude-3-5-haiku-20241022", "display_name": "Claude 3.5 Haiku", "created_at": "2024-10-22T00:00:00Z", "type": "model"}, + // Claude 3 + {"id": "claude-3-opus-20240229", "display_name": "Claude 3 Opus", "created_at": "2024-02-29T00:00:00Z", "type": "model"}, + {"id": "claude-3-sonnet-20240229", "display_name": "Claude 3 Sonnet", "created_at": "2024-02-29T00:00:00Z", "type": "model"}, + {"id": "claude-3-haiku-20240307", "display_name": "Claude 3 Haiku", "created_at": "2024-03-07T00:00:00Z", "type": "model"}, ], "has_more": false, "first_id": "claude-opus-4-6", - "last_id": "claude-haiku-4-5-20251001", + "last_id": "claude-3-haiku-20240307", }) }); @@ -328,6 +347,49 @@ fn backend_error_to_response(error: BackendError) -> Response { (StatusCode::INTERNAL_SERVER_ERROR, Json(err)).into_response() } +/// Inject degradation warnings as `x-anyllm-degradation` header if any features were dropped. +pub(crate) fn inject_degradation_header( + headers: &mut axum::http::HeaderMap, + warnings: &anyllm_translate::TranslationWarnings, +) { + if let Some(val) = warnings.as_header_value() { + if let Ok(hv) = axum::http::HeaderValue::from_str(&val) { + headers.insert("x-anyllm-degradation", hv); + } + } +} + +/// Embeddings passthrough: forwards OpenAI-format embedding requests directly to the backend. +/// No translation needed — embedding model names pass through unchanged. +/// Returns 501 for the Anthropic passthrough backend (no embeddings endpoint). +async fn embeddings( + State(state): State, + headers: axum::http::HeaderMap, + body: axum::body::Bytes, +) -> Response { + let content_type = headers + .get(axum::http::header::CONTENT_TYPE) + .and_then(|v| v.to_str().ok()) + .unwrap_or("application/json") + .to_string(); + + match state + .backend + .embeddings_passthrough(body, &content_type) + .await + { + Ok((status, resp_headers, resp_body)) => { + let mut response = (status, resp_body).into_response(); + // Forward content-type from backend response + for (k, v) in &resp_headers { + response.headers_mut().insert(k, v.clone()); + } + response + } + Err(e) => backend_error_to_response(e), + } +} + async fn messages( State(state): State, headers: axum::http::HeaderMap, @@ -358,6 +420,8 @@ async fn messages( ); } + let warnings = compute_request_warnings(&body); + if body.stream == Some(true) { if state.log_bodies() { tracing::debug!(model = %body.model, "streaming request initiated"); @@ -368,6 +432,7 @@ async fn messages( Ok((rate_limits, sse)) => { let mut response = sse.into_response(); rate_limits.inject_anthropic_response_headers(response.headers_mut()); + inject_degradation_header(response.headers_mut(), &warnings); return response; } Err(e) => { @@ -376,9 +441,9 @@ async fn messages( } } } - match &state.backend { BackendClient::OpenAI(client) + | BackendClient::AzureOpenAI(client) | BackendClient::Vertex(client) | BackendClient::GeminiOpenAI(client) => { let mut openai_req = mapping::message_map::anthropic_to_openai_request(&body); @@ -419,6 +484,7 @@ async fn messages( ); let mut response = (StatusCode::OK, Json(anthropic_resp)).into_response(); rate_limits.inject_anthropic_response_headers(response.headers_mut()); + inject_degradation_header(response.headers_mut(), &warnings); response } Err(e) => { @@ -476,6 +542,7 @@ async fn messages( ); let mut response = (StatusCode::OK, Json(anthropic_resp)).into_response(); rate_limits.inject_anthropic_response_headers(response.headers_mut()); + inject_degradation_header(response.headers_mut(), &warnings); response } Err(e) => { diff --git a/crates/proxy/src/server/streaming.rs b/crates/proxy/src/server/streaming.rs index 80eaac1..afef5a2 100644 --- a/crates/proxy/src/server/streaming.rs +++ b/crates/proxy/src/server/streaming.rs @@ -1,6 +1,6 @@ // SSE streaming infrastructure and the messages_stream handler. -use crate::backend::{BackendClient, RateLimitHeaders, find_double_newline, MAX_SSE_BUFFER_SIZE}; +use crate::backend::{find_double_newline, BackendClient, RateLimitHeaders, MAX_SSE_BUFFER_SIZE}; use crate::metrics::Metrics; use anyllm_translate::{anthropic, mapping, openai}; use axum::response::sse::{Event, KeepAlive, Sse}; @@ -165,6 +165,7 @@ pub(crate) async fn messages_stream( match &state.backend { BackendClient::OpenAI(client) + | BackendClient::AzureOpenAI(client) | BackendClient::Vertex(client) | BackendClient::GeminiOpenAI(client) => { let client = client.clone(); diff --git a/crates/proxy/tests/live_azure.rs b/crates/proxy/tests/live_azure.rs new file mode 100644 index 0000000..fb71a9b --- /dev/null +++ b/crates/proxy/tests/live_azure.rs @@ -0,0 +1,117 @@ +//! Live integration tests against Azure OpenAI endpoints. +//! +//! All tests are `#[ignore]` so they never run in CI or default `cargo test`. +//! +//! Run manually: +//! ```sh +//! AZURE_OPENAI_API_KEY=... \ +//! AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com \ +//! AZURE_OPENAI_DEPLOYMENT=your-deployment \ +//! cargo test --test live_azure -- --ignored --test-threads=1 +//! ``` + +use anyllm_proxy::config::{self, Config}; +use anyllm_proxy::server::routes; +use serde_json::{json, Value}; +use tokio::net::TcpListener; + +fn azure_test_config() -> Config { + let api_key = std::env::var("AZURE_OPENAI_API_KEY") + .expect("AZURE_OPENAI_API_KEY must be set for live Azure tests"); + let endpoint = std::env::var("AZURE_OPENAI_ENDPOINT") + .expect("AZURE_OPENAI_ENDPOINT must be set for live Azure tests"); + let deployment = std::env::var("AZURE_OPENAI_DEPLOYMENT") + .expect("AZURE_OPENAI_DEPLOYMENT must be set for live Azure tests"); + let api_version = + std::env::var("AZURE_OPENAI_API_VERSION").unwrap_or_else(|_| "2024-10-21".to_string()); + + let base_url = format!( + "{}/openai/deployments/{}/chat/completions?api-version={}", + endpoint.trim_end_matches('/'), + deployment, + api_version + ); + + Config { + backend: config::BackendKind::AzureOpenAI, + openai_api_key: String::new(), + openai_base_url: base_url, + listen_port: 0, + model_mapping: config::ModelMapping { + big_model: deployment.clone(), + small_model: deployment, + }, + tls: config::TlsConfig::default(), + backend_auth: config::BackendAuth::AzureApiKey(api_key), + log_bodies: true, + openai_api_format: config::OpenAIApiFormat::Chat, + } +} + +async fn spawn_test_server(config: Config) -> String { + let app = routes::app(config); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + format!("http://127.0.0.1:{}", addr.port()) +} + +/// Verify a basic non-streaming request through Azure OpenAI. +#[tokio::test] +#[ignore] +async fn azure_non_streaming_hello() { + let base = spawn_test_server(azure_test_config()).await; + let client = reqwest::Client::new(); + + let resp = client + .post(format!("{base}/v1/messages")) + .header("x-api-key", "test-key") + .header("content-type", "application/json") + .header("anthropic-version", "2023-06-01") + .json(&json!({ + "model": "claude-sonnet-4-20250514", + "max_tokens": 64, + "messages": [{"role": "user", "content": "Say hello in exactly one word."}] + })) + .send() + .await + .unwrap(); + + assert_eq!(resp.status(), 200, "body: {}", resp.text().await.unwrap()); +} + +/// Verify streaming through Azure OpenAI produces SSE events. +#[tokio::test] +#[ignore] +async fn azure_streaming_hello() { + let base = spawn_test_server(azure_test_config()).await; + let client = reqwest::Client::new(); + + let resp = client + .post(format!("{base}/v1/messages")) + .header("x-api-key", "test-key") + .header("content-type", "application/json") + .header("anthropic-version", "2023-06-01") + .json(&json!({ + "model": "claude-sonnet-4-20250514", + "max_tokens": 64, + "stream": true, + "messages": [{"role": "user", "content": "Say hello in exactly one word."}] + })) + .send() + .await + .unwrap(); + + assert_eq!(resp.status(), 200); + let body = resp.text().await.unwrap(); + assert!( + body.contains("event: message_start"), + "expected SSE events in: {body}" + ); + assert!( + body.contains("event: message_stop"), + "expected message_stop in: {body}" + ); +} diff --git a/crates/translator/src/error.rs b/crates/translator/src/error.rs index afb2394..4def573 100644 --- a/crates/translator/src/error.rs +++ b/crates/translator/src/error.rs @@ -10,4 +10,8 @@ pub enum TranslateError { /// A translation step failed (validation, unsupported feature with strict config, etc.). #[error("translation error: {0}")] Translation(String), + + /// A required field was missing from the input. + #[error("missing field: {0}")] + MissingField(String), } diff --git a/crates/translator/src/lib.rs b/crates/translator/src/lib.rs index f62212a..e8f98e3 100644 --- a/crates/translator/src/lib.rs +++ b/crates/translator/src/lib.rs @@ -57,6 +57,9 @@ pub mod util; pub use config::{LossyBehavior, TranslationConfig, TranslationConfigBuilder}; pub use error::TranslateError; pub use translate::{ - new_responses_stream_translator, new_stream_translator, translate_request, - translate_request_responses, translate_response, translate_response_responses, + compute_request_warnings, new_responses_stream_translator, new_reverse_stream_translator, + new_stream_translator, translate_anthropic_to_openai_response, translate_openai_to_anthropic_request, + translate_request, translate_request_responses, translate_response, + translate_response_responses, TranslationWarnings, }; +pub use mapping::reverse_streaming_map::ReverseStreamingTranslator; diff --git a/crates/translator/src/mapping/mod.rs b/crates/translator/src/mapping/mod.rs index 3567513..1ea1ce8 100644 --- a/crates/translator/src/mapping/mod.rs +++ b/crates/translator/src/mapping/mod.rs @@ -12,6 +12,12 @@ pub mod streaming_map; pub mod tools_map; /// Token usage field mapping between Anthropic and OpenAI formats. pub mod usage_map; +/// Degradation warning collection for client-visible feature-drop signals. +pub mod warnings; +/// Reverse message mapping: OpenAI Chat Completions -> Anthropic Messages. +pub mod reverse_message_map; +/// Reverse streaming: Anthropic SSE events -> OpenAI ChatCompletionChunk SSE. +pub mod reverse_streaming_map; /// Format an OpenAI refusal string as Anthropic text content. /// Anthropic has no refusal type, so we surface it as a bracketed text marker. diff --git a/crates/translator/src/mapping/reverse_message_map.rs b/crates/translator/src/mapping/reverse_message_map.rs new file mode 100644 index 0000000..eee9ffb --- /dev/null +++ b/crates/translator/src/mapping/reverse_message_map.rs @@ -0,0 +1,729 @@ +// Reverse message mapping: OpenAI Chat Completions -> Anthropic Messages +// +// Converts OpenAI-format requests to Anthropic format (for accepting OpenAI +// input) and Anthropic responses back to OpenAI format. + +use crate::anthropic; +use crate::error::TranslateError; +use crate::mapping::{tools_map, usage_map, warnings::TranslationWarnings}; +use crate::openai; +use crate::util; + +/// Convert an OpenAI ChatCompletionRequest to an Anthropic MessageCreateRequest. +/// +/// Returns an error if `max_tokens` and `max_completion_tokens` are both absent +/// (Anthropic requires `max_tokens`). +pub fn openai_to_anthropic_request( + req: &openai::ChatCompletionRequest, + warnings: &mut TranslationWarnings, +) -> Result { + // max_tokens is required in Anthropic; reject if absent + let max_tokens = req + .max_completion_tokens + .or(req.max_tokens) + .ok_or_else(|| { + TranslateError::MissingField("max_tokens or max_completion_tokens is required".into()) + })?; + + let mut system: Option = None; + let mut messages = Vec::new(); + + for msg in &req.messages { + match msg.role { + openai::ChatRole::System | openai::ChatRole::Developer => { + // Extract system messages into the Anthropic system field. + // Multiple system messages are concatenated. + let text = extract_text_content(&msg.content); + if !text.is_empty() { + match &mut system { + Some(anthropic::System::Text(existing)) => { + existing.push('\n'); + existing.push_str(&text); + } + None => { + system = Some(anthropic::System::Text(text)); + } + _ => {} + } + } + } + openai::ChatRole::User => { + let content = convert_openai_content_to_anthropic(&msg.content); + messages.push(anthropic::InputMessage { + role: anthropic::Role::User, + content, + }); + } + openai::ChatRole::Assistant => { + let content = convert_assistant_to_anthropic(msg); + messages.push(anthropic::InputMessage { + role: anthropic::Role::Assistant, + content, + }); + } + openai::ChatRole::Tool => { + // Tool role messages become Anthropic tool_result blocks + // on a user message (Anthropic requires tool results in user turn) + let text = extract_text_content(&msg.content); + let tool_use_id = msg.tool_call_id.clone().unwrap_or_default(); + let content_block = anthropic::ContentBlock::ToolResult { + tool_use_id, + content: if text.is_empty() { + None + } else { + Some(anthropic::ToolResultContent::Text(text)) + }, + is_error: None, + }; + messages.push(anthropic::InputMessage { + role: anthropic::Role::User, + content: anthropic::Content::Blocks(vec![content_block]), + }); + } + openai::ChatRole::Function => { + // Deprecated function role: treat as tool + let text = extract_text_content(&msg.content); + let tool_use_id = msg.name.clone().unwrap_or_default(); + let content_block = anthropic::ContentBlock::ToolResult { + tool_use_id, + content: if text.is_empty() { + None + } else { + Some(anthropic::ToolResultContent::Text(text)) + }, + is_error: None, + }; + messages.push(anthropic::InputMessage { + role: anthropic::Role::User, + content: anthropic::Content::Blocks(vec![content_block]), + }); + } + } + } + + // Map tools + let tools = req + .tools + .as_ref() + .map(|t| tools_map::openai_tools_to_anthropic(t)); + + // Map tool_choice + let tool_choice = req + .tool_choice + .as_ref() + .map(tools_map::openai_tool_choice_to_anthropic); + + // Map stop sequences + let stop_sequences = req.stop.as_ref().map(|s| match s { + openai::Stop::Single(s) => vec![s.clone()], + openai::Stop::Multiple(v) => v.clone(), + }); + + // Map user to metadata + let metadata = req.user.as_ref().map(|u| anthropic::Metadata { + user_id: Some(u.clone()), + }); + + // Record lossy fields as warnings + if req.presence_penalty.is_some() { + warnings.add("presence_penalty"); + } + if req.frequency_penalty.is_some() { + warnings.add("frequency_penalty"); + } + if req.response_format.is_some() { + warnings.add("response_format"); + } + if req.extra.contains_key("logprobs") { + warnings.add("logprobs"); + } + if req.extra.contains_key("n") { + warnings.add("n"); + } + if req.extra.contains_key("seed") { + warnings.add("seed"); + } + if req.stream_options.is_some() { + warnings.add("stream_options"); + } + + // Map parallel_tool_calls: false -> disable_parallel_tool_use: true + let tool_choice = match (tool_choice, req.parallel_tool_calls) { + (Some(anthropic::ToolChoice::Auto { .. }), Some(false)) => { + Some(anthropic::ToolChoice::Auto { + disable_parallel_tool_use: Some(true), + }) + } + (Some(anthropic::ToolChoice::Any { .. }), Some(false)) => { + Some(anthropic::ToolChoice::Any { + disable_parallel_tool_use: Some(true), + }) + } + (tc, _) => tc, + }; + + Ok(anthropic::MessageCreateRequest { + model: req.model.clone(), + max_tokens, + messages, + system, + temperature: req.temperature, + top_p: req.top_p, + top_k: None, + stop_sequences, + tools, + tool_choice, + metadata, + thinking: None, + stream: req.stream, + extra: serde_json::Map::new(), + }) +} + +/// Convert an Anthropic MessageResponse to an OpenAI ChatCompletionResponse. +pub fn anthropic_to_openai_response( + resp: &anthropic::MessageResponse, + model: &str, +) -> openai::ChatCompletionResponse { + let mut text_parts = Vec::new(); + let mut tool_calls = Vec::new(); + let mut reasoning_content: Option = None; + + for block in &resp.content { + match block { + anthropic::ContentBlock::Text { text } => { + text_parts.push(text.clone()); + } + anthropic::ContentBlock::ToolUse { id, name, input } => { + tool_calls.push(openai::ToolCall { + id: id.clone(), + call_type: "function".to_string(), + function: openai::FunctionCall { + name: name.clone(), + arguments: util::json::value_to_json_string(input), + }, + }); + } + anthropic::ContentBlock::Thinking { thinking, .. } => { + match &mut reasoning_content { + Some(existing) => { + existing.push_str(thinking); + } + None => { + reasoning_content = Some(thinking.clone()); + } + } + } + _ => {} + } + } + + let content = if text_parts.is_empty() { + None + } else { + Some(openai::ChatContent::Text(text_parts.join(""))) + }; + + let finish_reason = resp + .stop_reason + .as_ref() + .map(anthropic_stop_reason_to_openai); + + let usage = usage_map::anthropic_to_openai_usage(&resp.usage); + + let id = format!("chatcmpl-{}", util::ids::generate_uuid()); + + openai::ChatCompletionResponse { + id, + object: "chat.completion".to_string(), + model: model.to_string(), + choices: vec![openai::Choice { + index: 0, + message: openai::ChatMessage { + role: openai::ChatRole::Assistant, + content, + name: None, + tool_calls: if tool_calls.is_empty() { + None + } else { + Some(tool_calls) + }, + tool_call_id: None, + refusal: None, + reasoning_content, + }, + finish_reason, + logprobs: None, + }], + usage: Some(usage), + created: resp.created, + system_fingerprint: None, + service_tier: None, + } +} + +/// Map Anthropic stop_reason to OpenAI finish_reason. +pub fn anthropic_stop_reason_to_openai( + stop_reason: &anthropic::StopReason, +) -> openai::FinishReason { + match stop_reason { + anthropic::StopReason::EndTurn => openai::FinishReason::Stop, + anthropic::StopReason::MaxTokens => openai::FinishReason::Length, + anthropic::StopReason::ToolUse => openai::FinishReason::ToolCalls, + anthropic::StopReason::StopSequence => openai::FinishReason::Stop, + } +} + +/// Compute warnings for an OpenAI request about features that will be dropped. +pub fn compute_openai_request_warnings(req: &openai::ChatCompletionRequest) -> TranslationWarnings { + let mut w = TranslationWarnings::default(); + openai_to_anthropic_request(req, &mut w).ok(); + w +} + +// --- Helper functions --- + +fn extract_text_content(content: &Option) -> String { + match content { + Some(openai::ChatContent::Text(s)) => s.clone(), + Some(openai::ChatContent::Parts(parts)) => parts + .iter() + .filter_map(|p| match p { + openai::ChatContentPart::Text { text } => Some(text.as_str()), + _ => None, + }) + .collect::>() + .join(""), + None => String::new(), + } +} + +fn convert_openai_content_to_anthropic( + content: &Option, +) -> anthropic::Content { + match content { + Some(openai::ChatContent::Text(s)) => anthropic::Content::Text(s.clone()), + Some(openai::ChatContent::Parts(parts)) => { + let mut blocks = Vec::new(); + for part in parts { + match part { + openai::ChatContentPart::Text { text } => { + blocks.push(anthropic::ContentBlock::Text { text: text.clone() }); + } + openai::ChatContentPart::ImageUrl { image_url } => { + // Parse data URIs back to base64 + media_type + let source = url_to_image_source(&image_url.url); + blocks.push(anthropic::ContentBlock::Image { source }); + } + // InputAudio and File have no Anthropic equivalent; drop them + _ => {} + } + } + if blocks.is_empty() { + anthropic::Content::Text(String::new()) + } else { + anthropic::Content::Blocks(blocks) + } + } + None => anthropic::Content::Text(String::new()), + } +} + +fn convert_assistant_to_anthropic(msg: &openai::ChatMessage) -> anthropic::Content { + let mut blocks = Vec::new(); + + // Map reasoning_content to thinking block + if let Some(ref reasoning) = msg.reasoning_content { + if !reasoning.is_empty() { + blocks.push(anthropic::ContentBlock::Thinking { + thinking: reasoning.clone(), + signature: None, + }); + } + } + + // Map text content + match &msg.content { + Some(openai::ChatContent::Text(text)) => { + if !text.is_empty() { + blocks.push(anthropic::ContentBlock::Text { text: text.clone() }); + } + } + Some(openai::ChatContent::Parts(parts)) => { + for part in parts { + if let openai::ChatContentPart::Text { text } = part { + blocks.push(anthropic::ContentBlock::Text { text: text.clone() }); + } + } + } + None => {} + } + + // Map tool calls to tool_use blocks + if let Some(ref tool_calls) = msg.tool_calls { + for tc in tool_calls { + blocks.push(anthropic::ContentBlock::ToolUse { + id: tc.id.clone(), + name: tc.function.name.clone(), + input: util::json::parse_tool_arguments(&tc.function.arguments), + }); + } + } + + if blocks.is_empty() { + anthropic::Content::Text(String::new()) + } else if blocks.len() == 1 { + if let anthropic::ContentBlock::Text { ref text } = blocks[0] { + return anthropic::Content::Text(text.clone()); + } + anthropic::Content::Blocks(blocks) + } else { + anthropic::Content::Blocks(blocks) + } +} + +/// Parse a URL string into an Anthropic ImageSource. +/// Handles both data URIs (data:image/png;base64,...) and regular URLs. +fn url_to_image_source(url: &str) -> anthropic::ImageSource { + if let Some(rest) = url.strip_prefix("data:") { + // Parse data URI: data:media_type;base64,data + if let Some((meta, data)) = rest.split_once(',') { + let media_type = meta.strip_suffix(";base64").unwrap_or(meta); + return anthropic::ImageSource { + source_type: "base64".to_string(), + media_type: Some(media_type.to_string()), + data: Some(data.to_string()), + url: None, + }; + } + } + // Regular URL + anthropic::ImageSource { + source_type: "url".to_string(), + media_type: None, + data: None, + url: Some(url.to_string()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn make_basic_request() -> openai::ChatCompletionRequest { + serde_json::from_value(json!({ + "model": "claude-sonnet-4-20250514", + "messages": [ + {"role": "user", "content": "Hello"} + ], + "max_tokens": 100 + })) + .unwrap() + } + + #[test] + fn basic_message_conversion() { + let req = make_basic_request(); + let mut w = TranslationWarnings::default(); + let result = openai_to_anthropic_request(&req, &mut w).unwrap(); + assert_eq!(result.model, "claude-sonnet-4-20250514"); + assert_eq!(result.max_tokens, 100); + assert_eq!(result.messages.len(), 1); + assert_eq!(result.messages[0].role, anthropic::Role::User); + } + + #[test] + fn system_message_extraction() { + let req: openai::ChatCompletionRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-20250514", + "messages": [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hi"} + ], + "max_tokens": 100 + })) + .unwrap(); + let mut w = TranslationWarnings::default(); + let result = openai_to_anthropic_request(&req, &mut w).unwrap(); + assert!(matches!(result.system, Some(anthropic::System::Text(ref s)) if s == "You are helpful.")); + assert_eq!(result.messages.len(), 1); // system not in messages + } + + #[test] + fn developer_role_maps_to_system() { + let req: openai::ChatCompletionRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-20250514", + "messages": [ + {"role": "developer", "content": "Be concise."}, + {"role": "user", "content": "Hi"} + ], + "max_tokens": 100 + })) + .unwrap(); + let mut w = TranslationWarnings::default(); + let result = openai_to_anthropic_request(&req, &mut w).unwrap(); + assert!(matches!(result.system, Some(anthropic::System::Text(ref s)) if s == "Be concise.")); + } + + #[test] + fn missing_max_tokens_rejected() { + let req: openai::ChatCompletionRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-20250514", + "messages": [{"role": "user", "content": "Hi"}] + })) + .unwrap(); + let mut w = TranslationWarnings::default(); + let result = openai_to_anthropic_request(&req, &mut w); + assert!(result.is_err()); + } + + #[test] + fn max_completion_tokens_used_as_fallback() { + let req: openai::ChatCompletionRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-20250514", + "messages": [{"role": "user", "content": "Hi"}], + "max_completion_tokens": 200 + })) + .unwrap(); + let mut w = TranslationWarnings::default(); + let result = openai_to_anthropic_request(&req, &mut w).unwrap(); + assert_eq!(result.max_tokens, 200); + } + + #[test] + fn tool_call_conversion() { + let req: openai::ChatCompletionRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-20250514", + "messages": [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "tool_calls": [{ + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{\"loc\":\"NYC\"}"} + }] + }, + {"role": "tool", "tool_call_id": "call_1", "content": "Sunny, 72F"} + ], + "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + "max_tokens": 100 + })) + .unwrap(); + let mut w = TranslationWarnings::default(); + let result = openai_to_anthropic_request(&req, &mut w).unwrap(); + assert_eq!(result.messages.len(), 3); + assert!(result.tools.is_some()); + // Second message (assistant) should have tool_use block + match &result.messages[1].content { + anthropic::Content::Blocks(blocks) => { + assert!(matches!(&blocks[0], anthropic::ContentBlock::ToolUse { name, .. } if name == "get_weather")); + } + _ => panic!("expected blocks"), + } + // Third message (tool result) should be user with tool_result + assert_eq!(result.messages[2].role, anthropic::Role::User); + } + + #[test] + fn lossy_fields_generate_warnings() { + let req: openai::ChatCompletionRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-20250514", + "messages": [{"role": "user", "content": "Hi"}], + "max_tokens": 100, + "presence_penalty": 0.5, + "frequency_penalty": 0.3, + "logprobs": true, + "seed": 42 + })) + .unwrap(); + let mut w = TranslationWarnings::default(); + openai_to_anthropic_request(&req, &mut w).unwrap(); + let header = w.as_header_value().unwrap(); + assert!(header.contains("presence_penalty")); + assert!(header.contains("frequency_penalty")); + assert!(header.contains("logprobs")); + assert!(header.contains("seed")); + } + + #[test] + fn stop_sequences_mapping() { + let req: openai::ChatCompletionRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-20250514", + "messages": [{"role": "user", "content": "Hi"}], + "max_tokens": 100, + "stop": ["END", "STOP"] + })) + .unwrap(); + let mut w = TranslationWarnings::default(); + let result = openai_to_anthropic_request(&req, &mut w).unwrap(); + assert_eq!(result.stop_sequences, Some(vec!["END".into(), "STOP".into()])); + } + + // --- Response tests --- + + #[test] + fn basic_response_conversion() { + let resp = anthropic::MessageResponse { + id: "msg_123".to_string(), + response_type: "message".to_string(), + role: anthropic::Role::Assistant, + content: vec![anthropic::ContentBlock::Text { + text: "Hello!".to_string(), + }], + model: "claude-sonnet-4-20250514".to_string(), + stop_reason: Some(anthropic::StopReason::EndTurn), + stop_sequence: None, + usage: anthropic::Usage { + input_tokens: 10, + output_tokens: 5, + cache_creation_input_tokens: None, + cache_read_input_tokens: None, + }, + created: Some(1700000000), + }; + let result = anthropic_to_openai_response(&resp, "claude-sonnet-4-20250514"); + assert_eq!(result.object, "chat.completion"); + assert!(result.id.starts_with("chatcmpl-")); + assert_eq!(result.choices.len(), 1); + match &result.choices[0].message.content { + Some(openai::ChatContent::Text(s)) => assert_eq!(s, "Hello!"), + other => panic!("expected Text, got {:?}", other), + } + assert_eq!(result.choices[0].finish_reason, Some(openai::FinishReason::Stop)); + let usage = result.usage.unwrap(); + assert_eq!(usage.prompt_tokens, 10); + assert_eq!(usage.completion_tokens, 5); + } + + #[test] + fn tool_use_response_conversion() { + let resp = anthropic::MessageResponse { + id: "msg_456".to_string(), + response_type: "message".to_string(), + role: anthropic::Role::Assistant, + content: vec![anthropic::ContentBlock::ToolUse { + id: "call_1".to_string(), + name: "get_weather".to_string(), + input: json!({"location": "NYC"}), + }], + model: "claude-sonnet-4-20250514".to_string(), + stop_reason: Some(anthropic::StopReason::ToolUse), + stop_sequence: None, + usage: anthropic::Usage::default(), + created: None, + }; + let result = anthropic_to_openai_response(&resp, "claude-sonnet-4-20250514"); + let tc = result.choices[0].message.tool_calls.as_ref().unwrap(); + assert_eq!(tc.len(), 1); + assert_eq!(tc[0].id, "call_1"); + assert_eq!(tc[0].function.name, "get_weather"); + assert_eq!(result.choices[0].finish_reason, Some(openai::FinishReason::ToolCalls)); + } + + #[test] + fn thinking_block_maps_to_reasoning_content() { + let resp = anthropic::MessageResponse { + id: "msg_789".to_string(), + response_type: "message".to_string(), + role: anthropic::Role::Assistant, + content: vec![ + anthropic::ContentBlock::Thinking { + thinking: "Let me think...".to_string(), + signature: None, + }, + anthropic::ContentBlock::Text { + text: "The answer is 4.".to_string(), + }, + ], + model: "claude-sonnet-4-20250514".to_string(), + stop_reason: Some(anthropic::StopReason::EndTurn), + stop_sequence: None, + usage: anthropic::Usage::default(), + created: None, + }; + let result = anthropic_to_openai_response(&resp, "claude-sonnet-4-20250514"); + assert_eq!( + result.choices[0].message.reasoning_content.as_deref(), + Some("Let me think...") + ); + match &result.choices[0].message.content { + Some(openai::ChatContent::Text(s)) => assert_eq!(s, "The answer is 4."), + other => panic!("expected Text, got {:?}", other), + } + } + + #[test] + fn stop_reason_mapping() { + assert_eq!( + anthropic_stop_reason_to_openai(&anthropic::StopReason::EndTurn), + openai::FinishReason::Stop + ); + assert_eq!( + anthropic_stop_reason_to_openai(&anthropic::StopReason::MaxTokens), + openai::FinishReason::Length + ); + assert_eq!( + anthropic_stop_reason_to_openai(&anthropic::StopReason::ToolUse), + openai::FinishReason::ToolCalls + ); + assert_eq!( + anthropic_stop_reason_to_openai(&anthropic::StopReason::StopSequence), + openai::FinishReason::Stop + ); + } + + #[test] + fn data_uri_image_parsing() { + let source = url_to_image_source("data:image/png;base64,iVBORw0KGgo="); + assert_eq!(source.source_type, "base64"); + assert_eq!(source.media_type.as_deref(), Some("image/png")); + assert_eq!(source.data.as_deref(), Some("iVBORw0KGgo=")); + assert!(source.url.is_none()); + } + + #[test] + fn regular_url_image_source() { + let source = url_to_image_source("https://example.com/img.png"); + assert_eq!(source.source_type, "url"); + assert_eq!(source.url.as_deref(), Some("https://example.com/img.png")); + assert!(source.data.is_none()); + } + + #[test] + fn user_field_maps_to_metadata() { + let req: openai::ChatCompletionRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-20250514", + "messages": [{"role": "user", "content": "Hi"}], + "max_tokens": 100, + "user": "user-123" + })) + .unwrap(); + let mut w = TranslationWarnings::default(); + let result = openai_to_anthropic_request(&req, &mut w).unwrap(); + assert_eq!( + result.metadata.as_ref().and_then(|m| m.user_id.as_deref()), + Some("user-123") + ); + } + + #[test] + fn parallel_tool_calls_false_maps_to_disable() { + let req: openai::ChatCompletionRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-20250514", + "messages": [{"role": "user", "content": "Hi"}], + "max_tokens": 100, + "tools": [{"type": "function", "function": {"name": "test", "parameters": {"type": "object"}}}], + "tool_choice": "auto", + "parallel_tool_calls": false + })) + .unwrap(); + let mut w = TranslationWarnings::default(); + let result = openai_to_anthropic_request(&req, &mut w).unwrap(); + assert!(matches!( + result.tool_choice, + Some(anthropic::ToolChoice::Auto { disable_parallel_tool_use: Some(true) }) + )); + } +} diff --git a/crates/translator/src/mapping/reverse_streaming_map.rs b/crates/translator/src/mapping/reverse_streaming_map.rs new file mode 100644 index 0000000..dc13f67 --- /dev/null +++ b/crates/translator/src/mapping/reverse_streaming_map.rs @@ -0,0 +1,370 @@ +// Reverse streaming: Anthropic SSE events -> OpenAI ChatCompletionChunk SSE +// +// Consumes Anthropic StreamEvent items and emits OpenAI ChatCompletionChunk +// objects. This is the inverse of StreamingTranslator in streaming_map.rs. + +use crate::anthropic; +use crate::openai; +use crate::openai::streaming::{ChatCompletionChunk, ChunkChoice, ChunkDelta, ChunkFunctionCall, ChunkToolCall}; + +/// Sentinel value returned by `process_event` to signal the stream is done. +/// The caller should emit `data: [DONE]\n\n` when it sees this. +pub const DONE_SENTINEL: &str = "[DONE]"; + +/// State machine that converts Anthropic SSE events into OpenAI ChatCompletionChunk objects. +/// +/// Feed events via `process_event`, which returns zero or more chunks to send. +/// When `message_stop` is received, `is_done()` returns true and the caller +/// should emit `data: [DONE]\n\n`. +pub struct ReverseStreamingTranslator { + message_id: String, + model: String, + tool_call_index: i32, + input_tokens: Option, + output_tokens: Option, + created: u64, + done: bool, +} + +impl ReverseStreamingTranslator { + pub fn new(id: String, model: String) -> Self { + Self { + message_id: id, + model, + tool_call_index: -1, + input_tokens: None, + output_tokens: None, + created: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + done: false, + } + } + + pub fn is_done(&self) -> bool { + self.done + } + + /// Process a single Anthropic StreamEvent and return zero or more OpenAI chunks. + pub fn process_event(&mut self, event: &anthropic::StreamEvent) -> Vec { + match event { + anthropic::StreamEvent::MessageStart { message } => { + self.input_tokens = Some(message.usage.input_tokens); + if let Some(created) = message.created { + self.created = created; + } + // Emit first chunk with role + vec![self.make_chunk( + ChunkDelta { + role: Some(openai::ChatRole::Assistant), + ..Default::default() + }, + None, + )] + } + anthropic::StreamEvent::ContentBlockStart { content_block, .. } => { + match content_block { + anthropic::ContentBlock::ToolUse { id, name, .. } => { + self.tool_call_index += 1; + let tc = ChunkToolCall { + index: self.tool_call_index as u32, + id: Some(id.clone()), + call_type: Some("function".to_string()), + function: Some(ChunkFunctionCall { + name: Some(name.clone()), + arguments: Some(String::new()), + }), + }; + vec![self.make_chunk( + ChunkDelta { + tool_calls: Some(vec![tc]), + ..Default::default() + }, + None, + )] + } + // Text and Thinking blocks emit their content via deltas + _ => vec![], + } + } + anthropic::StreamEvent::ContentBlockDelta { delta, .. } => { + match delta { + anthropic::streaming::Delta::TextDelta { text } => { + vec![self.make_chunk( + ChunkDelta { + content: Some(text.clone()), + ..Default::default() + }, + None, + )] + } + anthropic::streaming::Delta::InputJsonDelta { partial_json } => { + if self.tool_call_index < 0 { + return vec![]; + } + let tc = ChunkToolCall { + index: self.tool_call_index as u32, + id: None, + call_type: None, + function: Some(ChunkFunctionCall { + name: None, + arguments: Some(partial_json.clone()), + }), + }; + vec![self.make_chunk( + ChunkDelta { + tool_calls: Some(vec![tc]), + ..Default::default() + }, + None, + )] + } + anthropic::streaming::Delta::ThinkingDelta { thinking } => { + vec![self.make_chunk( + ChunkDelta { + reasoning_content: Some(thinking.clone()), + ..Default::default() + }, + None, + )] + } + anthropic::streaming::Delta::SignatureDelta { .. } => vec![], + } + } + anthropic::StreamEvent::ContentBlockStop { .. } => vec![], + anthropic::StreamEvent::MessageDelta { delta, usage } => { + if let Some(u) = usage { + self.output_tokens = Some(u.output_tokens); + } + let finish_reason = delta.stop_reason.as_ref().map(|sr| { + match sr { + anthropic::StopReason::EndTurn => openai::FinishReason::Stop, + anthropic::StopReason::MaxTokens => openai::FinishReason::Length, + anthropic::StopReason::ToolUse => openai::FinishReason::ToolCalls, + anthropic::StopReason::StopSequence => openai::FinishReason::Stop, + } + }); + let mut chunks = vec![self.make_chunk( + ChunkDelta::default(), + finish_reason, + )]; + // Emit usage chunk if we have token counts + if let (Some(input), Some(output)) = (self.input_tokens, self.output_tokens) { + chunks.push(ChatCompletionChunk { + id: self.message_id.clone(), + object: "chat.completion.chunk".to_string(), + model: self.model.clone(), + choices: vec![], + usage: Some(openai::ChatUsage { + prompt_tokens: input, + completion_tokens: output, + total_tokens: input + output, + completion_tokens_details: None, + prompt_tokens_details: None, + }), + created: Some(self.created), + system_fingerprint: None, + }); + } + chunks + } + anthropic::StreamEvent::MessageStop {} => { + self.done = true; + vec![] + } + anthropic::StreamEvent::Ping {} => vec![], + anthropic::StreamEvent::Error { .. } => { + self.done = true; + vec![] + } + } + } + + fn make_chunk( + &self, + delta: ChunkDelta, + finish_reason: Option, + ) -> ChatCompletionChunk { + ChatCompletionChunk { + id: self.message_id.clone(), + object: "chat.completion.chunk".to_string(), + model: self.model.clone(), + choices: vec![ChunkChoice { + index: 0, + delta, + finish_reason, + logprobs: None, + }], + usage: None, + created: Some(self.created), + system_fingerprint: None, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::anthropic::messages::{ContentBlock, StopReason, Usage}; + use crate::anthropic::streaming::*; + + fn make_translator() -> ReverseStreamingTranslator { + ReverseStreamingTranslator::new("chatcmpl-test".to_string(), "gpt-4o".to_string()) + } + + #[test] + fn message_start_emits_role_chunk() { + let mut t = make_translator(); + let event = StreamEvent::MessageStart { + message: MessageStartData { + id: "msg_123".to_string(), + msg_type: "message".to_string(), + role: "assistant".to_string(), + content: vec![], + model: "claude-sonnet".to_string(), + stop_reason: None, + stop_sequence: None, + usage: Usage { input_tokens: 10, output_tokens: 0, cache_creation_input_tokens: None, cache_read_input_tokens: None }, + created: Some(1700000000), + }, + }; + let chunks = t.process_event(&event); + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].choices[0].delta.role, Some(openai::ChatRole::Assistant)); + assert!(chunks[0].choices[0].finish_reason.is_none()); + } + + #[test] + fn text_delta_emits_content_chunk() { + let mut t = make_translator(); + let event = StreamEvent::ContentBlockDelta { + index: 0, + delta: Delta::TextDelta { text: "Hello".to_string() }, + }; + let chunks = t.process_event(&event); + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].choices[0].delta.content.as_deref(), Some("Hello")); + } + + #[test] + fn tool_use_streaming() { + let mut t = make_translator(); + // Start tool use block + let start = StreamEvent::ContentBlockStart { + index: 0, + content_block: ContentBlock::ToolUse { + id: "call_123".to_string(), + name: "get_weather".to_string(), + input: serde_json::Value::Object(serde_json::Map::new()), + }, + }; + let chunks = t.process_event(&start); + assert_eq!(chunks.len(), 1); + let tc = &chunks[0].choices[0].delta.tool_calls.as_ref().unwrap()[0]; + assert_eq!(tc.id.as_deref(), Some("call_123")); + assert_eq!(tc.function.as_ref().unwrap().name.as_deref(), Some("get_weather")); + + // Delta with args + let delta = StreamEvent::ContentBlockDelta { + index: 0, + delta: Delta::InputJsonDelta { partial_json: "{\"loc".to_string() }, + }; + let chunks = t.process_event(&delta); + assert_eq!(chunks.len(), 1); + let tc = &chunks[0].choices[0].delta.tool_calls.as_ref().unwrap()[0]; + assert_eq!(tc.index, 0); + assert!(tc.id.is_none()); // Only first chunk has id + assert_eq!(tc.function.as_ref().unwrap().arguments.as_deref(), Some("{\"loc")); + } + + #[test] + fn thinking_delta_emits_reasoning_content() { + let mut t = make_translator(); + let event = StreamEvent::ContentBlockDelta { + index: 0, + delta: Delta::ThinkingDelta { thinking: "Let me think...".to_string() }, + }; + let chunks = t.process_event(&event); + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].choices[0].delta.reasoning_content.as_deref(), Some("Let me think...")); + } + + #[test] + fn message_delta_emits_finish_reason_and_usage() { + let mut t = make_translator(); + // Set input tokens via message_start + let start = StreamEvent::MessageStart { + message: MessageStartData { + id: "msg_1".to_string(), + msg_type: "message".to_string(), + role: "assistant".to_string(), + content: vec![], + model: "claude".to_string(), + stop_reason: None, + stop_sequence: None, + usage: Usage { input_tokens: 10, output_tokens: 0, cache_creation_input_tokens: None, cache_read_input_tokens: None }, + created: None, + }, + }; + t.process_event(&start); + + let event = StreamEvent::MessageDelta { + delta: MessageDeltaData { + stop_reason: Some(StopReason::EndTurn), + stop_sequence: None, + }, + usage: Some(DeltaUsage { output_tokens: 5 }), + }; + let chunks = t.process_event(&event); + assert_eq!(chunks.len(), 2); // finish chunk + usage chunk + assert_eq!(chunks[0].choices[0].finish_reason, Some(openai::FinishReason::Stop)); + let usage = chunks[1].usage.as_ref().unwrap(); + assert_eq!(usage.prompt_tokens, 10); + assert_eq!(usage.completion_tokens, 5); + assert_eq!(usage.total_tokens, 15); + } + + #[test] + fn message_stop_sets_done() { + let mut t = make_translator(); + assert!(!t.is_done()); + t.process_event(&StreamEvent::MessageStop {}); + assert!(t.is_done()); + } + + #[test] + fn ping_produces_no_chunks() { + let mut t = make_translator(); + let chunks = t.process_event(&StreamEvent::Ping {}); + assert!(chunks.is_empty()); + } + + #[test] + fn multiple_tool_calls_track_index() { + let mut t = make_translator(); + // First tool + let start1 = StreamEvent::ContentBlockStart { + index: 0, + content_block: ContentBlock::ToolUse { + id: "call_1".to_string(), + name: "fn_a".to_string(), + input: serde_json::Value::Object(serde_json::Map::new()), + }, + }; + let chunks = t.process_event(&start1); + assert_eq!(chunks[0].choices[0].delta.tool_calls.as_ref().unwrap()[0].index, 0); + + // Second tool + let start2 = StreamEvent::ContentBlockStart { + index: 1, + content_block: ContentBlock::ToolUse { + id: "call_2".to_string(), + name: "fn_b".to_string(), + input: serde_json::Value::Object(serde_json::Map::new()), + }, + }; + let chunks = t.process_event(&start2); + assert_eq!(chunks[0].choices[0].delta.tool_calls.as_ref().unwrap()[0].index, 1); + } +} diff --git a/crates/translator/src/translate.rs b/crates/translator/src/translate.rs index 841cfb1..dff6802 100644 --- a/crates/translator/src/translate.rs +++ b/crates/translator/src/translate.rs @@ -6,9 +6,20 @@ use crate::anthropic::{MessageCreateRequest, MessageResponse}; use crate::config::TranslationConfig; use crate::error::TranslateError; -use crate::mapping::{message_map, responses_message_map, responses_streaming_map, streaming_map}; +use crate::mapping::{ + message_map, responses_message_map, responses_streaming_map, reverse_message_map, + reverse_streaming_map, streaming_map, +}; use crate::openai::responses::{ResponsesRequest, ResponsesResponse}; use crate::openai::{ChatCompletionRequest, ChatCompletionResponse}; +pub use crate::mapping::warnings::TranslationWarnings; + +/// Compute degradation warnings for a request — features that will be dropped in translation. +/// +/// Call this before translating; inject the result as `x-anyllm-degradation` header. +pub fn compute_request_warnings(req: &MessageCreateRequest) -> TranslationWarnings { + message_map::compute_request_warnings(req) +} /// Translate an Anthropic request to an OpenAI Chat Completions request. /// @@ -60,6 +71,38 @@ pub fn translate_response_responses( responses_message_map::responses_to_anthropic_response(resp, original_model) } +/// Translate an OpenAI Chat Completions request to an Anthropic request. +/// +/// Returns an error if `max_tokens` / `max_completion_tokens` is absent. +/// Populates `warnings` with features dropped during translation. +pub fn translate_openai_to_anthropic_request( + req: &ChatCompletionRequest, + warnings: &mut TranslationWarnings, +) -> Result { + reverse_message_map::openai_to_anthropic_request(req, warnings) +} + +/// Translate an Anthropic response to an OpenAI Chat Completions response. +/// +/// `model` is used as the response's `model` field. +pub fn translate_anthropic_to_openai_response( + resp: &MessageResponse, + model: &str, +) -> ChatCompletionResponse { + reverse_message_map::anthropic_to_openai_response(resp, model) +} + +/// Create a new reverse streaming translator (Anthropic SSE -> OpenAI chunks). +/// +/// The returned translator is stateful: feed Anthropic StreamEvent items via +/// `process_event()`, which returns OpenAI ChatCompletionChunk objects. +pub fn new_reverse_stream_translator( + id: String, + model: String, +) -> reverse_streaming_map::ReverseStreamingTranslator { + reverse_streaming_map::ReverseStreamingTranslator::new(id, model) +} + /// Create a new streaming translator for OpenAI Responses API events. /// /// Same stateful pattern as `new_stream_translator`. diff --git a/crates/translator/src/util/ids.rs b/crates/translator/src/util/ids.rs index c804650..33e2be6 100644 --- a/crates/translator/src/util/ids.rs +++ b/crates/translator/src/util/ids.rs @@ -18,6 +18,11 @@ pub fn generate_tool_use_id() -> String { format!("toolu_{}", uuid::Uuid::new_v4().as_simple()) } +/// Generate a raw UUID v4 without hyphens (for custom prefix use). +pub fn generate_uuid() -> String { + uuid::Uuid::new_v4().as_simple().to_string() +} + #[cfg(test)] mod tests { use super::*; diff --git a/docs/ENV.md b/docs/ENV.md index f9703fa..c5ee94a 100644 --- a/docs/ENV.md +++ b/docs/ENV.md @@ -1,5 +1,36 @@ # Environment Variables +## Env Files + +Instead of setting variables in the shell, you can store them in a `.env` file and load it at startup. + +**Auto-load:** If `.anyllm.env` exists in the current directory, it is loaded automatically. + +**Explicit flag:** +```bash +anyllm_proxy --env-file ~/configs/deepseek.env +``` + +**File format** (`KEY=VALUE`, Docker `--env-file` compatible): +```env +# Comments are supported +OPENAI_API_KEY=sk-... +OPENAI_BASE_URL=https://api.deepseek.com/v1 +BIG_MODEL=deepseek-coder +SMALL_MODEL=deepseek-chat +export LISTEN_PORT=3000 # export prefix is also accepted +``` + +Rules: +- Lines starting with `#` are ignored. +- Values may be optionally quoted with `"double"` or `'single'` quotes. +- Environment variables already set in the shell take precedence over the file. +- Use `docker run --env-file ` to pass the same file to a container. + +The admin UI (Settings tab) has an **Export .env** button that generates a template from the current running configuration. + +--- + ## Core These are the variables most users need. @@ -12,6 +43,35 @@ These are the variables most users need. | `BIG_MODEL` | `gpt-4o` | OpenAI model used when the Anthropic request specifies a sonnet or opus model. | | `SMALL_MODEL` | `gpt-4o-mini` | OpenAI model used when the Anthropic request specifies a haiku model. | | `RUST_LOG` | `info` | Tracing filter. Examples: `debug`, `anyllm_proxy=trace`. | +| `DISABLE_ADMIN` | (unset) | Set to `1`, `true`, or `yes` to force-disable the admin web interface even when `--webui` is passed. Useful in automated/container environments. | + +## Azure OpenAI + +Set `BACKEND=azure` to route through Azure OpenAI Service. The request/response format is identical to standard OpenAI Chat Completions; only the URL scheme and auth header differ. + +| Variable | Default | Description | +|----------|---------|-------------| +| `AZURE_OPENAI_API_KEY` | (required) | Azure OpenAI API key. Sent as `api-key` header. | +| `AZURE_OPENAI_ENDPOINT` | (required) | Full Azure resource endpoint, e.g. `https://my-resource.openai.azure.com`. Accepts sovereign cloud URLs. | +| `AZURE_OPENAI_DEPLOYMENT` | (required) | Deployment name (the model deployment you created in Azure portal). | +| `AZURE_OPENAI_API_VERSION` | `2024-10-21` | Azure API version string appended as `?api-version=` query parameter. | + +The proxy constructs the full URL as: +``` +{AZURE_OPENAI_ENDPOINT}/openai/deployments/{AZURE_OPENAI_DEPLOYMENT}/chat/completions?api-version={AZURE_OPENAI_API_VERSION} +``` + +### Example + +```bash +BACKEND=azure \ +AZURE_OPENAI_API_KEY=abc123 \ +AZURE_OPENAI_ENDPOINT=https://my-resource.openai.azure.com \ +AZURE_OPENAI_DEPLOYMENT=gpt-4o \ +cargo run -p anyllm_proxy +``` + +--- ## mTLS Client Certificates @@ -46,3 +106,45 @@ TLS_CLIENT_CERT_PASSWORD=changeit \ TLS_CA_CERT=/etc/proxy/corp-ca.pem \ cargo run -p anyllm_proxy ``` + +--- + +## Admin Web UI + +The admin web interface is **opt-in**. Start the proxy with `--webui` or `--admin` to enable it. + +```bash +anyllm_proxy --webui +``` + +The dashboard binds to `localhost:3001` only (never externally accessible). It shows live request logs, latency percentiles, error rates, per-backend metrics, and lets you change log level and model mappings without restarting the server. The Settings tab also displays all active environment variables (secrets are masked). + +| Variable | Default | Description | +|----------|---------|-------------| +| `ADMIN_PORT` | `3001` | Port for the admin dashboard. Must differ from `LISTEN_PORT`. | +| `ADMIN_TOKEN` | (generated) | Bearer token for the admin API. If unset, a random UUID is generated at startup and written to `ADMIN_TOKEN_FILE`. | +| `ADMIN_TOKEN_FILE` | `.admin_token` | File path where the generated admin token is written. Permissions are set to `0600` on Unix. | +| `ADMIN_DB_PATH` | `admin.db` | SQLite database path for request logging and config overrides (model mappings, log level). Config overrides survive restarts. | +| `ADMIN_LOG_RETENTION_DAYS` | `7` | Days to retain request log entries before automatic purge. | +| `DISABLE_ADMIN` | (unset) | Set to `1`, `true`, or `yes` to force-disable the admin server even when `--webui` is passed. Useful in container deployments where the flag might be baked into the entrypoint. | + +### Token security + +The admin token is printed to `ADMIN_TOKEN_FILE` (default `.admin_token`) rather than stdout/stderr, because container log drivers capture stderr and persist it in centralized logging systems. On Unix, the file is created with mode `0600`. + +In production, set `ADMIN_TOKEN` explicitly: + +```bash +ADMIN_TOKEN=$(openssl rand -hex 32) anyllm_proxy --webui +``` + +### Example + +```bash +# Proxy + admin UI on a custom port with a fixed token +ADMIN_PORT=4000 \ +ADMIN_TOKEN=my-secret-token \ +ADMIN_DB_PATH=/var/lib/anyllm/admin.db \ +anyllm_proxy --webui +# Open: http://127.0.0.1:4000/admin/?token=my-secret-token +``` diff --git a/specs/20260325-120000-litellm-gap-fill/contracts/admin-keys.md b/specs/20260325-120000-litellm-gap-fill/contracts/admin-keys.md new file mode 100644 index 0000000..a7351fe --- /dev/null +++ b/specs/20260325-120000-litellm-gap-fill/contracts/admin-keys.md @@ -0,0 +1,93 @@ +# Contract: Admin Virtual Key Management + +All endpoints require admin token auth (`Authorization: Bearer {admin_token}`). Admin server is localhost-only. + +## POST /admin/api/keys + +Create a new virtual API key. + +### Request +```json +{ + "description": "Team Alpha dev key", + "expires_at": "2026-06-01T00:00:00Z", + "rpm_limit": 60, + "tpm_limit": 100000, + "spend_limit": 50.00 +} +``` + +All fields optional. + +### Response (201 Created) +```json +{ + "id": 1, + "key": "sk-vkA1B2C3D4...full-key-shown-once", + "key_prefix": "sk-vkA1B", + "description": "Team Alpha dev key", + "created_at": "2026-03-25T12:00:00Z", + "expires_at": "2026-06-01T00:00:00Z", + "rpm_limit": 60, + "tpm_limit": 100000, + "spend_limit": 50.00 +} +``` + +The `key` field contains the raw API key. It is shown exactly once at creation time and is not stored or retrievable afterward. + +## GET /admin/api/keys + +List all virtual keys (active, expired, and revoked). + +### Response (200 OK) +```json +{ + "keys": [ + { + "id": 1, + "key_prefix": "sk-vkA1B", + "description": "Team Alpha dev key", + "created_at": "2026-03-25T12:00:00Z", + "expires_at": "2026-06-01T00:00:00Z", + "revoked_at": null, + "rpm_limit": 60, + "tpm_limit": 100000, + "spend_limit": 50.00, + "total_spend": 2.34, + "total_requests": 142, + "total_tokens": 53200, + "status": "active" + } + ] +} +``` + +`status` is computed: `"active"`, `"expired"`, or `"revoked"`. + +## DELETE /admin/api/keys/{id} + +Revoke a virtual key. Takes effect immediately (no restart). + +### Response (200 OK) +```json +{ + "id": 1, + "revoked_at": "2026-03-25T14:00:00Z", + "status": "revoked" +} +``` + +### Error (404) +```json +{ + "error": "Key not found" +} +``` + +## Auth check order (proxy middleware) + +1. Check `PROXY_API_KEYS` env-var hashes (existing behavior, backward-compatible) +2. SHA-256 hash the incoming key, look up in DashMap +3. If found: check `revoked_at`, `expires_at`, rate limits +4. If not found in either: reject 401 diff --git a/specs/20260325-120000-litellm-gap-fill/contracts/chat-completions.md b/specs/20260325-120000-litellm-gap-fill/contracts/chat-completions.md new file mode 100644 index 0000000..3bc6763 --- /dev/null +++ b/specs/20260325-120000-litellm-gap-fill/contracts/chat-completions.md @@ -0,0 +1,111 @@ +# Contract: POST /v1/chat/completions + +Accepts OpenAI Chat Completions format, translates internally through the Anthropic pipeline, returns OpenAI format. + +## Request + +``` +POST /v1/chat/completions +Content-Type: application/json +x-api-key: {key} +``` + +### Body + +```json +{ + "model": "claude-sonnet-4-20250514", + "messages": [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"} + ], + "max_tokens": 1024, + "temperature": 0.7, + "stream": false +} +``` + +### Required fields +- `model` (string): Any model ID accepted by the proxy's model mapping +- `messages` (array): At least one message +- `max_tokens` or `max_completion_tokens` (integer): Required (Anthropic constraint); 400 if absent + +### Optional fields (translated) +- `temperature`, `top_p`, `stop`, `tools`, `tool_choice`, `user`, `stream` + +### Optional fields (dropped with `x-anyllm-degradation`) +- `presence_penalty`, `frequency_penalty`, `response_format`, `logprobs`, `top_logprobs`, `n`, `seed`, `stream_options` + +## Response (non-streaming) + +``` +HTTP/1.1 200 OK +Content-Type: application/json +x-request-id: {uuid} +x-anyllm-degradation: presence_penalty,frequency_penalty +``` + +```json +{ + "id": "chatcmpl-{uuid}", + "object": "chat.completion", + "created": 1711360000, + "model": "claude-sonnet-4-20250514", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello! How can I help?" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 12, + "completion_tokens": 8, + "total_tokens": 20 + } +} +``` + +## Response (streaming, `stream: true`) + +``` +HTTP/1.1 200 OK +Content-Type: text/event-stream +Cache-Control: no-cache +``` + +``` +data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":1711360000,"model":"claude-sonnet-4-20250514","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]} + +data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":1711360000,"model":"claude-sonnet-4-20250514","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]} + +data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":1711360000,"model":"claude-sonnet-4-20250514","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]} + +data: [DONE] +``` + +## Error responses + +All errors returned in OpenAI error format: + +```json +{ + "error": { + "message": "max_tokens is required", + "type": "invalid_request_error", + "param": "max_tokens", + "code": null + } +} +``` + +| Condition | Status | Type | +|---|---|---| +| Missing `max_tokens` | 400 | `invalid_request_error` | +| Empty messages | 400 | `invalid_request_error` | +| Invalid API key | 401 | `authentication_error` | +| Rate limited | 429 | `rate_limit_error` | +| Backend error | 502 | `server_error` | diff --git a/specs/20260325-120000-litellm-gap-fill/data-model.md b/specs/20260325-120000-litellm-gap-fill/data-model.md new file mode 100644 index 0000000..f92bf71 --- /dev/null +++ b/specs/20260325-120000-litellm-gap-fill/data-model.md @@ -0,0 +1,169 @@ +# Data Model: LiteLLM Gap Fill + +**Date**: 2026-03-25 | **Branch**: `20260325-120000-litellm-gap-fill` + +--- + +## 1. Virtual API Key (SQLite, new table) + +### Entity: `virtual_api_key` + +| Field | Type | Constraints | Description | +|---|---|---|---| +| `id` | INTEGER | PRIMARY KEY AUTOINCREMENT | Internal row ID | +| `key_hash` | TEXT | NOT NULL, UNIQUE, INDEXED | Hex-encoded SHA-256 of raw key | +| `key_prefix` | TEXT | NOT NULL | First 8 chars of raw key (display only) | +| `description` | TEXT | nullable | Human-readable label | +| `created_at` | TEXT | NOT NULL | ISO 8601 timestamp | +| `expires_at` | TEXT | nullable | ISO 8601 timestamp; NULL = no expiry | +| `revoked_at` | TEXT | nullable | ISO 8601 timestamp; NULL = active | +| `spend_limit` | REAL | nullable | Max USD spend; NULL = unlimited | +| `rpm_limit` | INTEGER | nullable | Max requests/minute; NULL = unlimited | +| `tpm_limit` | INTEGER | nullable | Max tokens/minute; NULL = unlimited | +| `total_spend` | REAL | NOT NULL DEFAULT 0 | Cumulative USD spent | +| `total_requests` | INTEGER | NOT NULL DEFAULT 0 | Cumulative request count | +| `total_tokens` | INTEGER | NOT NULL DEFAULT 0 | Cumulative token count | + +### Relationships +- No foreign keys to other tables. +- `key_hash` is the join key between SQLite persistence and in-memory `DashMap`. + +### State transitions +- **Created**: `revoked_at = NULL`, `expires_at = NULL or future` +- **Active**: `revoked_at = NULL AND (expires_at IS NULL OR expires_at > now())` +- **Expired**: `expires_at <= now() AND revoked_at IS NULL` +- **Revoked**: `revoked_at IS NOT NULL` (terminal; cannot be un-revoked) + +### Validation rules +- `key_hash` must be exactly 64 hex characters (SHA-256). +- `key_prefix` must be 8 characters, starting with `sk-vk`. +- `rpm_limit` and `tpm_limit` must be positive if set. +- `spend_limit` must be non-negative if set. + +--- + +## 2. In-Memory Key Cache (Rust structs) + +### `VirtualKeyMeta` + +```rust +struct VirtualKeyMeta { + id: i64, + description: Option, + expires_at: Option, // epoch seconds + rpm_limit: Option, + tpm_limit: Option, + spend_limit: Option, // USD + rate_state: Arc, +} +``` + +### `RateLimitState` + +```rust +struct RateLimitState { + rpm_window: Mutex>, // request timestamps (ms) + tpm_window: Mutex>, // (timestamp_ms, token_count) +} +``` + +### Cache structure +`DashMap<[u8; 32], VirtualKeyMeta>` keyed by SHA-256 hash bytes. Stored in `SharedState`. + +--- + +## 3. Backend Configuration (Rust enums, extended) + +### `BackendKind` (extended) + +Existing variants: `OpenAI`, `OpenAIResponses`, `Vertex`, `GeminiOpenAI`, `Anthropic` + +New variants: +- `Bedrock` -- AWS Bedrock with SigV4 auth +- `AzureOpenAI` -- Azure OpenAI with `api-key` header and deployment URL + +### `BedrockConfig` + +```rust +struct BedrockConfig { + region: String, // AWS_REGION + access_key_id: String, // AWS_ACCESS_KEY_ID + secret_access_key: String, // AWS_SECRET_ACCESS_KEY (redacted in logs) + session_token: Option, // AWS_SESSION_TOKEN + big_model: String, // e.g., "anthropic.claude-3-5-sonnet-20241022-v2:0" + small_model: String, // e.g., "anthropic.claude-3-5-haiku-20241022-v1:0" +} +``` + +### `AzureOpenAIConfig` + +```rust +struct AzureOpenAIConfig { + endpoint: String, // AZURE_OPENAI_ENDPOINT (full URL) + deployment: String, // AZURE_OPENAI_DEPLOYMENT + api_key: String, // AZURE_OPENAI_API_KEY (redacted in logs) + api_version: String, // AZURE_OPENAI_API_VERSION (default: "2024-10-21") +} +``` + +### `BackendAuth` (extended) + +Existing variants: `BearerToken(String)`, `GoogleApiKey(String)`, `None` + +New variants: +- `AzureApiKey(String)` -- Maps to `api-key: {value}` header +- `AwsSigV4(BedrockCredentials)` -- SigV4 signing applied per-request + +--- + +## 4. Reverse Translation Types (translator crate, new) + +### `ReverseStreamingTranslator` + +```rust +struct ReverseStreamingTranslator { + message_id: String, // from Anthropic message_start + model: String, // from Anthropic message_start + tool_call_index: i32, // tracks current tool_call slot + input_tokens: Option, // from message_start.usage + output_tokens: Option, // from message_delta.usage +} +``` + +### State transitions +1. `New` -> receives `message_start` -> emits first chunk with `role: "assistant"` +2. `TextContent` -> receives `content_block_delta(TextDelta)` -> emits `delta.content` +3. `ToolContent` -> receives `content_block_start(ToolUse)` -> emits `delta.tool_calls[index]` with id/name +4. `ToolContent` -> receives `content_block_delta(InputJsonDelta)` -> emits `delta.tool_calls[index].function.arguments` +5. `ThinkingContent` -> receives `content_block_delta(ThinkingDelta)` -> emits `delta.reasoning_content` +6. `Done` -> receives `message_delta` -> emits `finish_reason` + optional usage chunk +7. `Done` -> receives `message_stop` -> emits `data: [DONE]` + +--- + +## 5. Client Library Types (anyllm_client, extended) + +### `ClientBuilder` + +```rust +struct ClientBuilder { + base_url: Option, + api_key: Option, + timeout: Option, + read_timeout: Option, + max_retries: Option, + tls_config: Option, +} +``` + +### `ToolBuilder` + +```rust +struct ToolBuilder { + name: String, + description: Option, + input_schema: serde_json::Value, +} +``` + +These are convenience wrappers over existing `Tool` and `ToolChoice` types in the translator crate. diff --git a/specs/20260325-120000-litellm-gap-fill/plan.md b/specs/20260325-120000-litellm-gap-fill/plan.md new file mode 100644 index 0000000..99169b8 --- /dev/null +++ b/specs/20260325-120000-litellm-gap-fill/plan.md @@ -0,0 +1,168 @@ +# Implementation Plan: LiteLLM Gap Fill + Rust Client Library + +**Branch**: `20260325-120000-litellm-gap-fill` | **Date**: 2026-03-25 | **Spec**: [spec.md](./spec.md) +**Input**: Feature specification from `/specs/20260325-120000-litellm-gap-fill/spec.md` + +## Summary + +Close the highest-value feature gaps between anyllm-proxy and LiteLLM: accept OpenAI Chat Completions input, add AWS Bedrock and Azure OpenAI backends, implement virtual key management with per-key rate limiting, improve the Rust client library, and add optional OpenTelemetry export. Research is complete (see [research.md](./research.md)). + +## Technical Context + +**Language/Version**: Rust stable, Cargo workspace (3 crates) +**Primary Dependencies**: axum, reqwest, tokio, serde, tracing, rusqlite, sha2; NEW: aws-sigv4, aws-credential-types, dashmap; OPTIONAL: opentelemetry 0.31, tracing-opentelemetry 0.32 +**Storage**: SQLite (existing admin DB, extended with `virtual_api_key` table) +**Testing**: `cargo test` (~480 existing tests); new unit + integration tests per requirement +**Target Platform**: Linux/macOS server, single static binary +**Project Type**: Web service (HTTP proxy) +**Performance Goals**: Existing 100 concurrent request limit; virtual key auth adds one DashMap lookup per request +**Constraints**: Source files under 400 lines (excluding tests); translator crate must remain IO-free +**Scale/Scope**: 7 requirements (5 Tier 1, 2 Tier 2); ~15 new/modified files across 3 crates + +## Constitution Check + +*GATE: Must pass before Phase 0 research. Re-check after Phase 1 design.* + +| Principle | Status | Notes | +|---|---|---| +| I. Security First | PASS | SHA-256 key hashing (existing pattern), SigV4 via audited crate, no secrets in code/logs | +| II. Test Coverage | PASS | Each requirement has acceptance tests defined in spec; TDD approach | +| III. File Size Discipline | PASS | New modules scoped to single concerns; reverse streaming translator is a new file | +| IV. Code Quality | PASS | `cargo clippy -- -D warnings`, `cargo fmt --check` required per quality gates | +| V. Minimal and Correct Changes | PASS | Reuses existing types and patterns; Azure reuses OpenAI client code | +| Dependency Policy | REVIEW NEEDED | 3 new prod deps: `aws-sigv4`, `aws-credential-types`, `dashmap`. Justified: SigV4 cannot be safely hand-rolled; DashMap replaces what would be `RwLock` on the hot auth path. OTEL deps are optional (feature-gated). | + +**Post-design re-check**: All file counts estimated under 400 lines. `reverse_streaming_map.rs` is the largest new file (~250 lines estimated). Constitution compliant. + +## Project Structure + +### Documentation (this feature) + +```text +specs/20260325-120000-litellm-gap-fill/ +├── plan.md # This file +├── research.md # Phase 0 output (complete) +├── data-model.md # Phase 1 output (complete) +├── quickstart.md # Phase 1 output (complete) +├── contracts/ +│ ├── chat-completions.md # POST /v1/chat/completions contract +│ └── admin-keys.md # Virtual key admin API contract +└── tasks.md # Phase 2 output (NOT created by /speckit.plan) +``` + +### Source Code (repository root) + +```text +crates/translator/src/ +├── mapping/ +│ ├── message_map.rs # MODIFIED: add openai_to_anthropic_request, anthropic_to_openai_response +│ ├── reverse_streaming_map.rs # NEW: ReverseStreamingTranslator (Anthropic SSE -> OpenAI chunks) +│ ├── mod.rs # MODIFIED: pub mod reverse_streaming_map +│ └── [existing files unchanged] +├── translate.rs # MODIFIED: add reverse translation convenience wrappers +└── lib.rs # MODIFIED: re-exports + +crates/proxy/src/ +├── config/ +│ └── mod.rs # MODIFIED: BackendKind::Bedrock, BackendKind::AzureOpenAI, env var parsing +├── backend/ +│ ├── mod.rs # MODIFIED: BackendClient::Bedrock, BackendClient::AzureOpenAI variants +│ ├── openai_client.rs # MODIFIED: Azure URL construction + api-key header +│ ├── bedrock_client.rs # NEW: SigV4-signed reqwest client, event stream decoder +│ └── [existing files unchanged] +├── server/ +│ ├── routes.rs # MODIFIED: register POST /v1/chat/completions +│ ├── chat_completions.rs # NEW: handler for OpenAI-format input +│ └── [existing files unchanged] +├── admin/ +│ ├── routes.rs # MODIFIED: add key management endpoints +│ ├── db.rs # MODIFIED: virtual_api_key table DDL + CRUD +│ └── keys.rs # NEW: key generation, hashing, validation logic +├── middleware/ +│ └── auth.rs # MODIFIED: extend to check DashMap virtual keys +├── otel.rs # NEW: OpenTelemetry init (behind #[cfg(feature = "otel")]) +└── main.rs # MODIFIED: OTEL guard, DashMap init from DB + +crates/client/src/ +├── client.rs # MODIFIED: ClientBuilder, Stream return type +├── lib.rs # MODIFIED: re-exports, version 0.2.0 +└── tools.rs # NEW: ToolBuilder, ToolChoiceBuilder helpers +``` + +**Structure Decision**: Existing 3-crate workspace is preserved. No new crates. New functionality distributed across existing module boundaries. The translator crate remains IO-free. + +## Complexity Tracking + +| Violation | Why Needed | Simpler Alternative Rejected Because | +|-----------|------------|-------------------------------------| +| `dashmap` new dependency | Hot-path auth check for virtual keys needs concurrent reads without global lock | `RwLock` serializes all reads during any write; unacceptable for auth middleware on every request | +| `aws-sigv4` + `aws-credential-types` new dependencies | SigV4 request signing requires HMAC chain, canonical request construction, and session token handling | Manual implementation (~150 lines) is error-prone and unaudited; the official crate is ~22K SLoC and well-tested | +| Bedrock event stream decoder | AWS streaming uses binary framing, not SSE | `aws-smithy-eventstream` is already a transitive dep of `aws-sigv4`; alternatively a ~80-line manual parser fits in one file | + +## Implementation Phases + +### Phase A: Reverse Translation (R1 prerequisite) +1. `openai_to_anthropic_request` in `mapping/message_map.rs` +2. `anthropic_to_openai_response` in `mapping/message_map.rs` +3. `ReverseStreamingTranslator` in `mapping/reverse_streaming_map.rs` +4. Unit tests for all new mapping functions +5. Convenience wrappers in `translate.rs` + +### Phase B: Chat Completions Endpoint (R1) +1. `server/chat_completions.rs` handler (non-streaming + streaming) +2. Route registration in `routes.rs` +3. Integration tests (non-streaming, streaming, tool calls, error cases) + +### Phase C: Azure Backend (R3) +1. Config parsing: `BackendKind::AzureOpenAI`, env vars +2. URL construction in `openai_client.rs` (reuse existing client) +3. `api-key` auth header variant +4. Integration test (`#[ignore]`, requires Azure credentials) + +### Phase D: Bedrock Backend (R2) +1. `bedrock_client.rs`: SigV4 signing with `aws-sigv4` +2. Non-streaming `InvokeModel` path +3. Event stream binary decoder for streaming +4. Config parsing: `BackendKind::Bedrock`, env vars +5. Integration test (`#[ignore]`, requires AWS credentials) + +### Phase E: Virtual Key Management (R4) +1. SQLite schema in `admin/db.rs` +2. Key generation and hashing in `admin/keys.rs` +3. Admin API endpoints in `admin/routes.rs` +4. DashMap cache in `SharedState`, loaded from DB on startup +5. Auth middleware extension to check virtual keys +6. Unit + integration tests + +### Phase F: Per-Key Rate Limiting (R7, depends on E) +1. `RateLimitState` with sliding window in `admin/keys.rs` +2. RPM/TPM enforcement in auth middleware +3. HTTP 429 + `retry-after` header on limit exceeded +4. Unit tests for window behavior + +### Phase G: Client Library (R5) +1. `ClientBuilder` in `client/client.rs` +2. `Stream` return type for SSE +3. `ToolBuilder` in `client/tools.rs` +4. Rustdoc examples on all public types +5. Version bump to 0.2.0 + +### Phase H: OpenTelemetry (R6) +1. Feature-gated deps in `Cargo.toml` +2. `otel.rs` initialization module +3. `OpenTelemetryLayer` integration in `main.rs` +4. Span attributes for request ID, model, latency, token counts +5. Manual verification with local OTEL collector + +### Phase Order / Dependencies + +``` +A -> B (reverse translation before endpoint) +C (independent, can parallel with A/B) +D (independent, can parallel with A/B/C) +E -> F (virtual keys before rate limiting) +G (independent) +H (independent) +``` + +Phases A+B are the critical path (highest user value). C and D can proceed in parallel once A is done. diff --git a/specs/20260325-120000-litellm-gap-fill/quickstart.md b/specs/20260325-120000-litellm-gap-fill/quickstart.md new file mode 100644 index 0000000..a3b0331 --- /dev/null +++ b/specs/20260325-120000-litellm-gap-fill/quickstart.md @@ -0,0 +1,103 @@ +# Quickstart: LiteLLM Gap Fill Features + +## 1. OpenAI Chat Completions Input + +After this feature, any OpenAI-native client can use the proxy: + +```bash +# Start proxy backed by OpenAI +OPENAI_API_KEY=sk-... cargo run -p anyllm_proxy + +# Send an OpenAI-format request (NEW) +curl http://localhost:3000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "x-api-key: your-proxy-key" \ + -d '{ + "model": "claude-sonnet-4-20250514", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 100 + }' + +# Response is OpenAI format (not Anthropic) +``` + +## 2. AWS Bedrock Backend + +```bash +BACKEND=bedrock \ +AWS_REGION=us-east-1 \ +AWS_ACCESS_KEY_ID=AKIA... \ +AWS_SECRET_ACCESS_KEY=... \ +BIG_MODEL=anthropic.claude-3-5-sonnet-20241022-v2:0 \ +SMALL_MODEL=anthropic.claude-3-5-haiku-20241022-v1:0 \ +cargo run -p anyllm_proxy +``` + +## 3. Azure OpenAI Backend + +```bash +BACKEND=azure \ +AZURE_OPENAI_ENDPOINT=https://myresource.openai.azure.com \ +AZURE_OPENAI_DEPLOYMENT=my-gpt4o \ +AZURE_OPENAI_API_KEY=... \ +cargo run -p anyllm_proxy +``` + +## 4. Virtual Key Management + +```bash +# Create a key via admin API +curl -X POST http://localhost:3001/admin/api/keys \ + -H "Authorization: Bearer $ADMIN_TOKEN" \ + -H "Content-Type: application/json" \ + -d '{"description": "dev key", "rpm_limit": 60}' + +# Use the returned key +curl http://localhost:3000/v1/messages \ + -H "x-api-key: sk-vk..." \ + -H "Content-Type: application/json" \ + -d '{"model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}]}' + +# Revoke it (takes effect immediately) +curl -X DELETE http://localhost:3001/admin/api/keys/1 \ + -H "Authorization: Bearer $ADMIN_TOKEN" +``` + +## 5. OpenTelemetry Export + +```bash +# Build with OTEL feature +cargo build -p anyllm_proxy --features otel + +# Run with OTEL collector endpoint +OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318 \ +OTEL_SERVICE_NAME=anyllm-proxy \ +OPENAI_API_KEY=sk-... \ +cargo run -p anyllm_proxy --features otel +``` + +## 6. Rust Client Library + +```rust +use anyllm_client::{ClientBuilder, Tool, ToolChoice}; + +let client = ClientBuilder::new() + .base_url("http://localhost:3000") + .api_key("sk-vk...") + .timeout(Duration::from_secs(30)) + .max_retries(3) + .build()?; + +// Non-streaming +let response = client.messages(request).await?; + +// Streaming (returns impl Stream) +let mut stream = client.messages_stream(request).await?; +while let Some(event) = stream.next().await { + match event? { + StreamEvent::ContentBlockDelta { delta, .. } => print!("{}", delta.text()), + StreamEvent::MessageStop => break, + _ => {} + } +} +``` diff --git a/specs/20260325-120000-litellm-gap-fill/research.md b/specs/20260325-120000-litellm-gap-fill/research.md new file mode 100644 index 0000000..b8e1f0f --- /dev/null +++ b/specs/20260325-120000-litellm-gap-fill/research.md @@ -0,0 +1,237 @@ +# Research: LiteLLM Gap Fill + +**Date**: 2026-03-25 | **Branch**: `20260325-120000-litellm-gap-fill` + +--- + +## 1. OpenAI Chat Completions Input Endpoint (R1) + +### Decision +New handler at `POST /v1/chat/completions` with new mapping functions in the translator crate. Not middleware, not a new backend variant. + +### Rationale +- All OpenAI types (`ChatCompletionRequest`, `ChatCompletionResponse`, `ChatCompletionChunk`) already exist in the codebase. +- Several reverse mapping functions already exist: `openai_tools_to_anthropic`, `openai_tool_choice_to_anthropic`, `anthropic_to_openai_usage`. +- A dedicated handler isolates OpenAI compat from the Anthropic-in pipeline. + +### What exists (reusable) +| Function | Direction | +|---|---| +| `openai_tools_to_anthropic` | OpenAI -> Anthropic (exists) | +| `openai_tool_choice_to_anthropic` | OpenAI -> Anthropic (exists) | +| `anthropic_to_openai_usage` | Anthropic -> OpenAI (exists) | +| `StreamingTranslator` | OpenAI chunks -> Anthropic events (exists, forward only) | + +### What must be written +| Function | Location | +|---|---| +| `openai_to_anthropic_request` | `mapping/message_map.rs` | +| `anthropic_to_openai_response` | `mapping/message_map.rs` | +| `ReverseStreamingTranslator` | `mapping/reverse_streaming_map.rs` (new file) | +| `POST /v1/chat/completions` handler | `server/chat_completions.rs` (new file) | + +### Key field mappings (request) +- `messages[role=system]` -> `system` field +- `messages[role=user/assistant]` -> Anthropic `messages[]` +- `messages[role=tool]` -> Anthropic `tool_result` blocks +- `tool_calls` -> `tool_use` blocks (`arguments` JSON string -> `input` JSON object) +- `max_tokens` / `max_completion_tokens` -> `max_tokens` (required in Anthropic; reject 400 if absent) +- `stop` -> `stop_sequences` + +### Lossy fields (drop with `x-anyllm-degradation`) +`presence_penalty`, `frequency_penalty`, `response_format`, `logprobs`, `n`, `seed`, `stream_options` + +### Streaming reverse mapping +Anthropic `StreamEvent` -> OpenAI `ChatCompletionChunk`. The reverse translator tracks message ID, model, tool_call index. OpenAI has no `content_block_start/stop` envelope; tool calls use array index. + +### Open design decision +`max_tokens` is required in Anthropic but optional in OpenAI. Options: (a) reject 400 if absent, (b) supply configurable default (e.g., 4096). Recommend (a) for correctness. + +--- + +## 2. AWS Bedrock Backend (R2) + +### Decision +Use `aws-sigv4` v1.4 + `aws-credential-types` v1.2 for minimal SigV4 signing. No full AWS SDK. + +### Rationale +- `aws-sigv4` adds ~20-30 transitive crates vs ~80-120 for `aws-sdk-bedrockruntime`. +- The project uses reqwest for all HTTP; the full SDK would introduce a parallel hyper-based HTTP stack. +- Manual credential loading from env vars avoids pulling in `aws-config`. + +### Alternatives rejected +| Alternative | Why rejected | +|---|---| +| `aws-sdk-bedrockruntime` | 80-120 crate dependency explosion, hyper conflicts | +| `aws-sign-v4` (third-party) | Sparse maintenance, no session token support | +| `reqsign` | 474K SLoC transitive, wraps `aws-sigv4` anyway | +| Manual SigV4 | Error-prone, security risk | + +### Bedrock API shape +- Non-streaming: `POST /model/{modelId}/invoke` with SigV4 auth +- Request body is Anthropic Messages format + `anthropic_version: "bedrock-2023-05-31"`, model in URL not body +- Response body is raw Anthropic JSON (no Bedrock envelope) +- Streaming: `POST /model/{modelId}/invoke-with-response-stream`, returns AWS Event Stream binary framing +- Per-chunk payload is base64-encoded Anthropic SSE JSON after unwrapping binary frame + +### Streaming complexity +AWS Event Stream is binary framing (4-byte prelude + headers + payload + CRC32), NOT SSE. Requires either `aws-smithy-eventstream` crate or a manual frame parser (~60-100 lines). After decoding, the content is standard Anthropic streaming events usable by existing `StreamingTranslator`. + +### Env vars +`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_SESSION_TOKEN` (optional), `AWS_REGION` + +### Retryable errors +Existing `is_retryable()` covers 429/5xx/408, which maps correctly to Bedrock's `ThrottlingException`, `ModelTimeoutException`, `ServiceUnavailableException`, `InternalServerException`. + +--- + +## 3. Azure OpenAI Backend (R3) + +### Decision +Reuse existing `OpenAIClient` with Azure-specific URL construction and `api-key` auth header. Minimal code changes. + +### Rationale +- Azure Chat Completions request/response body is identical to standard OpenAI. +- Streaming SSE format is identical (`data: {...}\n\n` with `data: [DONE]` terminator). +- No changes needed to translator crate, streaming code, or SSE parser. +- The `model` field in JSON body is ignored by Azure (deployment in URL determines model). + +### URL format +``` +{AZURE_OPENAI_ENDPOINT}/openai/deployments/{AZURE_OPENAI_DEPLOYMENT}/chat/completions?api-version={AZURE_OPENAI_API_VERSION} +``` + +### Auth +`api-key: {key}` header (not `Authorization: Bearer`). New `BackendAuth::AzureApiKey` variant or reuse existing `RequestAuth::Header { name, value }`. + +### Env vars +| Variable | Required | Default | +|---|---|---| +| `AZURE_OPENAI_API_KEY` | Yes | none | +| `AZURE_OPENAI_ENDPOINT` | Yes | none (full URL, e.g., `https://myresource.openai.azure.com`) | +| `AZURE_OPENAI_DEPLOYMENT` | Yes | none | +| `AZURE_OPENAI_API_VERSION` | No | `2024-10-21` | + +### What changes in codebase +1. `config/mod.rs`: Add `BackendKind::AzureOpenAI`, parse env vars, construct URL +2. `backend/mod.rs`: Add `BackendClient::AzureOpenAI(OpenAIClient)` variant +3. `backend/openai_client.rs`: Azure arm in URL construction (pre-constructed from config) +4. Auth mapping for `api-key` header + +--- + +## 4. Virtual Key Management (R4) + +### Decision +SHA-256 hashed keys in SQLite, `DashMap<[u8;32], VirtualKeyMeta>` as in-memory cache, immediate invalidation on revocation. + +### Rationale +- `sha2` and `subtle` already in `Cargo.toml`; existing auth uses SHA-256 + constant-time compare. +- bcrypt/argon2 are wrong for high-entropy API tokens (50-300ms per check at 100 concurrent requests). +- In-memory DashMap avoids SQLite on the hot auth path. +- Follows existing two-phase pattern: SQLite persist, then in-memory apply. + +### Schema +```sql +CREATE TABLE IF NOT EXISTS virtual_api_key ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + key_hash TEXT NOT NULL UNIQUE, + key_prefix TEXT NOT NULL, + description TEXT, + created_at TEXT NOT NULL, + expires_at TEXT, + revoked_at TEXT, + spend_limit REAL, + rpm_limit INTEGER, + tpm_limit INTEGER, + total_spend REAL NOT NULL DEFAULT 0, + total_requests INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0 +); +``` + +### Key generation +Two UUID v4s concatenated (256 bits entropy), prefixed `sk-vk`. Zero new dependencies (`uuid` already in tree). + +### Invalidation +- Admin API writes SQLite first, then updates DashMap (insert on create, remove on revoke). +- On startup, load all non-revoked, non-expired keys from SQLite into DashMap. +- Auth check order: env-var keys (existing `ALLOWED_KEY_HASHES`) -> DashMap virtual keys -> reject. + +### New dependency +`dashmap` for concurrent HashMap. Only new production dependency required. + +--- + +## 5. Per-Key Rate Limiting (R7) + +### Decision +Per-key `Arc` stored inside `VirtualKeyMeta` in the DashMap. Sliding window with `Mutex`. + +### Rationale +- `VecDeque` supports O(1) front-drain for expiry. +- Separate locks for RPM and TPM avoids contention. +- In-memory only (no Redis); sufficient for single-process proxy. + +### Data structure +```rust +struct RateLimitState { + rpm_window: Mutex>, // timestamps + tpm_window: Mutex>, // (timestamp, tokens) +} +``` + +### Alternatives rejected +Token bucket (worse burst control), Redis (overkill), SQLite (too slow for hot path). + +--- + +## 6. Rust Client Library Improvements (R5) + +### Decision +Typed builder pattern, `Stream` return type for SSE, tool-call helpers, comprehensive rustdoc. + +### Rationale +- The existing `Client` struct has a simple config struct but no builder ergonomics. +- Streaming returns raw bytes; should return typed `AnthropicStreamEvent`. +- Tool definitions require manual JSON construction; should have typed builders. + +### Scope +- `ClientBuilder` with method chaining +- `impl Stream>` for streaming responses +- `Tool`, `ToolChoice` builder types +- Re-export all public types from crate root +- Version bump to 0.2.0 + +--- + +## 7. OpenTelemetry Export (R6) + +### Decision +Feature-gated `otel` in `anyllm_proxy/Cargo.toml`. Use `opentelemetry` 0.31 + `tracing-opentelemetry` 0.32 + `opentelemetry-otlp` 0.31 with `http-proto` + `reqwest-client` transport. + +### Rationale +- Reuses existing `reqwest` dependency for OTLP HTTP export; avoids `tonic`/`prost`/`h2` gRPC stack. +- `tracing-opentelemetry` bridges existing `#[tracing::instrument]` spans into OTEL spans without code changes. +- Feature-gated: zero runtime overhead when disabled. + +### Version compatibility +`tracing-opentelemetry` 0.32.x requires `opentelemetry` 0.31.x (deliberate +1 offset). Pin all three together. + +### Cargo feature config +```toml +[features] +otel = ["opentelemetry", "opentelemetry_sdk", "opentelemetry-otlp", "tracing-opentelemetry"] + +[dependencies] +opentelemetry = { version = "0.31", optional = true } +opentelemetry_sdk = { version = "0.31", optional = true } +opentelemetry-otlp = { version = "0.31", features = ["trace", "http-proto", "reqwest-client"], default-features = false, optional = true } +tracing-opentelemetry = { version = "0.32", optional = true } +``` + +### Initialization +Add `OpenTelemetryLayer` to existing `tracing_subscriber::registry()`. `OtelGuard` struct flushes on shutdown. Must fold into the single `.init()` call in `main.rs`. + +### Key env vars +`OTEL_EXPORTER_OTLP_ENDPOINT` (standard), `OTEL_SERVICE_NAME`, `OTEL_TRACES_SAMPLER`. diff --git a/specs/20260325-120000-litellm-gap-fill/spec.md b/specs/20260325-120000-litellm-gap-fill/spec.md new file mode 100644 index 0000000..a0e4b77 --- /dev/null +++ b/specs/20260325-120000-litellm-gap-fill/spec.md @@ -0,0 +1,110 @@ +# Feature Spec: LiteLLM Gap Fill + Rust Client Library Improvements + +**Branch**: `20260325-120000-litellm-gap-fill` +**Date**: 2026-03-25 +**Reference**: `docs/COMPARISON_LITELLM.md` + +## Overview + +Close the highest-value feature gaps between anyllm-proxy and LiteLLM while strengthening the +Rust client library (`anyllm_client`). The goal is not to replicate LiteLLM wholesale but to +eliminate blockers that prevent common OpenAI-native clients from using the proxy, add enterprise +backends (Bedrock, Azure), and provide a first-class Rust SDK experience. + +## Problem Statement + +1. **No OpenAI-format input**: Any client speaking `POST /v1/chat/completions` (OpenAI format) + cannot use the proxy without an additional translation layer. This is the single largest + adoption blocker. + +2. **Missing enterprise backends**: AWS Bedrock (SigV4 auth) and Azure OpenAI (separate URL + scheme, API version param, different auth header) are common enterprise targets with no support. + +3. **Static API key management**: Adding/revoking proxy auth keys requires a process restart. + No per-key metadata, expiry, or spend limits. + +4. **No per-key rate limiting**: Global concurrency limit only. No per-key RPM/TPM enforcement. + +5. **Weak Rust client library**: `anyllm_client` is a thin wrapper, not a first-class SDK. + Missing: typed builder API, streaming ergonomics, tool-call helpers, retry configuration, + and comprehensive documentation/examples. + +6. **No OpenTelemetry export**: Observability is limited to stdout tracing and SQLite logs. + No integration with Datadog, Honeycomb, or other OTEL collectors. + +## Requirements + +### Tier 1 (Must Have) + +**R1. `POST /v1/chat/completions` input endpoint** +- Accept OpenAI Chat Completions format requests on the existing proxy listener +- Translate internally to Anthropic format, forward to configured backend, translate response back +- Support both non-streaming and streaming (`stream: true`) responses +- Return OpenAI-format responses (not Anthropic format) +- Set `x-anyllm-degradation` if features are dropped during reverse translation + +**R2. AWS Bedrock backend** +- New `BACKEND=bedrock` option +- SigV4 request signing (AWS SDK for Rust or manual implementation) +- Support Claude-on-Bedrock model IDs (e.g., `anthropic.claude-3-5-sonnet-20241022-v2:0`) +- Map Anthropic request → Bedrock `InvokeModel` / `InvokeModelWithResponseStream` +- Required env vars: `AWS_REGION`, `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY` (+ optional `AWS_SESSION_TOKEN`) + +**R3. Azure OpenAI backend** +- New `BACKEND=azure` option +- Base URL: `https://{resource}.openai.azure.com/openai/deployments/{deployment}` +- Auth header: `api-key: {key}` (not `Authorization: Bearer`) +- Query param: `api-version=2024-02-01` +- Required env vars: `AZURE_OPENAI_API_KEY`, `AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_DEPLOYMENT`, `AZURE_OPENAI_API_VERSION` + +**R4. Virtual key management** +- Admin API endpoints: `POST /admin/keys`, `GET /admin/keys`, `DELETE /admin/keys/{id}` +- Per-key fields: id, created_at, expires_at (optional), description, spend_limit (optional) +- Keys stored in SQLite (extend existing admin DB) +- Key revocation takes effect immediately without proxy restart +- Backward-compatible: if `PROXY_API_KEYS` env var is set, those keys still work + +**R5. Rust client library improvements** +- Typed builder API with `ClientBuilder` pattern +- Streaming API: `impl Stream` return type +- Tool-call helpers: typed `Tool`, `ToolChoice` builders +- Comprehensive rustdoc with examples +- Re-export all public types from crate root +- `anyllm_client` version bump to 0.2.0 + +### Tier 2 (Should Have) + +**R6. OpenTelemetry export** +- Feature-gated with `features = ["otel"]` in `anyllm_proxy/Cargo.toml` +- Export spans to OTEL collector via `opentelemetry-otlp` +- Env var: `OTEL_EXPORTER_OTLP_ENDPOINT` (standard OTEL env var) +- Span attributes: request ID, model, backend, latency, token counts, degradation flags + +**R7. Per-key rate limiting** +- RPM (requests per minute) and TPM (tokens per minute) limits per virtual key +- Add `rpm_limit` and `tpm_limit` fields to virtual key schema +- In-memory rate limit state (atomic counters with 60s sliding window) +- Return HTTP 429 with standard `retry-after` header when limit exceeded +- Requires R4 (virtual keys) to be complete first + +### Out of Scope + +- Response caching (Redis dependency, significant scope) +- Cross-provider fallback chains (router redesign required) +- Real batch processing (async job queue, out of scope) +- Cost tracking / pricing database +- RBAC / OIDC / SAML +- Audio, image, reranking endpoints + +## Acceptance Criteria + +1. `cargo test` passes (including new tests for each requirement) +2. `cargo clippy -- -D warnings` clean +3. `cargo fmt --check` clean +4. All new source files under 400 lines +5. `POST /v1/chat/completions` works with curl against a running proxy backed by OpenAI +6. Bedrock backend connects and returns a response (tested with `#[ignore]` live test) +7. Azure backend connects and returns a response (tested with `#[ignore]` live test) +8. Virtual key CRUD via admin API, key revocation verified without restart +9. `anyllm_client` rustdoc builds without warnings (`cargo doc --no-deps`) +10. OTEL spans visible in a local collector when feature flag is enabled (manual verification) diff --git a/specs/20260325-120000-litellm-gap-fill/tasks.md b/specs/20260325-120000-litellm-gap-fill/tasks.md new file mode 100644 index 0000000..22d5ba3 --- /dev/null +++ b/specs/20260325-120000-litellm-gap-fill/tasks.md @@ -0,0 +1,300 @@ +# Tasks: LiteLLM Gap Fill + Rust Client Library + +**Input**: Design documents from `/specs/20260325-120000-litellm-gap-fill/` +**Prerequisites**: plan.md, spec.md, research.md, data-model.md, contracts/ + +**Organization**: Tasks grouped by user story (one per requirement). Each story is independently implementable and testable after the foundational phase. + +## Format: `[ID] [P?] [Story] Description` + +- **[P]**: Can run in parallel (different files, no dependencies) +- **[Story]**: Which user story this task belongs to (US1-US7) +- Exact file paths included in every task description + +## Path Conventions + +- **Translator crate**: `crates/translator/src/` +- **Proxy crate**: `crates/proxy/src/` +- **Client crate**: `crates/client/src/` +- **Integration tests**: `crates/proxy/tests/` + +--- + +## Phase 1: Setup + +**Purpose**: Add new dependencies and create empty module scaffolding + +- [x] T001 Add `dashmap = "6"` to `crates/proxy/Cargo.toml` dependencies +- [x] T002 [P] Add `aws-sigv4 = { version = "1.4", features = ["sign-http"] }` and `aws-credential-types = "1.2"` to `crates/proxy/Cargo.toml` dependencies +- [x] T003 [P] Add feature-gated OTEL dependencies to `crates/proxy/Cargo.toml`: `opentelemetry`, `opentelemetry_sdk`, `opentelemetry-otlp` (with `trace`, `http-proto`, `reqwest-client` features), `tracing-opentelemetry` all as optional under `[features] otel = [...]` +- [x] T004 Add `pub mod reverse_streaming_map;` to `crates/translator/src/mapping/mod.rs` + +--- + +## Phase 2: Foundational (Reverse Translation Mapping) + +**Purpose**: Pure translation functions required by US1. These are IO-free functions in the translator crate that convert OpenAI request types to Anthropic types and vice versa. MUST complete before US1 can begin. + +- [x] T005 Implement `openai_to_anthropic_request(req: &ChatCompletionRequest) -> Result` in `crates/translator/src/mapping/message_map.rs`. Must handle: system message extraction, user/assistant/tool message conversion, `tool_calls` -> `tool_use` blocks, `max_tokens`/`max_completion_tokens` -> `max_tokens` (reject if absent), `stop` -> `stop_sequences`, `temperature`/`top_p` passthrough. Drop unsupported fields (`presence_penalty`, `frequency_penalty`, `response_format`, `logprobs`, `n`, `seed`) and record them via `TranslationWarnings`. +- [x] T006 [P] Implement reverse stop_reason mapping helper in `crates/translator/src/mapping/message_map.rs`: `anthropic_stop_reason_to_openai(stop_reason: &StopReason) -> &str` mapping `end_turn`->`stop`, `max_tokens`->`length`, `tool_use`->`tool_calls`, `stop_sequence`->`stop`. +- [x] T007 Implement `anthropic_to_openai_response(resp: &MessageResponse, model: &str) -> ChatCompletionResponse` in `crates/translator/src/mapping/message_map.rs`. Must handle: text content concatenation, `tool_use` -> `tool_calls` (input object -> arguments string), thinking blocks -> `reasoning_content`, stop_reason mapping, usage mapping (reuse existing `anthropic_to_openai_usage`), generate `chatcmpl-` prefixed ID. +- [x] T008 Create `ReverseStreamingTranslator` struct in `crates/translator/src/mapping/reverse_streaming_map.rs`. Fields: `message_id`, `model`, `tool_call_index`, `input_tokens`, `output_tokens`. Implement `fn new(id: String, model: String) -> Self`. +- [x] T009 Implement `fn process_event(&mut self, event: &StreamEvent) -> Vec` on `ReverseStreamingTranslator` in `crates/translator/src/mapping/reverse_streaming_map.rs`. Map: `message_start` -> first chunk with `role: "assistant"`, `content_block_delta(TextDelta)` -> `delta.content`, `content_block_start(ToolUse)` -> `delta.tool_calls[index]` with id/name, `content_block_delta(InputJsonDelta)` -> `delta.tool_calls[index].function.arguments`, `content_block_delta(ThinkingDelta)` -> `delta.reasoning_content`, `message_delta` -> `finish_reason` chunk, `message_stop` -> `[DONE]` sentinel. +- [x] T010 Add unit tests for `openai_to_anthropic_request` in `crates/translator/src/mapping/message_map.rs` `#[cfg(test)]` module: basic message conversion, system message extraction, tool call conversion, missing max_tokens rejection, lossy field warnings. +- [x] T011 [P] Add unit tests for `anthropic_to_openai_response` in `crates/translator/src/mapping/message_map.rs` `#[cfg(test)]` module: text response, tool use response, thinking blocks, stop reason mapping, usage mapping. +- [x] T012 [P] Add unit tests for `ReverseStreamingTranslator` in `crates/translator/src/mapping/reverse_streaming_map.rs` `#[cfg(test)]` module: text streaming, tool call streaming with index tracking, thinking content, finish reason, `[DONE]` emission. +- [x] T013 Add convenience wrappers `translate_openai_to_anthropic_request` and `translate_anthropic_to_openai_response` in `crates/translator/src/translate.rs`. Re-export `ReverseStreamingTranslator` from `crates/translator/src/lib.rs`. + +**Checkpoint**: `cargo test -p anyllm_translate` passes with all new reverse translation tests. All mapping functions are pure (no IO). + +--- + +## Phase 3: User Story 1 - OpenAI Chat Completions Input (Priority: P1, MVP) + +**Goal**: Accept `POST /v1/chat/completions` in OpenAI format, translate through Anthropic pipeline, return OpenAI format. Highest-value feature: unblocks all OpenAI-native clients. + +**Independent Test**: `curl -X POST http://localhost:3000/v1/chat/completions -H "x-api-key: test" -H "Content-Type: application/json" -d '{"model":"claude-sonnet-4-20250514","messages":[{"role":"user","content":"Hello"}],"max_tokens":100}'` returns OpenAI-format JSON. + +### Implementation for User Story 1 + +- [x] T014 [US1] Create `crates/proxy/src/server/chat_completions.rs` with non-streaming handler: extract `Json`, call `openai_to_anthropic_request`, dispatch to `BackendClient`, call `anthropic_to_openai_response`, return `Json`. Set `x-anyllm-degradation` header from `TranslationWarnings`. Return OpenAI-shaped errors on validation failure (missing max_tokens -> 400 `invalid_request_error`). +- [x] T015 [US1] Add streaming handler in `crates/proxy/src/server/chat_completions.rs`: when `stream: true`, dispatch to backend streaming path, create `ReverseStreamingTranslator`, emit `text/event-stream` with `data: {chunk}\n\n` lines (no `event:` prefix, matching OpenAI SSE format). Terminate with `data: [DONE]\n\n`. +- [x] T016 [US1] Register `POST /v1/chat/completions` route in `crates/proxy/src/server/routes.rs` on the existing backend router. Apply same middleware (auth, request ID, size limit, concurrency limit) as the `/v1/messages` route. +- [ ] T017 [US1] Add integration tests in `crates/proxy/tests/` (new file `chat_completions.rs` or extend existing): non-streaming basic response, streaming basic response, tool call round-trip, missing max_tokens returns 400, degradation header set for lossy fields, empty messages returns 400. + +**Checkpoint**: `cargo test -p anyllm_proxy` passes. `POST /v1/chat/completions` works end-to-end with a mock or live backend. + +--- + +## Phase 4: User Story 2 - Azure OpenAI Backend (Priority: P1) + +**Goal**: `BACKEND=azure` routes requests through Azure OpenAI using deployment-scoped URLs and `api-key` header. Reuses existing OpenAI client code. + +**Independent Test**: `BACKEND=azure AZURE_OPENAI_ENDPOINT=https://... AZURE_OPENAI_DEPLOYMENT=gpt4o AZURE_OPENAI_API_KEY=... cargo run -p anyllm_proxy` starts and responds to `/v1/messages`. + +### Implementation for User Story 2 + +- [ ] T018 [US2] Add `BackendKind::AzureOpenAI` variant to the backend enum in `crates/proxy/src/config/mod.rs`. Parse `AZURE_OPENAI_API_KEY`, `AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_DEPLOYMENT`, `AZURE_OPENAI_API_VERSION` (default `"2024-10-21"`) from env. Construct full URL: `{endpoint}/openai/deployments/{deployment}/chat/completions?api-version={version}` at config load time. Validate URL with existing `validate_url` function. +- [ ] T019 [US2] Add `BackendAuth::AzureApiKey(String)` variant (or equivalent) in `crates/proxy/src/backend/mod.rs`. Map it to `RequestAuth::Header { name: "api-key", value: key }` in the auth application logic. Add `BackendClient::AzureOpenAI(OpenAIClient)` variant that constructs `OpenAIClient` with the pre-built Azure URL and `AzureApiKey` auth. +- [ ] T020 [US2] Modify `crates/proxy/src/backend/openai_client.rs` to accept Azure's pre-constructed URL. The `chat_completions_url` for Azure is the full URL from config (no `/v1/chat/completions` suffix appended). Ensure the `model` field in the request body is still populated (Azure ignores it but accepts it). +- [ ] T021 [US2] Add `#[ignore]` integration test in `crates/proxy/tests/` for Azure backend: send a request via the proxy configured with `BACKEND=azure`, verify response is valid Anthropic format. Requires `AZURE_OPENAI_API_KEY` env var to run. +- [ ] T022 [US2] Update `docs/ENV.md` with Azure-specific env vars and usage example. + +**Checkpoint**: `cargo build` clean. Azure config parsing tested. `#[ignore]` live test exists. + +--- + +## Phase 5: User Story 3 - AWS Bedrock Backend (Priority: P1) + +**Goal**: `BACKEND=bedrock` routes requests through AWS Bedrock using SigV4-signed requests. Non-streaming via `InvokeModel`, streaming via `InvokeModelWithResponseStream` with binary event stream decoding. + +**Independent Test**: `BACKEND=bedrock AWS_REGION=us-east-1 AWS_ACCESS_KEY_ID=... AWS_SECRET_ACCESS_KEY=... BIG_MODEL=anthropic.claude-3-5-sonnet-20241022-v2:0 cargo run -p anyllm_proxy` starts and responds. + +### Implementation for User Story 3 + +- [ ] T023 [US3] Create `crates/proxy/src/backend/bedrock_client.rs` with `BedrockClient` struct. Fields: `http_client: reqwest::Client`, `region: String`, `credentials: aws_credential_types::Credentials`, `big_model: String`, `small_model: String`. Implement `fn new(config: &BedrockConfig, http_client: reqwest::Client) -> Self`. +- [ ] T024 [US3] Implement non-streaming `send_request` on `BedrockClient` in `crates/proxy/src/backend/bedrock_client.rs`: build Bedrock URL (`https://bedrock-runtime.{region}.amazonaws.com/model/{model_id}/invoke`), serialize Anthropic request body with `anthropic_version: "bedrock-2023-05-31"` (model field omitted from body), sign request with `aws_sigv4::http_request::sign()`, send via reqwest, deserialize response as `MessageResponse`. +- [ ] T025 [US3] Implement AWS Event Stream binary frame decoder in `crates/proxy/src/backend/bedrock_client.rs` (or a submodule): parse 4-byte prelude length, 4-byte headers length, headers, payload, 4-byte CRC32 checksum. Extract `chunk.bytes` field, base64-decode to get Anthropic SSE JSON. Target: ~80 lines. +- [ ] T026 [US3] Implement streaming `send_request_stream` on `BedrockClient` in `crates/proxy/src/backend/bedrock_client.rs`: build URL with `/invoke-with-response-stream`, sign request, send via reqwest with streaming response, pipe response bytes through event stream decoder, yield Anthropic `StreamEvent` items compatible with existing `StreamingTranslator`. +- [ ] T027 [US3] Add `BackendKind::Bedrock` to `crates/proxy/src/config/mod.rs`. Parse `AWS_REGION`, `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_SESSION_TOKEN` (optional) from env. Store as `BedrockConfig`. Add `BackendClient::Bedrock(BedrockClient)` variant to `crates/proxy/src/backend/mod.rs` and wire through dispatch. +- [ ] T028 [US3] Add unit tests for event stream decoder in `crates/proxy/src/backend/bedrock_client.rs` `#[cfg(test)]` module: parse a known binary frame, extract payload, verify CRC, handle partial frames. +- [ ] T029 [US3] Add `#[ignore]` integration test in `crates/proxy/tests/` for Bedrock backend: non-streaming and streaming paths. Requires AWS credentials. +- [ ] T030 [US3] Update `docs/ENV.md` with Bedrock-specific env vars and usage example. + +**Checkpoint**: `cargo build` clean. Event stream decoder unit tests pass. `#[ignore]` live tests exist. + +--- + +## Phase 6: User Story 4 - Virtual Key Management (Priority: P1) + +**Goal**: Admin API for creating, listing, and revoking API keys stored in SQLite. Keys take effect immediately without proxy restart. + +**Independent Test**: `POST /admin/api/keys` returns a key; that key authenticates against `/v1/messages`; `DELETE /admin/api/keys/{id}` revokes it; subsequent requests with that key return 401. + +### Implementation for User Story 4 + +- [x] T031 [US4] Add `virtual_api_key` table DDL to `crates/proxy/src/admin/db.rs` in the existing `init_db` function (or equivalent). Schema per `data-model.md`: `id INTEGER PRIMARY KEY AUTOINCREMENT`, `key_hash TEXT NOT NULL UNIQUE`, `key_prefix TEXT NOT NULL`, `description TEXT`, `created_at TEXT NOT NULL`, `expires_at TEXT`, `revoked_at TEXT`, `spend_limit REAL`, `rpm_limit INTEGER`, `tpm_limit INTEGER`, `total_spend REAL NOT NULL DEFAULT 0`, `total_requests INTEGER NOT NULL DEFAULT 0`, `total_tokens INTEGER NOT NULL DEFAULT 0`. Add index on `key_hash`. +- [x] T032 [US4] Create `crates/proxy/src/admin/keys.rs` with key generation and hashing. `fn generate_virtual_key() -> (String, String, [u8; 32])` returns `(raw_key, key_prefix, key_hash)`. Use two UUID v4s concatenated with `sk-vk` prefix for the raw key. Hash with SHA-256 (reuse existing `sha2` dependency). `key_prefix` is first 8 chars. +- [x] T033 [US4] Add CRUD functions in `crates/proxy/src/admin/db.rs`: `insert_virtual_key(conn, key_hash, key_prefix, description, expires_at, rpm_limit, tpm_limit, spend_limit)`, `list_virtual_keys(conn) -> Vec`, `revoke_virtual_key(conn, id) -> Option`, `load_active_virtual_keys(conn) -> Vec`. +- [x] T034 [US4] Add `DashMap<[u8; 32], VirtualKeyMeta>` to `SharedState` in `crates/proxy/src/admin/state.rs` (or wherever `SharedState` is defined). On startup in `crates/proxy/src/main.rs`, call `load_active_virtual_keys` and populate the DashMap. +- [x] T035 [US4] Add admin API endpoints in `crates/proxy/src/admin/routes.rs`: `POST /admin/api/keys` (create key, insert to DB, insert to DashMap, return raw key once), `GET /admin/api/keys` (list from DB with computed status), `DELETE /admin/api/keys/{id}` (set `revoked_at` in DB, remove from DashMap, return confirmation). +- [x] T036 [US4] Extend auth middleware in `crates/proxy/src/server/middleware.rs` to check the DashMap after checking env-var keys. SHA-256 hash the incoming credential, look up in DashMap, verify `revoked_at` is None and `expires_at` is not past. If both checks fail, return 401. +- [x] T037 [US4] Add unit tests for key generation and hashing in `crates/proxy/src/admin/keys.rs` `#[cfg(test)]` module: key format, prefix extraction, hash determinism. +- [ ] T038 [US4] Add integration tests for virtual key admin API in `crates/proxy/tests/`: create key, list keys, use key for auth, revoke key, verify revoked key is rejected. + +**Checkpoint**: `cargo test -p anyllm_proxy` passes. Virtual key CRUD works. Key revocation is immediate. + +--- + +## Phase 7: User Story 5 - Rust Client Library Improvements (Priority: P1) + +**Goal**: `anyllm_client` becomes a first-class Rust SDK with builder pattern, typed streaming, and tool helpers. + +**Independent Test**: `cargo doc -p anyllm_client --no-deps` builds without warnings. `cargo test -p anyllm_client` passes. + +### Implementation for User Story 5 + +- [ ] T039 [P] [US5] Add `ClientBuilder` to `crates/client/src/client.rs` with method chaining: `fn new() -> Self`, `fn base_url(mut self, url: &str) -> Self`, `fn api_key(mut self, key: &str) -> Self`, `fn timeout(mut self, d: Duration) -> Self`, `fn read_timeout(mut self, d: Duration) -> Self`, `fn max_retries(mut self, n: u32) -> Self`, `fn tls_config(mut self, cfg: TlsConfig) -> Self`, `fn build(self) -> Result`. Implement `Client::builder() -> ClientBuilder` convenience method. +- [ ] T040 [P] [US5] Create `crates/client/src/tools.rs` with `ToolBuilder` and `ToolChoiceBuilder`. `ToolBuilder`: `fn new(name: &str) -> Self`, `fn description(mut self, desc: &str) -> Self`, `fn input_schema(mut self, schema: Value) -> Self`, `fn build(self) -> Tool`. `ToolChoiceBuilder`: `fn auto() -> ToolChoice`, `fn any() -> ToolChoice`, `fn none() -> ToolChoice`, `fn specific(name: &str) -> ToolChoice`. +- [ ] T041 [US5] Add streaming return type to `crates/client/src/client.rs`: `fn messages_stream(&self, req: MessageCreateRequest) -> Result>, ClientError>`. Parse SSE frames from the reqwest response byte stream, deserialize each `data:` line into `StreamEvent`. +- [ ] T042 [US5] Update `crates/client/src/lib.rs` to re-export all public types: `Client`, `ClientBuilder`, `ClientConfig`, `ClientError`, `Tool`, `ToolBuilder`, `ToolChoice`, `ToolChoiceBuilder`, `StreamEvent`, and all Anthropic request/response types from `anyllm_translate`. +- [ ] T043 [US5] Add rustdoc examples to all public types and methods in `crates/client/src/client.rs`, `crates/client/src/tools.rs`, and `crates/client/src/lib.rs`. Each builder method and each public function gets a `/// # Examples` block. +- [ ] T044 [US5] Bump `anyllm_client` version to `0.2.0` in `crates/client/Cargo.toml`. +- [ ] T045 [US5] Add unit tests for `ClientBuilder` (valid build, missing required fields), `ToolBuilder`, and `ToolChoiceBuilder` in their respective `#[cfg(test)]` modules. + +**Checkpoint**: `cargo doc -p anyllm_client --no-deps` builds clean. `cargo test -p anyllm_client` passes. + +--- + +## Phase 8: User Story 6 - Per-Key Rate Limiting (Priority: P2, depends on US4) + +**Goal**: RPM and TPM limits per virtual key with sliding window enforcement. Returns 429 with `retry-after` when exceeded. + +**Independent Test**: Create a key with `rpm_limit: 2`, send 3 requests, third returns 429. + +### Implementation for User Story 6 + +- [x] T046 [US6] Add `RateLimitState` struct to `crates/proxy/src/admin/keys.rs`: `rpm_window: Mutex>`, `tpm_window: Mutex>`. Add `fn check_rpm(&self, limit: u32) -> Result<(), Duration>` (returns Ok or Err with retry-after duration) and `fn record_rpm(&self)`. Same pattern for TPM: `fn check_tpm(&self, limit: u32, tokens: u32) -> Result<(), Duration>` and `fn record_tpm(&self, tokens: u32)`. Drain entries older than 60 seconds on each check. +- [x] T047 [US6] Add `rate_state: Arc` field to `VirtualKeyMeta` in DashMap. Initialize a new `RateLimitState` for each key loaded on startup and each key created via admin API. +- [x] T048 [US6] Extend auth middleware in `crates/proxy/src/server/middleware.rs`: after virtual key validation passes, check `rate_state.check_rpm(key.rpm_limit)`. If exceeded, return HTTP 429 with `retry-after: {seconds}` header and OpenAI-shaped rate limit error body. TPM check happens after the response (post-middleware or in the handler) since token count is only known after the backend responds. +- [ ] T049 [US6] Add post-response TPM recording: after the backend response is received and token count is known, call `rate_state.record_tpm(output_tokens)`. If TPM would be exceeded, the next request's pre-check catches it. +- [x] T050 [US6] Add unit tests for `RateLimitState` in `crates/proxy/src/admin/keys.rs` `#[cfg(test)]` module: window expiry, RPM enforcement, TPM enforcement, concurrent access safety. +- [ ] T051 [US6] Add integration test for rate limiting in `crates/proxy/tests/`: create key with `rpm_limit: 2`, send 2 requests (200), send 3rd request (429 with `retry-after` header). + +**Checkpoint**: `cargo test -p anyllm_proxy` passes. Rate limiting enforced per-key. + +--- + +## Phase 9: User Story 7 - OpenTelemetry Export (Priority: P2) + +**Goal**: Optional OTEL span export via feature flag. When enabled, all request spans are exported to an OTLP collector with request metadata as attributes. + +**Independent Test**: `cargo build -p anyllm_proxy --features otel` compiles. With a local OTEL collector running, spans appear in the collector UI. + +### Implementation for User Story 7 + +- [ ] T052 [US7] Create `crates/proxy/src/otel.rs` behind `#[cfg(feature = "otel")]`. Implement `fn init_otel() -> OtelGuard`: build `SdkTracerProvider` with `opentelemetry-otlp` `SpanExporter` (http-proto, reqwest-client), set global tracer provider, set `TraceContextPropagator`. Return `OtelGuard` struct whose `Drop` impl calls `provider.shutdown()`. +- [ ] T053 [US7] Modify tracing subscriber initialization in `crates/proxy/src/main.rs`: under `#[cfg(feature = "otel")]`, add `OpenTelemetryLayer::new(tracer)` to the existing `tracing_subscriber::registry()` chain. Store `OtelGuard` in a variable that lives for the duration of `main`. Ensure the non-otel path is unchanged via `#[cfg(not(feature = "otel"))]`. +- [ ] T054 [US7] Add span attributes to request handlers: in the existing request middleware or handler instrumentation, record `http.request.id`, `gen_ai.request.model`, `gen_ai.response.model`, `http.response.status_code`, `gen_ai.usage.input_tokens`, `gen_ai.usage.output_tokens` via `tracing::Span::current().record(...)`. Ensure the `#[tracing::instrument]` macros declare these fields. +- [ ] T055 [US7] Verify `cargo build -p anyllm_proxy` (without `otel` feature) still compiles and has no OTEL dependencies. Verify `cargo build -p anyllm_proxy --features otel` compiles clean. +- [ ] T056 [US7] Update `docs/ENV.md` with OTEL-related env vars: `OTEL_EXPORTER_OTLP_ENDPOINT`, `OTEL_SERVICE_NAME`, `OTEL_TRACES_SAMPLER`. Document the `--features otel` build flag. + +**Checkpoint**: Both `cargo build` (default) and `cargo build --features otel` compile. No runtime overhead when feature is off. + +--- + +## Phase 10: Polish and Cross-Cutting Concerns + +**Purpose**: Final validation, documentation updates, and CI adjustments + +- [ ] T057 [P] Update `docs/COMPARISON_LITELLM.md` to reflect closed gaps: `POST /v1/chat/completions` input, Bedrock backend, Azure backend, virtual key management, per-key rate limiting, OTEL export. Move items from "Major gap" to "Advantage" or "Parity" as appropriate. +- [ ] T058 [P] Update `CLAUDE.md` with new backend types, new env vars, new admin endpoints, new source files, and updated test counts. +- [ ] T059 [P] Update `README.md` with quickstart examples for new features (reference `quickstart.md` content). +- [ ] T060 Run `cargo clippy -- -D warnings` across all crates and fix any warnings. +- [ ] T061 Run `cargo fmt --check` and fix any formatting issues. +- [ ] T062 Run `cargo test` full suite and verify all tests pass (expect ~550+ tests). +- [ ] T063 Verify all new source files are under 400 lines (excluding `#[cfg(test)]` modules). + +--- + +## Dependencies and Execution Order + +### Phase Dependencies + +- **Phase 1 (Setup)**: No dependencies, start immediately +- **Phase 2 (Foundational)**: Depends on Phase 1 (T004 specifically) +- **Phase 3 (US1)**: Depends on Phase 2 completion +- **Phase 4 (US2)**: Depends on Phase 1 only (independent of Phase 2) +- **Phase 5 (US3)**: Depends on Phase 1 only (independent of Phase 2) +- **Phase 6 (US4)**: Depends on Phase 1 only (independent of Phase 2) +- **Phase 7 (US5)**: Depends on Phase 1 only (independent) +- **Phase 8 (US6)**: Depends on Phase 6 (US4) completion +- **Phase 9 (US7)**: Depends on Phase 1 (T003 specifically) +- **Phase 10 (Polish)**: Depends on all user stories + +### User Story Dependencies + +- **US1 (Chat Completions)**: Requires Phase 2 (reverse translation). Critical path. +- **US2 (Azure)**: Independent. Can start after Phase 1. +- **US3 (Bedrock)**: Independent. Can start after Phase 1. +- **US4 (Virtual Keys)**: Independent. Can start after Phase 1. +- **US5 (Client Library)**: Independent. Can start after Phase 1. +- **US6 (Rate Limiting)**: Depends on US4 completion. +- **US7 (OTEL)**: Independent. Can start after Phase 1. + +### Within Each User Story + +- Types/models before services +- Services before handlers/endpoints +- Core implementation before integration tests +- Unit tests alongside implementation + +### Parallel Opportunities + +After Phase 1 completes, up to 5 user stories can proceed in parallel: + +``` +Phase 1 (Setup) + | + +---> Phase 2 (Foundational) ---> Phase 3 (US1: Chat Completions) + | + +---> Phase 4 (US2: Azure) + | + +---> Phase 5 (US3: Bedrock) + | + +---> Phase 6 (US4: Virtual Keys) ---> Phase 8 (US6: Rate Limiting) + | + +---> Phase 7 (US5: Client Library) + | + +---> Phase 9 (US7: OTEL) +``` + +--- + +## Parallel Example: After Phase 1 + +``` +# These can all run simultaneously: +Agent 1: Phase 2 (T005-T013) -> Phase 3 (T014-T017) +Agent 2: Phase 4 (T018-T022) Azure backend +Agent 3: Phase 5 (T023-T030) Bedrock backend +Agent 4: Phase 6 (T031-T038) Virtual keys -> Phase 8 (T046-T051) Rate limiting +Agent 5: Phase 7 (T039-T045) Client library +Agent 6: Phase 9 (T052-T056) OTEL +``` + +--- + +## Implementation Strategy + +### MVP First (US1 Only) + +1. Complete Phase 1: Setup (T001-T004) +2. Complete Phase 2: Foundational reverse translation (T005-T013) +3. Complete Phase 3: US1 Chat Completions endpoint (T014-T017) +4. **STOP and VALIDATE**: `POST /v1/chat/completions` works with curl +5. This alone closes the single largest adoption gap + +### Incremental Delivery + +1. Setup + Foundational + US1 -> Chat Completions works (MVP) +2. Add US2 (Azure) -> Enterprise Azure users unblocked +3. Add US3 (Bedrock) -> Enterprise AWS users unblocked +4. Add US4 (Virtual Keys) -> Dynamic key management +5. Add US5 (Client Library) -> First-class Rust SDK +6. Add US6 (Rate Limiting) -> Per-key enforcement +7. Add US7 (OTEL) -> Observability integration +8. Polish phase -> Documentation and CI + +Each story adds value independently without breaking previous stories. + +--- + +## Notes + +- [P] tasks = different files, no dependencies on incomplete tasks +- [USn] label maps task to specific user story +- Translator crate must remain IO-free (no reqwest, no tokio, no file access) +- All new files must be under 400 lines (excluding `#[cfg(test)]` modules) +- Commit after each task or logical group +- `cargo clippy -- -D warnings` must stay clean throughout