Utils.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

{- |
Module      : Langchain.Runnable.Utils
Description : Utility wrappers for Runnable components in LangChain
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>

This module provides various utility wrappers for 'Runnable' components that enhance
their behavior with common patterns like:

* Configuration management
* Result caching
* Automatic retries
* Timeout handling

These utilities follow the decorator pattern, wrapping existing 'Runnable' instances
with additional functionality while preserving the original input/output types.

Note: This module is experimental and the API may change in future versions.
-}
module Langchain.Runnable.Utils
  ( -- * Configuration Management
    WithConfig (..)

    -- * Caching
  , Cached (..)
  , cached

    -- * Resilience Patterns
  , Retry (..)
  , WithTimeout (..)
  ) where

import Control.Concurrent
import Data.Map.Strict as Map
import Langchain.Error (llmError)
import Langchain.Runnable.Core

{- | Wrapper for 'Runnable' components with configurable behavior.

This wrapper allows attaching configuration data to a 'Runnable' instance.
The configuration data can be accessed and modified without changing the
underlying 'Runnable' implementation.

Example:

@
data LLMConfig = LLMConfig
  { temperature :: Float
  , maxTokens :: Int
  }

let
  baseModel = OpenAI defaultOpenAIConfig
  configuredModel = WithConfig
    { configuredRunnable = baseModel
    , runnableConfig = LLMConfig 0.7 100
    }

-- Later, modify the configuration without changing the model
let updatedModel = configuredModel { runnableConfig = LLMConfig 0.9 150 }

-- Use the model as a regular Runnable
result <- invoke updatedModel "Explain monads in Haskell"
@
-}
data WithConfig config r
  = (Runnable r) =>
  WithConfig
  { configuredRunnable :: r
  -- ^ The wrapped 'Runnable' instance
  , runnableConfig :: config
  -- ^ Configuration data for this 'Runnable'
  }

-- | Make WithConfig a Runnable that applies the configuration
instance (Runnable r) => Runnable (WithConfig config r) where
  type RunnableInput (WithConfig config r) = RunnableInput r
  type RunnableOutput (WithConfig config r) = RunnableOutput r

  invoke (WithConfig r1 _) = invoke r1

{- | Cache results of a 'Runnable' to avoid duplicate computations.

This wrapper stores previously computed results in a thread-safe cache.
When an input is encountered again, the cached result is returned instead
of recomputing it, which can significantly improve performance for expensive
operations or when the same inputs are frequently processed.

Note: The cached results are stored in-memory and will be lost when the program
terminates. For persistent caching, consider implementing a custom wrapper that
uses database storage.

The 'RunnableInput' type must be an instance of 'Ord' for map lookups.
-}
data Cached r
  = (Runnable r, Ord (RunnableInput r)) =>
  Cached
  { cachedRunnable :: r
  -- ^ The wrapped 'Runnable' instance
  , cacheMap :: MVar (Map.Map (RunnableInput r) (RunnableOutput r))
  -- ^ Thread-safe cache storage
  }

cached :: (Runnable r, Ord (RunnableInput r)) => r -> IO (Cached r)
cached r = do
  cache <- newMVar Map.empty
  return $ Cached r cache

-- | Make Cached a Runnable that uses a cache
instance (Runnable r, Ord (RunnableInput r)) => Runnable (Cached r) where
  type RunnableInput (Cached r) = RunnableInput r
  type RunnableOutput (Cached r) = RunnableOutput r

  invoke (Cached r cacheRef) input = do
    cache <- readMVar cacheRef
    case Map.lookup input cache of
      Just output -> return $ Right output -- Cache hit: return cached result
      Nothing -> do
        -- Cache miss: compute and store resul
        result <- invoke r input
        case result of
          Left err -> return $ Left err
          Right output -> do
            modifyMVar_ cacheRef $ \c -> return $ Map.insert input output c
            return $ Right output

{- | Add retry capability to any 'Runnable'.

This wrapper automatically retries failed operations up to a specified
number of times with a configurable delay between attempts. This is particularly
useful for network operations or external API calls that might fail transiently.

Example:

@
-- Create an LLM with automatic retry for network failures
let
  baseModel = OpenAI defaultConfig
  resilientModel = Retry
    { retryRunnable = baseModel
    , maxRetries = 3
    , retryDelay = 1000000  -- 1 second delay between retries
    }

-- If the API call fails, it will retry up to 3 times
result <- invoke resilientModel "Generate a story about a Haskell programmer"
@
-}
data Retry r
  = (Runnable r) =>
  Retry
  { retryRunnable :: r
  -- ^ The wrapped 'Runnable' instance
  , maxRetries :: Int
  -- ^ Maximum number of retry attempts
  , retryDelay :: Int
  -- ^ Delay between retry attempts in microseconds
  }

-- | Make Retry a Runnable that retries on failure
instance (Runnable r) => Runnable (Retry r) where
  type RunnableInput (Retry r) = RunnableInput r
  type RunnableOutput (Retry r) = RunnableOutput r

  invoke (Retry r maxRetries_ delay) input = retryWithCount 0
    where
      retryWithCount count = do
        result <- invoke r input
        case result of
          Left err ->
            if count < maxRetries_
              then do
                threadDelay delay
                retryWithCount (count + 1)
              else return $ Left err
          Right output -> return $ Right output

{- | Add timeout capability to any 'Runnable'.

This wrapper enforces a maximum execution time for the wrapped 'Runnable'.
If the operation takes longer than the specified timeout, it is cancelled and
an error is returned. This is useful for limiting the execution time of potentially
long-running operations.

Example:

@
-- Create an LLM with a 30-second timeout
let
  baseModel = OpenAI defaultConfig
  timeboxedModel = WithTimeout
    { timeoutRunnable = baseModel
    , timeoutMicroseconds = 30000000  -- 30 seconds
    }

-- If the API call takes longer than 30 seconds, it will be cancelled
result <- invoke timeboxedModel "Generate a detailed analysis of Haskell's type system"
@

Note: This implementation uses 'forkIO' and 'killThread', which may not always
cleanly terminate the underlying operation, especially for certain types of I/O.
For critical applications, consider implementing a more robust timeout mechanism.
-}
data WithTimeout r
  = (Runnable r) =>
  WithTimeout
  { timeoutRunnable :: r
  -- ^ The wrapped 'Runnable' instance
  , timeoutMicroseconds :: Int
  -- ^ Timeout duration in microseconds
  }

-- | Make WithTimeout a Runnable that times out
instance (Runnable r) => Runnable (WithTimeout r) where
  type RunnableInput (WithTimeout r) = RunnableInput r
  type RunnableOutput (WithTimeout r) = RunnableOutput r

  invoke (WithTimeout r timeout) input = do
    resultVar <- newEmptyMVar

    -- Fork a thread to run the computation
    tid <- forkIO $ do
      result <- invoke r input
      putMVar resultVar (Just result)

    -- Set up the timeout
    timeoutTid <- forkIO $ do
      threadDelay timeout
      putMVar resultVar Nothing

    -- Wait for either result or timeout
    result <- takeMVar resultVar

    -- Kill the other thread
    killThread tid
    killThread timeoutTid

    case result of
      Just r_ -> return r_
      Nothing -> return $ Left (llmError "Operation timed out" Nothing Nothing)