HTTP.hs

{-# LANGUAGE DataKinds, TypeOperators #-}

module GenAILib.HTTP
  ( BearerToken (..)
  , GenAIException
  , ollamaChat
  , ollamaChatJ
  , openaiV1Chat
  , openaiV1ChatJ
  , tokenFromText
  , tokenFromFile
  )
  where

import Control.Exception (Exception, throwIO)
import Data.Aeson (FromJSON, ToJSON, Value (Object), decode, encode, toJSON)
import Data.Maybe (fromMaybe)
import Data.Proxy (Proxy (..))
import Data.Text qualified as TS
import Data.Text.IO qualified as TS
import Data.Typeable (Typeable)
import Network.HTTP.Client (Manager, defaultManagerSettings,
  managerResponseTimeout, newManager, responseTimeoutMicro)
import Network.HTTP.Client.TLS (newTlsManagerWith, tlsManagerSettings)
import Servant.API ((:<|>) (..), (:>), Header', JSON, Post, ReqBody, Required,
  ToHttpApiData)
import Servant.Client (BaseUrl (..), ClientM, mkClientEnv, runClientM)
import Servant.Client.Core (clientIn)

import GenAILib.Ollama (OllamaResponse)
import GenAILib.Ollama qualified as Ollama
import GenAILib.OpenAI (OpenAIResponse)
import GenAILib.OpenAI qualified as OpenAI


defaultHttpTimeout :: Int
defaultHttpTimeout = 1_200_000_000  -- In microseconds, this is 2 minutes


-- | A type to contain the Authorization Bearer token data needed for some requests
newtype BearerToken = BearerToken TS.Text
  deriving ToHttpApiData


-- | Construct bearer token authorization header data from @Text@
tokenFromText :: TS.Text -> BearerToken
tokenFromText = BearerToken . ("Bearer " <>) . TS.strip


-- | Construct bearer token authorization header data from the contents of a file
tokenFromFile :: FilePath -> IO BearerToken
tokenFromFile path = tokenFromText <$> TS.readFile path


{- NOTE: The reason we've chosen Aeson's Value type for the response is we want
   to do one of these things with these responses:

   1. Display only the LLM response text (in the JSON Value's 'response' or 'messages.content' field)
   2. Pretty-print the entire JSON response ( | jq . )
   3. Dump the entire JSON response as-is ( | jq -r|--raw-output . )

   We don't need to model the entire Ollama response JSON document to achieve these goals.
-}
type CommonReqMethodRes = ReqBody '[JSON] Value :> Post '[JSON] Value

type API
  =    "api" :> "chat" :> CommonReqMethodRes
  :<|> "v1" :> "chat" :> "completions"
    :> Header' '[Required] "Authorization" BearerToken
    :> CommonReqMethodRes

-- A Proxy to our API
api :: Proxy API
api = Proxy

clientM :: Proxy ClientM
clientM = Proxy

clientOllamaChat :: Value -> ClientM Value
clientOpenAIV1Chat :: BearerToken -> Value -> ClientM Value
(clientOllamaChat :<|> clientOpenAIV1Chat) = api `clientIn` clientM


data GenAIException = PostMarshallingException | UnknownResponse String
  deriving (Show, Typeable)

instance Exception GenAIException


-- | Make a request to an Ollama server resulting in a JSON @Value@
ollamaChatJ :: ToJSON j => Maybe BaseUrl -> j -> IO Value
ollamaChatJ mBaseUrl j = do
  manager' <- newManager $ defaultManagerSettings
    { managerResponseTimeout = responseTimeoutMicro defaultHttpTimeout }
  doRequest clientOllamaChat manager' (fromMaybe Ollama.defaultUrl mBaseUrl) j

-- | Make a request to an Ollama server resulting in a high-level Ollama response
ollamaChat :: ToJSON j => Maybe BaseUrl -> j -> IO OllamaResponse
ollamaChat mBaseUrl j = maybe
  (throwIO . UnknownResponse $ "Unable to read JSON response as an OllamaResponse. Try using ollamaChatJ instead to inspect the bare JSON") pure . decode . encode =<< ollamaChatJ mBaseUrl j

-- | Make a request to an OpenAI server resulting in a JSON @Value@
openaiV1ChatJ :: ToJSON j => BearerToken -> Maybe BaseUrl -> j -> IO Value
openaiV1ChatJ token mBaseUrl j = do
  manager' <- newTlsManagerWith $ tlsManagerSettings
    { managerResponseTimeout = responseTimeoutMicro defaultHttpTimeout }
  doRequest (clientOpenAIV1Chat token) manager' (fromMaybe OpenAI.defaultUrl mBaseUrl) $ toJSON j

-- | Make a request to an OpenAI server resulting in a high-level OpenAI response
openaiV1Chat :: ToJSON j => BearerToken -> Maybe BaseUrl -> j -> IO OpenAIResponse
openaiV1Chat token mBaseUrl j = maybe
  (throwIO . UnknownResponse $ "Unable to read JSON response as an OpenAIResponse. Try using openaiV1ChatJ instead to inspect the bare JSON") pure . decode . encode =<< openaiV1ChatJ token mBaseUrl j


doRequest :: (FromJSON f, ToJSON t) => (f -> ClientM Value) -> Manager -> BaseUrl -> t -> IO Value
doRequest client' manager' baseUrl bodyData = do
  bodyDoc <- maybe (throwIO PostMarshallingException) pure $ decode . encode $ bodyData
  eres <- runClientM (client' bodyDoc) (mkClientEnv manager' baseUrl)
  case eres of
    Right v@(Object _) -> pure v
    Left err -> throwIO err
    Right somethingElse -> throwIO . UnknownResponse . show $ somethingElse