liter_llm 1.0.0.pre.rc.6

data/vendor/liter-llm/src/tower/cache_opendal.rs

@@ -0,0 +1,270 @@
+//! OpenDAL-backed cache store for the response cache.
+//!
+//! Implements [`CacheStore`] using an [`opendal::Operator`] for persistence.
+//! Supports any OpenDAL backend (S3, Redis, GCS, local filesystem, etc.).
+
+use std::collections::HashMap;
+use std::future::Future;
+use std::pin::Pin;
+use std::str::FromStr;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+use opendal::Operator;
+use serde::{Deserialize, Serialize};
+
+use super::cache::{CacheStore, CachedResponse};
+
+/// A cached entry stored via OpenDAL, including metadata for TTL and
+/// collision detection.
+#[derive(Serialize, Deserialize)]
+struct StoredEntry {
+    request_body: String,
+    response: CachedResponse,
+    /// Unix timestamp (seconds) when this entry expires.
+    expires_at: u64,
+}
+
+/// Cache store backed by an [`opendal::Operator`].
+///
+/// Entries are stored as JSON files under `{prefix}/{key}`. TTL is embedded
+/// in the stored entry and checked on read. Backend failures are non-fatal:
+/// they log a warning and behave as a cache miss / no-op.
+pub struct OpenDalCacheStore {
+    operator: Operator,
+    prefix: String,
+    ttl: Duration,
+}
+
+impl OpenDalCacheStore {
+    /// Create a new OpenDAL cache store.
+    ///
+    /// `operator` must be a fully configured OpenDAL operator.
+    /// `prefix` is prepended to all cache keys (e.g. `"llm-cache/"`).
+    /// `ttl` controls how long entries are valid.
+    pub fn new(operator: Operator, prefix: impl Into<String>, ttl: Duration) -> Self {
+        Self {
+            operator,
+            prefix: prefix.into(),
+            ttl,
+        }
+    }
+
+    /// Build an OpenDAL operator from a scheme name and config map.
+    ///
+    /// # Errors
+    /// Returns an error if the scheme is unknown or the config is invalid.
+    pub fn from_config(
+        scheme: &str,
+        config: HashMap<String, String>,
+        prefix: impl Into<String>,
+        ttl: Duration,
+    ) -> crate::error::Result<Self> {
+        let parsed_scheme =
+            opendal::Scheme::from_str(scheme).map_err(|e| crate::error::LiterLlmError::InternalError {
+                message: format!("unknown OpenDAL scheme '{scheme}': {e}"),
+            })?;
+        let operator =
+            Operator::via_iter(parsed_scheme, config).map_err(|e| crate::error::LiterLlmError::InternalError {
+                message: format!("failed to build OpenDAL operator for '{scheme}': {e}"),
+            })?;
+        Ok(Self::new(operator, prefix, ttl))
+    }
+
+    fn key_path(&self, key: u64) -> String {
+        format!("{}{key}", self.prefix)
+    }
+
+    fn now_secs() -> u64 {
+        SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap_or_default()
+            .as_secs()
+    }
+}
+
+impl CacheStore for OpenDalCacheStore {
+    fn get(&self, key: u64, request_body: &str) -> Pin<Box<dyn Future<Output = Option<CachedResponse>> + Send + '_>> {
+        let path = self.key_path(key);
+        let request_body = request_body.to_owned();
+        Box::pin(async move {
+            let bytes = match self.operator.read(&path).await {
+                Ok(b) => b,
+                Err(_) => return None,
+            };
+            let entry: StoredEntry = match serde_json::from_slice(bytes.to_bytes().as_ref()) {
+                Ok(e) => e,
+                Err(_) => return None,
+            };
+            // Check TTL
+            if Self::now_secs() > entry.expires_at {
+                // Lazily delete expired entry
+                let _ = self.operator.delete(&path).await;
+                return None;
+            }
+            // Verify request body matches (collision guard)
+            if entry.request_body != request_body {
+                return None;
+            }
+            Some(entry.response)
+        })
+    }
+
+    fn put(
+        &self,
+        key: u64,
+        request_body: String,
+        response: CachedResponse,
+    ) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
+        let path = self.key_path(key);
+        let entry = StoredEntry {
+            request_body,
+            response,
+            expires_at: Self::now_secs() + self.ttl.as_secs(),
+        };
+        Box::pin(async move {
+            let bytes = match serde_json::to_vec(&entry) {
+                Ok(b) => b,
+                Err(e) => {
+                    tracing::warn!("OpenDAL cache: failed to serialize entry: {e}");
+                    return;
+                }
+            };
+            if let Err(e) = self.operator.write(&path, bytes).await {
+                tracing::warn!("OpenDAL cache: failed to write {path}: {e}");
+            }
+        })
+    }
+
+    fn remove(&self, key: u64) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
+        let path = self.key_path(key);
+        Box::pin(async move {
+            if let Err(e) = self.operator.delete(&path).await {
+                tracing::warn!("OpenDAL cache: failed to delete {path}: {e}");
+            }
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::tower::cache::{CacheStore, CachedResponse};
+    use crate::types::{AssistantMessage, ChatCompletionResponse, Choice, FinishReason};
+
+    fn memory_store(ttl_secs: u64) -> OpenDalCacheStore {
+        let op = Operator::via_iter(opendal::Scheme::Memory, std::iter::empty::<(String, String)>())
+            .expect("memory backend should always build");
+        OpenDalCacheStore::new(op, "test/", Duration::from_secs(ttl_secs))
+    }
+
+    fn dummy_response() -> CachedResponse {
+        CachedResponse::Chat(ChatCompletionResponse {
+            id: "test-resp-001".into(),
+            object: "chat.completion".into(),
+            created: 1_700_000_000,
+            model: "gpt-4".into(),
+            choices: vec![Choice {
+                index: 0,
+                message: AssistantMessage {
+                    content: Some("Hello!".into()),
+                    name: None,
+                    tool_calls: None,
+                    refusal: None,
+                    function_call: None,
+                },
+                finish_reason: Some(FinishReason::Stop),
+            }],
+            usage: None,
+            system_fingerprint: None,
+            service_tier: None,
+        })
+    }
+
+    #[tokio::test]
+    async fn put_and_get_round_trip() {
+        let store = memory_store(300);
+        store.put(42, "request-body-a".into(), dummy_response()).await;
+        let cached = store.get(42, "request-body-a").await;
+        assert!(cached.is_some(), "expected a cached response after put");
+        match cached.unwrap() {
+            CachedResponse::Chat(resp) => {
+                assert_eq!(resp.id, "test-resp-001");
+                assert_eq!(resp.model, "gpt-4");
+            }
+            _ => panic!("expected CachedResponse::Chat variant"),
+        }
+    }
+
+    #[tokio::test]
+    async fn get_returns_none_for_missing_key() {
+        let store = memory_store(300);
+        let result = store.get(999, "any-body").await;
+        assert!(result.is_none(), "expected None for a key that was never stored");
+    }
+
+    #[tokio::test]
+    async fn get_returns_none_for_wrong_request_body() {
+        let store = memory_store(300);
+        store.put(1, "body-alpha".into(), dummy_response()).await;
+        // Same key but different request body should miss (collision guard).
+        let result = store.get(1, "body-beta").await;
+        assert!(result.is_none(), "expected None when request body does not match");
+    }
+
+    #[tokio::test]
+    async fn expired_entry_returns_none() {
+        let store = memory_store(0); // 0-second TTL = immediate expiry
+        store.put(1, "req".into(), dummy_response()).await;
+        // TTL is stored in whole seconds, so we must wait at least 1 full
+        // second for the wall-clock to advance past the `expires_at` timestamp.
+        tokio::time::sleep(Duration::from_millis(1100)).await;
+        let result = store.get(1, "req").await;
+        assert!(result.is_none(), "expected None for expired entry");
+    }
+
+    #[tokio::test]
+    async fn remove_deletes_entry() {
+        let store = memory_store(300);
+        store.put(7, "req".into(), dummy_response()).await;
+        // Confirm it exists first.
+        assert!(store.get(7, "req").await.is_some());
+        // Remove and verify it is gone.
+        store.remove(7).await;
+        assert!(store.get(7, "req").await.is_none(), "expected None after remove");
+    }
+
+    #[tokio::test]
+    async fn overwrite_replaces_previous_entry() {
+        let store = memory_store(300);
+        store.put(1, "req".into(), dummy_response()).await;
+
+        // Overwrite with a different response.
+        let replacement = CachedResponse::Chat(ChatCompletionResponse {
+            id: "test-resp-002".into(),
+            object: "chat.completion".into(),
+            created: 1_700_000_001,
+            model: "gpt-4o".into(),
+            choices: vec![],
+            usage: None,
+            system_fingerprint: None,
+            service_tier: None,
+        });
+        store.put(1, "req".into(), replacement).await;
+
+        match store.get(1, "req").await {
+            Some(CachedResponse::Chat(resp)) => assert_eq!(resp.id, "test-resp-002"),
+            _ => panic!("expected updated CachedResponse::Chat variant"),
+        }
+    }
+
+    #[test]
+    fn from_config_rejects_unknown_scheme() {
+        let result = OpenDalCacheStore::from_config(
+            "nonexistent_backend_xyz",
+            std::collections::HashMap::new(),
+            "prefix/",
+            Duration::from_secs(60),
+        );
+        assert!(result.is_err(), "expected error for unknown scheme");
+    }
+}

data/vendor/liter-llm/src/tower/cooldown.rs

@@ -0,0 +1,231 @@
+//! Deployment cooldown middleware.
+//!
+//! [`CooldownLayer`] wraps a service and implements a cooldown period after
+//! transient errors.  When the inner service returns a transient error (as
+//! determined by [`LiterLlmError::is_transient`]), the service is marked as
+//! cooling down for a configurable duration.  During the cooldown period,
+//! incoming requests are immediately rejected with
+//! [`LiterLlmError::ServiceUnavailable`] without calling the inner service.
+
+use std::sync::Arc;
+use std::task::{Context, Poll};
+use std::time::{Duration, Instant};
+
+use tokio::sync::RwLock;
+use tower::{Layer, Service};
+
+use super::types::{LlmRequest, LlmResponse};
+use crate::client::BoxFuture;
+use crate::error::{LiterLlmError, Result};
+
+// ---- State -----------------------------------------------------------------
+
+struct CooldownState {
+    /// `None` when not cooling down, `Some(start)` when a cooldown is active.
+    cooldown_start: Option<Instant>,
+}
+
+// ---- Layer -----------------------------------------------------------------
+
+/// Tower [`Layer`] that applies a cooldown period after transient errors.
+pub struct CooldownLayer {
+    duration: Duration,
+}
+
+impl CooldownLayer {
+    /// Create a new cooldown layer.
+    ///
+    /// After a transient error, the wrapped service will reject all requests
+    /// for `duration` before allowing traffic through again.
+    #[must_use]
+    pub fn new(duration: Duration) -> Self {
+        Self { duration }
+    }
+}
+
+impl<S> Layer<S> for CooldownLayer {
+    type Service = CooldownService<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        CooldownService {
+            inner,
+            duration: self.duration,
+            state: Arc::new(RwLock::new(CooldownState { cooldown_start: None })),
+        }
+    }
+}
+
+// ---- Service ---------------------------------------------------------------
+
+/// Tower service produced by [`CooldownLayer`].
+pub struct CooldownService<S> {
+    inner: S,
+    duration: Duration,
+    state: Arc<RwLock<CooldownState>>,
+}
+
+impl<S: Clone> Clone for CooldownService<S> {
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner.clone(),
+            duration: self.duration,
+            state: Arc::clone(&self.state),
+        }
+    }
+}
+
+impl<S> Service<LlmRequest> for CooldownService<S>
+where
+    S: Service<LlmRequest, Response = LlmResponse, Error = LiterLlmError> + Send + Clone + 'static,
+    S::Future: Send + 'static,
+{
+    type Response = LlmResponse;
+    type Error = LiterLlmError;
+    type Future = BoxFuture<'static, LlmResponse>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: LlmRequest) -> Self::Future {
+        let state = Arc::clone(&self.state);
+        let duration = self.duration;
+        // IMPORTANT: do NOT call self.inner.call(req) here — the inner service
+        // must only be invoked *after* the cooldown check passes inside the
+        // async block.  Calling it eagerly would send the request even when the
+        // service is in a cooldown period.
+        let mut inner = self.inner.clone();
+
+        Box::pin(async move {
+            // Check whether we are in a cooldown period.
+            {
+                let read = state.read().await;
+                if let Some(start) = read.cooldown_start {
+                    if start.elapsed() < duration {
+                        return Err(LiterLlmError::ServiceUnavailable {
+                            message: format!(
+                                "service is cooling down for {:.0}s after a transient error",
+                                duration.as_secs_f64()
+                            ),
+                        });
+                    }
+                    // Cooldown has expired — we need to reset it.
+                    // Drop the read lock first, then take the write lock.
+                    drop(read);
+                    let mut write = state.write().await;
+                    // Double-check under write lock (another task may have reset it).
+                    if let Some(s) = write.cooldown_start
+                        && s.elapsed() >= duration
+                    {
+                        write.cooldown_start = None;
+                    }
+                }
+            }
+
+            // Only call the inner service after cooldown check passes.
+            match inner.call(req).await {
+                Ok(resp) => Ok(resp),
+                Err(e) if e.is_transient() => {
+                    // Enter cooldown.
+                    let mut write = state.write().await;
+                    write.cooldown_start = Some(Instant::now());
+                    Err(e)
+                }
+                Err(e) => Err(e),
+            }
+        })
+    }
+}
+
+// ---- Tests -----------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use tower::{Layer as _, Service as _};
+
+    use super::*;
+    use crate::tower::service::LlmService;
+    use crate::tower::tests_common::{MockClient, chat_req};
+    use crate::tower::types::LlmRequest;
+
+    #[tokio::test]
+    async fn passes_through_on_success() {
+        let layer = CooldownLayer::new(Duration::from_secs(10));
+        let inner = LlmService::new(MockClient::ok());
+        let mut svc = layer.layer(inner);
+
+        let resp = svc.call(LlmRequest::Chat(chat_req("gpt-4"))).await;
+        assert!(resp.is_ok());
+    }
+
+    #[tokio::test]
+    async fn enters_cooldown_after_transient_error() {
+        let layer = CooldownLayer::new(Duration::from_secs(60));
+        let inner = LlmService::new(MockClient::failing_timeout());
+        let mut svc = layer.layer(inner);
+
+        // First call — transient error.
+        let err = svc
+            .call(LlmRequest::Chat(chat_req("gpt-4")))
+            .await
+            .expect_err("should fail");
+        assert!(matches!(err, LiterLlmError::Timeout));
+
+        // Second call — should be rejected with ServiceUnavailable (cooldown).
+        let err = svc
+            .call(LlmRequest::Chat(chat_req("gpt-4")))
+            .await
+            .expect_err("should be in cooldown");
+        assert!(
+            matches!(err, LiterLlmError::ServiceUnavailable { .. }),
+            "expected ServiceUnavailable during cooldown, got {err:?}"
+        );
+    }
+
+    #[tokio::test]
+    async fn cooldown_expires_after_duration() {
+        // Use a zero-second cooldown so it expires immediately.
+        let layer = CooldownLayer::new(Duration::from_millis(0));
+        let inner = LlmService::new(MockClient::failing_timeout());
+        let mut svc = layer.layer(inner);
+
+        // First call — transient error triggers cooldown.
+        svc.call(LlmRequest::Chat(chat_req("gpt-4")))
+            .await
+            .expect_err("should fail");
+
+        // With zero duration, cooldown is already expired.  The next call should
+        // reach the inner service (which will fail again with Timeout, not
+        // ServiceUnavailable).
+        let err = svc
+            .call(LlmRequest::Chat(chat_req("gpt-4")))
+            .await
+            .expect_err("should fail");
+        assert!(
+            matches!(err, LiterLlmError::Timeout),
+            "expected Timeout (cooldown expired), got {err:?}"
+        );
+    }
+
+    #[tokio::test]
+    async fn non_transient_error_does_not_trigger_cooldown() {
+        let layer = CooldownLayer::new(Duration::from_secs(60));
+        let inner = LlmService::new(MockClient::failing_auth());
+        let mut svc = layer.layer(inner);
+
+        // First call — non-transient error.
+        svc.call(LlmRequest::Chat(chat_req("gpt-4")))
+            .await
+            .expect_err("should fail");
+
+        // Second call — should reach inner service (not cooldown).
+        let err = svc
+            .call(LlmRequest::Chat(chat_req("gpt-4")))
+            .await
+            .expect_err("should fail with auth, not cooldown");
+        assert!(
+            matches!(err, LiterLlmError::BadRequest { .. }),
+            "expected BadRequest (auth), not ServiceUnavailable, got {err:?}"
+        );
+    }
+}