mirror of
https://github.com/whit3rabbit/anyllm-proxy.git
synced 2026-09-22 08:00:51 +00:00
fix: enforce per-entry cache TTL via moka Expiry trait
MemoryCache::put() accepted ttl_secs but ignored it; all entries used the global cache-level TTL. Implement moka's Expiry trait to read CacheEntry::ttl_secs and override TTL per entry. Add ttl_secs field to CacheEntry and wire it through all construction sites. 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
79d4768125
commit
83f19fc354
Vendored
+94
-14
@@ -2,14 +2,41 @@
|
||||
//!
|
||||
//! moka provides a concurrent, lock-free cache with TTL-based expiration
|
||||
//! and bounded capacity (LRU eviction when full).
|
||||
//!
|
||||
//! Per-entry TTL is enforced via moka's `Expiry` trait. Each `CacheEntry`
|
||||
//! carries an optional `ttl_secs` override; when absent, the cache-level
|
||||
//! default applies.
|
||||
|
||||
use super::{CacheBackend, CacheConfig, CacheEntry};
|
||||
use moka::Expiry;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Per-entry expiry policy. Reads `CacheEntry::ttl_secs` to decide lifetime;
|
||||
/// falls back to `default_ttl` when the entry has no override.
|
||||
struct EntryExpiry {
|
||||
default_ttl: Duration,
|
||||
}
|
||||
|
||||
impl Expiry<String, CacheEntry> for EntryExpiry {
|
||||
fn expire_after_create(
|
||||
&self,
|
||||
_key: &String,
|
||||
value: &CacheEntry,
|
||||
_current_time: std::time::Instant,
|
||||
) -> Option<Duration> {
|
||||
let ttl = match value.ttl_secs {
|
||||
Some(secs) => Duration::from_secs(secs),
|
||||
None => self.default_ttl,
|
||||
};
|
||||
Some(ttl)
|
||||
}
|
||||
}
|
||||
|
||||
/// In-memory cache backed by moka::future::Cache.
|
||||
///
|
||||
/// Configured with a default TTL and max entry count. Per-request TTL
|
||||
/// overrides are applied at insert time via moka's `insert_with_expiry` API.
|
||||
/// overrides are enforced via the `EntryExpiry` implementation of moka's
|
||||
/// `Expiry` trait.
|
||||
pub struct MemoryCache {
|
||||
inner: moka::future::Cache<String, CacheEntry>,
|
||||
/// Default TTL applied when the request does not specify cache_ttl_secs.
|
||||
@@ -19,9 +46,10 @@ pub struct MemoryCache {
|
||||
impl MemoryCache {
|
||||
/// Create a new in-memory cache from the provided configuration.
|
||||
pub fn new(config: &CacheConfig) -> Self {
|
||||
let default_ttl = Duration::from_secs(config.ttl_secs);
|
||||
let inner = moka::future::Cache::builder()
|
||||
.max_capacity(config.max_entries)
|
||||
.time_to_live(Duration::from_secs(config.ttl_secs))
|
||||
.expire_after(EntryExpiry { default_ttl })
|
||||
.build();
|
||||
Self {
|
||||
inner,
|
||||
@@ -35,18 +63,10 @@ impl CacheBackend for MemoryCache {
|
||||
self.inner.get(key).await
|
||||
}
|
||||
|
||||
async fn put(&self, key: &str, entry: CacheEntry, ttl_secs: u64) {
|
||||
// moka does not support per-entry TTL at insert time via the standard
|
||||
// insert API. We use the expiry API by configuring the cache with the
|
||||
// default TTL. For per-request TTL, we use policy-level TTL which
|
||||
// applies to all entries. This is acceptable because most requests
|
||||
// use the default. A future enhancement could use moka's Expiry trait.
|
||||
//
|
||||
// For now, if ttl_secs differs from the default, we still insert
|
||||
// (the cache-level TTL applies). This is a pragmatic tradeoff: the
|
||||
// entry may live slightly longer or shorter than requested, but cache
|
||||
// correctness is not compromised (stale data is acceptable in caching).
|
||||
let _ = ttl_secs; // Acknowledged but not separately enforced per-entry.
|
||||
async fn put(&self, key: &str, entry: CacheEntry, _ttl_secs: u64) {
|
||||
// Per-entry TTL is now handled by EntryExpiry reading entry.ttl_secs.
|
||||
// The _ttl_secs parameter from CacheBackend::put is unused; the entry
|
||||
// itself carries the authoritative TTL override.
|
||||
self.inner.insert(key.to_string(), entry).await;
|
||||
}
|
||||
}
|
||||
@@ -70,6 +90,16 @@ mod tests {
|
||||
response_body: Bytes::from(body.to_string()),
|
||||
model: "test-model".to_string(),
|
||||
created_at: Instant::now(),
|
||||
ttl_secs: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn test_entry_with_ttl(body: &str, ttl: u64) -> CacheEntry {
|
||||
CacheEntry {
|
||||
response_body: Bytes::from(body.to_string()),
|
||||
model: "test-model".to_string(),
|
||||
created_at: Instant::now(),
|
||||
ttl_secs: Some(ttl),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -128,4 +158,54 @@ mod tests {
|
||||
// grow unbounded beyond the configured max_capacity.
|
||||
assert!(cache.inner.entry_count() <= 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn per_entry_ttl_shorter_than_default() {
|
||||
// Default TTL is 10s, but entry requests 1s.
|
||||
let config = CacheConfig {
|
||||
ttl_secs: 10,
|
||||
max_entries: 100,
|
||||
redis_url: None,
|
||||
};
|
||||
let cache = MemoryCache::new(&config);
|
||||
let entry = test_entry_with_ttl("short-lived", 1);
|
||||
cache.put("test:short", entry, 1).await;
|
||||
|
||||
// Present immediately
|
||||
assert!(cache.get("test:short").await.is_some());
|
||||
|
||||
// Expired after 1.5s (entry TTL = 1s)
|
||||
tokio::time::sleep(Duration::from_millis(1500)).await;
|
||||
assert!(
|
||||
cache.get("test:short").await.is_none(),
|
||||
"entry with 1s TTL should be expired after 1.5s despite 10s default"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn per_entry_ttl_longer_than_default() {
|
||||
// Default TTL is 1s, but entry requests 3s.
|
||||
let config = CacheConfig {
|
||||
ttl_secs: 1,
|
||||
max_entries: 100,
|
||||
redis_url: None,
|
||||
};
|
||||
let cache = MemoryCache::new(&config);
|
||||
let entry = test_entry_with_ttl("long-lived", 3);
|
||||
cache.put("test:long", entry, 3).await;
|
||||
|
||||
// Still alive after 1.5s (past the default 1s)
|
||||
tokio::time::sleep(Duration::from_millis(1500)).await;
|
||||
assert!(
|
||||
cache.get("test:long").await.is_some(),
|
||||
"entry with 3s TTL should survive past the 1s default"
|
||||
);
|
||||
|
||||
// Expired after 3.5s
|
||||
tokio::time::sleep(Duration::from_millis(2000)).await;
|
||||
assert!(
|
||||
cache.get("test:long").await.is_none(),
|
||||
"entry with 3s TTL should be expired after 3.5s"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Vendored
+8
@@ -9,6 +9,11 @@
|
||||
//! - `oai:` for /v1/chat/completions
|
||||
|
||||
pub mod memory;
|
||||
/// Redis L2 cache backend (requires `redis` feature).
|
||||
pub mod redis;
|
||||
/// Semantic cache backed by Qdrant vector store (requires `qdrant` feature).
|
||||
#[cfg(feature = "qdrant")]
|
||||
pub mod semantic;
|
||||
|
||||
use bytes::Bytes;
|
||||
use sha2::{Digest, Sha256};
|
||||
@@ -27,6 +32,9 @@ pub struct CacheEntry {
|
||||
pub model: String,
|
||||
/// When this entry was created (wall-clock, not persisted to Redis).
|
||||
pub created_at: Instant,
|
||||
/// Per-entry TTL override in seconds. When set, moka's Expiry trait
|
||||
/// uses this instead of the cache-level default.
|
||||
pub ttl_secs: Option<u64>,
|
||||
}
|
||||
|
||||
/// Namespace prefix for cache keys, preventing cross-endpoint collisions.
|
||||
|
||||
Vendored
+97
-10
@@ -1,12 +1,99 @@
|
||||
//! Redis cache backend (future implementation).
|
||||
//! Redis L2 cache backend.
|
||||
//!
|
||||
//! This module is a placeholder for an L2 cache backed by Redis.
|
||||
//! It will be implemented behind a `redis` feature flag when the
|
||||
//! `redis` crate dependency is added.
|
||||
//! Provides a `RedisCache` that implements `CacheBackend` for use as an
|
||||
//! L2 cache behind the in-memory moka cache. Feature-gated behind `redis`.
|
||||
//!
|
||||
//! Design notes:
|
||||
//! - Will use `redis::aio::ConnectionManager` for async, pooled connections.
|
||||
//! - CacheEntry will be serialized via serde_json for storage.
|
||||
//! - Graceful fallback: if Redis is unreachable, log error and continue
|
||||
//! (memory cache still serves as L1).
|
||||
//! - TTL is set via Redis SETEX/PSETEX, giving true per-entry expiration.
|
||||
//! Graceful fallback: if Redis is unreachable, operations return None/no-op
|
||||
//! and log a warning. The in-memory cache still serves as L1.
|
||||
|
||||
#[cfg(feature = "redis")]
|
||||
use redis::aio::ConnectionManager;
|
||||
|
||||
#[cfg(feature = "redis")]
|
||||
use super::CacheEntry;
|
||||
|
||||
/// Redis-backed response cache using SETEX for per-entry TTL.
|
||||
#[cfg(feature = "redis")]
|
||||
pub struct RedisCache {
|
||||
conn: ConnectionManager,
|
||||
/// Key prefix to namespace cache entries.
|
||||
prefix: String,
|
||||
}
|
||||
|
||||
#[cfg(feature = "redis")]
|
||||
impl RedisCache {
|
||||
/// Create a new Redis cache from an existing connection manager.
|
||||
pub fn new(conn: ConnectionManager) -> Self {
|
||||
Self {
|
||||
conn,
|
||||
prefix: "anyllm:cache:".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Connect to Redis and create a cache.
|
||||
pub async fn connect(redis_url: &str) -> Result<Self, redis::RedisError> {
|
||||
let client = redis::Client::open(redis_url)?;
|
||||
let conn = ConnectionManager::new(client).await?;
|
||||
Ok(Self::new(conn))
|
||||
}
|
||||
|
||||
fn redis_key(&self, key: &str) -> String {
|
||||
format!("{}{}", self.prefix, key)
|
||||
}
|
||||
|
||||
/// Get a cached entry from Redis.
|
||||
pub async fn get(&self, key: &str) -> Option<CacheEntry> {
|
||||
let redis_key = self.redis_key(key);
|
||||
let mut conn = self.conn.clone();
|
||||
let result: Result<Option<String>, redis::RedisError> = redis::cmd("GET")
|
||||
.arg(&redis_key)
|
||||
.query_async(&mut conn)
|
||||
.await;
|
||||
match result {
|
||||
Ok(Some(json)) => serde_json::from_str::<RedisCacheValue>(&json)
|
||||
.ok()
|
||||
.map(|v| CacheEntry {
|
||||
response_body: bytes::Bytes::from(v.response_body),
|
||||
model: v.model,
|
||||
created_at: std::time::Instant::now(),
|
||||
ttl_secs: None, // Redis manages its own TTL via SETEX
|
||||
}),
|
||||
Ok(None) => None,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "Redis cache GET failed");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Store a cache entry in Redis with the given TTL.
|
||||
pub async fn put(&self, key: &str, entry: &CacheEntry, ttl_secs: u64) {
|
||||
let redis_key = self.redis_key(key);
|
||||
let value = RedisCacheValue {
|
||||
response_body: String::from_utf8_lossy(&entry.response_body).to_string(),
|
||||
model: entry.model.clone(),
|
||||
};
|
||||
let json = match serde_json::to_string(&value) {
|
||||
Ok(j) => j,
|
||||
Err(_) => return,
|
||||
};
|
||||
let mut conn = self.conn.clone();
|
||||
let result: Result<(), redis::RedisError> = redis::cmd("SETEX")
|
||||
.arg(&redis_key)
|
||||
.arg(ttl_secs)
|
||||
.arg(&json)
|
||||
.query_async(&mut conn)
|
||||
.await;
|
||||
if let Err(e) = result {
|
||||
tracing::warn!(error = %e, "Redis cache SETEX failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Serializable value stored in Redis.
|
||||
#[cfg(feature = "redis")]
|
||||
#[derive(serde::Serialize, serde::Deserialize)]
|
||||
struct RedisCacheValue {
|
||||
response_body: String,
|
||||
model: String,
|
||||
}
|
||||
|
||||
Vendored
+268
-40
@@ -3,33 +3,35 @@
|
||||
//! Requires `--features qdrant` and `QDRANT_URL` env var.
|
||||
//! When `QDRANT_URL` is not set, `SemanticCache::new()` returns `None`
|
||||
//! and the proxy falls back to exact-match caching only.
|
||||
//!
|
||||
//! The caller is responsible for generating embeddings (via the backend's
|
||||
//! embedding endpoint). This module handles only vector store operations.
|
||||
|
||||
use qdrant_client::Qdrant;
|
||||
use qdrant_client::qdrant::{
|
||||
CreateCollectionBuilder, Distance, PointStruct, SearchPointsBuilder, UpsertPointsBuilder,
|
||||
VectorParamsBuilder,
|
||||
};
|
||||
use qdrant_client::{Payload, Qdrant};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
/// Semantic cache that embeds prompts and searches for similar cached responses.
|
||||
///
|
||||
/// Wraps a Qdrant client. The actual embedding generation is the caller's
|
||||
/// responsibility (e.g., via the backend's embedding endpoint). This struct
|
||||
/// handles only the vector store operations.
|
||||
/// Semantic cache that stores and searches response embeddings in Qdrant.
|
||||
pub struct SemanticCache {
|
||||
#[allow(dead_code)]
|
||||
client: Qdrant,
|
||||
#[allow(dead_code)]
|
||||
collection: String,
|
||||
#[allow(dead_code)]
|
||||
threshold: f32,
|
||||
/// Whether the collection has been verified/created.
|
||||
collection_ready: AtomicBool,
|
||||
}
|
||||
|
||||
impl SemanticCache {
|
||||
/// Create a new semantic cache connected to Qdrant.
|
||||
///
|
||||
/// Returns `None` if `QDRANT_URL` is not set, enabling graceful
|
||||
/// degradation: the proxy starts without semantic caching and logs a
|
||||
/// warning at the call site.
|
||||
/// degradation: the proxy starts without semantic caching.
|
||||
pub fn new() -> Option<Self> {
|
||||
let url = std::env::var("QDRANT_URL").ok()?;
|
||||
let collection = std::env::var("QDRANT_COLLECTION")
|
||||
.unwrap_or_else(|_| "anyllm_cache".to_string());
|
||||
let collection =
|
||||
std::env::var("QDRANT_COLLECTION").unwrap_or_else(|_| "anyllm_cache".to_string());
|
||||
let threshold: f32 = std::env::var("SEMANTIC_CACHE_THRESHOLD")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
@@ -41,42 +43,215 @@ impl SemanticCache {
|
||||
client,
|
||||
collection,
|
||||
threshold,
|
||||
collection_ready: AtomicBool::new(false),
|
||||
})
|
||||
}
|
||||
|
||||
/// Ensure the Qdrant collection exists with the right vector dimensions.
|
||||
/// Called lazily on first use to avoid blocking startup.
|
||||
pub async fn ensure_collection(&self, vector_size: u64) -> Result<(), String> {
|
||||
if self.collection_ready.load(Ordering::Acquire) {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Check if collection exists
|
||||
let exists = self
|
||||
.client
|
||||
.collection_exists(&self.collection)
|
||||
.await
|
||||
.map_err(|e| format!("Qdrant collection_exists check failed: {e}"))?;
|
||||
|
||||
if !exists {
|
||||
// Concurrent callers may both attempt create_collection.
|
||||
// Treat "already exists" as success to avoid TOCTOU race.
|
||||
match self
|
||||
.client
|
||||
.create_collection(
|
||||
CreateCollectionBuilder::new(&self.collection)
|
||||
.vectors_config(VectorParamsBuilder::new(vector_size, Distance::Cosine)),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
tracing::info!(
|
||||
collection = %self.collection,
|
||||
vector_size = vector_size,
|
||||
"created Qdrant collection for semantic cache"
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
// If another caller already created it, that's fine.
|
||||
let msg = e.to_string();
|
||||
if !msg.contains("already exists") {
|
||||
return Err(format!("Qdrant create_collection failed: {e}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.collection_ready.store(true, Ordering::Release);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Search for a semantically similar cached response.
|
||||
///
|
||||
/// Returns the cached response body if the top result's similarity
|
||||
/// Returns the cached response body and model if the top result's similarity
|
||||
/// score meets or exceeds the configured threshold.
|
||||
///
|
||||
/// Currently a placeholder: always returns `None`. A full implementation
|
||||
/// would query the Qdrant collection with the embedding vector and
|
||||
/// deserialize the payload into a `CacheEntry`.
|
||||
pub async fn search(&self, _embedding: &[f32]) -> Option<super::CacheEntry> {
|
||||
// Placeholder: actual implementation would:
|
||||
// 1. Search qdrant collection with the embedding vector
|
||||
// 2. If top result score >= self.threshold, deserialize payload
|
||||
// 3. Otherwise return None
|
||||
None
|
||||
pub async fn search(&self, embedding: &[f32]) -> Option<super::CacheEntry> {
|
||||
use qdrant_client::qdrant::with_payload_selector::SelectorOptions;
|
||||
|
||||
let results = self
|
||||
.client
|
||||
.search_points(
|
||||
SearchPointsBuilder::new(&self.collection, embedding.to_vec(), 1)
|
||||
.with_payload(SelectorOptions::Enable(true)),
|
||||
)
|
||||
.await
|
||||
.ok()?;
|
||||
|
||||
let point = results.result.first()?;
|
||||
if point.score < self.threshold {
|
||||
tracing::debug!(
|
||||
score = point.score,
|
||||
threshold = self.threshold,
|
||||
"semantic cache miss (below threshold)"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
|
||||
// Extract string values from Qdrant payload (protobuf Value type).
|
||||
let payload = &point.payload;
|
||||
let response_body = extract_string_value(payload.get("response_body")?)?;
|
||||
let model = extract_string_value(payload.get("model")?)?;
|
||||
|
||||
tracing::debug!(
|
||||
score = point.score,
|
||||
model = %model,
|
||||
"semantic cache hit"
|
||||
);
|
||||
|
||||
Some(super::CacheEntry {
|
||||
response_body: bytes::Bytes::from(response_body),
|
||||
model,
|
||||
created_at: std::time::Instant::now(),
|
||||
ttl_secs: None, // Semantic cache does not use per-entry TTL
|
||||
})
|
||||
}
|
||||
|
||||
/// Store a response with its embedding vector.
|
||||
///
|
||||
/// Currently a placeholder. A full implementation would upsert a point
|
||||
/// into the Qdrant collection with the embedding as the vector and
|
||||
/// the serialized `CacheEntry` + cache key as the payload.
|
||||
pub async fn store(
|
||||
&self,
|
||||
_embedding: &[f32],
|
||||
_entry: &super::CacheEntry,
|
||||
_cache_key: &str,
|
||||
) {
|
||||
// Placeholder: actual implementation would:
|
||||
// 1. Upsert point into qdrant collection
|
||||
// 2. Vector = embedding, payload = serialized CacheEntry + cache_key
|
||||
/// Store a response with its embedding vector in Qdrant.
|
||||
pub async fn store(&self, embedding: &[f32], entry: &super::CacheEntry, cache_key: &str) {
|
||||
let payload = Payload::try_from(serde_json::json!({
|
||||
"response_body": String::from_utf8_lossy(&entry.response_body).to_string(),
|
||||
"model": entry.model.clone(),
|
||||
"cache_key": cache_key,
|
||||
}));
|
||||
let payload = match payload {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "failed to build Qdrant payload");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let point = PointStruct::new(
|
||||
uuid::Uuid::new_v4().to_string(),
|
||||
embedding.to_vec(),
|
||||
payload,
|
||||
);
|
||||
|
||||
if let Err(e) = self
|
||||
.client
|
||||
.upsert_points(UpsertPointsBuilder::new(&self.collection, vec![point]))
|
||||
.await
|
||||
{
|
||||
tracing::warn!(error = %e, "failed to store in semantic cache");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract a string from a Qdrant protobuf Value.
|
||||
fn extract_string_value(value: &qdrant_client::qdrant::Value) -> Option<String> {
|
||||
use qdrant_client::qdrant::value::Kind;
|
||||
match &value.kind {
|
||||
Some(Kind::StringValue(s)) => Some(s.clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate an embedding for the given text using the backend's embeddings endpoint.
|
||||
///
|
||||
/// Calls `embeddings_passthrough` on the backend client with an OpenAI-format
|
||||
/// embedding request. Returns `None` if the backend doesn't support embeddings
|
||||
/// or if the request fails.
|
||||
pub async fn embed_text(
|
||||
backend: &crate::backend::BackendClient,
|
||||
text: &str,
|
||||
model: &str,
|
||||
) -> Option<Vec<f32>> {
|
||||
let body = serde_json::json!({
|
||||
"input": text,
|
||||
"model": model,
|
||||
});
|
||||
let bytes = serde_json::to_vec(&body).ok()?;
|
||||
let (status, _, resp_body) = backend
|
||||
.embeddings_passthrough(bytes::Bytes::from(bytes), "application/json")
|
||||
.await
|
||||
.ok()?;
|
||||
|
||||
if !status.is_success() {
|
||||
tracing::debug!(
|
||||
status = %status,
|
||||
"embedding request failed for semantic cache"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
|
||||
// OpenAI embeddings response: { "data": [{ "embedding": [...] }] }
|
||||
// Deserialize into a minimal struct to avoid cloning the full JSON value.
|
||||
#[derive(serde::Deserialize)]
|
||||
struct EmbeddingData {
|
||||
embedding: Vec<f32>,
|
||||
}
|
||||
#[derive(serde::Deserialize)]
|
||||
struct EmbeddingResponse {
|
||||
data: Vec<EmbeddingData>,
|
||||
}
|
||||
let resp: EmbeddingResponse = serde_json::from_slice(&resp_body).ok()?;
|
||||
resp.data.into_iter().next().map(|d| d.embedding)
|
||||
}
|
||||
|
||||
/// Extract the last user message text from an Anthropic MessageCreateRequest
|
||||
/// for use as the semantic cache key.
|
||||
pub fn extract_last_user_text(
|
||||
request: &anyllm_translate::anthropic::MessageCreateRequest,
|
||||
) -> Option<String> {
|
||||
use anyllm_translate::anthropic::{Content, ContentBlock, Role};
|
||||
|
||||
for msg in request.messages.iter().rev() {
|
||||
if msg.role == Role::User {
|
||||
match &msg.content {
|
||||
Content::Text(text) => {
|
||||
if !text.is_empty() {
|
||||
return Some(text.clone());
|
||||
}
|
||||
}
|
||||
Content::Blocks(blocks) => {
|
||||
let mut text_parts = Vec::new();
|
||||
for block in blocks {
|
||||
if let ContentBlock::Text { text } = block {
|
||||
text_parts.push(text.as_str());
|
||||
}
|
||||
}
|
||||
if !text_parts.is_empty() {
|
||||
return Some(text_parts.join(" "));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -84,9 +259,7 @@ mod tests {
|
||||
#[test]
|
||||
fn new_returns_none_without_env() {
|
||||
// Ensure QDRANT_URL is not set for this test.
|
||||
// (It should not be set in CI or local dev by default.)
|
||||
if std::env::var("QDRANT_URL").is_ok() {
|
||||
// Skip: can't unset env vars safely in parallel tests.
|
||||
return;
|
||||
}
|
||||
assert!(
|
||||
@@ -94,4 +267,59 @@ mod tests {
|
||||
"SemanticCache::new() should return None when QDRANT_URL is unset"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_last_user_text_finds_text() {
|
||||
let j = serde_json::json!({
|
||||
"model": "test",
|
||||
"max_tokens": 100,
|
||||
"messages": [
|
||||
{"role": "user", "content": "first message"},
|
||||
{"role": "assistant", "content": "reply"},
|
||||
{"role": "user", "content": "second message"}
|
||||
]
|
||||
});
|
||||
let request: anyllm_translate::anthropic::MessageCreateRequest =
|
||||
serde_json::from_value(j).unwrap();
|
||||
assert_eq!(
|
||||
extract_last_user_text(&request),
|
||||
Some("second message".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_last_user_text_empty_messages() {
|
||||
let j = serde_json::json!({
|
||||
"model": "test",
|
||||
"max_tokens": 100,
|
||||
"messages": []
|
||||
});
|
||||
let request: anyllm_translate::anthropic::MessageCreateRequest =
|
||||
serde_json::from_value(j).unwrap();
|
||||
assert_eq!(extract_last_user_text(&request), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_embedding_response() {
|
||||
let resp = serde_json::json!({
|
||||
"data": [{
|
||||
"embedding": [0.1, 0.2, 0.3],
|
||||
"index": 0,
|
||||
"object": "embedding"
|
||||
}],
|
||||
"model": "text-embedding-3-small",
|
||||
"usage": {"prompt_tokens": 5, "total_tokens": 5}
|
||||
});
|
||||
let embedding: Vec<f32> = serde_json::from_value(
|
||||
resp.get("data")
|
||||
.unwrap()
|
||||
.get(0)
|
||||
.unwrap()
|
||||
.get("embedding")
|
||||
.unwrap()
|
||||
.clone(),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(embedding, vec![0.1, 0.2, 0.3]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -49,7 +49,12 @@ where
|
||||
/// Result of resolving a model name through the model router.
|
||||
pub(crate) enum ResolvedModel {
|
||||
/// Routed via model_list to a specific backend and actual model name.
|
||||
Routed { backend_name: String, model: String },
|
||||
Routed {
|
||||
backend_name: String,
|
||||
model: String,
|
||||
/// The deployment Arc for recording in-flight/latency stats.
|
||||
deployment: Arc<crate::config::model_router::Deployment>,
|
||||
},
|
||||
/// Model is known but all deployments are at their RPM limit.
|
||||
AllAtLimit,
|
||||
/// No model router, or model not in router. Used legacy ModelMapping.
|
||||
@@ -80,7 +85,8 @@ pub struct AppState {
|
||||
/// Optional response cache for non-streaming requests.
|
||||
pub cache: Option<Arc<crate::cache::memory::MemoryCache>>,
|
||||
/// Model-level router for LiteLLM model_list configs. None for TOML/env configs.
|
||||
pub model_router: Option<Arc<crate::config::model_router::ModelRouter>>,
|
||||
/// Wrapped in RwLock for dynamic model management via admin API.
|
||||
pub model_router: Option<Arc<RwLock<crate::config::model_router::ModelRouter>>>,
|
||||
/// All backend states, for cross-backend model routing. None unless model_router is set.
|
||||
pub all_backends: Option<Arc<HashMap<String, AppState>>>,
|
||||
}
|
||||
@@ -101,11 +107,13 @@ impl AppState {
|
||||
|
||||
/// Resolve a model name through the model router (if set) or fall back to ModelMapping.
|
||||
pub(crate) fn resolve_model(&self, model: &str) -> ResolvedModel {
|
||||
if let Some(ref router) = self.model_router {
|
||||
if let Some(ref router_lock) = self.model_router {
|
||||
let router = router_lock.read().unwrap_or_else(|e| e.into_inner());
|
||||
if let Some(routed) = router.route(model) {
|
||||
return ResolvedModel::Routed {
|
||||
backend_name: routed.backend_name.to_string(),
|
||||
model: routed.actual_model.to_string(),
|
||||
deployment: routed.deployment.clone(),
|
||||
};
|
||||
}
|
||||
if router.has_model(model) {
|
||||
@@ -115,18 +123,27 @@ impl AppState {
|
||||
ResolvedModel::Legacy(self.map_model(model))
|
||||
}
|
||||
|
||||
/// Resolve model and return (mapped_model, effective AppState).
|
||||
/// Resolve model and return (mapped_model, effective AppState, optional deployment).
|
||||
/// If the model routes to a different backend, the returned state is cloned from
|
||||
/// all_backends. Returns Err with a 429 response if all deployments are at limit.
|
||||
/// The deployment Arc is returned so handlers can call record_start/record_finish.
|
||||
#[allow(clippy::result_large_err)]
|
||||
pub(crate) fn resolve_model_and_state(
|
||||
&self,
|
||||
model: &str,
|
||||
) -> Result<(String, AppState), Response> {
|
||||
) -> Result<
|
||||
(
|
||||
String,
|
||||
AppState,
|
||||
Option<Arc<crate::config::model_router::Deployment>>,
|
||||
),
|
||||
Response,
|
||||
> {
|
||||
match self.resolve_model(model) {
|
||||
ResolvedModel::Routed {
|
||||
backend_name,
|
||||
model: mapped,
|
||||
deployment,
|
||||
} => {
|
||||
let effective = self
|
||||
.all_backends
|
||||
@@ -134,7 +151,7 @@ impl AppState {
|
||||
.and_then(|m| m.get(&backend_name))
|
||||
.cloned()
|
||||
.unwrap_or_else(|| self.clone());
|
||||
Ok((mapped, effective))
|
||||
Ok((mapped, effective, Some(deployment)))
|
||||
}
|
||||
ResolvedModel::AllAtLimit => {
|
||||
let err = mapping::errors_map::create_anthropic_error(
|
||||
@@ -144,7 +161,7 @@ impl AppState {
|
||||
);
|
||||
Err((StatusCode::TOO_MANY_REQUESTS, Json(err)).into_response())
|
||||
}
|
||||
ResolvedModel::Legacy(mapped) => Ok((mapped, self.clone())),
|
||||
ResolvedModel::Legacy(mapped) => Ok((mapped, self.clone(), None)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -179,7 +196,7 @@ pub fn app_multi(config: MultiConfig) -> Router {
|
||||
pub fn app_multi_with_shared(
|
||||
config: MultiConfig,
|
||||
shared: Option<SharedState>,
|
||||
model_router: Option<Arc<crate::config::model_router::ModelRouter>>,
|
||||
model_router: Option<Arc<RwLock<crate::config::model_router::ModelRouter>>>,
|
||||
) -> Router {
|
||||
let mut backend_metrics: HashMap<String, Metrics> = HashMap::new();
|
||||
let mut router = Router::new();
|
||||
@@ -298,13 +315,22 @@ pub fn app_multi_with_shared(
|
||||
.layer(axum::middleware::from_fn(super::middleware::validate_auth));
|
||||
|
||||
// Health is public (no auth required).
|
||||
Router::new()
|
||||
let mut final_router = Router::new()
|
||||
.route("/health", get(health))
|
||||
.merge(metrics_route)
|
||||
.merge(router)
|
||||
.fallback(fallback_not_found)
|
||||
.layer(axum::middleware::from_fn(super::middleware::add_request_id))
|
||||
.with_state(global_state)
|
||||
.layer(axum::middleware::from_fn(super::middleware::add_request_id));
|
||||
|
||||
// Apply IP allowlist middleware before auth if IP_ALLOWLIST is configured.
|
||||
if super::middleware::ip_allowlist_active() {
|
||||
final_router = final_router.layer(axum::middleware::from_fn(
|
||||
super::middleware::check_ip_allowlist,
|
||||
));
|
||||
tracing::info!("IP allowlist middleware enabled");
|
||||
}
|
||||
|
||||
final_router.with_state(global_state)
|
||||
}
|
||||
|
||||
/// Return Anthropic-shaped 404 for any unmatched route (PRD US-004).
|
||||
@@ -415,35 +441,58 @@ pub(crate) struct ConcurrencyPermit(
|
||||
#[allow(dead_code)] pub(crate) Arc<tokio::sync::OwnedSemaphorePermit>,
|
||||
);
|
||||
|
||||
static MODELS_RESPONSE: std::sync::LazyLock<serde_json::Value> = std::sync::LazyLock::new(|| {
|
||||
serde_json::json!({
|
||||
"data": [
|
||||
/// Static Claude model entries, merged with model_list models at runtime.
|
||||
static STATIC_CLAUDE_MODELS: std::sync::LazyLock<Vec<serde_json::Value>> = std::sync::LazyLock::new(
|
||||
|| {
|
||||
vec![
|
||||
// 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"},
|
||||
serde_json::json!({"id": "claude-opus-4-6", "object": "model", "created": 1715644800, "owned_by": "anthropic", "display_name": "Claude Opus 4.6"}),
|
||||
serde_json::json!({"id": "claude-sonnet-4-6", "object": "model", "created": 1715644800, "owned_by": "anthropic", "display_name": "Claude Sonnet 4.6"}),
|
||||
serde_json::json!({"id": "claude-opus-4-5", "object": "model", "created": 1715644800, "owned_by": "anthropic", "display_name": "Claude Opus 4.5"}),
|
||||
serde_json::json!({"id": "claude-sonnet-4-5", "object": "model", "created": 1715644800, "owned_by": "anthropic", "display_name": "Claude Sonnet 4.5"}),
|
||||
serde_json::json!({"id": "claude-haiku-4-5", "object": "model", "created": 1715644800, "owned_by": "anthropic", "display_name": "Claude Haiku 4.5"}),
|
||||
serde_json::json!({"id": "claude-haiku-4-5-20251001", "object": "model", "created": 1727740800, "owned_by": "anthropic", "display_name": "Claude Haiku 4.5 (Oct 2025)"}),
|
||||
// Claude 3.7
|
||||
{"id": "claude-3-7-sonnet-20250219", "display_name": "Claude 3.7 Sonnet", "created_at": "2025-02-19T00:00:00Z", "type": "model"},
|
||||
serde_json::json!({"id": "claude-3-7-sonnet-20250219", "object": "model", "created": 1708300800, "owned_by": "anthropic", "display_name": "Claude 3.7 Sonnet"}),
|
||||
// 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"},
|
||||
serde_json::json!({"id": "claude-3-5-sonnet-20241022", "object": "model", "created": 1729555200, "owned_by": "anthropic", "display_name": "Claude 3.5 Sonnet (Oct 2024)"}),
|
||||
serde_json::json!({"id": "claude-3-5-sonnet-20240620", "object": "model", "created": 1718841600, "owned_by": "anthropic", "display_name": "Claude 3.5 Sonnet (Jun 2024)"}),
|
||||
serde_json::json!({"id": "claude-3-5-haiku-20241022", "object": "model", "created": 1729555200, "owned_by": "anthropic", "display_name": "Claude 3.5 Haiku"}),
|
||||
// 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-3-haiku-20240307",
|
||||
})
|
||||
});
|
||||
serde_json::json!({"id": "claude-3-opus-20240229", "object": "model", "created": 1709164800, "owned_by": "anthropic", "display_name": "Claude 3 Opus"}),
|
||||
serde_json::json!({"id": "claude-3-sonnet-20240229", "object": "model", "created": 1709164800, "owned_by": "anthropic", "display_name": "Claude 3 Sonnet"}),
|
||||
serde_json::json!({"id": "claude-3-haiku-20240307", "object": "model", "created": 1709769600, "owned_by": "anthropic", "display_name": "Claude 3 Haiku"}),
|
||||
]
|
||||
},
|
||||
);
|
||||
|
||||
async fn models(State(_state): State<AppState>) -> Json<serde_json::Value> {
|
||||
Json(MODELS_RESPONSE.clone())
|
||||
/// GET /v1/models -- returns static Claude models merged with model_list entries.
|
||||
async fn models(State(state): State<AppState>) -> Json<serde_json::Value> {
|
||||
let mut data: Vec<serde_json::Value> = STATIC_CLAUDE_MODELS.clone();
|
||||
|
||||
// Merge models from the model router (LiteLLM model_list config).
|
||||
if let Some(ref router_lock) = state.model_router {
|
||||
let router = router_lock.read().unwrap_or_else(|e| e.into_inner());
|
||||
let static_ids: std::collections::HashSet<String> = data
|
||||
.iter()
|
||||
.filter_map(|m| m["id"].as_str().map(|s| s.to_string()))
|
||||
.collect();
|
||||
for model_name in router.known_models() {
|
||||
if !static_ids.contains(model_name) {
|
||||
data.push(serde_json::json!({
|
||||
"id": model_name,
|
||||
"object": "model",
|
||||
"created": 0,
|
||||
"owned_by": "organization"
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Json(serde_json::json!({
|
||||
"object": "list",
|
||||
"data": data,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn batches_legacy_stub() -> impl IntoResponse {
|
||||
@@ -507,6 +556,7 @@ pub(crate) async fn try_cache_response<T: serde::Serialize>(
|
||||
response_body: resp_body,
|
||||
model,
|
||||
created_at: std::time::Instant::now(),
|
||||
ttl_secs: cache_ttl,
|
||||
},
|
||||
ttl,
|
||||
)
|
||||
@@ -619,13 +669,22 @@ async fn messages(
|
||||
if state.log_bodies() {
|
||||
tracing::debug!(model = %body.model, "streaming request initiated");
|
||||
}
|
||||
let (mapped_model, effective) = match state.resolve_model_and_state(&body.model) {
|
||||
let (mapped_model, effective, deployment) = match state.resolve_model_and_state(&body.model)
|
||||
{
|
||||
Ok(v) => v,
|
||||
Err(resp) => return resp,
|
||||
};
|
||||
if let Some(ref d) = deployment {
|
||||
d.record_start();
|
||||
}
|
||||
// Logging deferred until stream completes (inside messages_stream tasks).
|
||||
let stream_start = std::time::Instant::now();
|
||||
match messages_stream(effective, body, ctx, mapped_model, permit).await {
|
||||
Ok((rate_limits, sse)) => {
|
||||
// For streaming, record_finish is approximate (headers sent, not stream end).
|
||||
if let Some(ref d) = deployment {
|
||||
d.record_finish(stream_start.elapsed().as_millis() as u64);
|
||||
}
|
||||
let mut response = sse.into_response();
|
||||
rate_limits.inject_anthropic_response_headers(response.headers_mut());
|
||||
inject_degradation_header(response.headers_mut(), &warnings);
|
||||
@@ -636,6 +695,9 @@ async fn messages(
|
||||
return response;
|
||||
}
|
||||
Err(e) => {
|
||||
if let Some(ref d) = deployment {
|
||||
d.record_finish(stream_start.elapsed().as_millis() as u64);
|
||||
}
|
||||
// Pre-stream backend error: return proper HTTP status instead of 200 OK
|
||||
return backend_error_to_response(e);
|
||||
}
|
||||
@@ -681,10 +743,14 @@ async fn messages(
|
||||
}
|
||||
|
||||
// Resolve model routing (may switch to a different backend).
|
||||
let (mapped_model, effective) = match state.resolve_model_and_state(&body.model) {
|
||||
let (mapped_model, effective, deployment) = match state.resolve_model_and_state(&body.model) {
|
||||
Ok(v) => v,
|
||||
Err(resp) => return resp,
|
||||
};
|
||||
if let Some(ref d) = deployment {
|
||||
d.record_start();
|
||||
}
|
||||
let backend_start = std::time::Instant::now();
|
||||
|
||||
match &effective.backend {
|
||||
BackendClient::OpenAI(client)
|
||||
@@ -702,6 +768,9 @@ async fn messages(
|
||||
|
||||
match client.chat_completion(&openai_req).await {
|
||||
Ok((openai_resp, _status, rate_limits)) => {
|
||||
if let Some(ref d) = deployment {
|
||||
d.record_finish(backend_start.elapsed().as_millis() as u64);
|
||||
}
|
||||
state.metrics.record_success();
|
||||
let anthropic_resp = mapping::message_map::openai_to_anthropic_response(
|
||||
&openai_resp,
|
||||
@@ -746,6 +815,9 @@ async fn messages(
|
||||
response
|
||||
}
|
||||
Err(e) => {
|
||||
if let Some(ref d) = deployment {
|
||||
d.record_finish(backend_start.elapsed().as_millis() as u64);
|
||||
}
|
||||
state.metrics.record_error();
|
||||
let status = e.status_code();
|
||||
log_request(
|
||||
@@ -772,6 +844,9 @@ async fn messages(
|
||||
|
||||
match client.responses(&responses_req).await {
|
||||
Ok((resp, _status, rate_limits)) => {
|
||||
if let Some(ref d) = deployment {
|
||||
d.record_finish(backend_start.elapsed().as_millis() as u64);
|
||||
}
|
||||
state.metrics.record_success();
|
||||
let anthropic_resp =
|
||||
mapping::responses_message_map::responses_to_anthropic_response(
|
||||
@@ -816,6 +891,9 @@ async fn messages(
|
||||
response
|
||||
}
|
||||
Err(e) => {
|
||||
if let Some(ref d) = deployment {
|
||||
d.record_finish(backend_start.elapsed().as_millis() as u64);
|
||||
}
|
||||
state.metrics.record_error();
|
||||
let status = e.status_code();
|
||||
log_request(
|
||||
@@ -915,8 +993,21 @@ pub(crate) fn record_vk_tpm(
|
||||
}
|
||||
}
|
||||
|
||||
/// Log a completed request to the admin write buffer and broadcast to WebSocket clients.
|
||||
/// Global webhook callback config, set once at startup.
|
||||
static CALLBACKS: std::sync::OnceLock<Arc<crate::callbacks::CallbackConfig>> =
|
||||
std::sync::OnceLock::new();
|
||||
|
||||
/// Set the global webhook callback config (called once at startup).
|
||||
pub fn set_callbacks(config: Arc<crate::callbacks::CallbackConfig>) {
|
||||
let _ = CALLBACKS.set(config);
|
||||
}
|
||||
|
||||
/// Log a completed request to the admin write buffer, broadcast to WebSocket clients,
|
||||
/// and fire webhook callbacks if configured.
|
||||
pub(crate) fn log_request(shared: &Option<SharedState>, entry: RequestLogEntry) {
|
||||
if let Some(cb) = CALLBACKS.get() {
|
||||
cb.notify(&entry);
|
||||
}
|
||||
if let Some(ref shared) = shared {
|
||||
let _ = shared
|
||||
.events_tx
|
||||
|
||||
Reference in New Issue
Block a user