From 045b75ba3b95ccba0733485296faca4eee8b014a Mon Sep 17 00:00:00 2001 From: rincel <0xrinegade@gmail.com> Date: Thu, 9 Jan 2025 10:08:13 +0300 Subject: [PATCH] metal mcp init --- .gitignore | 38 ++++ README.md | 53 +++++ package.json | 23 +++ src/index.ts | 12 ++ src/server.ts | 532 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/types.ts | 74 +++++++ tsconfig.json | 17 ++ 7 files changed, 749 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 package.json create mode 100644 src/index.ts create mode 100644 src/server.ts create mode 100644 src/types.ts create mode 100644 tsconfig.json diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1215557 --- /dev/null +++ b/.gitignore @@ -0,0 +1,38 @@ +# Dependencies +node_modules/ +package-lock.json + +# Build output +build/ +dist/ +*.tsbuildinfo + +# Environment +.env +.env.local +.env.*.local + +# IDE +.vscode/ +.idea/ +*.swp +*.swo + +# Logs +logs/ +*.log +npm-debug.log* +yarn-debug.log* +yarn-error.log* + +# OS +.DS_Store +Thumbs.db + +# Testing +coverage/ +.nyc_output/ + +# Cache +.npm/ +.eslintcache diff --git a/README.md b/README.md new file mode 100644 index 0000000..e48aa8f --- /dev/null +++ b/README.md @@ -0,0 +1,53 @@ +# Metal MCP Server + +An MCP server providing Metal Framework documentation search and code generation capabilities. + +## One-Line Installation + +```bash +npx @modelcontextprotocol/create-server metal-mcp && cd metal-mcp && npm install && npm run build +``` + +## Features + +### Tools + +1. `search_metal_docs` + - Search Metal Framework documentation and code examples using natural language queries + - Parameters: + - `query`: Natural language query about Metal Framework + - `limit`: Maximum number of results to return (default: 3) + +2. `generate_metal_code` + - Generate Metal Framework code for common tasks + - Parameters: + - `task`: Description of the Metal task to generate code for + - `language`: Programming language (objective-c, swift, or metal) + +### Resources + +1. `metal://docs/getting-started` + - Comprehensive guide for getting started with Metal Framework + +2. `metal://docs/best-practices` + - Best practices and optimization tips for Metal Framework + +## Usage + +After installation, add the server to your MCP configuration: + +```json +{ + "mcpServers": { + "metal": { + "command": "node", + "args": ["/path/to/metal-mcp/build/index.js"] + } + } +} +``` + +The server will provide Metal Framework expertise through the MCP protocol, allowing you to: +- Search Metal documentation with natural language queries +- Generate code snippets for common Metal tasks +- Access Metal best practices and getting started guides diff --git a/package.json b/package.json new file mode 100644 index 0000000..4b690b9 --- /dev/null +++ b/package.json @@ -0,0 +1,23 @@ +{ + "name": "metal-mcp-server", + "version": "0.1.0", + "type": "module", + "scripts": { + "build": "tsc && chmod +x build/index.js", + "start": "node build/index.js", + "dev": "ts-node-esm src/index.ts" + }, + "dependencies": { + "@modelcontextprotocol/sdk": "^1.1.0", + "@xenova/transformers": "^2.14.0", + "axios": "^1.6.5", + "cheerio": "^1.0.0-rc.12", + "hnswlib-node": "^2.0.0" + }, + "devDependencies": { + "@types/cheerio": "^0.22.35", + "@types/node": "^20.11.5", + "ts-node": "^10.9.2", + "typescript": "^5.3.3" + } +} diff --git a/src/index.ts b/src/index.ts new file mode 100644 index 0000000..455bcc3 --- /dev/null +++ b/src/index.ts @@ -0,0 +1,12 @@ +#!/usr/bin/env node +import { MetalExpertServer } from './server.js'; + +async function main() { + const server = new MetalExpertServer(); + await server.run(); +} + +main().catch((error) => { + console.error('Server error:', error); + process.exit(1); +}); diff --git a/src/server.ts b/src/server.ts new file mode 100644 index 0000000..8f2eb94 --- /dev/null +++ b/src/server.ts @@ -0,0 +1,532 @@ +import { EventEmitter } from 'events'; +import { pipeline } from '@xenova/transformers'; +import hnswlib from 'hnswlib-node'; +import axios from 'axios'; +import * as cheerio from 'cheerio'; +import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; +import { + Tool, + Resource, + McpError, + ErrorCode, + ListToolsRequestSchema, + ListResourcesRequestSchema, + CallToolRequestSchema, + ReadResourceRequestSchema +} from '@modelcontextprotocol/sdk/types.js'; + +interface DocItem { + title: string; + content: string; + url: string; + type: 'doc' | 'example'; +} + +const TOOLS: Tool[] = [ + { + name: 'search_metal_docs', + description: 'Search Metal Framework documentation and code examples using natural language queries', + inputSchema: { + type: 'object', + properties: { + query: { + type: 'string', + description: 'Natural language query about Metal Framework', + }, + limit: { + type: 'number', + description: 'Maximum number of results to return', + default: 3, + }, + }, + required: ['query'], + }, + }, + { + name: 'generate_metal_code', + description: 'Generate Metal Framework code for common tasks', + inputSchema: { + type: 'object', + properties: { + task: { + type: 'string', + description: 'Description of the Metal task to generate code for', + }, + language: { + type: 'string', + description: 'Programming language (objective-c, swift, or metal)', + default: 'swift', + }, + }, + required: ['task'], + }, + }, +]; + +const RESOURCES: Resource[] = [ + { + uri: 'metal://docs/getting-started', + name: 'Metal Getting Started Guide', + description: 'Comprehensive guide for getting started with Metal Framework', + }, + { + uri: 'metal://docs/best-practices', + name: 'Metal Best Practices', + description: 'Best practices and optimization tips for Metal Framework', + }, +]; + +export class MetalExpertServer { + private docs: DocItem[] = []; + private vectorStore: any; + private embedder: any = null; + private initialized = false; + private server: Server; + + constructor() { + this.vectorStore = new hnswlib.HierarchicalNSW('cosine', 384); + this.server = new Server( + { + name: 'metal-mcp-server', + version: '0.1.0', + }, + { + capabilities: { + tools: TOOLS.reduce((acc, tool) => { + acc[tool.name] = tool; + return acc; + }, {} as Record), + resources: RESOURCES.reduce((acc, resource) => { + acc[resource.uri] = resource; + return acc; + }, {} as Record) + } + } + ); + this.setupHandlers(); + } + + private setupHandlers() { + this.server.setRequestHandler(ListToolsRequestSchema, async () => ({ + tools: TOOLS + })); + + this.server.setRequestHandler(ListResourcesRequestSchema, async () => ({ + resources: RESOURCES + })); + + this.server.setRequestHandler(CallToolRequestSchema, this.handleToolRequest.bind(this)); + this.server.setRequestHandler(ReadResourceRequestSchema, this.handleResourceRequest.bind(this)); + } + + async run() { + const transport = new StdioServerTransport(); + await this.server.connect(transport); + console.error('Metal MCP server running on stdio'); + } + + private async initEmbedder() { + if (!this.embedder) { + this.embedder = await pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2'); + } + return this.embedder; + } + + private async initializeRAG() { + if (this.initialized) return; + + // Fetch Metal documentation + const response = await axios.get('https://developer.apple.com/documentation/metal'); + const $ = cheerio.load(response.data); + + // Extract documentation content + $('.documentation-content').each((i, elem) => { + const title = $(elem).find('h1').text(); + const content = $(elem).text(); + const url = $(elem).find('a').attr('href') || ''; + + this.docs.push({ + title, + content, + url, + type: 'doc' + }); + }); + + // Add code examples + this.docs.push({ + title: 'Basic Metal Setup', + content: ` + // Create a Metal device + id device = MTLCreateSystemDefaultDevice(); + + // Create a command queue + id commandQueue = [device newCommandQueue]; + + // Create a Metal library + id defaultLibrary = [device newDefaultLibrary]; + `, + url: 'examples/basic_setup', + type: 'example' + }); + + // Initialize vector store + this.vectorStore.initIndex(this.docs.length, true); + + // Get embeddings and add to vector store + const embedder = await this.initEmbedder(); + for (let i = 0; i < this.docs.length; i++) { + const embedding = await embedder(this.docs[i].content, { pooling: 'mean', normalize: true }); + this.vectorStore.addPoint(embedding.data, i); + } + + this.initialized = true; + } + + async handleToolRequest(request: { params: { name: string; arguments?: Record } }): Promise { + await this.initializeRAG(); + + switch (request.params.name) { + case 'search_metal_docs': { + const query = request.params.arguments?.query as string; + const limit = (request.params.arguments?.limit as number) ?? 3; + + // Get query embedding + const embedder = await this.initEmbedder(); + const queryEmbedding = await embedder(query, { pooling: 'mean', normalize: true }); + + // Search vector store + const results = this.vectorStore.searchKnn(queryEmbedding.data, limit); + + // Format results + const searchResults = results.neighbors.map((idx: number) => { + const doc = this.docs[idx]; + return { + title: doc.title, + content: doc.content, + url: doc.url, + type: doc.type, + }; + }); + + return { + content: [ + { + type: 'text', + text: JSON.stringify(searchResults, null, 2), + }, + ], + }; + } + + case 'generate_metal_code': { + const task = request.params.arguments?.task as string; + const language = (request.params.arguments?.language as string) ?? 'swift'; + + // Search for relevant examples + const embedder = await this.initEmbedder(); + const taskEmbedding = await embedder(task, { pooling: 'mean', normalize: true }); + const results = this.vectorStore.searchKnn(taskEmbedding.data, 3); + + // Use examples to generate contextual code + const examples = results.neighbors.map((idx: number) => this.docs[idx]); + + // Generate code based on task and examples + let code = ''; + if (task.includes('compute')) { + code = this.generateComputeCode(language); + } else if (task.includes('render')) { + code = this.generateRenderCode(language); + } else { + code = this.generateBasicCode(language); + } + + return { + content: [ + { + type: 'text', + text: code, + }, + ], + }; + } + + default: + throw new McpError( + ErrorCode.MethodNotFound, + `Unknown tool: ${request.params.name}` + ); + } + } + + async handleResourceRequest(request: { params: { uri: string } }): Promise { + await this.initializeRAG(); + + const uri = request.params.uri; + let content = ''; + + switch (uri) { + case 'metal://docs/getting-started': + content = this.docs.filter(d => d.type === 'doc' && d.title.includes('Getting Started')) + .map(d => d.content) + .join('\n\n'); + break; + + case 'metal://docs/best-practices': + content = this.docs.filter(d => d.type === 'doc' && d.title.includes('Best Practices')) + .map(d => d.content) + .join('\n\n'); + break; + + default: + throw new McpError( + ErrorCode.InvalidRequest, + `Unknown resource URI: ${uri}` + ); + } + + return { + contents: [ + { + uri, + mimeType: 'text/plain', + text: content, + }, + ], + }; + } + + async listTools(): Promise { + return TOOLS; + } + + async listResources(): Promise { + return RESOURCES; + } + + private generateComputeCode(language: string): string { + switch (language) { + case 'swift': + return ` +import Metal + +guard let device = MTLCreateSystemDefaultDevice() else { + fatalError("GPU not available") +} + +let commandQueue = device.makeCommandQueue()! + +// Create compute pipeline +let library = device.makeDefaultLibrary()! +let function = library.makeFunction(name: "compute_function")! +let pipeline = try! device.makeComputePipelineState(function: function) + +// Create command buffer and encoder +let commandBuffer = commandQueue.makeCommandBuffer()! +let computeEncoder = commandBuffer.makeComputeCommandEncoder()! + +computeEncoder.setComputePipelineState(pipeline) +// Set buffer bindings and dispatch here + +computeEncoder.endEncoding() +commandBuffer.commit() +`; + + case 'objective-c': + return ` +id device = MTLCreateSystemDefaultDevice(); +id commandQueue = [device newCommandQueue]; + +// Create compute pipeline +id library = [device newDefaultLibrary]; +id function = [library newFunctionWithName:@"compute_function"]; +id pipeline = [device newComputePipelineStateWithFunction:function error:nil]; + +// Create command buffer and encoder +id commandBuffer = [commandQueue commandBuffer]; +id computeEncoder = [commandBuffer computeCommandEncoder]; + +[computeEncoder setComputePipelineState:pipeline]; +// Set buffer bindings and dispatch here + +[computeEncoder endEncoding]; +[commandBuffer commit]; +`; + + case 'metal': + return ` +#include +using namespace metal; + +kernel void compute_function( + device float *input_buffer [[buffer(0)]], + device float *output_buffer [[buffer(1)]], + uint index [[thread_position_in_grid]]) +{ + // Compute work here + output_buffer[index] = input_buffer[index] * 2.0; +} +`; + + default: + throw new Error(`Unsupported language: ${language}`); + } + } + + private generateRenderCode(language: string): string { + switch (language) { + case 'swift': + return ` +import Metal +import MetalKit + +guard let device = MTLCreateSystemDefaultDevice() else { + fatalError("GPU not available") +} + +let commandQueue = device.makeCommandQueue()! + +// Create render pipeline +let library = device.makeDefaultLibrary()! +let vertexFunction = library.makeFunction(name: "vertex_function")! +let fragmentFunction = library.makeFunction(name: "fragment_function")! + +let pipelineDescriptor = MTLRenderPipelineDescriptor() +pipelineDescriptor.vertexFunction = vertexFunction +pipelineDescriptor.fragmentFunction = fragmentFunction +pipelineDescriptor.colorAttachments[0].pixelFormat = .bgra8Unorm + +let pipeline = try! device.makeRenderPipelineState(descriptor: pipelineDescriptor) + +// Create command buffer and encoder +let commandBuffer = commandQueue.makeCommandBuffer()! +let renderPassDescriptor = MTLRenderPassDescriptor() +let renderEncoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDescriptor)! + +renderEncoder.setRenderPipelineState(pipeline) +// Set vertex buffers and draw here + +renderEncoder.endEncoding() +commandBuffer.commit() +`; + + case 'objective-c': + return ` +id device = MTLCreateSystemDefaultDevice(); +id commandQueue = [device newCommandQueue]; + +// Create render pipeline +id library = [device newDefaultLibrary]; +id vertexFunction = [library newFunctionWithName:@"vertex_function"]; +id fragmentFunction = [library newFunctionWithName:@"fragment_function"]; + +MTLRenderPipelineDescriptor *pipelineDescriptor = [[MTLRenderPipelineDescriptor alloc] init]; +pipelineDescriptor.vertexFunction = vertexFunction; +pipelineDescriptor.fragmentFunction = fragmentFunction; +pipelineDescriptor.colorAttachments[0].pixelFormat = MTLPixelFormatBGRA8Unorm; + +id pipeline = [device newRenderPipelineStateWithDescriptor:pipelineDescriptor error:nil]; + +// Create command buffer and encoder +id commandBuffer = [commandQueue commandBuffer]; +MTLRenderPassDescriptor *renderPassDescriptor = [MTLRenderPassDescriptor renderPassDescriptor]; +id renderEncoder = [commandBuffer renderCommandEncoderWithDescriptor:renderPassDescriptor]; + +[renderEncoder setRenderPipelineState:pipeline]; +// Set vertex buffers and draw here + +[renderEncoder endEncoding]; +[commandBuffer commit]; +`; + + case 'metal': + return ` +#include +using namespace metal; + +struct VertexIn { + float3 position [[attribute(0)]]; + float2 texCoord [[attribute(1)]]; +}; + +struct VertexOut { + float4 position [[position]]; + float2 texCoord; +}; + +vertex VertexOut vertex_function( + VertexIn in [[stage_in]]) +{ + VertexOut out; + out.position = float4(in.position, 1.0); + out.texCoord = in.texCoord; + return out; +} + +fragment float4 fragment_function( + VertexOut in [[stage_in]], + texture2d texture [[texture(0)]], + sampler textureSampler [[sampler(0)]]) +{ + return texture.sample(textureSampler, in.texCoord); +} +`; + + default: + throw new Error(`Unsupported language: ${language}`); + } + } + + private generateBasicCode(language: string): string { + switch (language) { + case 'swift': + return ` +import Metal + +guard let device = MTLCreateSystemDefaultDevice() else { + fatalError("GPU not available") +} + +let commandQueue = device.makeCommandQueue()! +let library = device.makeDefaultLibrary()! + +// Create buffers +let buffer = device.makeBuffer(length: 1024, options: .storageModeShared)! + +// Create command buffer +let commandBuffer = commandQueue.makeCommandBuffer()! +commandBuffer.commit() +`; + + case 'objective-c': + return ` +id device = MTLCreateSystemDefaultDevice(); +id commandQueue = [device newCommandQueue]; +id library = [device newDefaultLibrary]; + +// Create buffers +id buffer = [device newBufferWithLength:1024 options:MTLResourceStorageModeShared]; + +// Create command buffer +id commandBuffer = [commandQueue commandBuffer]; +[commandBuffer commit]; +`; + + case 'metal': + return ` +#include +using namespace metal; + +struct Uniforms { + float4x4 modelMatrix; + float4x4 viewMatrix; + float4x4 projectionMatrix; +}; +`; + + default: + throw new Error(`Unsupported language: ${language}`); + } + } +} diff --git a/src/types.ts b/src/types.ts new file mode 100644 index 0000000..1d452d2 --- /dev/null +++ b/src/types.ts @@ -0,0 +1,74 @@ +export interface Server { + name: string; + version: string; +} + +export interface ServerCapabilities { + resources: Record; + tools: Record; +} + +export interface ToolRequest { + params: { + name: string; + arguments: Record; + }; +} + +export interface ResourceRequest { + params: { + uri: string; + }; +} + +export interface Content { + type: string; + text: string; +} + +export interface ToolResponse { + content: Content[]; +} + +export interface ResourceResponse { + contents: { + uri: string; + mimeType: string; + text: string; + }[]; +} + +export interface Tool { + name: string; + description: string; + inputSchema: { + type: string; + properties: Record; + required?: string[]; + }; +} + +export interface Resource { + uri: string; + name: string; + description?: string; + mimeType?: string; +} + +export interface ErrorResponse { + code: number; + message: string; +} + +export class McpError extends Error { + constructor(public code: number, message: string) { + super(message); + this.name = 'McpError'; + } +} + +export const ErrorCode = { + InvalidRequest: 400, + MethodNotFound: 404, + InternalError: 500, +} as const; diff --git a/tsconfig.json b/tsconfig.json new file mode 100644 index 0000000..da386cd --- /dev/null +++ b/tsconfig.json @@ -0,0 +1,17 @@ +{ + "compilerOptions": { + "target": "ES2020", + "module": "ES2020", + "moduleResolution": "node", + "esModuleInterop": true, + "strict": true, + "skipLibCheck": true, + "forceConsistentCasingInFileNames": true, + "outDir": "build", + "declaration": true, + "allowJs": true, + "resolveJsonModule": true + }, + "include": ["src/**/*"], + "exclude": ["node_modules", "build"] +}