{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

{- |
Module      : Langchain.LLM.Ollama
Description : Ollama integration for LangChain Haskell
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability   : experimental

Ollama implementation of LangChain's LLM interface , supporting:

- Text generation
- Chat interactions
- Streaming responses
- Callback integration

Example usage:

@
-- Create Ollama configuration
ollamaLLM = Ollama "llama3" [stdOutCallback]

-- Generate text
response <- generate ollamaLLM "Explain Haskell monads" Nothing
-- Right "Monads in Haskell..."

-- Chat interaction
let messages = UserMessage "What's the capital of France?" :| []
chatResponse <- chat ollamaLLM messages Nothing
-- Right "The capital of France is Paris."

-- Streaming
streamHandler = StreamHandler print (putStrLn "Done")
streamResult <- stream ollamaLLM messages streamHandler Nothing
@
-}
module Langchain.LLM.Ollama (Ollama (..)) where

import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Ollama.Chat as OllamaChat
import qualified Data.Ollama.Generate as OllamaGenerate
import Data.Text (Text)
import Langchain.Callback (Callback, Event (..))
import Langchain.LLM.Core
import qualified Langchain.Runnable.Core as Run

{- | Ollama LLM configuration
Contains:

- Model name (e.g., "llama3:latest")
- Callbacks for event tracking

Example:

>>> Ollama "nomic-embed" [logCallback]
Ollama "nomic-embed"
-}
data Ollama = Ollama
  { Ollama -> Text
modelName :: Text
  -- ^ The name of the Ollama model
  , Ollama -> [Callback]
callbacks :: [Callback]
  -- ^ Event handlers for LLM operations
  }

instance Show Ollama where
  show :: Ollama -> String
show (Ollama Text
modelName [Callback]
_) = String
"Ollama " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> String
forall a. Show a => a -> String
show Text
modelName

{- | Ollama implementation of the LLM typeclass
Note: Params argument is currently ignored (see TODOs).

Example instance usage:

@
-- Generate text with error handling
case generate ollamaLLM "Hello" Nothing of
  Left err -> putStrLn $ "Error: " ++ err
  Right res -> putStrLn res
@
-}
instance LLM Ollama where
  -- \| Generate text from a prompt
  --  Returns Left on API errors, Right on success.
  --
  --  Example:
  --  >>> generate (Ollama "llama3.2" []) "Hello" Nothing
  --  Right "Hello! How can I assist you today?"
  --
  generate :: Ollama -> Text -> Maybe Params -> IO (Either String Text)
generate (Ollama Text
model [Callback]
cbs) Text
prompt Maybe Params
_ = do
    (Callback -> IO ()) -> [Callback] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Callback
cb -> Callback
cb Event
LLMStart) [Callback]
cbs
    Either String GenerateResponse
eRes <-
      GenerateOps -> IO (Either String GenerateResponse)
OllamaGenerate.generate
        GenerateOps
OllamaGenerate.defaultGenerateOps
          { OllamaGenerate.modelName = model
          , OllamaGenerate.prompt = prompt
          , OllamaGenerate.stream = Nothing
          }
    case Either String GenerateResponse
eRes of
      Left String
err -> do
        (Callback -> IO ()) -> [Callback] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Callback
cb -> Callback
cb (String -> Event
LLMError String
err)) [Callback]
cbs
        Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String Text -> IO (Either String Text))
-> Either String Text -> IO (Either String Text)
forall a b. (a -> b) -> a -> b
$ String -> Either String Text
forall a b. a -> Either a b
Left (ShowS
forall a. Show a => a -> String
show String
err)
      Right GenerateResponse
res -> do
        (Callback -> IO ()) -> [Callback] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Callback
cb -> Callback
cb Event
LLMEnd) [Callback]
cbs
        Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String Text -> IO (Either String Text))
-> Either String Text -> IO (Either String Text)
forall a b. (a -> b) -> a -> b
$ Text -> Either String Text
forall a b. b -> Either a b
Right (GenerateResponse -> Text
OllamaGenerate.response_ GenerateResponse
res)

  -- \| Chat interaction with message history.
  --  Uses Ollama's chat API for multi-turn conversations.
  --
  --  Example:
  --  >>> let msgs = UserMessage "Hi" :| [AssistantMessage "Hello!"]
  --  >>> chat (Ollama "llama3" []) msgs Nothing
  --  Right "How are you today?"
  --
  chat :: Ollama -> ChatMessage -> Maybe Params -> IO (Either String Text)
chat (Ollama Text
model [Callback]
cbs) ChatMessage
messages Maybe Params
_ = do
    (Callback -> IO ()) -> [Callback] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Callback
cb -> Callback
cb Event
LLMStart) [Callback]
cbs
    Either String ChatResponse
eRes <-
      ChatOps -> IO (Either String ChatResponse)
OllamaChat.chat
        ChatOps
OllamaChat.defaultChatOps
          { OllamaChat.chatModelName = model
          , OllamaChat.messages = toOllamaMessages messages
          , OllamaChat.stream = Nothing
          }
    case Either String ChatResponse
eRes of
      Left String
err -> do
        (Callback -> IO ()) -> [Callback] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Callback
cb -> Callback
cb (String -> Event
LLMError String
err)) [Callback]
cbs
        Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String Text -> IO (Either String Text))
-> Either String Text -> IO (Either String Text)
forall a b. (a -> b) -> a -> b
$ String -> Either String Text
forall a b. a -> Either a b
Left (ShowS
forall a. Show a => a -> String
show String
err)
      Right ChatResponse
res -> do
        (Callback -> IO ()) -> [Callback] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Callback
cb -> Callback
cb Event
LLMEnd) [Callback]
cbs
        Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String Text -> IO (Either String Text))
-> Either String Text -> IO (Either String Text)
forall a b. (a -> b) -> a -> b
$ Text -> Either String Text
forall a b. b -> Either a b
Right (ChatResponse -> Text
chatRespToText ChatResponse
res)
    where
      chatRespToText :: ChatResponse -> Text
chatRespToText ChatResponse
resp = Text -> (Message -> Text) -> Maybe Message -> Text
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Text
"" Message -> Text
OllamaChat.content (ChatResponse -> Maybe Message
OllamaChat.message ChatResponse
resp)

  -- \| Streaming response handling.
  --  Processes tokens in real-time via StreamHandler.
  --
  --  Example:
  --  >>> let handler = StreamHandler (putStr . ("Token: " ++)) (putStrLn "Complete")
  --  >>> stream (Ollama "llama3" []) messages handler Nothing
  --  Token: H Token: i Complete
  --
  stream :: Ollama
-> ChatMessage
-> StreamHandler
-> Maybe Params
-> IO (Either String ())
stream (Ollama Text
model_ [Callback]
cbs) ChatMessage
messages StreamHandler {Text -> IO ()
onToken :: Text -> IO ()
onToken :: StreamHandler -> Text -> IO ()
onToken, IO ()
onComplete :: IO ()
onComplete :: StreamHandler -> IO ()
onComplete} Maybe Params
_ = do
    (Callback -> IO ()) -> [Callback] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Callback
cb -> Callback
cb Event
LLMStart) [Callback]
cbs
    Either String ChatResponse
eRes <-
      ChatOps -> IO (Either String ChatResponse)
OllamaChat.chat
        ChatOps
OllamaChat.defaultChatOps
          { OllamaChat.chatModelName = model_
          , OllamaChat.messages = toOllamaMessages messages
          , OllamaChat.stream = Just (onToken . chatRespToText, onComplete)
          }
    case Either String ChatResponse
eRes of
      Left String
err -> do
        (Callback -> IO ()) -> [Callback] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Callback
cb -> Callback
cb (String -> Event
LLMError String
err)) [Callback]
cbs
        Either String () -> IO (Either String ())
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String () -> IO (Either String ()))
-> Either String () -> IO (Either String ())
forall a b. (a -> b) -> a -> b
$ String -> Either String ()
forall a b. a -> Either a b
Left (ShowS
forall a. Show a => a -> String
show String
err)
      Right ChatResponse
_ -> do
        (Callback -> IO ()) -> [Callback] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Callback
cb -> Callback
cb Event
LLMEnd) [Callback]
cbs
        Either String () -> IO (Either String ())
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String () -> IO (Either String ()))
-> Either String () -> IO (Either String ())
forall a b. (a -> b) -> a -> b
$ () -> Either String ()
forall a b. b -> Either a b
Right ()
    where
      chatRespToText :: ChatResponse -> Text
chatRespToText OllamaChat.ChatResponse {Bool
Maybe Int64
Maybe Message
Text
UTCTime
message :: ChatResponse -> Maybe Message
model :: Text
createdAt :: UTCTime
message :: Maybe Message
done :: Bool
totalDuration :: Maybe Int64
loadDuration :: Maybe Int64
promptEvalCount :: Maybe Int64
promptEvalDuration :: Maybe Int64
evalCount :: Maybe Int64
evalDuration :: Maybe Int64
model :: ChatResponse -> Text
createdAt :: ChatResponse -> UTCTime
done :: ChatResponse -> Bool
totalDuration :: ChatResponse -> Maybe Int64
loadDuration :: ChatResponse -> Maybe Int64
promptEvalCount :: ChatResponse -> Maybe Int64
promptEvalDuration :: ChatResponse -> Maybe Int64
evalCount :: ChatResponse -> Maybe Int64
evalDuration :: ChatResponse -> Maybe Int64
..} = Text -> (Message -> Text) -> Maybe Message -> Text
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Text
"" Message -> Text
OllamaChat.content Maybe Message
message

{- | Convert LangChain messages to Ollama format.
Current limitations:
- Ignores 'messageData' field
- No tool call support (see TODO)

Example conversion:
>>> let msg = Message System "You are an assistant" defaultMessageData
>>> toOllamaMessages (msg :| [])
NonEmpty [OllamaChat.Message System "You are an assistant" Nothing Nothing]
-}
toOllamaMessages :: NonEmpty Message -> NonEmpty OllamaChat.Message
toOllamaMessages :: ChatMessage -> NonEmpty Message
toOllamaMessages = (Message -> Message) -> ChatMessage -> NonEmpty Message
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NonEmpty.map ((Message -> Message) -> ChatMessage -> NonEmpty Message)
-> (Message -> Message) -> ChatMessage -> NonEmpty Message
forall a b. (a -> b) -> a -> b
$ \Message {Text
MessageData
Role
messageData :: Message -> MessageData
role :: Role
content :: Text
messageData :: MessageData
role :: Message -> Role
content :: Message -> Text
..} ->
  Role -> Text -> Maybe [Text] -> Maybe [Value] -> Message
OllamaChat.Message (Role -> Role
toOllamaRole Role
role) Text
content Maybe [Text]
forall a. Maybe a
Nothing Maybe [Value]
forall a. Maybe a
Nothing
  where
    toOllamaRole :: Role -> Role
toOllamaRole Role
User = Role
OllamaChat.User
    toOllamaRole Role
System = Role
OllamaChat.System
    toOllamaRole Role
Assistant = Role
OllamaChat.Assistant
    toOllamaRole Role
Tool = Role
OllamaChat.Tool

instance Run.Runnable Ollama where
  type RunnableInput Ollama = ChatMessage
  type RunnableOutput Ollama = Text

  -- TODO: need to figure out a way to pass mbParams
  -- \| Runnable interface implementation.
  --  Currently delegates to 'chat' method with default parameters.
  --
  invoke :: Ollama
-> RunnableInput Ollama
-> IO (Either String (RunnableOutput Ollama))
invoke Ollama
model RunnableInput Ollama
input = Ollama -> ChatMessage -> Maybe Params -> IO (Either String Text)
forall m.
LLM m =>
m -> ChatMessage -> Maybe Params -> IO (Either String Text)
chat Ollama
model ChatMessage
RunnableInput Ollama
input Maybe Params
forall a. Maybe a
Nothing

{- $examples
Test case patterns:
1. Basic generation
   >>> generate (Ollama "test-model" []) "Hello" Nothing
   Right "Mock response"

2. Error handling
   >>> generate (Ollama "invalid-model" []) "Test" Nothing
   Left "API request failed"

3. Streaming interaction
   >>> let handler = StreamHandler print (pure ())
   >>> stream (Ollama "llama3" []) (UserMessage "Hi" :| []) handler Nothing
   Right ()
-}