{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module Data.Ollama.Generate
  ( -- * Generate Texts
    generate
  , defaultGenerateOps
  , GenerateOps (..)
  , GenerateResponse (..)
  ) where

import Data.Aeson
import Data.ByteString.Char8 qualified as BS
import Data.ByteString.Lazy.Char8 qualified as BSL
import Data.Maybe
import Data.Ollama.Common.Utils as CU
import Data.Text (Text)
import Data.Text qualified as T
import Data.Time (UTCTime)
import GHC.Int (Int64)
import Network.HTTP.Client

-- TODO: Add Options parameter
-- TODO: Add Context parameter

{- | 
  Input type for generate functions. This data type represents all possible configurations 
  that you can pass to the Ollama generate API.
  
  Example:
  
  > let ops = GenerateOps 
  >         { modelName = "llama3.2"
  >         , prompt = "What is the meaning of life?"
  >         , suffix = Nothing
  >         , images = Nothing
  >         , format = Just "text"
  >         , system = Nothing
  >         , template = Nothing
  >         , stream = Nothing
  >         , raw = Just False
  >         , keepAlive = Just "yes"
  >         }
-}
data GenerateOps = GenerateOps
  { GenerateOps -> Text
modelName :: Text
  -- ^ The name of the model to be used for generation.
  , GenerateOps -> Text
prompt :: Text
  -- ^ The prompt text that will be provided to the model for generating a response.
  , GenerateOps -> Maybe Text
suffix :: Maybe Text
  -- ^ An optional suffix to append to the generated text.
  , GenerateOps -> Maybe [Text]
images :: Maybe [Text]
  -- ^ Optional list of base64 encoded images to include with the request.
  , GenerateOps -> Maybe Text
format :: Maybe Text
  -- ^ An optional format specifier for the response.
  , GenerateOps -> Maybe Text
system :: Maybe Text
  -- ^ Optional system text that can be included in the generation context.
  , GenerateOps -> Maybe Text
template :: Maybe Text
  -- ^ An optional template to format the response.
  , GenerateOps -> Maybe (GenerateResponse -> IO (), IO ())
stream :: Maybe (GenerateResponse -> IO (), IO ())
  -- ^ An optional streaming function where the first function handles each chunk of response, and the second flushes the stream.
  , GenerateOps -> Maybe Bool
raw :: Maybe Bool
  -- ^ An optional flag to return the raw response.
  , GenerateOps -> Maybe Text
keepAlive :: Maybe Text
  -- ^ Optional text to specify keep-alive behavior.
  }

instance Show GenerateOps where
  show :: GenerateOps -> String
show GenerateOps {Maybe Bool
Maybe [Text]
Maybe (GenerateResponse -> IO (), IO ())
Maybe Text
Text
$sel:modelName:GenerateOps :: GenerateOps -> Text
$sel:prompt:GenerateOps :: GenerateOps -> Text
$sel:suffix:GenerateOps :: GenerateOps -> Maybe Text
$sel:images:GenerateOps :: GenerateOps -> Maybe [Text]
$sel:format:GenerateOps :: GenerateOps -> Maybe Text
$sel:system:GenerateOps :: GenerateOps -> Maybe Text
$sel:template:GenerateOps :: GenerateOps -> Maybe Text
$sel:stream:GenerateOps :: GenerateOps -> Maybe (GenerateResponse -> IO (), IO ())
$sel:raw:GenerateOps :: GenerateOps -> Maybe Bool
$sel:keepAlive:GenerateOps :: GenerateOps -> Maybe Text
modelName :: Text
prompt :: Text
suffix :: Maybe Text
images :: Maybe [Text]
format :: Maybe Text
system :: Maybe Text
template :: Maybe Text
stream :: Maybe (GenerateResponse -> IO (), IO ())
raw :: Maybe Bool
keepAlive :: Maybe Text
..} =
    String
"GenerateOps { "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"model : "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
modelName
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
", prompt : "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
prompt
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
", suffix : "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Maybe Text -> String
forall a. Show a => a -> String
show Maybe Text
suffix
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
", images : "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Maybe [Text] -> String
forall a. Show a => a -> String
show Maybe [Text]
images
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
", format : "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Maybe Text -> String
forall a. Show a => a -> String
show Maybe Text
format
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
", system : "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Maybe Text -> String
forall a. Show a => a -> String
show Maybe Text
system
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
", template : "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Maybe Text -> String
forall a. Show a => a -> String
show Maybe Text
template
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
", stream : "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"Stream functions"
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
", raw : "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Maybe Bool -> String
forall a. Show a => a -> String
show Maybe Bool
raw
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
", keepAlive : "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Maybe Text -> String
forall a. Show a => a -> String
show Maybe Text
keepAlive

instance Eq GenerateOps where
    == :: GenerateOps -> GenerateOps -> Bool
(==) GenerateOps
a GenerateOps
b = 
        GenerateOps -> Text
modelName GenerateOps
a Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Text
modelName GenerateOps
b Bool -> Bool -> Bool
&&
        GenerateOps -> Text
prompt GenerateOps
a Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Text
prompt GenerateOps
b Bool -> Bool -> Bool
&&
        GenerateOps -> Maybe Text
suffix GenerateOps
a Maybe Text -> Maybe Text -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Maybe Text
suffix GenerateOps
b Bool -> Bool -> Bool
&&
        GenerateOps -> Maybe [Text]
images GenerateOps
a Maybe [Text] -> Maybe [Text] -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Maybe [Text]
images GenerateOps
b Bool -> Bool -> Bool
&&
        GenerateOps -> Maybe Text
format GenerateOps
a Maybe Text -> Maybe Text -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Maybe Text
format GenerateOps
b Bool -> Bool -> Bool
&&
        GenerateOps -> Maybe Text
system GenerateOps
a Maybe Text -> Maybe Text -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Maybe Text
system GenerateOps
b Bool -> Bool -> Bool
&&
        GenerateOps -> Maybe Text
template GenerateOps
a Maybe Text -> Maybe Text -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Maybe Text
template GenerateOps
b Bool -> Bool -> Bool
&&
        GenerateOps -> Maybe Bool
raw GenerateOps
a Maybe Bool -> Maybe Bool -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Maybe Bool
raw GenerateOps
b Bool -> Bool -> Bool
&&
        GenerateOps -> Maybe Text
keepAlive GenerateOps
a Maybe Text -> Maybe Text -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Maybe Text
keepAlive GenerateOps
b

-- TODO: Add Context Param

-- | 
-- Result type for generate function containing the model's response and meta-information.
data GenerateResponse = GenerateResponse
  { GenerateResponse -> Text
model :: Text
  -- ^ The name of the model that generated the response.
  , GenerateResponse -> UTCTime
createdAt :: UTCTime
  -- ^ The timestamp when the response was created.
  , GenerateResponse -> Text
response_ :: Text
  -- ^ The generated response from the model.
  , GenerateResponse -> Bool
done :: Bool
  -- ^ A flag indicating whether the generation process is complete.
  , GenerateResponse -> Maybe Int64
totalDuration :: Maybe Int64
  -- ^ Optional total duration in milliseconds for the generation process.
  , GenerateResponse -> Maybe Int64
loadDuration :: Maybe Int64
  -- ^ Optional load duration in milliseconds for loading the model.
  , GenerateResponse -> Maybe Int64
promptEvalCount :: Maybe Int64
  -- ^ Optional count of prompt evaluations during the generation process.
  , GenerateResponse -> Maybe Int64
promptEvalDuration :: Maybe Int64
  -- ^ Optional duration in milliseconds for evaluating the prompt.
  , GenerateResponse -> Maybe Int64
evalCount :: Maybe Int64
  -- ^ Optional count of evaluations during the generation process.
  , GenerateResponse -> Maybe Int64
evalDuration :: Maybe Int64
  -- ^ Optional duration in milliseconds for evaluations during the generation process.
  }
  deriving (Int -> GenerateResponse -> ShowS
[GenerateResponse] -> ShowS
GenerateResponse -> String
(Int -> GenerateResponse -> ShowS)
-> (GenerateResponse -> String)
-> ([GenerateResponse] -> ShowS)
-> Show GenerateResponse
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> GenerateResponse -> ShowS
showsPrec :: Int -> GenerateResponse -> ShowS
$cshow :: GenerateResponse -> String
show :: GenerateResponse -> String
$cshowList :: [GenerateResponse] -> ShowS
showList :: [GenerateResponse] -> ShowS
Show, GenerateResponse -> GenerateResponse -> Bool
(GenerateResponse -> GenerateResponse -> Bool)
-> (GenerateResponse -> GenerateResponse -> Bool)
-> Eq GenerateResponse
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: GenerateResponse -> GenerateResponse -> Bool
== :: GenerateResponse -> GenerateResponse -> Bool
$c/= :: GenerateResponse -> GenerateResponse -> Bool
/= :: GenerateResponse -> GenerateResponse -> Bool
Eq)

instance ToJSON GenerateOps where
  toJSON :: GenerateOps -> Value
toJSON
    ( GenerateOps
        Text
model
        Text
prompt
        Maybe Text
suffix
        Maybe [Text]
images
        Maybe Text
format
        Maybe Text
system
        Maybe Text
template
        Maybe (GenerateResponse -> IO (), IO ())
stream
        Maybe Bool
raw
        Maybe Text
keepAlive
      ) =
      [Pair] -> Value
object
        [ Key
"model" Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Text
model
        , Key
"prompt" Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Text
prompt
        , Key
"suffix" Key -> Maybe Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Text
suffix
        , Key
"images" Key -> Maybe [Text] -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe [Text]
images
        , Key
"format" Key -> Maybe Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Text
format
        , Key
"system" Key -> Maybe Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Text
system
        , Key
"template" Key -> Maybe Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Text
template
        , Key
"stream" Key -> Maybe Bool -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= if Maybe (GenerateResponse -> IO (), IO ()) -> Bool
forall a. Maybe a -> Bool
isNothing Maybe (GenerateResponse -> IO (), IO ())
stream then Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
False else Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
True
        , Key
"raw" Key -> Maybe Bool -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Bool
raw
        , Key
"keep_alive" Key -> Maybe Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Text
keepAlive
        ]

instance FromJSON GenerateResponse where
  parseJSON :: Value -> Parser GenerateResponse
parseJSON = String
-> (Object -> Parser GenerateResponse)
-> Value
-> Parser GenerateResponse
forall a. String -> (Object -> Parser a) -> Value -> Parser a
withObject String
"GenerateResponse" ((Object -> Parser GenerateResponse)
 -> Value -> Parser GenerateResponse)
-> (Object -> Parser GenerateResponse)
-> Value
-> Parser GenerateResponse
forall a b. (a -> b) -> a -> b
$ \Object
v ->
    Text
-> UTCTime
-> Text
-> Bool
-> Maybe Int64
-> Maybe Int64
-> Maybe Int64
-> Maybe Int64
-> Maybe Int64
-> Maybe Int64
-> GenerateResponse
GenerateResponse
      (Text
 -> UTCTime
 -> Text
 -> Bool
 -> Maybe Int64
 -> Maybe Int64
 -> Maybe Int64
 -> Maybe Int64
 -> Maybe Int64
 -> Maybe Int64
 -> GenerateResponse)
-> Parser Text
-> Parser
     (UTCTime
      -> Text
      -> Bool
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> GenerateResponse)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
v Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"model"
      Parser
  (UTCTime
   -> Text
   -> Bool
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> GenerateResponse)
-> Parser UTCTime
-> Parser
     (Text
      -> Bool
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> GenerateResponse)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser UTCTime
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"created_at"
      Parser
  (Text
   -> Bool
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> GenerateResponse)
-> Parser Text
-> Parser
     (Bool
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> GenerateResponse)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"response"
      Parser
  (Bool
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> GenerateResponse)
-> Parser Bool
-> Parser
     (Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> GenerateResponse)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser Bool
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"done"
      Parser
  (Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> GenerateResponse)
-> Parser (Maybe Int64)
-> Parser
     (Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> GenerateResponse)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser (Maybe Int64)
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"total_duration"
      Parser
  (Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> GenerateResponse)
-> Parser (Maybe Int64)
-> Parser
     (Maybe Int64
      -> Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> GenerateResponse)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser (Maybe Int64)
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"load_duration"
      Parser
  (Maybe Int64
   -> Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> GenerateResponse)
-> Parser (Maybe Int64)
-> Parser
     (Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> GenerateResponse)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser (Maybe Int64)
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"prompt_eval_count"
      Parser
  (Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> GenerateResponse)
-> Parser (Maybe Int64)
-> Parser (Maybe Int64 -> Maybe Int64 -> GenerateResponse)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser (Maybe Int64)
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"prompt_eval_duration"
      Parser (Maybe Int64 -> Maybe Int64 -> GenerateResponse)
-> Parser (Maybe Int64) -> Parser (Maybe Int64 -> GenerateResponse)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser (Maybe Int64)
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"eval_count"
      Parser (Maybe Int64 -> GenerateResponse)
-> Parser (Maybe Int64) -> Parser GenerateResponse
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser (Maybe Int64)
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"eval_duration"

-- | 
-- A function to create a default 'GenerateOps' type with preset values.
-- 
-- Example:
--
-- > let ops = defaultGenerateOps
-- > generate ops
-- 
-- This will generate a response using the default configuration.
defaultGenerateOps :: GenerateOps
defaultGenerateOps :: GenerateOps
defaultGenerateOps =
  GenerateOps
    { $sel:modelName:GenerateOps :: Text
modelName = Text
"llama3.2"
    , $sel:prompt:GenerateOps :: Text
prompt = Text
"what is 2+2"
    , $sel:suffix:GenerateOps :: Maybe Text
suffix = Maybe Text
forall a. Maybe a
Nothing
    , $sel:images:GenerateOps :: Maybe [Text]
images = Maybe [Text]
forall a. Maybe a
Nothing
    , $sel:format:GenerateOps :: Maybe Text
format = Maybe Text
forall a. Maybe a
Nothing
    , $sel:system:GenerateOps :: Maybe Text
system = Maybe Text
forall a. Maybe a
Nothing
    , $sel:template:GenerateOps :: Maybe Text
template = Maybe Text
forall a. Maybe a
Nothing
    , $sel:stream:GenerateOps :: Maybe (GenerateResponse -> IO (), IO ())
stream = Maybe (GenerateResponse -> IO (), IO ())
forall a. Maybe a
Nothing
    , $sel:raw:GenerateOps :: Maybe Bool
raw = Maybe Bool
forall a. Maybe a
Nothing
    , $sel:keepAlive:GenerateOps :: Maybe Text
keepAlive = Maybe Text
forall a. Maybe a
Nothing
    }

-- | 
-- Generate function that returns either a 'GenerateResponse' type or an error message.
-- It takes a 'GenerateOps' configuration and performs a request to the Ollama generate API.
-- 
-- Examples:
--
-- Basic usage without streaming:
--
-- > let ops = GenerateOps 
-- >         { modelName = "llama3.2"
-- >         , prompt = "Tell me a joke."
-- >         , suffix = Nothing
-- >         , images = Nothing
-- >         , format = Nothing
-- >         , system = Nothing
-- >         , template = Nothing
-- >         , stream = Nothing
-- >         , raw = Nothing
-- >         , keepAlive = Nothing
-- >         }
-- > result <- generate ops
-- > case result of
-- >   Left errorMsg -> putStrLn ("Error: " ++ errorMsg)
-- >   Right response -> print response
--
-- Usage with streaming to print responses to the console:
--
-- > void $
-- >   generate
-- >     defaultGenerateOps
-- >       { modelName = "llama3.2"
-- >       , prompt = "what is functional programming?"
-- >       , stream = Just (T.putStr . response_, pure ())
-- >       }
--
-- In this example, the first function in the 'stream' tuple processes each chunk of response by printing it,
-- and the second function is a simple no-op flush.generate :: GenerateOps -> IO (Either String GenerateResponse)
generate :: GenerateOps -> IO (Either String GenerateResponse)
generate :: GenerateOps -> IO (Either String GenerateResponse)
generate GenerateOps
genOps = do
  let url :: Text
url = OllamaClient -> Text
CU.host OllamaClient
defaultOllama
  Manager
manager <-
    ManagerSettings -> IO Manager
newManager -- Setting response timeout to 5 minutes, since llm takes time
      ManagerSettings
defaultManagerSettings {managerResponseTimeout = responseTimeoutMicro (5 * 60 * 1000000)}
  Request
initialRequest <- String -> IO Request
forall (m :: * -> *). MonadThrow m => String -> m Request
parseRequest (String -> IO Request) -> String -> IO Request
forall a b. (a -> b) -> a -> b
$ Text -> String
T.unpack (Text
url Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"/api/generate")
  let reqBody :: GenerateOps
reqBody = GenerateOps
genOps
      request :: Request
request =
        Request
initialRequest
          { method = "POST"
          , requestBody = RequestBodyLBS $ encode reqBody
          }
  Request
-> Manager
-> (Response BodyReader -> IO (Either String GenerateResponse))
-> IO (Either String GenerateResponse)
forall a.
Request -> Manager -> (Response BodyReader -> IO a) -> IO a
withResponse Request
request Manager
manager ((Response BodyReader -> IO (Either String GenerateResponse))
 -> IO (Either String GenerateResponse))
-> (Response BodyReader -> IO (Either String GenerateResponse))
-> IO (Either String GenerateResponse)
forall a b. (a -> b) -> a -> b
$ \Response BodyReader
response -> do
    let streamResponse :: (GenerateResponse -> IO a) -> IO a -> IO (Either String b)
streamResponse GenerateResponse -> IO a
sendChunk IO a
flush = do
          ByteString
bs <- BodyReader -> BodyReader
brRead (BodyReader -> BodyReader) -> BodyReader -> BodyReader
forall a b. (a -> b) -> a -> b
$ Response BodyReader -> BodyReader
forall body. Response body -> body
responseBody Response BodyReader
response
          if ByteString -> Bool
BS.null ByteString
bs
            then String -> IO ()
putStrLn String
"" IO () -> IO (Either String b) -> IO (Either String b)
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Either String b -> IO (Either String b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String b
forall a b. a -> Either a b
Left String
"")
            else do
              let eRes :: Either String GenerateResponse
eRes = ByteString -> Either String GenerateResponse
forall a. FromJSON a => ByteString -> Either String a
eitherDecode (ByteString -> ByteString
BSL.fromStrict ByteString
bs) :: Either String GenerateResponse
              case Either String GenerateResponse
eRes of
                Left String
e -> Either String b -> IO (Either String b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String b
forall a b. a -> Either a b
Left String
e)
                Right GenerateResponse
r -> do
                  a
_ <- GenerateResponse -> IO a
sendChunk GenerateResponse
r
                  a
_ <- IO a
flush
                  if GenerateResponse -> Bool
done GenerateResponse
r then Either String b -> IO (Either String b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String b
forall a b. a -> Either a b
Left String
"") else (GenerateResponse -> IO a) -> IO a -> IO (Either String b)
streamResponse GenerateResponse -> IO a
sendChunk IO a
flush
    let genResponse :: ByteString -> IO (Either String GenerateResponse)
genResponse ByteString
op = do
          ByteString
bs <- BodyReader -> BodyReader
brRead (BodyReader -> BodyReader) -> BodyReader -> BodyReader
forall a b. (a -> b) -> a -> b
$ Response BodyReader -> BodyReader
forall body. Response body -> body
responseBody Response BodyReader
response
          if ByteString
bs ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
""
            then do
              let eRes0 :: Either String GenerateResponse
eRes0 = ByteString -> Either String GenerateResponse
forall a. FromJSON a => ByteString -> Either String a
eitherDecode (ByteString -> ByteString
BSL.fromStrict ByteString
op) :: Either String GenerateResponse
              case Either String GenerateResponse
eRes0 of
                Left String
e -> Either String GenerateResponse
-> IO (Either String GenerateResponse)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String GenerateResponse
forall a b. a -> Either a b
Left String
e)
                Right GenerateResponse
r -> Either String GenerateResponse
-> IO (Either String GenerateResponse)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (GenerateResponse -> Either String GenerateResponse
forall a b. b -> Either a b
Right GenerateResponse
r)
            else ByteString -> IO (Either String GenerateResponse)
genResponse (ByteString
op ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
bs)
    case GenerateOps -> Maybe (GenerateResponse -> IO (), IO ())
stream GenerateOps
genOps of
      Maybe (GenerateResponse -> IO (), IO ())
Nothing -> ByteString -> IO (Either String GenerateResponse)
genResponse ByteString
""
      Just (GenerateResponse -> IO ()
sendChunk, IO ()
flush) -> (GenerateResponse -> IO ())
-> IO () -> IO (Either String GenerateResponse)
forall {a} {a} {b}.
(GenerateResponse -> IO a) -> IO a -> IO (Either String b)
streamResponse GenerateResponse -> IO ()
sendChunk IO ()
flush