diff --git a/internal/translator/gemini_helper.go b/internal/translator/gemini_helper.go index e126bfe4a4..059869614b 100644 --- a/internal/translator/gemini_helper.go +++ b/internal/translator/gemini_helper.go @@ -429,7 +429,7 @@ func assistantMsgToGeminiParts(msg *openai.ChatCompletionAssistantMessageParam) // } // // ]. -func openAIToolsToGeminiTools(openaiTools []openai.Tool, parametersJSONSchemaAvailable bool) ([]genai.Tool, error) { +func openAIToolsToGeminiTools(openaiTools []openai.Tool) ([]genai.Tool, error) { if len(openaiTools) == 0 { return nil, nil } @@ -448,21 +448,7 @@ func openAIToolsToGeminiTools(openaiTools []openai.Tool, parametersJSONSchemaAva Description: tool.Function.Description, } - if parametersJSONSchemaAvailable { - functionDecl.ParametersJsonSchema = tool.Function.Parameters - } else if tool.Function.Parameters != nil { - paramsMap, ok := tool.Function.Parameters.(map[string]any) - if !ok { - return nil, fmt.Errorf("%w: tool %s parameters must be a JSON object", internalapi.ErrInvalidRequestBody, tool.Function.Name) - } - - if len(paramsMap) > 0 { - var err error - if functionDecl.Parameters, err = jsonSchemaToGemini(paramsMap); err != nil { - return nil, fmt.Errorf("%w: invalid JSON schema for parameters in tool %s: %w", internalapi.ErrInvalidRequestBody, tool.Function.Name, err) - } - } - } + functionDecl.ParametersJsonSchema = tool.Function.Parameters functionDecls = append(functionDecls, functionDecl) } case openai.ToolTypeImageGeneration: @@ -581,11 +567,6 @@ func openAIToolChoiceToGeminiToolConfig(toolChoice *openai.ChatCompletionToolCho // Check and Update these functions when new Gemini model versions are released // and their feature support changes. -// it only works with models after gemini2.5 according to https://ai.google.dev/gemini-api/docs/structured-output#json-schema, separate it as a small function to make it easier to maintain -func responseJSONSchemaAvailable(requestModel internalapi.RequestModel) bool { - return strings.Contains(requestModel, "gemini") && (strings.Contains(requestModel, "2.5") || strings.Contains(requestModel, "3")) -} - // mediaResolutionAvailable checks if the model supports media resolution settings. // Only Gemini 3.0+ models support this feature for controlling image/video quality. func mediaResolutionAvailable(requestModel internalapi.RequestModel) bool { @@ -683,16 +664,7 @@ func openAIReqToGeminiGenerationConfig(openAIReq *openai.ChatCompletionRequest, responseMode = responseModeJSON - if responseJSONSchemaAvailable(requestModel) { - gc.ResponseJsonSchema = schemaMap - } else { - convertedSchema, err := jsonSchemaToGemini(schemaMap) - if err != nil { - return nil, responseMode, fmt.Errorf("%w: invalid JSON schema: %w", internalapi.ErrInvalidRequestBody, err) - } - gc.ResponseSchema = convertedSchema - - } + gc.ResponseJsonSchema = schemaMap } } diff --git a/internal/translator/gemini_helper_test.go b/internal/translator/gemini_helper_test.go index 529efc726a..6b36d73115 100644 --- a/internal/translator/gemini_helper_test.go +++ b/internal/translator/gemini_helper_test.go @@ -1390,6 +1390,33 @@ func TestOpenAIReqToGeminiGenerationConfig(t *testing.T) { expectedResponseMode: responseModeJSON, requestModel: "gemini-2.5-flash", }, + { + name: "json schema with references", + input: &openai.ChatCompletionRequest{ + ResponseFormat: &openai.ChatCompletionResponseFormatUnion{ + OfJSONSchema: &openai.ChatCompletionResponseFormatJSONSchema{ + Type: openai.ChatCompletionResponseFormatTypeJSONSchema, + JSONSchema: openai.ChatCompletionResponseFormatJSONSchemaJSONSchema{ + Schema: json.RawMessage(`{"type":"object","properties":{"result":{"$ref":"#/$defs/result"}},"$defs":{"result":{"type":"string"}}}`), + }, + }, + }, + }, + expectedGenerationConfig: &genai.GenerationConfig{ + ResponseMIMEType: "application/json", + ResponseJsonSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "result": map[string]any{"$ref": "#/$defs/result"}, + }, + "$defs": map[string]any{ + "result": map[string]any{"type": "string"}, + }, + }, + }, + expectedResponseMode: responseModeJSON, + requestModel: "gemini-2.5-flash", + }, { name: "json schema (string)", input: &openai.ChatCompletionRequest{ @@ -1616,59 +1643,28 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { }, "required": []any{"a", "b"}, } - tests := []struct { - name string - openaiTools []openai.Tool - parametersJSONSchemaAvailable bool - expected []genai.Tool - expectedError string - }{ - { - name: "empty tools with parametersJSONSchemaAvailable=false", - openaiTools: nil, - parametersJSONSchemaAvailable: false, - expected: nil, + refParams := map[string]any{ + "type": "object", + "properties": map[string]any{ + "location": map[string]any{"$ref": "#/$defs/location"}, }, - { - name: "empty tools with parametersJSONSchemaAvailable=true", - openaiTools: nil, - parametersJSONSchemaAvailable: true, - expected: nil, + "$defs": map[string]any{ + "location": map[string]any{"type": "string"}, }, + } + tests := []struct { + name string + openaiTools []openai.Tool + expected []genai.Tool + expectedError string + }{ { - name: "single function tool with parameters - parametersJSONSchemaAvailable=false", - openaiTools: []openai.Tool{ - { - Type: openai.ToolTypeFunction, - Function: &openai.FunctionDefinition{ - Name: "add", - Description: "Add two numbers", - Parameters: funcParams, - }, - }, - }, - parametersJSONSchemaAvailable: false, - expected: []genai.Tool{ - { - FunctionDeclarations: []*genai.FunctionDeclaration{ - { - Name: "add", - Description: "Add two numbers", - Parameters: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "a": {Type: "integer"}, - "b": {Type: "integer"}, - }, - Required: []string{"a", "b"}, - }, - }, - }, - }, - }, + name: "empty tools", + openaiTools: nil, + expected: nil, }, { - name: "single function tool with parameters - parametersJSONSchemaAvailable=true", + name: "single function tool with parameters", openaiTools: []openai.Tool{ { Type: openai.ToolTypeFunction, @@ -1679,7 +1675,6 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { }, }, }, - parametersJSONSchemaAvailable: true, expected: []genai.Tool{ { FunctionDeclarations: []*genai.FunctionDeclaration{ @@ -1693,45 +1688,29 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { }, }, { - name: "multiple function tools with nil/empty parameters - parametersJSONSchemaAvailable=false", + name: "function tool preserves JSON Schema references", openaiTools: []openai.Tool{ { Type: openai.ToolTypeFunction, Function: &openai.FunctionDefinition{ - Name: "foo", - Description: "Foo function", - Parameters: map[string]any{}, // empty parameters - }, - }, - { - Type: openai.ToolTypeFunction, - Function: &openai.FunctionDefinition{ - Name: "bar", - Description: "Bar function", - Parameters: nil, // nil parameters + Name: "get_weather", + Parameters: refParams, }, }, }, - parametersJSONSchemaAvailable: false, expected: []genai.Tool{ { FunctionDeclarations: []*genai.FunctionDeclaration{ { - Name: "foo", - Description: "Foo function", - Parameters: nil, - }, - { - Name: "bar", - Description: "Bar function", - Parameters: nil, + Name: "get_weather", + ParametersJsonSchema: refParams, }, }, }, }, }, { - name: "multiple function tools with nil/empty parameters - parametersJSONSchemaAvailable=true", + name: "multiple function tools with nil/empty parameters", openaiTools: []openai.Tool{ { Type: openai.ToolTypeFunction, @@ -1748,7 +1727,6 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { }, }, }, - parametersJSONSchemaAvailable: true, expected: []genai.Tool{ { FunctionDeclarations: []*genai.FunctionDeclaration{ @@ -1767,41 +1745,7 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { }, }, { - name: "tool with invalid parameters schema - parametersJSONSchemaAvailable=false", - openaiTools: []openai.Tool{ - { - Type: openai.ToolTypeFunction, - Function: &openai.FunctionDefinition{ - Name: "bad", - Description: "Bad function", - Parameters: "invalid-json", - }, - }, - }, - parametersJSONSchemaAvailable: false, - expectedError: "tool bad parameters must be a JSON object", - }, - { - name: "tool with unresolvable $ref in parameters - parametersJSONSchemaAvailable=false", - openaiTools: []openai.Tool{ - { - Type: openai.ToolTypeFunction, - Function: &openai.FunctionDefinition{ - Name: "bad_ref", - Parameters: map[string]any{ - "type": "object", - "properties": map[string]any{ - "a": map[string]any{"$ref": "#/nonexistent/path"}, - }, - }, - }, - }, - }, - parametersJSONSchemaAvailable: false, - expectedError: "invalid JSON schema for parameters in tool bad_ref", - }, - { - name: "tool with invalid parameters schema - parametersJSONSchemaAvailable=true", + name: "tool with invalid parameters schema", openaiTools: []openai.Tool{ { Type: openai.ToolTypeFunction, @@ -1812,7 +1756,6 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { }, }, }, - parametersJSONSchemaAvailable: true, expected: []genai.Tool{ { FunctionDeclarations: []*genai.FunctionDeclaration{ @@ -1826,77 +1769,6 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { }, }, }, - { - name: "complex nested schema - parametersJSONSchemaAvailable=false", - openaiTools: []openai.Tool{ - { - Type: openai.ToolTypeFunction, - Function: &openai.FunctionDefinition{ - Name: "complex_tool", - Description: "Complex tool with nested parameters", - Parameters: map[string]any{ - "type": "object", - "properties": map[string]any{ - "user": map[string]any{ - "type": "object", - "properties": map[string]any{ - "name": map[string]any{"type": "string"}, - "age": map[string]any{"type": "integer"}, - }, - "required": []any{"name"}, - }, - "items": map[string]any{ - "type": "array", - "items": map[string]any{ - "type": "object", - "properties": map[string]any{ - "id": map[string]any{"type": "integer"}, - "name": map[string]any{"type": "string"}, - }, - }, - }, - }, - "required": []any{"user"}, - }, - }, - }, - }, - parametersJSONSchemaAvailable: false, - expected: []genai.Tool{ - { - FunctionDeclarations: []*genai.FunctionDeclaration{ - { - Name: "complex_tool", - Description: "Complex tool with nested parameters", - Parameters: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "user": { - Type: "object", - Properties: map[string]*genai.Schema{ - "name": {Type: "string"}, - "age": {Type: "integer"}, - }, - Required: []string{"name"}, - }, - "items": { - Type: "array", - Items: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "id": {Type: "integer"}, - "name": {Type: "string"}, - }, - }, - }, - }, - Required: []string{"user"}, - }, - }, - }, - }, - }, - }, { name: "non-function tool is ignored - both modes", openaiTools: []openai.Tool{ @@ -1904,11 +1776,10 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { Type: "retrieval", }, }, - parametersJSONSchemaAvailable: false, - expectedError: "unsupported tool type: retrieval", + expectedError: "unsupported tool type: retrieval", }, { - name: "mixed valid and invalid tools - parametersJSONSchemaAvailable=false", + name: "mixed valid and invalid tools", openaiTools: []openai.Tool{ { Type: "retrieval", // Should be ignored @@ -1927,8 +1798,7 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { }, }, }, - parametersJSONSchemaAvailable: false, - expectedError: "unsupported tool type: retrieval", + expectedError: "unsupported tool type: retrieval", }, { name: "tool with nil function - should not panic", @@ -1938,8 +1808,7 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { Function: nil, }, }, - parametersJSONSchemaAvailable: false, - expected: nil, + expected: nil, }, { name: "enterprise search tool only", @@ -1948,7 +1817,6 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { Type: openai.ToolTypeEnterpriseWebSearch, }, }, - parametersJSONSchemaAvailable: false, expected: []genai.Tool{ { EnterpriseWebSearch: &genai.EnterpriseWebSearch{}, @@ -1970,7 +1838,6 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { Type: openai.ToolTypeEnterpriseWebSearch, }, }, - parametersJSONSchemaAvailable: false, expected: []genai.Tool{ { EnterpriseWebSearch: &genai.EnterpriseWebSearch{}, @@ -1978,16 +1845,9 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { { FunctionDeclarations: []*genai.FunctionDeclaration{ { - Name: "get_weather", - Description: "Get current weather", - Parameters: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "a": {Type: "integer"}, - "b": {Type: "integer"}, - }, - Required: []string{"a", "b"}, - }, + Name: "get_weather", + Description: "Get current weather", + ParametersJsonSchema: funcParams, }, }, }, @@ -2000,7 +1860,6 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { Type: openai.ToolTypeGoogleSearch, }, }, - parametersJSONSchemaAvailable: false, expected: []genai.Tool{ { GoogleSearch: &genai.GoogleSearch{}, @@ -2017,7 +1876,6 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { }, }, }, - parametersJSONSchemaAvailable: false, expected: []genai.Tool{ { GoogleSearch: &genai.GoogleSearch{ @@ -2037,7 +1895,6 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { }, }, }, - parametersJSONSchemaAvailable: false, expected: []genai.Tool{ { GoogleSearch: &genai.GoogleSearch{ @@ -2060,7 +1917,6 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { }, }, }, - parametersJSONSchemaAvailable: false, expected: []genai.Tool{ { GoogleSearch: &genai.GoogleSearch{ @@ -2090,7 +1946,6 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { }, }, }, - parametersJSONSchemaAvailable: false, expected: []genai.Tool{ { GoogleSearch: &genai.GoogleSearch{ @@ -2100,16 +1955,9 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { { FunctionDeclarations: []*genai.FunctionDeclaration{ { - Name: "search_products", - Description: "Search for products", - Parameters: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "a": {Type: "integer"}, - "b": {Type: "integer"}, - }, - Required: []string{"a", "b"}, - }, + Name: "search_products", + Description: "Search for products", + ParametersJsonSchema: funcParams, }, }, }, @@ -2119,7 +1967,7 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) { for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - result, err := openAIToolsToGeminiTools(tc.openaiTools, tc.parametersJSONSchemaAvailable) + result, err := openAIToolsToGeminiTools(tc.openaiTools) if tc.expectedError != "" { require.ErrorContains(t, err, tc.expectedError) require.ErrorIs(t, err, internalapi.ErrInvalidRequestBody) @@ -3017,344 +2865,6 @@ func TestGeminiCandidatesToOpenAIChoices(t *testing.T) { } } -func TestOpenAIReqToGeminiGenerationConfigWithJsonSchemaToGemini(t *testing.T) { - tests := []struct { - name string - input *openai.ChatCompletionRequest - requestModel internalapi.RequestModel - expectedGenerationConfig *genai.GenerationConfig - expectedResponseMode geminiResponseMode - expectedErrMsg string - }{ - { - name: "json schema for older gemini model (uses jsonSchemaToGemini)", - input: &openai.ChatCompletionRequest{ - ResponseFormat: &openai.ChatCompletionResponseFormatUnion{ - OfJSONSchema: &openai.ChatCompletionResponseFormatJSONSchema{ - Type: openai.ChatCompletionResponseFormatTypeJSONSchema, - JSONSchema: openai.ChatCompletionResponseFormatJSONSchemaJSONSchema{ - Schema: json.RawMessage(`{ - "type": "object", - "properties": { - "name": {"type": "string"}, - "age": {"type": "number"} - }, - "required": ["name", "age"] - }`), - }, - }, - }, - }, - requestModel: "gemini-2.0-flash", // older model that doesn't support native JSON schema - expectedGenerationConfig: &genai.GenerationConfig{ - ResponseMIMEType: "application/json", - ResponseSchema: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "name": {Type: "string"}, - "age": {Type: "number"}, - }, - Required: []string{"name", "age"}, - }, - }, - expectedResponseMode: responseModeJSON, - }, - { - name: "complex json schema with nested objects (uses jsonSchemaToGemini)", - input: &openai.ChatCompletionRequest{ - ResponseFormat: &openai.ChatCompletionResponseFormatUnion{ - OfJSONSchema: &openai.ChatCompletionResponseFormatJSONSchema{ - Type: openai.ChatCompletionResponseFormatTypeJSONSchema, - JSONSchema: openai.ChatCompletionResponseFormatJSONSchemaJSONSchema{ - Schema: json.RawMessage(`{ - "type": "object", - "properties": { - "user": { - "type": "object", - "properties": { - "name": {"type": "string"}, - "contact": { - "type": "object", - "properties": { - "email": {"type": "string"}, - "phone": {"type": "string"} - } - } - } - }, - "items": { - "type": "array", - "items": { - "type": "object", - "properties": { - "id": {"type": "number"}, - "description": {"type": "string"} - } - } - } - }, - "required": ["user", "items"] - }`), - }, - }, - }, - }, - requestModel: "gemini-2.0-flash", // older model - expectedGenerationConfig: &genai.GenerationConfig{ - ResponseMIMEType: "application/json", - ResponseSchema: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "user": { - Type: "object", - Properties: map[string]*genai.Schema{ - "name": {Type: "string"}, - "contact": { - Type: "object", - Properties: map[string]*genai.Schema{ - "email": {Type: "string"}, - "phone": {Type: "string"}, - }, - }, - }, - }, - "items": { - Type: "array", - Items: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "id": {Type: "number"}, - "description": {Type: "string"}, - }, - }, - }, - }, - Required: []string{"user", "items"}, - }, - }, - expectedResponseMode: responseModeJSON, - }, - { - name: "json schema with anyOf (uses jsonSchemaToGemini)", - input: &openai.ChatCompletionRequest{ - ResponseFormat: &openai.ChatCompletionResponseFormatUnion{ - OfJSONSchema: &openai.ChatCompletionResponseFormatJSONSchema{ - Type: openai.ChatCompletionResponseFormatTypeJSONSchema, - JSONSchema: openai.ChatCompletionResponseFormatJSONSchemaJSONSchema{ - Schema: json.RawMessage(`{ - "type": "object", - "properties": { - "value": { - "anyOf": [ - {"type": "string"}, - {"type": "number"} - ] - } - } - }`), - }, - }, - }, - }, - requestModel: "gemini-2.0-flash", - expectedGenerationConfig: &genai.GenerationConfig{ - ResponseMIMEType: "application/json", - ResponseSchema: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "value": { - AnyOf: []*genai.Schema{ - {Type: "string"}, - {Type: "number"}, - }, - }, - }, - }, - }, - expectedResponseMode: responseModeJSON, - }, - { - name: "json schema with nullable type (uses jsonSchemaToGemini)", - input: &openai.ChatCompletionRequest{ - ResponseFormat: &openai.ChatCompletionResponseFormatUnion{ - OfJSONSchema: &openai.ChatCompletionResponseFormatJSONSchema{ - Type: openai.ChatCompletionResponseFormatTypeJSONSchema, - JSONSchema: openai.ChatCompletionResponseFormatJSONSchemaJSONSchema{ - Schema: json.RawMessage(`{ - "type": "object", - "properties": { - "optional_field": { - "anyOf": [ - {"type": "string"}, - {"type": "null"} - ] - } - } - }`), - }, - }, - }, - }, - requestModel: "gemini-2.0-flash", - expectedGenerationConfig: &genai.GenerationConfig{ - ResponseMIMEType: "application/json", - ResponseSchema: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "optional_field": { - AnyOf: []*genai.Schema{ - {Type: "string"}, - }, - Nullable: ptr.To(true), - }, - }, - }, - }, - expectedResponseMode: responseModeJSON, - }, - { - name: "json schema with $ref (uses jsonSchemaToGemini)", - input: &openai.ChatCompletionRequest{ - ResponseFormat: &openai.ChatCompletionResponseFormatUnion{ - OfJSONSchema: &openai.ChatCompletionResponseFormatJSONSchema{ - Type: openai.ChatCompletionResponseFormatTypeJSONSchema, - JSONSchema: openai.ChatCompletionResponseFormatJSONSchemaJSONSchema{ - Schema: json.RawMessage(`{ - "type": "object", - "properties": { - "user": {"$ref": "#/$defs/User"} - }, - "$defs": { - "User": { - "type": "object", - "properties": { - "name": {"type": "string"}, - "age": {"type": "number"} - } - } - } - }`), - }, - }, - }, - }, - requestModel: "gemini-2.0-flash", - expectedGenerationConfig: &genai.GenerationConfig{ - ResponseMIMEType: "application/json", - ResponseSchema: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "user": { - Type: "object", - Properties: map[string]*genai.Schema{ - "name": {Type: "string"}, - "age": {Type: "number"}, - }, - }, - }, - }, - }, - expectedResponseMode: responseModeJSON, - }, - { - name: "invalid json schema causes jsonSchemaToGemini error", - input: &openai.ChatCompletionRequest{ - ResponseFormat: &openai.ChatCompletionResponseFormatUnion{ - OfJSONSchema: &openai.ChatCompletionResponseFormatJSONSchema{ - Type: openai.ChatCompletionResponseFormatTypeJSONSchema, - JSONSchema: openai.ChatCompletionResponseFormatJSONSchemaJSONSchema{ - Schema: json.RawMessage(`{ - "type": "object", - "properties": { - "invalid_ref": {"$ref": "#/nonexistent/path"} - } - }`), - }, - }, - }, - }, - requestModel: "gemini-2.0-flash", - expectedErrMsg: "invalid JSON schema", - }, - { - name: "complex nested schema with arrays and refs (uses jsonSchemaToGemini)", - input: &openai.ChatCompletionRequest{ - ResponseFormat: &openai.ChatCompletionResponseFormatUnion{ - OfJSONSchema: &openai.ChatCompletionResponseFormatJSONSchema{ - Type: openai.ChatCompletionResponseFormatTypeJSONSchema, - JSONSchema: openai.ChatCompletionResponseFormatJSONSchemaJSONSchema{ - Schema: json.RawMessage(`{ - "type": "object", - "properties": { - "steps": { - "type": "array", - "items": {"$ref": "#/$defs/Step"} - }, - "final_answer": {"type": "string"} - }, - "required": ["steps", "final_answer"], - "$defs": { - "Step": { - "type": "object", - "properties": { - "explanation": {"type": "string"}, - "output": {"type": "string"} - }, - "required": ["explanation", "output"] - } - } - }`), - }, - }, - }, - }, - requestModel: "gemini-2.0-flash", - expectedGenerationConfig: &genai.GenerationConfig{ - ResponseMIMEType: "application/json", - ResponseSchema: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "steps": { - Type: "array", - Items: &genai.Schema{ - Type: "object", - Properties: map[string]*genai.Schema{ - "explanation": {Type: "string"}, - "output": {Type: "string"}, - }, - Required: []string{"explanation", "output"}, - }, - }, - "final_answer": {Type: "string"}, - }, - Required: []string{"steps", "final_answer"}, - }, - }, - expectedResponseMode: responseModeJSON, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - got, responseMode, err := openAIReqToGeminiGenerationConfig(tc.input, tc.requestModel) - if tc.expectedErrMsg != "" { - require.ErrorContains(t, err, tc.expectedErrMsg) - require.ErrorIs(t, err, internalapi.ErrInvalidRequestBody) - } else { - require.NoError(t, err) - - if diff := cmp.Diff(tc.expectedGenerationConfig, got, cmpopts.IgnoreUnexported(genai.GenerationConfig{}, genai.Schema{})); diff != "" { - t.Errorf("GenerationConfig mismatch (-want +got):\n%s", diff) - } - - if responseMode != tc.expectedResponseMode { - t.Errorf("geminiResponseMode mismatch: got %v, want %v", responseMode, tc.expectedResponseMode) - } - } - }) - } -} - func TestMapDetailMediaResolution(t *testing.T) { tests := []struct { name string diff --git a/internal/translator/jsonschema_helper.go b/internal/translator/jsonschema_helper.go deleted file mode 100644 index 4df3d7164b..0000000000 --- a/internal/translator/jsonschema_helper.go +++ /dev/null @@ -1,658 +0,0 @@ -// Copyright Envoy AI Gateway Authors -// SPDX-License-Identifier: Apache-2.0 -// The full text of the Apache license is available in the LICENSE file at -// the root of the repo. - -package translator - -import ( - "fmt" - "strings" - - "google.golang.org/genai" - - "github.com/envoyproxy/ai-gateway/internal/json" -) - -// Constants for safety limits -const ( - jsonSchemaMaxRecursionDepth = 100 - // jsonSchemaMaxNodes caps the total number of nodes visited while resolving $refs. The depth - // limit alone does not bound the work: sibling $refs to the same definition are expanded - // independently, so a small schema can expand exponentially. - jsonSchemaMaxNodes = 50000 -) - -// Errors for safety violations -var ( - errJSONSchemaMaxRecursionDepthExceeded = fmt.Errorf("maximum recursion depth exceeded") - errJSONSchemaMaxNodesExceeded = fmt.Errorf("maximum number of schema nodes exceeded") - errInvalidJSONSchema = fmt.Errorf("invalid JSON schema") -) - -// jsonSchemaDeepCopyMapStringAny creates a safe deep copy of a map[string]any. -func jsonSchemaDeepCopyMapStringAny(original map[string]any) (map[string]any, error) { - if original == nil { - return nil, nil - } - return jsonSchemaDeepCopyMapStringAnySafe(original, 0) -} - -// jsonSchemaDeepCopyMapStringAnySafe performs the actual deep copy with recursion depth checks. -func jsonSchemaDeepCopyMapStringAnySafe(original map[string]any, depth int) (map[string]any, error) { - if original == nil { - return nil, nil - } - - if depth >= jsonSchemaMaxRecursionDepth { - return nil, fmt.Errorf("%w: depth %d", errJSONSchemaMaxRecursionDepthExceeded, depth) - } - - copied := make(map[string]any, len(original)) - for key, value := range original { - copiedValue, err := jsonSchemaDeepCopyAnySafe(value, depth+1) - if err != nil { - return nil, err - } - copied[key] = copiedValue - } - return copied, nil -} - -// jsonSchemaDeepCopyAny performs a safe deep copy of any value. -func jsonSchemaDeepCopyAny(value any) (any, error) { - return jsonSchemaDeepCopyAnySafe(value, 0) -} - -// jsonSchemaDeepCopyAnySafe performs the actual deep copy with recursion depth checks. -func jsonSchemaDeepCopyAnySafe(value any, depth int) (any, error) { - if depth >= jsonSchemaMaxRecursionDepth { - return nil, fmt.Errorf("%w: depth %d", errJSONSchemaMaxRecursionDepthExceeded, depth) - } - - switch v := value.(type) { - case map[string]any: - return jsonSchemaDeepCopyMapStringAnySafe(v, depth) - case []any: - copiedSlice := make([]any, len(v)) - for i, elem := range v { - copiedElem, err := jsonSchemaDeepCopyAnySafe(elem, depth+1) - if err != nil { - return nil, err - } - copiedSlice[i] = copiedElem - } - return copiedSlice, nil - default: - // For primitive types (int, bool, string, etc.) and other value types, - // direct assignment performs a copy. - return value, nil - } -} - -// jsonSchemaRetrieveRef safely fetches a deeply-nested reference from a schema map. -// It includes input validation and protection against path traversal attacks. -func jsonSchemaRetrieveRef(path string, schema map[string]any) (any, error) { - if path == "" { - return nil, fmt.Errorf("%w: ref path cannot be empty", errInvalidJSONSchema) - } - - if schema == nil { - return nil, fmt.Errorf("%w: schema cannot be nil", errInvalidJSONSchema) - } - - // Validate path format and prevent path traversal - if !strings.HasPrefix(path, "#/") { - return nil, fmt.Errorf("%w: ref paths must start with '#/', got: %s", errInvalidJSONSchema, path) - } - - // Split and validate path components - components := strings.Split(path, "/") - if len(components) < 2 { - return nil, fmt.Errorf("%w: invalid ref path format: %s", errInvalidJSONSchema, path) - } - - // Remove the "#" prefix - components = components[1:] - - current := schema - for i, component := range components { - if component == "" { - return nil, fmt.Errorf("%w: ref path contains empty component at position %d", errInvalidJSONSchema, i+1) - } - - // Prevent directory traversal attempts - if strings.Contains(component, "..") || strings.Contains(component, "./") { - return nil, fmt.Errorf("%w: ref path contains invalid characters: %s", errInvalidJSONSchema, component) - } - - // Type assertion to ensure `current` is a map - val, exists := current[component] - if !exists { - return nil, fmt.Errorf("%w: reference '%s' not found: component '%s' does not exist", errInvalidJSONSchema, path, component) - } - - // Check if we're at the last component - if i == len(components)-1 { - // Create and return a deep copy to prevent mutation of the original schema - deepCopy, err := jsonSchemaDeepCopyAny(val) - if err != nil { - return nil, fmt.Errorf("failed to create deep copy: %w", err) - } - return deepCopy, nil - } - - // For intermediate components, ensure they are maps - nextMap, ok := val.(map[string]any) - if !ok { - return nil, fmt.Errorf("%w: reference '%s' invalid: intermediate component '%s' is not a map (got %T)", errInvalidJSONSchema, path, component, val) - } - current = nextMap - } - - // This should never be reached due to the loop structure above - return nil, fmt.Errorf("%w: unexpected end of ref path traversal for: %s", errInvalidJSONSchema, path) -} - -// jsonSchemaNodeBudget tracks the remaining number of nodes that can be visited while -// processing a schema. -type jsonSchemaNodeBudget struct { - remaining int -} - -func newJSONSchemaNodeBudget() *jsonSchemaNodeBudget { - return &jsonSchemaNodeBudget{remaining: jsonSchemaMaxNodes} -} - -// consume accounts for one visited node and fails once the budget is exhausted. -func (b *jsonSchemaNodeBudget) consume() error { - if b.remaining <= 0 { - return fmt.Errorf("%w: limit %d", errJSONSchemaMaxNodesExceeded, jsonSchemaMaxNodes) - } - b.remaining-- - return nil -} - -// jsonSchemaDereferenceHelper recursively dereferences JSON schema references. -func jsonSchemaDereferenceHelper( - obj any, - fullSchema map[string]any, - skipKeys []string, - processedRefs map[string]struct{}, - budget *jsonSchemaNodeBudget, - depth int, -) (any, error) { - // Check recursion depth - if depth >= jsonSchemaMaxRecursionDepth { - return nil, fmt.Errorf("%w: depth %d", errJSONSchemaMaxRecursionDepthExceeded, depth) - } - if err := budget.consume(); err != nil { - return nil, err - } - - // Handle dictionaries (maps) - if dict, ok := obj.(map[string]any); ok { - objOut := make(map[string]any) - - for k, v := range dict { - // Check if key should be skipped - shouldSkip := false - for _, skipKey := range skipKeys { - if k == skipKey { - shouldSkip = true - break - } - } - if shouldSkip { - objOut[k] = v - continue - } - - // Handle reference key "$ref" - if k == "$ref" { - refPath, isString := v.(string) - if !isString { - return nil, fmt.Errorf("%w: '$ref' value must be a string, got %T", errInvalidJSONSchema, v) - } - - // Check for circular references - if _, exists := processedRefs[refPath]; exists { - return nil, fmt.Errorf("%w: circular reference detected: %s", errInvalidJSONSchema, refPath) - } - - // Mark as being processed - processedRefs[refPath] = struct{}{} - - ref, err := jsonSchemaRetrieveRef(refPath, fullSchema) - if err != nil { - delete(processedRefs, refPath) // Clean up on error - return nil, fmt.Errorf("failed to retrieve reference %s: %w", refPath, err) - } - - fullRef, err := jsonSchemaDereferenceHelper(ref, fullSchema, skipKeys, processedRefs, budget, depth+1) - if err != nil { - delete(processedRefs, refPath) // Clean up on error - return nil, fmt.Errorf("failed to dereference %s: %w", refPath, err) - } - - // Remove from processed set after successful resolution - delete(processedRefs, refPath) - return fullRef, nil - } - - // Recurse on nested structures - if _, isDict := v.(map[string]any); isDict { - res, err := jsonSchemaDereferenceHelper(v, fullSchema, skipKeys, processedRefs, budget, depth+1) - if err != nil { - return nil, err - } - objOut[k] = res - } else if _, isList := v.([]any); isList { - res, err := jsonSchemaDereferenceHelper(v, fullSchema, skipKeys, processedRefs, budget, depth+1) - if err != nil { - return nil, err - } - objOut[k] = res - } else { - objOut[k] = v - } - } - return objOut, nil - } - - // Handle lists (slices) - if list, ok := obj.([]any); ok { - listOut := make([]any, len(list)) - for i, el := range list { - res, err := jsonSchemaDereferenceHelper(el, fullSchema, skipKeys, processedRefs, budget, depth+1) - if err != nil { - return nil, err - } - listOut[i] = res - } - return listOut, nil - } - - // Return non-dictionary and non-list types as-is - return obj, nil -} - -// jsonSchemaSkipKeys recursively traverses a schema to find keys that should be skipped. -func jsonSchemaSkipKeys( - obj any, - fullSchema map[string]any, - processedRefs map[string]struct{}, - budget *jsonSchemaNodeBudget, - depth int, -) ([]string, error) { - // Check recursion depth - if depth >= jsonSchemaMaxRecursionDepth { - return nil, fmt.Errorf("%w: depth %d", errJSONSchemaMaxRecursionDepthExceeded, depth) - } - if err := budget.consume(); err != nil { - return nil, err - } - - var keys []string - - // Handle dictionaries (maps) - if dict, ok := obj.(map[string]any); ok { - for k, v := range dict { - if k == "$ref" { - refPath, isString := v.(string) - if !isString { - return nil, fmt.Errorf("%w: '$ref' value must be a string, got %T", errInvalidJSONSchema, v) - } - - // Skip if reference has already been processed (circular reference protection) - if _, exists := processedRefs[refPath]; exists { - return nil, fmt.Errorf("%w: circular reference detected: %s", errInvalidJSONSchema, refPath) - } - processedRefs[refPath] = struct{}{} - - ref, err := jsonSchemaRetrieveRef(refPath, fullSchema) - if err != nil { - delete(processedRefs, refPath) // Clean up on error - return nil, err - } - - // Add the top-level key of the reference to the list - // This relies on the reference path format, e.g., "#/components/..." - components := strings.Split(refPath, "/") - if len(components) > 1 { - keys = append(keys, components[1]) - } - - // Recurse on the referenced schema - nestedKeys, err := jsonSchemaSkipKeys(ref, fullSchema, processedRefs, budget, depth+1) - if err != nil { - delete(processedRefs, refPath) // Clean up on error - return nil, err - } - keys = append(keys, nestedKeys...) - - // Clean up after processing - delete(processedRefs, refPath) - } else if _, isDict := v.(map[string]any); isDict { - nestedKeys, err := jsonSchemaSkipKeys(v, fullSchema, processedRefs, budget, depth+1) - if err != nil { - return nil, err - } - keys = append(keys, nestedKeys...) - } else if _, isList := v.([]any); isList { - nestedKeys, err := jsonSchemaSkipKeys(v, fullSchema, processedRefs, budget, depth+1) - if err != nil { - return nil, err - } - keys = append(keys, nestedKeys...) - } - } - } else if list, ok := obj.([]any); ok { - // Handle lists (slices) - for _, el := range list { - nestedKeys, err := jsonSchemaSkipKeys(el, fullSchema, processedRefs, budget, depth+1) - if err != nil { - return nil, err - } - keys = append(keys, nestedKeys...) - } - } - - return keys, nil -} - -// jsonSchemaDereference substitutes $refs in a JSON Schema object. -func jsonSchemaDereference(schemaObj map[string]any) (any, error) { - if schemaObj == nil { - return nil, fmt.Errorf("%w: schema object cannot be nil", errInvalidJSONSchema) - } - - // The budget is shared by both passes so the total work is bounded. - budget := newJSONSchemaNodeBudget() - processedRefs := make(map[string]struct{}) - skipKeys, err := jsonSchemaSkipKeys(schemaObj, schemaObj, processedRefs, budget, 0) - if err != nil { - return nil, fmt.Errorf("failed to determine skip keys: %w", err) - } - - // Reset processed refs for the actual dereferencing - processedRefs = make(map[string]struct{}) - - // Call the recursive helper function to perform the dereferencing - return jsonSchemaDereferenceHelper(schemaObj, schemaObj, skipKeys, processedRefs, budget, 0) -} - -// jsonSchemaToGapic formats a JSON schema for a gapic request with improved safety. -// This is modified from the langchain-google implementation with enhanced error handling. -func jsonSchemaToGapic(schema map[string]any, allowedSchemaFieldsSet map[string]struct{}) (map[string]any, error) { - if schema == nil { - return nil, fmt.Errorf("%w: schema cannot be nil", errInvalidJSONSchema) - } - if allowedSchemaFieldsSet == nil { - return nil, fmt.Errorf("%w: allowedSchemaFieldsSet cannot be nil", errInvalidJSONSchema) - } - - convertedSchema := make(map[string]any) - - for key, value := range schema { - switch key { - case "$defs": - // Skip $defs as they are handled separately - continue - - case "items": - subSchema, ok := value.(map[string]any) - if !ok { - return nil, fmt.Errorf("%w: 'items' must be a dict, got %T", errInvalidJSONSchema, value) - } - convertedItems, err := jsonSchemaToGapic(subSchema, allowedSchemaFieldsSet) - if err != nil { - return nil, fmt.Errorf("failed to convert items schema: %w", err) - } - convertedSchema["items"] = convertedItems - - case "properties": - properties, ok := value.(map[string]any) - if !ok { - return nil, fmt.Errorf("%w: 'properties' must be a dict, got %T", errInvalidJSONSchema, value) - } - - convertedProperties := make(map[string]any) - for pkey, pvalue := range properties { - pSubSchema, ok := pvalue.(map[string]any) - if !ok { - return nil, fmt.Errorf("%w: property '%s' must be a dict, got %T", errInvalidJSONSchema, pkey, pvalue) - } - - convertedProperty, err := jsonSchemaToGapic(pSubSchema, allowedSchemaFieldsSet) - if err != nil { - return nil, fmt.Errorf("failed to convert property '%s': %w", pkey, err) - } - convertedProperties[pkey] = convertedProperty - } - convertedSchema["properties"] = convertedProperties - - case "type": - convertedType, err := processJSONSchemaTypeField(value) - if err != nil { - return nil, err - } - // Merge the converted type result into the schema - for k, v := range convertedType { - convertedSchema[k] = v - } - - case "allOf": - convertedAllOf, err := processJSONSchemaAllOfField(value, allowedSchemaFieldsSet) - if err != nil { - return nil, err - } - return convertedAllOf, nil - - case "anyOf": - convertedAnyOf, err := processJSONSchemaAnyOfField(value, allowedSchemaFieldsSet) - if err != nil { - return nil, err - } - // Merge the anyOf result - for k, v := range convertedAnyOf { - convertedSchema[k] = v - } - - default: - // Check if the key is in the allowed set - if _, allowed := allowedSchemaFieldsSet[key]; allowed { - convertedSchema[key] = value - } - } - } - return convertedSchema, nil -} - -// processJSONSchemaTypeField handles the "type" field conversion with improved type safety -func processJSONSchemaTypeField(value any) (map[string]any, error) { - result := make(map[string]any) - - switch typeValue := value.(type) { - case []any: - if len(typeValue) != 2 { - return nil, fmt.Errorf("%w: if type is a list, length must be 2, got %d", errInvalidJSONSchema, len(typeValue)) - } - - hasNull := false - var nonNullType any - for _, t := range typeValue { - if t == "null" { - hasNull = true - } else { - nonNullType = t - } - } - - if !hasNull || nonNullType == nil { - return nil, fmt.Errorf("%w: if type is a list, it must contain one non-null type and 'null'", errInvalidJSONSchema) - } - - // Non-null types in JSON Schema should be strings, not maps - if _, ok := nonNullType.(map[string]any); ok { - return nil, fmt.Errorf("%w: unexpected map type in type array", errInvalidJSONSchema) - } - - result["type"] = fmt.Sprintf("%v", nonNullType) - result["nullable"] = true - - case string: - result["type"] = typeValue - - default: - return nil, fmt.Errorf("%w: 'type' must be a list or string, got %T", errInvalidJSONSchema, value) - } - - return result, nil -} - -// processJSONSchemaAllOfField handles the "allOf" field conversion -func processJSONSchemaAllOfField(value any, allowedSchemaFieldsSet map[string]struct{}) (map[string]any, error) { - allOfList, ok := value.([]any) - if !ok { - return nil, fmt.Errorf("%w: 'allOf' must be a list, got %T", errInvalidJSONSchema, value) - } - - if len(allOfList) == 0 { - return nil, fmt.Errorf("%w: 'allOf' cannot be empty", errInvalidJSONSchema) - } - - if len(allOfList) > 1 { - return nil, fmt.Errorf("%w: only one value for 'allOf' key is supported, got %d", errInvalidJSONSchema, len(allOfList)) - } - - subSchema, ok := allOfList[0].(map[string]any) - if !ok { - return nil, fmt.Errorf("%w: item in 'allOf' must be an object, got %T", errInvalidJSONSchema, allOfList[0]) - } - - return jsonSchemaToGapic(subSchema, allowedSchemaFieldsSet) -} - -// processJSONSchemaAnyOfField handles the "anyOf" field conversion -func processJSONSchemaAnyOfField(value any, allowedSchemaFieldsSet map[string]struct{}) (map[string]any, error) { - anyOfList, ok := value.([]any) - if !ok { - return nil, fmt.Errorf("%w: 'anyOf' must be a list, got %T", errInvalidJSONSchema, value) - } - - if len(anyOfList) == 0 { - return nil, fmt.Errorf("%w: 'anyOf' cannot be empty", errInvalidJSONSchema) - } - - result := make(map[string]any) - anyOfResults := make([]any, 0) - nullable := false - - for i, v := range anyOfList { - subSchema, ok := v.(map[string]any) - if !ok { - return nil, fmt.Errorf("%w: item %d in 'anyOf' must be a dict, got %T", errInvalidJSONSchema, i, v) - } - - if t, exists := subSchema["type"]; exists && t == "null" { - nullable = true - } else { - convertedSubSchema, err := jsonSchemaToGapic(subSchema, allowedSchemaFieldsSet) - if err != nil { - return nil, fmt.Errorf("failed to convert anyOf item %d: %w", i, err) - } - anyOfResults = append(anyOfResults, convertedSubSchema) - } - } - - if nullable { - result["nullable"] = true - } - result["anyOf"] = anyOfResults - - return result, nil -} - -// jsonSchemaMapToSchema safely converts a map[string]any to a genai.Schema struct. -// This function includes proper input validation and error handling. -func jsonSchemaMapToSchema(schemaMap map[string]any) (*genai.Schema, error) { - if schemaMap == nil { - return nil, fmt.Errorf("%w: schemaMap cannot be nil", errInvalidJSONSchema) - } - - // Marshal the map to JSON - jsonBytes, err := json.Marshal(schemaMap) - if err != nil { - return nil, fmt.Errorf("failed to marshal map to JSON: %w", err) - } - - // Unmarshal into the genai.Schema struct - var genSchema genai.Schema - if err := json.Unmarshal(jsonBytes, &genSchema); err != nil { - return nil, fmt.Errorf("failed to unmarshal JSON to Schema: %w", err) - } - - return &genSchema, nil -} - -// jsonSchemaToGemini converts a JSON schema to a Gemini Schema with comprehensive validation. -// This function combines dereferencing and GCP formatting with proper error handling. -func jsonSchemaToGemini(schema map[string]any) (*genai.Schema, error) { - if schema == nil { - return nil, fmt.Errorf("%w: schema cannot be nil", errInvalidJSONSchema) - } - - // Step 1: Dereference the schema to resolve all $ref pointers - dereferencedSchema, err := jsonSchemaDereference(schema) - if err != nil { - return nil, fmt.Errorf("failed to dereference schema: %w", err) - } - - // Ensure the dereferenced result is a map - dereferencedMap, ok := dereferencedSchema.(map[string]any) - if !ok { - return nil, fmt.Errorf("%w: dereferenced schema is not a map[string]any, got %T", errInvalidJSONSchema, dereferencedSchema) - } - - // Step 2: Define allowed schema fields for genai.Schema - // This list corresponds to the fields supported by the genai.Schema struct - allowedSchemaFieldsSet := map[string]struct{}{ - "anyOf": {}, - "default": {}, - "description": {}, - "enum": {}, - "example": {}, - "format": {}, - "items": {}, - "maxItems": {}, - "maxLength": {}, - "maxProperties": {}, - "maximum": {}, - "minItems": {}, - "minLength": {}, - "minProperties": {}, - "minimum": {}, - "nullable": {}, - "pattern": {}, - "properties": {}, - "propertyOrdering": {}, - "required": {}, - "title": {}, - "type": {}, - } - - // Step 3: Convert to GCP/Gapic compatible format - schemaMap, err := jsonSchemaToGapic(dereferencedMap, allowedSchemaFieldsSet) - if err != nil { - return nil, fmt.Errorf("failed to convert to Gapic format: %w", err) - } - - // Step 4: Convert to genai.Schema struct - retSchema, err := jsonSchemaMapToSchema(schemaMap) - if err != nil { - return nil, fmt.Errorf("failed to convert to genai.Schema: %w", err) - } - - return retSchema, nil -} diff --git a/internal/translator/jsonschema_helper_test.go b/internal/translator/jsonschema_helper_test.go deleted file mode 100644 index 78e8c5be21..0000000000 --- a/internal/translator/jsonschema_helper_test.go +++ /dev/null @@ -1,1051 +0,0 @@ -// Copyright Envoy AI Gateway Authors -// SPDX-License-Identifier: Apache-2.0 -// The full text of the Apache license is available in the LICENSE file at -// the root of the repo. - -package translator - -import ( - "fmt" - "testing" - "time" - - "github.com/google/go-cmp/cmp" - "github.com/google/go-cmp/cmp/cmpopts" - "github.com/stretchr/testify/require" - "google.golang.org/genai" - - "github.com/envoyproxy/ai-gateway/internal/json" -) - -func TestJsonSchemaDeepCopyMapStringAny(t *testing.T) { - tests := []struct { - name string - input map[string]any - expected map[string]any - }{ - { - name: "nil map", - input: nil, - expected: nil, - }, - { - name: "empty map", - input: map[string]any{}, - expected: map[string]any{}, - }, - { - name: "simple map", - input: map[string]any{ - "key1": "value1", - "key2": 42, - "key3": true, - }, - expected: map[string]any{ - "key1": "value1", - "key2": 42, - "key3": true, - }, - }, - { - name: "nested map", - input: map[string]any{ - "outer": map[string]any{ - "inner": "value", - }, - }, - expected: map[string]any{ - "outer": map[string]any{ - "inner": "value", - }, - }, - }, - { - name: "map with slice", - input: map[string]any{ - "items": []any{"a", "b", "c"}, - }, - expected: map[string]any{ - "items": []any{"a", "b", "c"}, - }, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - got, err := jsonSchemaDeepCopyMapStringAny(tc.input) - require.NoError(t, err) - require.Equal(t, tc.expected, got) - - // Verify it's a deep copy by modifying original - if tc.input != nil { - tc.input["new_key"] = "new_value" - require.NotContains(t, got, "new_key") - } - }) - } -} - -func TestJsonSchemaDeepCopyAny(t *testing.T) { - tests := []struct { - name string - input any - expected any - }{ - { - name: "string", - input: "test", - expected: "test", - }, - { - name: "int", - input: 42, - expected: 42, - }, - { - name: "bool", - input: true, - expected: true, - }, - { - name: "nil", - input: nil, - expected: nil, - }, - { - name: "map", - input: map[string]any{"key": "value"}, - expected: map[string]any{"key": "value"}, - }, - { - name: "slice", - input: []any{1, 2, 3}, - expected: []any{1, 2, 3}, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - got, err := jsonSchemaDeepCopyAny(tc.input) - require.NoError(t, err) - require.Equal(t, tc.expected, got) - }) - } -} - -func TestJsonSchemaRetrieveRef(t *testing.T) { - schema := map[string]any{ - "components": map[string]any{ - "schemas": map[string]any{ - "User": map[string]any{ - "type": "object", - "properties": map[string]any{ - "name": map[string]any{"type": "string"}, - }, - }, - }, - }, - } - - tests := []struct { - name string - path string - schema map[string]any - expectedResult any - expectedErrMsg string - }{ - { - name: "valid reference", - path: "#/components/schemas/User", - schema: schema, - expectedResult: map[string]any{ - "type": "object", - "properties": map[string]any{ - "name": map[string]any{"type": "string"}, - }, - }, - }, - { - name: "empty path", - path: "", - schema: schema, - expectedErrMsg: "ref path cannot be empty", - }, - { - name: "invalid path format", - path: "invalid/path", - schema: schema, - expectedErrMsg: "ref paths must start with '#/'", - }, - { - name: "path with empty component", - path: "#//components", - schema: schema, - expectedErrMsg: "ref path contains empty component", - }, - { - name: "non-existent reference", - path: "#/components/schemas/NonExistent", - schema: schema, - expectedErrMsg: "component 'NonExistent' does not exist", - }, - { - name: "reference to non-map intermediate", - path: "#/components/schemas/User/type/invalid", - schema: schema, - expectedErrMsg: "intermediate component 'type' is not a map", - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - got, err := jsonSchemaRetrieveRef(tc.path, tc.schema) - - if tc.expectedErrMsg != "" { - require.ErrorContains(t, err, tc.expectedErrMsg) - require.Nil(t, got) - } else { - require.NoError(t, err) - require.Equal(t, tc.expectedResult, got) - } - }) - } -} - -func TestJsonSchemaDereferenceHelper(t *testing.T) { - schema := map[string]any{ - "components": map[string]any{ - "schemas": map[string]any{ - "User": map[string]any{ - "type": "object", - "properties": map[string]any{ - "name": map[string]any{"type": "string"}, - }, - }, - }, - }, - } - - tests := []struct { - name string - obj any - fullSchema map[string]any - skipKeys []string - processedRefs map[string]struct{} - expectedResult any - expectedErrMsg string - }{ - { - name: "primitive value", - obj: "string", - fullSchema: schema, - skipKeys: []string{}, - expectedResult: "string", - }, - { - name: "map without refs", - obj: map[string]any{"type": "string"}, - fullSchema: schema, - skipKeys: []string{}, - expectedResult: map[string]any{"type": "string"}, - }, - { - name: "map with valid ref", - obj: map[string]any{ - "$ref": "#/components/schemas/User", - }, - fullSchema: schema, - skipKeys: []string{}, - expectedResult: map[string]any{ - "type": "object", - "properties": map[string]any{ - "name": map[string]any{"type": "string"}, - }, - }, - }, - { - name: "map with non-string ref", - obj: map[string]any{ - "$ref": 123, - }, - fullSchema: schema, - skipKeys: []string{}, - expectedErrMsg: "'$ref' value must be a string, got int", - }, - { - name: "recursive reference", - obj: map[string]any{ - "$ref": "#/components/schemas/User", - }, - fullSchema: schema, - skipKeys: []string{}, - processedRefs: map[string]struct{}{ - "#/components/schemas/User": {}, - }, - expectedErrMsg: "circular reference detected", - }, - { - name: "skip key", - obj: map[string]any{ - "skipme": "value", - "keepme": "value", - }, - fullSchema: schema, - skipKeys: []string{"skipme"}, - expectedResult: map[string]any{ - "skipme": "value", - "keepme": "value", - }, - }, - { - name: "slice", - obj: []any{"a", "b"}, - fullSchema: schema, - skipKeys: []string{}, - expectedResult: []any{"a", "b"}, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - // Create processedRefs if provided - processedRefs := tc.processedRefs - if processedRefs == nil { - processedRefs = make(map[string]struct{}) - } - - got, err := jsonSchemaDereferenceHelper(tc.obj, tc.fullSchema, tc.skipKeys, processedRefs, newJSONSchemaNodeBudget(), 0) - - if tc.expectedErrMsg != "" { - require.ErrorContains(t, err, tc.expectedErrMsg) - } else { - require.NoError(t, err) - require.Equal(t, tc.expectedResult, got) - } - }) - } -} - -func TestJsonSchemaSkipKeys(t *testing.T) { - schema := map[string]any{ - "components": map[string]any{ - "schemas": map[string]any{ - "User": map[string]any{ - "type": "object", - }, - }, - }, - } - - tests := []struct { - name string - obj any - fullSchema map[string]any - processedRefs map[string]struct{} - expectedResult []string - expectedErrMsg string - }{ - { - name: "primitive value", - obj: "string", - fullSchema: schema, - expectedResult: nil, - }, - { - name: "map with ref", - obj: map[string]any{ - "$ref": "#/components/schemas/User", - }, - fullSchema: schema, - expectedResult: []string{"components"}, - }, - { - name: "map with non-string ref", - obj: map[string]any{ - "$ref": 123, - }, - fullSchema: schema, - expectedErrMsg: "'$ref' value must be a string, got int", - }, - { - name: "recursive reference", - obj: map[string]any{ - "$ref": "#/components/schemas/User", - }, - fullSchema: schema, - processedRefs: map[string]struct{}{ - "#/components/schemas/User": {}, - }, - expectedErrMsg: "circular reference detected", - }, - { - name: "slice", - obj: []any{"a", "b"}, - fullSchema: schema, - expectedResult: nil, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - // Create processedRefs if provided - processedRefs := tc.processedRefs - if processedRefs == nil { - processedRefs = make(map[string]struct{}) - } - - got, err := jsonSchemaSkipKeys(tc.obj, tc.fullSchema, processedRefs, newJSONSchemaNodeBudget(), 0) - - if tc.expectedErrMsg != "" { - require.ErrorContains(t, err, tc.expectedErrMsg) - } else { - require.NoError(t, err) - if tc.expectedResult == nil { - require.Empty(t, got) - } else { - require.Equal(t, tc.expectedResult, got) - } - } - }) - } -} - -func TestJsonSchemaDereference(t *testing.T) { - tests := []struct { - name string - schemaObj map[string]any - expectedErrMsg string - }{ - { - name: "nil schema", - schemaObj: nil, - expectedErrMsg: "schema object cannot be nil", - }, - { - name: "valid schema", - schemaObj: map[string]any{ - "type": "object", - }, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - got, err := jsonSchemaDereference(tc.schemaObj) - - if tc.expectedErrMsg != "" { - require.ErrorContains(t, err, tc.expectedErrMsg) - } else { - require.NoError(t, err) - require.NotNil(t, got) - } - }) - } -} - -// exponentialRefSchema returns a schema where each definition references the one below it twice, -// so that fully expanding it yields about 2^levels nodes. -func exponentialRefSchema(levels int) map[string]any { - defs := map[string]any{"d0": map[string]any{"type": "string"}} - for i := 1; i <= levels; i++ { - ref := fmt.Sprintf("#/$defs/d%d", i-1) - defs[fmt.Sprintf("d%d", i)] = map[string]any{ - "type": "object", - "properties": map[string]any{ - "a": map[string]any{"$ref": ref}, - "b": map[string]any{"$ref": ref}, - }, - } - } - return map[string]any{"$defs": defs, "$ref": fmt.Sprintf("#/$defs/d%d", levels)} -} - -// exponentialAnyOfRefSchema is like exponentialRefSchema, but the duplicated $refs are elements -// of an anyOf list instead of object properties. -func exponentialAnyOfRefSchema(levels int) map[string]any { - defs := map[string]any{"d0": map[string]any{"type": "string"}} - for i := 1; i <= levels; i++ { - ref := fmt.Sprintf("#/$defs/d%d", i-1) - defs[fmt.Sprintf("d%d", i)] = map[string]any{ - "anyOf": []any{ - map[string]any{"$ref": ref}, - map[string]any{"$ref": ref}, - }, - } - } - return map[string]any{"$defs": defs, "$ref": fmt.Sprintf("#/$defs/d%d", levels)} -} - -func TestJsonSchemaDereferenceNodeLimit(t *testing.T) { - t.Run("small expansion is allowed", func(t *testing.T) { - _, err := jsonSchemaDereference(exponentialRefSchema(8)) - require.NoError(t, err) - }) - - t.Run("exponential expansion is rejected", func(t *testing.T) { - start := time.Now() - // Keep the number of levels low enough to stay under jsonSchemaMaxRecursionDepth (each - // level adds 3 frames), so the node limit is the only one that can be hit regardless of - // the map iteration order. - _, err := jsonSchemaDereference(exponentialRefSchema(25)) - require.ErrorIs(t, err, errJSONSchemaMaxNodesExceeded) - require.Less(t, time.Since(start), 5*time.Second) - }) - - t.Run("small anyOf expansion is allowed", func(t *testing.T) { - _, err := jsonSchemaDereference(exponentialAnyOfRefSchema(8)) - require.NoError(t, err) - }) - - t.Run("exponential anyOf expansion is rejected", func(t *testing.T) { - start := time.Now() - _, err := jsonSchemaDereference(exponentialAnyOfRefSchema(30)) - require.ErrorIs(t, err, errJSONSchemaMaxNodesExceeded) - require.Less(t, time.Since(start), 5*time.Second) - }) - - t.Run("expansion within skip keys budget is still bounded", func(t *testing.T) { - budget := newJSONSchemaNodeBudget() - schema := exponentialRefSchema(14) - _, err := jsonSchemaDereferenceHelper(schema, schema, []string{"$defs"}, map[string]struct{}{}, budget, 0) - require.ErrorIs(t, err, errJSONSchemaMaxNodesExceeded) - }) -} - -func TestJsonSchemaToGapic(t *testing.T) { - allowedFields := map[string]struct{}{ - "type": {}, - "description": {}, - "properties": {}, - "items": {}, - "nullable": {}, - "anyOf": {}, - } - - tests := []struct { - name string - schema map[string]any - allowedFields map[string]struct{} - expectedErrMsg string - }{ - { - name: "nil schema", - schema: nil, - allowedFields: allowedFields, - expectedErrMsg: "schema cannot be nil", - }, - { - name: "nil allowed fields", - schema: map[string]any{}, - allowedFields: nil, - expectedErrMsg: "allowedSchemaFieldsSet cannot be nil", - }, - { - name: "invalid items type", - schema: map[string]any{ - "items": "not a map", - }, - allowedFields: allowedFields, - expectedErrMsg: "'items' must be a dict, got string", - }, - { - name: "invalid properties type", - schema: map[string]any{ - "properties": "not a map", - }, - allowedFields: allowedFields, - expectedErrMsg: "'properties' must be a dict, got string", - }, - { - name: "invalid property type", - schema: map[string]any{ - "properties": map[string]any{ - "field": "not a map", - }, - }, - allowedFields: allowedFields, - expectedErrMsg: "property 'field' must be a dict, got string", - }, - { - name: "invalid type value", - schema: map[string]any{ - "type": 123, - }, - allowedFields: allowedFields, - expectedErrMsg: "'type' must be a list or string, got int", - }, - { - name: "type list with wrong length", - schema: map[string]any{ - "type": []any{"string", "number", "boolean"}, - }, - allowedFields: allowedFields, - expectedErrMsg: "if type is a list, length must be 2", - }, - { - name: "type list without null", - schema: map[string]any{ - "type": []any{"string", "number"}, - }, - allowedFields: allowedFields, - expectedErrMsg: "it must contain one non-null type and 'null'", - }, - { - name: "invalid allOf type", - schema: map[string]any{ - "allOf": "not a list", - }, - allowedFields: allowedFields, - expectedErrMsg: "'allOf' must be a list, got string", - }, - { - name: "empty allOf", - schema: map[string]any{ - "allOf": []any{}, - }, - allowedFields: allowedFields, - expectedErrMsg: "'allOf' cannot be empty", - }, - { - name: "allOf with too many items", - schema: map[string]any{ - "allOf": []any{ - map[string]any{"type": "string"}, - map[string]any{"type": "number"}, - }, - }, - allowedFields: allowedFields, - expectedErrMsg: "only one value for 'allOf' key is supported", - }, - { - name: "invalid allOf item type", - schema: map[string]any{ - "allOf": []any{"not a map"}, - }, - allowedFields: allowedFields, - expectedErrMsg: "item in 'allOf' must be an object", - }, - { - name: "invalid anyOf type", - schema: map[string]any{ - "anyOf": "not a list", - }, - allowedFields: allowedFields, - expectedErrMsg: "'anyOf' must be a list, got string", - }, - { - name: "empty anyOf", - schema: map[string]any{ - "anyOf": []any{}, - }, - allowedFields: allowedFields, - expectedErrMsg: "'anyOf' cannot be empty", - }, - { - name: "invalid anyOf item type", - schema: map[string]any{ - "anyOf": []any{"not a map"}, - }, - allowedFields: allowedFields, - expectedErrMsg: "item 0 in 'anyOf' must be a dict, got string", - }, - { - name: "valid schema", - schema: map[string]any{ - "type": "object", - "description": "A test object", - "properties": map[string]any{ - "name": map[string]any{ - "type": "string", - }, - }, - }, - allowedFields: allowedFields, - }, - { - name: "schema with $defs (skipped)", - schema: map[string]any{ - "type": "object", - "$defs": map[string]any{ - "User": map[string]any{"type": "object"}, - }, - }, - allowedFields: allowedFields, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - got, err := jsonSchemaToGapic(tc.schema, tc.allowedFields) - - if tc.expectedErrMsg != "" { - require.ErrorContains(t, err, tc.expectedErrMsg) - } else { - require.NoError(t, err) - require.NotNil(t, got) - } - }) - } -} - -func TestJsonSchemaMapToSchema(t *testing.T) { - tests := []struct { - name string - schemaMap map[string]any - expectedErrMsg string - }{ - { - name: "nil schema map", - schemaMap: nil, - expectedErrMsg: "schemaMap cannot be nil", - }, - { - name: "valid schema map", - schemaMap: map[string]any{ - "type": "string", - }, - }, - { - name: "schema map with invalid JSON values", - schemaMap: map[string]any{ - "invalid": func() {}, // functions can't be marshaled to JSON - }, - expectedErrMsg: "failed to marshal map to JSON", - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - got, err := jsonSchemaMapToSchema(tc.schemaMap) - - if tc.expectedErrMsg != "" { - require.ErrorContains(t, err, tc.expectedErrMsg) - require.Nil(t, got) - } else { - require.NoError(t, err) - require.NotNil(t, got) - } - }) - } -} - -func TestJsonSchemaToGemini(t *testing.T) { - tests := []struct { - name string - schema map[string]any - expectedErrMsg string - }{ - { - name: "nil schema", - schema: nil, - expectedErrMsg: "schema cannot be nil", - }, - { - name: "valid schema", - schema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "name": map[string]any{ - "type": "string", - }, - }, - }, - }, - { - name: "schema with invalid reference", - schema: map[string]any{ - "$ref": "#/invalid/ref", - }, - expectedErrMsg: "component 'invalid' does not exist", - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - got, err := jsonSchemaToGemini(tc.schema) - - if tc.expectedErrMsg != "" { - require.ErrorContains(t, err, tc.expectedErrMsg) - require.Nil(t, got) - } else { - require.NoError(t, err) - require.NotNil(t, got) - } - }) - } -} - -func TestJsonSchemaToGeminiIntegration(t *testing.T) { - trueBool := true - tests := []struct { - name string - input json.RawMessage - expectedResponseSchema *genai.Schema - expectedErrMsg string - }{ - { - name: "nested schema for ResponseSchema", - input: json.RawMessage(`{ - "type": "object", - "properties": { - "steps": { - "type": "array", - "items": { - "$ref": "#/$defs/step" - } - }, - "final_answer": { - "type": "string" - } - }, - "$defs": { - "step": { - "type": "object", - "properties": { - "explanation": { - "type": "string" - }, - "output": { - "type": "string" - } - }, - "required": [ - "explanation", - "output" - ], - "additionalProperties": false - } - }, - "required": [ - "steps", - "final_answer" - ], - "additionalProperties": false -}`), - - expectedResponseSchema: &genai.Schema{ - Properties: map[string]*genai.Schema{ - "final_answer": {Type: "string"}, - "steps": { - Items: &genai.Schema{ - Properties: map[string]*genai.Schema{ - "explanation": {Type: "string"}, - "output": {Type: "string"}, - }, - Type: "object", - Required: []string{"explanation", "output"}, - }, - Type: "array", - }, - }, - Type: "object", - Required: []string{"steps", "final_answer"}, - }, - }, - - { - name: "anyof list for ResponseSchema", - input: json.RawMessage(`{ - "type": "object", - "properties": { - "item": { - "anyOf": [ - { - "type": "object", - "description": "The user object to insert into the database", - "properties": { - "name": { - "type": "string", - "description": "The name of the user" - }, - "age": { - "type": "number", - "description": "The age of the user" - } - }, - "additionalProperties": false, - "required": [ - "name", - "age" - ] - }, - { - "type": "object", - "description": "The address object to insert into the database", - "properties": { - "number": { - "type": "string", - "description": "The number of the address. Eg. for 123 main st, this would be 123" - }, - "street": { - "type": "string", - "description": "The street name. Eg. for 123 main st, this would be main st" - }, - "city": { - "type": "string", - "description": "The city of the address" - } - }, - "additionalProperties": false, - "required": [ - "number", - "street", - "city" - ] - }, - { - "type": "object", - "description": "The email address object to insert into the database", - "properties": { - "company": { - "type": "string", - "description": "The company to use." - }, - "url": { - "type": "string", - "description": "The email address" - } - }, - "additionalProperties": false, - "required": [ - "company", - "url" - ] - } - ] - } - }, - "additionalProperties": false, - "required": [ - "item" - ] -}`), - - expectedResponseSchema: &genai.Schema{ - Properties: map[string]*genai.Schema{ - "item": { - AnyOf: []*genai.Schema{ - { - Description: "The user object to insert into the database", - Properties: map[string]*genai.Schema{ - "age": {Type: "number", Description: "The age of the user"}, - "name": {Type: "string", Description: "The name of the user"}, - }, - Type: "object", - Required: []string{"name", "age"}, - }, - { - Description: "The address object to insert into the database", - Properties: map[string]*genai.Schema{ - "city": {Type: "string", Description: "The city of the address"}, - "number": {Type: "string", Description: "The number of the address. Eg. for 123 main st, this would be 123"}, - "street": {Type: "string", Description: "The street name. Eg. for 123 main st, this would be main st"}, - }, - Type: "object", - Required: []string{"number", "street", "city"}, - }, - { - Description: "The email address object to insert into the database", - Properties: map[string]*genai.Schema{ - "company": {Type: "string", Description: "The company to use."}, - "url": {Type: "string", Description: "The email address"}, - }, - Required: []string{"company", "url"}, - Type: "object", - }, - }, - }, - }, - Type: "object", - Required: []string{"item"}, - }, - }, - - { - name: "anyof null for ResponseSchema", - input: json.RawMessage(`{ - "type": "object", - "description": "Data model identifying a single paragraph for paragraph re-ranking.", - "properties": { - "paragraph_id": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Paragraph Id" - }, - "document_id": { - "title": "Document Id", - "type": "string" - } - }, - "required": [ - "paragraph_id", - "document_id" - ], - "title": "ParagraphIdentifier", - "additionalProperties": false -}`), - - expectedResponseSchema: &genai.Schema{ - Description: "Data model identifying a single paragraph for paragraph re-ranking.", - Properties: map[string]*genai.Schema{ - "document_id": {Type: "string", Title: "Document Id"}, - "paragraph_id": { - AnyOf: []*genai.Schema{ - { - Type: "string", - }, - }, - Nullable: &trueBool, - Title: "Paragraph Id", - }, - }, - Required: []string{"paragraph_id", "document_id"}, - Title: "ParagraphIdentifier", - Type: "object", - }, - }, - - { - name: "allOf with not a list", - input: json.RawMessage(`{ - "type": "object", - "properties": { - "item": { - "allOf": "string" - } - }, - "additionalProperties": false, - "required": [ - "item" - ] -}`), - expectedErrMsg: "invalid JSON schema: 'allOf' must be a list", - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - var schemaMap map[string]any - err := json.Unmarshal([]byte(tc.input), &schemaMap) - require.NoError(t, err) - - got, err := jsonSchemaToGemini(schemaMap) - - if tc.expectedErrMsg != "" { - require.ErrorContains(t, err, tc.expectedErrMsg) - } else { - require.NoError(t, err) - - if diff := cmp.Diff(tc.expectedResponseSchema, got, cmpopts.IgnoreUnexported(genai.Schema{})); diff != "" { - t.Errorf("ResponseSchema mismatch (-want +got):\n%s", diff) - } - } - }) - } -} diff --git a/internal/translator/openai_gcpvertexai.go b/internal/translator/openai_gcpvertexai.go index 529f265ab8..8b0a29e27a 100644 --- a/internal/translator/openai_gcpvertexai.go +++ b/internal/translator/openai_gcpvertexai.go @@ -523,10 +523,8 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) openAIMessageToGeminiMes return nil, err } - // Some models support only partialJSONSchema. - parametersJSONSchemaAvailable := responseJSONSchemaAvailable(requestModel) // Convert OpenAI tools to Gemini tools. - tools, err := openAIToolsToGeminiTools(openAIReq.Tools, parametersJSONSchemaAvailable) + tools, err := openAIToolsToGeminiTools(openAIReq.Tools) if err != nil { return nil, fmt.Errorf("invalid tools: %w", err) } diff --git a/internal/translator/openai_gcpvertexai_test.go b/internal/translator/openai_gcpvertexai_test.go index 439c82d219..bebaad7ed8 100644 --- a/internal/translator/openai_gcpvertexai_test.go +++ b/internal/translator/openai_gcpvertexai_test.go @@ -112,7 +112,7 @@ func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_RequestBody(t *testing.T) { "name": "get_weather", "description": "Get the current weather in a given location", - "parameters": { + "parametersJsonSchema": { "type": "object", "properties": { "location": { @@ -158,7 +158,7 @@ func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_RequestBody(t *testing.T) { "name": "test_function", "description": "A test function", - "parameters": { + "parametersJsonSchema": { "type": "object", "properties": { "param1": { @@ -198,7 +198,7 @@ func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_RequestBody(t *testing.T) { "name": "test_function", "description": "A test function", - "parameters": { + "parametersJsonSchema": { "type": "object", "properties": { "param1": { @@ -270,7 +270,7 @@ func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_RequestBody(t *testing.T) { "name": "test_function", "description": "A test function", - "parameters": { + "parametersJsonSchema": { "type": "object", "properties": { "param1": { @@ -313,7 +313,7 @@ func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_RequestBody(t *testing.T) { "name": "test_function", "description": "A test function", - "parameters": { + "parametersJsonSchema": { "type": "object", "properties": { "param1": { diff --git a/internal/translator/tokenize_gcpvertexai.go b/internal/translator/tokenize_gcpvertexai.go index e22dfce61b..8ecee47362 100644 --- a/internal/translator/tokenize_gcpvertexai.go +++ b/internal/translator/tokenize_gcpvertexai.go @@ -47,10 +47,8 @@ func (o *ToGCPVertexAIV1Tokenize) tokenizeToGeminiCountToken(tokenizeChatReq *to return nil, fmt.Errorf("messages must produce at least one content entry") } - // Some models support only partialJSONSchema. - parametersJSONSchemaAvailable := responseJSONSchemaAvailable(requestModel) // Convert OpenAI tools to Gemini tools. - tools, err := openAIToolsToGeminiTools(tokenizeChatReq.Tools, parametersJSONSchemaAvailable) + tools, err := openAIToolsToGeminiTools(tokenizeChatReq.Tools) if err != nil { return nil, fmt.Errorf("error converting tools: %w", err) } diff --git a/tests/data-plane/testupstream_test.go b/tests/data-plane/testupstream_test.go index f784a110e8..cbd49272b8 100644 --- a/tests/data-plane/testupstream_test.go +++ b/tests/data-plane/testupstream_test.go @@ -419,42 +419,42 @@ func TestWithTestUpstream(t *testing.T) { backend: "gcp-vertexai", path: "/v1/chat/completions", method: http.MethodPost, - requestBody: `{"model":"gemini-1.5-pro","messages":[{"role":"system","content":"You are a helpful assistant."}]}`, + requestBody: `{"model":"gemini-3.5-flash","messages":[{"role":"system","content":"You are a helpful assistant."}]}`, expRequestBody: `{"contents":null,"tools":null,"generationConfig":{},"systemInstruction":{"parts":[{"text":"You are a helpful assistant."}]}}`, - expPath: "/v1/projects/gcp-project-name/locations/gcp-region/publishers/google/models/gemini-1.5-pro:generateContent", + expPath: "/v1/projects/gcp-project-name/locations/gcp-region/publishers/google/models/gemini-3.5-flash:generateContent", expRequestHeaders: map[string]string{"Authorization": "Bearer " + fakeGCPAuthToken}, responseStatus: strconv.Itoa(http.StatusOK), responseBody: `{"candidates":[{"content":{"parts":[{"text":"This is a test response from Gemini."}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":15,"candidatesTokenCount":10,"totalTokenCount":35,"cachedContentTokenCount":10,"thoughtsTokenCount":10}}`, expStatus: http.StatusOK, - expResponseBody: `{"choices":[{"finish_reason":"stop","index":0,"message":{"content":"This is a test response from Gemini.","role":"assistant"}}],"model":"gemini-1.5-pro","object":"chat.completion","usage":{"completion_tokens":20,"completion_tokens_details":{"reasoning_tokens":10},"prompt_tokens":15,"prompt_tokens_details":{"cached_tokens":10},"total_tokens":35}}`, + expResponseBody: `{"choices":[{"finish_reason":"stop","index":0,"message":{"content":"This is a test response from Gemini.","role":"assistant"}}],"model":"gemini-3.5-flash","object":"chat.completion","usage":{"completion_tokens":20,"completion_tokens_details":{"reasoning_tokens":10},"prompt_tokens":15,"prompt_tokens_details":{"cached_tokens":10},"total_tokens":35}}`, }, { name: "gcp-vertexai - /v1/chat/completions", backend: "gcp-vertexai", path: "/v1/chat/completions", method: http.MethodPost, - requestBody: `{"model":"gemini-1.5-pro","messages":[{"role":"system","content":"You are a helpful assistant."}]}`, + requestBody: `{"model":"gemini-3.5-flash","messages":[{"role":"system","content":"You are a helpful assistant."}]}`, expRequestBody: `{"contents":null,"tools":null,"generationConfig":{},"systemInstruction":{"parts":[{"text":"You are a helpful assistant."}]}}`, - expPath: "/v1/projects/gcp-project-name/locations/gcp-region/publishers/google/models/gemini-1.5-pro:generateContent", + expPath: "/v1/projects/gcp-project-name/locations/gcp-region/publishers/google/models/gemini-3.5-flash:generateContent", expRequestHeaders: map[string]string{"Authorization": "Bearer " + fakeGCPAuthToken}, responseStatus: strconv.Itoa(http.StatusOK), responseBody: `{"candidates":[{"content":{"parts":[{"text":"This is a test response from Gemini."}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":15,"candidatesTokenCount":10,"totalTokenCount":25}}`, expStatus: http.StatusOK, - expResponseBody: `{"choices":[{"finish_reason":"stop","index":0,"message":{"content":"This is a test response from Gemini.","role":"assistant"}}],"model":"gemini-1.5-pro","object":"chat.completion","usage":{"completion_tokens":10,"completion_tokens_details":{},"prompt_tokens":15,"total_tokens":25,"prompt_tokens_details":{}}}`, + expResponseBody: `{"choices":[{"finish_reason":"stop","index":0,"message":{"content":"This is a test response from Gemini.","role":"assistant"}}],"model":"gemini-3.5-flash","object":"chat.completion","usage":{"completion_tokens":10,"completion_tokens_details":{},"prompt_tokens":15,"total_tokens":25,"prompt_tokens_details":{}}}`, }, { name: "gcp-vertexai - /v1/chat/completions - tool use", backend: "gcp-vertexai", path: "/v1/chat/completions", method: http.MethodPost, - requestBody: `{"model":"gemini-1.5-pro","messages":[{"role":"user","content":"tell me the delivery date for order 123"}],"tools":[{"type":"function","function":{"name":"get_delivery_date","description":"Get the delivery date for a customer's order. Call this whenever you need to know the delivery date, for example when a customer asks 'Where is my package'","parameters":{"type":"object","properties":{"order_id":{"type":"string","description":"The customer's order ID."}},"required":["order_id"]}}}]}`, + requestBody: `{"model":"gemini-3.5-flash","messages":[{"role":"user","content":"tell me the delivery date for order 123"}],"tools":[{"type":"function","function":{"name":"get_delivery_date","description":"Get the delivery date for a customer's order. Call this whenever you need to know the delivery date, for example when a customer asks 'Where is my package'","parameters":{"type":"object","properties":{"order_id":{"type":"string","description":"The customer's order ID."}},"required":["order_id"]}}}]}`, expRequestBody: `{"contents":[{"parts":[{"text":"tell me the delivery date for order 123"}],"role":"user"}],"tools":[{"functionDeclarations":[{"description":"Get the delivery date for a customer's order. Call this whenever you need to know the delivery date, for example when a customer asks 'Where is my package'","name":"get_delivery_date","parameters":{"properties":{"order_id":{"description":"The customer's order ID.","type":"string"}},"required":["order_id"],"type":"object"}}]}],"generationConfig":{}}`, - expPath: "/v1/projects/gcp-project-name/locations/gcp-region/publishers/google/models/gemini-1.5-pro:generateContent", + expPath: "/v1/projects/gcp-project-name/locations/gcp-region/publishers/google/models/gemini-3.5-flash:generateContent", expRequestHeaders: map[string]string{"Authorization": "Bearer " + fakeGCPAuthToken}, responseStatus: strconv.Itoa(http.StatusOK), - responseBody: `{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"get_delivery_date","args":{"order_id":"123"}}}]},"finishReason":"STOP","avgLogprobs":0.000001220789272338152}],"usageMetadata":{"promptTokenCount":50,"candidatesTokenCount":11,"totalTokenCount":61,"trafficType":"ON_DEMAND","promptTokensDetails":[{"modality":"TEXT","tokenCount":50}],"candidatesTokensDetails":[{"modality":"TEXT","tokenCount":11}]},"modelVersion":"gemini-2.0-flash-001","createTime":"2025-07-11T22:15:44.956335Z","responseId":"EI5xaK-vOtqJm22IPmuCR14AI"}`, + responseBody: `{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"get_delivery_date","args":{"order_id":"123"}}}]},"finishReason":"STOP","avgLogprobs":0.000001220789272338152}],"usageMetadata":{"promptTokenCount":50,"candidatesTokenCount":11,"totalTokenCount":61,"trafficType":"ON_DEMAND","promptTokensDetails":[{"modality":"TEXT","tokenCount":50}],"candidatesTokensDetails":[{"modality":"TEXT","tokenCount":11}]},"modelVersion":"gemini-3.5-flash-001","createTime":"2025-07-11T22:15:44.956335Z","responseId":"EI5xaK-vOtqJm22IPmuCR14AI"}`, expStatus: http.StatusOK, - expResponseBody: `{"choices":[{"finish_reason":"tool_calls","index":0,"message":{"role":"assistant","tool_calls":[{"id":"703482f8-2e5b-4dcc-a872-d74bd66c3866","function":{"arguments":"{\"order_id\":\"123\"}","name":"get_delivery_date"},"type":"function"}]}}],"created":123, "id":"EI5xaK-vOtqJm22IPmuCR14AI","model":"gemini-2.0-flash-001","object":"chat.completion","usage":{"completion_tokens":11,"completion_tokens_details":{},"prompt_tokens":50,"total_tokens":61,"prompt_tokens_details":{}}}`, + expResponseBody: `{"choices":[{"finish_reason":"tool_calls","index":0,"message":{"role":"assistant","tool_calls":[{"id":"703482f8-2e5b-4dcc-a872-d74bd66c3866","function":{"arguments":"{\"order_id\":\"123\"}","name":"get_delivery_date"},"type":"function"}]}}],"created":123, "id":"EI5xaK-vOtqJm22IPmuCR14AI","model":"gemini-3.5-flash-001","object":"chat.completion","usage":{"completion_tokens":11,"completion_tokens_details":{},"prompt_tokens":50,"total_tokens":61,"prompt_tokens_details":{}}}`, }, { name: "gcp-anthropicai - /v1/chat/completions", @@ -616,9 +616,9 @@ data: [DONE] path: "/v1/chat/completions", responseType: "sse", method: http.MethodPost, - requestBody: `{"model":"gemini-1.5-pro","messages":[{"role":"system","content":"You are a helpful assistant."}], "stream": true}`, + requestBody: `{"model":"gemini-3.5-flash","messages":[{"role":"system","content":"You are a helpful assistant."}], "stream": true}`, expRequestBody: `{"contents":null,"tools":null,"generationConfig":{},"systemInstruction":{"parts":[{"text":"You are a helpful assistant."}]}}`, - expPath: "/v1/projects/gcp-project-name/locations/gcp-region/publishers/google/models/gemini-1.5-pro:streamGenerateContent", + expPath: "/v1/projects/gcp-project-name/locations/gcp-region/publishers/google/models/gemini-3.5-flash:streamGenerateContent", expRawQuery: "alt=sse", expRequestHeaders: map[string]string{"Authorization": "Bearer " + fakeGCPAuthToken}, responseStatus: strconv.Itoa(http.StatusOK), @@ -630,21 +630,21 @@ data: [DONE] {"responseId":"msg_123","createTime":"2024-11-15T09:00:00Z","candidates":[{"content":{"parts":[{"text":" today"}],"role":"model"}}]} {"responseId":"msg_123","createTime":"2024-11-15T09:00:00Z","candidates":[{"content":{"parts":[{"text":"?"}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":7,"totalTokenCount":17}}`, expStatus: http.StatusOK, - expResponseBody: `data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":"Hello","role":"assistant"}}],"created":123,"model":"gemini-1.5-pro","object":"chat.completion.chunk"} + expResponseBody: `data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":"Hello","role":"assistant"}}],"created":123,"model":"gemini-3.5-flash","object":"chat.completion.chunk"} -data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":"! How","role":"assistant"}}],"created":123,"model":"gemini-1.5-pro","object":"chat.completion.chunk"} +data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":"! How","role":"assistant"}}],"created":123,"model":"gemini-3.5-flash","object":"chat.completion.chunk"} -data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":" can I","role":"assistant"}}],"created":123,"model":"gemini-1.5-pro","object":"chat.completion.chunk"} +data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":" can I","role":"assistant"}}],"created":123,"model":"gemini-3.5-flash","object":"chat.completion.chunk"} -data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":" help","role":"assistant"}}],"created":123,"model":"gemini-1.5-pro","object":"chat.completion.chunk"} +data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":" help","role":"assistant"}}],"created":123,"model":"gemini-3.5-flash","object":"chat.completion.chunk"} -data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":" you","role":"assistant"}}],"created":123,"model":"gemini-1.5-pro","object":"chat.completion.chunk"} +data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":" you","role":"assistant"}}],"created":123,"model":"gemini-3.5-flash","object":"chat.completion.chunk"} -data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":" today","role":"assistant"}}],"created":123,"model":"gemini-1.5-pro","object":"chat.completion.chunk"} +data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":" today","role":"assistant"}}],"created":123,"model":"gemini-3.5-flash","object":"chat.completion.chunk"} -data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":"?","role":"assistant"},"finish_reason":"stop"}],"created":123,"model":"gemini-1.5-pro","object":"chat.completion.chunk"} +data: {"id":"msg_123","choices":[{"index":0,"delta":{"content":"?","role":"assistant"},"finish_reason":"stop"}],"created":123,"model":"gemini-3.5-flash","object":"chat.completion.chunk"} -data: {"id":"msg_123","choices":[],"created":123,"model":"gemini-1.5-pro","object":"chat.completion.chunk","usage":{"prompt_tokens":10,"completion_tokens":7,"total_tokens":17,"completion_tokens_details":{},"prompt_tokens_details":{}}} +data: {"id":"msg_123","choices":[],"created":123,"model":"gemini-3.5-flash","object":"chat.completion.chunk","usage":{"prompt_tokens":10,"completion_tokens":7,"total_tokens":17,"completion_tokens_details":{},"prompt_tokens_details":{}}} data: [DONE] `, @@ -791,8 +791,8 @@ data: [DONE] path: "/v1/chat/completions", responseType: "", method: http.MethodPost, - requestBody: `{"model":"gemini-1.5-pro","messages":[{"role":"system","content":"You are a helpful assistant."}]}`, - expPath: "/v1/projects/gcp-project-name/locations/gcp-region/publishers/google/models/gemini-1.5-pro:generateContent", + requestBody: `{"model":"gemini-3.5-flash","messages":[{"role":"system","content":"You are a helpful assistant."}]}`, + expPath: "/v1/projects/gcp-project-name/locations/gcp-region/publishers/google/models/gemini-3.5-flash:generateContent", responseStatus: "400", expStatus: http.StatusBadRequest, responseBody: `{"error":{"code":400,"message":"Invalid request: missing required field","status":"INVALID_ARGUMENT"}}`,