From 135207caa300618cdf7af628d4841b86d0800aec Mon Sep 17 00:00:00 2001
From: lujingxuan <lujingxuan@bytedance.com>
Date: Tue, 2 Jan 2024 15:02:31 +0800
Subject: [PATCH] fix initialize

---
 agents/executor_test.go          | 12 +++++++-----
 agents/initialize.go             |  2 --
 agents/openai_functions_agent.go |  7 ++++---
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/agents/executor_test.go b/agents/executor_test.go
index 670e5675..bd9861b1 100644
--- a/agents/executor_test.go
+++ b/agents/executor_test.go
@@ -119,18 +119,20 @@ func TestExecutorWithOpenAIFunctionAgent(t *testing.T) {
 
 	calculator := tools.Calculator{}
 
-	a, err := agents.Initialize(
-		llm,
-		[]tools.Tool{searchTool, calculator},
-		agents.OpenAIFunctionAgentDescription,
+	toolList := []tools.Tool{searchTool, calculator}
+
+	a := agents.NewOpenAIFunctionsAgent(llm,
+		toolList,
 		agents.NewOpenAIOption().WithSystemMessage("you are a helpful assistant"),
 		agents.NewOpenAIOption().WithExtraMessages([]prompts.MessageFormatter{
 			prompts.NewHumanMessagePromptTemplate("please be strict", nil),
 		}),
 	)
+
+	e := agents.NewExecutor(a, toolList)
 	require.NoError(t, err)
 
-	result, err := chains.Run(context.Background(), a, "what is HK singer Eason Chan's years old?") //nolint:lll
+	result, err := chains.Run(context.Background(), e, "what is HK singer Eason Chan's years old?") //nolint:lll
 	require.NoError(t, err)
 
 	require.True(t, strings.Contains(result, "47") || strings.Contains(result, "49"),
diff --git a/agents/initialize.go b/agents/initialize.go
index af121d62..fe509707 100644
--- a/agents/initialize.go
+++ b/agents/initialize.go
@@ -37,8 +37,6 @@ func Initialize(
 		agent = NewOneShotAgent(llm, tools, opts...)
 	case ConversationalReactDescription:
 		agent = NewConversationalAgent(llm, tools, opts...)
-	case OpenAIFunctionAgentDescription:
-		agent = NewOpenAIFunctionsAgent(llm, tools, opts...)
 	default:
 		return Executor{}, ErrUnknownAgentType
 	}
diff --git a/agents/openai_functions_agent.go b/agents/openai_functions_agent.go
index e4bec358..ba7b3c60 100644
--- a/agents/openai_functions_agent.go
+++ b/agents/openai_functions_agent.go
@@ -16,7 +16,7 @@ import (
 type OpenAIFunctionsAgent struct {
 	// LLM is the llm used to call with the values. The llm should have an
 	// input called "agent_scratchpad" for the agent to put its thoughts in.
-	LLM    llms.LanguageModel
+	LLM    llms.ChatLLM
 	Prompt prompts.FormatPrompter
 	// Chain chains.Chain
 	// Tools is a list of the tools the agent can use.
@@ -30,7 +30,7 @@ type OpenAIFunctionsAgent struct {
 var _ Agent = (*OpenAIFunctionsAgent)(nil)
 
 // NewOpenAIFunctionsAgent creates a new OpenAIFunctionsAgent.
-func NewOpenAIFunctionsAgent(llm llms.LanguageModel, tools []tools.Tool, opts ...CreationOption) *OpenAIFunctionsAgent {
+func NewOpenAIFunctionsAgent(llm llms.ChatLLM, tools []tools.Tool, opts ...CreationOption) *OpenAIFunctionsAgent {
 	options := openAIFunctionsDefaultOptions()
 	for _, opt := range opts {
 		opt(&options)
@@ -89,8 +89,9 @@ func (o *OpenAIFunctionsAgent) Plan(
 		return nil, nil, err
 	}
 
-	result, err := o.LLM.GeneratePrompt(
+	result, err := llms.GenerateChatPrompt(
 		ctx,
+		o.LLM,
 		[]schema.PromptValue{prompt},
 		llms.WithFunctions(o.functions()),
 		llms.WithStreamingFunc(stream),