GenAILib.hs

module GenAILib
  (
  -- * A library for interacting with LLM back-ends like Ollama and OpenAI
  --
  -- | An example
  --
  -- > import Control.Exception (Handler (..), catches)
  -- > import Data.Text.Lazy.IO qualified as TL
  -- > import GenAILib (ClientError, Request (..), jsonToText, mkRequest, numopt,
  -- >   stringopt, systemmsg, usermsg)
  -- > import GenAILib.HTTP (GenAIException, openaiV1Chat, openaiV1ChatJ,
  -- >   tokenFromFile)
  -- > import GenAILib.OpenAI (OpenAIRequest, getMessage)
  -- >
  -- >
  -- > main :: IO ()
  -- > main = do
  -- >   openaiJSON
  -- >   openaiData
  -- >
  -- > -- A simple example with no error handling that expects an Aeson Value
  -- > -- (OpenAIRequest is an instance of ToJSON) and displays the encoded JSON
  -- > -- response
  -- > openaiJSON :: IO ()
  -- > openaiJSON = do
  -- >   let req :: OpenAIRequest = mkRequest "gpt-3.5-turbo" $ usermsg "Why is the sky blue?"
  -- >   token <- tokenFromFile "path/to/openai/key"
  -- >   TL.putStrLn . jsonToText =<< openaiV1ChatJ token Nothing req
  -- >
  -- > -- Another example with exception handling that expects an OpenAIResponse
  -- > -- data structure, displaying that and also just the response text
  -- > openaiData :: IO ()
  -- > openaiData = do
  -- >   let req :: OpenAIRequest = mkRequest "gpt-3.5-turbo"
  -- >           (  systemmsg "Answer in the style of Bugs Bunny. Try to work the phrase \"What's up Doc?\" in somewhere."
  -- >           <> usermsg "Why is the sky blue?"
  -- >           <> numopt "temperature" 0.8
  -- >           <> stringopt "service_tier" "default"
  -- >           )
  -- >   token <- tokenFromFile "path/to/openai/key"
  -- >   res <- openaiV1Chat token Nothing req `catches`
  -- >     [ Handler (\(e :: GenAIException) -> error . show $ e)
  -- >     , Handler (\(e :: ClientError) -> error . show $ e)
  -- >     ]
  -- >   print res  -- The entire OpenAIResponse
  -- >   TL.putStrLn . getMessage $ res  -- Just the response Message Content

    Request (..)
  , RequestMod (..)
  , usermsg
  , numopt
  , boolopt
  , stringopt
  , systemmsg
  , optional
  , mkMessage
  , mkMessageWithRole
  , Model (..)
  , Message (..)
  , Role (..)
  , Content (..)
  , ImageData (..)
  , Stream (..)

  -- re-exporting
  , ClientError
  , module GenAILib.JSON
  )
  where

import Data.Aeson (FromJSON, ToJSON, Value (Number, String))
import Data.Aeson qualified as A
import Data.List.NonEmpty (NonEmpty)
import Data.Scientific (fromFloatDigits)
import Data.Text qualified as TS
import Data.Text.Lazy qualified as TL
import GHC.Generics (Generic)
import Servant.Client.Core (ClientError)

import GenAILib.JSON


-- | The name of an LLM model
--
--   This lets the back-end know which model we want to use by name.
newtype Model = Model TL.Text
  deriving (Generic, Show)

instance FromJSON Model
instance ToJSON Model


-- | Text values sent to and from LLMs like prompts and responses.
newtype Content = Content TL.Text
  deriving (Generic, Show)

instance FromJSON Content
instance ToJSON Content


-- | Which role is responsible for this content
--
--   Valid values are 'system', 'user', 'assistant' or 'tool'
newtype Role = Role TL.Text
  deriving (Generic, Show)

instance FromJSON Role
instance ToJSON Role


-- | Base64 encoded image data for working with models
newtype ImageData = ImageData TL.Text
  deriving (Generic, Show)

instance FromJSON ImageData
instance ToJSON ImageData


-- | Requests and responses to and from LLMs contain one or more @Message@s
--
--   These data structures always contain a @Role@ and some @Content@ and may
--   also optionally have an array of image data.
--
--   NOTE: The images as base64 encoded @Text@ are specific to Ollama at this
--   time and work is ongoing to support OpenAI's APIs
data Message = Message
  { role :: Role
  , content :: Content
  , images :: Maybe (NonEmpty ImageData)
  }
  deriving (Generic, Show)

instance FromJSON Message
instance ToJSON Message


-- | Convenience function to create a simple @Message@ with the role of 'user',
--   some @Content@ (prompt text) and no image data.
mkMessage :: TL.Text -> Message
mkMessage = mkMessageWithRole $ Role "user"


-- | Convenience function to create a simple @Message@ with the specified role,
--   some @Content@ (prompt text) and no image data.
mkMessageWithRole :: Role -> TL.Text -> Message
mkMessageWithRole role' content = Message role' (Content content) Nothing


-- | Whether the back-end should use streaming to send response data in "chunks"
--
--   NOTE: This is mostly unused at this time and the behavior is untested
newtype Stream = Stream Bool
  deriving (Generic, Show)

instance ToJSON Stream


class Request a where
  -- | Append a @Message@ to an existing request
  appendMessage :: Message -> a -> a

  -- | Construct a new @Request@
  --   Note: This request has no prompt data yet so at a minimum it must be used like this
  --
  --   > mkRequest "somefabulousmodel" $ usermsg "some fabulous prompt"
  mkRequest :: TL.Text -> RequestMod a -> a

  -- | Set an arbitrary LLM option given a name and JSON @Value@
  --   This is used internally by functions like boolopt, stringopt, ...
  setOpt :: TS.Text -> Value -> a -> a


newtype RequestMod r = RequestMod (r -> r)

instance Request r => Semigroup (RequestMod r) where
  RequestMod f1 <> RequestMod f2 = RequestMod (f2 . f1)


-- | Append a system prompt
systemmsg :: Request r => TL.Text -> RequestMod r
systemmsg systemPrompt = RequestMod . appendMessage
  $ mkMessageWithRole (Role "system") systemPrompt


-- | Append a user message (prompt)
usermsg :: Request r => TL.Text -> RequestMod r
usermsg prompt = RequestMod . appendMessage $ mkMessage prompt


-- | Apply a request modifier to a Maybe value
--
--   > stringopt `optional` (Just "foo")
optional :: (a -> RequestMod r) -> Maybe a -> RequestMod r
optional = maybe $ RequestMod id


-- | Set an option with a @Bool@ value
boolopt :: Request r => TS.Text -> Bool -> RequestMod r
boolopt keystr = RequestMod . setOpt keystr . A.Bool


-- | Set an option with a @Double@ value
numopt :: Request r => TS.Text -> Double -> RequestMod r
numopt keystr = RequestMod . setOpt keystr . Number . fromFloatDigits


-- | Set an option with a @Text@ value
stringopt :: Request r => TS.Text -> TS.Text -> RequestMod r
stringopt keystr = RequestMod . setOpt keystr . String