retrieve_from_aws_kb
Retrieve relevant information from AWS Knowledge Bases by querying with specific questions and Knowledge Base ID using RAG via Bedrock Agent Runtime.
Instructions
Performs retrieval from the AWS Knowledge Base using the provided query and Knowledge Base ID.
Input Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| knowledgeBaseId | Yes | The ID of the AWS Knowledge Base | |
| n | No | Number of results to retrieve | |
| query | Yes | The query to perform retrieval on |
Implementation Reference
- index.ts:149-204 (handler)Executes the tool logic for 'retrieve_from_aws_kb': parses arguments, resolves knowledge base ID from args or env, calls retrieveContext helper, formats and returns context text 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:48-106 (helper)Performs the actual AWS Bedrock Knowledge Base retrieval using RetrieveCommand, processes 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: [] }; } }
- index.ts:112-125 (schema)JSON schema defining the tool's input parameters: query (required string), knowledgeBaseId (string, conditional required), n (number, default 3).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:109-126 (registration)Defines the Tool object for 'retrieve_from_aws_kb' with name, description, and inputSchema.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 tool by including it in the response to ListToolsRequest.server.setRequestHandler(ListToolsRequestSchema, async () => ({ tools: [RETRIEVAL_TOOL], }));