wip: update agent to use the new ai-sdk

Kujtim Hoxha created 2 months ago

Change summary

agent.go                                | 12 ++++++------
agent_stream_test.go                    |  9 ++++++---
examples/streaming-agent-simple/main.go |  3 ++-
examples/streaming-agent/main.go        |  9 ++++++---
4 files changed, 20 insertions(+), 13 deletions(-)

Detailed changes

agent.go

@@ -174,12 +174,12 @@ type AgentStreamCall struct {
 	RepairToolCall RepairToolCallFunction
 
 	// Agent-level callbacks
-	OnAgentStart  func()                      // Called when agent starts
-	OnAgentFinish func(result *AgentResult)   // Called when agent finishes
-	OnStepStart   func(stepNumber int)        // Called when a step starts
-	OnStepFinish  func(stepResult StepResult) // Called when a step finishes
-	OnFinish      func(result *AgentResult)   // Called when entire agent completes
-	OnError       func(error)                 // Called when an error occurs
+	OnAgentStart  func()                            // Called when agent starts
+	OnAgentFinish func(result *AgentResult) error   // Called when agent finishes
+	OnStepStart   func(stepNumber int) error        // Called when a step starts
+	OnStepFinish  func(stepResult StepResult) error // Called when a step finishes
+	OnFinish      func(result *AgentResult)         // Called when entire agent completes
+	OnError       func(error)                       // Called when an error occurs
 
 	// Stream part callbacks - called for each corresponding stream part type
 	OnChunk          func(StreamPart) error                                                                // Called for each stream part (catch-all)

agent_stream_test.go

@@ -109,14 +109,17 @@ func TestStreamingAgentCallbacks(t *testing.T) {
 		OnAgentStart: func() {
 			callbacks["OnAgentStart"] = true
 		},
-		OnAgentFinish: func(result *AgentResult) {
+		OnAgentFinish: func(result *AgentResult) error {
 			callbacks["OnAgentFinish"] = true
+			return nil
 		},
-		OnStepStart: func(stepNumber int) {
+		OnStepStart: func(stepNumber int) error {
 			callbacks["OnStepStart"] = true
+			return nil
 		},
-		OnStepFinish: func(stepResult StepResult) {
+		OnStepFinish: func(stepResult StepResult) error {
 			callbacks["OnStepFinish"] = true
+			return nil
 		},
 		OnFinish: func(result *AgentResult) {
 			callbacks["OnFinish"] = true

examples/streaming-agent-simple/main.go

@@ -76,8 +76,9 @@ func main() {
 		},
 
 		// Show when each step completes
-		OnStepFinish: func(step ai.StepResult) {
+		OnStepFinish: func(step ai.StepResult) error {
 			fmt.Printf("\n[Step completed: %s]\n", step.FinishReason)
+			return nil
 		},
 	}
 

examples/streaming-agent/main.go

@@ -126,15 +126,18 @@ func main() {
 		OnAgentStart: func() {
 			fmt.Println("🎬 Agent started")
 		},
-		OnAgentFinish: func(result *ai.AgentResult) {
+		OnAgentFinish: func(result *ai.AgentResult) error {
 			fmt.Printf("🏁 Agent finished with %d steps, total tokens: %d\n", len(result.Steps), result.TotalUsage.TotalTokens)
+			return nil
 		},
-		OnStepStart: func(stepNumber int) {
+		OnStepStart: func(stepNumber int) error {
 			stepCount++
 			fmt.Printf("📝 Step %d started\n", stepNumber+1)
+			return nil
 		},
-		OnStepFinish: func(stepResult ai.StepResult) {
+		OnStepFinish: func(stepResult ai.StepResult) error {
 			fmt.Printf("✅ Step completed (reason: %s, tokens: %d)\n", stepResult.FinishReason, stepResult.Usage.TotalTokens)
+			return nil
 		},
 		OnFinish: func(result *ai.AgentResult) {
 			fmt.Printf("🎯 Final result ready with %d steps\n", len(result.Steps))