mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-14 05:18:54 +00:00
fix(gemini,anthropic): base64 validation, escaped model paths, prompt-block mapping, unified unsupported-param semantics
This commit is contained in:
+17
-7
@@ -6,6 +6,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
@@ -91,15 +92,21 @@ func (a *geminiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite
|
||||
}
|
||||
|
||||
hdr := a.requestHeaders(true)
|
||||
url := a.baseURL + geminiAPIPrefix + model + ":generateContent"
|
||||
// escapedModel guards against a model id smuggling extra path segments
|
||||
// or query syntax ("/", "?", "..") into the upstream URL — url.PathEscape
|
||||
// leaves ordinary ids (letters, digits, ":", ".", "-") untouched, per
|
||||
// RFC 3986's pchar grammar, so a normal model id is byte-identical after
|
||||
// escaping.
|
||||
escapedModel := url.PathEscape(model)
|
||||
endpoint := a.baseURL + geminiAPIPrefix + escapedModel + ":generateContent"
|
||||
if streaming {
|
||||
// Ask the upstream for an SSE response. The Content-Type check
|
||||
// below still decides which forwarding path runs.
|
||||
hdr.Set("Accept", "text/event-stream")
|
||||
url = a.baseURL + geminiAPIPrefix + model + ":streamGenerateContent?alt=sse"
|
||||
endpoint = a.baseURL + geminiAPIPrefix + escapedModel + ":streamGenerateContent?alt=sse"
|
||||
}
|
||||
|
||||
resp, err := upstreamJSON(ctx, a.client, http.MethodPost, url, hdr, body)
|
||||
resp, err := upstreamJSON(ctx, a.client, http.MethodPost, endpoint, hdr, body)
|
||||
if err != nil {
|
||||
return usage{}, err
|
||||
}
|
||||
@@ -189,16 +196,19 @@ func (a *geminiAdapter) embeddings(ctx context.Context, w http.ResponseWriter, r
|
||||
}
|
||||
|
||||
var body map[string]any
|
||||
var url string
|
||||
var endpoint string
|
||||
escapedModel := url.PathEscape(model)
|
||||
if single {
|
||||
body = geminiEmbedContentRequest(texts[0])
|
||||
url = a.baseURL + geminiAPIPrefix + model + ":embedContent"
|
||||
endpoint = a.baseURL + geminiAPIPrefix + escapedModel + ":embedContent"
|
||||
} else {
|
||||
// geminiBatchEmbedContentsRequest still takes the unescaped bare
|
||||
// model — it writes "models/{model}" into the JSON body, not a URL.
|
||||
body = geminiBatchEmbedContentsRequest(model, texts)
|
||||
url = a.baseURL + geminiAPIPrefix + model + ":batchEmbedContents"
|
||||
endpoint = a.baseURL + geminiAPIPrefix + escapedModel + ":batchEmbedContents"
|
||||
}
|
||||
|
||||
resp, err := upstreamJSON(ctx, a.client, http.MethodPost, url, a.requestHeaders(true), body)
|
||||
resp, err := upstreamJSON(ctx, a.client, http.MethodPost, endpoint, a.requestHeaders(true), body)
|
||||
if err != nil {
|
||||
return usage{}, err
|
||||
}
|
||||
|
||||
@@ -153,6 +153,79 @@ func TestGeminiAdapter_ChatCompletion_ModelsPrefixStripped(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"model": "claude-opus-5",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"logit_bias": {}
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"model": "claude-opus-5",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 4096
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"model": "claude-opus-5",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"logprobs": false
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"model": "claude-opus-5",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 4096
|
||||
}
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"model": "gemini-2.5-pro",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": "data:text/plain,Hello"}}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
{"errContains": "must be base64-encoded", "notSupported": false}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"model": "gemini-2.5-pro",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"logit_bias": {}
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"contents": [
|
||||
{"role": "user", "parts": [{"text": "hi"}]}
|
||||
]
|
||||
}
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"model": "gemini-2.5-pro",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"logprobs": false
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"contents": [
|
||||
{"role": "user", "parts": [{"text": "hi"}]}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"model": "gemini-2.5-pro",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"logprobs": true
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
{"errContains": "logprobs is not supported for gemini models", "notSupported": false}
|
||||
Vendored
+9
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"role": "model", "parts": [{"text": "ok"}]},
|
||||
"finishReason": "STOP"
|
||||
}
|
||||
],
|
||||
"usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 1}
|
||||
}
|
||||
Vendored
+14
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"id": "chatcmpl-gemini-1734000000",
|
||||
"object": "chat.completion",
|
||||
"created": 1734000000,
|
||||
"model": "gw-gemini-model",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "ok"},
|
||||
"finish_reason": "stop"
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}
|
||||
}
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"promptFeedback": {"blockReason": "SAFETY"},
|
||||
"usageMetadata": {"promptTokenCount": 15, "candidatesTokenCount": 0},
|
||||
"responseId": "resp_BLOCKED1"
|
||||
}
|
||||
+14
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"id": "chatcmpl-resp_BLOCKED1",
|
||||
"object": "chat.completion",
|
||||
"created": 1734000000,
|
||||
"model": "gw-gemini-model",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": null},
|
||||
"finish_reason": "content_filter"
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 15, "completion_tokens": 0, "total_tokens": 15}
|
||||
}
|
||||
+2
@@ -0,0 +1,2 @@
|
||||
data: {"promptFeedback":{"blockReason":"SAFETY"},"usageMetadata":{"promptTokenCount":9,"candidatesTokenCount":0},"responseId":"resp_BLOCKSTREAM"}
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
{"prompt": 9, "completion": 0}
|
||||
@@ -0,0 +1,4 @@
|
||||
[
|
||||
{"id":"chatcmpl-resp_BLOCKSTREAM","object":"chat.completion.chunk","created":1734000000,"model":"gw-gemini-model","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]},
|
||||
{"id":"chatcmpl-resp_BLOCKSTREAM","object":"chat.completion.chunk","created":1734000000,"model":"gw-gemini-model","choices":[{"index":0,"delta":{},"finish_reason":"content_filter"}]}
|
||||
]
|
||||
+38
-7
@@ -43,6 +43,36 @@ func unsupportedFieldError(field string) *translateError {
|
||||
return &translateError{msg: fmt.Sprintf("%s is not supported for anthropic models", field)}
|
||||
}
|
||||
|
||||
// isTruthy reports whether v — an OpenAI request field's already-decoded
|
||||
// JSON value — is one this translator's truthy-only unsupported-field
|
||||
// checks (logprobs, logit_bias) should treat as "the client actually asked
|
||||
// for this", rather than merely "the field key was present in the request
|
||||
// body". Controller ruling: reject only when truthy/non-empty across both
|
||||
// anthropic and gemini translators — a present-but-falsy value (false, 0,
|
||||
// "", an empty {} or [], or JSON null) must not 400 a request that would
|
||||
// otherwise translate cleanly. Shared by translate_anthropic.go and
|
||||
// translate_gemini.go so both providers apply the identical rule.
|
||||
func isTruthy(v any) bool {
|
||||
switch t := v.(type) {
|
||||
case nil:
|
||||
return false
|
||||
case bool:
|
||||
return t
|
||||
case float64:
|
||||
return t != 0
|
||||
case string:
|
||||
return t != ""
|
||||
case map[string]any:
|
||||
return len(t) > 0
|
||||
case []any:
|
||||
return len(t) > 0
|
||||
default:
|
||||
// An unrecognized-but-non-nil type: conservatively truthy, since
|
||||
// this translator has no basis to call it empty.
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// toFloat64 extracts a numeric value from v, which may be float64 (the
|
||||
// normal shape after json.Unmarshal into map[string]any) or a Go-native
|
||||
// int/int64/float32 (a test, or a caller that builds req by hand).
|
||||
@@ -310,21 +340,22 @@ func anthropicToolsFromOpenAI(toolsRaw []any) []any {
|
||||
// anthropicRequestFromOpenAI maps an OpenAI chat-completion request body
|
||||
// to an Anthropic Messages API request body, per this file's request
|
||||
// mapping table. It returns a *translateError — never wrapped — for a
|
||||
// field Anthropic has no equivalent for (n>1, logit_bias, logprobs) or a
|
||||
// content part this translator cannot map (a non-data-URI image URL, a
|
||||
// malformed tool_calls argument string). Consecutive role:"tool" messages
|
||||
// are merged into one Anthropic user turn carrying all of their
|
||||
// tool_result blocks — see anthropicToolResultBlock's doc comment.
|
||||
// field Anthropic has no equivalent for (n>1; a truthy logit_bias or
|
||||
// logprobs, per isTruthy's doc comment) or a content part this translator
|
||||
// cannot map (a non-data-URI image URL, a malformed tool_calls argument
|
||||
// string). Consecutive role:"tool" messages are merged into one Anthropic
|
||||
// user turn carrying all of their tool_result blocks — see
|
||||
// anthropicToolResultBlock's doc comment.
|
||||
func anthropicRequestFromOpenAI(req map[string]any) (map[string]any, error) {
|
||||
if v, ok := req["n"]; ok {
|
||||
if n, ok2 := toFloat64(v); ok2 && n > 1 {
|
||||
return nil, unsupportedFieldError("n")
|
||||
}
|
||||
}
|
||||
if _, ok := req["logit_bias"]; ok {
|
||||
if v, ok := req["logit_bias"]; ok && isTruthy(v) {
|
||||
return nil, unsupportedFieldError("logit_bias")
|
||||
}
|
||||
if _, ok := req["logprobs"]; ok {
|
||||
if v, ok := req["logprobs"]; ok && isTruthy(v) {
|
||||
return nil, unsupportedFieldError("logprobs")
|
||||
}
|
||||
|
||||
|
||||
+113
-31
@@ -11,16 +11,6 @@ import (
|
||||
// and embedding-family URL is built under, followed by "{model}:{method}".
|
||||
const geminiAPIPrefix = "/v1beta/models/"
|
||||
|
||||
// isTruthyMap reports whether v decodes to a non-empty JSON object. Per
|
||||
// lesson (i), an OpenAI field this translator rejects when present (
|
||||
// logit_bias) is only rejected when it actually carries content — a
|
||||
// present-but-empty {} or a JSON null must not 400 a request that would
|
||||
// otherwise translate cleanly.
|
||||
func isTruthyMap(v any) bool {
|
||||
m, ok := v.(map[string]any)
|
||||
return ok && len(m) > 0
|
||||
}
|
||||
|
||||
// geminiUnsupportedFieldError returns a *translateError for an OpenAI
|
||||
// request field Gemini's API has no equivalent for. Deliberately not
|
||||
// translate_anthropic.go's unsupportedFieldError: that helper's message
|
||||
@@ -55,14 +45,37 @@ func geminiSystemTextFromContentParts(parts []any) (string, error) {
|
||||
return strings.Join(texts, "\n"), nil
|
||||
}
|
||||
|
||||
// dataURIHasBase64Param reports whether rawParams — the ";"-joined
|
||||
// parameter segment of a data: URI header, after the media type and its
|
||||
// leading ";" have been cut away — contains a bare "base64" parameter.
|
||||
// Splitting on ";" rather than checking a suffix or prefix means a
|
||||
// "base64" token anywhere in the parameter list counts, regardless of
|
||||
// what other parameters (";charset=utf-8", say) surround it.
|
||||
func dataURIHasBase64Param(rawParams string) bool {
|
||||
if rawParams == "" {
|
||||
return false
|
||||
}
|
||||
for _, p := range strings.Split(rawParams, ";") {
|
||||
if p == "base64" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// geminiInlineDataFromDataURI parses an OpenAI image_url.url as a
|
||||
// "data:<media-type-and-params>,<data>" URI and returns Gemini's inlineData
|
||||
// object. Per lesson (d), the media type is everything before the first
|
||||
// ";" after "data:" — any parameters after it (";charset=utf-8;base64", for
|
||||
// instance) are stripped rather than rejected or included. Any URL that is
|
||||
// not a data: URI (an http(s) image URL, which Gemini's generateContent API
|
||||
// does not fetch) returns a *translateError, per the request mapping
|
||||
// table's image row.
|
||||
// instance) are stripped rather than rejected or included. Gemini's
|
||||
// inlineData.data is always base64, so a data: URI whose parameter list
|
||||
// does not include a "base64" token is rejected outright — review fix: a
|
||||
// non-base64 URI like "data:text/plain,Hello" was previously forwarded
|
||||
// with its raw text bytes labeled as base64 data, corrupting the upstream
|
||||
// request instead of failing loudly. Any URL that is not a data: URI (an
|
||||
// http(s) image URL, which Gemini's generateContent API does not fetch)
|
||||
// also returns a *translateError, per the request mapping table's image
|
||||
// row.
|
||||
func geminiInlineDataFromDataURI(uri string) (map[string]any, error) {
|
||||
rest, ok := strings.CutPrefix(uri, "data:")
|
||||
if !ok {
|
||||
@@ -72,7 +85,10 @@ func geminiInlineDataFromDataURI(uri string) (map[string]any, error) {
|
||||
if !ok {
|
||||
return nil, &translateError{msg: "malformed data: URI in image_url"}
|
||||
}
|
||||
mimeType, _, _ := strings.Cut(header, ";")
|
||||
mimeType, rawParams, _ := strings.Cut(header, ";")
|
||||
if !dataURIHasBase64Param(rawParams) {
|
||||
return nil, &translateError{msg: "data: URI in image_url must be base64-encoded"}
|
||||
}
|
||||
return map[string]any{"mimeType": mimeType, "data": data}, nil
|
||||
}
|
||||
|
||||
@@ -253,7 +269,9 @@ func geminiToolConfigFromOpenAI(tc any) map[string]any {
|
||||
// Gemini generateContent/streamGenerateContent request body, per this
|
||||
// file's request mapping table. It returns a *translateError — never
|
||||
// wrapped — for a field Gemini has no equivalent for (n>1, a truthy
|
||||
// logit_bias) or a content part this translator cannot map (a non-data-URI
|
||||
// logit_bias or logprobs — controller ruling: unified with anthropic's
|
||||
// same truthy-only semantics, see isTruthy's doc comment) or a content
|
||||
// part this translator cannot map (a non-data-URI
|
||||
// image URL, a malformed tool_calls argument string, a role:"tool" message
|
||||
// whose tool_call_id was never seen on a preceding assistant message).
|
||||
// Consecutive role:"tool" messages are merged into one Gemini {"role":
|
||||
@@ -267,9 +285,12 @@ func geminiRequestFromOpenAI(req map[string]any) (map[string]any, error) {
|
||||
return nil, geminiUnsupportedFieldError("n")
|
||||
}
|
||||
}
|
||||
if v, ok := req["logit_bias"]; ok && isTruthyMap(v) {
|
||||
if v, ok := req["logit_bias"]; ok && isTruthy(v) {
|
||||
return nil, geminiUnsupportedFieldError("logit_bias")
|
||||
}
|
||||
if v, ok := req["logprobs"]; ok && isTruthy(v) {
|
||||
return nil, geminiUnsupportedFieldError("logprobs")
|
||||
}
|
||||
|
||||
out := map[string]any{}
|
||||
|
||||
@@ -433,9 +454,20 @@ type geminiUsageMetadata struct {
|
||||
// zero fields) — the stream mapping table's "last usageMetadata wins" rule
|
||||
// depends on that distinction.
|
||||
type geminiGenerateContentResponse struct {
|
||||
ResponseID string `json:"responseId"`
|
||||
UsageMetadata *geminiUsageMetadata `json:"usageMetadata"`
|
||||
Candidates []geminiCandidate `json:"candidates"`
|
||||
ResponseID string `json:"responseId"`
|
||||
UsageMetadata *geminiUsageMetadata `json:"usageMetadata"`
|
||||
PromptFeedback geminiPromptFeedback `json:"promptFeedback"`
|
||||
Candidates []geminiCandidate `json:"candidates"`
|
||||
}
|
||||
|
||||
// geminiPromptFeedback is the "promptFeedback" object Gemini sets when the
|
||||
// prompt itself — not any generated candidate — was blocked. BlockReason
|
||||
// non-empty and Candidates empty together mean the model never produced
|
||||
// any output at all: review fix — this case was previously reported to the
|
||||
// client as an ordinary empty "stop" completion instead of a content-filter
|
||||
// rejection.
|
||||
type geminiPromptFeedback struct {
|
||||
BlockReason string `json:"blockReason"`
|
||||
}
|
||||
|
||||
// geminiFinishReason maps a Gemini finishReason to an OpenAI finish_reason,
|
||||
@@ -471,12 +503,33 @@ func geminiToolCallID(ordinal int) string {
|
||||
return "call_" + strconv.Itoa(ordinal)
|
||||
}
|
||||
|
||||
// geminiResponseIDOrFallback returns responseID when non-empty, or else a
|
||||
// stable id derived from created — review fix (folded minor): Gemini does
|
||||
// not always set "responseId", and chatCompletionIDPrefix alone (yielding
|
||||
// a bare "chatcmpl-" with nothing after it) is not a usable per-response
|
||||
// id. created is always available (the adapter fixes it once, before any
|
||||
// upstream bytes arrive), so it is deterministic and collision-free across
|
||||
// requests issued in different seconds.
|
||||
func geminiResponseIDOrFallback(responseID string, created int64) string {
|
||||
if responseID != "" {
|
||||
return responseID
|
||||
}
|
||||
return fmt.Sprintf("gemini-%d", created)
|
||||
}
|
||||
|
||||
// openAIResponseFromGemini maps a non-streaming Gemini
|
||||
// GenerateContentResponse body to an OpenAI chat.completion response, per
|
||||
// this file's response mapping table. Only candidates[0] is read, matching
|
||||
// the table. model and created are supplied by the caller — the gateway-
|
||||
// facing model id and a response timestamp, since Gemini's own response
|
||||
// carries neither.
|
||||
//
|
||||
// A prompt-level safety block — Candidates empty and promptFeedback.
|
||||
// blockReason set — is reported as finish_reason "content_filter" with a
|
||||
// null message content, rather than falling through to the "no candidates,
|
||||
// no finishReason" default of an ordinary "stop" completion (review fix:
|
||||
// that default previously misreported a blocked prompt as a normal empty
|
||||
// success).
|
||||
func openAIResponseFromGemini(body []byte, model string, created int64) (map[string]any, usage, error) {
|
||||
var resp geminiGenerateContentResponse
|
||||
if err := json.Unmarshal(body, &resp); err != nil {
|
||||
@@ -526,17 +579,24 @@ func openAIResponseFromGemini(body []byte, model string, created int64) (map[str
|
||||
message["tool_calls"] = toolCalls
|
||||
}
|
||||
|
||||
finishReason := geminiFinishReason(finishReasonRaw, len(toolCalls) > 0)
|
||||
promptBlocked := len(resp.Candidates) == 0 && resp.PromptFeedback.BlockReason != ""
|
||||
if promptBlocked {
|
||||
finishReason = "content_filter"
|
||||
message["content"] = nil
|
||||
}
|
||||
|
||||
u := usage{prompt: resp.UsageMetadata.safePrompt(), completion: resp.UsageMetadata.safeCompletion()}
|
||||
|
||||
out := map[string]any{
|
||||
"id": chatCompletionIDPrefix + resp.ResponseID,
|
||||
"id": chatCompletionIDPrefix + geminiResponseIDOrFallback(resp.ResponseID, created),
|
||||
"object": "chat.completion",
|
||||
"created": created,
|
||||
"model": model,
|
||||
"choices": []any{map[string]any{
|
||||
"index": 0,
|
||||
"message": message,
|
||||
"finish_reason": geminiFinishReason(finishReasonRaw, len(toolCalls) > 0),
|
||||
"finish_reason": finishReason,
|
||||
}},
|
||||
"usage": map[string]any{
|
||||
"prompt_tokens": u.prompt,
|
||||
@@ -565,11 +625,12 @@ func (m *geminiUsageMetadata) safeCompletion() int64 {
|
||||
}
|
||||
|
||||
// geminiStreamState is one streaming chat completion's translation state:
|
||||
// the OpenAI-shaped chunk envelope fields (id, captured from the first
|
||||
// chunk that carries a non-empty "responseId"; model and created, fixed at
|
||||
// construction), the running usage total, and the per-stream tool-call
|
||||
// ordinal counter lesson (a) requires. A caller constructs one per stream
|
||||
// and calls translate once per upstream sseEvent, in order.
|
||||
// the OpenAI-shaped chunk envelope fields (id — set at construction to a
|
||||
// stable fallback derived from created, per geminiResponseIDOrFallback,
|
||||
// and locked in place after the first translated chunk; model and created,
|
||||
// fixed at construction), the running usage total, and the per-stream
|
||||
// tool-call ordinal counter lesson (a) requires. A caller constructs one
|
||||
// per stream and calls translate once per upstream sseEvent, in order.
|
||||
type geminiStreamState struct {
|
||||
id string
|
||||
model string
|
||||
@@ -593,9 +654,13 @@ type geminiStreamState struct {
|
||||
|
||||
// newGeminiStreamState returns a geminiStreamState for one streaming chat
|
||||
// completion. model is the gateway-facing model id and created the
|
||||
// response timestamp every emitted chunk carries.
|
||||
// response timestamp every emitted chunk carries. id starts at the
|
||||
// created-derived fallback (review fix, folded minor): if the first chunk
|
||||
// never carries a "responseId" of its own, every chunk in the stream still
|
||||
// shares one stable, non-empty id instead of drifting between an empty
|
||||
// string and whatever a later chunk happens to report.
|
||||
func newGeminiStreamState(model string, created int64) *geminiStreamState {
|
||||
return &geminiStreamState{model: model, created: created}
|
||||
return &geminiStreamState{model: model, created: created, id: geminiResponseIDOrFallback("", created)}
|
||||
}
|
||||
|
||||
// usage returns the prompt/completion token counts captured so far, from
|
||||
@@ -658,20 +723,37 @@ func (st *geminiStreamState) translate(ev sseEvent) ([][]byte, error) {
|
||||
if err := json.Unmarshal(ev.data, &resp); err != nil {
|
||||
return nil, fmt.Errorf("%w: decode gemini stream chunk: %w", errUpstream, err)
|
||||
}
|
||||
if resp.ResponseID != "" {
|
||||
|
||||
// isFirstChunk gates both the role-delta chunk below and the id
|
||||
// adoption right after it: the id resolves exactly once, on the first
|
||||
// translate call, per the "first one wins" ruling. A responseId
|
||||
// arriving on any later chunk is ignored — st.id already holds either
|
||||
// that first chunk's real id or the construction-time fallback, and
|
||||
// stays there for the rest of the stream.
|
||||
isFirstChunk := !st.startedRole
|
||||
if isFirstChunk && resp.ResponseID != "" {
|
||||
st.id = resp.ResponseID
|
||||
}
|
||||
|
||||
if resp.UsageMetadata != nil {
|
||||
st.u = usage{prompt: resp.UsageMetadata.PromptTokenCount, completion: resp.UsageMetadata.CandidatesTokenCount}
|
||||
}
|
||||
|
||||
var chunks [][]byte
|
||||
if !st.startedRole {
|
||||
if isFirstChunk {
|
||||
st.startedRole = true
|
||||
chunks = append(chunks, st.chunk(map[string]any{"role": "assistant"}, nil))
|
||||
}
|
||||
|
||||
if len(resp.Candidates) == 0 {
|
||||
// A prompt-level safety block reports no candidates at all, only
|
||||
// promptFeedback.blockReason — review fix: this chunk previously
|
||||
// produced nothing beyond the role delta, so the stream ended with
|
||||
// only "[DONE]" and no finish_reason for the client to act on.
|
||||
if resp.PromptFeedback.BlockReason != "" {
|
||||
reason := "content_filter"
|
||||
chunks = append(chunks, st.chunk(map[string]any{}, &reason))
|
||||
}
|
||||
return chunks, nil
|
||||
}
|
||||
c := resp.Candidates[0]
|
||||
|
||||
@@ -236,6 +236,15 @@ func TestGeminiStreamGolden_TextThenFuncall(t *testing.T) {
|
||||
geminiRunStreamGolden(t, "stream_text_then_funcall")
|
||||
}
|
||||
|
||||
// TestGeminiStreamGolden_PromptBlocked covers a prompt-level safety block:
|
||||
// a single chunk with no "candidates" at all, only promptFeedback.
|
||||
// blockReason set — review fix (item 3). Proves the role chunk still
|
||||
// fires, followed by a finish chunk mapped to "content_filter" rather than
|
||||
// the stream silently ending with no finish signal at all.
|
||||
func TestGeminiStreamGolden_PromptBlocked(t *testing.T) {
|
||||
geminiRunStreamGolden(t, "stream_prompt_blocked")
|
||||
}
|
||||
|
||||
// TestGeminiStreamState_MalformedEventReturnsError proves a malformed
|
||||
// (invalid-JSON) SSE data payload returns an error wrapping errUpstream,
|
||||
// rather than being swallowed as (nil, nil) — lesson (c).
|
||||
@@ -272,6 +281,65 @@ func TestGeminiStreamState_UsageMetadataOnlyOverwritesWhenPresent(t *testing.T)
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeminiStreamState_IDLockedOnFirstChunk proves the id resolves once,
|
||||
// on the first translated chunk, and never changes after — review fix
|
||||
// (item 5, folded minor). The first chunk here carries no responseId, so
|
||||
// every emitted chunk (including one after a later chunk arrives with a
|
||||
// real responseId) must share the same construction-time fallback id.
|
||||
func TestGeminiStreamState_IDLockedOnFirstChunk(t *testing.T) {
|
||||
const created = int64(1734000000)
|
||||
st := newGeminiStreamState(geminiGoldenGatewayModel, created)
|
||||
wantID := "chatcmpl-gemini-1734000000"
|
||||
|
||||
chunks, err := st.translate(sseEvent{data: []byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"hi"}]}}]}`)})
|
||||
if err != nil {
|
||||
t.Fatalf("translate: %v", err)
|
||||
}
|
||||
for _, c := range chunks {
|
||||
assertGeminiChunkID(t, c, wantID)
|
||||
}
|
||||
|
||||
// A later chunk carries a real responseId — it must NOT override the
|
||||
// fallback already locked in on the first chunk.
|
||||
chunks, err = st.translate(sseEvent{data: []byte(`{"responseId":"resp_LATE","candidates":[{"content":{"role":"model","parts":[{"text":" there"}]},"finishReason":"STOP"}]}`)})
|
||||
if err != nil {
|
||||
t.Fatalf("translate: %v", err)
|
||||
}
|
||||
if len(chunks) == 0 {
|
||||
t.Fatalf("translate: want at least one chunk, got none")
|
||||
}
|
||||
for _, c := range chunks {
|
||||
assertGeminiChunkID(t, c, wantID)
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeminiStreamState_IDAdoptsFirstChunkResponseID proves the opposite
|
||||
// side of the same rule: when the very first chunk DOES carry a
|
||||
// responseId, that id — not the construction-time fallback — is what
|
||||
// locks in for the rest of the stream.
|
||||
func TestGeminiStreamState_IDAdoptsFirstChunkResponseID(t *testing.T) {
|
||||
st := newGeminiStreamState(geminiGoldenGatewayModel, geminiGoldenCreated)
|
||||
chunks, err := st.translate(sseEvent{data: []byte(`{"responseId":"resp_FIRST","candidates":[{"content":{"role":"model","parts":[{"text":"hi"}]}}]}`)})
|
||||
if err != nil {
|
||||
t.Fatalf("translate: %v", err)
|
||||
}
|
||||
for _, c := range chunks {
|
||||
assertGeminiChunkID(t, c, "chatcmpl-resp_FIRST")
|
||||
}
|
||||
}
|
||||
|
||||
// assertGeminiChunkID decodes chunk and asserts its "id" field equals want.
|
||||
func assertGeminiChunkID(t *testing.T, chunk []byte, want string) {
|
||||
t.Helper()
|
||||
var v map[string]any
|
||||
if err := json.Unmarshal(chunk, &v); err != nil {
|
||||
t.Fatalf("chunk is not valid JSON: %v (%s)", err, chunk)
|
||||
}
|
||||
if v["id"] != want {
|
||||
t.Errorf("chunk id = %v, want %q (chunk: %s)", v["id"], want, chunk)
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeminiRequestFromOpenAI_MaxCompletionTokensPrecedence proves
|
||||
// max_completion_tokens takes precedence over max_tokens when an OpenAI
|
||||
// request sends both — the newer field wins, mirroring the anthropic
|
||||
|
||||
Reference in New Issue
Block a user