697 lines
19 KiB
Go
697 lines
19 KiB
Go
// Copyright (c) 2026 Probo Inc <hello@probo.com>.
|
|
//
|
|
// Permission to use, copy, modify, and/or distribute this software for any
|
|
// purpose with or without fee is hereby granted, provided that the above
|
|
// copyright notice and this permission notice appear in all copies.
|
|
//
|
|
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
|
|
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
|
|
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
|
|
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
|
|
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
|
|
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
|
|
// PERFORMANCE OF THIS SOFTWARE.
|
|
|
|
package llm_test
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"go.opentelemetry.io/otel/codes"
|
|
sdktrace "go.opentelemetry.io/otel/sdk/trace"
|
|
"go.opentelemetry.io/otel/sdk/trace/tracetest"
|
|
"go.probo.inc/probo/pkg/llm"
|
|
)
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Helpers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
type mockProvider struct {
|
|
chatResp *llm.ChatCompletionResponse
|
|
chatErr error
|
|
streamResp llm.ChatCompletionStream
|
|
streamErr error
|
|
}
|
|
|
|
func (m *mockProvider) ChatCompletion(_ context.Context, _ *llm.ChatCompletionRequest) (*llm.ChatCompletionResponse, error) {
|
|
return m.chatResp, m.chatErr
|
|
}
|
|
|
|
func (m *mockProvider) ChatCompletionStream(_ context.Context, _ *llm.ChatCompletionRequest) (llm.ChatCompletionStream, error) {
|
|
return m.streamResp, m.streamErr
|
|
}
|
|
|
|
type mockStream struct {
|
|
events []llm.ChatCompletionStreamEvent
|
|
idx int
|
|
current llm.ChatCompletionStreamEvent
|
|
err error
|
|
closeErr error
|
|
closed bool
|
|
}
|
|
|
|
func (s *mockStream) Next() bool {
|
|
if s.idx >= len(s.events) {
|
|
return false
|
|
}
|
|
|
|
s.current = s.events[s.idx]
|
|
s.idx++
|
|
|
|
return true
|
|
}
|
|
|
|
func (s *mockStream) Event() llm.ChatCompletionStreamEvent { return s.current }
|
|
func (s *mockStream) Err() error { return s.err }
|
|
func (s *mockStream) Close() error { s.closed = true; return s.closeErr }
|
|
|
|
func newTestClient(provider llm.Provider) (*llm.Client, *tracetest.SpanRecorder) {
|
|
recorder := tracetest.NewSpanRecorder()
|
|
tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
|
|
|
|
client := llm.NewClient(
|
|
provider,
|
|
"test",
|
|
llm.WithTracerProvider(tp),
|
|
)
|
|
|
|
return client, recorder
|
|
}
|
|
|
|
func spanAttrMap(recorder *tracetest.SpanRecorder) map[string]any {
|
|
spans := recorder.Ended()
|
|
if len(spans) == 0 {
|
|
return nil
|
|
}
|
|
|
|
m := make(map[string]any)
|
|
for _, a := range spans[0].Attributes() {
|
|
m[string(a.Key)] = a.Value.AsInterface()
|
|
}
|
|
|
|
return m
|
|
}
|
|
|
|
//go:fix inline
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Message.Text
|
|
// ---------------------------------------------------------------------------
|
|
|
|
func TestMessageText(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
msg llm.Message
|
|
want string
|
|
}{
|
|
{
|
|
name: "single text part",
|
|
msg: llm.Message{
|
|
Parts: []llm.Part{llm.TextPart{Text: "hello"}},
|
|
},
|
|
want: "hello",
|
|
},
|
|
{
|
|
name: "multiple text parts concatenated",
|
|
msg: llm.Message{
|
|
Parts: []llm.Part{
|
|
llm.TextPart{Text: "hello"},
|
|
llm.TextPart{Text: " world"},
|
|
},
|
|
},
|
|
want: "hello world",
|
|
},
|
|
{
|
|
name: "image parts skipped",
|
|
msg: llm.Message{
|
|
Parts: []llm.Part{
|
|
llm.TextPart{Text: "before"},
|
|
llm.ImagePart{URL: "http://example.com/img.png"},
|
|
llm.TextPart{Text: "after"},
|
|
},
|
|
},
|
|
want: "beforeafter",
|
|
},
|
|
{
|
|
name: "no parts returns empty string",
|
|
msg: llm.Message{},
|
|
want: "",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
assert.Equal(t, tt.want, tt.msg.Text())
|
|
})
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Usage.Add
|
|
// ---------------------------------------------------------------------------
|
|
|
|
func TestUsageAdd(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
a := llm.Usage{InputTokens: 10, OutputTokens: 5}
|
|
b := llm.Usage{InputTokens: 20, OutputTokens: 15}
|
|
c := a.Add(b)
|
|
|
|
assert.Equal(t, 30, c.InputTokens)
|
|
assert.Equal(t, 20, c.OutputTokens)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Error types
|
|
// ---------------------------------------------------------------------------
|
|
|
|
func TestErrors(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
inner := errors.New("upstream")
|
|
|
|
t.Run("ErrRateLimit", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
t.Run("with retry after", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := &llm.ErrRateLimit{RetryAfter: 30 * time.Second, Err: inner}
|
|
assert.Contains(t, e.Error(), "retry after 30s")
|
|
assert.Contains(t, e.Error(), "upstream")
|
|
assert.ErrorIs(t, e, inner)
|
|
})
|
|
|
|
t.Run("without retry after", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := &llm.ErrRateLimit{Err: inner}
|
|
assert.Contains(t, e.Error(), "rate limited")
|
|
assert.NotContains(t, e.Error(), "retry after")
|
|
assert.ErrorIs(t, e, inner)
|
|
})
|
|
|
|
t.Run("errors.As", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var target *llm.ErrRateLimit
|
|
|
|
e := &llm.ErrRateLimit{RetryAfter: 5 * time.Second, Err: inner}
|
|
require.ErrorAs(t, e, &target)
|
|
assert.Equal(t, 5*time.Second, target.RetryAfter)
|
|
})
|
|
})
|
|
|
|
t.Run("ErrContextLength", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
t.Run("with max tokens", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := &llm.ErrContextLength{MaxTokens: 4096, Err: inner}
|
|
assert.Contains(t, e.Error(), "4096")
|
|
assert.ErrorIs(t, e, inner)
|
|
})
|
|
|
|
t.Run("without max tokens", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := &llm.ErrContextLength{Err: inner}
|
|
assert.Contains(t, e.Error(), "context length exceeded")
|
|
assert.NotContains(t, e.Error(), "max")
|
|
assert.ErrorIs(t, e, inner)
|
|
})
|
|
})
|
|
|
|
t.Run("ErrContentFilter", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := &llm.ErrContentFilter{Err: inner}
|
|
assert.Contains(t, e.Error(), "content filtered")
|
|
assert.ErrorIs(t, e, inner)
|
|
})
|
|
|
|
t.Run("ErrAuthentication", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
e := &llm.ErrAuthentication{Err: inner}
|
|
assert.Contains(t, e.Error(), "authentication failed")
|
|
assert.ErrorIs(t, e, inner)
|
|
})
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Client — ChatCompletion
|
|
// ---------------------------------------------------------------------------
|
|
|
|
func TestChatCompletion(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
t.Run("success", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
provider := &mockProvider{
|
|
chatResp: &llm.ChatCompletionResponse{
|
|
Model: "test-model",
|
|
Message: llm.Message{
|
|
Role: llm.RoleAssistant,
|
|
Parts: []llm.Part{llm.TextPart{Text: "Hello!"}},
|
|
},
|
|
Usage: llm.Usage{InputTokens: 10, OutputTokens: 5},
|
|
FinishReason: llm.FinishReasonStop,
|
|
},
|
|
}
|
|
|
|
client, recorder := newTestClient(provider)
|
|
temp := 0.7
|
|
resp, err := client.ChatCompletion(context.Background(), &llm.ChatCompletionRequest{
|
|
Model: "test-model",
|
|
Messages: []llm.Message{
|
|
{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "Hi"}}},
|
|
},
|
|
Temperature: &temp,
|
|
})
|
|
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "test-model", resp.Model)
|
|
assert.Equal(t, "Hello!", resp.Message.Text())
|
|
assert.Equal(t, 10, resp.Usage.InputTokens)
|
|
assert.Equal(t, 5, resp.Usage.OutputTokens)
|
|
assert.Equal(t, llm.FinishReasonStop, resp.FinishReason)
|
|
|
|
spans := recorder.Ended()
|
|
require.Len(t, spans, 1)
|
|
assert.Equal(t, "chat test-model", spans[0].Name())
|
|
|
|
attrs := spanAttrMap(recorder)
|
|
assert.Equal(t, "test", attrs["gen_ai.provider.name"])
|
|
assert.Equal(t, "test-model", attrs["gen_ai.request.model"])
|
|
assert.Equal(t, 0.7, attrs["gen_ai.request.temperature"])
|
|
assert.Equal(t, "test-model", attrs["gen_ai.response.model"])
|
|
assert.Equal(t, int64(10), attrs["gen_ai.usage.input_tokens"])
|
|
assert.Equal(t, int64(5), attrs["gen_ai.usage.output_tokens"])
|
|
})
|
|
|
|
t.Run("all span attributes", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
provider := &mockProvider{
|
|
chatResp: &llm.ChatCompletionResponse{
|
|
Model: "gpt-4",
|
|
FinishReason: llm.FinishReasonStop,
|
|
Usage: llm.Usage{InputTokens: 50, OutputTokens: 25},
|
|
Message: llm.Message{
|
|
Role: llm.RoleAssistant,
|
|
Parts: []llm.Part{llm.TextPart{Text: "ok"}},
|
|
},
|
|
},
|
|
}
|
|
|
|
client, recorder := newTestClient(provider)
|
|
maxTokens := 1024
|
|
topP := 0.9
|
|
resp, err := client.ChatCompletion(context.Background(), &llm.ChatCompletionRequest{
|
|
Model: "gpt-4",
|
|
Messages: []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "test"}}}},
|
|
MaxTokens: &maxTokens,
|
|
TopP: &topP,
|
|
StopSequences: []string{"END", "STOP"},
|
|
})
|
|
|
|
require.NoError(t, err)
|
|
require.NotNil(t, resp)
|
|
|
|
attrs := spanAttrMap(recorder)
|
|
assert.Equal(t, int64(1024), attrs["gen_ai.request.max_tokens"])
|
|
assert.Equal(t, 0.9, attrs["gen_ai.request.top_p"])
|
|
assert.Equal(t, []string{"END", "STOP"}, attrs["gen_ai.request.stop_sequences"])
|
|
})
|
|
|
|
t.Run("error sets span status", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
provider := &mockProvider{
|
|
chatErr: errors.New("provider error"),
|
|
}
|
|
|
|
client, recorder := newTestClient(provider)
|
|
_, err := client.ChatCompletion(context.Background(), &llm.ChatCompletionRequest{
|
|
Model: "test-model",
|
|
Messages: []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "Hi"}}}},
|
|
})
|
|
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "provider error")
|
|
|
|
spans := recorder.Ended()
|
|
require.Len(t, spans, 1)
|
|
assert.Equal(t, codes.Error, spans[0].Status().Code)
|
|
})
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Client — ChatCompletionStream
|
|
// ---------------------------------------------------------------------------
|
|
|
|
func TestChatCompletionStream(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
t.Run("success collects deltas", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
events := []llm.ChatCompletionStreamEvent{
|
|
{Delta: llm.MessageDelta{Content: "Hello"}},
|
|
{Delta: llm.MessageDelta{Content: " world"}},
|
|
{
|
|
FinishReason: new(llm.FinishReasonStop),
|
|
Usage: &llm.Usage{InputTokens: 8, OutputTokens: 4},
|
|
},
|
|
}
|
|
|
|
client, recorder := newTestClient(&mockProvider{
|
|
streamResp: &mockStream{events: events},
|
|
})
|
|
stream, err := client.ChatCompletionStream(context.Background(), &llm.ChatCompletionRequest{
|
|
Model: "test-model",
|
|
Messages: []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "Hi"}}}},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
var collected []string
|
|
|
|
for stream.Next() {
|
|
e := stream.Event()
|
|
if e.Delta.Content != "" {
|
|
collected = append(collected, e.Delta.Content)
|
|
}
|
|
}
|
|
|
|
require.NoError(t, stream.Err())
|
|
require.NoError(t, stream.Close())
|
|
|
|
assert.Equal(t, []string{"Hello", " world"}, collected)
|
|
|
|
spans := recorder.Ended()
|
|
require.Len(t, spans, 1)
|
|
assert.Equal(t, "chat test-model", spans[0].Name())
|
|
})
|
|
|
|
t.Run("provider error sets span status", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
client, recorder := newTestClient(&mockProvider{
|
|
streamErr: errors.New("stream open failed"),
|
|
})
|
|
_, err := client.ChatCompletionStream(context.Background(), &llm.ChatCompletionRequest{
|
|
Model: "test-model",
|
|
Messages: []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "Hi"}}}},
|
|
})
|
|
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "stream open failed")
|
|
|
|
spans := recorder.Ended()
|
|
require.Len(t, spans, 1)
|
|
assert.Equal(t, codes.Error, spans[0].Status().Code)
|
|
})
|
|
|
|
t.Run("inner stream error records error on span", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ms := &mockStream{
|
|
events: []llm.ChatCompletionStreamEvent{
|
|
{Delta: llm.MessageDelta{Content: "partial"}},
|
|
},
|
|
err: errors.New("connection reset"),
|
|
}
|
|
|
|
client, recorder := newTestClient(&mockProvider{streamResp: ms})
|
|
stream, err := client.ChatCompletionStream(context.Background(), &llm.ChatCompletionRequest{
|
|
Model: "test-model",
|
|
Messages: []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "Hi"}}}},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
for stream.Next() {
|
|
}
|
|
|
|
assert.ErrorContains(t, stream.Err(), "connection reset")
|
|
_ = stream.Close()
|
|
|
|
spans := recorder.Ended()
|
|
require.Len(t, spans, 1)
|
|
assert.Equal(t, codes.Error, spans[0].Status().Code)
|
|
})
|
|
|
|
t.Run("span finalized on close before exhausting", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
events := []llm.ChatCompletionStreamEvent{
|
|
{Delta: llm.MessageDelta{Content: "a"}},
|
|
{Delta: llm.MessageDelta{Content: "b"}},
|
|
{Delta: llm.MessageDelta{Content: "c"}},
|
|
}
|
|
|
|
client, recorder := newTestClient(&mockProvider{
|
|
streamResp: &mockStream{events: events},
|
|
})
|
|
stream, err := client.ChatCompletionStream(context.Background(), &llm.ChatCompletionRequest{
|
|
Model: "test-model",
|
|
Messages: []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "Hi"}}}},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Read only one event, then close early.
|
|
require.True(t, stream.Next())
|
|
require.NoError(t, stream.Close())
|
|
|
|
spans := recorder.Ended()
|
|
require.Len(t, spans, 1, "span should be ended by Close even without exhausting stream")
|
|
})
|
|
|
|
t.Run("close error traced on span", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ms := &mockStream{
|
|
events: []llm.ChatCompletionStreamEvent{
|
|
{Delta: llm.MessageDelta{Content: "partial"}},
|
|
},
|
|
closeErr: errors.New("broken pipe"),
|
|
}
|
|
|
|
client, recorder := newTestClient(&mockProvider{streamResp: ms})
|
|
stream, err := client.ChatCompletionStream(context.Background(), &llm.ChatCompletionRequest{
|
|
Model: "test-model",
|
|
Messages: []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "Hi"}}}},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
require.True(t, stream.Next())
|
|
closeErr := stream.Close()
|
|
require.ErrorContains(t, closeErr, "broken pipe")
|
|
|
|
spans := recorder.Ended()
|
|
require.Len(t, spans, 1)
|
|
assert.Equal(t, codes.Error, spans[0].Status().Code)
|
|
assert.Contains(t, spans[0].Status().Description, "broken pipe")
|
|
})
|
|
|
|
t.Run("stream span records usage and finish reason", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
events := []llm.ChatCompletionStreamEvent{
|
|
{Delta: llm.MessageDelta{Content: "done"}},
|
|
{
|
|
FinishReason: new(llm.FinishReasonLength),
|
|
Usage: &llm.Usage{InputTokens: 100, OutputTokens: 50},
|
|
},
|
|
}
|
|
|
|
client, recorder := newTestClient(&mockProvider{
|
|
streamResp: &mockStream{events: events},
|
|
})
|
|
stream, err := client.ChatCompletionStream(context.Background(), &llm.ChatCompletionRequest{
|
|
Model: "test-model",
|
|
Messages: []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "Hi"}}}},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
for stream.Next() {
|
|
}
|
|
|
|
require.NoError(t, stream.Err())
|
|
require.NoError(t, stream.Close())
|
|
|
|
spans := recorder.Ended()
|
|
require.Len(t, spans, 1)
|
|
|
|
attrs := make(map[string]any)
|
|
for _, a := range spans[0].Attributes() {
|
|
attrs[string(a.Key)] = a.Value.AsInterface()
|
|
}
|
|
|
|
assert.Equal(t, int64(100), attrs["gen_ai.usage.input_tokens"])
|
|
assert.Equal(t, int64(50), attrs["gen_ai.usage.output_tokens"])
|
|
assert.Equal(t, []string{"length"}, attrs["gen_ai.response.finish_reasons"])
|
|
})
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// StreamAccumulator
|
|
// ---------------------------------------------------------------------------
|
|
|
|
func TestStreamAccumulator(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
t.Run("text and single tool call", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
events := []llm.ChatCompletionStreamEvent{
|
|
{Model: "gpt-4o", Delta: llm.MessageDelta{Content: "Hello"}},
|
|
{Delta: llm.MessageDelta{Content: " world"}},
|
|
{Delta: llm.MessageDelta{
|
|
ToolCalls: []llm.ToolCallDelta{
|
|
{Index: 0, ID: "tc_1", Name: "get_weather"},
|
|
},
|
|
}},
|
|
{Delta: llm.MessageDelta{
|
|
ToolCalls: []llm.ToolCallDelta{
|
|
{Index: 0, Arguments: `{"city":`},
|
|
},
|
|
}},
|
|
{Delta: llm.MessageDelta{
|
|
ToolCalls: []llm.ToolCallDelta{
|
|
{Index: 0, Arguments: `"Paris"}`},
|
|
},
|
|
}},
|
|
{
|
|
FinishReason: new(llm.FinishReasonToolCalls),
|
|
Usage: &llm.Usage{InputTokens: 20, OutputTokens: 15},
|
|
},
|
|
}
|
|
|
|
acc := llm.NewStreamAccumulator(&mockStream{events: events})
|
|
for acc.Next() {
|
|
}
|
|
|
|
require.NoError(t, acc.Err())
|
|
|
|
resp := acc.Response()
|
|
assert.Equal(t, "gpt-4o", resp.Model)
|
|
assert.Equal(t, "Hello world", resp.Message.Text())
|
|
assert.Equal(t, llm.RoleAssistant, resp.Message.Role)
|
|
assert.Equal(t, llm.FinishReasonToolCalls, resp.FinishReason)
|
|
assert.Equal(t, 20, resp.Usage.InputTokens)
|
|
assert.Equal(t, 15, resp.Usage.OutputTokens)
|
|
|
|
require.Len(t, resp.Message.ToolCalls, 1)
|
|
tc := resp.Message.ToolCalls[0]
|
|
assert.Equal(t, "tc_1", tc.ID)
|
|
assert.Equal(t, "get_weather", tc.Function.Name)
|
|
assert.Equal(t, `{"city":"Paris"}`, tc.Function.Arguments)
|
|
})
|
|
|
|
t.Run("multiple tool calls at different indices", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
events := []llm.ChatCompletionStreamEvent{
|
|
{Delta: llm.MessageDelta{
|
|
ToolCalls: []llm.ToolCallDelta{
|
|
{Index: 0, ID: "tc_a", Name: "search"},
|
|
},
|
|
}},
|
|
{Delta: llm.MessageDelta{
|
|
ToolCalls: []llm.ToolCallDelta{
|
|
{Index: 0, Arguments: `{"q":"go"}`},
|
|
},
|
|
}},
|
|
{Delta: llm.MessageDelta{
|
|
ToolCalls: []llm.ToolCallDelta{
|
|
{Index: 1, ID: "tc_b", Name: "fetch"},
|
|
},
|
|
}},
|
|
{Delta: llm.MessageDelta{
|
|
ToolCalls: []llm.ToolCallDelta{
|
|
{Index: 1, Arguments: `{"url":"https://example.com"}`},
|
|
},
|
|
}},
|
|
{
|
|
FinishReason: new(llm.FinishReasonToolCalls),
|
|
Usage: &llm.Usage{InputTokens: 30, OutputTokens: 10},
|
|
},
|
|
}
|
|
|
|
acc := llm.NewStreamAccumulator(&mockStream{events: events})
|
|
for acc.Next() {
|
|
}
|
|
|
|
require.NoError(t, acc.Err())
|
|
|
|
resp := acc.Response()
|
|
require.Len(t, resp.Message.ToolCalls, 2)
|
|
|
|
assert.Equal(t, "tc_a", resp.Message.ToolCalls[0].ID)
|
|
assert.Equal(t, "search", resp.Message.ToolCalls[0].Function.Name)
|
|
assert.Equal(t, `{"q":"go"}`, resp.Message.ToolCalls[0].Function.Arguments)
|
|
|
|
assert.Equal(t, "tc_b", resp.Message.ToolCalls[1].ID)
|
|
assert.Equal(t, "fetch", resp.Message.ToolCalls[1].Function.Name)
|
|
assert.Equal(t, `{"url":"https://example.com"}`, resp.Message.ToolCalls[1].Function.Arguments)
|
|
})
|
|
|
|
t.Run("text only without tool calls", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
events := []llm.ChatCompletionStreamEvent{
|
|
{Delta: llm.MessageDelta{Content: "Just text."}},
|
|
{
|
|
FinishReason: new(llm.FinishReasonStop),
|
|
Usage: &llm.Usage{InputTokens: 5, OutputTokens: 3},
|
|
},
|
|
}
|
|
|
|
acc := llm.NewStreamAccumulator(&mockStream{events: events})
|
|
for acc.Next() {
|
|
}
|
|
|
|
require.NoError(t, acc.Err())
|
|
|
|
resp := acc.Response()
|
|
assert.Equal(t, "Just text.", resp.Message.Text())
|
|
assert.Equal(t, llm.FinishReasonStop, resp.FinishReason)
|
|
assert.Empty(t, resp.Message.ToolCalls)
|
|
})
|
|
|
|
t.Run("proxies events transparently", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
events := []llm.ChatCompletionStreamEvent{
|
|
{Delta: llm.MessageDelta{Content: "a"}},
|
|
{Delta: llm.MessageDelta{Content: "b"}},
|
|
{FinishReason: new(llm.FinishReasonStop)},
|
|
}
|
|
|
|
acc := llm.NewStreamAccumulator(&mockStream{events: events})
|
|
|
|
var seen []string
|
|
|
|
for acc.Next() {
|
|
e := acc.Event()
|
|
if e.Delta.Content != "" {
|
|
seen = append(seen, e.Delta.Content)
|
|
}
|
|
}
|
|
|
|
assert.Equal(t, []string{"a", "b"}, seen)
|
|
})
|
|
}
|