fix(tmc#353)(agents): add ChainCallOptions to Plan

disk0Dancer · May 23, 2024 · dae44c3 · dae44c3
1 parent 4c509c0
commit dae44c3
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 6 deletions.
diff --git a/agents/agents.go b/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

diff --git a/agents/conversational.go b/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

diff --git a/agents/executor.go b/agents/executor.go
@@ -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 {

diff --git a/agents/executor_test.go b/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

diff --git a/agents/mrkl.go b/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

diff --git a/agents/openai_functions_agent.go b/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 {