test: add tests for stream with and without tool calls

Author: andreynering

providertests/provider_test.go

@@ -2,6 +2,7 @@ package providertests
 
 import (
 	"context"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -82,3 +83,140 @@ func TestTool(t *testing.T) {
 		})
 	}
 }
+
+func TestStream(t *testing.T) {
+	for _, pair := range languageModelBuilders {
+		t.Run(pair.name, func(t *testing.T) {
+			r := newRecorder(t)
+
+			languageModel, err := pair.builder(r)
+			if err != nil {
+				t.Fatalf("failed to build language model: %v", err)
+			}
+
+			agent := ai.NewAgent(
+				languageModel,
+				ai.WithSystemPrompt("You are a helpful assistant"),
+			)
+
+			var collectedText strings.Builder
+			textDeltaCount := 0
+			stepCount := 0
+
+			streamCall := ai.AgentStreamCall{
+				Prompt: "Count from 1 to 3 in Spanish",
+				OnTextDelta: func(id, text string) error {
+					textDeltaCount++
+					collectedText.WriteString(text)
+					return nil
+				},
+				OnStepFinish: func(step ai.StepResult) error {
+					stepCount++
+					return nil
+				},
+			}
+
+			result, err := agent.Stream(t.Context(), streamCall)
+			if err != nil {
+				t.Fatalf("failed to stream: %v", err)
+			}
+
+			finalText := result.Response.Content.Text()
+			if finalText == "" {
+				t.Fatal("expected non-empty response")
+			}
+
+			if !strings.Contains(strings.ToLower(finalText), "uno") ||
+				!strings.Contains(strings.ToLower(finalText), "dos") ||
+				!strings.Contains(strings.ToLower(finalText), "tres") {
+				t.Fatalf("unexpected response: %q", finalText)
+			}
+
+			if textDeltaCount == 0 {
+				t.Fatal("expected at least one text delta callback")
+			}
+
+			if stepCount == 0 {
+				t.Fatal("expected at least one step finish callback")
+			}
+
+			if collectedText.String() == "" {
+				t.Fatal("expected collected text from deltas to be non-empty")
+			}
+		})
+	}
+}
+
+func TestStreamWithTools(t *testing.T) {
+	for _, pair := range languageModelBuilders {
+		t.Run(pair.name, func(t *testing.T) {
+			r := newRecorder(t)
+
+			languageModel, err := pair.builder(r)
+			if err != nil {
+				t.Fatalf("failed to build language model: %v", err)
+			}
+
+			type CalculatorInput struct {
+				A int `json:"a" description:"first number"`
+				B int `json:"b" description:"second number"`
+			}
+
+			calculatorTool := ai.NewAgentTool(
+				"add",
+				"Add two numbers",
+				func(ctx context.Context, input CalculatorInput, _ ai.ToolCall) (ai.ToolResponse, error) {
+					result := input.A + input.B
+					return ai.NewTextResponse(strings.TrimSpace(strconv.Itoa(result))), nil
+				},
+			)
+
+			agent := ai.NewAgent(
+				languageModel,
+				ai.WithSystemPrompt("You are a helpful assistant. Use the add tool to perform calculations."),
+				ai.WithTools(calculatorTool),
+			)
+
+			toolCallCount := 0
+			toolResultCount := 0
+			var collectedText strings.Builder
+
+			streamCall := ai.AgentStreamCall{
+				Prompt: "What is 15 + 27?",
+				OnTextDelta: func(id, text string) error {
+					collectedText.WriteString(text)
+					return nil
+				},
+				OnToolCall: func(toolCall ai.ToolCallContent) error {
+					toolCallCount++
+					if toolCall.ToolName != "add" {
+						t.Errorf("unexpected tool name: %s", toolCall.ToolName)
+					}
+					return nil
+				},
+				OnToolResult: func(result ai.ToolResultContent) error {
+					toolResultCount++
+					return nil
+				},
+			}
+
+			result, err := agent.Stream(t.Context(), streamCall)
+			if err != nil {
+				t.Fatalf("failed to stream: %v", err)
+			}
+
+			finalText := result.Response.Content.Text()
+			if !strings.Contains(finalText, "42") {
+				t.Fatalf("expected response to contain '42', got: %q", finalText)
+			}
+
+			if toolCallCount == 0 {
+				t.Fatal("expected at least one tool call")
+			}
+
+			if toolResultCount == 0 {
+				t.Fatal("expected at least one tool result")
+			}
+		})
+	}
+}