test: add tests for stream with and without tool calls

Andrey Nering created

Change summary

providertests/provider_test.go                                          | 138 
providertests/testdata/TestStream/anthropic-claude-sonnet.yaml          |  25 
providertests/testdata/TestStream/openai-gpt-4o.yaml                    |  25 
providertests/testdata/TestStreamWithTools/anthropic-claude-sonnet.yaml |  25 
providertests/testdata/TestStreamWithTools/openai-gpt-4o.yaml           |  25 
5 files changed, 238 insertions(+)

Detailed changes

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")
+			}
+		})
+	}
+}

providertests/testdata/TestStream/anthropic-claude-sonnet.yaml 🔗

@@ -0,0 +1,32 @@
+---
+version: 2
+interactions:
+- id: 0
+  request:
+    proto: HTTP/1.1
+    proto_major: 1
+    proto_minor: 1
+    content_length: 227
+    host: ""
+    body: "{\"max_tokens\":4096,\"messages\":[{\"content\":[{\"text\":\"Count from 1 to 3 in Spanish\",\"type\":\"text\"}],\"role\":\"user\"}],\"model\":\"claude-sonnet-4-20250514\",\"system\":[{\"text\":\"You are a helpful assistant\",\"type\":\"text\"}],\"stream\":true}"
+    headers:
+      Accept:
+      - application/json
+      Content-Type:
+      - application/json
+      User-Agent:
+      - Anthropic/Go 1.10.0
+    url: https://api.anthropic.com/v1/messages
+    method: POST
+  response:
+    proto: HTTP/2.0
+    proto_major: 2
+    proto_minor: 0
+    content_length: -1

providertests/testdata/TestStream/openai-gpt-4o.yaml 🔗

@@ -0,0 +1,32 @@
+---
+version: 2
+interactions:
+- id: 0
+  request:
+    proto: HTTP/1.1
+    proto_major: 1
+    proto_minor: 1
+    content_length: 200
+    host: ""
+    body: "{\"messages\":[{\"content\":\"You are a helpful assistant\",\"role\":\"system\"},{\"content\":\"Count from 1 to 3 in Spanish\",\"role\":\"user\"}],\"model\":\"gpt-4o\",\"stream_options\":{\"include_usage\":true},\"stream\":true}"
+    headers:
+      Accept:
+      - application/json
+      Content-Type:
+      - application/json
+      User-Agent:
+      - OpenAI/Go 2.3.0
+    url: https://api.openai.com/v1/chat/completions
+    method: POST
+  response:
+    proto: HTTP/2.0
+    proto_major: 2
+    proto_minor: 0
+    content_length: -1

providertests/testdata/TestStreamWithTools/anthropic-claude-sonnet.yaml 🔗

@@ -0,0 +1,61 @@
+---
+version: 2
+interactions:
+- id: 0
+  request:
+    proto: HTTP/1.1
+    proto_major: 1
+    proto_minor: 1
+    content_length: 553
+    host: ""
+    body: "{\"max_tokens\":4096,\"messages\":[{\"content\":[{\"text\":\"What is 15 + 27?\",\"type\":\"text\"}],\"role\":\"user\"}],\"model\":\"claude-sonnet-4-20250514\",\"system\":[{\"text\":\"You are a helpful assistant. Use the add tool to perform calculations.\",\"type\":\"text\"}],\"tool_choice\":{\"disable_parallel_tool_use\":false,\"type\":\"auto\"},\"tools\":[{\"input_schema\":{\"properties\":{\"a\":{\"description\":\"first number\",\"type\":\"integer\"},\"b\":{\"description\":\"second number\",\"type\":\"integer\"}},\"required\":[\"a\",\"b\"],\"type\":\"object\"},\"name\":\"add\",\"description\":\"Add two numbers\"}],\"stream\":true}"
+    headers:
+      Accept:
+      - application/json
+      Content-Type:
+      - application/json
+      User-Agent:
+      - Anthropic/Go 1.10.0
+    url: https://api.anthropic.com/v1/messages
+    method: POST
+  response:
+    proto: HTTP/2.0
+    proto_major: 2
+    proto_minor: 0
+    content_length: -1

providertests/testdata/TestStreamWithTools/openai-gpt-4o.yaml 🔗

@@ -0,0 +1,61 @@
+---
+version: 2
+interactions:
+- id: 0
+  request:
+    proto: HTTP/1.1
+    proto_major: 1
+    proto_minor: 1
+    content_length: 527
+    host: ""
+    body: "{\"messages\":[{\"content\":\"You are a helpful assistant. Use the add tool to perform calculations.\",\"role\":\"system\"},{\"content\":\"What is 15 + 27?\",\"role\":\"user\"}],\"model\":\"gpt-4o\",\"stream_options\":{\"include_usage\":true},\"tool_choice\":\"auto\",\"tools\":[{\"function\":{\"name\":\"add\",\"strict\":false,\"description\":\"Add two numbers\",\"parameters\":{\"properties\":{\"a\":{\"description\":\"first number\",\"type\":\"integer\"},\"b\":{\"description\":\"second number\",\"type\":\"integer\"}},\"required\":[\"a\",\"b\"],\"type\":\"object\"}},\"type\":\"function\"}],\"stream\":true}"
+    headers:
+      Accept:
+      - application/json
+      Content-Type:
+      - application/json
+      User-Agent:
+      - OpenAI/Go 2.3.0
+    url: https://api.openai.com/v1/chat/completions
+    method: POST
+  response:
+    proto: HTTP/2.0
+    proto_major: 2
+    proto_minor: 0
+    content_length: -1