import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import type { GoogleGenAI } from '@google/genai';
import { z } from 'zod';
import { TextModel } from '../types.js';
import { formatToolError } from '../utils/errors.js';
export function register(server: McpServer, ai: GoogleGenAI): void {
server.registerTool(
'generate_with_search',
{
title: 'Generate with Search',
description: 'Generate text with Google Search grounding for up-to-date, cited responses.',
inputSchema: {
prompt: z.string().min(1).describe('The prompt to search and answer'),
model: TextModel.default('gemini-2.5-flash').describe('Gemini model to use'),
temperature: z.number().min(0).max(2).optional().describe('Sampling temperature'),
maxOutputTokens: z.number().min(1).optional().describe('Maximum output tokens'),
},
annotations: {
readOnlyHint: true,
destructiveHint: false,
openWorldHint: true,
},
},
async ({ prompt, model, temperature, maxOutputTokens }) => {
try {
const response = await ai.models.generateContent({
model,
contents: prompt,
config: {
temperature,
maxOutputTokens,
tools: [{ googleSearch: {} }],
},
});
return { content: [{ type: 'text' as const, text: response.text ?? '' }] };
} catch (error) {
return formatToolError(error);
}
},
);
}
