Sanitize MCP FunctionDeclarations to workaround Vertex bug (#1330)
This commit is contained in:
parent
7b39dd8b28
commit
07880d43d2
|
@ -14,7 +14,8 @@ import {
|
||||||
afterEach,
|
afterEach,
|
||||||
Mocked,
|
Mocked,
|
||||||
} from 'vitest';
|
} from 'vitest';
|
||||||
import { discoverMcpTools } from './mcp-client.js';
|
import { discoverMcpTools, sanatizeParameters } from './mcp-client.js';
|
||||||
|
import { Schema, Type } from '@google/genai';
|
||||||
import { Config, MCPServerConfig } from '../config/config.js';
|
import { Config, MCPServerConfig } from '../config/config.js';
|
||||||
import { DiscoveredMCPTool } from './mcp-tool.js';
|
import { DiscoveredMCPTool } from './mcp-tool.js';
|
||||||
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
|
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
|
||||||
|
@ -536,3 +537,87 @@ describe('discoverMcpTools', () => {
|
||||||
expect(lastClientInstance?.onerror).toEqual(expect.any(Function));
|
expect(lastClientInstance?.onerror).toEqual(expect.any(Function));
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
|
describe('sanatizeParameters', () => {
|
||||||
|
it('should do nothing for an undefined schema', () => {
|
||||||
|
const schema = undefined;
|
||||||
|
sanatizeParameters(schema);
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should remove default when anyOf is present', () => {
|
||||||
|
const schema: Schema = {
|
||||||
|
anyOf: [{ type: Type.STRING }, { type: Type.NUMBER }],
|
||||||
|
default: 'hello',
|
||||||
|
};
|
||||||
|
sanatizeParameters(schema);
|
||||||
|
expect(schema.default).toBeUndefined();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should recursively sanatize items in anyOf', () => {
|
||||||
|
const schema: Schema = {
|
||||||
|
anyOf: [
|
||||||
|
{
|
||||||
|
anyOf: [{ type: Type.STRING }],
|
||||||
|
default: 'world',
|
||||||
|
},
|
||||||
|
{ type: Type.NUMBER },
|
||||||
|
],
|
||||||
|
};
|
||||||
|
sanatizeParameters(schema);
|
||||||
|
expect(schema.anyOf![0].default).toBeUndefined();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should recursively sanatize items in items', () => {
|
||||||
|
const schema: Schema = {
|
||||||
|
items: {
|
||||||
|
anyOf: [{ type: Type.STRING }],
|
||||||
|
default: 'world',
|
||||||
|
},
|
||||||
|
};
|
||||||
|
sanatizeParameters(schema);
|
||||||
|
expect(schema.items!.default).toBeUndefined();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should recursively sanatize items in properties', () => {
|
||||||
|
const schema: Schema = {
|
||||||
|
properties: {
|
||||||
|
prop1: {
|
||||||
|
anyOf: [{ type: Type.STRING }],
|
||||||
|
default: 'world',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
};
|
||||||
|
sanatizeParameters(schema);
|
||||||
|
expect(schema.properties!.prop1.default).toBeUndefined();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should handle complex nested schemas', () => {
|
||||||
|
const schema: Schema = {
|
||||||
|
properties: {
|
||||||
|
prop1: {
|
||||||
|
items: {
|
||||||
|
anyOf: [{ type: Type.STRING }],
|
||||||
|
default: 'world',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
prop2: {
|
||||||
|
anyOf: [
|
||||||
|
{
|
||||||
|
properties: {
|
||||||
|
nestedProp: {
|
||||||
|
anyOf: [{ type: Type.NUMBER }],
|
||||||
|
default: 123,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
};
|
||||||
|
sanatizeParameters(schema);
|
||||||
|
expect(schema.properties!.prop1.items!.default).toBeUndefined();
|
||||||
|
const nestedProp =
|
||||||
|
schema.properties!.prop2.anyOf![0].properties!.nestedProp;
|
||||||
|
expect(nestedProp?.default).toBeUndefined();
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
|
@ -11,7 +11,12 @@ import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/
|
||||||
import { parse } from 'shell-quote';
|
import { parse } from 'shell-quote';
|
||||||
import { MCPServerConfig } from '../config/config.js';
|
import { MCPServerConfig } from '../config/config.js';
|
||||||
import { DiscoveredMCPTool } from './mcp-tool.js';
|
import { DiscoveredMCPTool } from './mcp-tool.js';
|
||||||
import { CallableTool, FunctionDeclaration, mcpToTool } from '@google/genai';
|
import {
|
||||||
|
CallableTool,
|
||||||
|
FunctionDeclaration,
|
||||||
|
mcpToTool,
|
||||||
|
Schema,
|
||||||
|
} from '@google/genai';
|
||||||
import { ToolRegistry } from './tool-registry.js';
|
import { ToolRegistry } from './tool-registry.js';
|
||||||
|
|
||||||
export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
|
export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
|
||||||
|
@ -299,6 +304,8 @@ async function connectAndDiscover(
|
||||||
toolNameForModel.slice(0, 28) + '___' + toolNameForModel.slice(-32);
|
toolNameForModel.slice(0, 28) + '___' + toolNameForModel.slice(-32);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
sanatizeParameters(funcDecl.parameters);
|
||||||
|
|
||||||
// Ensure parameters is a valid JSON schema object, default to empty if not.
|
// Ensure parameters is a valid JSON schema object, default to empty if not.
|
||||||
const parameterSchema: Record<string, unknown> =
|
const parameterSchema: Record<string, unknown> =
|
||||||
funcDecl.parameters && typeof funcDecl.parameters === 'object'
|
funcDecl.parameters && typeof funcDecl.parameters === 'object'
|
||||||
|
@ -354,3 +361,24 @@ async function connectAndDiscover(
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export function sanatizeParameters(schema?: Schema) {
|
||||||
|
if (!schema) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if (schema.anyOf) {
|
||||||
|
// Vertex AI gets confused if both anyOf and default are set.
|
||||||
|
schema.default = undefined;
|
||||||
|
for (const item of schema.anyOf) {
|
||||||
|
sanatizeParameters(item);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (schema.items) {
|
||||||
|
sanatizeParameters(schema.items);
|
||||||
|
}
|
||||||
|
if (schema.properties) {
|
||||||
|
for (const item of Object.values(schema.properties)) {
|
||||||
|
sanatizeParameters(item);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in New Issue