OpenAI.hs

{-# LANGUAGE DuplicateRecordFields, OverloadedRecordDot #-}

module GenAILib.OpenAI
  ( OpenAIResponse (..)
  , OpenAIRequest (..)
  , defaultUrl
  , getMessage
  )
  where

import Data.Aeson (FromJSON, ToJSON, Value (Object, String), (.:),
  genericToEncoding, genericToJSON, parseJSON, toEncoding, toJSON)
import Data.Aeson.Key (fromText)
import Data.Aeson.KeyMap (singleton, union)
import Data.Aeson.Types (prependFailure, typeMismatch)
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 Data.Time.Clock.POSIX (posixSecondsToUTCTime)
import GHC.Generics (Generic)
import Servant.Client (BaseUrl, parseBaseUrl)

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


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

instance ToJSON OpenAIRequest where
  toJSON = flattenOptions . genericToJSON customJSONOptions
  toEncoding = genericToEncoding customJSONOptions

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

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

  setOpt :: TS.Text -> Value -> OpenAIRequest -> OpenAIRequest
  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


-- | Data returned for OpenAI requests
data OpenAIResponse = OpenAIResponse
  { model :: Model
  , created :: UTCTime
  , message :: Message
  , done :: Bool
  }
  deriving (Generic, Show)

-- Unfortunately, the OpenAI response JSON is a weird enough shape I decided to
-- take it apart it by-hand
{- HLINT ignore "Functor law" -}
instance FromJSON OpenAIResponse where
  parseJSON (Object o) = do
    choices <- o .: "choices"
    choice <- if null choices
      then fail "No elements in choices array, unable to decode this JSON"
      else pure $ head choices
    OpenAIResponse
      <$> (Model <$> o .: "model")
      <*> (posixSecondsToUTCTime <$> o .: "created")
      <*> (parseJSON =<< choice .: "message")
      <*> (isDone <$> choice .: "finish_reason")
  parseJSON invalid    = prependFailure "Parsing OpenAI chat completion response"
    $ typeMismatch "Object" invalid

isDone :: Value -> Bool
isDone (String s) = s == "stop"
isDone _ = False


-- | The default URL for OpenAI REST calls
defaultUrl :: BaseUrl
defaultUrl = fromMaybe (error "Unable to construct the default url for OpenAI!")
  $ parseBaseUrl "https://api.openai.com"


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