agents: add CallOptions to Plan() (issue #353)

.gitignore

@@ -19,3 +19,4 @@ service-account.json
 
 embeddings/cybertron/models/*
 examples/cybertron-embedding-example/models/*
+!/.idea/

agents/agents.go

@@ -3,6 +3,7 @@ package agents
 import (
 	"context"
 
+	"github.com/tmc/langchaingo/chains"
 	"github.com/tmc/langchaingo/schema"
 	"github.com/tmc/langchaingo/tools"
 )
@@ -11,7 +12,7 @@ import (
 type Agent interface {
 	// Plan 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
+	Plan(ctx context.Context, intermediateSteps []schema.AgentStep, inputs map[string]string, options ...chains.ChainCallOption) ([]schema.AgentAction, *schema.AgentFinish, error) //nolint:lll
 	GetInputKeys() []string
 	GetOutputKeys() []string
 	GetTools() []tools.Tool

agents/conversational.go

@@ -63,6 +63,7 @@ func (a *ConversationalAgent) Plan(
 	ctx context.Context,
 	intermediateSteps []schema.AgentStep,
 	inputs map[string]string,
+	options ...chains.ChainCallOption,
 ) ([]schema.AgentAction, *schema.AgentFinish, error) {
 	fullInputs := make(map[string]any, len(inputs))
 	for key, value := range inputs {
@@ -80,12 +81,14 @@ func (a *ConversationalAgent) Plan(
 		}
 	}
 
+	options = append(options, chains.WithStopWords([]string{"\nObservation:", "\n\tObservation:"}))
+	options = append(options, chains.WithStreamingFunc(stream))
+
 	output, err := chains.Predict(
 		ctx,
 		a.Chain,
 		fullInputs,
-		chains.WithStopWords([]string{"\nObservation:", "\n\tObservation:"}),
-		chains.WithStreamingFunc(stream),
+		options...,
 	)
 	if err != nil {
 		return nil, nil, err

agents/conversational_test.go

@@ -22,13 +22,10 @@ func TestConversationalWithMemory(t *testing.T) {
 	llm, err := openai.New(openai.WithModel("gpt-4"))
 	require.NoError(t, err)
 
-	executor, err := Initialize(
-		llm,
-		[]tools.Tool{tools.Calculator{}},
-		ConversationalReactDescription,
+	executor := NewExecutor(
+		NewConversationalAgent(llm, []tools.Tool{tools.Calculator{}}),
 		WithMemory(memory.NewConversationBuffer()),
 	)
-	require.NoError(t, err)
 
 	_, err = chains.Run(context.Background(), executor, "Hi! my name is Bob and the year I was born is 1987")
 	require.NoError(t, err)

agents/executor.go

@@ -47,7 +47,7 @@ func NewExecutor(agent Agent, opts ...Option) *Executor {
 	}
 }
 
-func (e *Executor) Call(ctx context.Context, inputValues map[string]any, _ ...chains.ChainCallOption) (map[string]any, error) { //nolint:lll
+func (e *Executor) Call(ctx context.Context, inputValues map[string]any, options ...chains.ChainCallOption) (map[string]any, error) { //nolint:lll
 	inputs, err := inputsToString(inputValues)
 	if err != nil {
 		return nil, err
@@ -57,7 +57,7 @@ func (e *Executor) Call(ctx context.Context, inputValues map[string]any, _ ...ch
 	steps := make([]schema.AgentStep, 0)
 	for i := 0; i < e.MaxIterations; i++ {
 		var finish map[string]any
-		steps, finish, err = e.doIteration(ctx, steps, nameToTool, inputs)
+		steps, finish, err = e.doIteration(ctx, steps, nameToTool, inputs, options...)
 		if finish != nil || err != nil {
 			return finish, err
 		}
@@ -79,8 +79,9 @@ func (e *Executor) doIteration( // nolint
 	steps []schema.AgentStep,
 	nameToTool map[string]tools.Tool,
 	inputs map[string]string,
+	options ...chains.ChainCallOption,
 ) ([]schema.AgentStep, map[string]any, error) {
-	actions, finish, err := e.Agent.Plan(ctx, steps, inputs)
+	actions, finish, err := e.Agent.Plan(ctx, steps, inputs, options...)
 	if errors.Is(err, ErrUnableToParseOutput) && e.ErrorHandler != nil {
 		formattedObservation := err.Error()
 		if e.ErrorHandler.Formatter != nil {

agents/executor_test.go

@@ -32,6 +32,7 @@ func (a *testAgent) Plan(
 	_ context.Context,
 	intermediateSteps []schema.AgentStep,
 	inputs map[string]string,
+	_ ...chains.ChainCallOption,
 ) ([]schema.AgentAction, *schema.AgentFinish, error) {
 	a.recordedIntermediateSteps = intermediateSteps
 	a.recordedInputs = inputs
@@ -90,15 +91,11 @@ func TestExecutorWithMRKLAgent(t *testing.T) {
 	require.NoError(t, err)
 
 	calculator := tools.Calculator{}
-
-	a, err := agents.Initialize(
-		llm,
-		[]tools.Tool{searchTool, calculator},
-		agents.ZeroShotReactDescription,
+	executor := agents.NewExecutor(
+		agents.NewOneShotAgent(llm, []tools.Tool{searchTool, calculator}),
 	)
-	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
+	result, err := chains.Run(context.Background(), executor, "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")

agents/mrkl.go

@@ -64,6 +64,7 @@ func (a *OneShotZeroAgent) Plan(
 	ctx context.Context,
 	intermediateSteps []schema.AgentStep,
 	inputs map[string]string,
+	options ...chains.ChainCallOption,
 ) ([]schema.AgentAction, *schema.AgentFinish, error) {
 	fullInputs := make(map[string]any, len(inputs))
 	for key, value := range inputs {
@@ -82,12 +83,14 @@ func (a *OneShotZeroAgent) Plan(
 		}
 	}
 
+	options = append(options, chains.WithStopWords([]string{"\nObservation:", "\n\tObservation:"}))
+	options = append(options, chains.WithStreamingFunc(stream))
+
 	output, err := chains.Predict(
 		ctx,
 		a.Chain,
 		fullInputs,
-		chains.WithStopWords([]string{"\nObservation:", "\n\tObservation:"}),
-		chains.WithStreamingFunc(stream),
+		options...,
 	)
 	if err != nil {
 		return nil, nil, err

agents/openai_functions_agent.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 
 	"github.com/tmc/langchaingo/callbacks"
+	"github.com/tmc/langchaingo/chains"
 	"github.com/tmc/langchaingo/llms"
 	"github.com/tmc/langchaingo/prompts"
 	"github.com/tmc/langchaingo/schema"
@@ -67,10 +68,12 @@ func (o *OpenAIFunctionsAgent) functions() []llms.FunctionDefinition {
 }
 
 // Plan decides what action to take or returns the final result of the input.
+// TODO(fix): add {options ...chains.ChainCallOption} to llm request.
 func (o *OpenAIFunctionsAgent) Plan(
 	ctx context.Context,
 	intermediateSteps []schema.AgentStep,
 	inputs map[string]string,
+	_ ...chains.ChainCallOption,
 ) ([]schema.AgentAction, *schema.AgentFinish, error) {
 	fullInputs := make(map[string]any, len(inputs))
 	for key, value := range inputs {