Cache.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}

-- | The cache storage effect and the memoizing interpreter (EP-6).
--
-- 'Cache' is a small @effectful@ dynamic effect with two operations —
-- 'lookupCache' and 'storeCache' — interpreted by a backend (memory, SQLite, …).
-- The /policy/ is separate: 'cachedLLM' re-interprets EP-1's @LLM@ effect so that
-- 'Shikumi.LLM.complete' is memoized — identical requests contact the provider
-- once. Keeping mechanism (storage) and policy (when to read/write) apart lets a
-- caller use the store directly, or compose the memoizer, independently.
module Shikumi.Cache
  ( -- * The storage effect
    Cache (..),
    lookupCache,
    storeCache,

    -- * The memoizing policy
    CacheConfig (..),
    defaultCacheConfig,
    cachedLLM,
    cachedLLMWith,

    -- * Re-exports
    CacheKey (..),
    cacheKey,
    currentKeyVersion,
    CachedResponse (..),
  )
where

import Baikai (Response, StopReason (ErrorReason))
import Control.Lens ((&), (.~), (^.))
import Control.Monad (when)
import Data.Generics.Labels ()
import Data.Maybe (isJust, isNothing)
import Data.Time.Clock (NominalDiffTime, diffUTCTime)
import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, (:>))
import Effectful.Dispatch.Dynamic (interpose, passthrough, send)
import Effectful.Error.Static (Error, throwError)
import GHC.Generics (Generic)
import Shikumi.Cache.Key (CacheKey (..), cacheKey, currentKeyVersion)
import Shikumi.Cache.Types (CachedResponse (..))
import Shikumi.Effect.Time (Time, getCurrentTime)
import Shikumi.Error (ShikumiError)
import Shikumi.LLM (LLM (..), complete)
import Shikumi.LLM.Continuation (validateRequestContinuation)

-- | The cache storage effect: look an entry up by key, or store one.
data Cache :: Effect where
  LookupCache :: CacheKey -> Cache m (Maybe CachedResponse)
  StoreCache :: CacheKey -> CachedResponse -> Cache m ()

type instance DispatchOf Cache = 'Dynamic

-- | Look up a cached response by key (a backend MISS is 'Nothing').
lookupCache :: (Cache :> es) => CacheKey -> Eff es (Maybe CachedResponse)
lookupCache = send . LookupCache

-- | Store a cached response under a key (idempotent upsert at every backend).
storeCache :: (Cache :> es) => CacheKey -> CachedResponse -> Eff es ()
storeCache k v = send (StoreCache k v)

-- | Policy knobs for 'cachedLLMWith', shared by every backend.
--
-- 'entryTTL' is the maximum age of a usable entry, measured against
-- 'CachedResponse.storedAt' at lookup time. 'Nothing' (the default) means
-- entries never expire, which is the uniform default across Memory, SQLite,
-- Redis, and Postgres. Expiry is enforced here, at the policy layer, so it
-- behaves identically no matter which backend interprets the 'Cache' effect; an
-- expired entry is treated as a MISS and overwritten by the fresh response.
newtype CacheConfig = CacheConfig
  { entryTTL :: Maybe NominalDiffTime
  }
  deriving stock (Eq, Show, Generic)

-- | No expiry.
defaultCacheConfig :: CacheConfig
defaultCacheConfig = CacheConfig {entryTTL = Nothing}

-- | Memoize EP-1's @LLM@ @complete@ through the 'Cache' effect: on a HIT (a key
-- present whose 'keyVersion' matches the live namespace and whose age satisfies
-- 'defaultCacheConfig') return the stored response without contacting the
-- provider; on a MISS delegate to the underlying @LLM@ handler, store the
-- result if it is successful, and return it. Cache backends are best-effort:
-- lookup failures are MISSes and store failures are no-ops. Entries never
-- expire unless 'entryTTL' is set via 'cachedLLMWith'. In-band error responses
-- are never cached. Under concurrent identical requests both callers may miss
-- and call the provider; this accepted check-then-act race is harmless because
-- stores are idempotent upserts keyed by content. The streaming op is passed
-- through unchanged — streams are not cached. Any evidence request bypasses
-- both reads and writes, since cached responses cannot establish a new provider
-- crossing. Compose defaults before this memoizer in request execution order.
cachedLLM ::
  (Cache :> es, LLM :> es, Time :> es, Error ShikumiError :> es) =>
  Eff es a ->
  Eff es a
cachedLLM = cachedLLMWith defaultCacheConfig

-- | A configured variant of 'cachedLLM'. See 'CacheConfig' for the shared TTL
-- policy.
cachedLLMWith ::
  (Cache :> es, LLM :> es, Time :> es, Error ShikumiError :> es) =>
  CacheConfig ->
  Eff es a ->
  Eff es a
cachedLLMWith cfg = interpose $ \env -> \case
  Complete model ctx opts -> do
    either throwError pure (validateRequestContinuation model ctx opts)
    if isJust (opts ^. #evidence)
      then complete model ctx opts
      else do
        let key = cacheKey model ctx opts
        hit <- lookupCache key
        now <- getCurrentTime
        case hit of
          Just cr
            | keyVersion cr == currentKeyVersion,
              fresh (entryTTL cfg) now (storedAt cr) ->
                pure (response cr & #evidence .~ Nothing)
          _ -> do
            resp <- complete model ctx opts
            stored <- getCurrentTime
            when (cacheable resp) $
              storeCache key (CachedResponse resp stored currentKeyVersion)
            pure resp
  other -> passthrough env other
  where
    fresh Nothing _ _ = True
    fresh (Just ttl) now written = diffUTCTime now written <= ttl

-- | Only successful responses are memoized. An in-band error response reports a
-- transient provider failure; caching it would replay the outage as the
-- permanent answer.
cacheable :: Response -> Bool
cacheable resp =
  (resp ^. #message . #stopReason) /= ErrorReason
    && isNothing (resp ^. #errorInfo)