support callback handler to llm chain (tmc#461)

* support callback handler to llm chain * fix golint --------- Co-authored-by: lvzhong <[email protected]>
t0mpl · Dec 29, 2023 · 9887348 · 9887348
1 parent bf59430
commit 9887348
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 9 deletions.
diff --git a/agents/conversational.go b/agents/conversational.go
@@ -45,7 +45,11 @@ func NewConversationalAgent(llm llms.LanguageModel, tools []tools.Tool, opts ...
 	}
 
 	return &ConversationalAgent{
-		Chain:            chains.NewLLMChain(llm, options.getConversationalPrompt(tools)),
+		Chain: chains.NewLLMChain(
+			llm,
+			options.getConversationalPrompt(tools),
+			chains.WithCallback(options.callbacksHandler),
+		),
 		Tools:            tools,
 		OutputKey:        options.outputKey,
 		CallbacksHandler: options.callbacksHandler,

diff --git a/agents/mrkl.go b/agents/mrkl.go
@@ -48,7 +48,11 @@ func NewOneShotAgent(llm llms.LanguageModel, tools []tools.Tool, opts ...Creatio
 	}
 
 	return &OneShotZeroAgent{
-		Chain:            chains.NewLLMChain(llm, options.getMrklPrompt(tools)),
+		Chain: chains.NewLLMChain(
+			llm,
+			options.getMrklPrompt(tools),
+			chains.WithCallback(options.callbacksHandler),
+		),
 		Tools:            tools,
 		OutputKey:        options.outputKey,
 		CallbacksHandler: options.callbacksHandler,

diff --git a/chains/llm.go b/chains/llm.go
@@ -29,14 +29,18 @@ var (
 )
 
 // NewLLMChain creates a new LLMChain with an LLM and a prompt.
-func NewLLMChain(llm llms.LanguageModel, prompt prompts.FormatPrompter) *LLMChain {
+func NewLLMChain(llm llms.LanguageModel, prompt prompts.FormatPrompter, opts ...ChainCallOption) *LLMChain {
+	opt := &chainCallOption{}
+	for _, o := range opts {
+		o(opt)
+	}
 	chain := &LLMChain{
-		Prompt:       prompt,
-		LLM:          llm,
-		OutputParser: outputparser.NewSimple(),
-		Memory:       memory.NewSimple(),
-
-		OutputKey: _llmChainDefaultOutputKey,
+		Prompt:           prompt,
+		LLM:              llm,
+		OutputParser:     outputparser.NewSimple(),
+		Memory:           memory.NewSimple(),
+		OutputKey:        _llmChainDefaultOutputKey,
+		CallbacksHandler: opt.CallbackHandler,
 	}
 
 	return chain