Ollama.hs

{-# LANGUAGE DuplicateRecordFields #-}

module GenAILib.Ollama
  ( Host (..)
  , OllamaRequest (..)
  , Stream (..)
  , System (..)
  , defaultHost
  , hostFromString
  )
  where

import Data.Aeson (ToJSON, Value, defaultOptions, genericToEncoding,
  genericToJSON, omitNothingFields, toEncoding, toJSON)
import Data.Aeson qualified as Aeson
import Data.String.Conv (toS)
import Data.Text.Lazy qualified as TL
import Formatting ((%), formatToString, int, text)
import GHC.Generics (Generic)
import Text.Read (readMaybe)

import GenAILib.Common (Model, Prompt, splitAtColon)


newtype Stream = Stream Bool
  deriving Generic

instance ToJSON Stream


{- ^ This data structure is used to construct the POST body for an Ollama REST call
-}
data OllamaRequest = OllamaRequest
  { model :: Model
  , system :: Maybe System
  , prompt :: Prompt
  , stream :: Stream
  , options :: Maybe Value
  }
  deriving Generic

customOptions :: Aeson.Options
customOptions = defaultOptions { omitNothingFields = True }

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


data Host = Host TL.Text Int

-- I don't normally do this because it harms debugging when we want to see the
-- actual Haskell data structure serialized BUT this plays nicer with
-- optparse-applicative's showDefault
instance Show Host where
  show (Host hostName port) = formatToString (text % ":" % int) hostName port


defaultHost :: Host
defaultHost = Host "localhost" 11434


hostFromString :: String -> Maybe Host
hostFromString combinedStr = do
  (hostName, portStr) <- splitAtColon combinedStr
  Host (toS hostName) <$> readMaybe portStr


newtype System = System TL.Text
  deriving Generic

instance ToJSON System