From 21827c57eba266ea6d87bf18f97f6d5447df6a46 Mon Sep 17 00:00:00 2001 From: Levi Neely Date: Tue, 21 Apr 2026 15:32:03 +0200 Subject: [PATCH] backend: extract stream parsers, shared HTTP, url.URL, comprehensive tests - Extract streamOpenAISSE, streamOllamaNDJSON from ChatStream goroutines into standalone functions taking io.Reader for direct testability - Remove ctx parameter from streamAnthropicSSE (caller handles cancellation) - Extract streamRequest into backend.go: shared HTTP execution, 429/error handling, and goroutine lifecycle for all three backends - Extract encodeOpenAIMessages, encodeOpenAITools from ChatStream - Replace string baseURL with *url.URL; validate in constructors - Constructors (NewOpenAI, NewOllama, NewAnthropic, NewCopilot) now return error - Remove unreachable json.Marshal and http.NewRequestWithContext error branches ChatStream coverage: 100% (OpenAI, Ollama, Anthropic) Stream parsers: 100% (OpenAI, Ollama), 96.6% (Anthropic) streamRequest: 100% Test files: - openai_encode_test.go: message/tool encoding, JSON round-trip - openai_stream_test.go: SSE parsing via io.Reader - ollama_stream_test.go: NDJSON parsing via io.Reader - anthropic_stream_test.go: SSE parsing, mapAnthropicStopReason - openai_test.go: HTTP integration, error paths, models, context length - backend_test.go: contract checks (all interface methods), stream text/usage/ tool calls/tool conversation, HTTP errors, rate limits for all backends --- pkg/backend/anthropic.go | 51 +- pkg/backend/anthropic_stream_test.go | 186 ++++++++ pkg/backend/backend.go | 29 ++ pkg/backend/backend_test.go | 682 ++++++++++++++++++++++++--- pkg/backend/copilot.go | 15 +- pkg/backend/new.go | 8 +- pkg/backend/ollama.go | 152 +++--- pkg/backend/ollama_stream_test.go | 142 ++++++ pkg/backend/openai.go | 250 +++++----- pkg/backend/openai_encode_test.go | 210 +++++++++ pkg/backend/openai_stream_test.go | 287 +++++++++++ pkg/backend/openai_test.go | 494 +++++++++++++++++++ pkg/backend/testhelper_test.go | 33 ++ 13 files changed, 2216 insertions(+), 323 deletions(-) create mode 100644 pkg/backend/anthropic_stream_test.go create mode 100644 pkg/backend/ollama_stream_test.go create mode 100644 pkg/backend/openai_encode_test.go create mode 100644 pkg/backend/openai_stream_test.go create mode 100644 pkg/backend/openai_test.go create mode 100644 pkg/backend/testhelper_test.go diff --git a/pkg/backend/anthropic.go b/pkg/backend/anthropic.go index a1bb68a..9d1da4a 100644 --- a/pkg/backend/anthropic.go +++ b/pkg/backend/anthropic.go @@ -8,6 +8,7 @@ import ( "fmt" "io" "net/http" + "net/url" "strings" ) @@ -18,13 +19,14 @@ type AnthropicBackend struct { apiKey string model string client *http.Client - BaseURL string // default: https://api.anthropic.com; override for testing + BaseURL *url.URL } -func NewAnthropic(apiKey string) *AnthropicBackend { - b := &AnthropicBackend{apiKey: apiKey, client: &http.Client{}, BaseURL: "https://api.anthropic.com"} +func NewAnthropic(apiKey string) (*AnthropicBackend, error) { + u, _ := url.Parse("https://api.anthropic.com") + b := &AnthropicBackend{apiKey: apiKey, client: &http.Client{}, BaseURL: u} b.model = b.DefaultModel() - return b + return b, nil } func (b *AnthropicBackend) Name() string { return "anthropic" } @@ -106,41 +108,14 @@ func (b *AnthropicBackend) ChatStream(ctx context.Context, messages []Message, t }) } - data, err := json.Marshal(areq) - if err != nil { - return nil, err - } + data, _ := json.Marshal(areq) - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.BaseURL+"/v1/messages", bytes.NewReader(data)) - if err != nil { - return nil, err - } + httpReq, _ := http.NewRequestWithContext(ctx, http.MethodPost, b.BaseURL.JoinPath("/v1/messages").String(), bytes.NewReader(data)) httpReq.Header.Set("Content-Type", "application/json") httpReq.Header.Set("X-Api-Key", b.apiKey) httpReq.Header.Set("Anthropic-Version", "2023-06-01") - resp, err := b.client.Do(httpReq) - if err != nil { - return nil, err - } - if resp.StatusCode == http.StatusTooManyRequests { - body, _ := io.ReadAll(resp.Body) - resp.Body.Close() - return nil, &RateLimitError{RetryAfter: parseRetryAfter(resp.Header.Get("Retry-After")), Message: string(body)} - } - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - resp.Body.Close() - return nil, fmt.Errorf("anthropic HTTP %d: %s", resp.StatusCode, body) - } - - ch := make(chan StreamEvent, 8) - go func() { - defer close(ch) - defer resp.Body.Close() - streamAnthropicSSE(ctx, resp.Body, ch) - }() - return ch, nil + return streamRequest(b.client, httpReq, "anthropic", streamAnthropicSSE) } // buildAnthropicMessages converts ollie messages to Anthropic wire format. @@ -202,7 +177,7 @@ func buildAnthropicMessages(messages []Message) (system string, out []anthropicM } // streamAnthropicSSE reads the Anthropic SSE stream and sends StreamEvents to ch. -func streamAnthropicSSE(ctx context.Context, body io.Reader, ch chan<- StreamEvent) { +func streamAnthropicSSE(body io.Reader, ch chan<- StreamEvent) { scanner := bufio.NewScanner(body) scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) @@ -261,11 +236,7 @@ func streamAnthropicSSE(ctx context.Context, body io.Reader, ch chan<- StreamEve switch v.Delta.Type { case "text_delta": if v.Delta.Text != "" { - select { - case ch <- StreamEvent{Content: v.Delta.Text}: - case <-ctx.Done(): - done = true - } + ch <- StreamEvent{Content: v.Delta.Text} } case "input_json_delta": if t := tools[v.Index]; t != nil { diff --git a/pkg/backend/anthropic_stream_test.go b/pkg/backend/anthropic_stream_test.go new file mode 100644 index 0000000..c9ed398 --- /dev/null +++ b/pkg/backend/anthropic_stream_test.go @@ -0,0 +1,186 @@ +package backend + +import ( + "fmt" + "strings" + "testing" +) + +func collectAnthropic(sse string) []StreamEvent { + ch := make(chan StreamEvent, 32) + go func() { + defer close(ch) + streamAnthropicSSE(strings.NewReader(sse), ch) + }() + var evs []StreamEvent + for ev := range ch { + evs = append(evs, ev) + } + return evs +} + +func TestAnthropicSSE_TextStream(t *testing.T) { + evs := collectAnthropic( + "event: message_start\ndata: {\"message\":{\"usage\":{\"input_tokens\":10}}}\n\n" + + "event: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n" + + "event: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\" world\"}}\n\n" + + "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":5}}\n\n" + + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ) + var text string + for _, ev := range evs { + text += ev.Content + } + if text != "hello world" { + t.Errorf("text = %q", text) + } + fin := evs[len(evs)-1] + if !fin.Done || fin.StopReason != "stop" { + t.Errorf("final = %+v", fin) + } + if fin.Usage.InputTokens != 10 || fin.Usage.OutputTokens != 5 { + t.Errorf("usage = %+v", fin.Usage) + } +} + +func TestAnthropicSSE_ToolCall(t *testing.T) { + evs := collectAnthropic( + "event: message_start\ndata: {\"message\":{\"usage\":{\"input_tokens\":5}}}\n\n" + + "event: content_block_start\ndata: {\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"c1\",\"name\":\"f\"}}\n\n" + + "event: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"a\\\":\"}}\n\n" + + "event: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"1}\"}}\n\n" + + "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":3}}\n\n" + + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ) + fin := evs[len(evs)-1] + if len(fin.ToolCalls) != 1 { + t.Fatalf("tool_calls len = %d", len(fin.ToolCalls)) + } + tc := fin.ToolCalls[0] + if tc.ID != "c1" || tc.Name != "f" || string(tc.Arguments) != `{"a":1}` { + t.Errorf("tc = %+v", tc) + } + if fin.StopReason != "tool_calls" { + t.Errorf("stop = %q", fin.StopReason) + } +} + +func TestAnthropicSSE_ParallelToolCalls(t *testing.T) { + evs := collectAnthropic( + "event: message_start\ndata: {\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n" + + "event: content_block_start\ndata: {\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"a\",\"name\":\"x\"}}\n\n" + + "event: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{}\"}}\n\n" + + "event: content_block_start\ndata: {\"index\":1,\"content_block\":{\"type\":\"tool_use\",\"id\":\"b\",\"name\":\"y\"}}\n\n" + + "event: content_block_delta\ndata: {\"index\":1,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{}\"}}\n\n" + + "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":1}}\n\n" + + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ) + fin := evs[len(evs)-1] + if len(fin.ToolCalls) != 2 { + t.Fatalf("len = %d", len(fin.ToolCalls)) + } + ids := map[string]bool{} + for _, tc := range fin.ToolCalls { + ids[tc.ID] = true + } + if !ids["a"] || !ids["b"] { + t.Errorf("ids = %v", ids) + } +} + +func TestAnthropicSSE_ErrorEvent(t *testing.T) { + evs := collectAnthropic( + "event: error\ndata: {\"error\":{\"message\":\"overloaded\"}}\n\n", + ) + if len(evs) != 1 || !evs[0].Done { + t.Fatal("expected Done") + } + if !strings.Contains(evs[0].StopReason, "overloaded") { + t.Errorf("stop = %q", evs[0].StopReason) + } +} + +func TestAnthropicSSE_TextBlockStartIgnored(t *testing.T) { + // content_block_start with type "text" should not create a tool accumulator. + evs := collectAnthropic( + "event: message_start\ndata: {\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n" + + "event: content_block_start\ndata: {\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n" + + "event: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n" + + "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1}}\n\n" + + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ) + fin := evs[len(evs)-1] + if len(fin.ToolCalls) != 0 { + t.Errorf("tool_calls = %+v; want empty", fin.ToolCalls) + } +} + +func TestAnthropicSSE_InputJSONDeltaUnknownIndex(t *testing.T) { + // input_json_delta for an index with no content_block_start should be silently ignored. + evs := collectAnthropic( + "event: message_start\ndata: {\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n" + + "event: content_block_delta\ndata: {\"index\":99,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"junk\"}}\n\n" + + "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1}}\n\n" + + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ) + fin := evs[len(evs)-1] + if len(fin.ToolCalls) != 0 { + t.Errorf("tool_calls = %+v; want empty", fin.ToolCalls) + } +} + +func TestAnthropicSSE_TrailingEventNoBlankLine(t *testing.T) { + // Event not terminated by a blank line should still be processed. + evs := collectAnthropic( + "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1}}\n\n" + + "event: message_stop\ndata: {\"type\":\"message_stop\"}", + ) + if len(evs) == 0 || !evs[len(evs)-1].Done { + t.Error("trailing event not processed") + } +} + +func TestAnthropicSSE_EmptyStream(t *testing.T) { + evs := collectAnthropic("") + if len(evs) != 0 { + t.Errorf("expected no events; got %d", len(evs)) + } +} + +func TestAnthropicSSE_ScannerError(t *testing.T) { + r := newErrAfterReader( + "event: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n", + fmt.Errorf("connection reset"), + ) + ch := make(chan StreamEvent, 32) + go func() { + defer close(ch) + streamAnthropicSSE(r, ch) + }() + var evs []StreamEvent + for ev := range ch { + evs = append(evs, ev) + } + fin := evs[len(evs)-1] + if !fin.Done || !strings.Contains(fin.StopReason, "stream read") { + t.Errorf("final = %+v", fin) + } +} + +// --- mapAnthropicStopReason --- + +func TestMapAnthropicStopReason(t *testing.T) { + tests := []struct{ in, want string }{ + {"end_turn", "stop"}, + {"stop_sequence", "stop"}, + {"max_tokens", "length"}, + {"tool_use", "tool_calls"}, + {"unknown", "unknown"}, + {"", ""}, + } + for _, tt := range tests { + if got := mapAnthropicStopReason(tt.in); got != tt.want { + t.Errorf("mapAnthropicStopReason(%q) = %q; want %q", tt.in, got, tt.want) + } + } +} diff --git a/pkg/backend/backend.go b/pkg/backend/backend.go index 1694971..367bb21 100644 --- a/pkg/backend/backend.go +++ b/pkg/backend/backend.go @@ -7,6 +7,8 @@ import ( "context" "encoding/json" "fmt" + "io" + "net/http" "time" ) @@ -92,3 +94,30 @@ type Backend interface { // Models returns the list of available model IDs from the provider. Models(ctx context.Context) []string } + +// streamRequest executes req via client, handles 429/non-200 errors, and +// spawns a goroutine that feeds resp.Body through parseFn into the returned +// channel. label is used in error messages (e.g. "openai", "ollama"). +func streamRequest(client *http.Client, req *http.Request, label string, parseFn func(io.Reader, chan<- StreamEvent)) (<-chan StreamEvent, error) { + resp, err := client.Do(req) + if err != nil { + return nil, err + } + if resp.StatusCode == http.StatusTooManyRequests { + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + return nil, &RateLimitError{RetryAfter: parseRetryAfter(resp.Header.Get("Retry-After")), Message: string(body)} + } + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + return nil, fmt.Errorf("%s HTTP %d: %s", label, resp.StatusCode, body) + } + ch := make(chan StreamEvent, 8) + go func() { + defer close(ch) + defer resp.Body.Close() + parseFn(resp.Body, ch) + }() + return ch, nil +} diff --git a/pkg/backend/backend_test.go b/pkg/backend/backend_test.go index 39f4905..7af252b 100644 --- a/pkg/backend/backend_test.go +++ b/pkg/backend/backend_test.go @@ -2,17 +2,51 @@ package backend_test import ( "context" + "encoding/json" + "fmt" + "io" "net/http" "net/http/httptest" + "net/url" + "strings" "testing" "time" "ollie/pkg/backend" ) -// stubBackend is a minimal Backend used to verify the contract. +func mustNewOpenAI(t *testing.T, name, baseURL, apiKey string) *backend.OpenAIBackend { + t.Helper() + b, err := backend.NewOpenAI(name, baseURL, apiKey) + if err != nil { + t.Fatal(err) + } + return b +} + +func mustNewOllama(t *testing.T, baseURL string) *backend.OllamaBackend { + t.Helper() + b, err := backend.NewOllama(baseURL) + if err != nil { + t.Fatal(err) + } + return b +} + +func mustNewAnthropic(t *testing.T, apiKey string) *backend.AnthropicBackend { + t.Helper() + b, err := backend.NewAnthropic(apiKey) + if err != nil { + t.Fatal(err) + } + return b +} + +// --- stub backend --- + type stubBackend struct { - model string + model string + streamFn func(context.Context, []backend.Message, []backend.Tool, backend.GenerationParams) (<-chan backend.StreamEvent, error) } func (s *stubBackend) Name() string { return "stub" } @@ -21,17 +55,25 @@ func (s *stubBackend) Model() string { return s.model } func (s *stubBackend) SetModel(m string) { s.model = m } func (s *stubBackend) ContextLength(_ context.Context) int { return 4096 } func (s *stubBackend) Models(_ context.Context) []string { return []string{"stub-default"} } -func (s *stubBackend) ChatStream(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { +func (s *stubBackend) ChatStream(ctx context.Context, msgs []backend.Message, tools []backend.Tool, params backend.GenerationParams) (<-chan backend.StreamEvent, error) { + if s.streamFn != nil { + return s.streamFn(ctx, msgs, tools, params) + } ch := make(chan backend.StreamEvent, 1) - ch <- backend.StreamEvent{Done: true, StopReason: "end_turn"} + ch <- backend.StreamEvent{Done: true, StopReason: "stop"} close(ch) return ch, nil } -// checkContract verifies Backend invariants. Call it with any Backend implementation. +// --- contract checks --- + +// checkContract verifies Backend invariants that hold for every implementation. func checkContract(t *testing.T, b backend.Backend) { t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + // Name / DefaultModel if b.Name() == "" { t.Error("Name() must be non-empty") } @@ -39,17 +81,24 @@ func checkContract(t *testing.T, b backend.Backend) { t.Error("DefaultModel() must be non-empty") } + // SetModel / Model b.SetModel("test-model") if got := b.Model(); got != "test-model" { - t.Errorf("Model() = %q after SetModel; want %q", got, "test-model") + t.Errorf("Model() = %q after SetModel; want test-model", got) } - msgs := []backend.Message{{Role: "user", Content: "ping"}} - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - ch, err := b.ChatStream(ctx, msgs, nil, backend.GenerationParams{}) + // ContextLength + if cl := b.ContextLength(ctx); cl < 0 { + t.Errorf("ContextLength() = %d; must be >= 0", cl) + } + + // Models + _ = b.Models(ctx) // must not panic + + // ChatStream: bare message → Done + ch, err := b.ChatStream(ctx, []backend.Message{{Role: "user", Content: "ping"}}, nil, backend.GenerationParams{}) if err != nil { - t.Fatalf("ChatStream returned error: %v", err) + t.Fatalf("ChatStream: %v", err) } var gotDone bool for ev := range ch { @@ -58,78 +107,601 @@ func checkContract(t *testing.T, b backend.Backend) { } } if !gotDone { - t.Error("ChatStream channel closed without a Done==true event") + t.Error("ChatStream channel closed without Done==true") } } -func TestStubBackendContract(t *testing.T) { +// checkStreamText verifies that text content arrives via StreamEvent.Content +// and the final event has Done==true with a non-empty StopReason. +func checkStreamText(t *testing.T, b backend.Backend) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + ch, err := b.ChatStream(ctx, []backend.Message{{Role: "user", Content: "hi"}}, nil, backend.GenerationParams{}) + if err != nil { + t.Fatalf("ChatStream: %v", err) + } + var text string + var final backend.StreamEvent + for ev := range ch { + text += ev.Content + if ev.Done { + final = ev + } + } + if text == "" { + t.Error("expected non-empty text content") + } + if !final.Done { + t.Error("no Done event") + } + if final.StopReason == "" { + t.Error("StopReason should be non-empty") + } +} + +// checkStreamUsage verifies that the Done event carries usage tokens. +func checkStreamUsage(t *testing.T, b backend.Backend) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + ch, err := b.ChatStream(ctx, []backend.Message{{Role: "user", Content: "hi"}}, nil, backend.GenerationParams{}) + if err != nil { + t.Fatalf("ChatStream: %v", err) + } + var final backend.StreamEvent + for ev := range ch { + if ev.Done { + final = ev + } + } + if final.Usage.InputTokens <= 0 { + t.Errorf("InputTokens = %d; want > 0", final.Usage.InputTokens) + } + if final.Usage.OutputTokens <= 0 { + t.Errorf("OutputTokens = %d; want > 0", final.Usage.OutputTokens) + } +} + +// checkStreamToolCalls verifies that when tools are provided and the model +// returns tool calls, they appear in the Done event's ToolCalls slice. +func checkStreamToolCalls(t *testing.T, b backend.Backend) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + tools := []backend.Tool{{ + Name: "get_weather", + Description: "Get weather", + Parameters: json.RawMessage(`{"type":"object","properties":{"city":{"type":"string"}}}`), + }} + ch, err := b.ChatStream(ctx, []backend.Message{{Role: "user", Content: "weather in paris"}}, tools, backend.GenerationParams{}) + if err != nil { + t.Fatalf("ChatStream: %v", err) + } + var final backend.StreamEvent + for ev := range ch { + if ev.Done { + final = ev + } + } + if len(final.ToolCalls) == 0 { + t.Fatal("expected tool calls in Done event") + } + tc := final.ToolCalls[0] + if tc.Name != "get_weather" { + t.Errorf("tool call name = %q; want get_weather", tc.Name) + } + if len(tc.Arguments) == 0 { + t.Error("tool call arguments empty") + } +} + +// checkStreamToolConversation verifies a multi-turn tool use conversation: +// user → assistant(tool_call) → tool(result) → assistant(text). +func checkStreamToolConversation(t *testing.T, b backend.Backend) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + msgs := []backend.Message{ + {Role: "user", Content: "weather?"}, + {Role: "assistant", ToolCalls: []backend.ToolCall{{ + ID: "call_1", Name: "get_weather", Arguments: json.RawMessage(`{"city":"paris"}`), + }}}, + {Role: "tool", Content: "sunny 22C", ToolCallID: "call_1"}, + } + ch, err := b.ChatStream(ctx, msgs, nil, backend.GenerationParams{}) + if err != nil { + t.Fatalf("ChatStream: %v", err) + } + var text string + var gotDone bool + for ev := range ch { + text += ev.Content + if ev.Done { + gotDone = true + } + } + if !gotDone { + t.Error("no Done event") + } + if text == "" { + t.Error("expected text response after tool result") + } +} + +// checkErrorHTTP verifies that non-200 HTTP responses return an error (not a channel). +func checkErrorHTTP(t *testing.T, b backend.Backend) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + _, err := b.ChatStream(ctx, []backend.Message{{Role: "user", Content: "hi"}}, nil, backend.GenerationParams{}) + if err == nil { + t.Error("expected error for non-200 HTTP") + } +} + +// checkErrorRateLimit verifies that 429 returns *RateLimitError. +func checkErrorRateLimit(t *testing.T, b backend.Backend) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + _, err := b.ChatStream(ctx, []backend.Message{{Role: "user", Content: "hi"}}, nil, backend.GenerationParams{}) + if err == nil { + t.Fatal("expected error for 429") + } + if _, ok := err.(*backend.RateLimitError); !ok { + t.Errorf("error type = %T; want *RateLimitError", err) + } +} + +// --- test server helpers --- + +// openAITextServer returns SSE with text content, usage, and stop reason. +func openAITextServer() *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, + "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n"+ + "data: {\"choices\":[{\"finish_reason\":\"stop\",\"delta\":{}}]}\n\n"+ + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5}}\n\n"+ + "data: [DONE]\n\n", + ) + })) +} + +// openAIToolServer returns SSE with a tool call. +func openAIToolServer() *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"paris\"}"}}]}}]}`+"\n\n"+ + `data: {"choices":[{"finish_reason":"tool_calls","delta":{}}]}`+"\n\n"+ + "data: [DONE]\n\n", + ) + })) +} + +// openAIErrorServer returns the given HTTP status code. +func openAIErrorServer(code int) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if code == 429 { + w.Header().Set("Retry-After", "30") + } + w.WriteHeader(code) + fmt.Fprintf(w, "error %d", code) + })) +} + +// ollamaTextServer returns NDJSON with text, usage, and stop. +func ollamaTextServer() *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/x-ndjson") + fmt.Fprint(w, + "{\"message\":{\"role\":\"assistant\",\"content\":\"hello\"},\"done\":false}\n"+ + "{\"done\":true,\"done_reason\":\"stop\",\"message\":{\"role\":\"assistant\",\"content\":\"\"},\"prompt_eval_count\":10,\"eval_count\":5}\n", + ) + })) +} + +// ollamaToolServer returns NDJSON with a tool call. +func ollamaToolServer() *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/x-ndjson") + fmt.Fprint(w, + `{"done":true,"done_reason":"stop","message":{"role":"assistant","content":"","tool_calls":[{"function":{"name":"get_weather","arguments":{"city":"paris"}}}]},"prompt_eval_count":10,"eval_count":5}`+"\n", + ) + })) +} + +// ollamaErrorServer returns the given HTTP status code. +func ollamaErrorServer(code int) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if code == 429 { + w.Header().Set("Retry-After", "30") + } + w.WriteHeader(code) + fmt.Fprintf(w, "error %d", code) + })) +} + +// anthropicTextServer returns Anthropic SSE with text, usage, and stop. +func anthropicTextServer() *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, + "event: message_start\ndata: {\"message\":{\"usage\":{\"input_tokens\":10}}}\n\n"+ + "event: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n"+ + "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":5}}\n\n"+ + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ) + })) +} + +// anthropicToolServer returns Anthropic SSE with a tool call. +func anthropicToolServer() *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, + "event: message_start\ndata: {\"message\":{\"usage\":{\"input_tokens\":10}}}\n\n"+ + "event: content_block_start\ndata: {\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_1\",\"name\":\"get_weather\"}}\n\n"+ + "event: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"city\\\":\\\"paris\\\"}\"}}\n\n"+ + "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":5}}\n\n"+ + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ) + })) +} + +// anthropicErrorServer returns the given HTTP status code. +func anthropicErrorServer(code int) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if code == 429 { + w.Header().Set("Retry-After", "30") + } + w.WriteHeader(code) + fmt.Fprintf(w, "error %d", code) + })) +} + +// --- stub contract --- + +func TestStubContract(t *testing.T) { checkContract(t, &stubBackend{model: "stub-default"}) } -func TestOllamaContract(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/x-ndjson") - w.Write([]byte(`{"done":true,"done_reason":"stop","message":{"role":"assistant","content":""}}` + "\n")) - })) - defer srv.Close() - checkContract(t, backend.NewOllama(srv.URL)) +func TestStubStreamText(t *testing.T) { + s := &stubBackend{model: "m"} + s.streamFn = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { + ch := make(chan backend.StreamEvent, 2) + ch <- backend.StreamEvent{Content: "hi"} + ch <- backend.StreamEvent{Done: true, StopReason: "stop", Usage: backend.Usage{InputTokens: 1, OutputTokens: 1}} + close(ch) + return ch, nil + } + checkStreamText(t, s) + checkStreamUsage(t, s) } +func TestStubStreamToolCalls(t *testing.T) { + s := &stubBackend{model: "m"} + s.streamFn = func(_ context.Context, _ []backend.Message, tools []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { + ch := make(chan backend.StreamEvent, 1) + ch <- backend.StreamEvent{ + Done: true, StopReason: "tool_calls", + ToolCalls: []backend.ToolCall{{Name: tools[0].Name, Arguments: json.RawMessage(`{"city":"paris"}`)}}, + } + close(ch) + return ch, nil + } + checkStreamToolCalls(t, s) +} + +// --- OpenAI --- + func TestOpenAIContract(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/event-stream") - w.Write([]byte("data: [DONE]\n\n")) - })) + srv := openAITextServer() defer srv.Close() - checkContract(t, backend.NewOpenAI("openai", srv.URL, "fake-key")) + checkContract(t, mustNewOpenAI(t, "openai", srv.URL, "k")) } +func TestOpenAIStreamText(t *testing.T) { + srv := openAITextServer() + defer srv.Close() + b := mustNewOpenAI(t, "openai", srv.URL, "k") + checkStreamText(t, b) + checkStreamUsage(t, b) +} + +func TestOpenAIStreamToolCalls(t *testing.T) { + srv := openAIToolServer() + defer srv.Close() + checkStreamToolCalls(t, mustNewOpenAI(t, "openai", srv.URL, "k")) +} + +func TestOpenAIStreamToolConversation(t *testing.T) { + srv := openAITextServer() + defer srv.Close() + checkStreamToolConversation(t, mustNewOpenAI(t, "openai", srv.URL, "k")) +} + +func TestOpenAIErrorHTTP(t *testing.T) { + for _, code := range []int{400, 401, 403, 500, 503} { + t.Run(fmt.Sprintf("%d", code), func(t *testing.T) { + srv := openAIErrorServer(code) + defer srv.Close() + checkErrorHTTP(t, mustNewOpenAI(t, "openai", srv.URL, "k")) + }) + } +} + +func TestOpenAIErrorRateLimit(t *testing.T) { + srv := openAIErrorServer(429) + defer srv.Close() + checkErrorRateLimit(t, mustNewOpenAI(t, "openai", srv.URL, "k")) +} + +// --- Copilot (OpenAI variant) --- + func TestCopilotContract(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/event-stream") - w.Write([]byte("data: [DONE]\n\n")) - })) + srv := openAITextServer() defer srv.Close() - // Copilot is an OpenAI backend; exercise it via NewOpenAI with the copilot name. - checkContract(t, backend.NewOpenAI("copilot", srv.URL, "fake-token")) + checkContract(t, mustNewOpenAI(t, "copilot", srv.URL, "k")) } +func TestCopilotStreamText(t *testing.T) { + srv := openAITextServer() + defer srv.Close() + b := mustNewOpenAI(t, "copilot", srv.URL, "k") + checkStreamText(t, b) + checkStreamUsage(t, b) +} + +// --- Ollama --- + +func TestOllamaContract(t *testing.T) { + srv := ollamaTextServer() + defer srv.Close() + checkContract(t, mustNewOllama(t, srv.URL)) +} + +func TestOllamaStreamText(t *testing.T) { + srv := ollamaTextServer() + defer srv.Close() + b := mustNewOllama(t, srv.URL) + checkStreamText(t, b) + checkStreamUsage(t, b) +} + +func TestOllamaStreamToolCalls(t *testing.T) { + srv := ollamaToolServer() + defer srv.Close() + checkStreamToolCalls(t, mustNewOllama(t, srv.URL)) +} + +func TestOllamaStreamToolConversation(t *testing.T) { + srv := ollamaTextServer() + defer srv.Close() + checkStreamToolConversation(t, mustNewOllama(t, srv.URL)) +} + +func TestOllamaErrorHTTP(t *testing.T) { + for _, code := range []int{400, 500, 503} { + t.Run(fmt.Sprintf("%d", code), func(t *testing.T) { + srv := ollamaErrorServer(code) + defer srv.Close() + checkErrorHTTP(t, mustNewOllama(t, srv.URL)) + }) + } +} + +func TestOllamaErrorRateLimit(t *testing.T) { + srv := ollamaErrorServer(429) + defer srv.Close() + checkErrorRateLimit(t, mustNewOllama(t, srv.URL)) +} + +// --- Anthropic --- + func TestAnthropicContract(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/event-stream") - w.Write([]byte( - "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":0}}\n\n" + - "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", - )) - })) + srv := anthropicTextServer() defer srv.Close() - b := backend.NewAnthropic("fake-key") - b.BaseURL = srv.URL + b := mustNewAnthropic(t, "k") + b.BaseURL, _ = url.Parse(srv.URL) checkContract(t, b) } +func TestAnthropicStreamText(t *testing.T) { + srv := anthropicTextServer() + defer srv.Close() + b := mustNewAnthropic(t, "k") + b.BaseURL, _ = url.Parse(srv.URL) + checkStreamText(t, b) + checkStreamUsage(t, b) +} + +func TestAnthropicStreamToolCalls(t *testing.T) { + srv := anthropicToolServer() + defer srv.Close() + b := mustNewAnthropic(t, "k") + b.BaseURL, _ = url.Parse(srv.URL) + checkStreamToolCalls(t, b) +} + +func TestAnthropicStreamToolConversation(t *testing.T) { + srv := anthropicTextServer() + defer srv.Close() + b := mustNewAnthropic(t, "k") + b.BaseURL, _ = url.Parse(srv.URL) + checkStreamToolConversation(t, b) +} + +func TestAnthropicErrorHTTP(t *testing.T) { + for _, code := range []int{400, 401, 500, 503} { + t.Run(fmt.Sprintf("%d", code), func(t *testing.T) { + srv := anthropicErrorServer(code) + defer srv.Close() + b := mustNewAnthropic(t, "k") + b.BaseURL, _ = url.Parse(srv.URL) + checkErrorHTTP(t, b) + }) + } +} + +func TestAnthropicErrorRateLimit(t *testing.T) { + srv := anthropicErrorServer(429) + defer srv.Close() + b := mustNewAnthropic(t, "k") + b.BaseURL, _ = url.Parse(srv.URL) + checkErrorRateLimit(t, b) +} + +// --- CodeWhisperer --- + func TestCodeWhispererContract(t *testing.T) { - // Static auth has no DefaultEndpoint, so ChatStream terminates immediately - // with Done==true via the endpoint-resolution error path. This verifies that - // the contract invariant — stream always terminates with Done — holds even - // on error paths. b, err := backend.NewCodeWhisperer("fake-token") if err != nil { - t.Fatalf("NewCodeWhisperer: %v", err) + t.Fatal(err) } checkContract(t, b) } -func TestRateLimitError(t *testing.T) { - e := &backend.RateLimitError{RetryAfter: 30 * time.Second, Message: "quota exceeded"} - want := "rate limited (retry after 30s): quota exceeded" - if got := e.Error(); got != want { - t.Errorf("Error() = %q; want %q", got, want) - } +// --- RateLimitError --- - e2 := &backend.RateLimitError{Message: "over limit"} - want2 := "rate limited: over limit" - if got := e2.Error(); got != want2 { - t.Errorf("Error() = %q; want %q", got, want2) +func TestRateLimitError(t *testing.T) { + tests := []struct { + err *backend.RateLimitError + want string + }{ + {&backend.RateLimitError{RetryAfter: 30 * time.Second, Message: "quota exceeded"}, "rate limited (retry after 30s): quota exceeded"}, + {&backend.RateLimitError{Message: "over limit"}, "rate limited: over limit"}, + } + for _, tt := range tests { + if got := tt.err.Error(); got != tt.want { + t.Errorf("Error() = %q; want %q", got, tt.want) + } + } +} + +// --- Anthropic: wire format edge cases --- + +func TestAnthropicNilToolParameters(t *testing.T) { + var gotBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, + "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1}}\n\n"+ + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ) + })) + defer srv.Close() + + b := mustNewAnthropic(t, "k") + b.BaseURL, _ = url.Parse(srv.URL) + ch, err := b.ChatStream(context.Background(), + []backend.Message{{Role: "user", Content: "hi"}}, + []backend.Tool{{Name: "f", Description: "d", Parameters: nil}}, + backend.GenerationParams{}, + ) + if err != nil { + t.Fatal(err) + } + for range ch { + } + // nil Parameters should become {"type":"object","properties":{}} + if !strings.Contains(string(gotBody), `"input_schema":{"type":"object","properties":{}}`) { + t.Errorf("body = %s", gotBody) + } +} + +func TestAnthropicMaxTokensDefault(t *testing.T) { + var gotBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, + "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1}}\n\n"+ + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ) + })) + defer srv.Close() + + b := mustNewAnthropic(t, "k") + b.BaseURL, _ = url.Parse(srv.URL) + ch, err := b.ChatStream(context.Background(), + []backend.Message{{Role: "user", Content: "hi"}}, + nil, + backend.GenerationParams{}, // MaxTokens == 0 + ) + if err != nil { + t.Fatal(err) + } + for range ch { + } + if !strings.Contains(string(gotBody), `"max_tokens":8192`) { + t.Errorf("body = %s; want max_tokens:8192", gotBody) + } +} + +func TestAnthropicMaxTokensExplicit(t *testing.T) { + var gotBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, + "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1}}\n\n"+ + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ) + })) + defer srv.Close() + + b := mustNewAnthropic(t, "k") + b.BaseURL, _ = url.Parse(srv.URL) + ch, err := b.ChatStream(context.Background(), + []backend.Message{{Role: "user", Content: "hi"}}, + nil, + backend.GenerationParams{MaxTokens: 256}, + ) + if err != nil { + t.Fatal(err) + } + for range ch { + } + if !strings.Contains(string(gotBody), `"max_tokens":256`) { + t.Errorf("body = %s; want max_tokens:256", gotBody) + } +} + +// --- error path: SSE stream error event (Anthropic) --- + +func TestAnthropicStreamError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "event: error\ndata: {\"error\":{\"message\":\"overloaded\"}}\n\n") + })) + defer srv.Close() + + b := mustNewAnthropic(t, "k") + b.BaseURL, _ = url.Parse(srv.URL) + ch, err := b.ChatStream(context.Background(), []backend.Message{{Role: "user", Content: "hi"}}, nil, backend.GenerationParams{}) + if err != nil { + t.Fatal(err) + } + var final backend.StreamEvent + for ev := range ch { + if ev.Done { + final = ev + } + } + if !final.Done { + t.Error("expected Done") + } + if !strings.Contains(final.StopReason, "overloaded") { + t.Errorf("StopReason = %q; want to contain 'overloaded'", final.StopReason) } } diff --git a/pkg/backend/copilot.go b/pkg/backend/copilot.go index ae6eee5..4802820 100644 --- a/pkg/backend/copilot.go +++ b/pkg/backend/copilot.go @@ -1,21 +1,22 @@ package backend -import "net/http" +import ( + "net/http" + "net/url" +) // NewCopilot returns an OpenAI-compatible backend configured for the GitHub // Copilot chat completions API. token is a short-lived Copilot bearer token. -// -// The required Copilot-specific headers (Openai-Intent, User-Agent) are -// injected on every request via the OpenAIBackend's extraHeaders mechanism. -func NewCopilot(token string) *OpenAIBackend { +func NewCopilot(token string) (*OpenAIBackend, error) { + u, _ := url.Parse("https://api.githubcopilot.com") return &OpenAIBackend{ name: "copilot", - baseURL: "https://api.githubcopilot.com", + baseURL: u, apiKey: token, client: &http.Client{}, extraHeaders: map[string]string{ "Openai-Intent": "conversation-edits", "User-Agent": "opencode/local", }, - } + }, nil } diff --git a/pkg/backend/new.go b/pkg/backend/new.go index 05bfc82..a2c9112 100644 --- a/pkg/backend/new.go +++ b/pkg/backend/new.go @@ -63,23 +63,23 @@ func New() (Backend, error) { switch which { case "ollama": - return NewOllama(os.Getenv("OLLIE_OLLAMA_URL")), nil + return NewOllama(os.Getenv("OLLIE_OLLAMA_URL")) case "openai", "openrouter": url := os.Getenv("OLLIE_OPENAI_URL") key := os.Getenv("OLLIE_OPENAI_KEY") - return NewOpenAI(openAIName(which, url), url, key), nil + return NewOpenAI(openAIName(which, url), url, key) case "anthropic": key := os.Getenv("OLLIE_ANTHROPIC_KEY") if key == "" { return nil, fmt.Errorf("OLLIE_ANTHROPIC_KEY is required for anthropic backend") } - return NewAnthropic(key), nil + return NewAnthropic(key) case "copilot": token := os.Getenv("OLLIE_COPILOT_TOKEN") if token == "" { return nil, fmt.Errorf("OLLIE_COPILOT_TOKEN is required for copilot backend") } - return NewCopilot(token), nil + return NewCopilot(token) case "kiro", "codewhisperer": return NewCodeWhisperer(os.Getenv("OLLIE_KIRO_TOKEN")) default: diff --git a/pkg/backend/ollama.go b/pkg/backend/ollama.go index d1c02ca..903e86b 100644 --- a/pkg/backend/ollama.go +++ b/pkg/backend/ollama.go @@ -8,25 +8,30 @@ import ( "fmt" "io" "net/http" + "net/url" "strings" ) // OllamaBackend speaks the Ollama /api/chat wire format. type OllamaBackend struct { - baseURL string + baseURL *url.URL model string client *http.Client ctxLength int ctxModel string } -func NewOllama(baseURL string) *OllamaBackend { +func NewOllama(baseURL string) (*OllamaBackend, error) { if baseURL == "" { baseURL = "http://localhost:11434" } - b := &OllamaBackend{baseURL: baseURL, client: &http.Client{}} + u, err := url.Parse(baseURL) + if err != nil { + return nil, fmt.Errorf("invalid base URL: %w", err) + } + b := &OllamaBackend{baseURL: u, client: &http.Client{}} b.model = b.DefaultModel() - return b + return b, nil } func (b *OllamaBackend) Name() string { return "ollama" } @@ -35,7 +40,7 @@ func (b *OllamaBackend) Model() string { return b.model } func (b *OllamaBackend) SetModel(m string) { b.model = m; b.ctxLength = 0 } func (b *OllamaBackend) Models(ctx context.Context) []string { - req, err := http.NewRequestWithContext(ctx, "GET", b.baseURL+"/api/tags", nil) + req, err := http.NewRequestWithContext(ctx, "GET", b.baseURL.JoinPath("/api/tags").String(), nil) if err != nil { return nil } @@ -67,7 +72,7 @@ func (b *OllamaBackend) ContextLength(ctx context.Context) int { return b.ctxLength } body, _ := json.Marshal(map[string]string{"name": b.model}) - req, err := http.NewRequestWithContext(ctx, "POST", b.baseURL+"/api/show", bytes.NewReader(body)) + req, err := http.NewRequestWithContext(ctx, "POST", b.baseURL.JoinPath("/api/show").String(), bytes.NewReader(body)) if err != nil { return 0 } @@ -167,94 +172,65 @@ func (b *OllamaBackend) ChatStream(ctx context.Context, messages []Message, tool }) } - data, err := json.Marshal(ollamaChatRequest{ + data, _ := json.Marshal(ollamaChatRequest{ Model: model, Messages: wireMessages, Tools: wireTools, Stream: true, }) - if err != nil { - return nil, err - } - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL+"/api/chat", bytes.NewReader(data)) - if err != nil { - return nil, err - } + httpReq, _ := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL.JoinPath("/api/chat").String(), bytes.NewReader(data)) httpReq.Header.Set("Content-Type", "application/json") - resp, err := b.client.Do(httpReq) - if err != nil { - return nil, err - } - if resp.StatusCode == http.StatusTooManyRequests { - body, _ := io.ReadAll(resp.Body) - resp.Body.Close() - retryAfter := parseRetryAfter(resp.Header.Get("Retry-After")) - return nil, &RateLimitError{RetryAfter: retryAfter, Message: string(body)} - } - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - resp.Body.Close() - return nil, fmt.Errorf("ollama HTTP %d: %s", resp.StatusCode, body) - } - - ch := make(chan StreamEvent, 8) - - go func() { - defer close(ch) - defer resp.Body.Close() - - scanner := bufio.NewScanner(resp.Body) - scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) - - var accumulated []ToolCall - - for scanner.Scan() { - line := scanner.Bytes() - if len(line) == 0 { - continue - } - - var wire ollamaChatResponse - if err := json.Unmarshal(line, &wire); err != nil { - ch <- StreamEvent{Done: true, StopReason: fmt.Sprintf("stream decode: %v", err)} - return - } - - for _, tc := range wire.Message.ToolCalls { - accumulated = append(accumulated, ToolCall{ - Name: tc.Function.Name, - Arguments: tc.Function.Arguments, - }) - } - - if wire.Done { - stopReason := "stop" - if wire.DoneReason != "" { - stopReason = wire.DoneReason - } - ch <- StreamEvent{ - Done: true, - StopReason: stopReason, - Usage: Usage{InputTokens: wire.PromptEvalCount, OutputTokens: wire.EvalCount}, - ToolCalls: accumulated, - } - return - } - - ev := StreamEvent{} - if wire.Message.Content != "" { - ev.Content = wire.Message.Content - } - if ev.Content != "" { - ch <- ev - } - } - if err := scanner.Err(); err != nil { - ch <- StreamEvent{Done: true, StopReason: fmt.Sprintf("stream read: %v", err)} - } - }() - - return ch, nil + return streamRequest(b.client, httpReq, "ollama", streamOllamaNDJSON) +} + +// streamOllamaNDJSON reads Ollama's newline-delimited JSON stream from r +// and sends StreamEvents to ch. +func streamOllamaNDJSON(r io.Reader, ch chan<- StreamEvent) { + scanner := bufio.NewScanner(r) + scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) + + var accumulated []ToolCall + + for scanner.Scan() { + line := scanner.Bytes() + if len(line) == 0 { + continue + } + + var wire ollamaChatResponse + if err := json.Unmarshal(line, &wire); err != nil { + ch <- StreamEvent{Done: true, StopReason: fmt.Sprintf("stream decode: %v", err)} + return + } + + for _, tc := range wire.Message.ToolCalls { + accumulated = append(accumulated, ToolCall{ + Name: tc.Function.Name, + Arguments: tc.Function.Arguments, + }) + } + + if wire.Done { + stopReason := "stop" + if wire.DoneReason != "" { + stopReason = wire.DoneReason + } + ch <- StreamEvent{ + Done: true, + StopReason: stopReason, + Usage: Usage{InputTokens: wire.PromptEvalCount, OutputTokens: wire.EvalCount}, + ToolCalls: accumulated, + } + return + } + + if wire.Message.Content != "" { + ch <- StreamEvent{Content: wire.Message.Content} + } + } + if err := scanner.Err(); err != nil { + ch <- StreamEvent{Done: true, StopReason: fmt.Sprintf("stream read: %v", err)} + } } diff --git a/pkg/backend/ollama_stream_test.go b/pkg/backend/ollama_stream_test.go new file mode 100644 index 0000000..235fec0 --- /dev/null +++ b/pkg/backend/ollama_stream_test.go @@ -0,0 +1,142 @@ +package backend + +import ( + "fmt" + "strings" + "testing" +) + +func collectOllama(ndjson string) []StreamEvent { + ch := make(chan StreamEvent, 32) + go func() { + defer close(ch) + streamOllamaNDJSON(strings.NewReader(ndjson), ch) + }() + var evs []StreamEvent + for ev := range ch { + evs = append(evs, ev) + } + return evs +} + +func TestOllamaNDJSON_TextStream(t *testing.T) { + evs := collectOllama( + "{\"message\":{\"role\":\"assistant\",\"content\":\"hello\"},\"done\":false}\n" + + "{\"message\":{\"role\":\"assistant\",\"content\":\" world\"},\"done\":false}\n" + + "{\"done\":true,\"done_reason\":\"stop\",\"message\":{\"role\":\"assistant\",\"content\":\"\"},\"prompt_eval_count\":10,\"eval_count\":5}\n", + ) + var text string + for _, ev := range evs { + text += ev.Content + } + if text != "hello world" { + t.Errorf("text = %q", text) + } + fin := evs[len(evs)-1] + if !fin.Done || fin.StopReason != "stop" { + t.Errorf("final = %+v", fin) + } + if fin.Usage.InputTokens != 10 || fin.Usage.OutputTokens != 5 { + t.Errorf("usage = %+v", fin.Usage) + } +} + +func TestOllamaNDJSON_ToolCalls(t *testing.T) { + evs := collectOllama( + `{"done":true,"done_reason":"stop","message":{"role":"assistant","content":"","tool_calls":[{"function":{"name":"f","arguments":{"k":"v"}}}]},"prompt_eval_count":1,"eval_count":1}` + "\n", + ) + fin := evs[len(evs)-1] + if len(fin.ToolCalls) != 1 { + t.Fatalf("tool_calls len = %d", len(fin.ToolCalls)) + } + if fin.ToolCalls[0].Name != "f" { + t.Errorf("name = %q", fin.ToolCalls[0].Name) + } + if string(fin.ToolCalls[0].Arguments) != `{"k":"v"}` { + t.Errorf("args = %s", fin.ToolCalls[0].Arguments) + } +} + +func TestOllamaNDJSON_DoneReasonEmpty(t *testing.T) { + evs := collectOllama( + "{\"done\":true,\"message\":{\"role\":\"assistant\",\"content\":\"\"}}\n", + ) + if evs[0].StopReason != "stop" { + t.Errorf("stop = %q; want stop (default)", evs[0].StopReason) + } +} + +func TestOllamaNDJSON_MalformedJSON(t *testing.T) { + evs := collectOllama("{not json}\n") + if len(evs) != 1 || !evs[0].Done { + t.Fatal("expected Done event") + } + if !strings.Contains(evs[0].StopReason, "stream decode") { + t.Errorf("stop = %q", evs[0].StopReason) + } +} + +func TestOllamaNDJSON_EmptyLines(t *testing.T) { + evs := collectOllama( + "\n\n" + + "{\"message\":{\"role\":\"assistant\",\"content\":\"x\"},\"done\":false}\n" + + "\n" + + "{\"done\":true,\"done_reason\":\"stop\",\"message\":{\"role\":\"assistant\",\"content\":\"\"}}\n", + ) + var text string + for _, ev := range evs { + text += ev.Content + } + if text != "x" { + t.Errorf("text = %q", text) + } +} + +func TestOllamaNDJSON_EmptyContent(t *testing.T) { + evs := collectOllama( + "{\"message\":{\"role\":\"assistant\",\"content\":\"\"},\"done\":false}\n" + + "{\"message\":{\"role\":\"assistant\",\"content\":\"hi\"},\"done\":false}\n" + + "{\"done\":true,\"done_reason\":\"stop\",\"message\":{\"role\":\"assistant\",\"content\":\"\"}}\n", + ) + // Empty content chunks should not produce events. + var contentEvents int + for _, ev := range evs { + if ev.Content != "" { + contentEvents++ + } + } + if contentEvents != 1 { + t.Errorf("content events = %d; want 1", contentEvents) + } +} + +func TestOllamaNDJSON_EOFWithoutDone(t *testing.T) { + evs := collectOllama( + "{\"message\":{\"role\":\"assistant\",\"content\":\"partial\"},\"done\":false}\n", + ) + // No done event — channel just closes. Content should still arrive. + var text string + for _, ev := range evs { + text += ev.Content + } + if text != "partial" { + t.Errorf("text = %q", text) + } +} + +func TestOllamaNDJSON_ScannerError(t *testing.T) { + r := newErrAfterReader("{\"message\":{\"role\":\"assistant\",\"content\":\"hi\"},\"done\":false}\n", fmt.Errorf("connection reset")) + ch := make(chan StreamEvent, 32) + go func() { + defer close(ch) + streamOllamaNDJSON(r, ch) + }() + var evs []StreamEvent + for ev := range ch { + evs = append(evs, ev) + } + fin := evs[len(evs)-1] + if !fin.Done || !strings.Contains(fin.StopReason, "stream read") { + t.Errorf("final = %+v", fin) + } +} diff --git a/pkg/backend/openai.go b/pkg/backend/openai.go index db3655f..b3c5d06 100644 --- a/pkg/backend/openai.go +++ b/pkg/backend/openai.go @@ -8,6 +8,7 @@ import ( "fmt" "io" "net/http" + "net/url" "strconv" "strings" "time" @@ -17,7 +18,7 @@ import ( // Compatible with OpenRouter, OpenAI, and any other OpenAI-compatible API. type OpenAIBackend struct { name string - baseURL string + baseURL *url.URL apiKey string model string client *http.Client @@ -32,13 +33,17 @@ type openAIModelInfo struct { ContextLength int `json:"context_length"` } -func NewOpenAI(name, baseURL, apiKey string) *OpenAIBackend { +func NewOpenAI(name, baseURL, apiKey string) (*OpenAIBackend, error) { if baseURL == "" { baseURL = "https://api.openai.com" } - b := &OpenAIBackend{name: name, baseURL: baseURL, apiKey: apiKey, client: &http.Client{}} + u, err := url.Parse(baseURL) + if err != nil { + return nil, fmt.Errorf("invalid base URL: %w", err) + } + b := &OpenAIBackend{name: name, baseURL: u, apiKey: apiKey, client: &http.Client{}} b.model = b.DefaultModel() - return b + return b, nil } func (b *OpenAIBackend) Name() string { return b.name } @@ -49,7 +54,7 @@ func (b *OpenAIBackend) fetchModels(ctx context.Context) []openAIModelInfo { if len(b.cachedModels) > 0 { return b.cachedModels } - req, err := http.NewRequestWithContext(ctx, "GET", b.baseURL+"/v1/models", nil) + req, err := http.NewRequestWithContext(ctx, "GET", b.baseURL.JoinPath("/v1/models").String(), nil) if err != nil { return nil } @@ -177,11 +182,11 @@ type openAIStreamResponse struct { Usage openAIUsage `json:"usage"` } -// -- implementation -- +// -- encoding -- -func (b *OpenAIBackend) ChatStream(ctx context.Context, messages []Message, tools []Tool, params GenerationParams) (<-chan StreamEvent, error) { - model := b.model - wireMessages := make([]openAIMessage, len(messages)) +// encodeOpenAIMessages converts canonical Messages to the OpenAI wire format. +func encodeOpenAIMessages(messages []Message) []openAIMessage { + wire := make([]openAIMessage, len(messages)) for i, m := range messages { wm := openAIMessage{ Role: m.Role, @@ -201,12 +206,16 @@ func (b *OpenAIBackend) ChatStream(ctx context.Context, messages []Message, tool if len(wm.ToolCalls) == 0 { wm.Content = &m.Content } - wireMessages[i] = wm + wire[i] = wm } + return wire +} - var wireTools []openAITool +// encodeOpenAITools converts canonical Tools to the OpenAI wire format. +func encodeOpenAITools(tools []Tool) []openAITool { + var wire []openAITool for _, t := range tools { - wireTools = append(wireTools, openAITool{ + wire = append(wire, openAITool{ Type: "function", Function: openAIToolFunction{ Name: t.Name, @@ -215,11 +224,105 @@ func (b *OpenAIBackend) ChatStream(ctx context.Context, messages []Message, tool }, }) } + return wire +} +// -- stream parsing -- + +// streamOpenAISSE reads an OpenAI-format SSE stream from r and sends +// StreamEvents to ch. It is the core parsing loop, separated from HTTP +// transport so it can be tested with an io.Reader directly. +func streamOpenAISSE(r io.Reader, ch chan<- StreamEvent) { + scanner := bufio.NewScanner(r) + scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) + + // Accumulate tool call fragments by index (OpenAI sends args incrementally). + type tcAccum struct { + id, name, args string + } + accum := make(map[int]*tcAccum) + + // finishReason and usage arrive in separate trailing chunks. + var finishReason string + var usage Usage + + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + if line == "data: [DONE]" { + ev := StreamEvent{ + Done: true, + StopReason: finishReason, + Usage: usage, + } + for _, a := range accum { + ev.ToolCalls = append(ev.ToolCalls, ToolCall{ + ID: a.id, + Name: a.name, + Arguments: json.RawMessage(a.args), + }) + } + ch <- ev + return + } + if !strings.HasPrefix(line, "data: ") { + continue + } + + var wire openAIStreamResponse + if err := json.Unmarshal([]byte(line[6:]), &wire); err != nil { + ch <- StreamEvent{Done: true, StopReason: fmt.Sprintf("stream decode: %v", err)} + return + } + + if wire.Usage.PromptTokens > 0 { + usage.InputTokens = wire.Usage.PromptTokens + } + if wire.Usage.CompletionTokens > 0 { + usage.OutputTokens = wire.Usage.CompletionTokens + } + + if len(wire.Choices) == 0 { + continue + } + + choice := wire.Choices[0] + if choice.FinishReason != "" { + finishReason = choice.FinishReason + } + if choice.Delta.Content != "" { + ch <- StreamEvent{Content: choice.Delta.Content} + } + + for _, dtc := range choice.Delta.ToolCalls { + idx := dtc.Index + if _, ok := accum[idx]; !ok { + accum[idx] = &tcAccum{id: dtc.ID, name: dtc.Function.Name} + } + a := accum[idx] + if dtc.ID != "" { + a.id = dtc.ID + } + if dtc.Function.Name != "" { + a.name = dtc.Function.Name + } + a.args += dtc.Function.Arguments + } + } + if err := scanner.Err(); err != nil { + ch <- StreamEvent{Done: true, StopReason: fmt.Sprintf("stream read: %v", err)} + } +} + +// -- implementation -- + +func (b *OpenAIBackend) ChatStream(ctx context.Context, messages []Message, tools []Tool, params GenerationParams) (<-chan StreamEvent, error) { req := openAIChatRequest{ - Model: model, - Messages: wireMessages, - Tools: wireTools, + Model: b.model, + Messages: encodeOpenAIMessages(messages), + Tools: encodeOpenAITools(tools), Stream: true, StreamOptions: &openAIStreamOptions{IncludeUsage: true}, MaxTokens: params.MaxTokens, @@ -228,15 +331,9 @@ func (b *OpenAIBackend) ChatStream(ctx context.Context, messages []Message, tool PresencePenalty: params.PresencePenalty, } - data, err := json.Marshal(req) - if err != nil { - return nil, err - } + data, _ := json.Marshal(req) - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL+"/v1/chat/completions", bytes.NewReader(data)) - if err != nil { - return nil, err - } + httpReq, _ := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL.JoinPath("/v1/chat/completions").String(), bytes.NewReader(data)) httpReq.Header.Set("Content-Type", "application/json") if b.apiKey != "" { httpReq.Header.Set("Authorization", "Bearer "+b.apiKey) @@ -245,112 +342,7 @@ func (b *OpenAIBackend) ChatStream(ctx context.Context, messages []Message, tool httpReq.Header.Set(k, v) } - resp, err := b.client.Do(httpReq) - if err != nil { - return nil, err - } - if resp.StatusCode == http.StatusTooManyRequests { - body, _ := io.ReadAll(resp.Body) - resp.Body.Close() - retryAfter := parseRetryAfter(resp.Header.Get("Retry-After")) - return nil, &RateLimitError{RetryAfter: retryAfter, Message: string(body)} - } - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - resp.Body.Close() - return nil, fmt.Errorf("openai HTTP %d: %s", resp.StatusCode, body) - } - - ch := make(chan StreamEvent, 8) - - go func() { - defer close(ch) - defer resp.Body.Close() - - scanner := bufio.NewScanner(resp.Body) - scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) - - // Accumulate tool call fragments by index (OpenAI sends args incrementally). - type tcAccum struct { - id, name, args string - } - accum := make(map[int]*tcAccum) - - // finishReason and usage arrive in separate trailing chunks. - var finishReason string - var usage Usage - - for scanner.Scan() { - line := strings.TrimSpace(scanner.Text()) - if line == "" { - continue - } - if line == "data: [DONE]" { - ev := StreamEvent{ - Done: true, - StopReason: finishReason, - Usage: usage, - } - for _, a := range accum { - ev.ToolCalls = append(ev.ToolCalls, ToolCall{ - ID: a.id, - Name: a.name, - Arguments: json.RawMessage(a.args), - }) - } - ch <- ev - return - } - if !strings.HasPrefix(line, "data: ") { - continue - } - - var wire openAIStreamResponse - if err := json.Unmarshal([]byte(line[6:]), &wire); err != nil { - ch <- StreamEvent{Done: true, StopReason: fmt.Sprintf("stream decode: %v", err)} - return - } - - if wire.Usage.PromptTokens > 0 { - usage.InputTokens = wire.Usage.PromptTokens - } - if wire.Usage.CompletionTokens > 0 { - usage.OutputTokens = wire.Usage.CompletionTokens - } - - if len(wire.Choices) == 0 { - continue - } - - choice := wire.Choices[0] - if choice.FinishReason != "" { - finishReason = choice.FinishReason - } - if choice.Delta.Content != "" { - ch <- StreamEvent{Content: choice.Delta.Content} - } - - for _, dtc := range choice.Delta.ToolCalls { - idx := dtc.Index - if _, ok := accum[idx]; !ok { - accum[idx] = &tcAccum{id: dtc.ID, name: dtc.Function.Name} - } - a := accum[idx] - if dtc.ID != "" { - a.id = dtc.ID - } - if dtc.Function.Name != "" { - a.name = dtc.Function.Name - } - a.args += dtc.Function.Arguments - } - } - if err := scanner.Err(); err != nil { - ch <- StreamEvent{Done: true, StopReason: fmt.Sprintf("stream read: %v", err)} - } - }() - - return ch, nil + return streamRequest(b.client, httpReq, "openai", streamOpenAISSE) } // parseRetryAfter parses the Retry-After header value, which may be an integer diff --git a/pkg/backend/openai_encode_test.go b/pkg/backend/openai_encode_test.go new file mode 100644 index 0000000..afbe11b --- /dev/null +++ b/pkg/backend/openai_encode_test.go @@ -0,0 +1,210 @@ +package backend + +import ( + "encoding/json" + "testing" +) + +// --- encodeOpenAIMessages --- + +func TestEncodeMessages_UserContent(t *testing.T) { + wire := encodeOpenAIMessages([]Message{{Role: "user", Content: "hello"}}) + if len(wire) != 1 { + t.Fatalf("len = %d; want 1", len(wire)) + } + if wire[0].Role != "user" || wire[0].Content == nil || *wire[0].Content != "hello" { + t.Errorf("got %+v", wire[0]) + } +} + +func TestEncodeMessages_SystemContent(t *testing.T) { + wire := encodeOpenAIMessages([]Message{{Role: "system", Content: "be helpful"}}) + if wire[0].Role != "system" || *wire[0].Content != "be helpful" { + t.Errorf("got %+v", wire[0]) + } +} + +func TestEncodeMessages_EmptyContent(t *testing.T) { + wire := encodeOpenAIMessages([]Message{{Role: "user", Content: ""}}) + if wire[0].Content == nil { + t.Error("content should be non-nil pointer to empty string") + } + if *wire[0].Content != "" { + t.Errorf("content = %q; want empty", *wire[0].Content) + } +} + +func TestEncodeMessages_AssistantWithToolCalls_NullContent(t *testing.T) { + wire := encodeOpenAIMessages([]Message{{ + Role: "assistant", + ToolCalls: []ToolCall{{ + ID: "c1", Name: "run", Arguments: json.RawMessage(`{}`), + }}, + }}) + if wire[0].Content != nil { + t.Error("content must be nil when tool_calls present") + } + if len(wire[0].ToolCalls) != 1 { + t.Fatalf("tool_calls len = %d", len(wire[0].ToolCalls)) + } + tc := wire[0].ToolCalls[0] + if tc.ID != "c1" || tc.Type != "function" || tc.Function.Name != "run" || tc.Function.Arguments != "{}" { + t.Errorf("tool_call = %+v", tc) + } +} + +func TestEncodeMessages_AssistantTextOnly_HasContent(t *testing.T) { + wire := encodeOpenAIMessages([]Message{{Role: "assistant", Content: "sure"}}) + if wire[0].Content == nil || *wire[0].Content != "sure" { + t.Errorf("got %+v", wire[0]) + } + if len(wire[0].ToolCalls) != 0 { + t.Error("tool_calls should be empty") + } +} + +func TestEncodeMessages_ToolResult(t *testing.T) { + wire := encodeOpenAIMessages([]Message{{Role: "tool", Content: "ok", ToolCallID: "c1"}}) + if wire[0].ToolCallID != "c1" { + t.Errorf("tool_call_id = %q", wire[0].ToolCallID) + } + if wire[0].Content == nil || *wire[0].Content != "ok" { + t.Errorf("content = %v", wire[0].Content) + } +} + +func TestEncodeMessages_MultipleToolCalls(t *testing.T) { + wire := encodeOpenAIMessages([]Message{{ + Role: "assistant", + ToolCalls: []ToolCall{ + {ID: "a", Name: "x", Arguments: json.RawMessage(`{}`)}, + {ID: "b", Name: "y", Arguments: json.RawMessage(`{"k":"v"}`)}, + }, + }}) + if len(wire[0].ToolCalls) != 2 { + t.Fatalf("len = %d; want 2", len(wire[0].ToolCalls)) + } + if wire[0].ToolCalls[1].Function.Arguments != `{"k":"v"}` { + t.Errorf("args = %q", wire[0].ToolCalls[1].Function.Arguments) + } +} + +func TestEncodeMessages_Empty(t *testing.T) { + wire := encodeOpenAIMessages(nil) + if len(wire) != 0 { + t.Errorf("len = %d; want 0", len(wire)) + } +} + +func TestEncodeMessages_FullConversation(t *testing.T) { + msgs := []Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "hi"}, + {Role: "assistant", ToolCalls: []ToolCall{{ID: "c1", Name: "f", Arguments: json.RawMessage(`{}`)}}}, + {Role: "tool", Content: "result", ToolCallID: "c1"}, + {Role: "assistant", Content: "done"}, + } + wire := encodeOpenAIMessages(msgs) + if len(wire) != 5 { + t.Fatalf("len = %d; want 5", len(wire)) + } + // system: has content + if wire[0].Content == nil { + t.Error("system content nil") + } + // assistant with tool_calls: null content + if wire[2].Content != nil { + t.Error("assistant[2] content should be nil") + } + // tool: has content + tool_call_id + if wire[3].ToolCallID != "c1" { + t.Error("tool missing tool_call_id") + } + // final assistant: has content + if wire[4].Content == nil || *wire[4].Content != "done" { + t.Error("assistant[4] content wrong") + } +} + +// --- encodeOpenAITools --- + +func TestEncodeTools_Empty(t *testing.T) { + wire := encodeOpenAITools(nil) + if wire != nil { + t.Errorf("got %v; want nil", wire) + } +} + +func TestEncodeTools_Single(t *testing.T) { + wire := encodeOpenAITools([]Tool{{ + Name: "f", Description: "desc", Parameters: json.RawMessage(`{"type":"object"}`), + }}) + if len(wire) != 1 { + t.Fatalf("len = %d", len(wire)) + } + if wire[0].Type != "function" { + t.Errorf("type = %q", wire[0].Type) + } + if wire[0].Function.Name != "f" || wire[0].Function.Description != "desc" { + t.Errorf("function = %+v", wire[0].Function) + } + if string(wire[0].Function.Parameters) != `{"type":"object"}` { + t.Errorf("parameters = %s", wire[0].Function.Parameters) + } +} + +func TestEncodeTools_Multiple(t *testing.T) { + wire := encodeOpenAITools([]Tool{ + {Name: "a", Description: "da", Parameters: json.RawMessage(`{}`)}, + {Name: "b", Description: "db", Parameters: json.RawMessage(`{}`)}, + }) + if len(wire) != 2 { + t.Fatalf("len = %d", len(wire)) + } + if wire[0].Function.Name != "a" || wire[1].Function.Name != "b" { + t.Error("names wrong") + } +} + +// --- JSON round-trip: verify wire format matches OpenAI spec --- + +func TestEncodeMessages_JSONRoundTrip(t *testing.T) { + wire := encodeOpenAIMessages([]Message{ + {Role: "assistant", ToolCalls: []ToolCall{{ID: "c1", Name: "f", Arguments: json.RawMessage(`{"a":1}`)}}}, + }) + data, err := json.Marshal(wire[0]) + if err != nil { + t.Fatal(err) + } + // content must be absent (null), not empty string + var raw map[string]json.RawMessage + json.Unmarshal(data, &raw) + if string(raw["content"]) != "null" { + t.Errorf("content JSON = %s; want null", raw["content"]) + } + // tool_calls[0].type must be "function" + var parsed openAIMessage + json.Unmarshal(data, &parsed) + if parsed.ToolCalls[0].Type != "function" { + t.Errorf("type = %q", parsed.ToolCalls[0].Type) + } + // arguments must be a string, not an object + if parsed.ToolCalls[0].Function.Arguments != `{"a":1}` { + t.Errorf("arguments = %q", parsed.ToolCalls[0].Function.Arguments) + } +} + +func TestEncodeRequest_ToolsOmittedWhenEmpty(t *testing.T) { + req := openAIChatRequest{ + Model: "m", + Messages: encodeOpenAIMessages([]Message{{Role: "user", Content: "hi"}}), + Tools: encodeOpenAITools(nil), + Stream: true, + } + data, _ := json.Marshal(req) + var raw map[string]json.RawMessage + json.Unmarshal(data, &raw) + if _, ok := raw["tools"]; ok { + t.Error("tools should be omitted when nil") + } +} diff --git a/pkg/backend/openai_stream_test.go b/pkg/backend/openai_stream_test.go new file mode 100644 index 0000000..6e15024 --- /dev/null +++ b/pkg/backend/openai_stream_test.go @@ -0,0 +1,287 @@ +package backend + +import ( + "fmt" + "strings" + "testing" +) + +// collect runs streamOpenAISSE on raw SSE text and returns all events. +func collect(sse string) []StreamEvent { + ch := make(chan StreamEvent, 32) + go func() { + defer close(ch) + streamOpenAISSE(strings.NewReader(sse), ch) + }() + var evs []StreamEvent + for ev := range ch { + evs = append(evs, ev) + } + return evs +} + +func lastEvent(evs []StreamEvent) StreamEvent { + return evs[len(evs)-1] +} + +// --- basic text streaming --- + +func TestSSE_TextDeltas(t *testing.T) { + evs := collect( + "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n" + + "data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n" + + "data: {\"choices\":[{\"finish_reason\":\"stop\",\"delta\":{}}]}\n\n" + + "data: [DONE]\n\n", + ) + var text string + for _, ev := range evs { + text += ev.Content + } + if text != "Hello world" { + t.Errorf("text = %q", text) + } + fin := lastEvent(evs) + if !fin.Done || fin.StopReason != "stop" { + t.Errorf("final = %+v", fin) + } +} + +// --- usage in trailing chunk --- + +func TestSSE_UsageChunk(t *testing.T) { + evs := collect( + "data: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n" + + "data: {\"choices\":[{\"finish_reason\":\"stop\",\"delta\":{}}]}\n\n" + + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":42,\"completion_tokens\":7}}\n\n" + + "data: [DONE]\n\n", + ) + fin := lastEvent(evs) + if fin.Usage.InputTokens != 42 || fin.Usage.OutputTokens != 7 { + t.Errorf("usage = %+v", fin.Usage) + } +} + +// --- usage inline with content --- + +func TestSSE_UsageInline(t *testing.T) { + evs := collect( + "data: {\"choices\":[{\"delta\":{\"content\":\"x\"}}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":3}}\n\n" + + "data: {\"choices\":[{\"finish_reason\":\"stop\",\"delta\":{}}]}\n\n" + + "data: [DONE]\n\n", + ) + fin := lastEvent(evs) + if fin.Usage.InputTokens != 10 || fin.Usage.OutputTokens != 3 { + t.Errorf("usage = %+v", fin.Usage) + } +} + +// --- tool call accumulation --- + +func TestSSE_ToolCallAccumulation(t *testing.T) { + evs := collect( + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","type":"function","function":{"name":"f","arguments":""}}]}}]}` + "\n\n" + + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"a\":"}}]}}]}` + "\n\n" + + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"1}"}}]}}]}` + "\n\n" + + `data: {"choices":[{"finish_reason":"tool_calls","delta":{}}]}` + "\n\n" + + "data: [DONE]\n\n", + ) + fin := lastEvent(evs) + if fin.StopReason != "tool_calls" { + t.Errorf("stop = %q", fin.StopReason) + } + if len(fin.ToolCalls) != 1 { + t.Fatalf("tool_calls len = %d", len(fin.ToolCalls)) + } + tc := fin.ToolCalls[0] + if tc.ID != "c1" || tc.Name != "f" || string(tc.Arguments) != `{"a":1}` { + t.Errorf("tc = %+v", tc) + } +} + +// --- parallel tool calls --- + +func TestSSE_ParallelToolCalls(t *testing.T) { + evs := collect( + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"a","type":"function","function":{"name":"x","arguments":"{}"}}]}}]}` + "\n\n" + + `data: {"choices":[{"delta":{"tool_calls":[{"index":1,"id":"b","type":"function","function":{"name":"y","arguments":"{}"}}]}}]}` + "\n\n" + + `data: {"choices":[{"finish_reason":"tool_calls","delta":{}}]}` + "\n\n" + + "data: [DONE]\n\n", + ) + fin := lastEvent(evs) + if len(fin.ToolCalls) != 2 { + t.Fatalf("len = %d", len(fin.ToolCalls)) + } + ids := map[string]bool{} + for _, tc := range fin.ToolCalls { + ids[tc.ID] = true + } + if !ids["a"] || !ids["b"] { + t.Errorf("ids = %v", ids) + } +} + +// --- text + tool calls mixed --- + +func TestSSE_TextThenToolCall(t *testing.T) { + evs := collect( + `data: {"choices":[{"delta":{"content":"thinking..."}}]}` + "\n\n" + + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","type":"function","function":{"name":"f","arguments":"{}"}}]}}]}` + "\n\n" + + `data: {"choices":[{"finish_reason":"tool_calls","delta":{}}]}` + "\n\n" + + "data: [DONE]\n\n", + ) + var text string + for _, ev := range evs { + text += ev.Content + } + if text != "thinking..." { + t.Errorf("text = %q", text) + } + fin := lastEvent(evs) + if len(fin.ToolCalls) != 1 || fin.ToolCalls[0].Name != "f" { + t.Errorf("tool_calls = %+v", fin.ToolCalls) + } +} + +// --- malformed JSON --- + +func TestSSE_MalformedJSON(t *testing.T) { + evs := collect("data: {not json}\n\n") + if len(evs) != 1 || !evs[0].Done { + t.Fatal("expected single Done event") + } + if !strings.Contains(evs[0].StopReason, "stream decode") { + t.Errorf("stop = %q", evs[0].StopReason) + } +} + +// --- empty choices array (usage-only chunk) --- + +func TestSSE_EmptyChoicesSkipped(t *testing.T) { + evs := collect( + "data: {\"choices\":[]}\n\n" + + "data: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n" + + "data: [DONE]\n\n", + ) + var text string + for _, ev := range evs { + text += ev.Content + } + if text != "x" { + t.Errorf("text = %q", text) + } +} + +// --- non-data lines (comments, event: lines) are ignored --- + +func TestSSE_NonDataLinesIgnored(t *testing.T) { + evs := collect( + ": this is a comment\n" + + "event: ping\n" + + "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\n" + + "data: [DONE]\n\n", + ) + var text string + for _, ev := range evs { + text += ev.Content + } + if text != "ok" { + t.Errorf("text = %q", text) + } +} + +// --- bare DONE with no prior content --- + +func TestSSE_BareDone(t *testing.T) { + evs := collect("data: [DONE]\n\n") + if len(evs) != 1 || !evs[0].Done { + t.Fatal("expected single Done event") + } + if evs[0].StopReason != "" { + t.Errorf("stop = %q; want empty", evs[0].StopReason) + } +} + +// --- stream ends without DONE (EOF) --- + +func TestSSE_EOFWithoutDone(t *testing.T) { + evs := collect("data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n") + // No Done event from scanner.Err (which is nil on clean EOF). + // Channel just closes. The content event should still arrive. + var text string + for _, ev := range evs { + text += ev.Content + } + if text != "partial" { + t.Errorf("text = %q", text) + } +} + +// --- finish_reason "length" (max tokens) --- + +func TestSSE_FinishReasonLength(t *testing.T) { + evs := collect( + "data: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n" + + "data: {\"choices\":[{\"finish_reason\":\"length\",\"delta\":{}}]}\n\n" + + "data: [DONE]\n\n", + ) + fin := lastEvent(evs) + if fin.StopReason != "length" { + t.Errorf("stop = %q; want length", fin.StopReason) + } +} + +// --- tool call ID/name arrive in later chunks --- + +func TestSSE_ToolCallIDUpdatedLater(t *testing.T) { + evs := collect( + // First chunk: index only, no ID yet + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"","type":"function","function":{"name":"f","arguments":""}}]}}]}` + "\n\n" + + // Second chunk: ID arrives + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"late_id","function":{"arguments":"{}"}}]}}]}` + "\n\n" + + "data: [DONE]\n\n", + ) + fin := lastEvent(evs) + if len(fin.ToolCalls) != 1 { + t.Fatalf("len = %d", len(fin.ToolCalls)) + } + if fin.ToolCalls[0].ID != "late_id" { + t.Errorf("id = %q; want late_id", fin.ToolCalls[0].ID) + } +} + +func TestSSE_ScannerError(t *testing.T) { + // Reader that returns data then an error (not EOF). + r := newErrAfterReader("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n", fmt.Errorf("connection reset")) + ch := make(chan StreamEvent, 32) + go func() { + defer close(ch) + streamOpenAISSE(r, ch) + }() + var evs []StreamEvent + for ev := range ch { + evs = append(evs, ev) + } + fin := evs[len(evs)-1] + if !fin.Done || !strings.Contains(fin.StopReason, "stream read") { + t.Errorf("final = %+v", fin) + } +} + +// --- whitespace-only lines between data lines --- + +func TestSSE_WhitespaceLines(t *testing.T) { + evs := collect( + " \n" + + "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n" + + "\t\n" + + "data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n" + + "data: [DONE]\n\n", + ) + var text string + for _, ev := range evs { + text += ev.Content + } + if text != "ab" { + t.Errorf("text = %q", text) + } +} diff --git a/pkg/backend/openai_test.go b/pkg/backend/openai_test.go new file mode 100644 index 0000000..a4a22f5 --- /dev/null +++ b/pkg/backend/openai_test.go @@ -0,0 +1,494 @@ +package backend + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// --- HTTP integration: request encoding --- + +func TestChatStream_RequestWireFormat(t *testing.T) { + var got openAIChatRequest + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + t.Errorf("method = %s; want POST", r.Method) + } + if r.URL.Path != "/v1/chat/completions" { + t.Errorf("path = %s; want /v1/chat/completions", r.URL.Path) + } + if ct := r.Header.Get("Content-Type"); ct != "application/json" { + t.Errorf("Content-Type = %q", ct) + } + if auth := r.Header.Get("Authorization"); auth != "Bearer test-key" { + t.Errorf("Authorization = %q", auth) + } + body, _ := io.ReadAll(r.Body) + json.Unmarshal(body, &got) + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer srv.Close() + + temp := 0.7 + b := mustOpenAI("openai", srv.URL, "test-key") + b.SetModel("gpt-4") + + ch, err := b.ChatStream(context.Background(), + []Message{ + {Role: "system", Content: "you are helpful"}, + {Role: "user", Content: "hello"}, + }, + []Tool{{Name: "get_weather", Description: "Get weather", Parameters: json.RawMessage(`{"type":"object"}`)}}, + GenerationParams{MaxTokens: 100, Temperature: &temp}, + ) + if err != nil { + t.Fatal(err) + } + for range ch { + } + + if got.Model != "gpt-4" { + t.Errorf("model = %q", got.Model) + } + if !got.Stream { + t.Error("stream should be true") + } + if got.StreamOptions == nil || !got.StreamOptions.IncludeUsage { + t.Error("stream_options.include_usage should be true") + } + if got.MaxTokens != 100 { + t.Errorf("max_tokens = %d", got.MaxTokens) + } + if got.Temperature == nil || *got.Temperature != 0.7 { + t.Errorf("temperature = %v", got.Temperature) + } + if len(got.Messages) != 2 { + t.Fatalf("messages len = %d", len(got.Messages)) + } + if len(got.Tools) != 1 || got.Tools[0].Function.Name != "get_weather" { + t.Errorf("tools = %+v", got.Tools) + } +} + +func TestChatStream_NoAuthWhenKeyEmpty(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if auth := r.Header.Get("Authorization"); auth != "" { + t.Errorf("Authorization = %q; want empty", auth) + } + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "") + ch, _ := b.ChatStream(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, GenerationParams{}) + for range ch { + } +} + +func TestChatStream_ExtraHeaders(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if v := r.Header.Get("X-Custom"); v != "val" { + t.Errorf("X-Custom = %q; want val", v) + } + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + b.extraHeaders = map[string]string{"X-Custom": "val"} + ch, _ := b.ChatStream(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, GenerationParams{}) + for range ch { + } +} + +func TestChatStream_OmitsZeroParams(t *testing.T) { + var raw []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + raw, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + ch, _ := b.ChatStream(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, GenerationParams{}) + for range ch { + } + + s := string(raw) + for _, field := range []string{"max_tokens", "temperature", "frequency_penalty", "presence_penalty"} { + if strings.Contains(s, field) { + t.Errorf("zero-value %s should be omitted", field) + } + } +} + +func TestChatStream_AllParamsPresent(t *testing.T) { + var raw []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + raw, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer srv.Close() + + temp, fp, pp := 0.5, 0.3, 0.1 + b := mustOpenAI("openai", srv.URL, "k") + ch, _ := b.ChatStream(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, + GenerationParams{MaxTokens: 200, Temperature: &temp, FrequencyPenalty: &fp, PresencePenalty: &pp}) + for range ch { + } + + s := string(raw) + for _, field := range []string{"max_tokens", "temperature", "frequency_penalty", "presence_penalty"} { + if !strings.Contains(s, field) { + t.Errorf("expected %s in request body", field) + } + } +} + +// --- HTTP error handling --- + +func TestChatStream_RateLimit(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Retry-After", "30") + w.WriteHeader(http.StatusTooManyRequests) + fmt.Fprint(w, "quota exceeded") + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + _, err := b.ChatStream(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, GenerationParams{}) + rle, ok := err.(*RateLimitError) + if !ok { + t.Fatalf("type = %T; want *RateLimitError", err) + } + if rle.RetryAfter != 30*time.Second { + t.Errorf("RetryAfter = %v", rle.RetryAfter) + } + if !strings.Contains(rle.Message, "quota exceeded") { + t.Errorf("Message = %q", rle.Message) + } +} + +func TestChatStream_RateLimitNoRetryAfter(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusTooManyRequests) + fmt.Fprint(w, "slow down") + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + _, err := b.ChatStream(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, GenerationParams{}) + rle, ok := err.(*RateLimitError) + if !ok { + t.Fatalf("type = %T", err) + } + if rle.RetryAfter != 0 { + t.Errorf("RetryAfter = %v; want 0", rle.RetryAfter) + } +} + +func TestChatStream_HTTPError(t *testing.T) { + for _, code := range []int{400, 401, 403, 500, 502, 503} { + t.Run(fmt.Sprintf("HTTP_%d", code), func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(code) + fmt.Fprintf(w, "error %d", code) + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + _, err := b.ChatStream(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, GenerationParams{}) + if err == nil { + t.Fatal("expected error") + } + if !strings.Contains(err.Error(), fmt.Sprintf("%d", code)) { + t.Errorf("error = %q; want to contain %d", err.Error(), code) + } + }) + } +} + +func TestChatStream_ConnectionRefused(t *testing.T) { + b := mustOpenAI("openai", "http://127.0.0.1:1", "k") + _, err := b.ChatStream(context.Background(), []Message{{Role: "user", Content: "hi"}}, nil, GenerationParams{}) + if err == nil { + t.Fatal("expected error") + } +} + +func TestChatStream_ContextCancelled(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Block forever; context cancellation should unblock the client. + <-r.Context().Done() + })) + defer srv.Close() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // cancel immediately + + b := mustOpenAI("openai", srv.URL, "k") + _, err := b.ChatStream(ctx, []Message{{Role: "user", Content: "hi"}}, nil, GenerationParams{}) + if err == nil { + t.Fatal("expected error") + } +} + +// --- /v1/models endpoint --- + +func TestModels_Success(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/models" { + t.Errorf("path = %s", r.URL.Path) + } + if r.Method != "GET" { + t.Errorf("method = %s; want GET", r.Method) + } + if auth := r.Header.Get("Authorization"); auth != "Bearer k" { + t.Errorf("Authorization = %q", auth) + } + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"id": "gpt-4", "context_length": 8192}, + {"id": "gpt-3.5-turbo", "context_length": 4096}, + }, + }) + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + models := b.Models(context.Background()) + if len(models) != 2 || models[0] != "gpt-4" { + t.Errorf("models = %v", models) + } +} + +func TestModels_Cached(t *testing.T) { + calls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{{"id": "m1"}}, + }) + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + b.Models(context.Background()) + b.Models(context.Background()) + if calls != 1 { + t.Errorf("API calls = %d; want 1", calls) + } +} + +func TestModels_HTTPError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(500) + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + if models := b.Models(context.Background()); len(models) != 0 { + t.Errorf("models = %v; want empty", models) + } +} + +func TestModels_MalformedJSON(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, "not json") + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + if models := b.Models(context.Background()); len(models) != 0 { + t.Errorf("models = %v; want empty", models) + } +} + +func TestModels_NoAuthWhenKeyEmpty(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if auth := r.Header.Get("Authorization"); auth != "" { + t.Errorf("Authorization = %q; want empty", auth) + } + json.NewEncoder(w).Encode(map[string]any{"data": []map[string]any{}}) + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "") + b.Models(context.Background()) +} + +// --- ContextLength --- + +func TestContextLength_Lookup(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{{"id": "gpt-4", "context_length": 8192}}, + }) + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + b.SetModel("gpt-4") + if cl := b.ContextLength(context.Background()); cl != 8192 { + t.Errorf("context_length = %d", cl) + } +} + +func TestContextLength_Cached(t *testing.T) { + calls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{{"id": "gpt-4", "context_length": 8192}}, + }) + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + b.SetModel("gpt-4") + b.ContextLength(context.Background()) + b.ContextLength(context.Background()) + if calls != 1 { + t.Errorf("API calls = %d; want 1", calls) + } +} + +func TestContextLength_InvalidatedBySetModel(t *testing.T) { + calls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"id": "a", "context_length": 100}, + {"id": "b", "context_length": 200}, + }, + }) + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + b.SetModel("a") + if cl := b.ContextLength(context.Background()); cl != 100 { + t.Errorf("cl = %d; want 100", cl) + } + b.SetModel("b") + if cl := b.ContextLength(context.Background()); cl != 200 { + t.Errorf("cl = %d; want 200", cl) + } +} + +func TestContextLength_UnknownModel(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{"data": []map[string]any{}}) + })) + defer srv.Close() + + b := mustOpenAI("openai", srv.URL, "k") + b.SetModel("nonexistent") + if cl := b.ContextLength(context.Background()); cl != 0 { + t.Errorf("cl = %d; want 0", cl) + } +} + +// --- constructor / defaults --- + +func TestNewOpenAI_DefaultBaseURL(t *testing.T) { + b := mustOpenAI("openai", "", "k") + if b.baseURL.String() != "https://api.openai.com" { + t.Errorf("baseURL = %q", b.baseURL) + } +} + +func TestNewOpenAI_DefaultModel(t *testing.T) { + b := mustOpenAI("openai", "http://x", "k") + if b.Model() != b.DefaultModel() { + t.Errorf("model = %q; want %q", b.Model(), b.DefaultModel()) + } +} + +func TestDefaultModel(t *testing.T) { + tests := []struct { + name string + want string + }{ + {"openrouter", "deepseek/deepseek-v3.2"}, + {"anthropic", "claude-sonnet-4-5"}, + {"openai", "qwen3.5:9b"}, + {"unknown", "qwen3.5:9b"}, + } + for _, tt := range tests { + if got := (&OpenAIBackend{name: tt.name}).DefaultModel(); got != tt.want { + t.Errorf("DefaultModel(%q) = %q; want %q", tt.name, got, tt.want) + } + } +} + +func TestSetModel_ClearsCache(t *testing.T) { + b := &OpenAIBackend{ctxLength: 100, ctxModel: "old"} + b.SetModel("new") + if b.ctxLength != 0 { + t.Error("SetModel should clear ctxLength") + } + if b.model != "new" { + t.Errorf("model = %q", b.model) + } +} + +func TestName(t *testing.T) { + b := mustOpenAI("myname", "http://x", "k") + if b.Name() != "myname" { + t.Errorf("Name() = %q", b.Name()) + } +} + +// --- parseRetryAfter --- + +func TestParseRetryAfter_Empty(t *testing.T) { + if d := parseRetryAfter(""); d != 0 { + t.Errorf("got %v", d) + } +} + +func TestParseRetryAfter_Seconds(t *testing.T) { + if d := parseRetryAfter("120"); d != 120*time.Second { + t.Errorf("got %v", d) + } +} + +func TestParseRetryAfter_HTTPDate(t *testing.T) { + future := time.Now().Add(60 * time.Second).UTC().Format(http.TimeFormat) + d := parseRetryAfter(future) + if d < 50*time.Second || d > 70*time.Second { + t.Errorf("got %v; want ~60s", d) + } +} + +func TestParseRetryAfter_PastDate(t *testing.T) { + past := time.Now().Add(-60 * time.Second).UTC().Format(http.TimeFormat) + if d := parseRetryAfter(past); d != 0 { + t.Errorf("got %v; want 0 for past date", d) + } +} + +func TestParseRetryAfter_Invalid(t *testing.T) { + if d := parseRetryAfter("not-a-number"); d != 0 { + t.Errorf("got %v", d) + } +} + +func TestParseRetryAfter_Whitespace(t *testing.T) { + if d := parseRetryAfter(" 60 "); d != 60*time.Second { + t.Errorf("got %v", d) + } +} diff --git a/pkg/backend/testhelper_test.go b/pkg/backend/testhelper_test.go new file mode 100644 index 0000000..aa563f2 --- /dev/null +++ b/pkg/backend/testhelper_test.go @@ -0,0 +1,33 @@ +package backend + +import ( + "strings" +) + +// errAfterReader returns data on the first reads, then the given error. +type errAfterReader struct { + r *strings.Reader + err error + init bool +} + +func newErrAfterReader(data string, err error) *errAfterReader { + return &errAfterReader{r: strings.NewReader(data), err: err} +} + +func (r *errAfterReader) Read(p []byte) (int, error) { + n, err := r.r.Read(p) + if err != nil { + return n, r.err + } + return n, nil +} + +// mustOpenAI creates an OpenAIBackend for testing; panics on error. +func mustOpenAI(name, baseURL, apiKey string) *OpenAIBackend { + b, err := NewOpenAI(name, baseURL, apiKey) + if err != nil { + panic(err) + } + return b +}