improvements

Author: julien-c

packages/mcp-client/src/Agent.ts

@@ -30,6 +30,18 @@ const taskCompletionTool: ChatCompletionInputTool = {
 		},
 	},
 };
+const askQuestionTool: ChatCompletionInputTool = {
+	type: "function",
+	function: {
+		name: "ask_question",
+		description: "Ask a question to the user to get more info required to solve or clarify their problem.",
+		parameters: {
+			type: "object",
+			properties: {},
+		},
+	},
+};
+const exitLoopTools = [taskCompletionTool, askQuestionTool];
 
 export class Agent extends McpClient {
 	private readonly servers: StdioServerParameters[];
@@ -66,16 +78,31 @@ export class Agent extends McpClient {
 			content: input,
 		});
 
-		let i = 0;
-		while (
-			!(this.messages.at(-1)?.role === "tool" && this.messages.at(-1)?.name === taskCompletionTool.function.name) &&
-			i < MAX_NUM_TURNS
-		) {
+		let total = 0;
+		let nextShouldCallTools = true;
+		while (true) {
 			yield* this.processSingleTurnWithTools(this.messages, {
-				taskCompletionTool,
+				exitLoopTools,
+				exitIfNoTool: total > 0 && nextShouldCallTools,
 			});
-			i++;
-			debug("current role", this.messages.at(-1)?.role);
+			total++;
+			// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+			const currentLast = this.messages.at(-1)!;
+			debug("current role", currentLast.role);
+			if (
+				currentLast.role === "tool" &&
+				currentLast.name &&
+				exitLoopTools.map((t) => t.function.name).includes(currentLast.name)
+			) {
+				return;
+			}
+			if (currentLast.role !== "tool" && total > MAX_NUM_TURNS) {
+				return;
+			}
+			if (currentLast.role !== "tool" && nextShouldCallTools) {
+				return;
+			}
+			nextShouldCallTools = currentLast.role !== "tool";
 		}
 	}
 }

packages/mcp-client/src/McpClient.ts

@@ -71,20 +71,23 @@ export class McpClient {
 
 	async *processSingleTurnWithTools(
 		messages: ChatCompletionInputMessage[],
-		opts: { taskCompletionTool?: ChatCompletionInputTool } = {}
+		opts: { exitLoopTools?: ChatCompletionInputTool[]; exitIfNoTool?: boolean } = {}
 	): AsyncGenerator<ChatCompletionOutput | ChatCompletionInputMessageTool> {
 		debug("start of single turn");
 
 		const response = await this.client.chatCompletion({
 			provider: this.provider,
 			model: this.model,
 			messages,
-			tools: opts.taskCompletionTool ? [...this.availableTools, opts.taskCompletionTool] : this.availableTools,
+			tools: opts.exitLoopTools ? [...opts.exitLoopTools, ...this.availableTools] : this.availableTools,
 			tool_choice: "auto",
 		});
 
 		const toolCalls = response.choices[0].message.tool_calls;
 		if (!toolCalls || toolCalls.length === 0) {
+			if (opts.exitIfNoTool) {
+				return;
+			}
 			messages.push({
 				role: response.choices[0].message.role,
 				content: response.choices[0].message.content,
@@ -101,7 +104,7 @@ export class McpClient {
 				content: "",
 				name: toolName,
 			};
-			if (toolName === opts.taskCompletionTool?.function.name) {
+			if (opts.exitLoopTools?.map((t) => t.function.name).includes(toolName)) {
 				messages.push(message);
 				return yield message;
 			}