#
tokens: 5175/50000 6/6 files
lines: on (toggle) GitHub
raw markdown copy reset
# Directory Structure

```
├── .gitignore
├── Dockerfile
├── index.ts
├── package.json
├── pnpm-lock.yaml
├── README.md
└── tsconfig.json
```

# Files

--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------

```
1 | build/
2 | *.log
3 | *.pyc
4 | *.swp
5 | node_modules/
6 | dist/
7 | 
```

--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------

```markdown
  1 | # AWS Knowledge Base Retrieval MCP Server
  2 | 
  3 | An MCP server implementation for retrieving information from the AWS Knowledge Base using the Bedrock Agent Runtime.
  4 | 
  5 | ## Features
  6 | 
  7 | - **RAG (Retrieval-Augmented Generation)**: Retrieve context from the AWS Knowledge Base based on a query and a Knowledge Base ID.
  8 | - **Supports multiple results retrieval**: Option to retrieve a customizable number of results.
  9 | 
 10 | ## Tools
 11 | 
 12 | - **retrieve_from_aws_kb**
 13 |   - Perform retrieval operations using the AWS Knowledge Base.
 14 |   - Inputs:
 15 |     - `query` (string): The search query for retrieval.
 16 |     - `knowledgeBaseId` (string): The ID of the AWS Knowledge Base.
 17 |     - `n` (number, optional): Number of results to retrieve (default: 3).
 18 |   - Response format:
 19 |     - The response now returns two separate content items:
 20 |       - A text item containing the raw context from the knowledge base.
 21 |       - A JSON item containing the structured RAG sources with metadata (id, fileName, snippet, and score).
 22 |     - This separation allows for more flexible processing of the results.
 23 | 
 24 | ## Configuration
 25 | 
 26 | ### Setting up AWS Credentials
 27 | 
 28 | You have two options for configuring AWS credentials:
 29 | 
 30 | #### Option 1: IAM Access Keys
 31 | 
 32 | 1. Obtain AWS access key ID, secret access key, and region from the AWS Management Console.
 33 | 2. Ensure these credentials have appropriate permissions for Bedrock Agent Runtime operations.
 34 | 3. Set the environment variables as shown in the configuration examples below.
 35 | 4. For temporary credentials, you can also provide a session token using the `AWS_SESSION_TOKEN` environment variable.
 36 | 
 37 | #### Option 2: AWS SSO (Single Sign-On)
 38 | 
 39 | The server now supports AWS SSO credentials:
 40 | 
 41 | 1. Configure AWS CLI with your SSO profile: `aws configure sso`
 42 | 2. Set only the AWS_REGION environment variable in the MCP server configuration.
 43 | 3. The server will use the default credential provider chain, which includes SSO credentials.
 44 | 
 45 | ### Optional: Configure Default Knowledge Base IDs
 46 | 
 47 | You can optionally specify one or more knowledge base IDs to use by default:
 48 | 
 49 | 1. Create an array of knowledge base IDs in JSON format.
 50 | 2. Set this as the AWS_KB_IDS environment variable in your configuration.
 51 | 3. When this is configured, the `knowledgeBaseId` parameter becomes optional in the tool.
 52 | 
 53 | ### Usage with Claude Desktop
 54 | 
 55 | Add this to your `claude_desktop_config.json`:
 56 | 
 57 | #### Docker with IAM Access Keys
 58 | 
 59 | ```json
 60 | {
 61 |   "mcpServers": {
 62 |     "aws-kb-retrieval": {
 63 |       "command": "docker",
 64 |       "args": [ "run", "-i", "--rm", "-e", "AWS_ACCESS_KEY_ID", "-e", "AWS_SECRET_ACCESS_KEY", "-e", "AWS_REGION", "-e", "AWS_KB_IDS", "mcp/aws-kb-retrieval-server" ],
 65 |       "env": {
 66 |         "AWS_ACCESS_KEY_ID": "YOUR_ACCESS_KEY_HERE",
 67 |         "AWS_SECRET_ACCESS_KEY": "YOUR_SECRET_ACCESS_KEY_HERE",
 68 |         "AWS_SESSION_TOKEN": "YOUR_OPTIONAL_SESSION_ID_FOR_SSO_TEMPORARY_CREDENTIALS_HERE",
 69 |         "AWS_REGION": "YOUR_AWS_REGION_HERE",
 70 |         "AWS_KB_IDS": "[\"kb-12345\", \"kb-67890\"]"
 71 |       }
 72 |     }
 73 |   }
 74 | }
 75 | ```
 76 | 
 77 | #### Docker with AWS SSO
 78 | 
 79 | ```json
 80 | {
 81 |   "mcpServers": {
 82 |     "aws-kb-retrieval": {
 83 |       "command": "docker",
 84 |       "args": [ "run", "-i", "--rm", "-e", "AWS_REGION", "-e", "AWS_KB_IDS", "-v", "${HOME}/.aws:/root/.aws", "mcp/aws-kb-retrieval-server" ],
 85 |       "env": {
 86 |         "AWS_ACCESS_KEY_ID": "YOUR_ACCESS_KEY_HERE",
 87 |         "AWS_SECRET_ACCESS_KEY": "YOUR_SECRET_ACCESS_KEY_HERE",
 88 |         "AWS_SESSION_TOKEN": "YOUR_OPTIONAL_SESSION_ID_FOR_SSO_TEMPORARY_CREDENTIALS_HERE",
 89 |         "AWS_REGION": "YOUR_AWS_REGION_HERE",
 90 |         "AWS_KB_IDS": "[\"kb-12345\", \"kb-67890\"]"
 91 |       }
 92 |     }
 93 |   }
 94 | }
 95 | ```
 96 | 
 97 | #### NPX with IAM Access Keys
 98 | 
 99 | ```json
100 | {
101 |   "mcpServers": {
102 |     "aws-kb-retrieval": {
103 |       "command": "npx",
104 |       "args": [
105 |         "-y",
106 |         "@modelcontextprotocol/server-aws-kb-retrieval"
107 |       ],
108 |       "env": {
109 |         "AWS_ACCESS_KEY_ID": "YOUR_ACCESS_KEY_HERE",
110 |         "AWS_SECRET_ACCESS_KEY": "YOUR_SECRET_ACCESS_KEY_HERE",
111 |         "AWS_SESSION_TOKEN": "YOUR_OPTIONAL_SESSION_ID_FOR_SSO_TEMPORARY_CREDENTIALS_HERE",
112 |         "AWS_REGION": "YOUR_AWS_REGION_HERE",
113 |         "AWS_KB_IDS": "[\"kb-12345\", \"kb-67890\"]"
114 |       }
115 |     }
116 |   }
117 | }
118 | ```
119 | 
120 | #### NPX with AWS SSO
121 | 
122 | ```json
123 | {
124 |   "mcpServers": {
125 |     "aws-kb-retrieval": {
126 |       "command": "npx",
127 |       "args": [
128 |         "-y",
129 |         "@modelcontextprotocol/server-aws-kb-retrieval"
130 |       ],
131 |       "env": {
132 |         "AWS_ACCESS_KEY_ID": "YOUR_ACCESS_KEY_HERE",
133 |         "AWS_SECRET_ACCESS_KEY": "YOUR_SECRET_ACCESS_KEY_HERE",
134 |         "AWS_SESSION_TOKEN": "YOUR_OPTIONAL_SESSION_ID_FOR_SSO_TEMPORARY_CREDENTIALS_HERE",
135 |         "AWS_REGION": "YOUR_AWS_REGION_HERE",
136 |         "AWS_KB_IDS": "[\"kb-12345\", \"kb-67890\"]"
137 |       }
138 |     }
139 |   }
140 | }
141 | ```
142 | 
143 | #### Local Repository (from cloned/built repo)
144 | 
145 | ```json
146 | {
147 |   "mcpServers": {
148 |     "aws-kb": {
149 |       "command": "node",
150 |       "args": [
151 |         "/path/to/mcp-aws-kb/dist/index.js"
152 |       ],
153 |       "env": {
154 |         "AWS_ACCESS_KEY_ID": "YOUR_ACCESS_KEY_HERE",
155 |         "AWS_SECRET_ACCESS_KEY": "YOUR_SECRET_ACCESS_KEY_HERE",
156 |         "AWS_SESSION_TOKEN": "YOUR_OPTIONAL_SESSION_ID_FOR_SSO_TEMPORARY_CREDENTIALS_HERE",
157 |         "AWS_REGION": "YOUR_AWS_REGION_HERE",
158 |         "AWS_KB_IDS": "[\"kb-12345\", \"kb-67890\"]"
159 |       },
160 |       "disabled": false,
161 |       "autoApprove": [
162 |         "retrieve_from_aws_kb"
163 |       ],
164 |       "timeout": 120
165 |     }
166 |   }
167 | }
168 | ```
169 | 
170 | ## Building
171 | 
172 | Docker:
173 | 
174 | ```sh
175 | docker build -t mcp/aws-kb-retrieval -f src/aws-kb-retrieval-server/Dockerfile .
176 | ```
177 | 
178 | ## License
179 | 
180 | 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.
181 | 
182 | 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`.
183 | 
```

--------------------------------------------------------------------------------
/Dockerfile:
--------------------------------------------------------------------------------

```dockerfile
 1 | FROM node:22.12-alpine AS builder
 2 | 
 3 | COPY src/aws-kb-retrieval-server /app
 4 | COPY tsconfig.json /tsconfig.json
 5 | 
 6 | WORKDIR /app
 7 | 
 8 | RUN --mount=type=cache,target=/root/.npm npm install
 9 | 
10 | FROM node:22-alpine AS release
11 | 
12 | WORKDIR /app
13 | 
14 | COPY --from=builder /app/dist /app/dist
15 | COPY --from=builder /app/package.json /app/package.json
16 | COPY --from=builder /app/package-lock.json /app/package-lock.json
17 | 
18 | ENV NODE_ENV=production
19 | 
20 | RUN npm ci --ignore-scripts --omit-dev
21 | 
22 | ENTRYPOINT ["node", "dist/index.js"]
```

--------------------------------------------------------------------------------
/tsconfig.json:
--------------------------------------------------------------------------------

```json
 1 | {
 2 |   "compilerOptions": {
 3 |     "target": "ES2020",
 4 |     "module": "NodeNext",
 5 |     "moduleResolution": "NodeNext",
 6 |     "esModuleInterop": true,
 7 |     "forceConsistentCasingInFileNames": true,
 8 |     "strict": true,
 9 |     "skipLibCheck": true,
10 |     "outDir": "./dist",
11 |     "rootDir": ".",
12 |     "declaration": true,
13 |     "composite": true,
14 |     "incremental": true,
15 |     "tsBuildInfoFile": "./dist/.tsbuildinfo",
16 |     "lib": ["ES2020", "DOM"]
17 |   },
18 |   "include": [
19 |     "./**/*.ts"
20 |   ],
21 |   "exclude": [
22 |     "node_modules",
23 |     "dist"
24 |   ]
25 | }
26 | 
```

--------------------------------------------------------------------------------
/package.json:
--------------------------------------------------------------------------------

```json
 1 | {
 2 |   "name": "@modelcontextprotocol/server-aws-kb-retrieval",
 3 |   "version": "0.7.0",
 4 |   "description": "MCP server for AWS Knowledge Base retrieval using Bedrock Agent Runtime",
 5 |   "license": "MIT",
 6 |   "author": "Anthropic, PBC (https://anthropic.com)",
 7 |   "homepage": "https://modelcontextprotocol.io",
 8 |   "bugs": "https://github.com/modelcontextprotocol/servers/issues",
 9 |   "type": "module",
10 |   "bin": {
11 |     "mcp-server-aws-kb-retrieval": "dist/index.js"
12 |   },
13 |   "files": [
14 |     "dist"
15 |   ],
16 |   "scripts": {
17 |     "build": "tsc && shx chmod +x dist/*.js",
18 |     "prepare": "npm run build",
19 |     "watch": "tsc --watch"
20 |   },
21 |   "dependencies": {
22 |     "@modelcontextprotocol/sdk": "1.7.0",
23 |     "@aws-sdk/client-bedrock-agent-runtime": "^3.765.0"
24 |   },
25 |   "devDependencies": {
26 |     "@types/node": "^22.13.10",
27 |     "shx": "^0.3.4",
28 |     "typescript": "^5.8.2"
29 |   },
30 |   "packageManager": "[email protected]+sha512.47870716bea1572b53df34ad8647b42962bc790ce2bf4562ba0f643237d7302a3d6a8ecef9e4bdfc01d23af1969aa90485d4cebb0b9638fa5ef1daef656f6c1b"
31 | }
32 | 
```

--------------------------------------------------------------------------------
/index.ts:
--------------------------------------------------------------------------------

```typescript
  1 | #!/usr/bin/env node
  2 | import { Server } from "@modelcontextprotocol/sdk/server/index.js";
  3 | import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
  4 | import {
  5 |   CallToolRequestSchema,
  6 |   ListToolsRequestSchema,
  7 |   Tool,
  8 | } from "@modelcontextprotocol/sdk/types.js";
  9 | import {
 10 |   BedrockAgentRuntimeClient,
 11 |   RetrieveCommand,
 12 |   RetrieveCommandInput,
 13 | } from "@aws-sdk/client-bedrock-agent-runtime";
 14 | 
 15 | // AWS client initialization
 16 | const clientConfig: Record<string, any> = {
 17 |   region: process.env.AWS_REGION,
 18 | };
 19 | 
 20 | // Only add explicit credentials if access keys are provided
 21 | // Otherwise use the default credential provider chain (supports SSO)
 22 | if (process.env.AWS_ACCESS_KEY_ID && process.env.AWS_SECRET_ACCESS_KEY) {
 23 |   clientConfig.credentials = {
 24 |     accessKeyId: process.env.AWS_ACCESS_KEY_ID,
 25 |     secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY,
 26 |   };
 27 | 
 28 |   // Add session token if provided (for temporary credentials)
 29 |   if (process.env.AWS_SESSION_TOKEN) {
 30 |     clientConfig.credentials.sessionToken = process.env.AWS_SESSION_TOKEN;
 31 |   }
 32 | }
 33 | 
 34 | const bedrockClient = new BedrockAgentRuntimeClient(clientConfig);
 35 | 
 36 | // Get configured knowledgebase IDs from environment variable
 37 | const configuredKnowledgeBaseIds = process.env.AWS_KB_IDS ?
 38 |   JSON.parse(process.env.AWS_KB_IDS) :
 39 |   [];
 40 | 
 41 | interface RAGSource {
 42 |   id: string;
 43 |   fileName: string;
 44 |   snippet: string;
 45 |   score: number;
 46 | }
 47 | 
 48 | async function retrieveContext(
 49 |   query: string,
 50 |   knowledgeBaseId: string,
 51 |   n: number = 3
 52 | ): Promise<{
 53 |   context: string;
 54 |   isRagWorking: boolean;
 55 |   ragSources: RAGSource[];
 56 | }> {
 57 |   try {
 58 |     if (!knowledgeBaseId) {
 59 |       console.error("knowledgeBaseId is not provided");
 60 |       return {
 61 |         context: "",
 62 |         isRagWorking: false,
 63 |         ragSources: [],
 64 |       };
 65 |     }
 66 | 
 67 |     const input: RetrieveCommandInput = {
 68 |       knowledgeBaseId: knowledgeBaseId,
 69 |       retrievalQuery: { text: query },
 70 |       retrievalConfiguration: {
 71 |         vectorSearchConfiguration: { numberOfResults: n },
 72 |       },
 73 |     };
 74 | 
 75 |     const command = new RetrieveCommand(input);
 76 |     const response = await bedrockClient.send(command);
 77 |     const rawResults = response?.retrievalResults || [];
 78 |     const ragSources: RAGSource[] = rawResults
 79 |       .filter((res) => res?.content?.text)
 80 |       .map((result, index) => {
 81 |         const uri = result?.location?.s3Location?.uri || "";
 82 |         const fileName = uri.split("/").pop() || `Source-${index}.txt`;
 83 |         return {
 84 |           id: (result.metadata?.["x-amz-bedrock-kb-chunk-id"] as string) || `chunk-${index}`,
 85 |           fileName: fileName.replace(/_/g, " ").replace(".txt", ""),
 86 |           snippet: result.content?.text || "",
 87 |           score: (result.score as number) || 0,
 88 |         };
 89 |       })
 90 |       .slice(0, 3);
 91 | 
 92 |     const context = rawResults
 93 |       .filter((res): res is { content: { text: string } } => res?.content?.text !== undefined)
 94 |       .map(res => res.content.text)
 95 |       .join("\n\n");
 96 | 
 97 |     return {
 98 |       context,
 99 |       isRagWorking: true,
100 |       ragSources,
101 |     };
102 |   } catch (error) {
103 |     console.error("RAG Error:", error);
104 |     return { context: "", isRagWorking: false, ragSources: [] };
105 |   }
106 | }
107 | 
108 | // Define the retrieval tool
109 | const RETRIEVAL_TOOL: Tool = {
110 |   name: "retrieve_from_aws_kb",
111 |   description: "Performs retrieval from the AWS Knowledge Base using the provided query and Knowledge Base ID.",
112 |   inputSchema: {
113 |     type: "object",
114 |     properties: {
115 |       query: { type: "string", description: "The query to perform retrieval on" },
116 |       knowledgeBaseId: {
117 |         type: "string",
118 |         description: configuredKnowledgeBaseIds.length > 0
119 |           ? "The ID of the AWS Knowledge Base (optional if configured via AWS_KB_IDS)"
120 |           : "The ID of the AWS Knowledge Base"
121 |       },
122 |       n: { type: "number", default: 3, description: "Number of results to retrieve" },
123 |     },
124 |     required: configuredKnowledgeBaseIds.length > 0 ? ["query"] : ["query", "knowledgeBaseId"],
125 |   },
126 | };
127 | 
128 | // Server setup
129 | const server = new Server(
130 |   {
131 |     name: "aws-kb-retrieval-server",
132 |     version: "0.2.0",
133 |   },
134 |   {
135 |     capabilities: {
136 |       tools: {},
137 |     },
138 |   },
139 | );
140 | 
141 | // Request handlers
142 | server.setRequestHandler(ListToolsRequestSchema, async () => ({
143 |   tools: [RETRIEVAL_TOOL],
144 | }));
145 | 
146 | server.setRequestHandler(CallToolRequestSchema, async (request) => {
147 |   const { name, arguments: args } = request.params;
148 | 
149 |   if (name === "retrieve_from_aws_kb") {
150 |     const { query, knowledgeBaseId, n = 3 } = args as Record<string, any>;
151 | 
152 |     // Determine which knowledge base ID to use
153 |     let actualKnowledgeBaseId = knowledgeBaseId;
154 | 
155 |     // If no ID provided but we have configured IDs, use the first one
156 |     if (!actualKnowledgeBaseId && configuredKnowledgeBaseIds.length > 0) {
157 |       actualKnowledgeBaseId = configuredKnowledgeBaseIds[0];
158 |       console.error(`Using configured knowledge base ID: ${actualKnowledgeBaseId}`);
159 |     }
160 | 
161 |     // If still no ID available, return an error
162 |     if (!actualKnowledgeBaseId) {
163 |       return {
164 |         content: [{
165 |           type: "text",
166 |           text: "No knowledge base ID provided. Either include a knowledgeBaseId in your request or configure AWS_KB_IDS in the environment."
167 |         }],
168 |         isError: true,
169 |       };
170 |     }
171 | 
172 |     try {
173 |       const result = await retrieveContext(query, actualKnowledgeBaseId, n);
174 |       if (result.isRagWorking) {
175 |         // Format RAG sources for readability
176 |         const formattedSources = result.ragSources.map((source, index) => {
177 |           return `Source ${index + 1}: ${source.fileName} (score: ${source.score.toFixed(3)})\n${source.snippet}`;
178 |         }).join('\n\n');
179 | 
180 |         return {
181 |           content: [
182 |             {
183 |               type: "text",
184 |               text: result.context
185 |             },
186 |             {
187 |               type: "json",
188 |               json: {
189 |                 ragSources: result.ragSources
190 |               }
191 |             }
192 |           ],
193 |         };
194 |       } else {
195 |         return {
196 |           content: [{ type: "text", text: "Retrieval failed or returned no results." }],
197 |         };
198 |       }
199 |     } catch (error) {
200 |       return {
201 |         content: [{ type: "text", text: `Error occurred: ${error}` }],
202 |       };
203 |     }
204 |   } else {
205 |     return {
206 |       content: [{ type: "text", text: `Unknown tool: ${name}` }],
207 |       isError: true,
208 |     };
209 |   }
210 | });
211 | 
212 | // Server startup
213 | async function runServer() {
214 |   const transport = new StdioServerTransport();
215 |   await server.connect(transport);
216 |   console.error("AWS KB Retrieval Server running on stdio");
217 | }
218 | 
219 | runServer().catch((error) => {
220 |   console.error("Fatal error running server:", error);
221 |   process.exit(1);
222 | });
223 | 
```