Repository: tmc/langchaingo
Pull request #95 (2023-05-19 14:20:01)

agents: refactor to better match upstream and increase consistency
This refactors the agents and tool packages to better match the typescript and python version and increase consistency with other packages in this version. 

In the current agent package all of the logic for executing the mrkl agent is a part of the mrkl agent. This is not how it is done upstream. In both the typescript and python version there is a executor that inherits from the base chain that does the agent steps and runs the tools. The agent classes are also stateless and take previous steps and the inputs and returns the action of the agent. This therefore moves the execution logic of the mrkl agent to an executor chain in the executor package.

https://github.com/hwchase17/langchainjs/blob/main/langchain/src/agents/executor.ts
https://github.com/hwchase17/langchain/blob/master/langchain/agents/agent.py line 620

Instead of implementing the agent executor interface the mrkl agent then implements the agent interface. 

```go
// Agent is the interface all agents must implement.
type Agent interface {
	// Given an input and previous steps decide what to do next. Returns either actions or a finish.
	Plan(ctx context.Context, intermediateSteps []schema.AgentStep, inputs map[string]string) ([]schema.AgentAction, *schema.AgentFinish, error) //nolint:lll
	GetInputKeys() []string
	GetOutputKeys() []string
}
```

In both the typescript and python version there is different classes for agents doing multiple steps and single steps and a class called agent that inherits from the agent class doing a single step, but I dont think having multiple interfaces for the agents will work well in a Go version. 

The function creating an executor is moved to the executor package and is renamed to "Initialize" to match upstream.

Tool is refactored to be a interface instead of a struct to better match other packages and a calculator tool is also added.

A test running the mrkl agent executor and a test in the mrkl package is added.

### PR Checklist

- [x] Read the [Contributing documentation](https://github.com/tmc/langchaingo/blob/main/CONTRIBUTING.md).
- [x] Read the [Code of conduct documentation](https://github.com/tmc/langchaingo/blob/main/CODE_OF_CONDUCT.md).
- [x] Name your Pull Request title clearly, concisely, and prefixed with the name of primarily affected package you changed according to [Good commit messages](https://go.dev/doc/contribute#commit_messages) (such as `memory: add interfaces for X, Y` or `util: add whizzbang helpers`).
- [x] Check that there isn't already a PR that solves the problem the same way to avoid creating a duplicate.
- [x] Provide a description in this PR that addresses **what** the PR is solving, or reference the issue that it solves (e.g. `fixes #123`).
- [x] Describes source of new concepts.
- [x] References existing implementations as appropriate.
- [x] Contains test coverage for new functions.
- [x] Passes all [`golangci-lint`](https://golangci-lint.run/) checks.
- [x] I am an existing contributor.



Patch:
diff --git a/exp/agent/agents.go b/exp/agent/agents.go
index 32c86454..105438c0 100644
--- a/exp/agent/agents.go
+++ b/exp/agent/agents.go
@@ -2,57 +2,16 @@
 package agent
 
 import (
-	"errors"
+	"context"
 
-	"github.com/tmc/langchaingo/exp/agent/executor"
-	"github.com/tmc/langchaingo/exp/agent/mrkl"
-	"github.com/tmc/langchaingo/exp/tools"
-	"github.com/tmc/langchaingo/llms"
+	"github.com/tmc/langchaingo/schema"
 )
 
-// AgentType is a string type representing the type of agent to create.
-type AgentType string
-
-const (
-	// ZeroShotReactDescription is an AgentType constant that represents the "zeroShotReactDescription" agent type.
-	ZeroShotReactDescription AgentType = "zeroShotReactDescription"
-)
-
-// AgentOption is a type alias for a map of string keys to any value, representing the options for the agent.
-type AgentOption map[string]any
-
-// Options is a function type that can be used to modify the AgentOption.
-type Options func(p *AgentOption)
-
-func defaultOptions() AgentOption {
-	return AgentOption{
-		"verbose": false,
-	}
-}
-
-// WithVerbosity is a function that sets the verbosity option for the agent.
-func WithVerbosity() Options {
-	return func(p *AgentOption) {
-		(*p)["verbose"] = true
-	}
-}
-
-// New is a function that creates a new agent with the specified LLM model, tools, agent type, and options.
-// It returns an AgentExecutor interface or an error if there is any issue during the creation process.
-func New(
-	llm llms.LLM,
-	tools []tools.Tool,
-	agentType AgentType,
-	opts ...Options,
-) (executor.AgentExecutor, error) {
-	options := defaultOptions()
-	for _, opt := range opts {
-		opt(&options)
-	}
-	switch agentType {
-	case ZeroShotReactDescription:
-		return mrkl.NewOneShotAgent(llm, tools, options)
-	default:
-		return nil, errors.New("Unknown agent type")
-	}
+// Agent is the interface all agents must implement.
+type Agent interface {
+	// Given an input and previous steps decide what to do next. Returns
+	// either actions or a finish.
+	Plan(ctx context.Context, intermediateSteps []schema.AgentStep, inputs map[string]string) ([]schema.AgentAction, *schema.AgentFinish, error) //nolint:lll
+	GetInputKeys() []string
+	GetOutputKeys() []string
 }
diff --git a/exp/agent/doc.go b/exp/agent/doc.go
index d5e52f47..9c7c9115 100644
--- a/exp/agent/doc.go
+++ b/exp/agent/doc.go
@@ -1,16 +1,7 @@
-// Package agent provides an implementation of the AgentExecutor interface for
-// langchaingo agents. The package includes the necessary types, structures,
-// and functions to create and configure agents.
+// Package agent contains the standard interface all agents must implement.
 //
-// The agent package provides the New function to create a new agent, and
-// supports different agent types and options to customize the agent's behavior.
-//
-// AgentType is a string type representing the type of agent to create.
-// The package currently supports the "zeroShotReactDescription" agent type.
-//
-// AgentOption is a type alias for a map of string keys to any value,
-// representing the options for the agent.
-//
-// Options is a function type that can be used to modify the AgentOption,
-// such as the WithVerbosity function, which sets the verbosity option for the agent.
+// An Agent is a wrapper around a model, which takes in user input and returns
+// a response corresponding to an “action” to take and a corresponding
+// “action input”. Alternatively the agent can return a finish with the
+// finished answer to the query.
 package agent
diff --git a/exp/agent/executor/agent_executor.go b/exp/agent/executor/agent_executor.go
deleted file mode 100644
index 6bd43b6c..00000000
--- a/exp/agent/executor/agent_executor.go
+++ /dev/null
@@ -1,15 +0,0 @@
-package executor
-
-import (
-	"context"
-
-	"github.com/tmc/langchaingo/schema"
-)
-
-// AgentExecutor is an interface that defines methods for executing agent
-// queries. Implementations of this interface are responsible for executing
-// a query and returning the result in the form of an AgentFinish object
-// or an error if the execution fails.
-type AgentExecutor interface {
-	Run(ctx context.Context, query string) (*schema.AgentFinish, error)
-}
diff --git a/exp/agent/executor/doc.go b/exp/agent/executor/doc.go
index 7604a166..ee9818e5 100644
--- a/exp/agent/executor/doc.go
+++ b/exp/agent/executor/doc.go
@@ -1,19 +1,21 @@
-// Package executor provides interfaces and implementations for executing
-// agent queries.
+// Package executor provides a standard chain for executing agent queries.
 //
-// An agent query is a string that describes a task to be performed by an agent,
-// which is a software program that acts on behalf of a user or another program.
-// The query can specify various parameters, such as input data, processing
-// options, and output format.
+// The Executor is an Agent and set of Tools. The agent executor is
+// responsible for calling the agent, getting back and action and action input,
+// calling the tool that the action references with the corresponding input,
+// getting the output of the tool, and then passing all that information back
+// into the Agent to get the next action it should take.
 //
-// To execute an agent query, an application needs to create an object that
-// implements the AgentExecutor interface, which defines a method for running
-// a query and returning the result. The AgentFinish type from the langchaingo
-// schema package is used to represent the result of a successful execution.
-// If the execution fails, an error is returned.
+// The package also contains functions to initialize executors with agents and
+// supports different agent types and options to customize the agent's behavior.
 //
-// The executor package provides a sample implementation of the AgentExecutor
-// interface, called BasicExecutor, which executes queries using a simple
-// command-line interface. This implementation is intended for testing and
-// demonstration purposes only and should not be used in production environments.
+// AgentType is a string type representing the type of agent to create.
+// The package currently supports the "zeroShotReactDescription" agent type.
+//
+// Options is a type alias for a map of string keys to any value,
+// representing the options for the agent and the executor.
+//
+// Option is a function type that can be used to modify the Options,
+// such as the WithVerbosity function, which sets the verbosity option
+// for the agent.
 package executor
diff --git a/exp/agent/executor/errors.go b/exp/agent/executor/errors.go
new file mode 100644
index 00000000..4f2f6330
--- /dev/null
+++ b/exp/agent/executor/errors.go
@@ -0,0 +1,18 @@
+package executor
+
+import "errors"
+
+var (
+	// ErrExecutorInputNotString is returned if an input to the executor call
+	// function is not a string.
+	ErrExecutorInputNotString = errors.New("input to executor not string")
+	// ErrAgentNoReturn is returned if the agent returns no actions and no
+	// finish.
+	ErrAgentNoReturn = errors.New("no actions or finish was returned by the agent")
+	// ErrNotFinished is returned if the agent does not give a finish before
+	// the number of iterations is larger then max iterations.
+	ErrNotFinished = errors.New("agent not finished before max iterations")
+	// ErrUnknownAgentType is returned if the type given to the initializer is
+	// invalid.
+	ErrUnknownAgentType = errors.New("unknown agent type")
+)
diff --git a/exp/agent/executor/executor.go b/exp/agent/executor/executor.go
new file mode 100644
index 00000000..a05ea175
--- /dev/null
+++ b/exp/agent/executor/executor.go
@@ -0,0 +1,112 @@
+package executor
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/tmc/langchaingo/chains"
+	"github.com/tmc/langchaingo/exp/agent"
+	"github.com/tmc/langchaingo/exp/tools"
+	"github.com/tmc/langchaingo/memory"
+	"github.com/tmc/langchaingo/schema"
+)
+
+// Executor is the chain responsible for running agents.
+type Executor struct {
+	Agent agent.Agent
+	Tools []tools.Tool
+
+	MaxIterations int
+}
+
+var _ chains.Chain = Executor{}
+
+// New creates a new agent executor with a agent, the tools the agent can use
+// and the max number of iterations.
+func New(agent agent.Agent, tools []tools.Tool, maxIterations int) Executor {
+	return Executor{
+		Agent:         agent,
+		Tools:         tools,
+		MaxIterations: maxIterations,
+	}
+}
+
+func (e Executor) Call(ctx context.Context, inputValues map[string]any, _ ...chains.ChainCallOption) (map[string]any, error) { //nolint:lll
+	inputs, err := inputsToString(inputValues)
+	if err != nil {
+		return nil, err
+	}
+
+	nameToTool := make(map[string]tools.Tool, len(e.Tools))
+	for _, tool := range e.Tools {
+		nameToTool[tool.Name()] = tool
+	}
+
+	steps := make([]schema.AgentStep, 0)
+	iterations := 0
+	for iterations < e.MaxIterations {
+		actions, finish, err := e.Agent.Plan(ctx, steps, inputs)
+		if err != nil {
+			return nil, err
+		}
+
+		if len(actions) == 0 && finish == nil {
+			return nil, ErrAgentNoReturn
+		}
+
+		if finish != nil {
+			return finish.ReturnValues, nil
+		}
+
+		for _, action := range actions {
+			tool, ok := nameToTool[action.Tool]
+			if !ok {
+				steps = append(steps, schema.AgentStep{
+					Action:      action,
+					Observation: fmt.Sprintf("%s is not a valid tool, try another one", action.Tool),
+				})
+			}
+
+			observation, err := tool.Call(action.ToolInput)
+			if err != nil {
+				return nil, err
+			}
+
+			steps = append(steps, schema.AgentStep{
+				Action:      action,
+				Observation: observation,
+			})
+		}
+	}
+
+	return nil, ErrNotFinished
+}
+
+// GetInputKeys gets the input keys the agent of the executor expects.
+// Often "input".
+func (e Executor) GetInputKeys() []string {
+	return e.Agent.GetInputKeys()
+}
+
+// GetOutputKeys gets the output keys the agent of the executor returns.
+func (e Executor) GetOutputKeys() []string {
+	return e.Agent.GetOutputKeys()
+}
+
+func (e Executor) GetMemory() schema.Memory {
+	return memory.NewSimple()
+}
+
+func inputsToString(inputValues map[string]any) (map[string]string, error) {
+	inputs := make(map[string]string, len(inputValues))
+	for key, value := range inputValues {
+		valueStr, ok := value.(string)
+		if !ok {
+			return nil, fmt.Errorf("%w: %s", ErrExecutorInputNotString, key)
+		}
+
+		inputs[key] = valueStr
+	}
+
+	return inputs, nil
+}
diff --git a/exp/agent/executor/initialize.go b/exp/agent/executor/initialize.go
new file mode 100644
index 00000000..b541ce48
--- /dev/null
+++ b/exp/agent/executor/initialize.go
@@ -0,0 +1,75 @@
+package executor
+
+import (
+	"github.com/tmc/langchaingo/exp/agent"
+	"github.com/tmc/langchaingo/exp/agent/mrkl"
+	"github.com/tmc/langchaingo/exp/tools"
+	"github.com/tmc/langchaingo/llms"
+)
+
+// AgentType is a string type representing the type of agent to create.
+type AgentType string
+
+const (
+	// ZeroShotReactDescription is an AgentType constant that represents
+	// the "zeroShotReactDescription" agent type.
+	ZeroShotReactDescription AgentType = "zeroShotReactDescription"
+)
+
+// options is a type alias for a map of string keys to any value,
+// representing the options for the agent and the executor.
+type options map[string]any
+
+// Option is a function type that can be used to modify the creation of the
+// executor and agent.
+type Option func(p *options)
+
+func defaultOptions() options {
+	return options{
+		"verbose":       false,
+		"maxIterations": 3,
+	}
+}
+
+// WithVerbosity is a function that sets the verbosity option for the agent.
+func WithVerbosity() Option {
+	return func(p *options) {
+		(*p)["verbose"] = true
+	}
+}
+
+// WithMaxIterations is a function that sets the max iterations for the executor.
+func WithMaxIterations(maxIterations int) Option {
+	return func(p *options) {
+		(*p)["maxIterations"] = maxIterations
+	}
+}
+
+// Initialize is a function that creates a new executor with the specified LLM
+// model, tools, agent type, and options. It returns an Executor or an error
+// if there is any issue during the creation process.
+func Initialize(
+	llm llms.LLM,
+	tools []tools.Tool,
+	agentType AgentType,
+	opts ...Option,
+) (Executor, error) {
+	options := defaultOptions()
+	for _, opt := range opts {
+		opt(&options)
+	}
+	var agent agent.Agent
+
+	switch agentType {
+	case ZeroShotReactDescription:
+		agent = mrkl.NewOneShotAgent(llm, tools, options)
+	default:
+		return Executor{}, ErrUnknownAgentType
+	}
+
+	return Executor{
+		Agent:         agent,
+		Tools:         tools,
+		MaxIterations: options["maxIterations"].(int),
+	}, nil
+}
diff --git a/exp/agent/mrkl/doc.go b/exp/agent/mrkl/doc.go
index 2008a76a..e6155740 100644
--- a/exp/agent/mrkl/doc.go
+++ b/exp/agent/mrkl/doc.go
@@ -1,11 +1,4 @@
-// Package mrkl provides an implementation of the OneShotZeroAgent,
-// which is responsible for executing queries and returning results using
-// the LLM model, tools, and an internal chain. The package also includes
-// related structures and methods to support the agent's functionality.
-//
-// The OneShotZeroAgent struct implements the AgentExecutor interface
-// and provides methods to run queries and manage the agent's internal state.
-//
-// The package also includes utility functions for checking options,
-// planning, and running tools to perform various operations during query execution.
+// Package mrkl provides and implementation of the agent interface that
+// uses the ReAct Framework (based on the descriptions of tools) to
+// decide what action to take. This agent is optimized to be used with LLMs.
 package mrkl
diff --git a/exp/agent/mrkl/mrkl.go b/exp/agent/mrkl/mrkl.go
index 25fc9dbb..cfa970eb 100644
--- a/exp/agent/mrkl/mrkl.go
+++ b/exp/agent/mrkl/mrkl.go
@@ -4,203 +4,154 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"regexp"
 	"strings"
 	"time"
 
 	"github.com/tmc/langchaingo/chains"
-	"github.com/tmc/langchaingo/exp/agent/executor"
+	"github.com/tmc/langchaingo/exp/agent"
 	"github.com/tmc/langchaingo/exp/tools"
 	"github.com/tmc/langchaingo/llms"
 	"github.com/tmc/langchaingo/schema"
 )
 
-// OneShotZeroAgent is a struct that represents an agent responsible for executing a query
-// and returning the result using the LLM model, tools, and an internal chain.
+var (
+	// ErrUnableToParseOutput is returned if the output of the llm is unparsable.
+	ErrUnableToParseOutput = errors.New("unable to parse agent output")
+	// ErrInvalidChainReturnType is returned if the internal chain of the agent
+	// returns a value in the "text" filed that is not a string.
+	ErrInvalidChainReturnType = errors.New("agent chain did not return a string")
+	// ErrInvalidOptions is returned if the options given to the NewOneShotAgent
+	// function is invalid.
+	ErrInvalidOptions = errors.New("options given are invalid")
+)
+
+const (
+	_finalAnswerAction = "Final Answer:"
+	_defaultOutputKey  = "output"
+)
+
+// OneShotZeroAgent is a struct that represents an agent responsible for deciding
+// what to do or give the final output if the task is finished given a set of inputs
+// and previous steps taken.
+//
+// This agent is optimized to be used with LLMs.
 type OneShotZeroAgent struct {
-	llm        llms.LLM
-	query      string
-	chain      chains.Chain
-	tools      []tools.Tool
-	verbose    bool
-	maxRetries int
+	// Chain is the chain used to call with the values. The chain should have an
+	// input called "agent_scratchpad" for the agent to put it's thoughts in.
+	Chain chains.Chain
+	// Tools is a list of the tools the agent can use.
+	Tools []tools.Tool
+	// Output key is the key where the final output is placed.
+	OutputKey string
 }
 
+var _ agent.Agent = (*OneShotZeroAgent)(nil)
+
 // OneShotZeroAgentOptions is a type alias for a map of string keys to any value,
 // representing the options for the OneShotZeroAgent.
 type OneShotZeroAgentOptions map[string]any
 
-var _ executor.AgentExecutor = (*OneShotZeroAgent)(nil)
-
-const FinalAnswerAction = "Final Answer:"
-
 func checkOptions(opts OneShotZeroAgentOptions) OneShotZeroAgentOptions {
-	if _, ok := opts["verbose"].(bool); !ok {
-		opts["verbose"] = false
-	}
-	if _, ok := opts["maxRetries"].(int); !ok {
-		opts["maxRetries"] = 3
+	if _, ok := opts["outputKey"].(string); !ok {
+		opts["outputKey"] = _defaultOutputKey
 	}
 	return opts
 }
 
 // NewOneShotAgent creates a new OneShotZeroAgent with the given LLM model, tools,
-// and options. It returns a pointer to the created agent and an error if there is any
-// issue during the creation process.
-func NewOneShotAgent(llm llms.LLM, tools []tools.Tool, opts map[string]any) (*OneShotZeroAgent, error) {
-	// Validate opts
+// and options. It returns a pointer to the created agent.
+func NewOneShotAgent(llm llms.LLM, tools []tools.Tool, opts map[string]any) *OneShotZeroAgent {
+	// Validate opts.
 	opts = checkOptions(opts)
 
 	return &OneShotZeroAgent{
-		llm:        llm,
-		query:      "",
-		chain:      chains.NewLLMChain(llm, createPrompt()),
-		tools:      tools,
-		verbose:    opts["verbose"].(bool),
-		maxRetries: opts["maxRetries"].(int),
-	}, nil
-}
-
-// Run is an implementation of the AgentExecutor interface. It takes a query as input
-// and executes it, returning an AgentFinish object containing the result, or an error
-// if the execution fails.
-func (a *OneShotZeroAgent) Run(ctx context.Context, query string) (*schema.AgentFinish, error) {
-	var attempts int
-	a.query = query
-
-	// Call the chain
-	resp, _ := a.chain.Call(ctx, map[string]interface{}{
-		"today":             time.Now().Format("January 02, 2006"),
-		"tool_names":        toolNames(a.tools),
-		"tool_descriptions": toolDescriptions(a.tools),
-		"input":             a.query,
-		"agent_scratchpad":  "",
-		"stop":              []string{"\nObservation:", "\n\tObservation:"},
-	})
-
-	// Validate the response
-	output, ok := resp["text"].(string)
-	if !ok {
-		return nil, errors.New("Agent did not return a string")
+		Chain:     chains.NewLLMChain(llm, CreatePrompt(tools)),
+		Tools:     tools,
+		OutputKey: opts["outputKey"].(string),
 	}
+}
 
-	for output != "" || attempts < a.maxRetries {
-		var err error
-		action, finish := a.plan(output)
-		if finish != nil {
-			return finish, nil
-		}
-		output, err = a.nextStep(ctx, *action)
-		if err != nil {
-			return nil, err
-		}
-
-		attempts++
+// Plan decides what action to take or returns the final result of the input.
+func (a *OneShotZeroAgent) Plan(
+	ctx context.Context,
+	intermediateSteps []schema.AgentStep,
+	inputs map[string]string,
+) ([]schema.AgentAction, *schema.AgentFinish, error) {
+	fullInputs := make(map[string]any, len(inputs))
+	for key, value := range inputs {
+		fullInputs[key] = value
 	}
 
-	return nil, fmt.Errorf("Agent did not finish after %d attempts", attempts)
-}
+	fullInputs["agent_scratchpad"] = a.constructScratchPad(intermediateSteps)
+	fullInputs["today"] = time.Now().Format("January 02, 2006")
 
-func (a *OneShotZeroAgent) nextStep(ctx context.Context, action schema.AgentAction) (string, error) {
-	var scratchpad []string
-	// Perform your desired operation with the text value
-	observation, err := runTool(action.Tool, action.ToolInput.(string), &a.tools)
+	resp, err := chains.Call(
+		ctx,
+		a.Chain,
+		fullInputs,
+		chains.WithStopWords([]string{"\nObservation:", "\n\tObservation:"}),
+	)
 	if err != nil {
-		return "", err
-	}
-	scratchpad = append(scratchpad, action.Log+observation)
-	if a.verbose {
-		fmt.Println(getCurrentThought(scratchpad))
+		return nil, nil, err
 	}
 
-	// Update resp using a.chain.Call()
-	newResp, err := a.chain.Call(ctx, map[string]interface{}{
-		"input":            a.query,
-		"agent_scratchpad": strings.Join(scratchpad, "\n"),
-		"stop":             []string{"\nObservation:", "\n\tObservation:"},
-	})
-	if err != nil {
-		return "", err
+	output, ok := resp["text"].(string)
+	if !ok {
+		return nil, nil, ErrInvalidChainReturnType
 	}
 
-	// Use the updated resp in the next iteration
-	return newResp["text"].(string), nil
+	return a.parseOutput(output)
 }
 
-func (a *OneShotZeroAgent) plan(info string) (*schema.AgentAction, *schema.AgentFinish) {
-	action := getAgentAction(info)
-	if aswer := getFinalAnswer(action.Log); aswer != "" {
-		return nil, &schema.AgentFinish{
-			ReturnValues: map[string]any{
-				"answer": aswer,
-			},
-			Log: action.Log,
-		}
-	}
-	return &action, nil
-}
+func (a *OneShotZeroAgent) GetInputKeys() []string {
+	chainInputs := a.Chain.GetInputKeys()
 
-func getCurrentThought(scratchpad []string) string {
-	if len(scratchpad) == 0 {
-		return ""
-	}
-	lastThought := scratchpad[len(scratchpad)-1]
-	if len(scratchpad) == 1 {
-		lastThought = "Thought:" + lastThought
+	// Remove inputs given in plan.
+	agentInput := make([]string, 0, len(chainInputs))
+	for _, v := range chainInputs {
+		if v == "agent_scratchpad" || v == "today" {
+			continue
+		}
+		agentInput = append(agentInput, v)
 	}
 
-	return lastThought
+	return agentInput
 }
 
-func getFinalAnswer(text string) string {
-	finalAnswerPrefix := "Final Answer:"
-	startIndex := strings.Index(text, finalAnswerPrefix)
-
-	if startIndex == -1 {
-		return ""
-	}
-	startIndex += len(finalAnswerPrefix)
-	text = text[startIndex:]
-	trimmed := strings.Join(strings.Fields(text), " ")
-	return trimmed
+func (a *OneShotZeroAgent) GetOutputKeys() []string {
+	return []string{a.OutputKey}
 }
 
-func getAgentAction(input string) schema.AgentAction {
-	var agentAction schema.AgentAction
-	agentAction.Log = input
-	fields := strings.Split(input, "\n")
-	for _, field := range fields {
-		if strings.HasPrefix(field, "Action: ") {
-			agentAction.Tool = strings.TrimPrefix(field, "Action: ")
-		} else if strings.HasPrefix(field, "Action Input: ") {
-			agentAction.ToolInput = strings.TrimPrefix(field, "Action Input: ")
-		}
+func (a *OneShotZeroAgent) constructScratchPad(steps []schema.AgentStep) string {
+	var scratchPad string
+	for _, step := range steps {
+		scratchPad += step.Action.Log
+		scratchPad += "Observation: " + step.Observation
 	}
-	return agentAction
-}
-
-func runTool(action string, actionInput string, tools *[]tools.Tool) (string, error) {
-	// Sanitize the action
-	action = strings.ToLower(strings.Trim(action, " "))
 
-	// Sanitize the action input
-	actionInput = strings.TrimSpace(actionInput)
+	return scratchPad
+}
 
-	// Find the tool that matches the action
-	var observation string
-	for _, tool := range *tools {
-		if tool.Name != strings.Trim(action, " ") {
-			continue
-		}
+func (a *OneShotZeroAgent) parseOutput(output string) ([]schema.AgentAction, *schema.AgentFinish, error) {
+	if strings.Contains(output, _finalAnswerAction) {
+		splits := strings.Split(output, _finalAnswerAction)
 
-		// Run the tool
-		toolOutput, err := tool.Run(actionInput)
-		if err != nil {
-			return "", err
-		}
+		return nil, &schema.AgentFinish{
+			ReturnValues: map[string]any{
+				a.OutputKey: splits[len(splits)-1],
+			},
+			Log: output,
+		}, nil
+	}
 
-		// Add the tool's output to the observation
-		observation = "\nObservation: " + toolOutput + "\n"
-		break
+	r := regexp.MustCompile(`Action:\s*(.+)\s*Action Input:\s*(.+)`)
+	matches := r.FindStringSubmatch(output)
+	if len(matches) == 0 {
+		return nil, nil, fmt.Errorf("%w: %s", ErrUnableToParseOutput, output)
 	}
-	return observation, nil
+
+	return []schema.AgentAction{
+		{Tool: strings.TrimSpace(matches[1]), ToolInput: strings.TrimSpace(matches[2]), Log: output},
+	}, nil, nil
 }
diff --git a/exp/agent/mrkl/prompt.go b/exp/agent/mrkl/prompt.go
index b004a1db..178338b5 100644
--- a/exp/agent/mrkl/prompt.go
+++ b/exp/agent/mrkl/prompt.go
@@ -90,7 +90,7 @@ func toolNames(tools []tools.Tool) string {
 		if i > 0 {
 			tn.WriteString(", ")
 		}
-		tn.WriteString(tool.Name)
+		tn.WriteString(tool.Name())
 	}
 
 	return tn.String()
@@ -99,17 +99,16 @@ func toolNames(tools []tools.Tool) string {
 func toolDescriptions(tools []tools.Tool) string {
 	var ts strings.Builder
 	for _, tool := range tools {
-		ts.WriteString(fmt.Sprintf("- %s: %s\n", tool.Name, tool.Description))
+		ts.WriteString(fmt.Sprintf("- %s: %s\n", tool.Name(), tool.Description()))
 	}
 
 	return ts.String()
 }
 
-// createPrompt is a function that takes a slice of tools and a variadic list of prompt
-// template options, and returns a prompt template with the specified options.
-func createPrompt(
-	options ...PromptTemplateOption,
-) prompts.PromptTemplate {
+// CreatePrompt is a function that takes a slice of tools and a variadic list of prompt
+// template options, and returns a prompt template used by the agent with the specified
+// options.
+func CreatePrompt(tools []tools.Tool, options ...PromptTemplateOption) prompts.PromptTemplate {
 	opts := &CreatePromptOptions{}
 	for _, option := range options {
 		option(opts)
@@ -121,6 +120,10 @@ func createPrompt(
 	return prompts.PromptTemplate{
 		Template:       template,
 		TemplateFormat: prompts.TemplateFormatGoTemplate,
-		InputVariables: []string{"input", "tools", "tool_names", "tools_list", "agent_scratchpad"},
+		InputVariables: []string{"input", "agent_scratchpad", "today"},
+		PartialVariables: map[string]any{
+			"tool_names":        toolNames(tools),
+			"tool_descriptions": toolDescriptions(tools),
+		},
 	}
 }
diff --git a/exp/tools/calculator.go b/exp/tools/calculator.go
new file mode 100644
index 00000000..1fe23c4b
--- /dev/null
+++ b/exp/tools/calculator.go
@@ -0,0 +1,33 @@
+package tools
+
+import (
+	"go.starlark.net/lib/math"
+	"go.starlark.net/starlark"
+)
+
+// Calculator is a tool that can do math.
+type Calculator struct{}
+
+var _ Tool = Calculator{}
+
+// Description returns a string describing the calculator tool.
+func (c Calculator) Description() string {
+	return `Useful for getting the result of a math expression. 
+	The input to this tool should be a valid mathematical expression that could be executed by a starlark evaluator.`
+}
+
+// Name returns the name of the tool.
+func (c Calculator) Name() string {
+	return "calculator"
+}
+
+// Call evaluates the input using a starlak evaluator and returns the
+// result as a string.
+func (c Calculator) Call(input string) (string, error) {
+	v, err := starlark.Eval(&starlark.Thread{Name: "main"}, "input", input, math.Module.Members)
+	if err != nil {
+		return "I don't know how to do that.", nil
+	}
+
+	return v.String(), nil
+}
diff --git a/exp/tools/doc.go b/exp/tools/doc.go
new file mode 100644
index 00000000..f6eae88d
--- /dev/null
+++ b/exp/tools/doc.go
@@ -0,0 +1,2 @@
+// Package tools defines a standard interface for tools to be used by agents.
+package tools
diff --git a/exp/tools/serpapi/internal/serpapi_client.go b/exp/tools/serpapi/internal/serpapi_client.go
index 63781e58..7ce3b72a 100644
--- a/exp/tools/serpapi/internal/serpapi_client.go
+++ b/exp/tools/serpapi/internal/serpapi_client.go
@@ -8,13 +8,12 @@ import (
 	"log"
 	"net/http"
 	"net/url"
-	"os"
 	"strings"
 )
 
 var (
 	ErrMissingToken = errors.New("missing the OpenAI API key, set it in the SERPAPI_API_KEY environment variable")
-	NoGoodResult    = "No good search result found"
+	ErrNoGoodResult = errors.New("no good search results found")
 )
 
 type SerpapiClient struct {
@@ -22,17 +21,11 @@ type SerpapiClient struct {
 	baseURL string
 }
 
-func New() (*SerpapiClient, error) {
-	apiKey := os.Getenv("SERPAPI_API_KEY")
-	if apiKey == "" {
-		return nil, ErrMissingToken
-	}
-
+func New(apiKey string) *SerpapiClient {
 	return &SerpapiClient{
-			apiKey:  apiKey,
-			baseURL: "https://serpapi.com/search",
-		},
-		nil
+		apiKey:  apiKey,
+		baseURL: "https://serpapi.com/search",
+	}
 }
 
 func (s *SerpapiClient) Search(query string) (string, error) {
@@ -107,5 +100,5 @@ func processResponse(res map[string]interface{}) (string, error) {
 		}
 	}
 
-	return NoGoodResult, nil
+	return "", ErrNoGoodResult
 }
diff --git a/exp/tools/serpapi/serpapi.go b/exp/tools/serpapi/serpapi.go
new file mode 100644
index 00000000..c2fe4c11
--- /dev/null
+++ b/exp/tools/serpapi/serpapi.go
@@ -0,0 +1,55 @@
+package serpapi
+
+import (
+	"errors"
+	"os"
+	"strings"
+
+	"github.com/tmc/langchaingo/exp/tools"
+	"github.com/tmc/langchaingo/exp/tools/serpapi/internal"
+)
+
+var ErrMissingToken = errors.New("missing the serpapi API key, set it in the SERPAPI_API_KEY environment variable")
+
+type Tool struct {
+	client *internal.SerpapiClient
+}
+
+var _ tools.Tool = Tool{}
+
+// New creates a new serpapi tool to search on internet.
+func New() (*Tool, error) {
+	apiKey := os.Getenv("SERPAPI_API_KEY")
+	if apiKey == "" {
+		return nil, ErrMissingToken
+	}
+
+	return &Tool{
+		client: internal.New(apiKey),
+	}, nil
+}
+
+func (t Tool) Name() string {
+	return "Google Search"
+}
+
+func (t Tool) Description() string {
+	return `
+	"A wrapper around Google Search. "
+	"Useful for when you need to answer questions about current events. "
+	"Always one of the first options when you need to find information on internet"
+	"Input should be a search query."`
+}
+
+func (t Tool) Call(input string) (string, error) {
+	result, err := t.client.Search(input)
+	if err != nil {
+		if errors.Is(err, internal.ErrNoGoodResult) {
+			return "No good Google Search Result was found", nil
+		}
+
+		return "", err
+	}
+
+	return strings.Join(strings.Fields(result), " "), nil
+}
diff --git a/exp/tools/serpapi/serpapi_tool.go b/exp/tools/serpapi/serpapi_tool.go
deleted file mode 100644
index 62815135..00000000
--- a/exp/tools/serpapi/serpapi_tool.go
+++ /dev/null
@@ -1,33 +0,0 @@
-package serpapi
-
-import (
-	"strings"
-
-	"github.com/tmc/langchaingo/exp/tools"
-	"github.com/tmc/langchaingo/exp/tools/serpapi/internal"
-)
-
-// Create a new tool for serpapi to search on internet.
-func New() (*tools.Tool, error) {
-	client, err := internal.New()
-	if err != nil {
-		return nil, err
-	}
-	return &tools.Tool{
-		Name: "Google Search",
-		Description: `"A wrapper around Google Search. "
-        "Useful for when you need to answer questions about current events. "
-		"Always one of the first options when you need to find information on internet"
-        "Input should be a search query."`,
-		Run: func(input string) (string, error) {
-			result, err := client.Search(input)
-			if err != nil {
-				return "", err
-			}
-			if len(result) == 0 {
-				return "No good Google Search Result was found", nil
-			}
-			return strings.Join(strings.Fields(result), " "), nil
-		},
-	}, nil
-}
diff --git a/exp/tools/tool.go b/exp/tools/tool.go
index c1e9a51e..9ff769ca 100644
--- a/exp/tools/tool.go
+++ b/exp/tools/tool.go
@@ -1,17 +1,8 @@
-// Package tools defines the types for tools to be used by the llms agents.
 package tools
 
-// Tool is a tool for the llm agent to interact with different application.
-type Tool struct {
-	Name        string
-	Description string
-	Run         func(string) (string, error)
-}
-
-func NewTool(name, description string, run func(string) (string, error)) *Tool {
-	return &Tool{
-		Name:        name,
-		Description: description,
-		Run:         run,
-	}
+// Tool is a tool for the llm agent to interact with different applications.
+type Tool interface {
+	Name() string
+	Description() string
+	Call(string) (string, error)
 }
diff --git a/schema/schema.go b/schema/schema.go
index e93adae7..0529529c 100644
--- a/schema/schema.go
+++ b/schema/schema.go
@@ -3,10 +3,16 @@ package schema
 // AgentAction is the agent's action to take.
 type AgentAction struct {
 	Tool      string
-	ToolInput any
+	ToolInput string
 	Log       string
 }
 
+// AgentStep is a step of the agent.
+type AgentStep struct {
+	Action      AgentAction
+	Observation string
+}
+
 // AgentFinish is the agent's return value.
 type AgentFinish struct {
 	ReturnValues map[string]any


Test patch:
diff --git a/exp/agent/executor/executor_test.go b/exp/agent/executor/executor_test.go
new file mode 100644
index 00000000..147f7b99
--- /dev/null
+++ b/exp/agent/executor/executor_test.go
@@ -0,0 +1,44 @@
+package executor_test
+
+import (
+	"context"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/tmc/langchaingo/chains"
+	"github.com/tmc/langchaingo/exp/agent/executor"
+	"github.com/tmc/langchaingo/exp/tools"
+	"github.com/tmc/langchaingo/exp/tools/serpapi"
+	"github.com/tmc/langchaingo/llms/openai"
+)
+
+func TestMRKL(t *testing.T) {
+	if openaiKey := os.Getenv("OPENAI_API_KEY"); openaiKey == "" {
+		t.Skip("OPENAI_API_KEY not set")
+	}
+	if serpapiKey := os.Getenv("SERPAPI_API_KEY"); serpapiKey == "" {
+		t.Skip("SERPAPI_API_KEY not set")
+	}
+
+	llm, err := openai.New()
+	require.NoError(t, err)
+
+	searchTool, err := serpapi.New()
+	require.NoError(t, err)
+
+	calculator := tools.Calculator{}
+
+	a, err := executor.Initialize(
+		llm,
+		[]tools.Tool{searchTool, calculator},
+		executor.ZeroShotReactDescription,
+	)
+	require.NoError(t, err)
+
+	result, err := chains.Run(context.Background(), a, "If a person lived three times as long as Jacklyn Zeman, how long would they live") //nolint:lll
+	require.NoError(t, err)
+
+	require.True(t, strings.Contains(result, "210"), "correct answer 210 not in response")
+}
diff --git a/exp/agent/mrkl/markl_test.go b/exp/agent/mrkl/markl_test.go
new file mode 100644
index 00000000..3323274a
--- /dev/null
+++ b/exp/agent/mrkl/markl_test.go
@@ -0,0 +1,38 @@
+package mrkl
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/tmc/langchaingo/schema"
+)
+
+func TestOutputParser(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		input           string
+		expectedActions []schema.AgentAction
+		expectedFinish  *schema.AgentFinish
+		expectedErr     error
+	}{
+		{
+			input: "Action:  foo Action Input: bar",
+			expectedActions: []schema.AgentAction{{
+				Tool:      "foo",
+				ToolInput: "bar",
+				Log:       "Action:  foo Action Input: bar",
+			}},
+			expectedFinish: nil,
+			expectedErr:    nil,
+		},
+	}
+
+	a := OneShotZeroAgent{}
+	for _, tc := range testCases {
+		actions, finish, err := a.parseOutput(tc.input)
+		require.ErrorIs(t, tc.expectedErr, err)
+		require.Equal(t, tc.expectedActions, actions)
+		require.Equal(t, tc.expectedFinish, finish)
+	}
+}