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

{- |
Module      : Data.Ollama.Generate
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability   : experimental
Description : Text generation functionality for the Ollama client.

This module provides functions and types for generating text using an Ollama model. It includes APIs
for sending generation requests, both in IO ('generate') and monadic ('generateM') contexts, with
support for streaming and non-streaming responses. The 'GenerateOps' type configures the generation
request, allowing customization of the model, prompt, images, format, and other parameters. The
'defaultGenerateOps' provides a convenient starting point for configuration.

The module supports advanced features like Base64-encoded images, custom templates, and model-specific
options (e.g., temperature). It also includes validation to ensure required fields are non-empty.

Example:

>>> let ops = defaultGenerateOps { modelName = "gemma3", prompt = "Write a poem." }
>>> generate ops Nothing
Right (GenerateResponse ...)
-}
module Data.Ollama.Generate
  ( -- * Generate Texts
    generate
  , generateM

    -- * Configuration
  , defaultGenerateOps
  , GenerateOps (..)
  , validateGenerateOps

    -- * Response and Configuration Types
  , GenerateResponse (..)
  , Format (..)
  , OllamaConfig (..)
  , defaultOllamaConfig
  , ModelOptions (..)
  , defaultModelOptions

    -- * Error Types
  , OllamaError (..)
  ) where

import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.Aeson
import Data.Maybe
import Data.Ollama.Common.Config (OllamaConfig (..), defaultOllamaConfig)
import Data.Ollama.Common.Error (OllamaError (..))
import Data.Ollama.Common.Types (Format (..), GenerateResponse (..), ModelOptions (..))
import Data.Ollama.Common.Utils as CU
import Data.Text (Text)
import Data.Text qualified as T

{- | Validates 'GenerateOps' to ensure required fields are non-empty.

Checks that the 'modelName' and 'prompt' fields are not empty. Returns 'Right' with the validated
'GenerateOps' or 'Left' with an 'OllamaError' if validation fails.

Example:

>>> validateGenerateOps defaultGenerateOps
Left (InvalidRequest "Prompt cannot be empty")
--
-- @since 0.2.0.0
-}
validateGenerateOps :: GenerateOps -> Either OllamaError GenerateOps
validateGenerateOps :: GenerateOps -> Either OllamaError GenerateOps
validateGenerateOps GenerateOps
ops
  | Text -> Bool
T.null (GenerateOps -> Text
modelName GenerateOps
ops) = OllamaError -> Either OllamaError GenerateOps
forall a b. a -> Either a b
Left (OllamaError -> Either OllamaError GenerateOps)
-> OllamaError -> Either OllamaError GenerateOps
forall a b. (a -> b) -> a -> b
$ String -> OllamaError
InvalidRequest String
"Model name cannot be empty"
  | Text -> Bool
T.null (GenerateOps -> Text
prompt GenerateOps
ops) = OllamaError -> Either OllamaError GenerateOps
forall a b. a -> Either a b
Left (OllamaError -> Either OllamaError GenerateOps)
-> OllamaError -> Either OllamaError GenerateOps
forall a b. (a -> b) -> a -> b
$ String -> OllamaError
InvalidRequest String
"Prompt cannot be empty"
  | Bool
otherwise = GenerateOps -> Either OllamaError GenerateOps
forall a b. b -> Either a b
Right GenerateOps
ops

-- | Configuration for a text generation request.
data GenerateOps = GenerateOps
  { GenerateOps -> Text
modelName :: !Text
  -- ^ The name of the model to use for generation (e.g., "gemma3").
  , GenerateOps -> Text
prompt :: !Text
  -- ^ The prompt text to provide to the model for generating a response.
  , GenerateOps -> Maybe Text
suffix :: Maybe Text
  -- ^ Optional suffix to append to the generated text (not supported by all models).
  , GenerateOps -> Maybe [Text]
images :: !(Maybe [Text])
  -- ^ Optional list of Base64-encoded images to include with the request.
  , GenerateOps -> Maybe Format
format :: !(Maybe Format)
  -- ^ Optional format specifier for the response (e.g., JSON).
  --
  -- @since 0.1.3.0
  , GenerateOps -> Maybe Text
system :: !(Maybe Text)
  -- ^ Optional system text to include in the generation context.
  , GenerateOps -> Maybe Text
template :: !(Maybe Text)
  -- ^ Optional template to format the response.
  , GenerateOps -> Maybe (GenerateResponse -> IO (), IO ())
stream :: !(Maybe (GenerateResponse -> IO (), IO ()))
  -- ^ Optional streaming functions: the first handles each response chunk, the second flushes the stream.
  , GenerateOps -> Maybe Bool
raw :: !(Maybe Bool)
  -- ^ Optional flag to return the raw response.
  , GenerateOps -> Maybe Int
keepAlive :: !(Maybe Int)
  -- ^ Optional override for how long (in minutes) the model stays loaded in memory (default: 5 minutes).
  , GenerateOps -> Maybe ModelOptions
options :: !(Maybe ModelOptions)
  -- ^ Optional model parameters (e.g., temperature) as specified in the Modelfile.
  --
  -- @since 0.1.3.0
  , GenerateOps -> Maybe Bool
think :: !(Maybe Bool)
  -- ^ Optional flag to enable thinking mode.
  --
  -- @since 0.2.0.0
  }

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

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 Int
keepAlive GenerateOps
a Maybe Int -> Maybe Int -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Maybe Int
keepAlive GenerateOps
b
      Bool -> Bool -> Bool
&& GenerateOps -> Maybe ModelOptions
options GenerateOps
a Maybe ModelOptions -> Maybe ModelOptions -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Maybe ModelOptions
options GenerateOps
b
      Bool -> Bool -> Bool
&& GenerateOps -> Maybe Bool
think GenerateOps
a Maybe Bool -> Maybe Bool -> Bool
forall a. Eq a => a -> a -> Bool
== GenerateOps -> Maybe Bool
think GenerateOps
b

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 Int
keepAlive
        Maybe ModelOptions
options
        Maybe Bool
think
      ) =
      [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 Int -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Int
keepAlive
        , Key
"options" Key -> Maybe ModelOptions -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe ModelOptions
options
        , Key
"think" 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
think
        ]

{- | Default configuration for text generation.

Provides a default 'GenerateOps' with the "gemma3" model and an empty prompt. Other fields are set
to 'Nothing' or default values. Can be customized by modifying fields as needed.

Example:

>>> let ops = defaultGenerateOps { modelName = "customModel", prompt = "Hello!" }
>>> generate ops Nothing
-}
defaultGenerateOps :: GenerateOps
defaultGenerateOps :: GenerateOps
defaultGenerateOps =
  GenerateOps
    { modelName :: Text
modelName = Text
"gemma3"
    , prompt :: Text
prompt = Text
""
    , suffix :: Maybe Text
suffix = Maybe Text
forall a. Maybe a
Nothing
    , images :: Maybe [Text]
images = Maybe [Text]
forall a. Maybe a
Nothing
    , format :: Maybe Format
format = Maybe Format
forall a. Maybe a
Nothing
    , system :: Maybe Text
system = Maybe Text
forall a. Maybe a
Nothing
    , template :: Maybe Text
template = Maybe Text
forall a. Maybe a
Nothing
    , stream :: Maybe (GenerateResponse -> IO (), IO ())
stream = Maybe (GenerateResponse -> IO (), IO ())
forall a. Maybe a
Nothing
    , raw :: Maybe Bool
raw = Maybe Bool
forall a. Maybe a
Nothing
    , keepAlive :: Maybe Int
keepAlive = Maybe Int
forall a. Maybe a
Nothing
    , options :: Maybe ModelOptions
options = Maybe ModelOptions
forall a. Maybe a
Nothing
    , think :: Maybe Bool
think = Maybe Bool
forall a. Maybe a
Nothing
    }

{- | Generates text using the specified model and configuration.

Validates the 'GenerateOps' configuration and sends a POST request to the "/api//generate" endpoint.
Supports both streaming and non-streaming responses based on the 'stream' field in 'GenerateOps'.
Returns 'Right' with a 'GenerateResponse' on success or 'Left' with an 'OllamaError' on failure.

Example:

>>> let ops = defaultGenerateOps { modelName = "gemma3", prompt = "Write a short poem." }
>>> generate ops Nothing
Right (GenerateResponse ...)
-}
generate :: GenerateOps -> Maybe OllamaConfig -> IO (Either OllamaError GenerateResponse)
generate :: GenerateOps
-> Maybe OllamaConfig -> IO (Either OllamaError GenerateResponse)
generate GenerateOps
ops Maybe OllamaConfig
mbConfig =
  case GenerateOps -> Either OllamaError GenerateOps
validateGenerateOps GenerateOps
ops of
    Left OllamaError
err -> Either OllamaError GenerateResponse
-> IO (Either OllamaError GenerateResponse)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either OllamaError GenerateResponse
 -> IO (Either OllamaError GenerateResponse))
-> Either OllamaError GenerateResponse
-> IO (Either OllamaError GenerateResponse)
forall a b. (a -> b) -> a -> b
$ OllamaError -> Either OllamaError GenerateResponse
forall a b. a -> Either a b
Left OllamaError
err
    Right GenerateOps
_ -> Text
-> ByteString
-> Maybe GenerateOps
-> Maybe OllamaConfig
-> (Response BodyReader
    -> IO (Either OllamaError GenerateResponse))
-> IO (Either OllamaError GenerateResponse)
forall payload response.
ToJSON payload =>
Text
-> ByteString
-> Maybe payload
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError response))
-> IO (Either OllamaError response)
withOllamaRequest Text
"/api//generate" ByteString
"POST" (GenerateOps -> Maybe GenerateOps
forall a. a -> Maybe a
Just GenerateOps
ops) Maybe OllamaConfig
mbConfig Response BodyReader -> IO (Either OllamaError GenerateResponse)
handler
  where
    handler :: Response BodyReader -> IO (Either OllamaError GenerateResponse)
handler = case GenerateOps -> Maybe (GenerateResponse -> IO (), IO ())
stream GenerateOps
ops of
      Maybe (GenerateResponse -> IO (), IO ())
Nothing -> Response BodyReader -> IO (Either OllamaError GenerateResponse)
forall a.
FromJSON a =>
Response BodyReader -> IO (Either OllamaError a)
commonNonStreamingHandler
      Just (GenerateResponse -> IO ()
sc, IO ()
fl) -> (GenerateResponse -> IO ())
-> IO ()
-> Response BodyReader
-> IO (Either OllamaError GenerateResponse)
forall a.
(HasDone a, FromJSON a) =>
(a -> IO ())
-> IO () -> Response BodyReader -> IO (Either OllamaError a)
commonStreamHandler GenerateResponse -> IO ()
sc IO ()
fl

{- | MonadIO version of 'generate' for use in monadic contexts.

Lifts the 'generate' function into a 'MonadIO' context, allowing it to be used in monadic computations.

Example:

>>> import Control.Monad.IO.Class
>>> let ops = defaultGenerateOps { modelName = "gemma3", prompt = "Hello!" }
>>> runReaderT (generateM ops Nothing) someContext
Right (GenerateResponse ...)
-}
generateM ::
  MonadIO m =>
  GenerateOps -> Maybe OllamaConfig -> m (Either OllamaError GenerateResponse)
generateM :: forall (m :: * -> *).
MonadIO m =>
GenerateOps
-> Maybe OllamaConfig -> m (Either OllamaError GenerateResponse)
generateM GenerateOps
ops Maybe OllamaConfig
mbCfg = IO (Either OllamaError GenerateResponse)
-> m (Either OllamaError GenerateResponse)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either OllamaError GenerateResponse)
 -> m (Either OllamaError GenerateResponse))
-> IO (Either OllamaError GenerateResponse)
-> m (Either OllamaError GenerateResponse)
forall a b. (a -> b) -> a -> b
$ GenerateOps
-> Maybe OllamaConfig -> IO (Either OllamaError GenerateResponse)
generate GenerateOps
ops Maybe OllamaConfig
mbCfg