Просмотр исходного кода

fix: openAI function to gemini function field adjusted to whitelist mode

Seefs 1 месяц назад
Родитель
Сommit
ea802f2297
1 измененных файлов с 138 добавлено и 71 удалено
  1. 138 71
      relay/channel/gemini/relay-gemini.go

+ 138 - 71
relay/channel/gemini/relay-gemini.go

@@ -655,102 +655,84 @@ func getSupportedMimeTypesList() []string {
 	return keys
 }
 
+var geminiOpenAPISchemaAllowedFields = 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":             {},
+}
+
+const geminiFunctionSchemaMaxDepth = 64
+
 // cleanFunctionParameters recursively removes unsupported fields from Gemini function parameters.
 func cleanFunctionParameters(params interface{}) interface{} {
+	return cleanFunctionParametersWithDepth(params, 0)
+}
+
+func cleanFunctionParametersWithDepth(params interface{}, depth int) interface{} {
 	if params == nil {
 		return nil
 	}
 
+	if depth >= geminiFunctionSchemaMaxDepth {
+		return cleanFunctionParametersShallow(params)
+	}
+
 	switch v := params.(type) {
 	case map[string]interface{}:
-		// Create a copy to avoid modifying the original
-		cleanedMap := make(map[string]interface{})
+		// Keep only Gemini-supported OpenAPI schema subset fields (per official SDK Schema).
+		cleanedMap := make(map[string]interface{}, len(v))
 		for k, val := range v {
-			cleanedMap[k] = val
-		}
-
-		// Remove unsupported root-level fields
-		delete(cleanedMap, "default")
-		delete(cleanedMap, "exclusiveMaximum")
-		delete(cleanedMap, "exclusiveMinimum")
-		delete(cleanedMap, "$schema")
-		delete(cleanedMap, "additionalProperties")
-		delete(cleanedMap, "propertyNames")
-
-		// Check and clean 'format' for string types
-		if propType, typeExists := cleanedMap["type"].(string); typeExists && propType == "string" {
-			if formatValue, formatExists := cleanedMap["format"].(string); formatExists {
-				if formatValue != "enum" && formatValue != "date-time" {
-					delete(cleanedMap, "format")
-				}
+			if _, ok := geminiOpenAPISchemaAllowedFields[k]; ok {
+				cleanedMap[k] = val
 			}
 		}
 
+		normalizeGeminiSchemaTypeAndNullable(cleanedMap)
+
 		// Clean properties
 		if props, ok := cleanedMap["properties"].(map[string]interface{}); ok && props != nil {
 			cleanedProps := make(map[string]interface{})
 			for propName, propValue := range props {
-				cleanedProps[propName] = cleanFunctionParameters(propValue)
+				cleanedProps[propName] = cleanFunctionParametersWithDepth(propValue, depth+1)
 			}
 			cleanedMap["properties"] = cleanedProps
 		}
 
 		// Recursively clean items in arrays
 		if items, ok := cleanedMap["items"].(map[string]interface{}); ok && items != nil {
-			cleanedMap["items"] = cleanFunctionParameters(items)
-		}
-		// Also handle items if it's an array of schemas
-		if itemsArray, ok := cleanedMap["items"].([]interface{}); ok {
-			cleanedItemsArray := make([]interface{}, len(itemsArray))
-			for i, item := range itemsArray {
-				cleanedItemsArray[i] = cleanFunctionParameters(item)
-			}
-			cleanedMap["items"] = cleanedItemsArray
-		}
-
-		// Recursively clean other schema composition keywords
-		for _, field := range []string{"allOf", "anyOf", "oneOf"} {
-			if nested, ok := cleanedMap[field].([]interface{}); ok {
-				cleanedNested := make([]interface{}, len(nested))
-				for i, item := range nested {
-					cleanedNested[i] = cleanFunctionParameters(item)
-				}
-				cleanedMap[field] = cleanedNested
-			}
-		}
-
-		// Recursively clean patternProperties
-		if patternProps, ok := cleanedMap["patternProperties"].(map[string]interface{}); ok {
-			cleanedPatternProps := make(map[string]interface{})
-			for pattern, schema := range patternProps {
-				cleanedPatternProps[pattern] = cleanFunctionParameters(schema)
-			}
-			cleanedMap["patternProperties"] = cleanedPatternProps
-		}
-
-		// Recursively clean definitions
-		if definitions, ok := cleanedMap["definitions"].(map[string]interface{}); ok {
-			cleanedDefinitions := make(map[string]interface{})
-			for defName, defSchema := range definitions {
-				cleanedDefinitions[defName] = cleanFunctionParameters(defSchema)
-			}
-			cleanedMap["definitions"] = cleanedDefinitions
+			cleanedMap["items"] = cleanFunctionParametersWithDepth(items, depth+1)
 		}
-
-		// Recursively clean $defs (newer JSON Schema draft)
-		if defs, ok := cleanedMap["$defs"].(map[string]interface{}); ok {
-			cleanedDefs := make(map[string]interface{})
-			for defName, defSchema := range defs {
-				cleanedDefs[defName] = cleanFunctionParameters(defSchema)
-			}
-			cleanedMap["$defs"] = cleanedDefs
+		// OpenAPI tuple-style items is not supported by Gemini SDK Schema; keep first to avoid API rejection.
+		if itemsArray, ok := cleanedMap["items"].([]interface{}); ok && len(itemsArray) > 0 {
+			cleanedMap["items"] = cleanFunctionParametersWithDepth(itemsArray[0], depth+1)
 		}
 
-		// Clean conditional keywords
-		for _, field := range []string{"if", "then", "else", "not"} {
-			if nested, ok := cleanedMap[field]; ok {
-				cleanedMap[field] = cleanFunctionParameters(nested)
+		// Recursively clean anyOf
+		if nested, ok := cleanedMap["anyOf"].([]interface{}); ok && nested != nil {
+			cleanedNested := make([]interface{}, len(nested))
+			for i, item := range nested {
+				cleanedNested[i] = cleanFunctionParametersWithDepth(item, depth+1)
 			}
+			cleanedMap["anyOf"] = cleanedNested
 		}
 
 		return cleanedMap
@@ -759,7 +741,7 @@ func cleanFunctionParameters(params interface{}) interface{} {
 		// Handle arrays of schemas
 		cleanedArray := make([]interface{}, len(v))
 		for i, item := range v {
-			cleanedArray[i] = cleanFunctionParameters(item)
+			cleanedArray[i] = cleanFunctionParametersWithDepth(item, depth+1)
 		}
 		return cleanedArray
 
@@ -769,6 +751,91 @@ func cleanFunctionParameters(params interface{}) interface{} {
 	}
 }
 
+func cleanFunctionParametersShallow(params interface{}) interface{} {
+	switch v := params.(type) {
+	case map[string]interface{}:
+		cleanedMap := make(map[string]interface{}, len(v))
+		for k, val := range v {
+			if _, ok := geminiOpenAPISchemaAllowedFields[k]; ok {
+				cleanedMap[k] = val
+			}
+		}
+		normalizeGeminiSchemaTypeAndNullable(cleanedMap)
+		// Stop recursion and avoid retaining huge nested structures.
+		delete(cleanedMap, "properties")
+		delete(cleanedMap, "items")
+		delete(cleanedMap, "anyOf")
+		return cleanedMap
+	case []interface{}:
+		// Prefer an empty list over deep recursion on attacker-controlled inputs.
+		return []interface{}{}
+	default:
+		return params
+	}
+}
+
+func normalizeGeminiSchemaTypeAndNullable(schema map[string]interface{}) {
+	rawType, ok := schema["type"]
+	if !ok || rawType == nil {
+		return
+	}
+
+	normalize := func(t string) (string, bool) {
+		switch strings.ToLower(strings.TrimSpace(t)) {
+		case "object":
+			return "OBJECT", false
+		case "array":
+			return "ARRAY", false
+		case "string":
+			return "STRING", false
+		case "integer":
+			return "INTEGER", false
+		case "number":
+			return "NUMBER", false
+		case "boolean":
+			return "BOOLEAN", false
+		case "null":
+			return "", true
+		default:
+			return t, false
+		}
+	}
+
+	switch t := rawType.(type) {
+	case string:
+		normalized, isNull := normalize(t)
+		if isNull {
+			schema["nullable"] = true
+			delete(schema, "type")
+			return
+		}
+		schema["type"] = normalized
+	case []interface{}:
+		nullable := false
+		var chosen string
+		for _, item := range t {
+			if s, ok := item.(string); ok {
+				normalized, isNull := normalize(s)
+				if isNull {
+					nullable = true
+					continue
+				}
+				if chosen == "" {
+					chosen = normalized
+				}
+			}
+		}
+		if nullable {
+			schema["nullable"] = true
+		}
+		if chosen != "" {
+			schema["type"] = chosen
+		} else {
+			delete(schema, "type")
+		}
+	}
+}
+
 func removeAdditionalPropertiesWithDepth(schema interface{}, depth int) interface{} {
 	if depth >= 5 {
 		return schema