nvidia-nim-mcp 1.0.0

package/dist/client.js

@@ -0,0 +1,178 @@
+import axios from "axios";
+import axiosRetry from "axios-retry";
+import { getConfig } from "./config.js";
+import { logger } from "./logger.js";
+const config = getConfig();
+// ─── Rate Limiter ─────────────────────────────────────────────────────────────
+class RateLimiter {
+    requests = [];
+    maxRequests;
+    windowMs = 60_000;
+    constructor(maxRequests) {
+        this.maxRequests = maxRequests;
+    }
+    async acquire() {
+        const now = Date.now();
+        this.requests = this.requests.filter((t) => now - t < this.windowMs);
+        if (this.requests.length >= this.maxRequests) {
+            const oldest = this.requests[0];
+            const waitMs = this.windowMs - (now - oldest) + 10;
+            logger.warn(`Rate limit reached. Waiting ${waitMs}ms`);
+            await new Promise((res) => setTimeout(res, waitMs));
+            return this.acquire();
+        }
+        this.requests.push(now);
+    }
+}
+// ─── Client ───────────────────────────────────────────────────────────────────
+export class NIMClient {
+    http;
+    rateLimiter;
+    constructor() {
+        this.rateLimiter = new RateLimiter(config.MAX_REQUESTS_PER_MINUTE);
+        this.http = axios.create({
+            baseURL: config.NVIDIA_NIM_BASE_URL,
+            timeout: config.REQUEST_TIMEOUT_MS,
+            headers: {
+                Authorization: `Bearer ${config.NVIDIA_API_KEY}`,
+                "Content-Type": "application/json",
+                "User-Agent": `${config.MCP_SERVER_NAME}/${config.MCP_SERVER_VERSION}`,
+            },
+        });
+        axiosRetry(this.http, {
+            retries: config.MAX_RETRIES,
+            retryDelay: (retryCount) => {
+                const delay = config.RETRY_DELAY_MS * Math.pow(2, retryCount - 1);
+                logger.warn(`Retry attempt ${retryCount}, waiting ${delay}ms`);
+                return delay;
+            },
+            retryCondition: (error) => {
+                const status = error.response?.status;
+                // Retry on 429, 500, 502, 503, 504
+                return (axiosRetry.isNetworkError(error) ||
+                    status === 429 ||
+                    (status !== undefined && status >= 500));
+            },
+            onRetry: (retryCount, error) => {
+                logger.warn(`Retrying request (${retryCount}/${config.MAX_RETRIES})`, {
+                    status: error.response?.status,
+                    message: error.message,
+                });
+            },
+        });
+        // Request interceptor for logging
+        this.http.interceptors.request.use((req) => {
+            logger.debug("NIM API request", {
+                method: req.method?.toUpperCase(),
+                url: req.url,
+                model: req.data?.model,
+            });
+            return req;
+        });
+        // Response interceptor for logging
+        this.http.interceptors.response.use((res) => {
+            logger.debug("NIM API response", {
+                status: res.status,
+                usage: res.data?.usage,
+            });
+            return res;
+        }, (err) => {
+            logger.error("NIM API error", {
+                status: err.response?.status,
+                data: err.response?.data,
+                message: err.message,
+            });
+            return Promise.reject(this.normalizeError(err));
+        });
+    }
+    normalizeError(error) {
+        const status = error.response?.status;
+        const data = error.response?.data;
+        const apiMsg = data?.detail ||
+            data?.message ||
+            error.message;
+        if (status === 401)
+            return new Error(`Authentication failed: ${apiMsg}`);
+        if (status === 403)
+            return new Error(`Authorization failed: ${apiMsg}`);
+        if (status === 404)
+            return new Error(`Model not found: ${apiMsg}`);
+        if (status === 422)
+            return new Error(`Invalid request: ${apiMsg}`);
+        if (status === 429)
+            return new Error(`Rate limit exceeded: ${apiMsg}`);
+        if (status && status >= 500)
+            return new Error(`NVIDIA NIM server error (${status}): ${apiMsg}`);
+        return new Error(`NIM API error: ${apiMsg}`);
+    }
+    async chatCompletion(request) {
+        await this.rateLimiter.acquire();
+        // Apply defaults
+        const payload = {
+            temperature: config.DEFAULT_TEMPERATURE,
+            top_p: config.DEFAULT_TOP_P,
+            max_tokens: Math.min(request.max_tokens ?? config.DEFAULT_MAX_TOKENS, config.MAX_TOKENS_PER_REQUEST),
+            ...request,
+            stream: false,
+        };
+        const { data } = await this.http.post("/chat/completions", payload);
+        return data;
+    }
+    async *chatCompletionStream(request) {
+        await this.rateLimiter.acquire();
+        const payload = {
+            temperature: config.DEFAULT_TEMPERATURE,
+            top_p: config.DEFAULT_TOP_P,
+            max_tokens: Math.min(request.max_tokens ?? config.DEFAULT_MAX_TOKENS, config.MAX_TOKENS_PER_REQUEST),
+            ...request,
+            stream: true,
+        };
+        const response = await this.http.post("/chat/completions", payload, {
+            responseType: "stream",
+        });
+        const stream = response.data;
+        for await (const chunk of stream) {
+            const lines = chunk
+                .toString("utf8")
+                .split("\n")
+                .filter((l) => l.startsWith("data: "));
+            for (const line of lines) {
+                const json = line.slice(6).trim();
+                if (json === "[DONE]")
+                    return;
+                try {
+                    const parsed = JSON.parse(json);
+                    const content = parsed.choices?.[0]?.delta?.content;
+                    if (content)
+                        yield content;
+                }
+                catch {
+                    // skip malformed chunks
+                }
+            }
+        }
+    }
+    async embeddings(request) {
+        await this.rateLimiter.acquire();
+        const { data } = await this.http.post("/embeddings", request);
+        return data;
+    }
+    async rerank(request) {
+        await this.rateLimiter.acquire();
+        const { data } = await this.http.post("/ranking", request);
+        return data;
+    }
+    async listModels() {
+        await this.rateLimiter.acquire();
+        try {
+            const { data } = await this.http.get("/models");
+            return data.data.map((m) => m.id);
+        }
+        catch {
+            // If the API doesn't support listing, return catalog
+            return Object.keys(NIM_MODELS_EXPORT);
+        }
+    }
+}
+// Re-export for convenience
+import { NIM_MODELS as NIM_MODELS_EXPORT } from "./models.js";

package/dist/config.d.ts

@@ -0,0 +1,24 @@
+export declare class ConfigError extends Error {
+    constructor(message: string);
+}
+export declare function getConfig(): {
+    NVIDIA_API_KEY?: string;
+    NVIDIA_NIM_BASE_URL?: string;
+    NVIDIA_CHAT_MODEL?: string;
+    NVIDIA_EMBEDDINGS_MODEL?: string;
+    NVIDIA_RERANKING_MODEL?: string;
+    MCP_SERVER_NAME?: string;
+    MCP_SERVER_VERSION?: string;
+    MCP_SERVER_PORT?: number;
+    LOG_LEVEL?: "error" | "warn" | "info" | "debug" | "silly";
+    MAX_REQUESTS_PER_MINUTE?: number;
+    MAX_TOKENS_PER_REQUEST?: number;
+    REQUEST_TIMEOUT_MS?: number;
+    MAX_RETRIES?: number;
+    RETRY_DELAY_MS?: number;
+    DEFAULT_MODEL?: string;
+    DEFAULT_TEMPERATURE?: number;
+    DEFAULT_TOP_P?: number;
+    DEFAULT_MAX_TOKENS?: number;
+};
+export type Config = ReturnType<typeof getConfig>;

package/dist/config.js

@@ -0,0 +1,50 @@
+import { z } from "zod";
+const ConfigSchema = z.object({
+    // NVIDIA API
+    NVIDIA_API_KEY: z.string().min(1, "NVIDIA_API_KEY is required"),
+    NVIDIA_NIM_BASE_URL: z
+        .string()
+        .url()
+        .default("https://integrate.api.nvidia.com/v1"),
+    NVIDIA_CHAT_MODEL: z.string().optional(),
+    NVIDIA_EMBEDDINGS_MODEL: z.string().optional(),
+    NVIDIA_RERANKING_MODEL: z.string().optional(),
+    // Server
+    MCP_SERVER_NAME: z.string().default("nvidia-nim-mcp"),
+    MCP_SERVER_VERSION: z.string().default("1.0.0"),
+    MCP_SERVER_PORT: z.coerce.number().int().positive().default(8080),
+    LOG_LEVEL: z
+        .enum(["error", "warn", "info", "debug", "silly"])
+        .default("info"),
+    // Rate Limiting
+    MAX_REQUESTS_PER_MINUTE: z.coerce.number().int().positive().default(40),
+    MAX_TOKENS_PER_REQUEST: z.coerce.number().int().positive().default(4096),
+    REQUEST_TIMEOUT_MS: z.coerce.number().int().positive().default(120000),
+    // Retry
+    MAX_RETRIES: z.coerce.number().int().min(0).max(10).default(3),
+    RETRY_DELAY_MS: z.coerce.number().int().positive().default(1000),
+    // Defaults
+    DEFAULT_MODEL: z.string().default("z-ai/glm5"),
+    DEFAULT_TEMPERATURE: z.coerce.number().min(0).max(2).default(0.3),
+    DEFAULT_TOP_P: z.coerce.number().min(0).max(1).default(0.95),
+    DEFAULT_MAX_TOKENS: z.coerce.number().int().positive().default(4096),
+});
+export class ConfigError extends Error {
+    constructor(message) {
+        super(message);
+        this.name = 'ConfigError';
+    }
+}
+function loadConfig() {
+    const result = ConfigSchema.safeParse(process.env);
+    if (!result.success) {
+        const errors = result.error.issues
+            .map((e) => `  - ${e.path.join(".")}: ${e.message}`)
+            .join("\n");
+        throw new ConfigError(`Configuration validation failed:\n${errors}`);
+    }
+    return result.data;
+}
+export function getConfig() {
+    return loadConfig();
+}

package/dist/handlers.d.ts

@@ -0,0 +1,26 @@
+import { NIMClient } from "./client.js";
+export type ToolResult = {
+    content: Array<{
+        type: "text";
+        text: string;
+    }>;
+    isError?: false;
+} | {
+    content: Array<{
+        type: "text";
+        text: string;
+    }>;
+    isError: true;
+};
+export declare class ToolHandlers {
+    private readonly client;
+    constructor(client: NIMClient);
+    handle(toolName: string, rawArgs: unknown): Promise<ToolResult>;
+    private chatCompletion;
+    private textGeneration;
+    private createEmbeddings;
+    private rerankPassages;
+    private functionCalling;
+    private listModels;
+    private getModelInfo;
+}

package/dist/handlers.js

@@ -0,0 +1,255 @@
+import { NIM_MODELS, getModelsByCategory, getModel } from "./models.js";
+import { getConfig } from "./config.js";
+import { logger } from "./logger.js";
+import { ChatCompletionSchema, TextGenerationSchema, EmbeddingsSchema, RerankSchema, FunctionCallingSchema, ListModelsSchema, ModelInfoSchema, } from "./tools.js";
+const config = getConfig();
+function ok(data) {
+    return {
+        content: [{ type: "text", text: JSON.stringify(data, null, 2) }],
+    };
+}
+function err(message) {
+    return {
+        content: [{ type: "text", text: `Error: ${message}` }],
+        isError: true,
+    };
+}
+export class ToolHandlers {
+    client;
+    constructor(client) {
+        this.client = client;
+    }
+    async handle(toolName, rawArgs) {
+        const start = Date.now();
+        logger.info(`Executing tool: ${toolName}`);
+        try {
+            let result;
+            switch (toolName) {
+                case "chat_completion":
+                    result = await this.chatCompletion(rawArgs);
+                    break;
+                case "text_generation":
+                    result = await this.textGeneration(rawArgs);
+                    break;
+                case "create_embeddings":
+                    result = await this.createEmbeddings(rawArgs);
+                    break;
+                case "rerank_passages":
+                    result = await this.rerankPassages(rawArgs);
+                    break;
+                case "function_calling":
+                    result = await this.functionCalling(rawArgs);
+                    break;
+                case "list_models":
+                    result = await this.listModels(rawArgs);
+                    break;
+                case "get_model_info":
+                    result = await this.getModelInfo(rawArgs);
+                    break;
+                default:
+                    result = err(`Unknown tool: ${toolName}`);
+            }
+            logger.info(`Tool ${toolName} completed in ${Date.now() - start}ms`);
+            return result;
+        }
+        catch (e) {
+            const msg = e instanceof Error ? e.message : String(e);
+            logger.error(`Tool ${toolName} failed: ${msg}`, { error: e });
+            return err(msg);
+        }
+    }
+    // ─── Handlers ─────────────────────────────────────────────────────────────
+    async chatCompletion(rawArgs) {
+        const args = ChatCompletionSchema.parse(rawArgs);
+        const model = args.model ?? config.DEFAULT_MODEL;
+        const messages = args.messages.map(msg => ({
+            role: msg.role,
+            content: msg.content
+        }));
+        if (args.system_prompt) {
+            messages.unshift({ role: "system", content: args.system_prompt });
+        }
+        const response = await this.client.chatCompletion({
+            model,
+            messages,
+            temperature: args.temperature,
+            top_p: args.top_p,
+            max_tokens: args.max_tokens,
+            stop: args.stop,
+            frequency_penalty: args.frequency_penalty,
+            presence_penalty: args.presence_penalty,
+            seed: args.seed,
+        });
+        return ok({
+            id: response.id,
+            model: response.model,
+            content: response.choices[0].message.content,
+            finish_reason: response.choices[0].finish_reason,
+            usage: response.usage,
+        });
+    }
+    async textGeneration(rawArgs) {
+        const args = TextGenerationSchema.parse(rawArgs);
+        const model = args.model ?? config.DEFAULT_MODEL;
+        const messages = [];
+        if (args.system_prompt) {
+            messages.push({ role: "system", content: args.system_prompt });
+        }
+        messages.push({ role: "user", content: args.prompt });
+        const response = await this.client.chatCompletion({
+            model,
+            messages,
+            temperature: args.temperature,
+            max_tokens: args.max_tokens,
+            stop: args.stop,
+        });
+        return ok({
+            model: response.model,
+            text: response.choices[0].message.content,
+            finish_reason: response.choices[0].finish_reason,
+            usage: response.usage,
+        });
+    }
+    async createEmbeddings(rawArgs) {
+        const args = EmbeddingsSchema.parse(rawArgs);
+        const model = args.model ?? "nvidia/nv-embed-v1";
+        const response = await this.client.embeddings({
+            model,
+            input: args.input,
+            encoding_format: args.encoding_format,
+            truncate: args.truncate,
+        });
+        const inputs = Array.isArray(args.input) ? args.input : [args.input];
+        return ok({
+            model: response.model,
+            embeddings: response.data.map((d, i) => ({
+                index: d.index,
+                text: inputs[i],
+                dimensions: d.embedding.length,
+                embedding: d.embedding,
+            })),
+            usage: response.usage,
+        });
+    }
+    async rerankPassages(rawArgs) {
+        const args = RerankSchema.parse(rawArgs);
+        const model = args.model ?? "nvidia/nv-rerankqa-mistral-4b-v3";
+        // Normalize passages
+        const passages = args.passages.map((p) => typeof p === "string" ? { text: p } : { text: p.text });
+        const response = await this.client.rerank({
+            model,
+            query: args.query,
+            passages,
+            truncate: args.truncate,
+        });
+        let rankings = response.rankings;
+        if (args.top_k) {
+            rankings = rankings.slice(0, args.top_k);
+        }
+        return ok({
+            query: args.query,
+            model,
+            rankings: rankings.map((r) => ({
+                rank: rankings.indexOf(r) + 1,
+                original_index: r.index,
+                score: r.logit,
+                text: r.passage.text,
+            })),
+            usage: response.usage,
+        });
+    }
+    async functionCalling(rawArgs) {
+        const args = FunctionCallingSchema.parse(rawArgs);
+        const model = args.model ?? config.DEFAULT_MODEL;
+        const response = await this.client.chatCompletion({
+            model,
+            messages: args.messages.map(msg => ({
+                role: msg.role,
+                content: msg.content
+            })),
+            tools: args.tools.map(tool => ({
+                type: tool.type,
+                function: {
+                    name: tool.function.name,
+                    description: tool.function.description,
+                    parameters: tool.function.parameters
+                }
+            })),
+            tool_choice: args.tool_choice,
+            temperature: args.temperature,
+            max_tokens: args.max_tokens,
+        });
+        const choice = response.choices[0];
+        return ok({
+            model: response.model,
+            finish_reason: choice.finish_reason,
+            message: choice.message.content,
+            tool_calls: choice.message.tool_calls?.map((tc) => ({
+                id: tc.id,
+                function: tc.function.name,
+                arguments: (() => {
+                    try {
+                        return JSON.parse(tc.function.arguments);
+                    }
+                    catch {
+                        return tc.function.arguments;
+                    }
+                })(),
+            })),
+            usage: response.usage,
+        });
+    }
+    async listModels(rawArgs) {
+        const args = ListModelsSchema.parse(rawArgs);
+        let models;
+        if (args.category === "all") {
+            models = Object.values(NIM_MODELS);
+        }
+        else {
+            models = getModelsByCategory(args.category);
+        }
+        return ok({
+            total: models.length,
+            category: args.category,
+            models: models.map((m) => ({
+                id: m.id,
+                name: m.name,
+                category: m.category,
+                description: m.description,
+                context_length: m.contextLength,
+                supports_streaming: m.supportsStreaming,
+                supports_function_calling: m.supportsFunctionCalling,
+                supports_vision: m.supportsVision ?? false,
+            })),
+        });
+    }
+    async getModelInfo(rawArgs) {
+        const args = ModelInfoSchema.parse(rawArgs);
+        const model = getModel(args.model_id);
+        if (!model) {
+            return err(`Model '${args.model_id}' not found in catalog. Use list_models to see available models.`);
+        }
+        return ok({
+            id: model.id,
+            name: model.name,
+            category: model.category,
+            description: model.description,
+            context_length: model.contextLength,
+            supports_streaming: model.supportsStreaming,
+            supports_function_calling: model.supportsFunctionCalling,
+            supports_vision: model.supportsVision ?? false,
+            recommended_use_cases: getUseCases(model.category),
+        });
+    }
+}
+function getUseCases(category) {
+    const useCases = {
+        language: ["Chat", "Summarization", "Translation", "Q&A", "Content generation"],
+        embedding: ["Semantic search", "RAG", "Clustering", "Similarity comparison"],
+        reranking: ["RAG re-ranking", "Search quality improvement", "Relevance scoring"],
+        vision: ["Image understanding", "Visual Q&A", "Document analysis"],
+        code: ["Code generation", "Code review", "Debugging", "Documentation"],
+        multimodal: ["Multi-modal tasks", "Image + text processing"],
+    };
+    return useCases[category] ?? [];
+}

package/dist/index.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/index.js

@@ -0,0 +1,107 @@
+#!/usr/bin/env node
+import dotenv from "dotenv";
+dotenv.config();
+import { Server } from "@modelcontextprotocol/sdk/server/index.js";
+import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import { CallToolRequestSchema, ListToolsRequestSchema, InitializeRequestSchema, } from "@modelcontextprotocol/sdk/types.js";
+import { getConfig } from "./config.js";
+import { logger } from "./logger.js";
+import { NIMClient } from "./client.js";
+import { ToolHandlers } from "./handlers.js";
+import { TOOL_DEFINITIONS } from "./tools.js";
+const config = getConfig();
+// ─── Graceful Shutdown ────────────────────────────────────────────────────────
+function setupShutdownHandlers(server) {
+    const shutdown = async (signal) => {
+        logger.info(`Received ${signal}, shutting down gracefully...`);
+        try {
+            await server.close();
+            logger.info("Server closed successfully");
+            process.exit(0);
+        }
+        catch (err) {
+            logger.error("Error during shutdown", { error: err });
+            process.exit(1);
+        }
+    };
+    process.on("SIGINT", () => shutdown("SIGINT"));
+    process.on("SIGTERM", () => shutdown("SIGTERM"));
+    process.on("uncaughtException", (error) => {
+        logger.error("Uncaught exception", { error });
+        process.exit(1);
+    });
+    process.on("unhandledRejection", (reason) => {
+        logger.error("Unhandled rejection", { reason });
+        process.exit(1);
+    });
+}
+// ─── Main ─────────────────────────────────────────────────────────────────────
+async function main() {
+    logger.info("Starting NVIDIA NIM MCP Server", {
+        name: config.MCP_SERVER_NAME,
+        version: config.MCP_SERVER_VERSION,
+        baseUrl: config.NVIDIA_NIM_BASE_URL,
+        defaultModel: config.DEFAULT_MODEL,
+        logLevel: config.LOG_LEVEL,
+        maxRPM: config.MAX_REQUESTS_PER_MINUTE,
+        maxTokens: config.MAX_TOKENS_PER_REQUEST,
+        timeout: config.REQUEST_TIMEOUT_MS,
+        maxRetries: config.MAX_RETRIES,
+    });
+    // Initialize NIM client
+    const nimClient = new NIMClient();
+    const toolHandlers = new ToolHandlers(nimClient);
+    // Create MCP server
+    const server = new Server({
+        name: config.MCP_SERVER_NAME,
+        version: config.MCP_SERVER_VERSION,
+    }, {
+        capabilities: {
+            tools: {},
+        },
+    });
+    setupShutdownHandlers(server);
+    // ─── Handler: Initialize ─────────────────────────────────────────────────
+    server.setRequestHandler(InitializeRequestSchema, async (request) => {
+        logger.info("Client initialized", {
+            clientName: request.params.clientInfo?.name,
+            clientVersion: request.params.clientInfo?.version,
+        });
+        return {
+            protocolVersion: "2024-11-05",
+            capabilities: { tools: {} },
+            serverInfo: {
+                name: config.MCP_SERVER_NAME,
+                version: config.MCP_SERVER_VERSION,
+            },
+        };
+    });
+    // ─── Handler: List Tools ──────────────────────────────────────────────────
+    server.setRequestHandler(ListToolsRequestSchema, async () => {
+        logger.debug("Listing tools");
+        return { tools: TOOL_DEFINITIONS };
+    });
+    // ─── Handler: Call Tool ───────────────────────────────────────────────────
+    server.setRequestHandler(CallToolRequestSchema, async (request) => {
+        const { name, arguments: args } = request.params;
+        logger.info(`Tool call received: ${name}`, {
+            hasArgs: args !== undefined,
+        });
+        const result = await toolHandlers.handle(name, args ?? {});
+        if (result.isError) {
+            logger.warn(`Tool ${name} returned an error`, {
+                error: result.content[0]?.text,
+            });
+        }
+        return result;
+    });
+    // ─── Start Transport ──────────────────────────────────────────────────────
+    const transport = new StdioServerTransport();
+    logger.info("Connecting MCP server via stdio transport...");
+    await server.connect(transport);
+    logger.info("NVIDIA NIM MCP Server is ready and listening");
+}
+main().catch((err) => {
+    logger.error("Fatal error during startup", { error: err });
+    process.exit(1);
+});

package/dist/logger.d.ts

@@ -0,0 +1,9 @@
+import winston from "winston";
+import { Config } from "./config.js";
+export interface LoggerOptions {
+    logLevel: Config['LOG_LEVEL'];
+    isProduction: boolean;
+    mcpServerVersion: string;
+}
+export declare let logger: winston.Logger;
+export declare function initLogger(options: LoggerOptions): winston.Logger;