{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module Data.Ollama.Chat
  ( -- * Chat APIs
    chat
  , chatJson
  , Message (..)
  , Role (..)
  , defaultChatOps
  , ChatOps (..)
  , ChatResponse (..)
  ) 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.List.NonEmpty as NonEmpty
import Data.Maybe (fromMaybe, isNothing)
import Data.Ollama.Common.Utils as CU
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding qualified as T
import Data.Time (UTCTime)
import GHC.Generics
import GHC.Int (Int64)
import Network.HTTP.Client

-- | Enumerated roles that can participate in a chat.
data Role = System | User | Assistant | Tool
  deriving (Int -> Role -> ShowS
[Role] -> ShowS
Role -> String
(Int -> Role -> ShowS)
-> (Role -> String) -> ([Role] -> ShowS) -> Show Role
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Role -> ShowS
showsPrec :: Int -> Role -> ShowS
$cshow :: Role -> String
show :: Role -> String
$cshowList :: [Role] -> ShowS
showList :: [Role] -> ShowS
Show, Role -> Role -> Bool
(Role -> Role -> Bool) -> (Role -> Role -> Bool) -> Eq Role
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Role -> Role -> Bool
== :: Role -> Role -> Bool
$c/= :: Role -> Role -> Bool
/= :: Role -> Role -> Bool
Eq)

instance ToJSON Role where
  toJSON :: Role -> Value
toJSON Role
System = Text -> Value
String Text
"system"
  toJSON Role
User = Text -> Value
String Text
"user"
  toJSON Role
Assistant = Text -> Value
String Text
"assistant"
  toJSON Role
Tool = Text -> Value
String Text
"tool"

instance FromJSON Role where
  parseJSON :: Value -> Parser Role
parseJSON (String Text
"system") = Role -> Parser Role
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Role
System
  parseJSON (String Text
"user") = Role -> Parser Role
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Role
User
  parseJSON (String Text
"assistant") = Role -> Parser Role
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Role
Assistant
  parseJSON (String Text
"tool") = Role -> Parser Role
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Role
Tool
  parseJSON Value
_ = String -> Parser Role
forall a. String -> Parser a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Invalid Role value"

-- TODO : Add tool_calls parameter

-- | Represents a message within a chat, including its role and content.
data Message = Message
  { Message -> Role
role :: Role
  -- ^ The role of the entity sending the message (e.g., 'User', 'Assistant').
  , Message -> Text
content :: Text
  -- ^ The textual content of the message.
  , Message -> Maybe [Text]
images :: Maybe [Text]
  -- ^ Optional list of base64 encoded images that accompany the message.
  }
  deriving (Int -> Message -> ShowS
[Message] -> ShowS
Message -> String
(Int -> Message -> ShowS)
-> (Message -> String) -> ([Message] -> ShowS) -> Show Message
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Message -> ShowS
showsPrec :: Int -> Message -> ShowS
$cshow :: Message -> String
show :: Message -> String
$cshowList :: [Message] -> ShowS
showList :: [Message] -> ShowS
Show, Message -> Message -> Bool
(Message -> Message -> Bool)
-> (Message -> Message -> Bool) -> Eq Message
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Message -> Message -> Bool
== :: Message -> Message -> Bool
$c/= :: Message -> Message -> Bool
/= :: Message -> Message -> Bool
Eq, (forall x. Message -> Rep Message x)
-> (forall x. Rep Message x -> Message) -> Generic Message
forall x. Rep Message x -> Message
forall x. Message -> Rep Message x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Message -> Rep Message x
from :: forall x. Message -> Rep Message x
$cto :: forall x. Rep Message x -> Message
to :: forall x. Rep Message x -> Message
Generic, [Message] -> Value
[Message] -> Encoding
Message -> Bool
Message -> Value
Message -> Encoding
(Message -> Value)
-> (Message -> Encoding)
-> ([Message] -> Value)
-> ([Message] -> Encoding)
-> (Message -> Bool)
-> ToJSON Message
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> (a -> Bool)
-> ToJSON a
$ctoJSON :: Message -> Value
toJSON :: Message -> Value
$ctoEncoding :: Message -> Encoding
toEncoding :: Message -> Encoding
$ctoJSONList :: [Message] -> Value
toJSONList :: [Message] -> Value
$ctoEncodingList :: [Message] -> Encoding
toEncodingList :: [Message] -> Encoding
$comitField :: Message -> Bool
omitField :: Message -> Bool
ToJSON, Maybe Message
Value -> Parser [Message]
Value -> Parser Message
(Value -> Parser Message)
-> (Value -> Parser [Message]) -> Maybe Message -> FromJSON Message
forall a.
(Value -> Parser a)
-> (Value -> Parser [a]) -> Maybe a -> FromJSON a
$cparseJSON :: Value -> Parser Message
parseJSON :: Value -> Parser Message
$cparseJSONList :: Value -> Parser [Message]
parseJSONList :: Value -> Parser [Message]
$comittedField :: Maybe Message
omittedField :: Maybe Message
FromJSON)

-- TODO: Add Options parameter
data ChatOps = ChatOps
  { ChatOps -> Text
chatModelName :: Text
  -- ^ The name of the chat model to be used.
  , ChatOps -> NonEmpty Message
messages :: NonEmpty Message
  -- ^ A non-empty list of messages forming the conversation context.
  , ChatOps -> Maybe Text
tools :: Maybe Text
  -- ^ Optional tools that may be used in the chat.
  , ChatOps -> Maybe Text
format :: Maybe Text
  -- ^ An optional format for the chat response.
  , ChatOps -> Maybe (ChatResponse -> IO (), IO ())
stream :: Maybe (ChatResponse -> IO (), IO ())
  -- ^ Optional streaming functions where the first handles each chunk of the response, and the second flushes the stream.
  , ChatOps -> Maybe Text
keepAlive :: Maybe Text
  -- ^ Optional text to specify keep-alive behavior.
  , ChatOps -> Maybe Text
hostUrl :: Maybe Text
  -- ^ Override default Ollama host url. Default url = "http://127.0.0.1:11434"
  , ChatOps -> Maybe Int
responseTimeOut :: Maybe Int
  -- ^ Override default response timeout in minutes. Default = 15 minutes
  }

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 Text
tools = Maybe Text
t, format :: ChatOps -> Maybe Text
format = Maybe Text
f, keepAlive :: ChatOps -> Maybe Text
keepAlive = Maybe Text
ka}) =
    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 Text -> String
forall a. Show a => a -> String
show Maybe Text
t
        formatStr :: String
formatStr = Maybe Text -> String
forall a. Show a => a -> String
show Maybe Text
f
        keepAliveStr :: String
keepAliveStr = Maybe Text -> String
forall a. Show a => a -> String
show Maybe Text
ka
     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

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 Text
tools ChatOps
a Maybe Text -> Maybe Text -> Bool
forall a. Eq a => a -> a -> Bool
== ChatOps -> Maybe Text
tools ChatOps
b
      Bool -> Bool -> Bool
&& ChatOps -> Maybe Text
format ChatOps
a Maybe Text -> Maybe Text -> Bool
forall a. Eq a => a -> a -> Bool
== ChatOps -> Maybe Text
format ChatOps
b
      Bool -> Bool -> Bool
&& ChatOps -> Maybe Text
keepAlive ChatOps
a Maybe Text -> Maybe Text -> Bool
forall a. Eq a => a -> a -> Bool
== ChatOps -> Maybe Text
keepAlive ChatOps
b

data ChatResponse = ChatResponse
  { ChatResponse -> Text
model :: Text
  -- ^ The name of the model that generated this response.
  , ChatResponse -> UTCTime
createdAt :: UTCTime
  -- ^ The timestamp when the response was created.
  , ChatResponse -> Maybe Message
message :: Maybe Message
  -- ^ The message content of the response, if any.
  , ChatResponse -> Bool
done :: Bool
  -- ^ Indicates whether the chat process has completed.
  , ChatResponse -> Maybe Int64
totalDuration :: Maybe Int64
  -- ^ Optional total duration in milliseconds for the chat process.
  , ChatResponse -> Maybe Int64
loadDuration :: Maybe Int64
  -- ^ Optional load duration in milliseconds for loading the model.
  , ChatResponse -> Maybe Int64
promptEvalCount :: Maybe Int64
  -- ^ Optional count of prompt evaluations during the chat process.
  , ChatResponse -> Maybe Int64
promptEvalDuration :: Maybe Int64
  -- ^ Optional duration in milliseconds for evaluating the prompt.
  , ChatResponse -> Maybe Int64
evalCount :: Maybe Int64
  -- ^ Optional count of evaluations during the chat process.
  , ChatResponse -> Maybe Int64
evalDuration :: Maybe Int64
  -- ^ Optional duration in milliseconds for evaluations during the chat process.
  }
  deriving (Int -> ChatResponse -> ShowS
[ChatResponse] -> ShowS
ChatResponse -> String
(Int -> ChatResponse -> ShowS)
-> (ChatResponse -> String)
-> ([ChatResponse] -> ShowS)
-> Show ChatResponse
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ChatResponse -> ShowS
showsPrec :: Int -> ChatResponse -> ShowS
$cshow :: ChatResponse -> String
show :: ChatResponse -> String
$cshowList :: [ChatResponse] -> ShowS
showList :: [ChatResponse] -> ShowS
Show, ChatResponse -> ChatResponse -> Bool
(ChatResponse -> ChatResponse -> Bool)
-> (ChatResponse -> ChatResponse -> Bool) -> Eq ChatResponse
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ChatResponse -> ChatResponse -> Bool
== :: ChatResponse -> ChatResponse -> Bool
$c/= :: ChatResponse -> ChatResponse -> Bool
/= :: ChatResponse -> ChatResponse -> Bool
Eq)

instance ToJSON ChatOps where
  toJSON :: ChatOps -> Value
toJSON (ChatOps Text
model_ NonEmpty Message
messages_ Maybe Text
tools_ Maybe Text
format_ Maybe (ChatResponse -> IO (), IO ())
stream_ Maybe Text
keepAlive_ Maybe Text
_ Maybe Int
_) =
    [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 Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Text
tools_
      , 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
"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 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 ChatResponse where
  parseJSON :: Value -> Parser ChatResponse
parseJSON = String
-> (Object -> Parser ChatResponse) -> Value -> Parser ChatResponse
forall a. String -> (Object -> Parser a) -> Value -> Parser a
withObject String
"ChatResponse" ((Object -> Parser ChatResponse) -> Value -> Parser ChatResponse)
-> (Object -> Parser ChatResponse) -> Value -> Parser ChatResponse
forall a b. (a -> b) -> a -> b
$ \Object
v ->
    Text
-> UTCTime
-> Maybe Message
-> Bool
-> Maybe Int64
-> Maybe Int64
-> Maybe Int64
-> Maybe Int64
-> Maybe Int64
-> Maybe Int64
-> ChatResponse
ChatResponse
      (Text
 -> UTCTime
 -> Maybe Message
 -> Bool
 -> Maybe Int64
 -> Maybe Int64
 -> Maybe Int64
 -> Maybe Int64
 -> Maybe Int64
 -> Maybe Int64
 -> ChatResponse)
-> Parser Text
-> Parser
     (UTCTime
      -> Maybe Message
      -> Bool
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> ChatResponse)
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
   -> Maybe Message
   -> Bool
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> ChatResponse)
-> Parser UTCTime
-> Parser
     (Maybe Message
      -> Bool
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> ChatResponse)
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
  (Maybe Message
   -> Bool
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> ChatResponse)
-> Parser (Maybe Message)
-> Parser
     (Bool
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> ChatResponse)
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 Message)
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"message"
      Parser
  (Bool
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> Maybe Int64
   -> ChatResponse)
-> Parser Bool
-> Parser
     (Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> ChatResponse)
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
   -> ChatResponse)
-> Parser (Maybe Int64)
-> Parser
     (Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> Maybe Int64
      -> ChatResponse)
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
   -> ChatResponse)
-> Parser (Maybe Int64)
-> Parser
     (Maybe Int64
      -> Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> ChatResponse)
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 -> ChatResponse)
-> Parser (Maybe Int64)
-> Parser
     (Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> ChatResponse)
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 -> ChatResponse)
-> Parser (Maybe Int64)
-> Parser (Maybe Int64 -> Maybe Int64 -> ChatResponse)
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 -> ChatResponse)
-> Parser (Maybe Int64) -> Parser (Maybe Int64 -> ChatResponse)
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 -> ChatResponse)
-> Parser (Maybe Int64) -> Parser ChatResponse
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 default configuration for initiating a chat with a model.
This can be used as a starting point and modified as needed.

Example:

> let ops = defaultChatOps { chatModelName = "customModel" }
> chat ops
-}
defaultChatOps :: ChatOps
defaultChatOps :: ChatOps
defaultChatOps =
  ChatOps
    { chatModelName :: Text
chatModelName = Text
"llama3.2"
    , messages :: NonEmpty Message
messages = Role -> Text -> Maybe [Text] -> Message
Message Role
User Text
"What is 2+2?" Maybe [Text]
forall a. Maybe a
Nothing Message -> [Message] -> NonEmpty Message
forall a. a -> [a] -> NonEmpty a
:| []
    , tools :: Maybe Text
tools = Maybe Text
forall a. Maybe a
Nothing
    , format :: Maybe Text
format = Maybe Text
forall a. Maybe a
Nothing
    , stream :: Maybe (ChatResponse -> IO (), IO ())
stream = Maybe (ChatResponse -> IO (), IO ())
forall a. Maybe a
Nothing
    , keepAlive :: Maybe Text
keepAlive = Maybe Text
forall a. Maybe a
Nothing
    , hostUrl :: Maybe Text
hostUrl = Maybe Text
forall a. Maybe a
Nothing
    , responseTimeOut :: Maybe Int
responseTimeOut = Maybe Int
forall a. Maybe a
Nothing
    }

{- |
Initiates a chat session with the specified 'ChatOps' configuration and returns either
a 'ChatResponse' or an error message.

This function sends a request to the Ollama chat API with the given options.

Example:

> let ops = defaultChatOps
> result <- chat ops
> case result of
>   Left errorMsg -> putStrLn ("Error: " ++ errorMsg)
>   Right response -> print response
-}
chat :: ChatOps -> IO (Either String ChatResponse)
chat :: ChatOps -> IO (Either String ChatResponse)
chat ChatOps
cOps = do
  let url :: Text
url = Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe Text
defaultOllamaUrl (ChatOps -> Maybe Text
hostUrl ChatOps
cOps)
      responseTimeout :: Int
responseTimeout = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
15 (ChatOps -> Maybe Int
responseTimeOut ChatOps
cOps)
  Manager
manager <-
    ManagerSettings -> IO Manager
newManager
      ManagerSettings
defaultManagerSettings -- Setting response timeout to 5 minutes, since llm takes time
        { 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/chat") :: IO (Either HttpException Request)
  case Either HttpException Request
eInitialRequest of
    Left HttpException
e -> Either String ChatResponse -> IO (Either String ChatResponse)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String ChatResponse -> IO (Either String ChatResponse))
-> Either String ChatResponse -> IO (Either String ChatResponse)
forall a b. (a -> b) -> a -> b
$ String -> Either String ChatResponse
forall a b. a -> Either a b
Left (String -> Either String ChatResponse)
-> String -> Either String ChatResponse
forall a b. (a -> b) -> a -> b
$ String
"Failed to parse host url: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> HttpException -> String
forall a. Show a => a -> String
show HttpException
e
    Right Request
initialRequest -> do
      let reqBody :: ChatOps
reqBody = ChatOps
cOps
          request :: Request
request =
            Request
initialRequest
              { method = "POST"
              , requestBody = RequestBodyLBS $ encode reqBody
              }
      Either HttpException (Either String ChatResponse)
eRes <-
        IO (Either String ChatResponse)
-> IO (Either HttpException (Either String ChatResponse))
forall e a. Exception e => IO a -> IO (Either e a)
try (Request
-> Manager
-> (Response BodyReader -> IO (Either String ChatResponse))
-> IO (Either String ChatResponse)
forall a.
Request -> Manager -> (Response BodyReader -> IO a) -> IO a
withResponse Request
request Manager
manager ((Response BodyReader -> IO (Either String ChatResponse))
 -> IO (Either String ChatResponse))
-> (Response BodyReader -> IO (Either String ChatResponse))
-> IO (Either String ChatResponse)
forall a b. (a -> b) -> a -> b
$ ChatOps -> Response BodyReader -> IO (Either String ChatResponse)
handleRequest ChatOps
cOps) ::
          IO (Either HttpException (Either String ChatResponse))
      case Either HttpException (Either String ChatResponse)
eRes of
        Left HttpException
e -> Either String ChatResponse -> IO (Either String ChatResponse)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String ChatResponse -> IO (Either String ChatResponse))
-> Either String ChatResponse -> IO (Either String ChatResponse)
forall a b. (a -> b) -> a -> b
$ String -> Either String ChatResponse
forall a b. a -> Either a b
Left (String -> Either String ChatResponse)
-> String -> Either String ChatResponse
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 ChatResponse
r -> Either String ChatResponse -> IO (Either String ChatResponse)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Either String ChatResponse
r

handleRequest :: ChatOps -> Response BodyReader -> IO (Either String ChatResponse)
handleRequest :: ChatOps -> Response BodyReader -> IO (Either String ChatResponse)
handleRequest ChatOps
cOps Response BodyReader
response = do
  let streamResponse :: (ChatResponse -> IO a) -> IO a -> IO (Either String b)
streamResponse ChatResponse -> 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 ChatResponse
eRes = ByteString -> Either String ChatResponse
forall a. FromJSON a => ByteString -> Either String a
eitherDecode (ByteString -> ByteString
BSL.fromStrict ByteString
bs) :: Either String ChatResponse
            case Either String ChatResponse
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 ChatResponse
r -> do
                a
_ <- ChatResponse -> IO a
sendChunk ChatResponse
r
                a
_ <- IO a
flush
                if ChatResponse -> Bool
done ChatResponse
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 (ChatResponse -> IO a) -> IO a -> IO (Either String b)
streamResponse ChatResponse -> IO a
sendChunk IO a
flush
  let genResponse :: ByteString -> IO (Either String ChatResponse)
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 -> Bool
BS.null ByteString
bs
          then do
            let eRes :: Either String ChatResponse
eRes = ByteString -> Either String ChatResponse
forall a. FromJSON a => ByteString -> Either String a
eitherDecode (ByteString -> ByteString
BSL.fromStrict ByteString
op) :: Either String ChatResponse
            case Either String ChatResponse
eRes of
              Left String
e -> Either String ChatResponse -> IO (Either String ChatResponse)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ChatResponse
forall a b. a -> Either a b
Left String
e)
              Right ChatResponse
r -> Either String ChatResponse -> IO (Either String ChatResponse)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ChatResponse -> Either String ChatResponse
forall a b. b -> Either a b
Right ChatResponse
r)
          else ByteString -> IO (Either String ChatResponse)
genResponse (ByteString
op ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
bs)
  case ChatOps -> Maybe (ChatResponse -> IO (), IO ())
stream ChatOps
cOps of
    Maybe (ChatResponse -> IO (), IO ())
Nothing -> ByteString -> IO (Either String ChatResponse)
genResponse ByteString
""
    Just (ChatResponse -> IO ()
sendChunk, IO ()
flush) -> (ChatResponse -> IO ()) -> IO () -> IO (Either String ChatResponse)
forall {a} {a} {b}.
(ChatResponse -> IO a) -> IO a -> IO (Either String b)
streamResponse ChatResponse -> IO ()
sendChunk IO ()
flush

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

 This function simply calls chat 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
  > }
  > let msg0 = Ollama.Message User "Sort given list: [4, 2 , 3, 67]. Also tell whether list was already sorted or not." Nothing
  > eRes3 <-
  >  chatJson
  >   defaultChatOps
  >    { Chat.chatModelName = "llama3.2"
  >      , Chat.messages = msg0 :| []
  >   }
  >      expectedJsonStrucutre
  >      (Just 2)
  > print eRes3
 Output:
  > 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.
-}
chatJson ::
  (FromJSON jsonResult, ToJSON jsonResult) =>
  ChatOps ->
  -- | Haskell type that you want your result in
  jsonResult ->
  -- | Max retries
  Maybe Int ->
  IO (Either String jsonResult)
chatJson :: forall jsonResult.
(FromJSON jsonResult, ToJSON jsonResult) =>
ChatOps -> jsonResult -> Maybe Int -> IO (Either String jsonResult)
chatJson cOps :: ChatOps
cOps@ChatOps {Maybe Int
Maybe (ChatResponse -> IO (), IO ())
Maybe Text
NonEmpty Message
Text
chatModelName :: ChatOps -> Text
messages :: ChatOps -> NonEmpty Message
tools :: ChatOps -> Maybe Text
format :: ChatOps -> Maybe Text
stream :: ChatOps -> Maybe (ChatResponse -> IO (), IO ())
keepAlive :: ChatOps -> Maybe Text
hostUrl :: ChatOps -> Maybe Text
responseTimeOut :: ChatOps -> Maybe Int
chatModelName :: Text
messages :: NonEmpty Message
tools :: Maybe Text
format :: Maybe Text
stream :: Maybe (ChatResponse -> IO (), IO ())
keepAlive :: Maybe Text
hostUrl :: Maybe Text
responseTimeOut :: Maybe Int
..} jsonResult
jsonStructure Maybe Int
mMaxRetries = do
  let lastMessage :: Message
lastMessage = NonEmpty Message -> Message
forall a. NonEmpty a -> a
NonEmpty.last NonEmpty Message
messages
      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
<> Message -> Text
content Message
lastMessage
          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 ChatResponse
chatResponse <-
    ChatOps -> IO (Either String ChatResponse)
chat
      ChatOps
cOps
        { messages =
            NonEmpty.fromList $
              lastMessage {content = jsonHelperPrompt} : NonEmpty.init messages
        }
  case Either String ChatResponse
chatResponse 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 ChatResponse
r -> do
      let mMessage :: Maybe Message
mMessage = ChatResponse -> Maybe Message
message ChatResponse
r
      case Maybe Message
mMessage of
        Maybe Message
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
"Something went wrong"
        Just Message
res -> 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
$ Message -> Text
content Message
res) 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 ChatOps -> jsonResult -> Maybe Int -> IO (Either String jsonResult)
forall jsonResult.
(FromJSON jsonResult, ToJSON jsonResult) =>
ChatOps -> jsonResult -> Maybe Int -> IO (Either String jsonResult)
chatJson ChatOps
cOps 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