litellm/litellm-rust/crates/inference/src/caching.rs

use std::{
    future::Future,
    marker::PhantomData,
    sync::Arc,
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use bytes::{Bytes, BytesMut};
use futures_util::{StreamExt, TryStreamExt, stream};
use litellm_cache_response::{
    CacheOptions, CachePolicy, ResponseCacheRequest, ResponseCacheService, ResponseEnvelope,
    ScopedCache, cache_key,
};
use litellm_host::{
    call::{CallOutput, OutputOf},
    interceptors::{
        Cost, ExecutionFacts, Interceptors, ProviderIdentity, ResultSource, WireRequest,
    },
    lifecycle::{CallEvent, ExecutionEvent},
    observation::ObservationSender,
    protocol::Protocol,
};
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::Value;
use tokio_util::codec::Decoder;

use crate::Error;

pub trait Cachable: Protocol<Error = Error> {
    const SURFACE: &'static str;

    fn reusable(_response: &Self::Response) -> bool {
        true
    }
}

pub struct CacheRequest {
    pub identity: ProviderIdentity,
    pub input: Value,
}

impl CacheRequest {
    pub fn from_wire(identity: ProviderIdentity, wire: Option<&WireRequest>) -> Self {
        Self {
            input: wire.map_or(Value::Null, |wire| {
                serde_json::json!({
                    "provider": identity.provider,
                    "model": identity.model,
                    "url": wire.url,
                    "headers": wire.headers,
                    "body": wire.body,
                })
            }),
            identity,
        }
    }
}

pub trait StreamCachable: Cachable {
    const TERMINAL_EVENT: &'static str;

    fn replay(data: Bytes) -> Option<OutputOf<Self>>;
    fn bytes(chunk: &Self::Chunk) -> &[u8];
}

#[derive(Serialize, Deserialize)]
#[serde(tag = "kind", content = "value")]
pub enum CachedOutput<R> {
    Response(R),
    Stream(String),
}

struct CacheSession {
    service: Arc<dyn ResponseCacheService>,
    request: ResponseCacheRequest,
}

impl CacheSession {
    fn prepare<P: Cachable>(
        service: Option<Arc<dyn ResponseCacheService>>,
        options: Option<CacheOptions>,
        request: &CacheRequest,
    ) -> Option<Self> {
        let options = options.filter(|options| options.policy.enabled())?;
        let service = service?;
        let input = request.input.clone();
        let request = options.request(&service.config().namespace, P::SURFACE, input);
        Some(Self { service, request })
    }

    async fn lookup<P: Cachable>(&self) -> Option<CachedOutput<P::Response>>
    where
        P::Response: DeserializeOwned,
    {
        if !self.request.controls.reads() {
            return None;
        }
        match self.service.lookup(&self.request, now()).await {
            Ok(Some(value)) => {
                serde_json::from_value::<ResponseEnvelope<CachedOutput<P::Response>>>(value)
                    .ok()
                    .and_then(|entry| entry.decode(P::SURFACE))
            }
            Ok(None) => None,
            Err(_) => {
                tracing::warn!("response cache lookup failed");
                None
            }
        }
    }

    async fn store(&self, entry: Value) {
        if !self.request.controls.writes() {
            return;
        }
        if self
            .service
            .store(&self.request, entry, now())
            .await
            .is_err()
        {
            tracing::warn!("response cache write failed");
        }
    }

    async fn store_response<P: Cachable>(&self, response: &P::Response)
    where
        P::Response: Serialize,
    {
        if !self.request.controls.writes() || !P::reusable(response) {
            return;
        }
        if let Ok(value) = serde_json::to_value(response)
            && let Ok(entry) = serde_json::to_value(ResponseEnvelope::new(
                P::SURFACE,
                CachedOutput::Response(value),
            ))
        {
            self.store(entry).await;
        }
    }
}

pub async fn execute_unary<P, F, Fut>(
    request: CacheRequest,
    cache: Option<Arc<dyn ResponseCacheService>>,
    options: Option<CacheOptions>,
    interceptors: &impl Interceptors<Error>,
    observers: Option<&ObservationSender>,
    provider: F,
) -> Result<P::Response, Error>
where
    P: Cachable,
    P::Response: Serialize + DeserializeOwned,
    F: FnOnce() -> Fut,
    Fut: Future<Output = Result<P::Response, Error>>,
{
    let identity = request.identity.clone();
    crate::diagnostic::provider(&identity.model, &identity.provider);
    let session = CacheSession::prepare::<P>(cache, options, &request);
    let hit = match &session {
        Some(session) => session.lookup::<P>().await.and_then(|entry| match entry {
            CachedOutput::Response(response) => Some((response, cache_key(&session.request.key))),
            CachedOutput::Stream(_) => None,
        }),
        None => None,
    };
    let (response, source) = match hit {
        Some((response, key)) => (response, ResultSource::Cache { key }),
        None => (provider().await?, ResultSource::Provider),
    };
    let from_provider = source == ResultSource::Provider;
    let cost = Cost::settle(&source, None);
    publish(
        ExecutionFacts {
            provider: identity,
            source,
            cost,
        },
        interceptors,
        observers,
    )
    .await?;
    if from_provider && let Some(session) = session {
        session.store_response::<P>(&response).await;
    }
    Ok(response)
}

pub async fn execute_streaming<P, F, Fut>(
    request: CacheRequest,
    cache: Option<Arc<dyn ResponseCacheService>>,
    options: Option<CacheOptions>,
    interceptors: &impl Interceptors<Error>,
    observers: Option<&ObservationSender>,
    provider: F,
) -> Result<OutputOf<P>, Error>
where
    P: StreamCachable,
    P::Response: Serialize + DeserializeOwned,
    F: FnOnce() -> Fut,
    Fut: Future<Output = Result<OutputOf<P>, Error>>,
{
    let identity = request.identity.clone();
    crate::diagnostic::provider(&identity.model, &identity.provider);
    let session = CacheSession::prepare::<P>(cache, options, &request);
    let cache = CallCache::<P> {
        session,
        protocol: PhantomData,
    };
    let hit = cache.lookup().await;
    let (output, source) = match hit {
        Some(hit) => hit,
        None => (provider().await?, ResultSource::Provider),
    };
    publish(
        ExecutionFacts {
            provider: identity,
            cost: Cost::settle(&source, None),
            source: source.clone(),
        },
        interceptors,
        observers,
    )
    .await?;
    Ok(cache.finish(output, &source).await)
}

pub struct CallCache<P> {
    session: Option<CacheSession>,
    protocol: PhantomData<P>,
}

impl<P: StreamCachable> CallCache<P> {
    pub fn from_wire(
        cache: Option<&ScopedCache>,
        policy: CachePolicy,
        identity: &ProviderIdentity,
        wire: &WireRequest,
    ) -> Self {
        let session = cache.and_then(|cache| {
            if !policy.enabled() {
                return None;
            }
            let options = cache.options(Some(policy));
            let request = CacheRequest::from_wire(identity.clone(), Some(wire));
            Some(CacheSession {
                request: options.request(
                    &cache.service.config().namespace,
                    P::SURFACE,
                    request.input,
                ),
                service: cache.service.clone(),
            })
        });
        Self {
            session,
            protocol: PhantomData,
        }
    }

    pub async fn lookup(&self) -> Option<(OutputOf<P>, ResultSource)>
    where
        P::Response: DeserializeOwned,
    {
        let session = self.session.as_ref()?;
        let output = match session.lookup::<P>().await? {
            CachedOutput::Response(response) => CallOutput::Complete(response),
            CachedOutput::Stream(data) => P::replay(Bytes::from(data))?,
        };
        Some((
            output,
            ResultSource::Cache {
                key: cache_key(&session.request.key),
            },
        ))
    }

    pub async fn finish(self, output: OutputOf<P>, source: &ResultSource) -> OutputOf<P>
    where
        P::Response: Serialize,
    {
        let Some(session) = self.session.filter(|session| {
            *source == ResultSource::Provider && session.request.controls.writes()
        }) else {
            return output;
        };
        match output {
            CallOutput::Complete(response) => {
                session.store_response::<P>(&response).await;
                CallOutput::Complete(response)
            }
            CallOutput::Stream { head, chunks } => CallOutput::Stream {
                head,
                chunks: capture_stream::<P>(chunks, session),
            },
        }
    }
}

fn capture_stream<P: StreamCachable>(
    chunks: futures_util::stream::BoxStream<'static, Result<P::Chunk, Error>>,
    session: CacheSession,
) -> futures_util::stream::BoxStream<'static, Result<P::Chunk, Error>> {
    stream::try_unfold(
        (chunks, Some(Vec::<u8>::new()), session),
        |(mut chunks, captured, session)| async move {
            match chunks.try_next().await? {
                Some(chunk) => {
                    let captured = captured.and_then(|mut data| {
                        let bytes = P::bytes(&chunk);
                        if data.len().saturating_add(bytes.len())
                            > session.service.config().max_entry_bytes
                        {
                            return None;
                        }
                        data.extend_from_slice(bytes);
                        Some(data)
                    });
                    Ok(Some((chunk, (chunks, captured, session))))
                }
                None => {
                    if let Some(data) = captured
                        && let Ok(text) = String::from_utf8(data)
                        && successful_stream(&text, P::TERMINAL_EVENT)
                        && let Ok(entry) = serde_json::to_value(ResponseEnvelope::new(
                            P::SURFACE,
                            CachedOutput::<Value>::Stream(text),
                        ))
                    {
                        session.store(entry).await;
                    }
                    Ok::<_, Error>(None)
                }
            }
        },
    )
    .boxed()
}

fn now() -> Duration {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
}

fn successful_stream(text: &str, terminal: &str) -> bool {
    let mut pending = BytesMut::from(text.as_bytes());
    let mut codec = litellm_framer::sse::SseCodec::default();
    let mut complete = false;
    loop {
        let event = match codec.decode(&mut pending) {
            Ok(Some(event)) => event,
            Ok(None) => return complete && pending.is_empty(),
            Err(_) => return false,
        };
        let Ok(value) = serde_json::from_str::<Value>(&event.data) else {
            return false;
        };
        let Some(kind) = value.get("type").and_then(Value::as_str) else {
            return false;
        };
        if matches!(kind, "error" | "response.failed" | "response.incomplete") {
            return false;
        }
        complete |= kind == terminal;
    }
}

async fn publish(
    facts: ExecutionFacts,
    interceptors: &impl Interceptors<Error>,
    observers: Option<&ObservationSender>,
) -> Result<(), Error> {
    if let Some(observers) = observers {
        observers.emit(CallEvent::Execution(ExecutionEvent::ResultReady {
            facts: facts.clone(),
        }));
    }
    let from_provider = facts.source == ResultSource::Provider;
    interceptors.result_ready(facts).await.map_err(|error| {
        if from_provider {
            Error::post_call(error)
        } else {
            error
        }
    })
}

#[cfg(test)]
mod tests {
    use litellm_core_utils::decline::Decline;
    use litellm_host::interceptors::{RawResponse, RequestContext, WireRequest};

    use super::*;

    struct RejectResult;

    impl Interceptors<Error> for RejectResult {
        async fn before_provider_request(
            &self,
            wire: WireRequest,
            _: RequestContext,
        ) -> Result<WireRequest, Error> {
            Ok(wire)
        }

        async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> {
            Ok(())
        }

        async fn result_ready(&self, _: ExecutionFacts) -> Result<(), Error> {
            Err(Error::Declined(Decline::new("result rejected")))
        }
    }

    #[rstest::rstest]
    #[case::provider(ResultSource::Provider, false)]
    #[case::cache(ResultSource::Cache { key: "cached".into() }, true)]
    #[tokio::test]
    async fn result_rejection_only_declines_without_a_provider_call(
        #[case] source: ResultSource,
        #[case] may_decline: bool,
    ) {
        let facts = ExecutionFacts {
            provider: ProviderIdentity {
                model: "model".into(),
                provider: "provider".into(),
            },
            cost: Cost::settle(&source, None),
            source,
        };
        let error = publish(facts, &RejectResult, None).await.unwrap_err();
        assert_eq!(error.declined().is_some(), may_decline);
        assert_eq!(error.dispatched(), !may_decline);
        assert!(
            litellm_error::Report::new(&error)
                .to_string()
                .contains("result rejected")
        );
    }
}