Simpler loop

Author: julien-c

packages/mcp-client/cli.ts

@@ -35,6 +35,10 @@ if (process.env.EXPERIMENTAL_HF_MCP_SERVER) {
 	});
 }
 
+let SAMPLE_INPUT = process.env.USE_SAMPLE_INPUT
+	? `generate a haiku about Hugging Face and save it to a file name hf.txt on my Desktop`
+	: undefined;
+
 async function main() {
 	if (!process.env.HF_TOKEN) {
 		console.error(`a valid HF_TOKEN must be present in the env`);

packages/mcp-client/src/McpClient.ts

@@ -84,31 +84,24 @@ export class McpClient {
 			tool_choice: "auto",
 		});
 
-		const firstChunkResult = await stream.next();
-		if (firstChunkResult.done) {
-			return;
-		}
-		const firstChunk = firstChunkResult.value;
-		const firstToolCalls = firstChunk.choices[0]?.delta.tool_calls;
-		if ((!firstToolCalls || firstToolCalls.length === 0) && opts.exitIfFirstChunkNoTool) {
-			return;
-		}
-		yield firstChunk;
-		debug(firstChunk.choices[0]);
 		const message = {
-			role: firstChunk.choices[0].delta.role,
-			content: firstChunk.choices[0].delta.content,
+			role: "unknown",
+			content: "",
 		} satisfies ChatCompletionInputMessage;
-
 		const finalToolCalls: Record<number, ChatCompletionStreamOutputDeltaToolCall> = {};
+		let numOfChunks = 0;
 
 		for await (const chunk of stream) {
 			yield chunk;
 			debug(chunk.choices[0]);
+			numOfChunks++;
 			const delta = chunk.choices[0]?.delta;
 			if (!delta) {
 				continue;
 			}
+			if (delta.role) {
+				message.role = delta.role;
+			}
 			if (delta.content) {
 				message.content += delta.content;
 			}
@@ -124,6 +117,10 @@ export class McpClient {
 					finalToolCalls[toolCall.index].function.arguments += toolCall.function.arguments;
 				}
 			}
+			if (opts.exitIfFirstChunkNoTool && numOfChunks <= 2 && Object.keys(finalToolCalls).length === 0) {
+				/// If no tool is present in chunk number 1 or 2, exit.
+				return;
+			}
 		}
 
 		messages.push(message);