HTTP.hs

{-# LANGUAGE DataKinds, TypeOperators #-}

module GenAILib.HTTP
  ( display
  , doChat
  , doCompletion
  )
  where

import Data.Aeson (FromJSON, ToJSON, Value (Object, String), decode, encode)
import Data.Aeson.KeyMap qualified as A
import Data.ByteString.Lazy qualified as BL
import Data.Maybe (fromMaybe)
import Data.Proxy (Proxy (..))
import Data.String.Conv (toS)
import Data.Text.IO qualified as TS
import Formatting ((%+), formatToString, string)
import Network.HTTP.Client (defaultManagerSettings, managerResponseTimeout,
  newManager, responseTimeoutMicro)
import Servant.API ((:<|>) (..), (:>), JSON, Post, ReqBody)
import Servant.Client (BaseUrl (..), Scheme (Http), ClientM,
  mkClientEnv, runClientM)
import Servant.Client.Core (clientIn)
import System.Exit (exitFailure)

import GenAILib.Common (RawOutput (..))
import GenAILib.Ollama (Host (..))
import GenAILib.System.Log (debugM, emergencyM, errorM, lname)


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


{- 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 API
  =    "api" :> "generate" :> ReqBody '[JSON] Value :> Post '[JSON] Value
  :<|> "api" :> "chat" :> ReqBody '[JSON] Value :> Post '[JSON] Value

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

clientM :: Proxy ClientM
clientM = Proxy

completion :: Value -> ClientM Value
chat :: Value -> ClientM Value
(completion :<|> chat) = api `clientIn` clientM


doCompletion :: ToJSON j => Host -> j -> IO Value
doCompletion = doRequest completion

doChat :: ToJSON j => Host -> j -> IO Value
doChat = doRequest chat

doRequest :: (FromJSON f, ToJSON t) => (f -> ClientM Value) -> Host -> t -> IO Value
doRequest client' (Host hostName port) bodyData = do
  debugM lname . toS . encode $ bodyData
  let
    bodyDoc = fromMaybe (error "Something went wrong with JSON-ifying the POST data")
      . decode . encode $ bodyData
  manager' <- newManager $ defaultManagerSettings
    { managerResponseTimeout = responseTimeoutMicro defaultHttpTimeout }
  eres <- runClientM (client' bodyDoc)
    (mkClientEnv manager' (BaseUrl Http (toS hostName) port ""))
  case eres of
    Left err -> logAndExit $ show err
    Right v@(Object _) -> pure v
    Right somethingElse -> logAndExit $ show somethingElse


logAndExit :: String -> IO a
logAndExit msg = do
  errorM lname $ formatToString ("Can't continue:" %+ string) msg
  exitFailure


display :: RawOutput -> Value -> IO ()

display (RawOutput True) v = BL.putStr $ encode v <> "\n"

display (RawOutput False) (Object keyMap) = case A.lookup "response" keyMap of
  Just (String t) -> TS.putStrLn t
  Just somethingElse -> emergencyM lname $ show somethingElse
  Nothing -> emergencyM lname $ show keyMap

display _ somethingElse = emergencyM lname $ show somethingElse