Skip to content

Commit ad5f75e

Browse files
sukumargaonkaryuzisun
sukumargaonkar
and
yuzisun
authored
feat: enhance openAIToolsToGeminiTools to handle response JSON schema availability (#1512)
**Description** older GCP models (model < 2.5, eg gemini-flash-2.0) do not support the ParametersJsonSchema field for tool parameters. Instead, they require the Parameters field to be converted to gemini specific schema. This PR ensures that the correct schema format is used based on the model version. --------- Signed-off-by: Sukumar Gaonkar <sgaonkar4@bloomberg.net> Co-authored-by: Dan Sun <dsun20@bloomberg.net>
1 parent 1d23373 commit ad5f75e

File tree

5 files changed

+265
-40
lines changed

5 files changed

+265
-40
lines changed

‎internal/extproc/translator/gemini_helper.go‎

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -321,21 +321,43 @@ func assistantMsgToGeminiParts(msg openai.ChatCompletionAssistantMessageParam) (
321321
// }
322322
//
323323
// ].
324-
func openAIToolsToGeminiTools(openaiTools []openai.Tool) ([]genai.Tool, error) {
324+
func openAIToolsToGeminiTools(openaiTools []openai.Tool, parametersJSONSchemaAvailable bool) ([]genai.Tool, error) {
325325
if len(openaiTools) == 0 {
326326
return nil, nil
327327
}
328328
var functionDecls []*genai.FunctionDeclaration
329+
329330
for _, tool := range openaiTools {
330-
if tool.Type == openai.ToolTypeFunction {
331+
switch tool.Type {
332+
case openai.ToolTypeFunction:
331333
if tool.Function != nil {
334+
332335
functionDecl := &genai.FunctionDeclaration{
333-
Name: tool.Function.Name,
334-
Description: tool.Function.Description,
335-
ParametersJsonSchema: tool.Function.Parameters,
336+
Name: tool.Function.Name,
337+
Description: tool.Function.Description,
338+
}
339+
340+
if parametersJSONSchemaAvailable {
341+
functionDecl.ParametersJsonSchema = tool.Function.Parameters
342+
} else if tool.Function.Parameters != nil {
343+
paramsMap, ok := tool.Function.Parameters.(map[string]any)
344+
if !ok {
345+
return nil, fmt.Errorf("invalid JSON schema for parameters in tool %s: expected map[string]any, got %T", tool.Function.Name, tool.Function.Parameters)
346+
}
347+
348+
if len(paramsMap) > 0 {
349+
var err error
350+
if functionDecl.Parameters, err = jsonSchemaToGemini(paramsMap); err != nil {
351+
return nil, fmt.Errorf("invalid JSON schema for parameters in tool %s: %w", tool.Function.Name, err)
352+
}
353+
}
336354
}
337355
functionDecls = append(functionDecls, functionDecl)
338356
}
357+
case openai.ToolTypeImageGeneration:
358+
return nil, fmt.Errorf("tool-type image generation not supported yet when translating OpenAI req to Gemini")
359+
default:
360+
return nil, fmt.Errorf("unsupported tool type: %s", tool.Type)
339361
}
340362
}
341363
if len(functionDecls) == 0 {

‎internal/extproc/translator/gemini_helper_test.go‎

Lines changed: 224 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -941,18 +941,26 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) {
941941
"required": []any{"a", "b"},
942942
}
943943
tests := []struct {
944-
name string
945-
openaiTools []openai.Tool
946-
expected []genai.Tool
947-
expectedError string
944+
name string
945+
openaiTools []openai.Tool
946+
parametersJSONSchemaAvailable bool
947+
expected []genai.Tool
948+
expectedError string
948949
}{
949950
{
950-
name: "empty tools",
951-
openaiTools: nil,
952-
expected: nil,
951+
name: "empty tools with parametersJSONSchemaAvailable=false",
952+
openaiTools: nil,
953+
parametersJSONSchemaAvailable: false,
954+
expected: nil,
953955
},
954956
{
955-
name: "single function tool with parameters",
957+
name: "empty tools with parametersJSONSchemaAvailable=true",
958+
openaiTools: nil,
959+
parametersJSONSchemaAvailable: true,
960+
expected: nil,
961+
},
962+
{
963+
name: "single function tool with parameters - parametersJSONSchemaAvailable=false",
956964
openaiTools: []openai.Tool{
957965
{
958966
Type: openai.ToolTypeFunction,
@@ -963,6 +971,39 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) {
963971
},
964972
},
965973
},
974+
parametersJSONSchemaAvailable: false,
975+
expected: []genai.Tool{
976+
{
977+
FunctionDeclarations: []*genai.FunctionDeclaration{
978+
{
979+
Name: "add",
980+
Description: "Add two numbers",
981+
Parameters: &genai.Schema{
982+
Type: "object",
983+
Properties: map[string]*genai.Schema{
984+
"a": {Type: "integer"},
985+
"b": {Type: "integer"},
986+
},
987+
Required: []string{"a", "b"},
988+
},
989+
},
990+
},
991+
},
992+
},
993+
},
994+
{
995+
name: "single function tool with parameters - parametersJSONSchemaAvailable=true",
996+
openaiTools: []openai.Tool{
997+
{
998+
Type: openai.ToolTypeFunction,
999+
Function: &openai.FunctionDefinition{
1000+
Name: "add",
1001+
Description: "Add two numbers",
1002+
Parameters: funcParams,
1003+
},
1004+
},
1005+
},
1006+
parametersJSONSchemaAvailable: true,
9661007
expected: []genai.Tool{
9671008
{
9681009
FunctionDeclarations: []*genai.FunctionDeclaration{
@@ -976,44 +1017,96 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) {
9761017
},
9771018
},
9781019
{
979-
name: "multiple function tools",
1020+
name: "multiple function tools with nil/empty parameters - parametersJSONSchemaAvailable=false",
9801021
openaiTools: []openai.Tool{
9811022
{
9821023
Type: openai.ToolTypeFunction,
9831024
Function: &openai.FunctionDefinition{
9841025
Name: "foo",
9851026
Description: "Foo function",
1027+
Parameters: map[string]any{}, // empty parameters
9861028
},
9871029
},
9881030
{
9891031
Type: openai.ToolTypeFunction,
9901032
Function: &openai.FunctionDefinition{
9911033
Name: "bar",
9921034
Description: "Bar function",
1035+
Parameters: nil, // nil parameters
9931036
},
9941037
},
9951038
},
1039+
parametersJSONSchemaAvailable: false,
9961040
expected: []genai.Tool{
9971041
{
9981042
FunctionDeclarations: []*genai.FunctionDeclaration{
9991043
{
1000-
Name: "foo",
1001-
Description: "Foo function",
1002-
Parameters: nil,
1003-
ParametersJsonSchema: nil,
1044+
Name: "foo",
1045+
Description: "Foo function",
1046+
Parameters: nil,
10041047
},
10051048
{
1006-
Name: "bar",
1007-
Description: "Bar function",
1008-
Parameters: nil,
1009-
ParametersJsonSchema: nil,
1049+
Name: "bar",
1050+
Description: "Bar function",
1051+
Parameters: nil,
10101052
},
10111053
},
10121054
},
10131055
},
10141056
},
10151057
{
1016-
name: "tool with invalid parameters schema",
1058+
name: "multiple function tools with nil/empty parameters - parametersJSONSchemaAvailable=true",
1059+
openaiTools: []openai.Tool{
1060+
{
1061+
Type: openai.ToolTypeFunction,
1062+
Function: &openai.FunctionDefinition{
1063+
Name: "foo",
1064+
Description: "Foo function",
1065+
},
1066+
},
1067+
{
1068+
Type: openai.ToolTypeFunction,
1069+
Function: &openai.FunctionDefinition{
1070+
Name: "bar",
1071+
Description: "Bar function",
1072+
},
1073+
},
1074+
},
1075+
parametersJSONSchemaAvailable: true,
1076+
expected: []genai.Tool{
1077+
{
1078+
FunctionDeclarations: []*genai.FunctionDeclaration{
1079+
{
1080+
Name: "foo",
1081+
Description: "Foo function",
1082+
ResponseJsonSchema: nil,
1083+
},
1084+
{
1085+
Name: "bar",
1086+
Description: "Bar function",
1087+
ResponseJsonSchema: nil,
1088+
},
1089+
},
1090+
},
1091+
},
1092+
},
1093+
{
1094+
name: "tool with invalid parameters schema - parametersJSONSchemaAvailable=false",
1095+
openaiTools: []openai.Tool{
1096+
{
1097+
Type: openai.ToolTypeFunction,
1098+
Function: &openai.FunctionDefinition{
1099+
Name: "bad",
1100+
Description: "Bad function",
1101+
Parameters: "invalid-json",
1102+
},
1103+
},
1104+
},
1105+
parametersJSONSchemaAvailable: false,
1106+
expectedError: "invalid JSON schema for parameters in tool bad: expected map[string]any",
1107+
},
1108+
{
1109+
name: "tool with invalid parameters schema - parametersJSONSchemaAvailable=true",
10171110
openaiTools: []openai.Tool{
10181111
{
10191112
Type: openai.ToolTypeFunction,
@@ -1024,32 +1117,140 @@ func TestOpenAIToolsToGeminiTools(t *testing.T) {
10241117
},
10251118
},
10261119
},
1120+
parametersJSONSchemaAvailable: true,
10271121
expected: []genai.Tool{
10281122
{
10291123
FunctionDeclarations: []*genai.FunctionDeclaration{
10301124
{
1031-
Description: "Bad function",
1032-
Name: "bad",
1125+
Description: "Bad function",
1126+
Name: "bad",
1127+
// ai-gateway does not validate schema, upstream will be expected to validate the bad json param
10331128
ParametersJsonSchema: "invalid-json",
10341129
},
10351130
},
10361131
},
10371132
},
10381133
},
10391134
{
1040-
name: "non-function tool is ignored",
1135+
name: "complex nested schema - parametersJSONSchemaAvailable=false",
1136+
openaiTools: []openai.Tool{
1137+
{
1138+
Type: openai.ToolTypeFunction,
1139+
Function: &openai.FunctionDefinition{
1140+
Name: "complex_tool",
1141+
Description: "Complex tool with nested parameters",
1142+
Parameters: map[string]any{
1143+
"type": "object",
1144+
"properties": map[string]any{
1145+
"user": map[string]any{
1146+
"type": "object",
1147+
"properties": map[string]any{
1148+
"name": map[string]any{"type": "string"},
1149+
"age": map[string]any{"type": "integer"},
1150+
},
1151+
"required": []any{"name"},
1152+
},
1153+
"items": map[string]any{
1154+
"type": "array",
1155+
"items": map[string]any{
1156+
"type": "object",
1157+
"properties": map[string]any{
1158+
"id": map[string]any{"type": "integer"},
1159+
"name": map[string]any{"type": "string"},
1160+
},
1161+
},
1162+
},
1163+
},
1164+
"required": []any{"user"},
1165+
},
1166+
},
1167+
},
1168+
},
1169+
parametersJSONSchemaAvailable: false,
1170+
expected: []genai.Tool{
1171+
{
1172+
FunctionDeclarations: []*genai.FunctionDeclaration{
1173+
{
1174+
Name: "complex_tool",
1175+
Description: "Complex tool with nested parameters",
1176+
Parameters: &genai.Schema{
1177+
Type: "object",
1178+
Properties: map[string]*genai.Schema{
1179+
"user": {
1180+
Type: "object",
1181+
Properties: map[string]*genai.Schema{
1182+
"name": {Type: "string"},
1183+
"age": {Type: "integer"},
1184+
},
1185+
Required: []string{"name"},
1186+
},
1187+
"items": {
1188+
Type: "array",
1189+
Items: &genai.Schema{
1190+
Type: "object",
1191+
Properties: map[string]*genai.Schema{
1192+
"id": {Type: "integer"},
1193+
"name": {Type: "string"},
1194+
},
1195+
},
1196+
},
1197+
},
1198+
Required: []string{"user"},
1199+
},
1200+
},
1201+
},
1202+
},
1203+
},
1204+
},
1205+
{
1206+
name: "non-function tool is ignored - both modes",
10411207
openaiTools: []openai.Tool{
10421208
{
10431209
Type: "retrieval",
10441210
},
10451211
},
1046-
expected: nil,
1212+
parametersJSONSchemaAvailable: false,
1213+
expectedError: "unsupported tool type: retrieval",
1214+
},
1215+
{
1216+
name: "mixed valid and invalid tools - parametersJSONSchemaAvailable=false",
1217+
openaiTools: []openai.Tool{
1218+
{
1219+
Type: "retrieval", // Should be ignored
1220+
},
1221+
{
1222+
Type: openai.ToolTypeFunction,
1223+
Function: &openai.FunctionDefinition{
1224+
Name: "valid_tool",
1225+
Description: "Valid function tool",
1226+
Parameters: map[string]any{
1227+
"type": "object",
1228+
"properties": map[string]any{
1229+
"param": map[string]any{"type": "string"},
1230+
},
1231+
},
1232+
},
1233+
},
1234+
},
1235+
parametersJSONSchemaAvailable: false,
1236+
expectedError: "unsupported tool type: retrieval",
1237+
},
1238+
{
1239+
name: "tool with nil function - should not panic",
1240+
openaiTools: []openai.Tool{
1241+
{
1242+
Type: openai.ToolTypeFunction,
1243+
Function: nil,
1244+
},
1245+
},
1246+
parametersJSONSchemaAvailable: false,
1247+
expected: nil,
10471248
},
10481249
}
10491250

10501251
for _, tc := range tests {
10511252
t.Run(tc.name, func(t *testing.T) {
1052-
result, err := openAIToolsToGeminiTools(tc.openaiTools)
1253+
result, err := openAIToolsToGeminiTools(tc.openaiTools, tc.parametersJSONSchemaAvailable)
10531254
if tc.expectedError != "" {
10541255
require.ErrorContains(t, err, tc.expectedError)
10551256
} else {

‎internal/extproc/translator/openai_gcpvertexai.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -272,8 +272,10 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) openAIMessageToGeminiMes
272272
return nil, err
273273
}
274274

275+
// Some models support only partialJSONSchema.
276+
parametersJSONSchemaAvailable := responseJSONSchemaAvailable(requestModel)
275277
// Convert OpenAI tools to Gemini tools.
276-
tools, err := openAIToolsToGeminiTools(openAIReq.Tools)
278+
tools, err := openAIToolsToGeminiTools(openAIReq.Tools, parametersJSONSchemaAvailable)
277279
if err != nil {
278280
return nil, fmt.Errorf("error converting tools: %w", err)
279281
}

0 commit comments

Comments
 (0)