mirror of
https://github.com/whit3rabbit/anyllm-proxy.git
synced 2026-09-21 16:00:49 +00:00
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) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.6
parent
f1df50ff37
commit
a4e655c8bb
@@ -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"] }
|
||||
|
||||
+179
-80
@@ -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<InternalError> 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<String>,
|
||||
api_key: Option<String>,
|
||||
timeout: Option<std::time::Duration>,
|
||||
read_timeout: Option<std::time::Duration>,
|
||||
max_retries: Option<u32>,
|
||||
}
|
||||
|
||||
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<Client, ClientError> {
|
||||
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<Result<StreamEvent, ClientError>>,
|
||||
}
|
||||
}
|
||||
|
||||
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::<ChatCompletionChunk>(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<StreamEvent, ClientError>;
|
||||
|
||||
fn poll_next(
|
||||
self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Option<Self::Item>> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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<Result<StreamEvent, ClientError>>,
|
||||
}
|
||||
}
|
||||
|
||||
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::<ChatCompletionChunk>(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<StreamEvent, ClientError>;
|
||||
|
||||
fn poll_next(
|
||||
self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Option<Self::Item>> {
|
||||
self.project().inner.poll_next(cx)
|
||||
}
|
||||
}
|
||||
@@ -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<String>,
|
||||
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");
|
||||
}
|
||||
}
|
||||
+31
-2
@@ -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"
|
||||
|
||||
@@ -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<u32>,
|
||||
tpm_limit: Option<u32>,
|
||||
spend_limit: Option<f64>,
|
||||
) -> rusqlite::Result<i64> {
|
||||
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<Vec<VirtualKeyRow>> {
|
||||
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<i64>>(7)?.map(|v| v as u32),
|
||||
tpm_limit: row.get::<_, Option<i64>>(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<Option<VirtualKeyRow>> {
|
||||
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<i64>>(7)?.map(|v| v as u32),
|
||||
tpm_limit: row.get::<_, Option<i64>>(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<Vec<VirtualKeyRow>> {
|
||||
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<i64>>(7)?.map(|v| v as u32),
|
||||
tpm_limit: row.get::<_, Option<i64>>(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::*;
|
||||
|
||||
@@ -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<String>,
|
||||
/// Epoch seconds; None = no expiry.
|
||||
pub expires_at: Option<i64>,
|
||||
pub rpm_limit: Option<u32>,
|
||||
pub tpm_limit: Option<u32>,
|
||||
pub rate_state: Arc<RateLimitState>,
|
||||
}
|
||||
|
||||
/// Sliding window rate limit state per virtual key.
|
||||
#[derive(Debug)]
|
||||
pub struct RateLimitState {
|
||||
pub rpm_window: Mutex<VecDeque<u64>>,
|
||||
pub tpm_window: Mutex<VecDeque<(u64, u32)>>,
|
||||
}
|
||||
|
||||
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<String>,
|
||||
pub created_at: String,
|
||||
pub expires_at: Option<String>,
|
||||
pub revoked_at: Option<String>,
|
||||
pub rpm_limit: Option<u32>,
|
||||
pub tpm_limit: Option<u32>,
|
||||
pub spend_limit: Option<f64>,
|
||||
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());
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
@@ -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<String>) -> 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<serde_json::Value> {
|
||||
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<SharedState>) -> Json<serde_json::Val
|
||||
Json(serde_json::json!({ "backends": backends }))
|
||||
}
|
||||
|
||||
// --- Virtual API Key Management ---
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct CreateKeyRequest {
|
||||
description: Option<String>,
|
||||
expires_at: Option<String>,
|
||||
rpm_limit: Option<u32>,
|
||||
tpm_limit: Option<u32>,
|
||||
spend_limit: Option<f64>,
|
||||
}
|
||||
|
||||
/// POST /admin/api/keys -- create a new virtual API key.
|
||||
async fn create_key(
|
||||
State(shared): State<SharedState>,
|
||||
Json(body): Json<CreateKeyRequest>,
|
||||
) -> 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<SharedState>) -> 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<serde_json::Value> = 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<SharedState>,
|
||||
Path(id): Path<i64>,
|
||||
) -> 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::*;
|
||||
|
||||
@@ -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<tokio::sync::Mutex<()>>,
|
||||
/// 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<DashMap<[u8; 32], VirtualKeyMeta>>,
|
||||
}
|
||||
|
||||
/// 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()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -42,6 +42,10 @@ pub(crate) async fn send_with_retry<E: RetryableError>(
|
||||
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<AnthropicClientError> 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)),
|
||||
|
||||
@@ -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: <https://platform.openai.com/docs/api-reference/embeddings/create>
|
||||
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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<String>,
|
||||
region: Option<String>,
|
||||
// Azure-specific
|
||||
endpoint: Option<String>,
|
||||
deployment: Option<String>,
|
||||
api_version: Option<String>,
|
||||
// Optional env var name for Google access token (Vertex)
|
||||
access_token: Option<String>,
|
||||
// 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);
|
||||
}
|
||||
}
|
||||
|
||||
+333
-202
@@ -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 <path> takes priority; otherwise auto-load .anyllm.env if present.
|
||||
let args: Vec<String> = 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<dyn Fn(&str) -> 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<Mutex> 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<String, anyllm_proxy::metrics::Metrics> =
|
||||
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<dyn Fn(&str) -> 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<Mutex> 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<String, anyllm_proxy::metrics::Metrics> =
|
||||
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<tokio::task::JoinHandle<()>> =
|
||||
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.
|
||||
|
||||
@@ -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<AppState>,
|
||||
headers: axum::http::HeaderMap,
|
||||
permit: Option<axum::Extension<ConcurrencyPermit>>,
|
||||
body: Result<Json<openai::ChatCompletionRequest>, 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<ConcurrencyPermit>,
|
||||
) -> 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::<Result<String, std::convert::Infallible>>(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::<openai::ChatCompletionChunk>(
|
||||
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
|
||||
}
|
||||
@@ -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<Arc<DashMap<[u8; 32], VirtualKeyMeta>>> = OnceLock::new();
|
||||
|
||||
/// Initialize the global virtual keys reference. Called once from main.
|
||||
pub fn set_virtual_keys(keys: Arc<DashMap<[u8; 32], VirtualKeyMeta>>) {
|
||||
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<Vec<[u8; 32]>> = 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.
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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<GlobalState> {
|
||||
} 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<serde_json::Value> = 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<AppState>,
|
||||
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<AppState>,
|
||||
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) => {
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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}"
|
||||
);
|
||||
}
|
||||
@@ -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),
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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<anthropic::MessageCreateRequest, TranslateError> {
|
||||
// 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<anthropic::System> = 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<String> = 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<openai::ChatContent>) -> 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::<Vec<_>>()
|
||||
.join(""),
|
||||
None => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_openai_content_to_anthropic(
|
||||
content: &Option<openai::ChatContent>,
|
||||
) -> 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) })
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -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<u32>,
|
||||
output_tokens: Option<u32>,
|
||||
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<ChatCompletionChunk> {
|
||||
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<openai::FinishReason>,
|
||||
) -> 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);
|
||||
}
|
||||
}
|
||||
@@ -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<MessageCreateRequest, TranslateError> {
|
||||
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`.
|
||||
|
||||
@@ -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::*;
|
||||
|
||||
+102
@@ -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 <path>` 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
|
||||
```
|
||||
|
||||
@@ -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
|
||||
@@ -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` |
|
||||
@@ -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<String>,
|
||||
expires_at: Option<u64>, // epoch seconds
|
||||
rpm_limit: Option<u32>,
|
||||
tpm_limit: Option<u32>,
|
||||
spend_limit: Option<f64>, // USD
|
||||
rate_state: Arc<RateLimitState>,
|
||||
}
|
||||
```
|
||||
|
||||
### `RateLimitState`
|
||||
|
||||
```rust
|
||||
struct RateLimitState {
|
||||
rpm_window: Mutex<VecDeque<u64>>, // request timestamps (ms)
|
||||
tpm_window: Mutex<VecDeque<(u64, u32)>>, // (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<String>, // 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<u32>, // from message_start.usage
|
||||
output_tokens: Option<u32>, // 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<String>,
|
||||
api_key: Option<String>,
|
||||
timeout: Option<Duration>,
|
||||
read_timeout: Option<Duration>,
|
||||
max_retries: Option<u32>,
|
||||
tls_config: Option<TlsConfig>,
|
||||
}
|
||||
```
|
||||
|
||||
### `ToolBuilder`
|
||||
|
||||
```rust
|
||||
struct ToolBuilder {
|
||||
name: String,
|
||||
description: Option<String>,
|
||||
input_schema: serde_json::Value,
|
||||
}
|
||||
```
|
||||
|
||||
These are convenience wrappers over existing `Tool` and `ToolChoice` types in the translator crate.
|
||||
@@ -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<HashMap>` 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<HashMap>` 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.
|
||||
@@ -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,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
```
|
||||
@@ -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<RateLimitState>` stored inside `VirtualKeyMeta` in the DashMap. Sliding window with `Mutex<VecDeque>`.
|
||||
|
||||
### 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<VecDeque<u64>>, // timestamps
|
||||
tpm_window: Mutex<VecDeque<(u64, u32)>>, // (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<Item = Result<StreamEvent>>` 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`.
|
||||
@@ -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<Item = AnthropicStreamEvent>` 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)
|
||||
@@ -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<MessageCreateRequest, TranslateError>` 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<ChatCompletionChunk>` 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<ChatCompletionRequest>`, call `openai_to_anthropic_request`, dispatch to `BackendClient`, call `anthropic_to_openai_response`, return `Json<ChatCompletionResponse>`. 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<VirtualKeyRow>`, `revoke_virtual_key(conn, id) -> Option<VirtualKeyRow>`, `load_active_virtual_keys(conn) -> Vec<VirtualKeyRow>`.
|
||||
- [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<Client, ClientError>`. 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<impl Stream<Item = Result<StreamEvent, ClientError>>, 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<VecDeque<u64>>`, `tpm_window: Mutex<VecDeque<(u64, u32)>>`. 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<RateLimitState>` 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
|
||||
Reference in New Issue
Block a user