Ollama.hs

{-# LANGUAGE DuplicateRecordFields, OverloadedRecordDot #-}

module GenAILib.Ollama
  ( OllamaRequest (..)
  , defaultUrl
  , KeepAlive (..)
  , keepalive
  , OllamaResponse (..)
  , getMessage
  -- , stream'
  )
  where

import Data.Aeson (FromJSON, ToJSON, Value (Object),
  genericToEncoding, genericToJSON, toEncoding, toJSON)
import Data.Aeson.Key (fromText)
import Data.Aeson.KeyMap (singleton, union)
import Data.Maybe (fromMaybe)
import Data.Sequence (Seq, (|>), empty)
import Data.Text qualified as TS
import Data.Text.Lazy qualified as TL
import Data.Time (UTCTime)
import GHC.Generics (Generic)
import Servant.Client (BaseUrl, parseBaseUrl)

import GenAILib (Content (..), Message (..), Model (..), Request (..),
  RequestMod (..), Stream (..))
import GenAILib.JSON (customJSONOptions)


-- | Data for making Ollama requests
data OllamaRequest = OllamaRequest
  { model :: Model
  , messages :: Seq Message
  , options :: Maybe Value
  , stream :: Stream
  , keep_alive :: Maybe KeepAlive
  }
  deriving (Generic, Show)

instance ToJSON OllamaRequest where
  toJSON     = genericToJSON customJSONOptions
  toEncoding = genericToEncoding customJSONOptions

instance Request OllamaRequest where
  appendMessage :: Message -> OllamaRequest -> OllamaRequest
  appendMessage newMessage oldReq = oldReq { messages = oldReq.messages |> newMessage }

  mkRequest :: TL.Text -> RequestMod OllamaRequest -> OllamaRequest
  mkRequest modelName (RequestMod optSetter) = optSetter
    $ OllamaRequest (Model modelName) empty Nothing (Stream False) Nothing

  setOpt :: TS.Text -> Value -> OllamaRequest -> OllamaRequest
  setOpt keystr val req = req { options = Just . Object $ combinedKm }
    where
      newKm = singleton (fromText keystr) val
      combinedKm = case options req of
        (Just (Object existingKm)) -> existingKm `union` newKm
        _ -> newKm


-- | How long an Ollama model is kept in memory before clean-up by the Ollama
--   daemon
--
--   This is specific to Ollama.
newtype KeepAlive = KeepAlive TL.Text
  deriving (Generic, Show)

instance ToJSON KeepAlive


-- FIXME streaming does not work yet for Ollama via this library
-- Enable or disable streaming
-- stream' :: Bool -> OllamaRequest -> OllamaRequest
-- stream' b req = req { stream = Stream b }


-- | Set the keep_alive duration
--
--   This takes string duration values like '35m' or '1h'
keepalive :: TL.Text -> RequestMod OllamaRequest
keepalive ka = RequestMod (\req -> req { keep_alive = Just (KeepAlive ka) })


-- | Data returned for Ollama requests
data OllamaResponse = OllamaResponse
  { model :: Model
  , created_at :: UTCTime
  , message :: Message
  , done :: Bool
  -- Doubtful anyone really needs these fields
  -- , total_duration :: Integer
  -- , load_duration :: Integer
  -- , prompt_eval_count :: Integer
  -- , prompt_eval_duration :: Integer
  -- , eval_count :: Integer
  -- , eval_duration :: Integer
  }
  deriving (Generic, Show)

instance FromJSON OllamaResponse


-- | The default URL for making Ollama REST calls
defaultUrl :: BaseUrl
defaultUrl = fromMaybe (error "Unable to construct the default url for Ollama!")
  $ parseBaseUrl "localhost:11434"


-- | Convenience function to extract the response text
getMessage :: OllamaResponse -> TL.Text
getMessage res = assistantResponse
  where (Content assistantResponse) = content . message $ res