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

{- |
Module:      Langchain.LLM.Core
Copyright:   (c) 2025 Tushar Adhatrao
License:     MIT
Maintainer:  Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability:   experimental

This module provides the core types and typeclasses for the Langchain library in Haskell,
which is designed to facilitate interaction with language models (LLMs). It defines a standardized
interface that allows different LLM implementations to be used interchangeably, promoting code reuse
and modularity.

The main components include:

* The 'LLM' typeclass, which defines the interface for language models.
* Data types such as 'Params' for configuring model invocations, 'Message' for conversation messages,
  and 'StreamHandler' for handling streaming responses.
* Default values like 'defaultParams' and 'defaultMessageData' for convenience.

This module is intended to be used as the foundation for building applications that interact with LLMs,
providing a consistent API across different model implementations.
-}
module Langchain.LLM.Core
  ( -- * LLM Typeclass
    LLM (..)

    -- * Parameters
  , Message (..)
  , Role (..)
  , ChatMessage
  , MessageData (..)
  , Params (..)
  , StreamHandler (..)

    -- * Default Values
  , defaultParams
  , defaultMessageData
  ) where

import Data.Aeson
import Data.List.NonEmpty
import Data.Text (Text)
import GHC.Generics

{- | Parameters for configuring language model invocations.
These parameters control aspects such as randomness, length, and stopping conditions of generated output.
This type corresponds to standard parameters in Python Langchain:
https://python.langchain.com/docs/concepts/chat_models/#standard-parameters

Example usage:

@
myParams :: Params
myParams = defaultParams
  { temperature = Just 0.7
  , maxTokens = Just 100
  }
@
-}
data Params = Params
  { Params -> Maybe Double
temperature :: Maybe Double
  -- ^ Sampling temperature. Higher values increase randomness (creativity), while lower values make output more focused.
  , Params -> Maybe Integer
maxTokens :: Maybe Integer
  , --- ^ Maximum number of tokens to generate in the response.
    Params -> Maybe Double
topP :: Maybe Double
  -- ^ Nucleus sampling parameter. Considers tokens whose cumulative probability mass is at least @topP@.
  , Params -> Maybe Int
n :: Maybe Int
  -- ^ Number of responses to generate (e.g., for sampling multiple outputs).
  , Params -> Maybe [Text]
stop :: Maybe [Text]
  -- ^ Sequences where generation should stop (e.g., ["\n"] stops at newlines).
  }
  deriving (Int -> Params -> ShowS
[Params] -> ShowS
Params -> String
(Int -> Params -> ShowS)
-> (Params -> String) -> ([Params] -> ShowS) -> Show Params
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Params -> ShowS
showsPrec :: Int -> Params -> ShowS
$cshow :: Params -> String
show :: Params -> String
$cshowList :: [Params] -> ShowS
showList :: [Params] -> ShowS
Show, Params -> Params -> Bool
(Params -> Params -> Bool)
-> (Params -> Params -> Bool) -> Eq Params
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Params -> Params -> Bool
== :: Params -> Params -> Bool
$c/= :: Params -> Params -> Bool
/= :: Params -> Params -> Bool
Eq)

{- | Callbacks for handling streaming responses from a language model.
This allows real-time processing of tokens as they are generated and an action
upon completion.

@
printHandler :: StreamHandler
printHandler = StreamHandler
  { onToken = putStrLn . ("Token: " ++)
  , onComplete = putStrLn "Streaming complete"
  }
@
-}
data StreamHandler = StreamHandler
  { StreamHandler -> Text -> IO ()
onToken :: Text -> IO ()
  -- ^ Action to perform for each token received
  , StreamHandler -> IO ()
onComplete :: IO ()
  -- ^ Action to perform when streaming is complete
  }

-- | Enumeration of possible roles in a conversation.
data Role
  = -- | System role, typically for instructions or context
    System
  | -- | User role, for user inputs
    User
  | -- | Assistant role, for model responses
    Assistant
  | -- | Tool role, for tool outputs or interactions
    Tool
  deriving (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, 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, (forall x. Role -> Rep Role x)
-> (forall x. Rep Role x -> Role) -> Generic Role
forall x. Rep Role x -> Role
forall x. Role -> Rep Role x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Role -> Rep Role x
from :: forall x. Role -> Rep Role x
$cto :: forall x. Rep Role x -> Role
to :: forall x. Rep Role x -> Role
Generic, [Role] -> Value
[Role] -> Encoding
Role -> Bool
Role -> Value
Role -> Encoding
(Role -> Value)
-> (Role -> Encoding)
-> ([Role] -> Value)
-> ([Role] -> Encoding)
-> (Role -> Bool)
-> ToJSON Role
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> (a -> Bool)
-> ToJSON a
$ctoJSON :: Role -> Value
toJSON :: Role -> Value
$ctoEncoding :: Role -> Encoding
toEncoding :: Role -> Encoding
$ctoJSONList :: [Role] -> Value
toJSONList :: [Role] -> Value
$ctoEncodingList :: [Role] -> Encoding
toEncodingList :: [Role] -> Encoding
$comitField :: Role -> Bool
omitField :: Role -> Bool
ToJSON, Maybe Role
Value -> Parser [Role]
Value -> Parser Role
(Value -> Parser Role)
-> (Value -> Parser [Role]) -> Maybe Role -> FromJSON Role
forall a.
(Value -> Parser a)
-> (Value -> Parser [a]) -> Maybe a -> FromJSON a
$cparseJSON :: Value -> Parser Role
parseJSON :: Value -> Parser Role
$cparseJSONList :: Value -> Parser [Role]
parseJSONList :: Value -> Parser [Role]
$comittedField :: Maybe Role
omittedField :: Maybe Role
FromJSON)

{- | Represents a message in a conversation, including the sender's role, content,
and additional metadata.
https://python.langchain.com/docs/concepts/messages/

@
userMsg :: Message
userMsg = Message
  { role = User
  , content = "Explain functional programming"
  , messageData = defaultMessageData
  }
@
-}
data Message = Message
  { Message -> Role
role :: Role
  -- ^ The role of the message sender
  , Message -> Text
content :: Text
  -- ^ The content of the message
  , Message -> MessageData
messageData :: MessageData
  -- ^ Additional data associated with the message
  }
  deriving (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, 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)

{- | Additional data for a message, such as a name or tool calls.
This type is designed for extensibility, allowing new fields to be added without
breaking changes. Use 'defaultMessageData' for typical usage.
-}
data MessageData = MessageData
  { MessageData -> Maybe Text
name :: Maybe Text
  -- ^ Optional name associated with the message
  , MessageData -> Maybe [Text]
toolCalls :: Maybe [Text]
  -- ^ Optional list of tool calls invoked by the message
  }
  deriving (MessageData -> MessageData -> Bool
(MessageData -> MessageData -> Bool)
-> (MessageData -> MessageData -> Bool) -> Eq MessageData
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: MessageData -> MessageData -> Bool
== :: MessageData -> MessageData -> Bool
$c/= :: MessageData -> MessageData -> Bool
/= :: MessageData -> MessageData -> Bool
Eq, Int -> MessageData -> ShowS
[MessageData] -> ShowS
MessageData -> String
(Int -> MessageData -> ShowS)
-> (MessageData -> String)
-> ([MessageData] -> ShowS)
-> Show MessageData
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MessageData -> ShowS
showsPrec :: Int -> MessageData -> ShowS
$cshow :: MessageData -> String
show :: MessageData -> String
$cshowList :: [MessageData] -> ShowS
showList :: [MessageData] -> ShowS
Show)

-- | JSON serialization for MessageData.
instance ToJSON MessageData where
  toJSON :: MessageData -> Value
toJSON MessageData {Maybe [Text]
Maybe Text
name :: MessageData -> Maybe Text
toolCalls :: MessageData -> Maybe [Text]
name :: Maybe Text
toolCalls :: Maybe [Text]
..} =
    [Pair] -> Value
object
      [ Key
"name" 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
name
      , Key
"tool_calls" 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]
toolCalls
      -- Add more fields as they are added
      ]

-- | JSON deserialization for MessageData.
instance FromJSON MessageData where
  parseJSON :: Value -> Parser MessageData
parseJSON = String
-> (Object -> Parser MessageData) -> Value -> Parser MessageData
forall a. String -> (Object -> Parser a) -> Value -> Parser a
withObject String
"MessageData" ((Object -> Parser MessageData) -> Value -> Parser MessageData)
-> (Object -> Parser MessageData) -> Value -> Parser MessageData
forall a b. (a -> b) -> a -> b
$ \Object
v ->
    Maybe Text -> Maybe [Text] -> MessageData
MessageData
      (Maybe Text -> Maybe [Text] -> MessageData)
-> Parser (Maybe Text) -> Parser (Maybe [Text] -> MessageData)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
v Object -> Key -> Parser (Maybe Text)
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"name"
      Parser (Maybe [Text] -> MessageData)
-> Parser (Maybe [Text]) -> Parser MessageData
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 [Text])
forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"tool_calls"

-- | Type alias for NonEmpty Message
type ChatMessage = NonEmpty Message

{- | Default message data with all fields set to Nothing.
Use this for standard messages without additional metadata
-}
defaultMessageData :: MessageData
defaultMessageData :: MessageData
defaultMessageData =
  MessageData
    { name :: Maybe Text
name = Maybe Text
forall a. Maybe a
Nothing
    , toolCalls :: Maybe [Text]
toolCalls = Maybe [Text]
forall a. Maybe a
Nothing
    }

{- | Typeclass defining the interface for language models.
This provides methods for invoking the model, chatting with it, and streaming
responses.

@
data TestLLM = TestLLM
  { responseText :: Text
  , shouldSucceed :: Bool
  }

instance LLM TestLLM where
  generate m _ _ = pure $ if shouldSucceed m
    then Right (responseText m)
    else Left "Test error"
@


@
ollamaLLM = Ollama "llama3.2:latest" [stdOutCallback]
response <- generate ollamaLLM "What is Haskell?" Nothing
@
-}
class LLM m where
  -- | Invoke the language model with a single prompt.
  --        Suitable for simple queries; returns either an error or generated text.

  {- === Using 'generate'
  To invoke an LLM with a single prompt:
  
  @
  let myLLM = ... -- assume this is an instance of LLM
  result <- generate myLLM "What is the meaning of life?" Nothing
  case result of
    Left err -> putStrLn $ "Error: " ++ err
    Right response -> putStrLn response
  @

  -}
  generate :: m -- ^ The type of the language model instance.
    -> Text -- ^ The prompt to send to the model.
    -> Maybe Params -- ^ Optional configuration parameters.
    -> IO (Either String Text)

  -- | Chat with the language model using a sequence of messages.
  -- Suitable for multi-turn conversations; returns either an error or the response.
  --
  chat :: m -- ^ The type of the language model instance.
    -> ChatMessage -- ^ A non-empty list of messages to send to the model.
    -> Maybe Params -- ^ Optional configuration parameters.
    -> IO (Either String Text) -- ^ The result of the chat, either an error or the response text.

  -- | Stream responses from the language model for a sequence of messages.
  -- Uses callbacks to process tokens in real-time; returns either an error or unit.
  stream :: m -> ChatMessage -> StreamHandler -> Maybe Params -> IO (Either String ())

{- | Default parameters with all fields set to Nothing.
Use this when no specific configuration is needed for the language model.

>>> generate myLLM "Hello" (Just defaultParams)
-}
defaultParams :: Params
defaultParams :: Params
defaultParams =
  Params
    { temperature :: Maybe Double
temperature = Maybe Double
forall a. Maybe a
Nothing
    , maxTokens :: Maybe Integer
maxTokens = Maybe Integer
forall a. Maybe a
Nothing
    , topP :: Maybe Double
topP = Maybe Double
forall a. Maybe a
Nothing
    , n :: Maybe Int
n = Maybe Int
forall a. Maybe a
Nothing
    , stop :: Maybe [Text]
stop = Maybe [Text]
forall a. Maybe a
Nothing
    }