Support arbitrary (local or remote) endpoint

Author: julien-c

packages/mcp-client/cli.ts

@@ -11,6 +11,7 @@ import { version as packageVersion } from "./package.json";
 
 const MODEL_ID = process.env.MODEL_ID ?? "Qwen/Qwen2.5-72B-Instruct";
 const PROVIDER = (process.env.PROVIDER as InferenceProvider) ?? "nebius";
+const ENDPOINT_URL = process.env.ENDPOINT_URL;
 
 const SERVERS: StdioServerParameters[] = [
 	{
@@ -48,12 +49,21 @@ async function main() {
 		process.exit(1);
 	}
 
-	const agent = new Agent({
-		provider: PROVIDER,
-		model: MODEL_ID,
-		apiKey: process.env.HF_TOKEN,
-		servers: SERVERS,
-	});
+	const agent = new Agent(
+		ENDPOINT_URL
+			? {
+					baseUrl: ENDPOINT_URL,
+					model: MODEL_ID,
+					apiKey: process.env.HF_TOKEN,
+					servers: SERVERS,
+			  }
+			: {
+					provider: PROVIDER,
+					model: MODEL_ID,
+					apiKey: process.env.HF_TOKEN,
+					servers: SERVERS,
+			  }
+	);
 
 	const rl = readline.createInterface({ input: stdin, output: stdout });
 	let abortController = new AbortController();

packages/mcp-client/src/Agent.ts

@@ -49,18 +49,28 @@ export class Agent extends McpClient {
 
 	constructor({
 		provider,
+		baseUrl,
 		model,
 		apiKey,
 		servers,
 		prompt,
-	}: {
-		provider: InferenceProvider;
+	}: (
+		| {
+				provider: InferenceProvider;
+				baseUrl?: undefined;
+		  }
+		| {
+				baseUrl: string;
+				provider?: undefined;
+		  }
+	) & {
 		model: string;
 		apiKey: string;
 		servers: StdioServerParameters[];
 		prompt?: string;
 	}) {
-		super({ provider, model, apiKey });
+		super(provider ? { provider, baseUrl, model, apiKey } : { provider, baseUrl, model, apiKey });
+		/// ^This shenanigan is just here to please an overzealous TS type-checker.
 		this.servers = servers;
 		this.messages = [
 			{

packages/mcp-client/src/McpClient.ts

@@ -2,7 +2,7 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import type { StdioServerParameters } from "@modelcontextprotocol/sdk/client/stdio.js";
 import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
 import { InferenceClient } from "@huggingface/inference";
-import type { InferenceProvider } from "@huggingface/inference";
+import type { InferenceClientEndpoint, InferenceProvider } from "@huggingface/inference";
 import type {
 	ChatCompletionInputMessage,
 	ChatCompletionInputTool,
@@ -22,14 +22,32 @@ export interface ChatCompletionInputMessageTool extends ChatCompletionInputMessa
 }
 
 export class McpClient {
-	protected client: InferenceClient;
-	protected provider: string;
+	protected client: InferenceClient | InferenceClientEndpoint;
+	protected provider: string | undefined;
+
 	protected model: string;
 	private clients: Map<ToolName, Client> = new Map();
 	public readonly availableTools: ChatCompletionInputTool[] = [];
 
-	constructor({ provider, model, apiKey }: { provider: InferenceProvider; model: string; apiKey: string }) {
-		this.client = new InferenceClient(apiKey);
+	constructor({
+		provider,
+		baseUrl,
+		model,
+		apiKey,
+	}: (
+		| {
+				provider: InferenceProvider;
+				baseUrl?: undefined;
+		  }
+		| {
+				baseUrl: string;
+				provider?: undefined;
+		  }
+	) & {
+		model: string;
+		apiKey: string;
+	}) {
+		this.client = baseUrl ? new InferenceClient(apiKey).endpoint(baseUrl) : new InferenceClient(apiKey);
 		this.provider = provider;
 		this.model = model;
 	}