Add Smol agent implem

Author: julien-c

packages/mcp-client/cli.ts

@@ -5,15 +5,9 @@ import { homedir } from "node:os";
 import { Agent } from "./src";
 import type { StdioServerParameters } from "@modelcontextprotocol/sdk/client/stdio.js";
 
-async function* boboGenerator(message: string) {
-	for (let i = 0; i < 6; i++) {
-		yield "bobo ";
-		await new Promise((resolve) => setTimeout(resolve, 400));
-	}
-}
-
 const ANSI = {
 	BLUE: "\x1b[34m",
+	GREEN: "\x1b[32m",
 	RED: "\x1b[31m",
 	RESET: "\x1b[0m",
 };
@@ -64,8 +58,18 @@ async function main() {
 
 	while (true) {
 		const input = await rl.question("> ");
-		for await (const bobo of boboGenerator(input)) {
-			console.log(bobo);
+		for await (const response of agent.processSingleTurn(input)) {
+			if ("choices" in response) {
+				stdout.write(response.choices[0].message.content ?? "");
+				stdout.write("\n\n");
+			} else {
+				/// Tool call info
+				stdout.write(ANSI.GREEN);
+				stdout.write(`Tool[${response.name}] ${response.tool_call_id}\n`);
+				stdout.write(response.content);
+				stdout.write(ANSI.RESET);
+				stdout.write("\n\n");
+			}
 		}
 	}
 }

packages/mcp-client/src/Agent.ts

@@ -1,6 +1,7 @@
 import type { InferenceProvider } from "@huggingface/inference";
+import type { ChatCompletionInputMessageTool } from "./McpClient";
 import { McpClient } from "./McpClient";
-import type { ChatCompletionInputMessage } from "@huggingface/tasks";
+import type { ChatCompletionInputMessage, ChatCompletionOutput } from "@huggingface/tasks";
 import type { StdioServerParameters } from "@modelcontextprotocol/sdk/client/stdio";
 
 const DEFAULT_SYSTEM_PROMPT = `
@@ -40,18 +41,14 @@ export class Agent extends McpClient {
 		return this.addMcpServers(this.servers);
 	}
 
-	async processUserMessage(query: string) {
+	async *processSingleTurn(input: string): AsyncGenerator<ChatCompletionOutput | ChatCompletionInputMessageTool> {
 		this.messages.push({
 			role: "user",
-			content: query,
+			content: input,
 		});
 
-		const stream = await this.client.chatCompletionStream({
-			provider: this.provider,
-			model: this.model,
-			messages,
-			tools: this.availableTools,
-			tool_choice: "auto",
-		});
+		while (this.messages.at(-1)?.role !== "assistant") {
+			yield* this.processSingleTurnWithTools(this.messages);
+		}
 	}
 }

packages/mcp-client/src/McpClient.ts

@@ -13,10 +13,17 @@ import { debug } from "./utils";
 
 type ToolName = string;
 
+export interface ChatCompletionInputMessageTool extends ChatCompletionInputMessage {
+	role: "tool";
+	tool_call_id: string;
+	content: string;
+	name?: string;
+}
+
 export class McpClient {
-	private client: InferenceClient;
-	private provider: string;
-	private model: string;
+	protected client: InferenceClient;
+	protected provider: string;
+	protected model: string;
 	private clients: Map<ToolName, Client> = new Map();
 	public readonly availableTools: ChatCompletionInputTool[] = [];
 
@@ -62,13 +69,10 @@ export class McpClient {
 		);
 	}
 
-	async processQuery(query: string): Promise<ChatCompletionOutput> {
-		const messages: ChatCompletionInputMessage[] = [
-			{
-				role: "user",
-				content: query,
-			},
-		];
+	async *processSingleTurnWithTools(
+		messages: ChatCompletionInputMessage[]
+	): AsyncGenerator<ChatCompletionOutput | ChatCompletionInputMessageTool> {
+		debug("start of single turn");
 
 		const response = await this.client.chatCompletion({
 			provider: this.provider,
@@ -80,43 +84,41 @@ export class McpClient {
 
 		const toolCalls = response.choices[0].message.tool_calls;
 		if (!toolCalls || toolCalls.length === 0) {
-			return response;
+			messages.push({
+				role: response.choices[0].message.role,
+				content: response.choices[0].message.content,
+			});
+			return yield response;
 		}
 		for (const toolCall of toolCalls) {
 			const toolName = toolCall.function.name;
 			const toolArgs = JSON.parse(toolCall.function.arguments);
 
+			const message: ChatCompletionInputMessageTool = {
+				role: "tool",
+				tool_call_id: toolCall.id,
+				content: "",
+				name: toolName,
+			};
 			/// Get the appropriate session for this tool
 			const client = this.clients.get(toolName);
 			if (client) {
 				const result = await client.callTool({ name: toolName, arguments: toolArgs });
-				messages.push({
-					tool_call_id: toolCall.id,
-					role: "tool",
-					name: toolName,
-					content: (result.content as Array<{ text: string }>)[0].text,
-				});
+				message.content = (result.content as Array<{ text: string }>)[0].text;
 			} else {
-				messages.push({
-					tool_call_id: toolCall.id,
-					role: "tool",
-					name: toolName,
-					content: `Error: No session found for tool: ${toolName}`,
-				});
+				message.content = `Error: No session found for tool: ${toolName}`;
 			}
+			messages.push(message);
+			yield message;
 		}
-
-		const enrichedResponse = await this.client.chatCompletion({
-			provider: this.provider,
-			model: this.model,
-			messages,
-		});
-
-		return enrichedResponse;
 	}
 
 	async cleanup(): Promise<void> {
 		const clients = new Set(this.clients.values());
 		await Promise.all([...clients].map((client) => client.close()));
 	}
+
+	async [Symbol.dispose](): Promise<void> {
+		return this.cleanup();
+	}
 }

packages/mcp-client/src/utils.ts

@@ -1,5 +1,5 @@
 export function debug(...args: unknown[]): void {
 	if (process.env.DEBUG) {
-		console.log(args);
+		console.debug(args);
 	}
 }