216 lines
10 KiB
Go
216 lines
10 KiB
Go
package openai
|
|
|
|
import (
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"iop/packages/go/config"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
func TestGeminiIngressAuthenticatesAndStreamsThroughChatRoute(t *testing.T) {
|
|
fake := &fakeRunService{events: bufferedRunEvents(
|
|
&iop.RunEvent{Type: "delta", Delta: "hello"},
|
|
&iop.RunEvent{Type: "complete", Usage: &iop.Usage{InputTokens: 3, OutputTokens: 2}, Metadata: map[string]string{"finish_reason": "length"}},
|
|
)}
|
|
srv := NewServer(config.EdgeOpenAIConf{BearerToken: "iop-principal", Adapter: "ollama", Target: "provider-model"}, fake, nil)
|
|
body := `{
|
|
"systemInstruction":{"role":"user","parts":[{"text":"be concise"}]},
|
|
"contents":[{"role":"user","parts":[{"text":"say hello"}]}],
|
|
"generationConfig":{"candidateCount":1,"maxOutputTokens":16,"temperature":0.2,"topK":8,"topP":0.9,"thinkingConfig":{"includeThoughts":true,"thinkingBudget":-1}}
|
|
}`
|
|
req := httptest.NewRequest(http.MethodPost, "/gemini/gemini-direct/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", strings.NewReader(body))
|
|
req.Header.Set("X-Goog-Api-Key", "iop-principal")
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK || w.Header().Get("Content-Type") != "text/event-stream" {
|
|
t.Fatalf("status/header: %d %q body=%s", w.Code, w.Header().Get("Content-Type"), w.Body.String())
|
|
}
|
|
if fake.req.ModelGroupKey != "gemini-direct" || fake.req.Target != "provider-model" {
|
|
t.Fatalf("route bypassed Chat admission: %+v", fake.req)
|
|
}
|
|
if !strings.Contains(fake.req.Prompt, "system: be concise") || !strings.Contains(fake.req.Prompt, "user: say hello") {
|
|
t.Fatalf("prompt conversion mismatch: %q", fake.req.Prompt)
|
|
}
|
|
options := fake.req.Input["options"].(map[string]any)
|
|
if options["max_tokens"] != 16 {
|
|
t.Fatalf("generation config mismatch: input=%+v", fake.req.Input)
|
|
}
|
|
if _, exists := options["top_k"]; exists {
|
|
t.Fatalf("deprecated Gemini sampling option reached Chat upstream: input=%+v", fake.req.Input)
|
|
}
|
|
response := w.Body.String()
|
|
for _, want := range []string{`"text":"hello"`, `"finishReason":"MAX_TOKENS"`} {
|
|
if !strings.Contains(response, want) {
|
|
t.Fatalf("missing %s in %s", want, response)
|
|
}
|
|
}
|
|
if strings.Contains(response, "iop-principal") {
|
|
t.Fatal("principal token leaked to response")
|
|
}
|
|
}
|
|
|
|
func TestGeminiIngressRejectsAuthenticationAndShapeBeforeDispatch(t *testing.T) {
|
|
base := `{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`
|
|
for _, tc := range []struct {
|
|
name, path, body, bearer, key string
|
|
status int
|
|
}{
|
|
{"missing key", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", base, "", "", 401},
|
|
{"conflicting auth", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", base, "Bearer other", "iop-principal", 401},
|
|
{"wrong query", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent", base, "", "iop-principal", 400},
|
|
{"duplicate member", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", `{"contents":[],"contents":[]}`, "", "iop-principal", 400},
|
|
{"two candidates", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", `{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"candidateCount":2}}`, "", "iop-principal", 400},
|
|
{"invalid thinking budget", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", `{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":-2}}}`, "", "iop-principal", 400},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
fake := &fakeRunService{}
|
|
srv := NewServer(config.EdgeOpenAIConf{BearerToken: "iop-principal", Adapter: "ollama"}, fake, nil)
|
|
req := httptest.NewRequest(http.MethodPost, tc.path, strings.NewReader(tc.body))
|
|
if tc.bearer != "" {
|
|
req.Header.Set("Authorization", tc.bearer)
|
|
}
|
|
if tc.key != "" {
|
|
req.Header.Set("X-Goog-Api-Key", tc.key)
|
|
}
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
if w.Code != tc.status {
|
|
t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
|
|
}
|
|
var failure geminiErrorResponse
|
|
if json.Unmarshal(w.Body.Bytes(), &failure) != nil || failure.Error.Code != tc.status {
|
|
t.Fatalf("not a Gemini error: %s", w.Body.String())
|
|
}
|
|
if fake.req.ModelGroupKey != "" {
|
|
t.Fatalf("unexpected dispatch: %+v", fake.req)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgePreservesToolsAndThoughtSignature(t *testing.T) {
|
|
body := []byte(`{
|
|
"contents":[
|
|
{"role":"model","parts":[{"text":"thinking","thought":true},{"functionCall":{"name":"lookup","args":{"q":"x"}},"thoughtSignature":"opaque"}]},
|
|
{"role":"user","parts":[{"functionResponse":{"name":"lookup","response":{"value":1}}}]},
|
|
{"role":"user","parts":[{"text":"continue"}]}
|
|
],
|
|
"tools":[{"functionDeclarations":[{"name":"lookup","description":"find","parametersJsonSchema":{"type":"object"}}]}],
|
|
"toolConfig":{"functionCallingConfig":{"mode":"ANY"}}
|
|
}`)
|
|
converted, err := prepareGeminiChatBridge(body, "preset-gemini-hybrid")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
text := string(converted)
|
|
for _, want := range []string{`"model":"preset-gemini-hybrid"`, `"reasoning_content":"thinking"`, `"thought_signature":"opaque"`, `"tool_call_id":"gemini_call_0_1"`, `"tool_choice":"required"`} {
|
|
if !strings.Contains(text, want) {
|
|
t.Fatalf("missing %s in %s", want, text)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgeUsesProviderNativeThinkingEnvelope(t *testing.T) {
|
|
converted, err := prepareGeminiChatBridge([]byte(`{
|
|
"contents":[{"role":"user","parts":[{"text":"hello"}]}],
|
|
"generationConfig":{"temperature":1,"topK":50,"topP":1,"thinkingConfig":{"includeThoughts":true,"thinkingBudget":-1}}
|
|
}`), "gemini-3.6-flash")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var body map[string]any
|
|
if json.Unmarshal(converted, &body) != nil {
|
|
t.Fatal("converted body is invalid")
|
|
}
|
|
for _, forbidden := range []string{"temperature", "top_k", "top_p", "think", "include_reasoning", "thinking_token_budget"} {
|
|
if _, exists := body[forbidden]; exists {
|
|
t.Fatalf("unsupported field %q in converted body: %s", forbidden, converted)
|
|
}
|
|
}
|
|
extra := body["extra_body"].(map[string]any)
|
|
google := extra["google"].(map[string]any)
|
|
thinking := google["thinking_config"].(map[string]any)
|
|
if thinking["include_thoughts"] != true || thinking["thinking_budget"] != float64(-1) {
|
|
t.Fatalf("thinking envelope mismatch: %+v", thinking)
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgeAcceptsOfficialAgyPlannerStructuredOutput(t *testing.T) {
|
|
converted, err := prepareGeminiChatBridge([]byte(`{
|
|
"contents":[{"role":"user","parts":[{"text":"plan"}]}],
|
|
"generationConfig":{
|
|
"responseMimeType":"application/json",
|
|
"responseJsonSchema":{"type":"OBJECT","properties":{"steps":{"type":"ARRAY"}},"required":["steps"]}
|
|
},
|
|
"tools":[{"functionDeclarations":[{
|
|
"name":"read_file","parameters":{"type":"OBJECT","properties":{"path":{"type":"STRING"}}},
|
|
"responseJsonSchema":{"type":"object"}
|
|
}]}]
|
|
}`), "gemini-3.6-flash")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var body map[string]any
|
|
if json.Unmarshal(converted, &body) != nil {
|
|
t.Fatal("converted body is invalid")
|
|
}
|
|
format := body["response_format"].(map[string]any)
|
|
jsonSchema := format["json_schema"].(map[string]any)
|
|
schema := jsonSchema["schema"].(map[string]any)
|
|
if format["type"] != "json_schema" || jsonSchema["name"] != "agy_response" || jsonSchema["strict"] != true || schema["type"] != "object" {
|
|
t.Fatalf("structured output mismatch: %+v", format)
|
|
}
|
|
tools := body["tools"].([]any)
|
|
function := tools[0].(map[string]any)["function"].(map[string]any)
|
|
if function["name"] != "read_file" {
|
|
t.Fatalf("function conversion mismatch: %+v", function)
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBridgeRejectsConflictingStructuredSchemas(t *testing.T) {
|
|
for _, body := range []string{
|
|
`{"contents":[{"role":"user","parts":[{"text":"x"}]}],"generationConfig":{"responseMimeType":"application/json","responseSchema":{"type":"object"},"responseJsonSchema":{"type":"object"}}}`,
|
|
`{"contents":[{"role":"user","parts":[{"text":"x"}]}],"tools":[{"functionDeclarations":[{"name":"f","parameters":{"type":"object"},"parametersJsonSchema":{"type":"object"}}]}]}`,
|
|
} {
|
|
if _, err := prepareGeminiChatBridge([]byte(body), "gemini-3.6-flash"); err == nil {
|
|
t.Fatal("conflicting schema alternatives must fail")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestGeminiStreamBridgeEmitsBoundedToolOnce(t *testing.T) {
|
|
w := httptest.NewRecorder()
|
|
stream := newGeminiBridgeStream(w, "gemini-3.6-flash")
|
|
input := "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"name\":\"lookup\",\"arguments\":\"{\\\"q\\\":\"}}]}}]}\n\n" +
|
|
"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"\\\"x\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n" +
|
|
"data: [DONE]\n\n"
|
|
if err := stream.Feed([]byte(input)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if strings.Count(w.Body.String(), `"functionCall"`) != 1 || !strings.Contains(w.Body.String(), `"finishReason":"STOP"`) {
|
|
t.Fatalf("tool terminal mismatch: %s", w.Body.String())
|
|
}
|
|
usageWriter := httptest.NewRecorder()
|
|
usageStream := newGeminiBridgeStream(usageWriter, "m")
|
|
if err := usageStream.Feed([]byte("data: {\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":2,\"total_tokens\":5,\"prompt_tokens_details\":{\"cached_tokens\":1},\"completion_tokens_details\":{\"reasoning_tokens\":1}}}\n\ndata: [DONE]\n\n")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for _, want := range []string{`"promptTokenCount":3`, `"candidatesTokenCount":2`, `"cachedContentTokenCount":1`, `"thoughtsTokenCount":1`, `"totalTokenCount":5`} {
|
|
if !strings.Contains(usageWriter.Body.String(), want) {
|
|
t.Fatalf("missing usage %s in %s", want, usageWriter.Body.String())
|
|
}
|
|
}
|
|
oversize := newGeminiBridgeStream(httptest.NewRecorder(), "m")
|
|
state := &geminiBridgeToolState{name: "lookup"}
|
|
state.arguments.WriteString(strings.Repeat("x", geminiToolArgumentLimit))
|
|
oversize.tools[0] = state
|
|
chunk := `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"x"}}]}}]}` + "\n\n"
|
|
if err := oversize.Feed([]byte(chunk)); err == nil {
|
|
t.Fatal("oversize tool arguments must fail")
|
|
}
|
|
}
|