refactor(MCP): enhance schema validation for gemini (#4153)
This commit is contained in:
SuYao
@@ -1,5 +1,17 @@
|
||||
import { Tool, ToolUnion, ToolUseBlock } from '@anthropic-ai/sdk/resources'
|
||||
import { FunctionCall, FunctionDeclaration, SchemaType, Tool as geminiToool } from '@google/generative-ai'
|
||||
import { FunctionCall, FunctionDeclaration, SchemaType, Tool as geminiTool } from '@google/generative-ai'
|
||||
import {
|
||||
ArraySchema,
|
||||
BaseSchema,
|
||||
BooleanSchema,
|
||||
EnumStringSchema,
|
||||
FunctionDeclarationSchema,
|
||||
FunctionDeclarationSchemaProperty,
|
||||
IntegerSchema,
|
||||
NumberSchema,
|
||||
ObjectSchema,
|
||||
SimpleStringSchema
|
||||
} from '@google/generative-ai'
|
||||
import { nanoid } from '@reduxjs/toolkit'
|
||||
import store from '@renderer/store'
|
||||
import { addMCPServer } from '@renderer/store/mcp'
|
||||
@@ -8,54 +20,143 @@ import { ChatCompletionMessageToolCall, ChatCompletionTool } from 'openai/resour
|
||||
|
||||
import { ChunkCallbackData } from '../providers'
|
||||
|
||||
const supportedAttributes = [
|
||||
'type',
|
||||
'nullable',
|
||||
'required',
|
||||
// 'format',
|
||||
'description',
|
||||
'properties',
|
||||
'items',
|
||||
'enum',
|
||||
'anyOf'
|
||||
]
|
||||
const ensureValidSchema = (obj: Record<string, any>): FunctionDeclarationSchemaProperty => {
|
||||
// Filter out unsupported keys for Gemini
|
||||
const filteredObj = filterUnsupportedKeys(obj)
|
||||
|
||||
function filterPropertieAttributes(tool: MCPTool, filterNestedObj = false) {
|
||||
// Handle base schema properties
|
||||
const baseSchema = {
|
||||
description: filteredObj.description,
|
||||
nullable: filteredObj.nullable
|
||||
} as BaseSchema
|
||||
|
||||
// Handle string type
|
||||
if (filteredObj.type?.toLowerCase() === SchemaType.STRING) {
|
||||
if (filteredObj.enum && Array.isArray(filteredObj.enum)) {
|
||||
return {
|
||||
...baseSchema,
|
||||
type: SchemaType.STRING,
|
||||
format: 'enum',
|
||||
enum: filteredObj.enum as string[]
|
||||
} as EnumStringSchema
|
||||
}
|
||||
return {
|
||||
...baseSchema,
|
||||
type: SchemaType.STRING,
|
||||
format: filteredObj.format === 'date-time' ? 'date-time' : undefined
|
||||
} as SimpleStringSchema
|
||||
}
|
||||
|
||||
// Handle number type
|
||||
if (filteredObj.type?.toLowerCase() === SchemaType.NUMBER) {
|
||||
return {
|
||||
...baseSchema,
|
||||
type: SchemaType.NUMBER,
|
||||
format: ['float', 'double'].includes(filteredObj.format) ? (filteredObj.format as 'float' | 'double') : undefined
|
||||
} as NumberSchema
|
||||
}
|
||||
|
||||
// Handle integer type
|
||||
if (filteredObj.type?.toLowerCase() === SchemaType.INTEGER) {
|
||||
return {
|
||||
...baseSchema,
|
||||
type: SchemaType.INTEGER,
|
||||
format: ['int32', 'int64'].includes(filteredObj.format) ? (filteredObj.format as 'int32' | 'int64') : undefined
|
||||
} as IntegerSchema
|
||||
}
|
||||
|
||||
// Handle boolean type
|
||||
if (filteredObj.type?.toLowerCase() === SchemaType.BOOLEAN) {
|
||||
return {
|
||||
...baseSchema,
|
||||
type: SchemaType.BOOLEAN
|
||||
} as BooleanSchema
|
||||
}
|
||||
|
||||
// Handle array type
|
||||
if (filteredObj.type?.toLowerCase() === SchemaType.ARRAY) {
|
||||
return {
|
||||
...baseSchema,
|
||||
type: SchemaType.ARRAY,
|
||||
items: filteredObj.items
|
||||
? ensureValidSchema(filteredObj.items as Record<string, any>)
|
||||
: ({ type: SchemaType.STRING } as SimpleStringSchema),
|
||||
minItems: filteredObj.minItems,
|
||||
maxItems: filteredObj.maxItems
|
||||
} as ArraySchema
|
||||
}
|
||||
|
||||
// Handle object type (default)
|
||||
const properties = filteredObj.properties
|
||||
? Object.fromEntries(
|
||||
Object.entries(filteredObj.properties).map(([key, value]) => [
|
||||
key,
|
||||
ensureValidSchema(value as Record<string, any>)
|
||||
])
|
||||
)
|
||||
: { _empty: { type: SchemaType.STRING } as SimpleStringSchema } // Ensure properties is never empty
|
||||
|
||||
return {
|
||||
...baseSchema,
|
||||
type: SchemaType.OBJECT,
|
||||
properties,
|
||||
required: Array.isArray(filteredObj.required) ? filteredObj.required : undefined
|
||||
} as ObjectSchema
|
||||
}
|
||||
|
||||
function filterUnsupportedKeys(obj: Record<string, any>): Record<string, any> {
|
||||
const supportedBaseKeys = ['description', 'nullable']
|
||||
const supportedStringKeys = [...supportedBaseKeys, 'type', 'format', 'enum']
|
||||
const supportedNumberKeys = [...supportedBaseKeys, 'type', 'format']
|
||||
const supportedBooleanKeys = [...supportedBaseKeys, 'type']
|
||||
const supportedArrayKeys = [...supportedBaseKeys, 'type', 'items', 'minItems', 'maxItems']
|
||||
const supportedObjectKeys = [...supportedBaseKeys, 'type', 'properties', 'required']
|
||||
|
||||
const filtered: Record<string, any> = {}
|
||||
|
||||
let keysToKeep: string[]
|
||||
|
||||
if (obj.type?.toLowerCase() === SchemaType.STRING) {
|
||||
keysToKeep = supportedStringKeys
|
||||
} else if (obj.type?.toLowerCase() === SchemaType.NUMBER) {
|
||||
keysToKeep = supportedNumberKeys
|
||||
} else if (obj.type?.toLowerCase() === SchemaType.INTEGER) {
|
||||
keysToKeep = supportedNumberKeys
|
||||
} else if (obj.type?.toLowerCase() === SchemaType.BOOLEAN) {
|
||||
keysToKeep = supportedBooleanKeys
|
||||
} else if (obj.type?.toLowerCase() === SchemaType.ARRAY) {
|
||||
keysToKeep = supportedArrayKeys
|
||||
} else {
|
||||
// Default to object type
|
||||
keysToKeep = supportedObjectKeys
|
||||
}
|
||||
|
||||
// copy supported keys
|
||||
for (const key of keysToKeep) {
|
||||
if (obj[key] !== undefined) {
|
||||
filtered[key] = obj[key]
|
||||
}
|
||||
}
|
||||
|
||||
return filtered
|
||||
}
|
||||
|
||||
function filterPropertieAttributes(tool: MCPTool, filterNestedObj: boolean = false): Record<string, object> {
|
||||
const properties = tool.inputSchema.properties
|
||||
if (!properties) {
|
||||
return {}
|
||||
}
|
||||
const getSubMap = (obj: Record<string, any>, keys: string[]) => {
|
||||
const filtered = Object.fromEntries(Object.entries(obj).filter(([key]) => keys.includes(key)))
|
||||
|
||||
if (filterNestedObj) {
|
||||
return {
|
||||
...filtered,
|
||||
...(obj.type === 'object' && obj.properties
|
||||
? {
|
||||
properties: Object.fromEntries(
|
||||
Object.entries(obj.properties).map(([k, v]) => [
|
||||
k,
|
||||
(v as any).type === 'object' ? getSubMap(v as Record<string, any>, keys) : v
|
||||
])
|
||||
)
|
||||
}
|
||||
: {}),
|
||||
...(obj.type === 'array' && obj.items?.type === 'object'
|
||||
? {
|
||||
items: getSubMap(obj.items, keys)
|
||||
}
|
||||
: {})
|
||||
}
|
||||
}
|
||||
|
||||
return filtered
|
||||
// For OpenAI, we don't need to validate as strictly
|
||||
if (!filterNestedObj) {
|
||||
return properties
|
||||
}
|
||||
|
||||
for (const [key, val] of Object.entries(properties)) {
|
||||
properties[key] = getSubMap(val, supportedAttributes)
|
||||
}
|
||||
return properties
|
||||
const processedProperties = Object.fromEntries(
|
||||
Object.entries(properties).map(([key, value]) => [key, ensureValidSchema(value as Record<string, any>)])
|
||||
)
|
||||
|
||||
return processedProperties
|
||||
}
|
||||
|
||||
export function mcpToolsToOpenAITools(mcpTools: MCPTool[]): Array<ChatCompletionTool> {
|
||||
@@ -132,7 +233,7 @@ export async function callMCPTool(tool: MCPTool): Promise<any> {
|
||||
if (tool.serverName === 'mcp-auto-install') {
|
||||
if (resp.data) {
|
||||
const mcpServer: MCPServer = {
|
||||
id: nanoid(),
|
||||
id: `f${nanoid()}`,
|
||||
name: resp.data.name,
|
||||
description: resp.data.description,
|
||||
baseUrl: resp.data.baseUrl,
|
||||
@@ -183,7 +284,7 @@ export function anthropicToolUseToMcpTool(mcpTools: MCPTool[] | undefined, toolU
|
||||
return tool
|
||||
}
|
||||
|
||||
export function mcpToolsToGeminiTools(mcpTools: MCPTool[] | undefined): geminiToool[] {
|
||||
export function mcpToolsToGeminiTools(mcpTools: MCPTool[] | undefined): geminiTool[] {
|
||||
if (!mcpTools || mcpTools.length === 0) {
|
||||
// No tools available
|
||||
return []
|
||||
@@ -195,18 +296,19 @@ export function mcpToolsToGeminiTools(mcpTools: MCPTool[] | undefined): geminiTo
|
||||
const functionDeclaration: FunctionDeclaration = {
|
||||
name: tool.id,
|
||||
description: tool.description,
|
||||
...(Object.keys(properties).length > 0
|
||||
? {
|
||||
parameters: {
|
||||
type: SchemaType.OBJECT,
|
||||
properties
|
||||
}
|
||||
}
|
||||
: {})
|
||||
parameters: {
|
||||
type: SchemaType.OBJECT,
|
||||
properties:
|
||||
Object.keys(properties).length > 0
|
||||
? Object.fromEntries(
|
||||
Object.entries(properties).map(([key, value]) => [key, ensureValidSchema(value as Record<string, any>)])
|
||||
)
|
||||
: { _empty: { type: SchemaType.STRING } as SimpleStringSchema }
|
||||
} as FunctionDeclarationSchema
|
||||
}
|
||||
functions.push(functionDeclaration)
|
||||
}
|
||||
const tool: geminiToool = {
|
||||
const tool: geminiTool = {
|
||||
functionDeclarations: functions
|
||||
}
|
||||
return [tool]
|
||||
|
||||
Reference in New Issue
Block a user