feat(job): add finalizers and save conversation after job is result is released (#114)

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

core/agent/agent.go

@@ -504,7 +504,9 @@ func (a *Agent) consumeJob(job *types.Job, role string) {
 
 		xlog.Debug("Finish job with reasoning", "reasoning", reasoning, "agent", a.Character.Name, "conversation", fmt.Sprintf("%+v", conv))
 		job.Result.Conversation = conv
-		a.saveCurrentConversation(conv)
+		job.Result.AddFinalizer(func(conv []openai.ChatCompletionMessage) {
+			a.saveCurrentConversation(conv)
+		})
 		job.Result.SetResponse(reasoning)
 		job.Result.Finish(nil)
 		return
@@ -669,9 +671,11 @@ func (a *Agent) consumeJob(job *types.Job, role string) {
 				}
 
 				conv = append(conv, msg)
-				a.saveCurrentConversation(conv)
 				job.Result.SetResponse(msg.Content)
 				job.Result.Conversation = conv
+				job.Result.AddFinalizer(func(conv []openai.ChatCompletionMessage) {
+					a.saveCurrentConversation(conv)
+				})
 				job.Result.Finish(nil)
 				return
 			}
@@ -751,7 +755,9 @@ func (a *Agent) consumeJob(job *types.Job, role string) {
 		conv = append(conv, msg)
 		job.Result.Conversation = conv
 		job.Result.SetResponse(msg.Content)
-		a.saveCurrentConversation(conv)
+		job.Result.AddFinalizer(func(conv []openai.ChatCompletionMessage) {
+			a.saveCurrentConversation(conv)
+		})
 		job.Result.Finish(nil)
 		return
 	}
@@ -780,7 +786,9 @@ func (a *Agent) consumeJob(job *types.Job, role string) {
 	job.Result.SetResponse(msg.Content)
 	xlog.Info("Response from LLM", "response", msg.Content, "agent", a.Character.Name)
 	job.Result.Conversation = conv
-	a.saveCurrentConversation(conv)
+	job.Result.AddFinalizer(func(conv []openai.ChatCompletionMessage) {
+		a.saveCurrentConversation(conv)
+	})
 	job.Result.Finish(nil)
 }
 

core/types/job.go

@@ -3,7 +3,6 @@ package types
 import (
 	"context"
 	"log"
-	"sync"
 
 	"github.com/google/uuid"
 	"github.com/sashabaranov/go-openai"
@@ -29,18 +28,6 @@ type Job struct {
 	cancel  context.CancelFunc
 }
 
-// JobResult is the result of a job
-type JobResult struct {
-	sync.Mutex
-	// The result of a job
-	State        []ActionState
-	Conversation []openai.ChatCompletionMessage
-
-	Response string
-	Error    error
-	ready    chan bool
-}
-
 type JobOption func(*Job)
 
 func WithConversationHistory(history []openai.ChatCompletionMessage) JobOption {

core/types/result.go

@@ -1,5 +1,25 @@
 package types
 
+import (
+	"sync"
+
+	"github.com/sashabaranov/go-openai"
+)
+
+// JobResult is the result of a job
+type JobResult struct {
+	sync.Mutex
+	// The result of a job
+	State        []ActionState
+	Conversation []openai.ChatCompletionMessage
+
+	Finalizers []func([]openai.ChatCompletionMessage)
+
+	Response string
+	Error    error
+	ready    chan bool
+}
+
 // SetResult sets the result of a job
 func (j *JobResult) SetResult(text ActionState) {
 	j.Lock()
@@ -10,11 +30,24 @@ func (j *JobResult) SetResult(text ActionState) {
 
 // SetResult sets the result of a job
 func (j *JobResult) Finish(e error) {
+	j.Lock()
+	j.Error = e
+	j.Unlock()
+
+	close(j.ready)
+
+	for _, f := range j.Finalizers {
+		f(j.Conversation)
+	}
+	j.Finalizers = []func([]openai.ChatCompletionMessage){}
+}
+
+// AddFinalizer adds a finalizer to the job result
+func (j *JobResult) AddFinalizer(f func([]openai.ChatCompletionMessage)) {
 	j.Lock()
 	defer j.Unlock()
 
-	j.Error = e
-	close(j.ready)
+	j.Finalizers = append(j.Finalizers, f)
 }
 
 // SetResult sets the result of a job