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

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

import Control.Exception (try)
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.Ollama.Common.Types (Format(..))
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding 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 Format
format :: Maybe Format
  -- ^ An optional format specifier for the response.
  -- ^ Since 0.1.3.0
  , 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.
  , GenerateOps -> Maybe Text
hostUrl :: Maybe Text
  -- ^ Override default Ollama host url. Default url = "http://127.0.0.1:11434"
  , GenerateOps -> Maybe Int
responseTimeOut :: Maybe Int
  -- ^ Override default response timeout in minutes. Default = 15 minutes
  , GenerateOps -> Maybe Value
options :: Maybe Value
  -- ^ additional model parameters listed in the documentation for the Modelfile such as temperature
  -- ^ Since 0.1.3.0
  }

instance Show GenerateOps where
  show :: GenerateOps -> String
show GenerateOps {Maybe Bool
Maybe Int
Maybe [Text]
Maybe (GenerateResponse -> IO (), IO ())
Maybe Value
Maybe Text
Maybe Format
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 Format
$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
$sel:hostUrl:GenerateOps :: GenerateOps -> Maybe Text
$sel:responseTimeOut:GenerateOps :: GenerateOps -> Maybe Int
$sel:options:GenerateOps :: GenerateOps -> Maybe Value
modelName :: Text
prompt :: Text
suffix :: Maybe Text
images :: Maybe [Text]
format :: Maybe Format
system :: Maybe Text
template :: Maybe Text
stream :: Maybe (GenerateResponse -> IO (), IO ())
raw :: Maybe Bool
keepAlive :: Maybe Text
hostUrl :: Maybe Text
responseTimeOut :: Maybe Int
options :: Maybe Value
..} =
    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 Format -> String
forall a. Show a => a -> String
show Maybe Format
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
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
", options : "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Maybe Value -> String
forall a. Show a => a -> String
show Maybe Value
options

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 Format
format GenerateOps
a Maybe Format -> Maybe Format -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Maybe Format
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
      Bool -> Bool -> Bool
&& GenerateOps -> Maybe Value
options GenerateOps
a Maybe Value -> Maybe Value -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Maybe Value
options 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 Format
format
        Maybe Text
system
        Maybe Text
template
        Maybe (GenerateResponse -> IO (), IO ())
stream
        Maybe Bool
raw
        Maybe Text
keepAlive
        Maybe Text
_ -- Host url
        Maybe Int
_ -- Response timeout
        Maybe Value
options 
      ) =
      [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 Format -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Format
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
        , Key
"options" Key -> Maybe Value -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Value
options
        ]

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 Format
format = Maybe Format
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
    , $sel:hostUrl:GenerateOps :: Maybe Text
hostUrl = Maybe Text
forall a. Maybe a
Nothing
    , $sel:responseTimeOut:GenerateOps :: Maybe Int
responseTimeOut = Maybe Int
forall a. Maybe a
Nothing
    , $sel:options:GenerateOps :: Maybe Value
options = Maybe Value
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 = Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe Text
defaultOllamaUrl (GenerateOps -> Maybe Text
hostUrl GenerateOps
genOps)
      responseTimeout :: Int
responseTimeout = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
15 (GenerateOps -> Maybe Int
responseTimeOut GenerateOps
genOps)
  Manager
manager <-
    ManagerSettings -> IO Manager
newManager -- Setting response timeout to 5 minutes, since llm takes time
      ManagerSettings
defaultManagerSettings
        { managerResponseTimeout = responseTimeoutMicro (responseTimeout * 60 * 1000000)
        }
  Either HttpException Request
eInitialRequest <-
    IO Request -> IO (Either HttpException Request)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO Request -> IO (Either HttpException Request))
-> IO Request -> IO (Either HttpException Request)
forall a b. (a -> b) -> a -> b
$ 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") :: IO (Either HttpException Request)
  case Either HttpException Request
eInitialRequest of
    Left HttpException
e -> do
      Either String GenerateResponse
-> IO (Either String GenerateResponse)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String GenerateResponse
 -> IO (Either String GenerateResponse))
-> Either String GenerateResponse
-> IO (Either String GenerateResponse)
forall a b. (a -> b) -> a -> b
$ String -> Either String GenerateResponse
forall a b. a -> Either a b
Left (String -> Either String GenerateResponse)
-> String -> Either String GenerateResponse
forall a b. (a -> b) -> a -> b
$ HttpException -> String
forall a. Show a => a -> String
show HttpException
e
    Right Request
initialRequest -> do
      let reqBody :: GenerateOps
reqBody = GenerateOps
genOps
          request :: Request
request =
            Request
initialRequest
              { method = "POST"
              , requestBody = RequestBodyLBS $ encode reqBody
              }
      Either HttpException (Either String GenerateResponse)
eRes <-
        IO (Either String GenerateResponse)
-> IO (Either HttpException (Either String GenerateResponse))
forall e a. Exception e => IO a -> IO (Either e a)
try (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
$ GenerateOps
-> Response BodyReader -> IO (Either String GenerateResponse)
handleRequest GenerateOps
genOps) ::
          IO (Either HttpException (Either String GenerateResponse))
      case Either HttpException (Either String GenerateResponse)
eRes of
        Left HttpException
e -> do
          Either String GenerateResponse
-> IO (Either String GenerateResponse)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String GenerateResponse
 -> IO (Either String GenerateResponse))
-> Either String GenerateResponse
-> IO (Either String GenerateResponse)
forall a b. (a -> b) -> a -> b
$ String -> Either String GenerateResponse
forall a b. a -> Either a b
Left (String -> Either String GenerateResponse)
-> String -> Either String GenerateResponse
forall a b. (a -> b) -> a -> b
$ String
"HTTP error occured: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> HttpException -> String
forall a. Show a => a -> String
show HttpException
e
        Right Either String GenerateResponse
r -> Either String GenerateResponse
-> IO (Either String GenerateResponse)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Either String GenerateResponse
r

handleRequest :: GenerateOps -> Response BodyReader -> IO (Either String GenerateResponse)
handleRequest :: GenerateOps
-> Response BodyReader -> IO (Either String GenerateResponse)
handleRequest GenerateOps
genOps Response BodyReader
response = do
  let streamResponse :: (GenerateResponse -> IO a)
-> IO a -> IO (Either String GenerateResponse)
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 GenerateResponse)
-> IO (Either String GenerateResponse)
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> 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
"")
          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 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 -> do
                a
_ <- GenerateResponse -> IO a
sendChunk GenerateResponse
r
                a
_ <- IO a
flush
                if GenerateResponse -> Bool
done GenerateResponse
r then 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 (GenerateResponse -> IO a)
-> IO a -> IO (Either String GenerateResponse)
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}.
(GenerateResponse -> IO a)
-> IO a -> IO (Either String GenerateResponse)
streamResponse GenerateResponse -> IO ()
sendChunk IO ()
flush

{- |
 generateJson is a higher level function that takes generateOps (similar to generate) and also takes
 a Haskell type (that has To and From JSON instance) and returns the response in provided type.

 This function simply calls generate with extra prompt appended to it, telling LLM to return the
 response in certain JSON format and serializes the response. This function will be helpful when you
 want to use the LLM to do something programmatic.

 For Example:
  > let expectedJsonStrucutre = Example {
  >   sortedList = ["sorted List here"]
  > , wasListAlreadSorted = False
  > }
  > eRes2 <- generateJson
  >     defaultGenerateOps
  >      { modelName = "llama3.2"
  >     , prompt = "Sort given list: [4, 2 , 3, 67]. Also tell whether list was already sorted or not."
  >       }
  >     expectedJsonStrucutre
  >     Nothing
  > case eRes2 of
  >   Left e -> putStrLn e
  >   Right r -> print ("JSON response: " :: String, r)

Output:
  > ("JSON response: ",Example {sortedList = ["1","2","3","4"], wasListAlreadSorted = False})

Note: While Passing the type, construct the type that will help LLM understand the field better.
 For example, in the above example, the sortedList's value is written as "Sorted List here". This
 will help LLM understand context better.

 You can also provide number of retries in case the LLM field to return the response in correct JSON
 in first attempt.
-}
generateJson ::
  (ToJSON jsonResult, FromJSON jsonResult) =>
  GenerateOps ->
  -- | Haskell type that you want your result in
  jsonResult ->
  -- | Max retries
  Maybe Int ->
  IO (Either String jsonResult)
generateJson :: forall jsonResult.
(ToJSON jsonResult, FromJSON jsonResult) =>
GenerateOps
-> jsonResult -> Maybe Int -> IO (Either String jsonResult)
generateJson genOps :: GenerateOps
genOps@GenerateOps {Maybe Bool
Maybe Int
Maybe [Text]
Maybe (GenerateResponse -> IO (), IO ())
Maybe Value
Maybe Text
Maybe Format
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 Format
$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
$sel:hostUrl:GenerateOps :: GenerateOps -> Maybe Text
$sel:responseTimeOut:GenerateOps :: GenerateOps -> Maybe Int
$sel:options:GenerateOps :: GenerateOps -> Maybe Value
modelName :: Text
prompt :: Text
suffix :: Maybe Text
images :: Maybe [Text]
format :: Maybe Format
system :: Maybe Text
template :: Maybe Text
stream :: Maybe (GenerateResponse -> IO (), IO ())
raw :: Maybe Bool
keepAlive :: Maybe Text
hostUrl :: Maybe Text
responseTimeOut :: Maybe Int
options :: Maybe Value
..} jsonResult
jsonStructure Maybe Int
mMaxRetries = do
  let jsonHelperPrompt :: Text
jsonHelperPrompt =
        Text
"You are an AI that returns only JSON object. \n"
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"* Your output should be a JSON object that matches the following schema: \n"
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> ByteString -> Text
T.decodeUtf8 (ByteString -> ByteString
BSL.toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ jsonResult -> ByteString
forall a. ToJSON a => a -> ByteString
encode jsonResult
jsonStructure)
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
prompt
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n"
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"# How to treat the task:\n"
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"* Stricly follow the schema for the output.\n"
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"* Never return anything other than a JSON object.\n"
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"* Do not talk to the user.\n"
  Either String GenerateResponse
generatedResponse <- GenerateOps -> IO (Either String GenerateResponse)
generate GenerateOps
genOps {prompt = jsonHelperPrompt}
  case Either String GenerateResponse
generatedResponse of
    Left String
err -> Either String jsonResult -> IO (Either String jsonResult)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String jsonResult -> IO (Either String jsonResult))
-> Either String jsonResult -> IO (Either String jsonResult)
forall a b. (a -> b) -> a -> b
$ String -> Either String jsonResult
forall a b. a -> Either a b
Left String
err
    Right GenerateResponse
r -> do
      case ByteString -> Maybe jsonResult
forall a. FromJSON a => ByteString -> Maybe a
decode (ByteString -> ByteString
BSL.fromStrict (ByteString -> ByteString)
-> (Text -> ByteString) -> Text -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
T.encodeUtf8 (Text -> ByteString) -> Text -> ByteString
forall a b. (a -> b) -> a -> b
$ GenerateResponse -> Text
response_ GenerateResponse
r) of
        Maybe jsonResult
Nothing -> do
          case Maybe Int
mMaxRetries of
            Maybe Int
Nothing -> Either String jsonResult -> IO (Either String jsonResult)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String jsonResult -> IO (Either String jsonResult))
-> Either String jsonResult -> IO (Either String jsonResult)
forall a b. (a -> b) -> a -> b
$ String -> Either String jsonResult
forall a b. a -> Either a b
Left String
"Decoding Failed :("
            Just Int
n ->
              if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1
                then Either String jsonResult -> IO (Either String jsonResult)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String jsonResult -> IO (Either String jsonResult))
-> Either String jsonResult -> IO (Either String jsonResult)
forall a b. (a -> b) -> a -> b
$ String -> Either String jsonResult
forall a b. a -> Either a b
Left String
"Decoding failed :("
                else GenerateOps
-> jsonResult -> Maybe Int -> IO (Either String jsonResult)
forall jsonResult.
(ToJSON jsonResult, FromJSON jsonResult) =>
GenerateOps
-> jsonResult -> Maybe Int -> IO (Either String jsonResult)
generateJson GenerateOps
genOps jsonResult
jsonStructure (Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
        Just jsonResult
resultInType -> Either String jsonResult -> IO (Either String jsonResult)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String jsonResult -> IO (Either String jsonResult))
-> Either String jsonResult -> IO (Either String jsonResult)
forall a b. (a -> b) -> a -> b
$ jsonResult -> Either String jsonResult
forall a b. b -> Either a b
Right jsonResult
resultInType


{- |
   Example usage of 'Ollama.generate' with a JSON schema format and options field.

   In this example we pass a JSON schema that expects an object with an integer field @age@.
   The options field is supplied as a JSON value.

   >>> import Data.Aeson (Value, object, (.=))
   >>> import Ollama (GenerateOps, defaultGenerateOps, SchemaFormat)
   >>> let x :: Value
   ...     x = object [ "type" .= ("object" :: String)
   ...                , "properties" .= object [ "age" .= object ["type" .= ("integer" :: String)] ]
   ...                ]
   >>> let opts :: Value
   ...     opts = object ["option" .= ("some value" :: String)]
   >>> generate defaultGenerateOps
   ...   { modelName = "llama3.2"
   ...   , prompt = "Ollama is 22 years old and is busy saving the world. Respond using JSON"
   ...   , format = Just (SchemaFormat x)
   ...   , options = opts
   ...   }
   Right (GenerateResponse {model = "llama3.2", createdAt = 2025-03-25 09:34:15.853417157 UTC,
                            response_ = "{\n    \"age\": 22\n}", done = True,
                            totalDuration = Just 6625631744, loadDuration = Just 2578791966,
                            promptEvalCount = Just 43, promptEvalDuration = Just 2983000000,
                            evalCount = Just 10, evalDuration = Just 1061000000})
-}