use async_trait::async_trait;
use std::time::Duration;
pub const DEFAULT_CLASSIFY_DEADLINE: Duration = Duration::from_secs(5);
#[derive(Clone)]
#[non_exhaustive]
pub struct ClassifyInput {
pub virtual_model: String,
pub path: String,
pub body: bytes::Bytes,
pub api_key: Option<String>,
}
impl std::fmt::Debug for ClassifyInput {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ClassifyInput")
.field("virtual_model", &self.virtual_model)
.field("path", &self.path)
.field("body_len", &self.body.len())
.field("api_key", &self.api_key.as_ref().map(|_| "<redacted>"))
.finish()
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct CacheStats {
pub read: u64,
pub creation_5m: u64,
pub creation_1h: u64,
pub creation_24h: u64,
}
impl CacheStats {
pub fn creation_total(&self) -> u64 {
self.creation_5m
.saturating_add(self.creation_1h)
.saturating_add(self.creation_24h)
}
pub fn is_zero(&self) -> bool {
*self == CacheStats::default()
}
}
#[async_trait]
pub trait CacheClassifier: Send + Sync {
async fn classify(&self, req: &ClassifyInput) -> CacheStats;
}
#[derive(Debug, Clone, Copy, Default)]
pub struct NoopCacheClassifier;
#[async_trait]
impl CacheClassifier for NoopCacheClassifier {
async fn classify(&self, _req: &ClassifyInput) -> CacheStats {
CacheStats::default()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cache_stats_default_is_all_zero() {
let s = CacheStats::default();
assert_eq!(s.read, 0);
assert_eq!(s.creation_5m, 0);
assert_eq!(s.creation_1h, 0);
assert_eq!(s.creation_24h, 0);
assert_eq!(s.creation_total(), 0);
assert!(s.is_zero());
}
#[test]
fn creation_total_sums_tiers() {
let s = CacheStats {
read: 100,
creation_5m: 1,
creation_1h: 2,
creation_24h: 3,
};
assert_eq!(s.creation_total(), 6);
assert!(!s.is_zero());
}
#[tokio::test]
async fn noop_classifier_returns_zero_stats() {
let classifier = NoopCacheClassifier;
let input = ClassifyInput {
virtual_model: "my-virtual-model".to_string(),
path: "/v1/chat/completions".to_string(),
body: bytes::Bytes::from_static(
br#"{"model":"my-virtual-model","messages":[{"role":"user","content":"hi"}]}"#,
),
api_key: Some("test-token".to_string()),
};
let stats = classifier.classify(&input).await;
assert_eq!(stats, CacheStats::default());
assert!(stats.is_zero());
}
#[test]
fn classify_input_debug_redacts_api_key() {
let secret = "totally-not-a-real-token";
let input = ClassifyInput {
virtual_model: "m".to_string(),
path: "/v1/chat/completions".to_string(),
body: bytes::Bytes::from_static(b"{}"),
api_key: Some(secret.to_string()),
};
let dbg = format!("{input:?}");
assert!(!dbg.contains(secret), "Debug leaked the api_key: {dbg}");
assert!(dbg.contains("<redacted>"), "Debug should redact: {dbg}");
let none_input = ClassifyInput {
api_key: None,
..input
};
assert!(format!("{none_input:?}").contains("api_key: None"));
}
}