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

{- |
Module      : Data.Ollama.Chat
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability   : experimental
Description : Chat functionality for interacting with the Ollama API.

This module provides functions and types for initiating and managing chat interactions with an Ollama model.
It includes APIs for sending chat requests, constructing messages with different roles, and configuring chat
operations. The module supports both streaming and non-streaming responses, as well as optional tools and
structured output formats.

The primary functions are 'chat' and 'chatM' for sending chat requests, and helper functions like
'systemMessage', 'userMessage', 'assistantMessage', and 'toolMessage' for constructing messages.
The 'ChatOps' type allows customization of chat parameters, and 'defaultChatOps' provides a convenient
starting point for configuration.

Example:

>>> let ops = defaultChatOps { chatModelName = "customModel", messages = userMessage "Hello!" :| [] }
>>> chat ops Nothing
Either OllamaError ChatResponse
-}
module Data.Ollama.Chat
  ( -- * Chat APIs
    chat
  , chatM

    -- * Message Types
  , Message (..)
  , Role (..)
  , systemMessage
  , userMessage
  , assistantMessage
  , toolMessage
  , genMessage

    -- * Chat Configuration
  , defaultChatOps
  , ChatOps (..)

    -- * Response Types
  , ChatResponse (..)
  , Format (..)

    -- * Configuration and Error Types
  , OllamaConfig (..)
  , defaultOllamaConfig
  , OllamaError (..)
  , ModelOptions (..)
  , defaultModelOptions

    -- * Tool and Function Types
  , InputTool (..)
  , FunctionDef (..)
  , FunctionParameters (..)
  , OutputFunction (..)
  , ToolCall (..)
  ) where

import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.Aeson
import Data.List.NonEmpty as NonEmpty
import Data.Maybe (isNothing)
import Data.Ollama.Common.Config
import Data.Ollama.Common.Error (OllamaError (..))
import Data.Ollama.Common.Types
  ( ChatResponse (..)
  , Format (..)
  , FunctionDef (..)
  , FunctionParameters (..)
  , InputTool (..)
  , Message (..)
  , ModelOptions (..)
  , OutputFunction (..)
  , Role (..)
  , ToolCall (..)
  )
import Data.Ollama.Common.Utils as CU
import Data.Text (Text)
import Data.Text qualified as T

{- | Constructs a 'Message' with the specified role and content.

Creates a 'Message' with the given 'Role' and textual content, setting optional fields
('images', 'tool_calls', 'thinking') to 'Nothing'.

Example:

>>> genMessage User "What's the weather like?"
Message {role = User, content = "What's the weather like?", images = Nothing, tool_calls = Nothing, thinking = Nothing}
-}
genMessage :: Role -> Text -> Message
genMessage :: Role -> Text -> Message
genMessage Role
r Text
c =
  Message
    { role :: Role
role = Role
r
    , content :: Text
content = Text
c
    , images :: Maybe [Text]
images = Maybe [Text]
forall a. Maybe a
Nothing
    , tool_calls :: Maybe [ToolCall]
tool_calls = Maybe [ToolCall]
forall a. Maybe a
Nothing
    , thinking :: Maybe Text
thinking = Maybe Text
forall a. Maybe a
Nothing
    }

{- | Creates a 'Message' with the 'System' role.

Example:

>>> systemMessage "You are a helpful assistant."
Message {role = System, content = "You are a helpful assistant.", images = Nothing, tool_calls = Nothing, thinking = Nothing}
-}
systemMessage :: Text -> Message
systemMessage :: Text -> Message
systemMessage Text
c = Role -> Text -> Message
genMessage Role
System Text
c

{- | Creates a 'Message' with the 'User' role.

Example:

>>> userMessage "What's 2+2?"
Message {role = User, content = "What's 2+2?", images = Nothing, tool_calls = Nothing, thinking = Nothing}
-}
userMessage :: Text -> Message
userMessage :: Text -> Message
userMessage Text
c = Role -> Text -> Message
genMessage Role
User Text
c

{- | Creates a 'Message' with the 'Assistant' role.

Example:

>>> assistantMessage "2+2 equals 4."
Message {role = Assistant, content = "2+2 equals 4.", images = Nothing, tool_calls = Nothing, thinking = Nothing}
-}
assistantMessage :: Text -> Message
assistantMessage :: Text -> Message
assistantMessage Text
c = Role -> Text -> Message
genMessage Role
Assistant Text
c

{- | Creates a 'Message' with the 'Tool' role.

Example:

>>> toolMessage "Tool output: success"
Message {role = Tool, content = "Tool output: success", images = Nothing, tool_calls = Nothing, thinking = Nothing}
-}
toolMessage :: Text -> Message
toolMessage :: Text -> Message
toolMessage Text
c = Role -> Text -> Message
genMessage Role
Tool Text
c

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

Checks that the 'chatModelName' is not empty and that no 'Message' in 'messages' has empty content.
Returns 'Right' with the validated 'ChatOps' or 'Left' with an 'OllamaError' if validation fails.
--
-- @since 0.2.0.0
-}
validateChatOps :: ChatOps -> Either OllamaError ChatOps
validateChatOps :: ChatOps -> Either OllamaError ChatOps
validateChatOps ChatOps
ops
  | Text -> Bool
T.null (ChatOps -> Text
chatModelName ChatOps
ops) = OllamaError -> Either OllamaError ChatOps
forall a b. a -> Either a b
Left (OllamaError -> Either OllamaError ChatOps)
-> OllamaError -> Either OllamaError ChatOps
forall a b. (a -> b) -> a -> b
$ String -> OllamaError
InvalidRequest String
"Chat model name cannot be empty"
  | (Message -> Bool) -> NonEmpty Message -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Text -> Bool
T.null (Text -> Bool) -> (Message -> Text) -> Message -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Message -> Text
content) (ChatOps -> NonEmpty Message
messages ChatOps
ops) =
      OllamaError -> Either OllamaError ChatOps
forall a b. a -> Either a b
Left (OllamaError -> Either OllamaError ChatOps)
-> OllamaError -> Either OllamaError ChatOps
forall a b. (a -> b) -> a -> b
$ String -> OllamaError
InvalidRequest String
"Messages cannot have empty content"
  | Bool
otherwise = ChatOps -> Either OllamaError ChatOps
forall a b. b -> Either a b
Right ChatOps
ops

{- | Configuration for initiating a chat with an Ollama model.

Defines the parameters for a chat request, including the model name, messages, and optional settings
for tools, response format, streaming, timeout, and model options.
-}
data ChatOps = ChatOps
  { ChatOps -> Text
chatModelName :: !Text
  -- ^ The name of the chat model to be used (e.g., "gemma3").
  , ChatOps -> NonEmpty Message
messages :: !(NonEmpty Message)
  -- ^ A non-empty list of messages forming the conversation context.
  , ChatOps -> Maybe [InputTool]
tools :: !(Maybe [InputTool])
  -- ^ Optional tools that may be used in the chat.
  , ChatOps -> Maybe Format
format :: !(Maybe Format)
  -- ^ Optional format for the chat response (e.g., JSON or JSON schema).
  --
  -- @since 0.1.3.0
  , ChatOps -> Maybe (ChatResponse -> IO (), IO ())
stream :: !(Maybe (ChatResponse -> IO (), IO ()))
  -- ^ Optional streaming functions: the first handles each response chunk, the second flushes the stream.
  , ChatOps -> Maybe Int
keepAlive :: !(Maybe Int)
  -- ^ Optional override for the response timeout in minutes (default: 15 minutes).
  , ChatOps -> Maybe ModelOptions
options :: !(Maybe ModelOptions)
  -- ^ Optional model parameters (e.g., temperature) as specified in the Modelfile.
  --
  -- @since 0.1.3.0
  , ChatOps -> Maybe Bool
think :: !(Maybe Bool)
  -- ^ Optional flag to enable thinking mode.
  --
  -- @since 0.2.0.0
  }

instance Show ChatOps where
  show :: ChatOps -> String
show
    ( ChatOps
        { chatModelName :: ChatOps -> Text
chatModelName = Text
m
        , messages :: ChatOps -> NonEmpty Message
messages = NonEmpty Message
ms
        , tools :: ChatOps -> Maybe [InputTool]
tools = Maybe [InputTool]
t
        , format :: ChatOps -> Maybe Format
format = Maybe Format
f
        , keepAlive :: ChatOps -> Maybe Int
keepAlive = Maybe Int
ka
        , think :: ChatOps -> Maybe Bool
think = Maybe Bool
th
        }
      ) =
      let messagesStr :: String
messagesStr = [Message] -> String
forall a. Show a => a -> String
show (NonEmpty Message -> [Message]
forall a. NonEmpty a -> [a]
toList NonEmpty Message
ms)
          toolsStr :: String
toolsStr = Maybe [InputTool] -> String
forall a. Show a => a -> String
show Maybe [InputTool]
t
          formatStr :: String
formatStr = Maybe Format -> String
forall a. Show a => a -> String
show Maybe Format
f
          keepAliveStr :: String
keepAliveStr = Maybe Int -> String
forall a. Show a => a -> String
show Maybe Int
ka
          thinkStr :: String
thinkStr = Maybe Bool -> String
forall a. Show a => a -> String
show Maybe Bool
th
       in Text -> String
T.unpack Text
m
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\nMessages:\n"
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
messagesStr
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\n"
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
toolsStr
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\n"
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
formatStr
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\n"
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
keepAliveStr
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\n"
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
thinkStr

instance Eq ChatOps where
  == :: ChatOps -> ChatOps -> Bool
(==) ChatOps
a ChatOps
b =
    ChatOps -> Text
chatModelName ChatOps
a Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== ChatOps -> Text
chatModelName ChatOps
b
      Bool -> Bool -> Bool
&& ChatOps -> NonEmpty Message
messages ChatOps
a NonEmpty Message -> NonEmpty Message -> Bool
forall a. Eq a => a -> a -> Bool
== ChatOps -> NonEmpty Message
messages ChatOps
b
      Bool -> Bool -> Bool
&& ChatOps -> Maybe [InputTool]
tools ChatOps
a Maybe [InputTool] -> Maybe [InputTool] -> Bool
forall a. Eq a => a -> a -> Bool
== ChatOps -> Maybe [InputTool]
tools ChatOps
b
      Bool -> Bool -> Bool
&& ChatOps -> Maybe Format
format ChatOps
a Maybe Format -> Maybe Format -> Bool
forall a. Eq a => a -> a -> Bool
== ChatOps -> Maybe Format
format ChatOps
b
      Bool -> Bool -> Bool
&& ChatOps -> Maybe Int
keepAlive ChatOps
a Maybe Int -> Maybe Int -> Bool
forall a. Eq a => a -> a -> Bool
== ChatOps -> Maybe Int
keepAlive ChatOps
b

instance ToJSON ChatOps where
  toJSON :: ChatOps -> Value
toJSON (ChatOps Text
model_ NonEmpty Message
messages_ Maybe [InputTool]
tools_ Maybe Format
format_ Maybe (ChatResponse -> IO (), IO ())
stream_ 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
"messages" Key -> NonEmpty Message -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= NonEmpty Message
messages_
      , Key
"tools" Key -> Maybe [InputTool] -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe [InputTool]
tools_
      , 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
"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 (ChatResponse -> IO (), IO ()) -> Bool
forall a. Maybe a -> Bool
isNothing Maybe (ChatResponse -> 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
"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 initiating a chat.

Provides a default 'ChatOps' with the "gemma3" model and a sample user message ("What is 2+2?").
Can be customized by modifying fields as needed.

Example:

>>> let ops = defaultChatOps { chatModelName = "customModel", messages = userMessage "Hello!" :| [] }
>>> chat ops Nothing
Either OllamaError ChatResponse
-}
defaultChatOps :: ChatOps
defaultChatOps :: ChatOps
defaultChatOps =
  ChatOps
    { chatModelName :: Text
chatModelName = Text
"gemma3"
    , messages :: NonEmpty Message
messages = Text -> Message
userMessage Text
"What is 2+2?" Message -> [Message] -> NonEmpty Message
forall a. a -> [a] -> NonEmpty a
:| []
    , tools :: Maybe [InputTool]
tools = Maybe [InputTool]
forall a. Maybe a
Nothing
    , format :: Maybe Format
format = Maybe Format
forall a. Maybe a
Nothing
    , stream :: Maybe (ChatResponse -> IO (), IO ())
stream = Maybe (ChatResponse -> IO (), IO ())
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
    }

{- | Sends a chat request to the Ollama API.

Validates the 'ChatOps' configuration and sends a POST request to the "/api//chat" endpoint.
Supports both streaming and non-streaming responses based on the 'stream' field in 'ChatOps'.
Returns an 'Either' containing an 'OllamaError' on failure or a 'ChatResponse' on success.

Example:

>>> let ops = defaultChatOps { chatModelName = "gemma3", messages = userMessage "What's the capital of France?" :| [] }
>>> chat ops Nothing
Either OllamaError ChatResponse
-}
chat :: ChatOps -> Maybe OllamaConfig -> IO (Either OllamaError ChatResponse)
chat :: ChatOps
-> Maybe OllamaConfig -> IO (Either OllamaError ChatResponse)
chat ChatOps
ops Maybe OllamaConfig
mbConfig =
  case ChatOps -> Either OllamaError ChatOps
validateChatOps ChatOps
ops of
    Left OllamaError
err -> Either OllamaError ChatResponse
-> IO (Either OllamaError ChatResponse)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either OllamaError ChatResponse
 -> IO (Either OllamaError ChatResponse))
-> Either OllamaError ChatResponse
-> IO (Either OllamaError ChatResponse)
forall a b. (a -> b) -> a -> b
$ OllamaError -> Either OllamaError ChatResponse
forall a b. a -> Either a b
Left OllamaError
err
    Right ChatOps
_ -> Text
-> ByteString
-> Maybe ChatOps
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError ChatResponse))
-> IO (Either OllamaError ChatResponse)
forall payload response.
ToJSON payload =>
Text
-> ByteString
-> Maybe payload
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError response))
-> IO (Either OllamaError response)
withOllamaRequest Text
"/api//chat" ByteString
"POST" (ChatOps -> Maybe ChatOps
forall a. a -> Maybe a
Just ChatOps
ops) Maybe OllamaConfig
mbConfig Response BodyReader -> IO (Either OllamaError ChatResponse)
handler
  where
    handler :: Response BodyReader -> IO (Either OllamaError ChatResponse)
handler = case ChatOps -> Maybe (ChatResponse -> IO (), IO ())
stream ChatOps
ops of
      Maybe (ChatResponse -> IO (), IO ())
Nothing -> Response BodyReader -> IO (Either OllamaError ChatResponse)
forall a.
FromJSON a =>
Response BodyReader -> IO (Either OllamaError a)
commonNonStreamingHandler
      Just (ChatResponse -> IO ()
sc, IO ()
fl) -> (ChatResponse -> IO ())
-> IO ()
-> Response BodyReader
-> IO (Either OllamaError ChatResponse)
forall a.
(HasDone a, FromJSON a) =>
(a -> IO ())
-> IO () -> Response BodyReader -> IO (Either OllamaError a)
commonStreamHandler ChatResponse -> IO ()
sc IO ()
fl

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

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

Example:

>>> import Control.Monad.IO.Class
>>> let ops = defaultChatOps { chatModelName = "gemma3", messages = userMessage "Hello!" :| [] }
>>> runReaderT (chatM ops Nothing) someContext
Either OllamaError ChatResponse
-}
chatM :: MonadIO m => ChatOps -> Maybe OllamaConfig -> m (Either OllamaError ChatResponse)
chatM :: forall (m :: * -> *).
MonadIO m =>
ChatOps
-> Maybe OllamaConfig -> m (Either OllamaError ChatResponse)
chatM ChatOps
ops Maybe OllamaConfig
mbCfg = IO (Either OllamaError ChatResponse)
-> m (Either OllamaError ChatResponse)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either OllamaError ChatResponse)
 -> m (Either OllamaError ChatResponse))
-> IO (Either OllamaError ChatResponse)
-> m (Either OllamaError ChatResponse)
forall a b. (a -> b) -> a -> b
$ ChatOps
-> Maybe OllamaConfig -> IO (Either OllamaError ChatResponse)
chat ChatOps
ops Maybe OllamaConfig
mbCfg