mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
Billing: - Passthrough SSE streams now bill the usage the upstream reports. The gateway falls back to a request-size estimate only on a 2xx stream or a truncated JSON body. - A unified stream that fails after it starts is estimated, not billed at zero. - Billing uses the built-in LiteLLM price table and honours modelMeta free: true. - Pricing entries are validated. Cost arithmetic rounds the total and saturates on overflow. - Gemini thinking tokens count as completion tokens. Limits and storage: - A rejected request rolls back its counts on the other scopes, so one user cannot lock out a group. - Redis EXPIRE replies are checked. Keys recreated without a TTL are repaired at once. - A mutating pipeline is not retried after reply bytes arrive. - History, model totals and target usage report a store outage. Routing and translation: - Discovered ids with a provider-like prefix are listed only in owner/id form, so listing and routing agree. - Anthropic and Gemini discovery follows pagination. - The developer role maps to system. - stream_options.include_usage is honoured on translated streams. - tool_choice none keeps tools when history holds tool blocks. - refusal maps to content_filter. Data URI parameters are stripped. Streaming, retry and failover: - SSE accepts CR, LF and CRLF line ends. This closes a CR injection path. - Retry-After is never jittered. A cancelled wait returns the context error. - The watchdog timer stops at end of body. - Oversized bodies return 413. - A failed stream ends with an error event. - Failover restores pre-existing headers. Client cancels do not count as provider successes. Also: "admin" is a reserved provider name. metrics.path must start with "/". Pricing warnings go to the right gateway instance. The admin targets view sends access as null or [] so the two states differ. Upstream organisation headers are stripped. The users-file reload error is logged once per error. Dead code is removed, and test-only helpers move to test files. Claude-Session: https://claude.ai/code/session_01V5jYJNxpDFdJkyyExpADrv
766 lines
27 KiB
Go
766 lines
27 KiB
Go
package traefikllmgateway
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// TestNewGeminiAdapter_KeylessIsAConstructorError proves an empty resolved
|
|
// API key is a constructor error for gemini-type providers, mirroring
|
|
// ruling (c) and anthropic's equivalent test.
|
|
func TestNewGeminiAdapter_KeylessIsAConstructorError(t *testing.T) {
|
|
if _, err := newGeminiAdapter("p1", "https://generativelanguage.googleapis.com", ""); err == nil {
|
|
t.Fatalf("newGeminiAdapter: want error for empty API key, got nil")
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_InjectAuth proves injectAuth sets the x-goog-api-key
|
|
// header unconditionally.
|
|
func TestGeminiAdapter_InjectAuth(t *testing.T) {
|
|
a, err := newGeminiAdapter("p1", "https://generativelanguage.googleapis.com", "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
req := httptest.NewRequest(http.MethodPost, "/", nil)
|
|
a.injectAuth(req)
|
|
|
|
if got := req.Header.Get("x-goog-api-key"); got != "AIza-test" { //nolint:gosec // test-only literal, not a real credential
|
|
t.Errorf("x-goog-api-key = %q, want %q", got, "AIza-test")
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_NonStreaming proves the adapter builds
|
|
// the {base}/v1beta/models/{model}:generateContent URL from req["model"],
|
|
// sends the x-goog-api-key header, translates the OpenAI request to
|
|
// Gemini's wire format, and translates the Gemini response back to an
|
|
// OpenAI-shaped chat.completion body on the wire, with usage extracted
|
|
// from usageMetadata.
|
|
func TestGeminiAdapter_ChatCompletion_NonStreaming(t *testing.T) {
|
|
const geminiResp = `{"candidates":[{"content":{"role":"model","parts":[{"text":"hi there"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":4},"responseId":"resp_01ABC"}`
|
|
|
|
var gotBody []byte
|
|
var gotAPIKey, gotPath, gotMethod string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
gotMethod = r.Method
|
|
gotBody, _ = io.ReadAll(r.Body)
|
|
gotAPIKey = r.Header.Get("x-goog-api-key")
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(geminiResp))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
req := map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
}
|
|
u, err := a.chatCompletion(context.Background(), rec, req)
|
|
if err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
|
|
if gotPath != "/v1beta/models/gemini-2.5-pro:generateContent" {
|
|
t.Errorf("path = %q, want /v1beta/models/gemini-2.5-pro:generateContent", gotPath)
|
|
}
|
|
if gotMethod != http.MethodPost {
|
|
t.Errorf("method = %q, want POST", gotMethod)
|
|
}
|
|
if u.prompt != 10 || u.completion != 4 {
|
|
t.Errorf("usage = %+v, want {prompt:10 completion:4}", u)
|
|
}
|
|
if gotAPIKey != "AIza-test" { //nolint:gosec // test-only literal, not a real credential
|
|
t.Errorf("x-goog-api-key = %q, want %q", gotAPIKey, "AIza-test")
|
|
}
|
|
|
|
sent := decodeJSONBody(t, gotBody)
|
|
if _, hasModel := sent["model"]; hasModel {
|
|
t.Errorf("upstream body carries a \"model\" field = %v, want none (Gemini's model is URL-only)", sent["model"])
|
|
}
|
|
contents, ok := sent["contents"].([]any)
|
|
if !ok || len(contents) != 1 {
|
|
t.Fatalf("upstream body contents = %v, want a single-element array", sent["contents"])
|
|
}
|
|
|
|
if got := rec.Header().Get("Content-Type"); got != "application/json" {
|
|
t.Errorf("Content-Type = %q, want application/json", got)
|
|
}
|
|
var out map[string]any
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
|
t.Fatalf("client body is not valid JSON: %v (%s)", err, rec.Body.String())
|
|
}
|
|
if out["object"] != "chat.completion" {
|
|
t.Errorf("object = %v, want %q", out["object"], "chat.completion")
|
|
}
|
|
if out["id"] != "chatcmpl-resp_01ABC" {
|
|
t.Errorf("id = %v, want %q", out["id"], "chatcmpl-resp_01ABC")
|
|
}
|
|
choices, ok := out["choices"].([]any)
|
|
if !ok || len(choices) != 1 {
|
|
t.Fatalf("choices = %v, want a single-element array", out["choices"])
|
|
}
|
|
choice := choices[0].(map[string]any)
|
|
message := choice["message"].(map[string]any)
|
|
if message["content"] != "hi there" {
|
|
t.Errorf("message.content = %v, want %q", message["content"], "hi there")
|
|
}
|
|
if choice["finish_reason"] != "stop" {
|
|
t.Errorf("finish_reason = %v, want %q", choice["finish_reason"], "stop")
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_ModelsPrefixStripped proves a
|
|
// req["model"] that already carries a "models/" prefix is not doubled into
|
|
// the URL.
|
|
func TestGeminiAdapter_ChatCompletion_ModelsPrefixStripped(t *testing.T) {
|
|
var gotPath string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
_, err = a.chatCompletion(context.Background(), rec, map[string]any{
|
|
"model": "models/gemini-2.5-pro",
|
|
"messages": []any{map[string]any{"role": "user", "content": "hi"}},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
if gotPath != "/v1beta/models/gemini-2.5-pro:generateContent" {
|
|
t.Errorf("path = %q, want /v1beta/models/gemini-2.5-pro:generateContent (no doubled models/ prefix)", gotPath)
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_ModelEscaped proves a model id
|
|
// containing "/" is percent-escaped into the upstream URL rather than
|
|
// smuggling an extra path segment in unescaped — review fix (item 2). Go's
|
|
// net/http server auto-decodes r.URL.Path, which would make an escaped and
|
|
// an unescaped "/" indistinguishable, so this asserts on r.URL.EscapedPath()
|
|
// (and cross-checks RequestURI), the wire form the upstream actually
|
|
// receives.
|
|
func TestGeminiAdapter_ChatCompletion_ModelEscaped(t *testing.T) {
|
|
var gotEscapedPath, gotRequestURI string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotEscapedPath = r.URL.EscapedPath()
|
|
gotRequestURI = r.RequestURI
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
_, err = a.chatCompletion(context.Background(), rec, map[string]any{
|
|
"model": "org/gemini-2.5-pro",
|
|
"messages": []any{map[string]any{"role": "user", "content": "hi"}},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
|
|
const wantEscapedPath = "/v1beta/models/org%2Fgemini-2.5-pro:generateContent"
|
|
if gotEscapedPath != wantEscapedPath {
|
|
t.Errorf("EscapedPath() = %q, want %q (model id's \"/\" must stay percent-escaped, not become an extra path segment)", gotEscapedPath, wantEscapedPath)
|
|
}
|
|
if !strings.Contains(gotRequestURI, "org%2Fgemini-2.5-pro") {
|
|
t.Errorf("RequestURI = %q, want it to contain the escaped model id %q", gotRequestURI, "org%2Fgemini-2.5-pro")
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_Embeddings_ModelEscaped mirrors the chatCompletion
|
|
// escaping test for the embeddings :embedContent URL.
|
|
func TestGeminiAdapter_Embeddings_ModelEscaped(t *testing.T) {
|
|
var gotEscapedPath string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotEscapedPath = r.URL.EscapedPath()
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"embedding":{"values":[0.1]}}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
_, err = a.embeddings(context.Background(), rec, map[string]any{
|
|
"model": "org/embed-004",
|
|
"input": "hi",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("embeddings: %v", err)
|
|
}
|
|
|
|
const wantEscapedPath = "/v1beta/models/org%2Fembed-004:embedContent"
|
|
if gotEscapedPath != wantEscapedPath {
|
|
t.Errorf("EscapedPath() = %q, want %q", gotEscapedPath, wantEscapedPath)
|
|
}
|
|
}
|
|
|
|
// geminiStreamFrames are the raw SSE "data:" lines a fake streaming
|
|
// upstream emits for the tests below: a role/text chunk, a text-only
|
|
// continuation, then a finish chunk carrying usageMetadata. Gemini's SSE
|
|
// stream has no "event:" field to set, unlike Anthropic's.
|
|
var geminiStreamFrames = []string{
|
|
`{"candidates":[{"content":{"role":"model","parts":[{"text":"Hi!"}]}}],"responseId":"resp_STREAM"}`,
|
|
`{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":6,"candidatesTokenCount":3}}`,
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_Streaming proves the adapter requests
|
|
// the :streamGenerateContent?alt=sse URL, forwards a streaming Gemini
|
|
// response as OpenAI-shaped SSE chunks on the wire, flushing per event, and
|
|
// terminates with "[DONE]" once the upstream stream ends — mirroring
|
|
// ruling (b): there is no distinct terminal event of Gemini's own to key
|
|
// off of.
|
|
func TestGeminiAdapter_ChatCompletion_Streaming(t *testing.T) {
|
|
var gotPath, gotQuery, gotAccept string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
gotQuery = r.URL.RawQuery
|
|
gotAccept = r.Header.Get("Accept")
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
f, ok := w.(http.Flusher)
|
|
if !ok {
|
|
t.Errorf("httptest ResponseWriter does not implement http.Flusher")
|
|
return
|
|
}
|
|
for _, data := range geminiStreamFrames {
|
|
fmt.Fprintf(w, "data: %s\n\n", data)
|
|
f.Flush()
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := newRecordingResponseWriter()
|
|
|
|
req := map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"stream": true,
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
}
|
|
u, err := a.chatCompletion(context.Background(), rec, req)
|
|
if err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
|
|
if gotPath != "/v1beta/models/gemini-2.5-pro:streamGenerateContent" {
|
|
t.Errorf("path = %q, want /v1beta/models/gemini-2.5-pro:streamGenerateContent", gotPath)
|
|
}
|
|
if gotQuery != "alt=sse" {
|
|
t.Errorf("query = %q, want %q", gotQuery, "alt=sse")
|
|
}
|
|
if gotAccept != "text/event-stream" {
|
|
t.Errorf("Accept = %q, want text/event-stream", gotAccept)
|
|
}
|
|
if u.prompt != 6 || u.completion != 3 {
|
|
t.Errorf("usage = %+v, want {prompt:6 completion:3}", u)
|
|
}
|
|
|
|
body := rec.body.String()
|
|
if got := rec.header.Get("Content-Type"); got != "text/event-stream" {
|
|
t.Errorf("Content-Type = %q, want text/event-stream", got)
|
|
}
|
|
|
|
wantSubstrings := []string{
|
|
`"delta":{"role":"assistant"}`,
|
|
`"delta":{"content":"Hi!"}`,
|
|
`"finish_reason":"stop"`,
|
|
"data: [DONE]\n\n",
|
|
}
|
|
for _, want := range wantSubstrings {
|
|
if !strings.Contains(body, want) {
|
|
t.Errorf("streamed body missing %q; got:\n%s", want, body)
|
|
}
|
|
}
|
|
// 3 chunks (role, content delta, finish) plus one more flush for the
|
|
// terminal [DONE]: 4 flushes, proving the adapter flushes per event
|
|
// rather than buffering the whole stream.
|
|
const wantFlushes = 4
|
|
if rec.flushed < wantFlushes {
|
|
t.Errorf("flushed = %d, want >= %d", rec.flushed, wantFlushes)
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_Streaming_IncludeUsage_EmitsFinalUsageChunk
|
|
// is the M1 regression test: a client that sets
|
|
// stream_options.include_usage on a request routed to a gemini-type
|
|
// provider must still get a final usage-only chunk on the wire, the same
|
|
// way an OpenAI-type provider's own forwardStream relays it from a real
|
|
// upstream usage chunk (provider_openai.go) — Gemini's own stream has no
|
|
// such chunk, so the adapter must synthesize one from the
|
|
// last-usageMetadata-wins total it already accumulated.
|
|
func TestGeminiAdapter_ChatCompletion_Streaming_IncludeUsage_EmitsFinalUsageChunk(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
f, ok := w.(http.Flusher)
|
|
if !ok {
|
|
t.Fatalf("httptest ResponseWriter does not implement http.Flusher")
|
|
}
|
|
for _, data := range geminiStreamFrames {
|
|
fmt.Fprintf(w, "data: %s\n\n", data)
|
|
f.Flush()
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := newRecordingResponseWriter()
|
|
|
|
req := map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"stream": true,
|
|
"stream_options": map[string]any{"include_usage": true},
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
}
|
|
u, err := a.chatCompletion(context.Background(), rec, req)
|
|
if err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
if u.prompt != 6 || u.completion != 3 {
|
|
t.Fatalf("usage = %+v, want {prompt:6 completion:3}", u)
|
|
}
|
|
|
|
chunks := sseDataChunks(t, rec.body.String())
|
|
if len(chunks) < 2 {
|
|
t.Fatalf("got %d SSE frames, want at least a usage chunk followed by [DONE]", len(chunks))
|
|
}
|
|
if chunks[len(chunks)-1] != "[DONE]" {
|
|
t.Fatalf("last SSE frame = %q, want [DONE]", chunks[len(chunks)-1])
|
|
}
|
|
usageFrame := chunks[len(chunks)-2]
|
|
var got struct {
|
|
Usage *struct {
|
|
PromptTokens int64 `json:"prompt_tokens"`
|
|
CompletionTokens int64 `json:"completion_tokens"`
|
|
TotalTokens int64 `json:"total_tokens"`
|
|
} `json:"usage"`
|
|
Choices []any `json:"choices"`
|
|
}
|
|
if err := json.Unmarshal([]byte(usageFrame), &got); err != nil {
|
|
t.Fatalf("decode usage chunk %q: %v", usageFrame, err)
|
|
}
|
|
if len(got.Choices) != 0 {
|
|
t.Errorf("usage chunk choices = %v, want empty", got.Choices)
|
|
}
|
|
if got.Usage == nil {
|
|
t.Fatalf("usage chunk has no usage object: %q", usageFrame)
|
|
}
|
|
if got.Usage.PromptTokens != 6 || got.Usage.CompletionTokens != 3 || got.Usage.TotalTokens != 9 {
|
|
t.Errorf("usage chunk = %+v, want {prompt_tokens:6 completion_tokens:3 total_tokens:9}", *got.Usage)
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_Streaming_NoIncludeUsage_NoFinalUsageChunk
|
|
// proves the M1 fix is opt-in: a client that never sets
|
|
// stream_options.include_usage gets the stream unchanged from before the
|
|
// fix — no extra usage-only chunk on the wire.
|
|
func TestGeminiAdapter_ChatCompletion_Streaming_NoIncludeUsage_NoFinalUsageChunk(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
f, ok := w.(http.Flusher)
|
|
if !ok {
|
|
t.Fatalf("httptest ResponseWriter does not implement http.Flusher")
|
|
}
|
|
for _, data := range geminiStreamFrames {
|
|
fmt.Fprintf(w, "data: %s\n\n", data)
|
|
f.Flush()
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := newRecordingResponseWriter()
|
|
|
|
req := map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"stream": true,
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
}
|
|
if _, err := a.chatCompletion(context.Background(), rec, req); err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
|
|
for _, chunk := range sseDataChunks(t, rec.body.String()) {
|
|
if strings.Contains(chunk, `"usage"`) {
|
|
t.Errorf("SSE frame %q carries a usage field though the client never set stream_options.include_usage", chunk)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_Streaming_CaseInsensitiveContentType
|
|
// proves the streaming-vs-JSON-fallback gate matches an upstream's
|
|
// Content-Type case-insensitively, per lesson (g) — tested with
|
|
// "Text/Event-Stream; charset=utf-8".
|
|
func TestGeminiAdapter_ChatCompletion_Streaming_CaseInsensitiveContentType(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "Text/Event-Stream; charset=utf-8")
|
|
w.WriteHeader(http.StatusOK)
|
|
f, ok := w.(http.Flusher)
|
|
if !ok {
|
|
t.Errorf("httptest ResponseWriter does not implement http.Flusher")
|
|
return
|
|
}
|
|
for _, data := range geminiStreamFrames {
|
|
fmt.Fprintf(w, "data: %s\n\n", data)
|
|
f.Flush()
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := newRecordingResponseWriter()
|
|
|
|
req := map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"stream": true,
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
}
|
|
u, err := a.chatCompletion(context.Background(), rec, req)
|
|
if err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
if u.prompt != 6 || u.completion != 3 {
|
|
t.Errorf("usage = %+v, want {prompt:6 completion:3} (proves the streaming path, not JSON fallback, was taken)", u)
|
|
}
|
|
if got := rec.header.Get("Content-Type"); got != "text/event-stream" {
|
|
t.Errorf("Content-Type on client response = %q, want %q (sseWriter's own header; JSON fallback would copy the upstream's mixed-case value verbatim)", got, "text/event-stream")
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_Upstream429 proves a non-2xx upstream
|
|
// response returns a *providerHTTPError with nothing written to the
|
|
// client.
|
|
func TestGeminiAdapter_ChatCompletion_Upstream429(t *testing.T) {
|
|
const errBody = `{"error":{"code":429,"message":"rate limited","status":"RESOURCE_EXHAUSTED"}}`
|
|
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
_, _ = w.Write([]byte(errBody))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
u, err := a.chatCompletion(context.Background(), rec, map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"messages": []any{map[string]any{"role": "user", "content": "hi"}},
|
|
})
|
|
if err == nil {
|
|
t.Fatalf("chatCompletion: want error, got nil (usage=%+v)", u)
|
|
}
|
|
var httpErr *providerHTTPError
|
|
if !errors.As(err, &httpErr) {
|
|
t.Fatalf("err = %v (%T), want *providerHTTPError", err, err)
|
|
}
|
|
if httpErr.status != http.StatusTooManyRequests {
|
|
t.Errorf("status = %d, want %d", httpErr.status, http.StatusTooManyRequests)
|
|
}
|
|
if rec.Body.Len() != 0 {
|
|
t.Errorf("client response body = %q, want empty", rec.Body.String())
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_TranslateErrorSkipsUpstream proves a
|
|
// request field geminiRequestFromOpenAI rejects (n>1, here) never reaches
|
|
// the upstream at all.
|
|
func TestGeminiAdapter_ChatCompletion_TranslateErrorSkipsUpstream(t *testing.T) {
|
|
called := false
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
called = true
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
_, err = a.chatCompletion(context.Background(), rec, map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"messages": []any{map[string]any{"role": "user", "content": "hi"}},
|
|
"n": float64(2),
|
|
})
|
|
var terr *translateError
|
|
if !errors.As(err, &terr) {
|
|
t.Fatalf("err = %v (%T), want *translateError", err, err)
|
|
}
|
|
if called {
|
|
t.Errorf("upstream was called, want the translate error to short-circuit before any request")
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_Embeddings_Single proves embeddings builds the
|
|
// :embedContent URL and translates the response to an OpenAI embeddings
|
|
// list body for a single-string "input".
|
|
func TestGeminiAdapter_Embeddings_Single(t *testing.T) {
|
|
var gotPath string
|
|
var gotBody []byte
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
gotBody, _ = io.ReadAll(r.Body)
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"embedding":{"values":[0.1,0.2,0.3]}}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
u, err := a.embeddings(context.Background(), rec, map[string]any{
|
|
"model": "text-embedding-004",
|
|
"input": "hello world",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("embeddings: %v", err)
|
|
}
|
|
if gotPath != "/v1beta/models/text-embedding-004:embedContent" {
|
|
t.Errorf("path = %q, want /v1beta/models/text-embedding-004:embedContent", gotPath)
|
|
}
|
|
sent := decodeJSONBody(t, gotBody)
|
|
if _, hasRequests := sent["requests"]; hasRequests {
|
|
t.Errorf("single-input body carries \"requests\", want the :embedContent shape")
|
|
}
|
|
if !u.estimated {
|
|
t.Errorf("usage.estimated = false, want true (Gemini reports no embeddings usage)")
|
|
}
|
|
|
|
var out map[string]any
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
|
t.Fatalf("client body is not valid JSON: %v (%s)", err, rec.Body.String())
|
|
}
|
|
if out["object"] != "list" {
|
|
t.Errorf("object = %v, want %q", out["object"], "list")
|
|
}
|
|
data, ok := out["data"].([]any)
|
|
if !ok || len(data) != 1 {
|
|
t.Fatalf("data = %v, want a single-element array", out["data"])
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_Embeddings_Batch proves an array "input" selects
|
|
// Gemini's :batchEmbedContents endpoint and translates its response to an
|
|
// OpenAI embeddings list body with one entry per input string.
|
|
func TestGeminiAdapter_Embeddings_Batch(t *testing.T) {
|
|
var gotPath string
|
|
var gotBody []byte
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
gotBody, _ = io.ReadAll(r.Body)
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"embeddings":[{"values":[0.1,0.2]},{"values":[0.3,0.4]}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
_, err = a.embeddings(context.Background(), rec, map[string]any{
|
|
"model": "text-embedding-004",
|
|
"input": []any{"hello", "world!!"},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("embeddings: %v", err)
|
|
}
|
|
if gotPath != "/v1beta/models/text-embedding-004:batchEmbedContents" {
|
|
t.Errorf("path = %q, want /v1beta/models/text-embedding-004:batchEmbedContents", gotPath)
|
|
}
|
|
sent := decodeJSONBody(t, gotBody)
|
|
requests, ok := sent["requests"].([]any)
|
|
if !ok || len(requests) != 2 {
|
|
t.Fatalf("requests = %v, want a 2-element array", sent["requests"])
|
|
}
|
|
first := requests[0].(map[string]any)
|
|
if first["model"] != "models/text-embedding-004" {
|
|
t.Errorf("requests[0].model = %v, want %q", first["model"], "models/text-embedding-004")
|
|
}
|
|
|
|
var out map[string]any
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
|
t.Fatalf("client body is not valid JSON: %v (%s)", err, rec.Body.String())
|
|
}
|
|
data, ok := out["data"].([]any)
|
|
if !ok || len(data) != 2 {
|
|
t.Fatalf("data = %v, want a 2-element array", out["data"])
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ListModels proves listModels parses Gemini's
|
|
// "models": [{"name": "models/..."}] Models API response into a flat
|
|
// slice of bare ids, stripping the "models/" prefix.
|
|
func TestGeminiAdapter_ListModels(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/v1beta/models" {
|
|
t.Errorf("path = %q, want /v1beta/models", r.URL.Path)
|
|
return
|
|
}
|
|
if r.Method != http.MethodGet {
|
|
t.Errorf("method = %q, want GET", r.Method)
|
|
return
|
|
}
|
|
if got := r.Header.Get("x-goog-api-key"); got != "AIza-test" {
|
|
t.Errorf("x-goog-api-key = %q, want %q", got, "AIza-test")
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-flash"}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
|
|
got, err := a.listModels(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("listModels: %v", err)
|
|
}
|
|
want := []string{"gemini-2.5-pro", "gemini-2.5-flash"}
|
|
if len(got) != len(want) || got[0] != want[0] || got[1] != want[1] {
|
|
t.Errorf("got %#v, want %#v", got, want)
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ListModels_FollowsPagination is the M7 regression
|
|
// test: a non-empty nextPageToken must make listModels fetch a second
|
|
// page via ?pageToken=<token> and return the union of both pages.
|
|
func TestGeminiAdapter_ListModels_FollowsPagination(t *testing.T) {
|
|
var gotPageTokens []string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPageTokens = append(gotPageTokens, r.URL.Query().Get("pageToken"))
|
|
w.Header().Set("Content-Type", "application/json")
|
|
if r.URL.Query().Get("pageToken") == "" {
|
|
_, _ = w.Write([]byte(`{"models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-flash"}],"nextPageToken":"page2"}`))
|
|
return
|
|
}
|
|
_, _ = w.Write([]byte(`{"models":[{"name":"models/gemini-2.0-flash"}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
|
|
got, err := a.listModels(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("listModels: %v", err)
|
|
}
|
|
want := []string{"gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.0-flash"}
|
|
if len(got) != len(want) {
|
|
t.Fatalf("got %#v, want %#v", got, want)
|
|
}
|
|
for i := range want {
|
|
if got[i] != want[i] {
|
|
t.Errorf("got[%d] = %q, want %q", i, got[i], want[i])
|
|
}
|
|
}
|
|
if len(gotPageTokens) != 2 || gotPageTokens[0] != "" || gotPageTokens[1] != "page2" {
|
|
t.Errorf("pageToken sequence = %#v, want [\"\", \"page2\"] (first page unqualified, second page's pageToken = first page's nextPageToken)", gotPageTokens)
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ListModels_PageCapStopsRunawayPagination proves the
|
|
// M7 fix's other half: an upstream that always returns a nextPageToken
|
|
// does not page forever — listModels stops after
|
|
// geminiListModelsPageCap requests and returns successfully with
|
|
// whatever it accumulated.
|
|
func TestGeminiAdapter_ListModels_PageCapStopsRunawayPagination(t *testing.T) {
|
|
var requests int
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
requests++
|
|
w.Header().Set("Content-Type", "application/json")
|
|
fmt.Fprintf(w, `{"models":[{"name":"models/model-%d"}],"nextPageToken":"tok-%d"}`, requests, requests)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
|
|
got, err := a.listModels(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("listModels: %v", err)
|
|
}
|
|
if requests != geminiListModelsPageCap {
|
|
t.Errorf("upstream received %d requests, want exactly %d (the page cap)", requests, geminiListModelsPageCap)
|
|
}
|
|
if len(got) != geminiListModelsPageCap {
|
|
t.Errorf("got %d ids, want %d (one per page up to the cap)", len(got), geminiListModelsPageCap)
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_HTTPClient proves httpClient returns the adapter's own
|
|
// *http.Client — the one newGeminiAdapter constructed via
|
|
// newAdapterHTTPClient — not a fresh or nil client.
|
|
func TestGeminiAdapter_HTTPClient(t *testing.T) {
|
|
a, err := newGeminiAdapter("p1", "https://generativelanguage.googleapis.com", "AIza-test")
|
|
require.NoError(t, err)
|
|
require.NotNil(t, a.httpClient())
|
|
assert.Same(t, a.client, a.httpClient())
|
|
}
|