retrieve_from_aws_kb
Retrieve information from AWS Knowledge Bases using queries to access documentation and resources for AWS services and solutions.
Instructions
Performs retrieval from the AWS Knowledge Base using the provided query and Knowledge Base ID.
Input Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| query | Yes | The query to perform retrieval on | |
| knowledgeBaseId | Yes | The ID of the AWS Knowledge Base | |
| n | No | Number of results to retrieve |
Implementation Reference
- index.ts:149-204 (handler)Handles the tool call for 'retrieve_from_aws_kb', resolves knowledge base ID from arguments or environment, invokes the retrieval helper, handles errors, and formats the MCP response with text context and JSON sources.if (name === "retrieve_from_aws_kb") { const { query, knowledgeBaseId, n = 3 } = args as Record<string, any>; // Determine which knowledge base ID to use let actualKnowledgeBaseId = knowledgeBaseId; // If no ID provided but we have configured IDs, use the first one if (!actualKnowledgeBaseId && configuredKnowledgeBaseIds.length > 0) { actualKnowledgeBaseId = configuredKnowledgeBaseIds[0]; console.error(`Using configured knowledge base ID: ${actualKnowledgeBaseId}`); } // If still no ID available, return an error if (!actualKnowledgeBaseId) { return { content: [{ type: "text", text: "No knowledge base ID provided. Either include a knowledgeBaseId in your request or configure AWS_KB_IDS in the environment." }], isError: true, }; } try { const result = await retrieveContext(query, actualKnowledgeBaseId, n); if (result.isRagWorking) { // Format RAG sources for readability const formattedSources = result.ragSources.map((source, index) => { return `Source ${index + 1}: ${source.fileName} (score: ${source.score.toFixed(3)})\n${source.snippet}`; }).join('\n\n'); return { content: [ { type: "text", text: result.context }, { type: "json", json: { ragSources: result.ragSources } } ], }; } else { return { content: [{ type: "text", text: "Retrieval failed or returned no results." }], }; } } catch (error) { return { content: [{ type: "text", text: `Error occurred: ${error}` }], }; } } else {
- index.ts:109-126 (schema)Defines the tool metadata including name, description, input schema with properties for query, optional knowledgeBaseId, and number of results.const RETRIEVAL_TOOL: Tool = { name: "retrieve_from_aws_kb", description: "Performs retrieval from the AWS Knowledge Base using the provided query and Knowledge Base ID.", inputSchema: { type: "object", properties: { query: { type: "string", description: "The query to perform retrieval on" }, knowledgeBaseId: { type: "string", description: configuredKnowledgeBaseIds.length > 0 ? "The ID of the AWS Knowledge Base (optional if configured via AWS_KB_IDS)" : "The ID of the AWS Knowledge Base" }, n: { type: "number", default: 3, description: "Number of results to retrieve" }, }, required: configuredKnowledgeBaseIds.length > 0 ? ["query"] : ["query", "knowledgeBaseId"], }, };
- index.ts:142-144 (registration)Registers the 'retrieve_from_aws_kb' tool by including it in the tools list returned by ListToolsRequestSchema.server.setRequestHandler(ListToolsRequestSchema, async () => ({ tools: [RETRIEVAL_TOOL], }));
- index.ts:48-106 (helper)Core helper function that executes the AWS Bedrock RetrieveCommand, processes retrieval results into context string and RAG sources with metadata.async function retrieveContext( query: string, knowledgeBaseId: string, n: number = 3 ): Promise<{ context: string; isRagWorking: boolean; ragSources: RAGSource[]; }> { try { if (!knowledgeBaseId) { console.error("knowledgeBaseId is not provided"); return { context: "", isRagWorking: false, ragSources: [], }; } const input: RetrieveCommandInput = { knowledgeBaseId: knowledgeBaseId, retrievalQuery: { text: query }, retrievalConfiguration: { vectorSearchConfiguration: { numberOfResults: n }, }, }; const command = new RetrieveCommand(input); const response = await bedrockClient.send(command); const rawResults = response?.retrievalResults || []; const ragSources: RAGSource[] = rawResults .filter((res) => res?.content?.text) .map((result, index) => { const uri = result?.location?.s3Location?.uri || ""; const fileName = uri.split("/").pop() || `Source-${index}.txt`; return { id: (result.metadata?.["x-amz-bedrock-kb-chunk-id"] as string) || `chunk-${index}`, fileName: fileName.replace(/_/g, " ").replace(".txt", ""), snippet: result.content?.text || "", score: (result.score as number) || 0, }; }) .slice(0, 3); const context = rawResults .filter((res): res is { content: { text: string } } => res?.content?.text !== undefined) .map(res => res.content.text) .join("\n\n"); return { context, isRagWorking: true, ragSources, }; } catch (error) { console.error("RAG Error:", error); return { context: "", isRagWorking: false, ragSources: [] }; } }