Common.hs

{-# LANGUAGE DuplicateRecordFields #-}

module GenAILib.Common
  where

import Control.Arrow ((&&&))
import Control.Monad (when)
import Data.Aeson (Key, ToJSON, Value (Number, String), object, toJSON)
import Data.Aeson.Key (fromString)
import Data.Aeson.Types (Pair)
import Data.Maybe (mapMaybe)
import Data.String.Conv (toS)
import Data.Text.Lazy qualified as TL
import GHC.Generics (Generic)


newtype Model = Model TL.Text
  deriving Generic

instance ToJSON Model


newtype Prompt = Prompt TL.Text
  deriving Generic

instance ToJSON Prompt


splitAtColon :: String -> Maybe (String, String)
splitAtColon combinedStr = do
  let (leftSide, rightSide) = takeWhile (/= ':') &&& dropWhile (/= ':') $ combinedStr
  when (null rightSide) Nothing
  pure (leftSide, tail rightSide)


newtype RawOutput = RawOutput Bool


newtype LLMOptions = LLMOptions { v :: Value }
  deriving Generic

instance ToJSON LLMOptions


-- For debugging the LLM options parsing, unused most of the time
debugOptsJSON :: LLMOptions -> Value
debugOptsJSON = toJSON

convertOptions :: [String] -> LLMOptions
convertOptions = LLMOptions . object . pairs'

pairs' :: [String] -> [Pair]
pairs' = map convertTypes . mapMaybe splitAtColon

convertTypes :: (String, String) -> (Key, Value)
-- Turns out the only non-number option we're interested in for Ollama is "stop"
convertTypes (keystr@"stop", valstr) = (fromString keystr, String . toS $ valstr)
convertTypes (keystr, valstr) = (fromString keystr, Number . read $ valstr)