From b34948846924bcf334d586ba595062270f7501b9 Mon Sep 17 00:00:00 2001 From: Skirano Date: Mon, 2 Dec 2024 15:11:07 -0500 Subject: [PATCH] added aws kb server --- src/aws-kb-retrieval-server/README.md | 53 +++++++ src/aws-kb-retrieval-server/index.ts | 166 ++++++++++++++++++++++ src/aws-kb-retrieval-server/package.json | 30 ++++ src/aws-kb-retrieval-server/tsconfig.json | 17 +++ 4 files changed, 266 insertions(+) create mode 100644 src/aws-kb-retrieval-server/README.md create mode 100644 src/aws-kb-retrieval-server/index.ts create mode 100644 src/aws-kb-retrieval-server/package.json create mode 100644 src/aws-kb-retrieval-server/tsconfig.json diff --git a/src/aws-kb-retrieval-server/README.md b/src/aws-kb-retrieval-server/README.md new file mode 100644 index 00000000..ac2bdb43 --- /dev/null +++ b/src/aws-kb-retrieval-server/README.md @@ -0,0 +1,53 @@ +# AWS Knowledge Base Retrieval MCP Server + +An MCP server implementation for retrieving information from the AWS Knowledge Base using the Bedrock Agent Runtime. + +## Features + +- **RAG (Retrieval-Augmented Generation)**: Retrieve context from the AWS Knowledge Base based on a query and a Knowledge Base ID. +- **Supports multiple results retrieval**: Option to retrieve a customizable number of results. + +## Tools + +- **retrieve_from_aws_kb** + - Perform retrieval operations using the AWS Knowledge Base. + - Inputs: + - `query` (string): The search query for retrieval. + - `knowledgeBaseId` (string): The ID of the AWS Knowledge Base. + - `n` (number, optional): Number of results to retrieve (default: 3). + +## Configuration + +### Setting up AWS Credentials + +1. Obtain AWS access key ID, secret access key, and region from the AWS Management Console. +2. Ensure these credentials have appropriate permissions for Bedrock Agent Runtime operations. + +### Usage with Claude Desktop + +Add this to your `claude_desktop_config.json`: + +```json +{ + "mcpServers": { + "aws-kb-retrieval": { + "command": "npx", + "args": [ + "-y", + "@modelcontextprotocol/server-aws-kb-retrieval" + ], + "env": { + "AWS_ACCESS_KEY_ID": "YOUR_ACCESS_KEY_HERE", + "AWS_SECRET_ACCESS_KEY": "YOUR_SECRET_ACCESS_KEY_HERE", + "AWS_REGION": "YOUR_AWS_REGION_HERE" + } + } + } +} +``` + +## License + +This MCP server is licensed under the MIT License. This means you are free to use, modify, and distribute the software, subject to the terms and conditions of the MIT License. For more details, please see the LICENSE file in the project repository. + +This README assumes that your server package is named `@modelcontextprotocol/server-aws-kb-retrieval`. Adjust the package name and installation details if they differ in your setup. Also, ensure that your server script is correctly built and that all dependencies are properly managed in your `package.json`. diff --git a/src/aws-kb-retrieval-server/index.ts b/src/aws-kb-retrieval-server/index.ts new file mode 100644 index 00000000..f60a544e --- /dev/null +++ b/src/aws-kb-retrieval-server/index.ts @@ -0,0 +1,166 @@ +#!/usr/bin/env node +import { Server } from "@modelcontextprotocol/sdk/server/index.js"; +import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; +import { + CallToolRequestSchema, + ListToolsRequestSchema, + Tool, +} from "@modelcontextprotocol/sdk/types.js"; +import { + BedrockAgentRuntimeClient, + RetrieveCommand, + RetrieveCommandInput, +} from "@aws-sdk/client-bedrock-agent-runtime"; + +// AWS client initialization +const bedrockClient = new BedrockAgentRuntimeClient({ + region: process.env.AWS_REGION, + credentials: { + accessKeyId: process.env.AWS_ACCESS_KEY_ID!, + secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY!, + }, +}); + +interface RAGSource { + id: string; + fileName: string; + snippet: string; + score: number; +} + +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: [] }; + } +} + +// Define the retrieval tool +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: "The ID of the AWS Knowledge Base" }, + n: { type: "number", default: 3, description: "Number of results to retrieve" }, + }, + required: ["query", "knowledgeBaseId"], + }, +}; + +// Server setup +const server = new Server( + { + name: "aws-kb-retrieval-server", + version: "0.2.0", + }, + { + capabilities: { + tools: {}, + }, + }, +); + +// Request handlers +server.setRequestHandler(ListToolsRequestSchema, async () => ({ + tools: [RETRIEVAL_TOOL], +})); + +server.setRequestHandler(CallToolRequestSchema, async (request) => { + const { name, arguments: args } = request.params; + + if (name === "retrieve_from_aws_kb") { + const { query, knowledgeBaseId, n = 3 } = args as Record; + try { + const result = await retrieveContext(query, knowledgeBaseId, n); + if (result.isRagWorking) { + return { + content: [ + { type: "text", text: `Context: ${result.context}` }, + { type: "text", text: `RAG Sources: ${JSON.stringify(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 { + return { + content: [{ type: "text", text: `Unknown tool: ${name}` }], + isError: true, + }; + } +}); + +// Server startup +async function runServer() { + const transport = new StdioServerTransport(); + await server.connect(transport); + console.error("AWS KB Retrieval Server running on stdio"); +} + +runServer().catch((error) => { + console.error("Fatal error running server:", error); + process.exit(1); +}); diff --git a/src/aws-kb-retrieval-server/package.json b/src/aws-kb-retrieval-server/package.json new file mode 100644 index 00000000..39ba7bd4 --- /dev/null +++ b/src/aws-kb-retrieval-server/package.json @@ -0,0 +1,30 @@ +{ + "name": "@modelcontextprotocol/server-aws-kb-retrieval", + "version": "0.1.0", + "description": "MCP server for AWS Knowledge Base retrieval using Bedrock Agent Runtime", + "license": "MIT", + "author": "Anthropic, PBC (https://anthropic.com)", + "homepage": "https://modelcontextprotocol.io", + "bugs": "https://github.com/modelcontextprotocol/servers/issues", + "type": "module", + "bin": { + "mcp-server-aws-kb-retrieval": "dist/index.js" + }, + "files": [ + "dist" + ], + "scripts": { + "build": "tsc && shx chmod +x dist/*.js", + "prepare": "npm run build", + "watch": "tsc --watch" + }, + "dependencies": { + "@modelcontextprotocol/sdk": "0.5.0", + "@aws-sdk/client-bedrock-agent-runtime": "^3.0.0" + }, + "devDependencies": { + "@types/node": "^20.10.0", + "shx": "^0.3.4", + "typescript": "^5.6.2" + } +} diff --git a/src/aws-kb-retrieval-server/tsconfig.json b/src/aws-kb-retrieval-server/tsconfig.json new file mode 100644 index 00000000..98b13da0 --- /dev/null +++ b/src/aws-kb-retrieval-server/tsconfig.json @@ -0,0 +1,17 @@ +{ + "extends": "../../tsconfig.json", + "compilerOptions": { + "outDir": "./dist", + "rootDir": ".", + "composite": true, + "incremental": true, + "tsBuildInfoFile": "./dist/.tsbuildinfo" + }, + "include": [ + "./**/*.ts" + ], + "exclude": [ + "node_modules", + "dist" + ] +}