From b3df0fc9026df52083307c16bce2b3bf4af5f255 Mon Sep 17 00:00:00 2001
From: Steve Dudenhoeffer <steve@stevedudenhoeffer.com>
Date: Tue, 18 Mar 2025 01:03:46 -0400
Subject: [PATCH] Add OnFunctionFinished callback and enhance function handling

Introduces the OnFunctionFinished callback to process function results, errors, and context from OnNewFunction. Updates OnNewFunction to return a result and modifies function calls to include argument handling and post-execution callbacks. Enhances the "answer" function with a new "sources" parameter to provide source tracking for answers.
---
 go.mod               |  2 +-
 pkg/answer/answer.go | 30 ++++++++++++++++++++++++++----
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/go.mod b/go.mod
index 237ab8e..ad374f7 100644
--- a/go.mod
+++ b/go.mod
@@ -8,7 +8,7 @@ replace github.com/rocketlaunchr/google-search => github.com/chrisjoyce911/googl
 
 require (
 	gitea.stevedudenhoeffer.com/steve/go-extractor v0.0.0-20250315044602-7c0e44a22f2c
-	gitea.stevedudenhoeffer.com/steve/go-llm v0.0.0-20250317041832-2737a5b2be93
+	gitea.stevedudenhoeffer.com/steve/go-llm v0.0.0-20250318034148-e5a046a70bd1
 	github.com/Edw590/go-wolfram v0.0.0-20241010091529-fb9031908c5d
 	github.com/advancedlogic/GoOse v0.0.0-20231203033844-ae6b36caf275
 	github.com/joho/godotenv v1.5.1
diff --git a/pkg/answer/answer.go b/pkg/answer/answer.go
index 953e63a..9d10f71 100644
--- a/pkg/answer/answer.go
+++ b/pkg/answer/answer.go
@@ -56,7 +56,13 @@ type Options struct {
 	// OnNewFunction is a callback that, if non-nil, will be called when a new function is called by the LLM.
 	// The "answer" and "no_answer" functions are not included in this callback.
 	// Return an error to stop the function from being called.
-	OnNewFunction func(ctx context.Context, funcName string, question string, parameter string) error
+	OnNewFunction func(ctx context.Context, funcName string, question string, parameter string) (any, error)
+
+	// OnFunctionFinished is a callback that, if non-nil, will be called when a function has finished executing.  The
+	// function name is passed in, as well as the question, the parameter, all similar to OnNewFunction.  The result of
+	// the function is also passed in, as well as any error that occurred.  Finally, the result passed from the
+	// OnNewFunction that preceded this function is passed in as well.
+	OnFunctionFinished func(ctx context.Context, funcName string, question string, parameter string, result string, err error, newFunctionResult any) error
 
 	// SystemPrompt is the prompt to use when asking the system to answer a question.
 	// If this is empty, DefaultPrompt will be used.
@@ -208,7 +214,8 @@ func (o Options) Answer(ctx context.Context, q Question) (string, error) {
 		"answer",
 		"You definitively answer a question, if you call this it means you know the answer and do not need to search for it or use any other function to find it",
 		func(ctx *gollm.Context, args struct {
-			Answer string `json:"answer" description:"the answer to the question"`
+			Answer  string   `json:"answer" description:"the answer to the question"`
+			Sources []string `json:"sources" description:"the sources used to find the answer (e.g.: urls of sites from search)"`
 		}) (string, error) {
 			answer = args.Answer
 			return args.Answer, nil
@@ -354,8 +361,10 @@ func (o Options) Answer(ctx context.Context, q Question) (string, error) {
 
 		for _, call := range choice.Calls {
 			go func(call gollm.ToolCall) {
+				var arg any
+				var err error
 				if o.OnNewFunction != nil {
-					err := o.OnNewFunction(ctx, call.FunctionCall.Name, q.Question, call.FunctionCall.Arguments)
+					arg, err = o.OnNewFunction(ctx, call.FunctionCall.Name, q.Question, call.FunctionCall.Arguments)
 					if err != nil {
 						callsOutput <- gollm.ToolCallResponse{
 							ID:    call.ID,
@@ -364,7 +373,20 @@ func (o Options) Answer(ctx context.Context, q Question) (string, error) {
 						return
 					}
 				}
-				callsOutput <- fnCall(call)
+				callRes := fnCall(call)
+
+				if o.OnFunctionFinished != nil {
+					err = o.OnFunctionFinished(ctx, call.FunctionCall.Name, q.Question, call.FunctionCall.Arguments, callRes.Result, callRes.Error, arg)
+					if err != nil {
+						callsOutput <- gollm.ToolCallResponse{
+							ID:    call.ID,
+							Error: err,
+						}
+						return
+					}
+				}
+
+				callsOutput <- callRes
 			}(call)
 		}