Core.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE RecordWildCards #-}

{- |
Module      : Langchain.Agents.Core
Description : Core implementation of LangChain agents
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>

Agents use LLMs as reasoning engines to determine actions dynamically. 
-}
module Langchain.Agents.Core
  ( AgentAction (..)
  , AgentFinish (..)
  , AgentStep (..)
  , Agent (..)
  , AnyTool (..)
  , AgentState (..)
  , runAgent
  , runAgentLoop
  , executeTool
  , runSingleStep
  , customAnyTool
  ) where

import Control.Exception (SomeException, try)
import Data.List (find)
import qualified Data.Map.Strict as Map
import Data.Text (Text)
import qualified Data.Text as T
import Langchain.LLM.Core (Message (Message), Role (..), defaultMessageData)
import Langchain.Memory.Core (BaseMemory (..))
import Langchain.PromptTemplate (PromptTemplate)
import Langchain.Tool.Core (Tool (..))

{- |
Represents an action to be taken by the agent
-}
data AgentAction = AgentAction
  { actionToolName :: Text
  -- ^ Tool name
  , actionInput :: Text
  -- ^ Input
  , actionLog :: Text
  -- ^ Execution log
  }
  deriving (Eq, Show)

-- | Represents that agent has finished work with final value
data AgentFinish = AgentFinish
  { returnValues :: Map.Map Text Text 
  , finishLog :: Text
  }
  deriving (Show, Eq)

-- | Type that will be return from LLM 
-- Could be either Continue, making another call to LLM or Finish with final value
data AgentStep
  = Continue AgentAction
  | Finish AgentFinish
  deriving (Eq, Show)

-- | Type for maintaining state of the agent 
data (BaseMemory m) => AgentState m = AgentState
  { agentMemory :: m -- ^ Memory for storing chat history
  , agentToolResults :: [(Text, Text)] -- ^ Tool results
  , agentSteps :: [AgentAction] -- ^ Agent steps happened so far
  }
  deriving (Eq, Show)

-- | A type that helps various types of Tools
-- It encapsulates the Tool, and conversion functions 
-- to and from Text for Tool input and output since Agent takes and returns Text. 
-- If Tool takes or returns Text type itself you can use `id` at these places.
data AnyTool = forall a. Tool a => AnyTool
  { anyTool :: a
  , textToInput :: Text -> Input a
  , outputToText :: Output a -> Text
  }

-- | Typeclass for Agent
class Agent a where
  planNextAction :: BaseMemory m => a -> AgentState m -> IO (Either String AgentStep)
  agentPrompt :: a -> IO PromptTemplate
  agentTools :: a -> IO [AnyTool]

-- | Function that *starts* the agent process.
runAgent :: (Agent a, BaseMemory m) => a -> AgentState m -> Text -> IO (Either String AgentFinish)
runAgent agent initialState@AgentState {..} initialInput = do
  memWithInput <- addUserMessage agentMemory initialInput
  case memWithInput of
    Left err -> return $ Left err
    Right updatedMem ->
      let newState = initialState {agentMemory = updatedMem}
       in runAgentLoop agent newState 0 10

-- | Helper function for runAgent
runAgentLoop ::
  (Agent a, BaseMemory m) => a -> AgentState m -> Int -> Int -> IO (Either String AgentFinish)
runAgentLoop agent agentState@AgentState {..} currIter maxIter
  | currIter > maxIter = return $ Left "Max iterations excedded"
  | otherwise = do
      eStepResult <- runSingleStep agent agentState
      case eStepResult of
        Left err -> return $ Left err
        Right (Finish agentFinish) -> return $ Right agentFinish
        Right (Continue act@AgentAction {..}) -> do
          toolList <- agentTools agent
          toolResult <- executeTool toolList actionToolName actionInput
          case toolResult of
            Left err -> return $ Left err
            Right result -> do
              -- Add the tool result to memory as a tool message
              let toolMsg = Message Tool result defaultMessageData
              updatedMemResult <- addMessage agentMemory toolMsg
              case updatedMemResult of
                Left err -> return $ Left err
                Right updatedMem ->
                  let updatedState =
                        agentState
                          { agentMemory = updatedMem
                          , agentToolResults = agentToolResults ++ [(actionToolName, result)]
                          , agentSteps = agentSteps ++ [act]
                          }
                   in runAgentLoop agent updatedState (currIter + 1) maxIter

-- | Alias for planNextAction
runSingleStep :: (Agent a, BaseMemory m) => a -> AgentState m -> IO (Either String AgentStep)
runSingleStep = planNextAction

{- |
Execute a single tool call
Handles tool lookup and input/output conversion.
-}
executeTool :: [AnyTool] -> Text -> Text -> IO (Either String Text)
executeTool tools toolName_ input = do
  case find (\(AnyTool t _ _) -> toolName t == toolName_) tools of
    Nothing -> return $ Left $ "Tool not found: " <> T.unpack toolName_
    Just (AnyTool {..}) -> do
      resultE <- try $ do
        let typedInput = textToInput input
        result <- runTool anyTool typedInput
        return $ outputToText result
      case resultE of
        Left ex -> return $ Left $ "Tool execution error: " <> show (ex :: SomeException)
        Right output -> return $ Right output

{- |
Helper for creating custom tool wrappers
Requires conversion functions between Text and tool-specific types.
-}
customAnyTool :: Tool a => a -> (Text -> Input a) -> (Output a -> Text) -> AnyTool
customAnyTool tool inputConv outputConv = AnyTool tool inputConv outputConv