query
Search uploaded documents using RAG to find answers with citations. Ask questions to retrieve information from your knowledge base.
Instructions
Query the FileSearchStore using RAG (Retrieval-Augmented Generation) to get answers based on uploaded documents. The AI will search through the documents and provide relevant answers with citations.
Input Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| query | Yes | The question or query to search for in the knowledge base |
Implementation Reference
- The execute method that implements the core logic of the 'query' tool: ensures the FileSearchStore exists, performs RAG query using Gemini client, and returns the result with citations.
async execute(args: QueryArgs): Promise<MCPToolResponse<QueryResult>> { const { geminiClient, storeDisplayName, defaultModel } = this.context; // Ensure store exists const store = await geminiClient.ensureStore(storeDisplayName); // Query the store using the default model from environment variable const result = await geminiClient.queryStore({ storeName: store.name, query: args.query, model: defaultModel, }); return { success: true, message: `Query completed successfully. Found ${String(result.citations.length)} citation(s).`, data: { text: result.text, citations: result.citations, query: args.query, model: defaultModel, storeName: store.name, }, }; } - Type definitions for input args and output result, the tool name 'query', description, and Zod input schema for validation.
type QueryArgs = { query: string; }; type QueryResult = { text: string; citations: string[]; query: string; model: string; storeName: string; }; export class QueryTool extends BaseTool<QueryArgs> { readonly name = "query"; readonly description = "Query the FileSearchStore using RAG (Retrieval-Augmented Generation) to get answers based on uploaded documents. The AI will search through the documents and provide relevant answers with citations."; getInputSchema() { return z.object({ query: z .string() .min(1) .describe("The question or query to search for in the knowledge base"), }); } - src/server/tool-registry.ts:24-37 (registration)Initializes the ToolRegistry by instantiating the QueryTool (and others) and storing instances in a map keyed by tool name.
initialize(context: ToolContext): void { // Manual tool registration for safety and explicit review const tools: Tool[] = [ new UploadFileTool(context), new UploadContentTool(context), new QueryTool(context), ]; for (const tool of tools) { this.toolInstances.set(tool.name, tool); } console.log(`✅ ToolRegistry initialized with ${String(this.toolInstances.size)} tools`); } - src/server/tool-registry.ts:43-57 (registration)Registers all tools, including 'query', with the MCP server using the tool's name, description, input schema, and bound handler function.
setupToolHandlers(): void { for (const tool of this.toolInstances.values()) { // Pass Zod schema directly to MCP SDK // SDK handles JSON Schema conversion internally for both stdio and HTTP transports this.server.registerTool( tool.name, { description: tool.description, inputSchema: tool.getInputSchema().shape, }, tool.handler.bind(tool) as never, ); this.registeredTools.push(tool.name); } }