{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}

{- |
Module      : Langchain.Agents.React
Description : Implementation of ReAct agent combining reasoning and action
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>

Implements the ReAct pattern where the agent alternates between:

1. Reasoning (generating thoughts)
2. Acting (executing tools)

Example agent interaction:

> agent <- createReactAgent llm [wikipediaTool, calculatorTool]
> result <- runAgentExecutor executor "What's the population of Paris?"
> -- Agent might:
> -- 1. Use Wikipedia tool to find current population data
> -- 2. Use calculator tool to verify numbers
> -- 3. Return final answer
-}
module Langchain.Agents.React
  ( ReactAgentOutputParser (..)
  , parseReactOutput
  , ReactAgent (..)
  , createReactAgent
  , formatToolDescriptions
  , formatToolNames
  , getLastUserInput
  ) where

import qualified Data.List.NonEmpty as NE
import qualified Data.Map.Strict as Map
import Data.Text (Text)
import qualified Data.Text as T
import Langchain.Agents.Core
import Langchain.LLM.Core
import Langchain.Memory.Core
import Langchain.OutputParser.Core
import Langchain.PromptTemplate
import Langchain.Tool.Core

{- |
Output parser for ReAct agent responses
Handles two primary formats:

1. Final answers containing "Final Answer:"
2. Action requests with "Action:" and "Action Input:"

Example parsing:

> parseReactOutput "Final Answer: 42"
> -- Right (Finish ...)
>
> parseReactOutput "Action: calculator\nAction Input: 5+3"
> -- Right (Continue ...)
-}
newtype ReactAgentOutputParser = ReactAgentOutputParser AgentStep

instance OutputParser ReactAgentOutputParser where
  parse :: Text -> Either String ReactAgentOutputParser
parse = Text -> Either String ReactAgentOutputParser
parseReactOutput

-- | Parses the output from a React agent
parseReactOutput :: Text -> Either String ReactAgentOutputParser
parseReactOutput :: Text -> Either String ReactAgentOutputParser
parseReactOutput Text
text
  | Text -> Text -> Bool
T.isInfixOf Text
"Final Answer:" Text
text =
      -- Extract the final answer
      let answer :: Text
answer = Text -> Text -> Text
extractAfter Text
"Final Answer:" Text
text
       in ReactAgentOutputParser -> Either String ReactAgentOutputParser
forall a b. b -> Either a b
Right (ReactAgentOutputParser -> Either String ReactAgentOutputParser)
-> ReactAgentOutputParser -> Either String ReactAgentOutputParser
forall a b. (a -> b) -> a -> b
$
            AgentStep -> ReactAgentOutputParser
ReactAgentOutputParser (AgentStep -> ReactAgentOutputParser)
-> AgentStep -> ReactAgentOutputParser
forall a b. (a -> b) -> a -> b
$
              AgentFinish -> AgentStep
Finish (AgentFinish -> AgentStep) -> AgentFinish -> AgentStep
forall a b. (a -> b) -> a -> b
$
                AgentFinish
                  { returnValues :: Map Text Text
returnValues = Text -> Text -> Map Text Text
forall k a. k -> a -> Map k a
Map.singleton Text
"output" Text
answer
                  , finishLog :: Text
finishLog = Text
text
                  }
  | Text -> Text -> Bool
T.isInfixOf Text
"Action:" Text
text Bool -> Bool -> Bool
&& Text -> Text -> Bool
T.isInfixOf Text
"Action Input:" Text
text =
      -- Extract action and action input
      let actionName :: Text
actionName = Text -> Text -> Text
extractAfter Text
"Action:" (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ (Char -> Bool) -> Text -> Text
T.takeWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'\n') (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ (Char -> Bool) -> Text -> Text
T.dropWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'A') Text
text
          actionInput_ :: Text
actionInput_ =
            Text -> Text -> Text
extractAfter Text
"Action Input:" (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ (Char -> Bool) -> Text -> Text
T.takeWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'\n') (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ (Text, Text) -> Text
forall a b. (a, b) -> b
snd ((Text, Text) -> Text) -> (Text, Text) -> Text
forall a b. (a -> b) -> a -> b
$ HasCallStack => Text -> Text -> (Text, Text)
Text -> Text -> (Text, Text)
T.breakOn Text
"Action Input:" Text
text
       in ReactAgentOutputParser -> Either String ReactAgentOutputParser
forall a b. b -> Either a b
Right (ReactAgentOutputParser -> Either String ReactAgentOutputParser)
-> ReactAgentOutputParser -> Either String ReactAgentOutputParser
forall a b. (a -> b) -> a -> b
$
            AgentStep -> ReactAgentOutputParser
ReactAgentOutputParser (AgentStep -> ReactAgentOutputParser)
-> AgentStep -> ReactAgentOutputParser
forall a b. (a -> b) -> a -> b
$
              AgentAction -> AgentStep
Continue (AgentAction -> AgentStep) -> AgentAction -> AgentStep
forall a b. (a -> b) -> a -> b
$
                AgentAction
                  { actionToolName :: Text
actionToolName = Text -> Text
T.strip Text
actionName
                  , actionInput :: Text
actionInput = Text -> Text
T.strip Text
actionInput_
                  , actionLog :: Text
actionLog = Text
text
                  }
  | Bool
otherwise = String -> Either String ReactAgentOutputParser
forall a b. a -> Either a b
Left (String -> Either String ReactAgentOutputParser)
-> String -> Either String ReactAgentOutputParser
forall a b. (a -> b) -> a -> b
$ String
"Could not parse agent output: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
text

{- |
Core ReAct agent configuration.
Contains:

- LLM for reasoning
- Available tools
- Prompt template for interaction

Example creation:

> agent <- createReactAgent
>   openAIGPT
>   [ AnyTool wikipediaTool
>   , AnyTool calculatorTool
>   ]
-}
data (LLM llm) => ReactAgent llm = ReactAgent
  { forall llm. LLM llm => ReactAgent llm -> llm
reactLLM :: llm
  , forall llm. LLM llm => ReactAgent llm -> [AnyTool]
reactTools :: [AnyTool]
  , forall llm. LLM llm => ReactAgent llm -> PromptTemplate
reactPromptTemplate :: PromptTemplate
  }

-- Helper function to extract text after a marker
extractAfter :: Text -> Text -> Text
extractAfter :: Text -> Text -> Text
extractAfter Text
marker Text
text =
  let afterMarker :: Text
afterMarker = (Text, Text) -> Text
forall a b. (a, b) -> b
snd ((Text, Text) -> Text) -> (Text, Text) -> Text
forall a b. (a -> b) -> a -> b
$ HasCallStack => Text -> Text -> (Text, Text)
Text -> Text -> (Text, Text)
T.breakOn Text
marker Text
text
   in if Text -> Bool
T.null Text
afterMarker
        then Text
""
        else Text -> Text
T.strip (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ (Char -> Bool) -> Text -> Text
T.dropWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
':') Text
afterMarker

{- |
Creates a ReAct agent with standard prompt structure
The prompt instructs the LLM to:

1. List available tools
2. Follow thought-action-observation pattern
3. Provide final answers

Example prompt excerpt:

> "Use the following format:
> Thought: ...
> Action: [tool_name]
> Action Input: ..."
-}
createReactAgent ::
  (LLM llm) =>
  llm ->
  [AnyTool] ->
  IO (Either String (ReactAgent llm))
createReactAgent :: forall llm.
LLM llm =>
llm -> [AnyTool] -> IO (Either String (ReactAgent llm))
createReactAgent llm
llm [AnyTool]
tools = do
  let reactPrompt :: PromptTemplate
reactPrompt =
        Text -> PromptTemplate
PromptTemplate (Text -> PromptTemplate) -> Text -> PromptTemplate
forall a b. (a -> b) -> a -> b
$
          [Text] -> Text
T.unlines
            [ Text
"You are an AI assistant designed to help with tasks."
            , Text
"You have access to the following tools:"
            , Text
"{tools_description}"
            , Text
""
            , Text
"Use the following format:"
            , Text
""
            , Text
"Thought: you should always think about what to do"
            , Text
"Action: the action to take, should be one of [{tool_names}]"
            , Text
"Action Input: the input to the action"
            , Text
"Observation: the result of the action"
            , Text
"... (this Thought/Action/Action Input/Observation can repeat N times)"
            , Text
"Thought: I now know the final answer"
            , Text
"Final Answer: the final answer to the original input question"
            ]
  Either String (ReactAgent llm)
-> IO (Either String (ReactAgent llm))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String (ReactAgent llm)
 -> IO (Either String (ReactAgent llm)))
-> Either String (ReactAgent llm)
-> IO (Either String (ReactAgent llm))
forall a b. (a -> b) -> a -> b
$
    ReactAgent llm -> Either String (ReactAgent llm)
forall a b. b -> Either a b
Right (ReactAgent llm -> Either String (ReactAgent llm))
-> ReactAgent llm -> Either String (ReactAgent llm)
forall a b. (a -> b) -> a -> b
$
      ReactAgent
        { reactLLM :: llm
reactLLM = llm
llm
        , reactTools :: [AnyTool]
reactTools = [AnyTool]
tools
        , reactPromptTemplate :: PromptTemplate
reactPromptTemplate = PromptTemplate
reactPrompt
        }

instance (LLM llm) => Agent (ReactAgent llm) where
  -- \|
  --  Core reasoning loop implementing ReAct pattern
  --
  --  1. Retrieve chat history
  --  2. Format tool information
  --  3. Construct reasoning prompt
  --  4. Execute LLM call
  --  5. Parse response into action/answer
  --
  --  Uses depth-first planning with backtracking
  --
  planNextAction :: forall m.
BaseMemory m =>
ReactAgent llm -> AgentState m -> IO (Either String AgentStep)
planNextAction ReactAgent {llm
[AnyTool]
PromptTemplate
reactLLM :: forall llm. LLM llm => ReactAgent llm -> llm
reactTools :: forall llm. LLM llm => ReactAgent llm -> [AnyTool]
reactPromptTemplate :: forall llm. LLM llm => ReactAgent llm -> PromptTemplate
reactLLM :: llm
reactTools :: [AnyTool]
reactPromptTemplate :: PromptTemplate
..} AgentState m
state = do
    let mem :: m
mem = AgentState m -> m
forall m. BaseMemory m => AgentState m -> m
agentMemory AgentState m
state
    Either String ChatMessage
msgResult <- m -> IO (Either String ChatMessage)
forall m. BaseMemory m => m -> IO (Either String ChatMessage)
messages m
mem
    case Either String ChatMessage
msgResult of
      Left String
err -> Either String AgentStep -> IO (Either String AgentStep)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String AgentStep -> IO (Either String AgentStep))
-> Either String AgentStep -> IO (Either String AgentStep)
forall a b. (a -> b) -> a -> b
$ String -> Either String AgentStep
forall a b. a -> Either a b
Left String
err
      Right ChatMessage
msgs -> do
        -- Format the tools descriptions
        let toolDescs :: Text
toolDescs = [AnyTool] -> Text
formatToolDescriptions [AnyTool]
reactTools
            userQuery :: Text
userQuery = ChatMessage -> Text
getLastUserInput ChatMessage
msgs
        -- Build the prompt variables
        let promptVars :: Map Text Text
promptVars =
              [(Text, Text)] -> Map Text Text
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
                [ (Text
"tools_description", Text
toolDescs)
                , (Text
"tool_names", [AnyTool] -> Text
formatToolNames [AnyTool]
reactTools)
                ]

        -- Render the prompt
        case PromptTemplate -> Map Text Text -> Either String Text
renderPrompt PromptTemplate
reactPromptTemplate Map Text Text
promptVars of
          Left String
err -> Either String AgentStep -> IO (Either String AgentStep)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String AgentStep -> IO (Either String AgentStep))
-> Either String AgentStep -> IO (Either String AgentStep)
forall a b. (a -> b) -> a -> b
$ String -> Either String AgentStep
forall a b. a -> Either a b
Left String
err
          Right Text
renderedPrompt -> do
            -- Call the LLM
            let m :: ChatMessage
m =
                  ( ChatMessage
msgs
                      ChatMessage -> ChatMessage -> ChatMessage
forall a. NonEmpty a -> NonEmpty a -> NonEmpty a
`NE.append` [Message] -> ChatMessage
forall a. HasCallStack => [a] -> NonEmpty a
NE.fromList
                        [ (Role -> Text -> MessageData -> Message
Message Role
System Text
renderedPrompt MessageData
defaultMessageData)
                        , (Role -> Text -> MessageData -> Message
Message Role
User Text
userQuery MessageData
defaultMessageData)
                        ]
                  )
            Either String Text
response <-
              llm -> ChatMessage -> Maybe Params -> IO (Either String Text)
forall m.
LLM m =>
m -> ChatMessage -> Maybe Params -> IO (Either String Text)
chat
                llm
reactLLM
                ChatMessage
m
                Maybe Params
forall a. Maybe a
Nothing
            case Either String Text
response of
              Left String
err -> Either String AgentStep -> IO (Either String AgentStep)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String AgentStep -> IO (Either String AgentStep))
-> Either String AgentStep -> IO (Either String AgentStep)
forall a b. (a -> b) -> a -> b
$ String -> Either String AgentStep
forall a b. a -> Either a b
Left String
err
              Right Text
llmOutput -> do
                -- Parse the output
                case Text -> Either String ReactAgentOutputParser
forall a. OutputParser a => Text -> Either String a
parse Text
llmOutput of
                  Left String
err -> Either String AgentStep -> IO (Either String AgentStep)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String AgentStep -> IO (Either String AgentStep))
-> Either String AgentStep -> IO (Either String AgentStep)
forall a b. (a -> b) -> a -> b
$ String -> Either String AgentStep
forall a b. a -> Either a b
Left (String -> Either String AgentStep)
-> String -> Either String AgentStep
forall a b. (a -> b) -> a -> b
$ String
"Failed to parse LLM output: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
err
                  Right (ReactAgentOutputParser AgentStep
step) -> Either String AgentStep -> IO (Either String AgentStep)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String AgentStep -> IO (Either String AgentStep))
-> Either String AgentStep -> IO (Either String AgentStep)
forall a b. (a -> b) -> a -> b
$ AgentStep -> Either String AgentStep
forall a b. b -> Either a b
Right AgentStep
step

  agentPrompt :: ReactAgent llm -> IO PromptTemplate
agentPrompt ReactAgent {llm
[AnyTool]
PromptTemplate
reactLLM :: forall llm. LLM llm => ReactAgent llm -> llm
reactTools :: forall llm. LLM llm => ReactAgent llm -> [AnyTool]
reactPromptTemplate :: forall llm. LLM llm => ReactAgent llm -> PromptTemplate
reactLLM :: llm
reactTools :: [AnyTool]
reactPromptTemplate :: PromptTemplate
..} = PromptTemplate -> IO PromptTemplate
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PromptTemplate
reactPromptTemplate
  agentTools :: ReactAgent llm -> IO [AnyTool]
agentTools ReactAgent {llm
[AnyTool]
PromptTemplate
reactLLM :: forall llm. LLM llm => ReactAgent llm -> llm
reactTools :: forall llm. LLM llm => ReactAgent llm -> [AnyTool]
reactPromptTemplate :: forall llm. LLM llm => ReactAgent llm -> PromptTemplate
reactLLM :: llm
reactTools :: [AnyTool]
reactPromptTemplate :: PromptTemplate
..} = [AnyTool] -> IO [AnyTool]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [AnyTool]
reactTools

{- |
Formats tool descriptions for LLM consumption
Creates a list like:

> "Tool: wikipedia
>  Description: Search Wikipedia..."
-}
formatToolDescriptions :: [AnyTool] -> Text
formatToolDescriptions :: [AnyTool] -> Text
formatToolDescriptions [AnyTool]
tools = Text -> [Text] -> Text
T.intercalate Text
"\n\n" ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$ (AnyTool -> Text) -> [AnyTool] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map AnyTool -> Text
formatTool [AnyTool]
tools
  where
    formatTool :: AnyTool -> Text
formatTool (AnyTool a
tool Text -> Input a
_ Output a -> Text
_) =
      [Text] -> Text
T.concat [Text
"Tool: ", a -> Text
forall a. Tool a => a -> Text
toolName a
tool, Text
"\nDescription: ", a -> Text
forall a. Tool a => a -> Text
toolDescription a
tool]

{- |
Creates comma-separated tool names for prompt inclusion
Example output: "wikipedia, calculator, weather"
-}
formatToolNames :: [AnyTool] -> Text
formatToolNames :: [AnyTool] -> Text
formatToolNames [AnyTool]
tools = Text -> [Text] -> Text
T.intercalate Text
", " ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$ (AnyTool -> Text) -> [AnyTool] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (\(AnyTool a
tool Text -> Input a
_ Output a -> Text
_) -> a -> Text
forall a. Tool a => a -> Text
toolName a
tool) [AnyTool]
tools

{- |
Extracts latest user query from chat history
Handles cases where:

- Multiple user messages exist
- No user input found
-}
getLastUserInput :: ChatMessage -> Text
getLastUserInput :: ChatMessage -> Text
getLastUserInput ChatMessage
msgs =
  let userMsgs :: [Message]
userMsgs = (Message -> Bool) -> [Message] -> [Message]
forall a. (a -> Bool) -> [a] -> [a]
filter (\Message
m -> Message -> Role
role Message
m Role -> Role -> Bool
forall a. Eq a => a -> a -> Bool
== Role
User) ([Message] -> [Message]) -> [Message] -> [Message]
forall a b. (a -> b) -> a -> b
$ ChatMessage -> [Message]
forall a. NonEmpty a -> [a]
NE.toList ChatMessage
msgs
   in if [Message] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Message]
userMsgs
        then Text
""
        else Message -> Text
content (Message -> Text) -> Message -> Text
forall a b. (a -> b) -> a -> b
$ [Message] -> Message
forall a. HasCallStack => [a] -> a
last [Message]
userMsgs