package llm import ( "context" "encoding/json" "net/http" "net/http/httptest" "testing" "time" ) // mockServer returns an OpenAI-compatible endpoint that always replies with the // given assistant content. func mockServer(t *testing.T, response string) *httptest.Server { t.Helper() return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/chat/completions" { t.Errorf("path = %q, want /chat/completions", r.URL.Path) } if ct := r.Header.Get("Content-Type"); ct != "application/json" { t.Errorf("Content-Type = %q, want application/json", ct) } w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(map[string]any{ "choices": []map[string]any{ {"message": map[string]any{"role": "assistant", "content": response}}, }, }) })) } func TestClient_Complete(t *testing.T) { srv := mockServer(t, "hello world") defer srv.Close() c := New(srv.URL, "", "test-model", 10*time.Second) got, err := c.Complete(context.Background(), "you are helpful", "say hello") if err != nil { t.Fatalf("Complete: %v", err) } if got != "hello world" { t.Errorf("got %q, want %q", got, "hello world") } } // TestClient_SendsMaxTokens guards Tapir's ADR-004 change to the copied client: // it MUST send a positive max_tokens, or thinking models return empty content. func TestClient_SendsMaxTokens(t *testing.T) { var body chatRequest srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _ = json.NewDecoder(r.Body).Decode(&body) _ = json.NewEncoder(w).Encode(map[string]any{ "choices": []map[string]any{{"message": map[string]any{"content": "ok"}}}, }) })) defer srv.Close() c := New(srv.URL, "", "test-model", 10*time.Second) if _, err := c.Complete(context.Background(), "sys", "user"); err != nil { t.Fatalf("Complete: %v", err) } if body.MaxTokens <= 0 { t.Errorf("max_tokens = %d, want > 0 (thinking models return empty content without it)", body.MaxTokens) } } // TestClient_WithMaxTokens overrides the completion budget — the summarizer caps // it small so prompt + max_tokens fits a small-context model's window (8k). func TestClient_WithMaxTokens(t *testing.T) { var body chatRequest srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _ = json.NewDecoder(r.Body).Decode(&body) _ = json.NewEncoder(w).Encode(map[string]any{ "choices": []map[string]any{{"message": map[string]any{"content": "ok"}}}, }) })) defer srv.Close() c := New(srv.URL, "", "test-model", 10*time.Second, WithMaxTokens(1500)) if _, err := c.Complete(context.Background(), "sys", "user"); err != nil { t.Fatalf("Complete: %v", err) } if body.MaxTokens != 1500 { t.Errorf("max_tokens = %d, want 1500", body.MaxTokens) } } // TestClient_UsageHookRecordsTokens: the usage hook fires with the model and the // prompt/completion token counts parsed from the response usage block. func TestClient_UsageHookRecordsTokens(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _ = json.NewEncoder(w).Encode(map[string]any{ "choices": []map[string]any{{"message": map[string]any{"content": "ok"}}}, "usage": map[string]any{"prompt_tokens": 123, "completion_tokens": 45}, }) })) defer srv.Close() var gotModel string var gotPrompt, gotCompletion int c := New(srv.URL, "", "test-model", 10*time.Second, WithUsageHook(func(model string, p, comp int) { gotModel, gotPrompt, gotCompletion = model, p, comp })) if _, err := c.Complete(context.Background(), "sys", "user"); err != nil { t.Fatalf("Complete: %v", err) } if gotModel != "test-model" || gotPrompt != 123 || gotCompletion != 45 { t.Errorf("usage hook got (%q, %d, %d), want (test-model, 123, 45)", gotModel, gotPrompt, gotCompletion) } } func TestClient_ReturnsErrorOnNon200(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Error(w, "overloaded", http.StatusServiceUnavailable) })) defer srv.Close() c := New(srv.URL, "", "test-model", 10*time.Second) if _, err := c.Complete(context.Background(), "sys", "user"); err == nil { t.Fatal("want error on non-200, got nil") } } func TestClient_SendsAuthHeader(t *testing.T) { var gotAuth string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotAuth = r.Header.Get("Authorization") _ = json.NewEncoder(w).Encode(map[string]any{ "choices": []map[string]any{{"message": map[string]any{"content": "ok"}}}, }) })) defer srv.Close() c := New(srv.URL, "my-key", "test-model", 10*time.Second) if _, err := c.Complete(context.Background(), "sys", "user"); err != nil { t.Fatalf("Complete: %v", err) } if gotAuth != "Bearer my-key" { t.Errorf("Authorization = %q, want %q", gotAuth, "Bearer my-key") } } func TestClient_Retries429(t *testing.T) { calls := 0 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls++ if calls == 1 { w.Header().Set("Retry-After", "0") w.WriteHeader(http.StatusTooManyRequests) return } _ = json.NewEncoder(w).Encode(map[string]any{ "choices": []map[string]any{{"message": map[string]any{"content": "retried"}}}, }) })) defer srv.Close() c := New(srv.URL, "", "test-model", 10*time.Second) got, err := c.Complete(context.Background(), "sys", "user") if err != nil { t.Fatalf("Complete: %v", err) } if got != "retried" { t.Errorf("got %q, want %q", got, "retried") } if calls != 2 { t.Errorf("calls = %d, want 2", calls) } }