fix(tools): improve json parsing (#1356)

* fix(tools): improve json parsing

* lint

Author: nsarrazin

src/lib/server/textGeneration/tools.ts

@@ -1,6 +1,5 @@
 import { ToolResultStatus, type ToolCall, type ToolResult } from "$lib/types/Tool";
 import { v4 as uuidV4 } from "uuid";
-import JSON5 from "json5";
 import type { BackendTool, BackendToolContext } from "../tools";
 import {
 	MessageToolUpdateType,
@@ -15,7 +14,7 @@ import directlyAnswer from "../tools/directlyAnswer";
 import websearch from "../tools/web/search";
 import { z } from "zod";
 import { logger } from "../logger";
-import { toolHasName } from "../tools/utils";
+import { extractJson, toolHasName } from "../tools/utils";
 import type { MessageFile } from "$lib/types/Message";
 import { mergeAsyncGenerators } from "$lib/utils/mergeAsyncGenerators";
 import { MetricsServer } from "../metrics";
@@ -143,33 +142,23 @@ export async function* runTools(
 		// look for a code blocks of ```json and parse them
 		// if they're valid json, add them to the calls array
 		if (output.generated_text) {
-			if (!output.generated_text.endsWith("```")) {
-				output.generated_text = output.generated_text + "```";
-			}
-			const codeBlocks = Array.from(output.generated_text.matchAll(/```json\n(.*?)```/gs))
-				.map(([, block]) => block)
-				// remove trailing comma
-				.map((block) => block.trim().replace(/,$/, ""));
-			if (codeBlocks.length === 0) continue;
-			// grab only the capture group from the regex match
-			for (const block of codeBlocks) {
-				// make it an array if it's not already
-				let call = JSON5.parse(block);
-				if (!Array.isArray(call)) {
-					call = [call];
-				}
-
-				try {
-					calls.push(...call.filter(isExternalToolCall).map(externalToToolCall).filter(Boolean));
-				} catch (e) {
-					logger.error(e, "Error while parsing tool calls, please retry");
-					// error parsing the calls
-					yield {
-						type: MessageUpdateType.Status,
-						status: MessageUpdateStatus.Error,
-						message: "Error while parsing tool calls, please retry",
-					};
-				}
+			logger.info(output.generated_text);
+			try {
+				const rawCalls = await extractJson(output.generated_text);
+				const newCalls = rawCalls
+					.filter(isExternalToolCall)
+					.map(externalToToolCall)
+					.filter((call) => call !== undefined) as ToolCall[];
+
+				calls.push(...newCalls);
+			} catch (e) {
+				logger.error(e, "Error while parsing tool calls, please retry");
+				// error parsing the calls
+				yield {
+					type: MessageUpdateType.Status,
+					status: MessageUpdateStatus.Error,
+					message: "Error while parsing tool calls, please retry",
+				};
 			}
 		}
 	}

src/lib/server/tools/utils.ts

@@ -1,6 +1,7 @@
 import { env } from "$env/dynamic/private";
 import { Client } from "@gradio/client";
 import { SignJWT } from "jose";
+import JSON5 from "json5";
 
 export type GradioImage = {
 	path: string;
@@ -52,3 +53,41 @@ export async function getIpToken(ip: string, username?: string) {
 }
 
 export { toolHasName } from "$lib/utils/tools";
+
+export async function extractJson(text: string): Promise<unknown[]> {
+	const calls: string[] = [];
+
+	let codeBlocks = Array.from(text.matchAll(/```json\n(.*?)```/gs))
+		.map(([, block]) => block)
+		// remove trailing comma
+		.map((block) => block.trim().replace(/,$/, ""));
+
+	// if there is no code block, try to find the first json object
+	// by trimming the string and trying to parse with JSON5
+	if (codeBlocks.length === 0) {
+		const start = [text.indexOf("["), text.indexOf("{")]
+			.filter((i) => i !== -1)
+			.reduce((a, b) => Math.max(a, b), -Infinity);
+		const end = [text.lastIndexOf("]"), text.lastIndexOf("}")]
+			.filter((i) => i !== -1)
+			.reduce((a, b) => Math.min(a, b), Infinity);
+
+		if (start === -Infinity || end === Infinity) {
+			return [""];
+		}
+
+		const json = text.substring(start, end + 1);
+		codeBlocks = [json];
+	}
+
+	// grab only the capture group from the regex match
+	for (const block of codeBlocks) {
+		// make it an array if it's not already
+		let call = JSON5.parse(block);
+		if (!Array.isArray(call)) {
+			call = [call];
+		}
+		calls.push(call);
+	}
+	return calls.flat();
+}