retrieve_from_aws_kb

Instructions

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: [] };
    }
  }