internal/gateway/generation_stream_test.go
Rama main · réplica de lectura del repositorio original.
package gateway
import (
"context"
"io"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/enterpilot/gomodel/internal/core"
"github.com/enterpilot/gomodel/internal/streaming"
)
// tracingProvider parks a completion in the request's generation trace the
// way telemetry does for a successful provider call, recording the outcome.
type tracingProvider struct {
providerTypeResolverStub
stream string
outcomes []core.GenerationOutcome
}
func (p *tracingProvider) park(ctx context.Context) {
core.GenerationTraceFromContext(ctx).Park(func(outcome core.GenerationOutcome) {
p.outcomes = append(p.outcomes, outcome)
})
}
func (p *tracingProvider) ChatCompletion(ctx context.Context, _ *core.ChatRequest) (*core.ChatResponse, error) {
p.park(ctx)
return p.chatResponse, nil
}
func (p *tracingProvider) StreamChatCompletion(ctx context.Context, _ *core.ChatRequest) (io.ReadCloser, error) {
p.park(ctx)
return io.NopCloser(strings.NewReader(p.stream)), nil
}
func (p *tracingProvider) StreamResponses(ctx context.Context, _ *core.ResponsesRequest) (io.ReadCloser, error) {
p.park(ctx)
return io.NopCloser(strings.NewReader(p.stream)), nil
}
func tracingWorkflow() *core.Workflow {
return &core.Workflow{ProviderType: "openai"}
}
func TestExecuteChatCompletionFinishesGenerationTrace(t *testing.T) {
resp := &core.ChatResponse{ID: "chatcmpl-1", Model: "gpt-5", Choices: []core.Choice{{FinishReason: "stop"}}}
provider := &tracingProvider{chatResponse: resp}
orchestrator := NewInferenceOrchestrator(InferenceConfig{Provider: provider})
req := &core.ChatRequest{Model: "gpt-5"}
ctx := core.WithGenerationTracing(t.Context(), false)
_, err := orchestrator.ExecuteChatCompletion(ctx, tracingWorkflow(), req, "req-1", "/v1/chat/completions")
require.NoError(t, err)
require.Len(t, provider.outcomes, 1)
assert.Same(t, req, provider.outcomes[0].Request)
assert.Same(t, resp, provider.outcomes[0].Response)
}
func TestExecuteChatCompletionFinishesGenerationTraceOnRejectedResponse(t *testing.T) {
provider := &tracingProvider{chatResponse: &core.ChatResponse{ID: "chatcmpl-empty"}}
orchestrator := NewInferenceOrchestrator(InferenceConfig{Provider: provider})
ctx := core.WithGenerationTracing(t.Context(), false)
_, err := orchestrator.ExecuteChatCompletion(ctx, tracingWorkflow(), &core.ChatRequest{Model: "gpt-5"}, "req-1", "/v1/chat/completions")
require.Error(t, err)
require.Len(t, provider.outcomes, 1, "the parked call still ends")
assert.Nil(t, provider.outcomes[0].Response)
}
func TestStreamChatCompletionFinishesGenerationTraceWhenStreamCloses(t *testing.T) {
provider := &tracingProvider{stream: strings.Join([]string{
`data: {"id":"chatcmpl-1","model":"gpt-5","choices":[{"index":0,"delta":{"role":"assistant","content":"Hel"}}]}`,
`data: {"id":"chatcmpl-1","model":"gpt-5","choices":[{"index":0,"delta":{"content":"lo"}}]}`,
`data: {"id":"chatcmpl-1","model":"gpt-5","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"lookup","arguments":"{\"q\""}}]}}]}`,
`data: {"id":"chatcmpl-1","model":"gpt-5","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":":1}"}}]},"finish_reason":"tool_calls"}]}`,
`data: {"id":"chatcmpl-1","model":"gpt-5","choices":[],"usage":{"prompt_tokens":9,"completion_tokens":4,"total_tokens":13}}`,
`data: [DONE]`,
``,
}, "\n\n")}
orchestrator := NewInferenceOrchestrator(InferenceConfig{Provider: provider})
req := &core.ChatRequest{Model: "gpt-5", Stream: true}
ctx := core.WithGenerationTracing(t.Context(), true)
result, err := orchestrator.StreamChatCompletion(ctx, tracingWorkflow(), req)
require.NoError(t, err)
observer := result.GenerationObserver()
require.NotNil(t, observer)
stream := streaming.NewObservedSSEStream(result.Stream, observer)
_, err = io.Copy(io.Discard, stream)
require.NoError(t, err)
require.Empty(t, provider.outcomes, "the trace waits for the stream to close")
require.NoError(t, stream.Close())
require.Len(t, provider.outcomes, 1)
assert.Same(t, req, provider.outcomes[0].Request)
resp, ok := provider.outcomes[0].Response.(*core.ChatResponse)
require.True(t, ok)
assert.Equal(t, "chatcmpl-1", resp.ID)
assert.Equal(t, "gpt-5", resp.Model)
assert.Equal(t, core.Usage{PromptTokens: 9, CompletionTokens: 4, TotalTokens: 13}, resp.Usage)
require.Len(t, resp.Choices, 1)
assert.Equal(t, "tool_calls", resp.Choices[0].FinishReason)
assert.Equal(t, "assistant", resp.Choices[0].Message.Role)
assert.Equal(t, "Hello", resp.Choices[0].Message.Content)
require.Len(t, resp.Choices[0].Message.ToolCalls, 1)
assert.Equal(t, core.ToolCall{ID: "call_1", Type: "function", Function: core.FunctionCall{Name: "lookup", Arguments: `{"q":1}`}}, resp.Choices[0].Message.ToolCalls[0])
}
func TestStreamResponsesFinishesGenerationTraceWithCompletedResponse(t *testing.T) {
provider := &tracingProvider{stream: strings.Join([]string{
`event: response.output_text.delta`,
`data: {"type":"response.output_text.delta","delta":"Hi"}`,
``,
`event: response.completed`,
`data: {"type":"response.completed","response":{"id":"resp_1","model":"gpt-5","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Hi"}]}],"usage":{"input_tokens":5,"output_tokens":1,"total_tokens":6}}}`,
``,
``,
}, "\n")}
orchestrator := NewInferenceOrchestrator(InferenceConfig{Provider: provider})
req := &core.ResponsesRequest{Model: "gpt-5", Stream: true}
ctx := core.WithGenerationTracing(t.Context(), false)
result, err := orchestrator.StreamResponses(ctx, tracingWorkflow(), req)
require.NoError(t, err)
stream := streaming.NewObservedSSEStream(result.Stream, result.GenerationObserver())
_, err = io.Copy(io.Discard, stream)
require.NoError(t, err)
require.NoError(t, stream.Close())
require.Len(t, provider.outcomes, 1)
resp, ok := provider.outcomes[0].Response.(*core.ResponsesResponse)
require.True(t, ok)
assert.Equal(t, "resp_1", resp.ID)
require.NotNil(t, resp.Usage)
assert.Equal(t, 5, resp.Usage.InputTokens)
assert.Equal(t, 1, resp.Usage.OutputTokens)
}
func TestGenerationStreamObserverBoundsCapturedContent(t *testing.T) {
_, trace := core.StartGenerationTrace(core.WithGenerationTracing(t.Context(), true))
observer := newChatGenerationObserver(trace, &core.ChatRequest{})
chunk := strings.Repeat("x", 1024)
for range 2 * core.GenerationContentLimit / len(chunk) {
observer.OnJSONEvent(map[string]any{"choices": []any{map[string]any{
"index": float64(0),
"delta": map[string]any{
"content": chunk,
"tool_calls": []any{map[string]any{"index": float64(0), "function": map[string]any{"arguments": chunk}}},
},
}}})
}
message := observer.chatResponse().Choices[0].Message
assert.Len(t, message.Content, core.GenerationContentLimit+1, "text stops growing one byte past the export limit")
assert.Len(t, message.ToolCalls[0].Function.Arguments, core.GenerationContentLimit+1)
}
func TestStreamResultWithoutTracingHasNoGenerationObserver(t *testing.T) {
provider := &tracingProvider{stream: "data: [DONE]\n\n"}
orchestrator := NewInferenceOrchestrator(InferenceConfig{Provider: provider})
result, err := orchestrator.StreamChatCompletion(t.Context(), tracingWorkflow(), &core.ChatRequest{Model: "gpt-5", Stream: true})
require.NoError(t, err)
assert.Nil(t, result.GenerationObserver())
}
func TestGenerationStreamObserverDecodesOnlyWhatItNeeds(t *testing.T) {
_, trace := core.StartGenerationTrace(core.WithGenerationTracing(t.Context(), false))
_, capturing := core.StartGenerationTrace(core.WithGenerationTracing(t.Context(), true))
contentChunk := []byte(`{"choices":[{"index":0,"delta":{"content":"Hi"},"finish_reason":null}],"usage":null}`)
finishChunk := []byte(`{"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`)
spacedFinishChunk := []byte(`{"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}`)
usageChunk := []byte(`{"choices":[],"usage":{"prompt_tokens":1}}`)
chat := newChatGenerationObserver(trace, &core.ChatRequest{})
assert.False(t, chat.WantsJSONEvent(contentChunk))
assert.True(t, chat.WantsJSONEvent(finishChunk))
assert.True(t, chat.WantsJSONEvent(spacedFinishChunk))
assert.True(t, chat.WantsJSONEvent(usageChunk))
assert.True(t, newChatGenerationObserver(capturing, &core.ChatRequest{}).WantsJSONEvent(contentChunk))
responses := newResponsesGenerationObserver(capturing, &core.ResponsesRequest{})
assert.False(t, responses.WantsJSONEvent([]byte(`{"type":"response.output_text.delta","delta":"Hi"}`)))
assert.True(t, responses.WantsJSONEvent([]byte(`{"type":"response.incomplete","response":{}}`)))
}