lkn
/
ollie-core
Archived
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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
This commit is contained in:
Levi Neely 2026-04-21 15:32:03 +02:00
parent 549e1ccb45
commit 21827c57eb
13 changed files with 2216 additions and 323 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
pkg/backend/anthropic.go
View File

@ -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 {

186
pkg/backend/anthropic_stream_test.go Normal file
View File

@ -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)
}
}
}

29
pkg/backend/backend.go
View File

@ -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
}

682
pkg/backend/backend_test.go
View File

@ -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)
}
}

15
pkg/backend/copilot.go
View File

@ -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
}

8
pkg/backend/new.go
View File

@ -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:

152
pkg/backend/ollama.go
View File

@ -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)}
}
}

142
pkg/backend/ollama_stream_test.go Normal file
View File

@ -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)
}
}

250
pkg/backend/openai.go
View File

@ -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

210
pkg/backend/openai_encode_test.go Normal file
View File

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

287
pkg/backend/openai_stream_test.go Normal file
View File

@ -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)
}
}

494
pkg/backend/openai_test.go Normal file
View File

@ -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)
}
}

33
pkg/backend/testhelper_test.go Normal file
View File

@ -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
}