lean_ctx/proxy/
ocla_cache_bridge.rs

1//! Adapter between proxy request fields and the OCLA response cache.
2
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5
6use axum::http::StatusCode;
7
8use crate::core::ocla::response_cache::{CachedResponse, ResponseCache, ResponseCacheKey};
9
10/// Shared proxy adapter for the OCLA response cache.
11#[derive(Clone, Debug)]
12pub struct OclaCacheBridge {
13    cache: Arc<ResponseCache>,
14}
15
16#[derive(Clone, Debug, PartialEq, Eq)]
17pub struct CachedProviderResponse {
18    pub body: Vec<u8>,
19    pub status: StatusCode,
20}
21
22impl OclaCacheBridge {
23    /// Creates a bridge backed by a shared OCLA response cache.
24    pub fn new(cache: Arc<ResponseCache>) -> Self {
25        Self { cache }
26    }
27
28    /// Returns a cached response body for matching request fields.
29    pub fn try_cache_hit(
30        &self,
31        model: &str,
32        prompt_hash: &str,
33        temp: f32,
34        max_tokens: u32,
35    ) -> Option<CachedProviderResponse> {
36        self.cache
37            .get(&cache_key(model, prompt_hash, temp, max_tokens))
38            .map(|response| CachedProviderResponse {
39                body: response.body,
40                status: StatusCode::from_u16(response.status).unwrap_or(StatusCode::OK),
41            })
42    }
43
44    /// Records a response for matching request fields.
45    pub fn record_response(
46        &self,
47        model: &str,
48        prompt_hash: &str,
49        temp: f32,
50        max_tokens: u32,
51        status: StatusCode,
52        body: &[u8],
53        tokens: u64,
54    ) {
55        self.cache.put(
56            cache_key(model, prompt_hash, temp, max_tokens),
57            CachedResponse {
58                body: body.to_vec(),
59                status: status.as_u16(),
60                tokens,
61                created_at: Instant::now(),
62                ttl: Duration::ZERO,
63            },
64        );
65    }
66}
67
68/// Computes the stable blake3 hash used by the bridge for a request body.
69pub fn prompt_hash(request_body: &[u8]) -> String {
70    blake3::hash(request_body).to_hex().to_string()
71}
72
73fn cache_key(model: &str, prompt_hash: &str, temp: f32, max_tokens: u32) -> ResponseCacheKey {
74    let prompt_digest = blake3::hash(prompt_hash.as_bytes());
75    let mut hash_bytes = [0; 8];
76    hash_bytes.copy_from_slice(&prompt_digest.as_bytes()[..8]);
77    ResponseCacheKey::new(
78        model,
79        u64::from_be_bytes(hash_bytes),
80        temp,
81        u64::from(max_tokens),
82    )
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[test]
90    fn prompt_hash_is_stable_and_content_sensitive() {
91        assert_eq!(prompt_hash(b"prompt"), prompt_hash(b"prompt"));
92        assert_ne!(prompt_hash(b"prompt"), prompt_hash(b"different"));
93    }
94
95    #[test]
96    fn bridge_serves_recorded_response_and_misses_other_requests() {
97        let bridge = OclaCacheBridge::new(Arc::new(ResponseCache::new(4, Duration::from_mins(1))));
98        let hash = prompt_hash(br#"{"prompt":"hello"}"#);
99
100        assert!(bridge.try_cache_hit("model", &hash, 0.2, 128).is_none());
101        bridge.record_response(
102            "model",
103            &hash,
104            0.2,
105            128,
106            StatusCode::INTERNAL_SERVER_ERROR,
107            b"answer",
108            7,
109        );
110
111        assert_eq!(
112            bridge.try_cache_hit("model", &hash, 0.2, 128),
113            Some(CachedProviderResponse {
114                body: b"answer".to_vec(),
115                status: StatusCode::INTERNAL_SERVER_ERROR,
116            })
117        );
118        assert!(bridge.try_cache_hit("model", &hash, 0.3, 128).is_none());
119        assert!(
120            bridge
121                .try_cache_hit("other-model", &hash, 0.2, 128)
122                .is_none()
123        );
124    }
125}
lean_ctx/proxy/ocla_cache_bridge.rs

lean_ctx/proxy/
ocla_cache_bridge.rs
