improvements

Author: julien-c

packages/mcp-client/cli.ts

@@ -4,13 +4,11 @@ import { join } from "node:path";
 import { homedir } from "node:os";
 import { Agent } from "./src";
 import type { StdioServerParameters } from "@modelcontextprotocol/sdk/client/stdio.js";
+import { ANSI } from "./src/utils";
+import type { InferenceProvider } from "@huggingface/inference";
 
-const ANSI = {
-	BLUE: "\x1b[34m",
-	GREEN: "\x1b[32m",
-	RED: "\x1b[31m",
-	RESET: "\x1b[0m",
-};
+const MODEL_ID = process.env.MODEL_ID ?? "Qwen/Qwen2.5-72B-Instruct";
+const PROVIDER = (process.env.PROVIDER as InferenceProvider) ?? "together";
 
 const SERVERS: StdioServerParameters[] = [
 	{
@@ -35,8 +33,8 @@ async function main() {
 	}
 
 	const agent = new Agent({
-		provider: "together",
-		model: "Qwen/Qwen2.5-72B-Instruct",
+		provider: PROVIDER,
+		model: MODEL_ID,
 		apiKey: process.env.HF_TOKEN,
 		servers: SERVERS,
 	});
@@ -58,7 +56,7 @@ async function main() {
 
 	while (true) {
 		const input = await rl.question("> ");
-		for await (const response of agent.processSingleTurn(input)) {
+		for await (const response of agent.run(input)) {
 			if ("choices" in response) {
 				stdout.write(response.choices[0].message.content ?? "");
 				stdout.write("\n\n");

packages/mcp-client/src/Agent.ts

@@ -2,7 +2,9 @@ import type { InferenceProvider } from "@huggingface/inference";
 import type { ChatCompletionInputMessageTool } from "./McpClient";
 import { McpClient } from "./McpClient";
 import type { ChatCompletionInputMessage, ChatCompletionOutput } from "@huggingface/tasks";
+import type { ChatCompletionInputTool } from "@huggingface/tasks/src/tasks/chat-completion/inference";
 import type { StdioServerParameters } from "@modelcontextprotocol/sdk/client/stdio";
+import { debug } from "./utils";
 
 const DEFAULT_SYSTEM_PROMPT = `
 You are an agent - please keep going until the user’s query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved, or if you need more info from the user to solve the problem.
@@ -12,6 +14,23 @@ If you are not sure about anything pertaining to the user’s request, use your
 You MUST plan extensively before each function call, and reflect extensively on the outcomes of the previous function calls. DO NOT do this entire process by making function calls only, as this can impair your ability to solve the problem and think insightfully.
 `.trim();
 
+/**
+ * Max number of tool calling + chat completion steps in response to a single user query.
+ */
+const MAX_NUM_TURNS = 10;
+
+const taskCompletionTool: ChatCompletionInputTool = {
+	type: "function",
+	function: {
+		name: "task_complete",
+		description: "Call this tool when the task given by the user is complete",
+		parameters: {
+			type: "object",
+			properties: {},
+		},
+	},
+};
+
 export class Agent extends McpClient {
 	private readonly servers: StdioServerParameters[];
 	protected messages: ChatCompletionInputMessage[];
@@ -41,14 +60,22 @@ export class Agent extends McpClient {
 		return this.addMcpServers(this.servers);
 	}
 
-	async *processSingleTurn(input: string): AsyncGenerator<ChatCompletionOutput | ChatCompletionInputMessageTool> {
+	async *run(input: string): AsyncGenerator<ChatCompletionOutput | ChatCompletionInputMessageTool> {
 		this.messages.push({
 			role: "user",
 			content: input,
 		});
 
-		while (this.messages.at(-1)?.role !== "assistant") {
-			yield* this.processSingleTurnWithTools(this.messages);
+		let i = 0;
+		while (
+			!(this.messages.at(-1)?.role === "tool" && this.messages.at(-1)?.name === taskCompletionTool.function.name) &&
+			i < MAX_NUM_TURNS
+		) {
+			yield* this.processSingleTurnWithTools(this.messages, {
+				taskCompletionTool,
+			});
+			i++;
+			debug("current role", this.messages.at(-1)?.role);
 		}
 	}
 }

packages/mcp-client/src/McpClient.ts

@@ -70,15 +70,16 @@ export class McpClient {
 	}
 
 	async *processSingleTurnWithTools(
-		messages: ChatCompletionInputMessage[]
+		messages: ChatCompletionInputMessage[],
+		opts: { taskCompletionTool?: ChatCompletionInputTool } = {}
 	): AsyncGenerator<ChatCompletionOutput | ChatCompletionInputMessageTool> {
 		debug("start of single turn");
 
 		const response = await this.client.chatCompletion({
 			provider: this.provider,
 			model: this.model,
 			messages,
-			tools: this.availableTools,
+			tools: opts.taskCompletionTool ? [...this.availableTools, opts.taskCompletionTool] : this.availableTools,
 			tool_choice: "auto",
 		});
 
@@ -100,6 +101,10 @@ export class McpClient {
 				content: "",
 				name: toolName,
 			};
+			if (toolName === opts.taskCompletionTool?.function.name) {
+				messages.push(message);
+				return yield message;
+			}
 			/// Get the appropriate session for this tool
 			const client = this.clients.get(toolName);
 			if (client) {

packages/mcp-client/src/utils.ts

@@ -3,3 +3,10 @@ export function debug(...args: unknown[]): void {
 		console.debug(args);
 	}
 }
+
+export const ANSI = {
+	BLUE: "\x1b[34m",
+	GREEN: "\x1b[32m",
+	RED: "\x1b[31m",
+	RESET: "\x1b[0m",
+};