diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index ff3ff6572d4..72fb338e7d8 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -6,4 +6,8 @@ pub enum Error { InvalidEntry, #[error("flushing Redis requires an explicit namespace")] UnscopedFlush, + #[error("cache backend does not support this operation")] + UnsupportedOperation, + #[error("semantic cache requires request messages")] + MissingPrompt, } diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index ce9f93b6dc4..b9c720fa3ee 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -5,6 +5,7 @@ mod capabilities; mod codec; mod dual; mod error; +mod semantic; pub use base_cache::{ BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext, @@ -19,3 +20,4 @@ pub use capabilities::{ pub use codec::{CacheCodec, JsonCodec}; pub use dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy}; pub use error::Error; +pub use semantic::{SemanticCacheContext, SemanticCacheScope}; diff --git a/litellm-rust/crates/cache/src/semantic.rs b/litellm-rust/crates/cache/src/semantic.rs new file mode 100644 index 00000000000..61f9023fa4c --- /dev/null +++ b/litellm-rust/crates/cache/src/semantic.rs @@ -0,0 +1,72 @@ +use std::time::Duration; + +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +use crate::CacheContext; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum SemanticCacheScope { + #[default] + Key, + EndUser, +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct SemanticCacheContext { + pub input: Option, + pub messages: Vec, + pub metadata: Map, + pub scope: SemanticCacheScope, + pub ttl: Option, +} + +impl CacheContext for SemanticCacheContext { + fn ttl(&self) -> Option { + self.ttl + } + + fn with_ttl(&self, ttl: Option) -> Self { + Self { + ttl, + ..self.clone() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn with_ttl_keeps_request_fields() { + let context = SemanticCacheContext { + input: Some("query".to_owned()), + messages: vec![serde_json::json!({"role": "user", "content": "hi"})], + metadata: Map::from_iter([("user".to_owned(), Value::from("u1"))]), + scope: SemanticCacheScope::EndUser, + ttl: None, + }; + + let updated = context.with_ttl(Some(Duration::from_secs(5))); + + assert_eq!(updated.ttl(), Some(Duration::from_secs(5))); + assert_eq!(updated.input, context.input); + assert_eq!(updated.messages, context.messages); + assert_eq!(updated.metadata, context.metadata); + assert_eq!(updated.scope, SemanticCacheScope::EndUser); + } + + #[test] + fn scope_serializes_like_python_cache_scope() { + assert_eq!( + serde_json::to_value(SemanticCacheScope::EndUser).unwrap(), + Value::from("end_user") + ); + assert_eq!( + serde_json::from_value::(Value::from("key")).unwrap(), + SemanticCacheScope::Key + ); + } +}