Simplify tool error handling (#1168)

* feat: simplify tool error handling

* feat: resolve linting and type errors

---------

Co-authored-by: Mishig <[email protected]>

src/lib/server/textGeneration/tools.ts

@@ -19,6 +19,7 @@ import { toolHasName } from "../tools/utils";
 import type { MessageFile } from "$lib/types/Message";
 import { mergeAsyncGenerators } from "$lib/utils/mergeAsyncGenerators";
 import { MetricsServer } from "../metrics";
+import { stringifyError } from "$lib/utils/stringifyError";
 
 function makeFilesPrompt(files: MessageFile[], fileMessageIndex: number): string {
 	if (files.length === 0) {
@@ -48,7 +49,7 @@ export function pickTools(
 	});
 }
 
-async function* runTool(
+async function* callTool(
 	ctx: BackendToolContext,
 	tools: BackendTool[],
 	call: ToolCall
@@ -72,47 +73,38 @@ async function* runTool(
 		uuid,
 		call,
 	};
+
 	try {
-		try {
-			const toolResult = yield* tool.call(call.parameters, ctx);
-			if (toolResult.status === ToolResultStatus.Error) {
-				yield {
-					type: MessageUpdateType.Tool,
-					subtype: MessageToolUpdateType.Error,
-					uuid,
-					message: toolResult.message,
-				};
-			} else {
-				yield {
-					type: MessageUpdateType.Tool,
-					subtype: MessageToolUpdateType.Result,
-					uuid,
-					result: { ...toolResult, call } as ToolResult,
-				};
-			}
+		const toolResult = yield* tool.call(call.parameters, ctx);
 
-			MetricsServer.getMetrics().tool.toolUseDuration.observe(
-				{ tool: call.name },
-				Date.now() - startTime
-			);
-
-			return { ...toolResult, call } as ToolResult;
-		} catch (e) {
-			MetricsServer.getMetrics().tool.toolUseCountError.inc({ tool: call.name });
-			yield {
-				type: MessageUpdateType.Tool,
-				subtype: MessageToolUpdateType.Error,
-				uuid,
-				message: e instanceof Error ? e.message : String(e),
-			};
-		}
-	} catch (cause) {
+		yield {
+			type: MessageUpdateType.Tool,
+			subtype: MessageToolUpdateType.Result,
+			uuid,
+			result: { ...toolResult, call } as ToolResult,
+		};
+
+		MetricsServer.getMetrics().tool.toolUseDuration.observe(
+			{ tool: call.name },
+			Date.now() - startTime
+		);
+
+		return { ...toolResult, call } as ToolResult;
+	} catch (error) {
 		MetricsServer.getMetrics().tool.toolUseCountError.inc({ tool: call.name });
-		console.error(Error(`Failed while running tool ${call.name}`), { cause });
+		logger.error(error, `Failed while running tool ${call.name}`);
+
+		yield {
+			type: MessageUpdateType.Tool,
+			subtype: MessageToolUpdateType.Error,
+			uuid,
+			message: stringifyError(error),
+		};
+
 		return {
 			call,
 			status: ToolResultStatus.Error,
-			message: cause instanceof Error ? cause.message : String(cause),
+			message: stringifyError(error),
 		};
 	}
 }
@@ -160,12 +152,12 @@ export async function* runTools(
 					calls.push(
 						...JSON5.parse(block).filter(isExternalToolCall).map(externalToToolCall).filter(Boolean)
 					);
-				} catch (cause) {
+				} catch (e) {
 					// error parsing the calls
 					yield {
 						type: MessageUpdateType.Status,
 						status: MessageUpdateStatus.Error,
-						message: cause instanceof Error ? cause.message : String(cause),
+						message: stringifyError(e),
 					};
 				}
 			}
@@ -179,7 +171,7 @@ export async function* runTools(
 
 	const toolContext: BackendToolContext = { conv, messages, preprompt, assistant, ip, username };
 	const toolResults: (ToolResult | undefined)[] = yield* mergeAsyncGenerators(
-		calls.map((call) => runTool(toolContext, tools, call))
+		calls.map((call) => callTool(toolContext, tools, call))
 	);
 	return toolResults.filter((result): result is ToolResult => result !== undefined);
 }

src/lib/server/tools/calculator.ts

@@ -1,4 +1,3 @@
-import { ToolResultStatus } from "$lib/types/Tool";
 import type { BackendTool } from ".";
 import vm from "node:vm";
 
@@ -22,14 +21,10 @@ const calculator: BackendTool = {
 			const query = blocks[blocks.length - 1].replace(/[^-()\d/*+.]/g, "");
 
 			return {
-				status: ToolResultStatus.Success,
 				outputs: [{ calculator: `${query} = ${vm.runInNewContext(query)}` }],
 			};
-		} catch (e) {
-			return {
-				status: ToolResultStatus.Error,
-				message: "Invalid expression",
-			};
+		} catch (cause) {
+			throw Error("Invalid expression", { cause });
 		}
 	},
 };

src/lib/server/tools/directlyAnswer.ts

@@ -1,4 +1,3 @@
-import { ToolResultStatus } from "$lib/types/Tool";
 import type { BackendTool } from ".";
 
 const directlyAnswer: BackendTool = {
@@ -10,7 +9,6 @@ const directlyAnswer: BackendTool = {
 	parameterDefinitions: {},
 	async *call() {
 		return {
-			status: ToolResultStatus.Success,
 			outputs: [],
 			display: false,
 		};

src/lib/server/tools/documentParser.ts

@@ -1,5 +1,4 @@
 import type { BackendTool } from ".";
-import { ToolResultStatus } from "$lib/types/Tool";
 import { callSpace, getIpToken } from "./utils";
 import { downloadFile } from "$lib/server/files/downloadFile";
 
@@ -29,18 +28,8 @@ const documentParser: BackendTool = {
 
 		const message = messages[fileMessageIndex];
 		const files = message?.files ?? [];
-		if (!files || files.length === 0) {
-			return {
-				status: ToolResultStatus.Error,
-				message: "User did not provide a pdf to parse",
-			};
-		}
-		if (fileIndex >= files.length) {
-			return {
-				status: ToolResultStatus.Error,
-				message: "Model provided an invalid file index",
-			};
-		}
+		if (!files || files.length === 0) throw Error("User did not provide a pdf to parse");
+		if (fileIndex >= files.length) throw Error("Model provided an invalid file index");
 
 		const file = files[fileIndex];
 		const fileBlob = await downloadFile(files[fileIndex].value, conv._id)
@@ -62,7 +51,6 @@ const documentParser: BackendTool = {
 			documentMarkdown = documentMarkdown.slice(0, 30_000) + "\n\n... (truncated)";
 		}
 		return {
-			status: ToolResultStatus.Success,
 			outputs: [{ [file.name]: documentMarkdown }],
 			display: false,
 		};

src/lib/server/tools/images/editing.ts

@@ -1,6 +1,5 @@
 import type { BackendTool } from "..";
 import { uploadFile } from "../../files/uploadFile";
-import { ToolResultStatus } from "$lib/types/Tool";
 import { MessageUpdateType } from "$lib/types/MessageUpdate";
 import { callSpace, getIpToken, type GradioImage } from "../utils";
 import { downloadFile } from "$lib/server/files/downloadFile";
@@ -44,23 +43,10 @@ const imageEditing: BackendTool = {
 
 		const message = messages[fileMessageIndex];
 		const images = message?.files ?? [];
-		if (!images || images.length === 0) {
-			return {
-				status: ToolResultStatus.Error,
-				message: "User did not provide an image to edit.",
-			};
-		}
-		if (fileIndex >= images.length) {
-			return {
-				status: ToolResultStatus.Error,
-				message: "Model provided an invalid file index",
-			};
-		}
+		if (!images || images.length === 0) throw Error("User did not provide an image to edit");
+		if (fileIndex >= images.length) throw Error("Model provided an invalid file index");
 		if (!images[fileIndex].mime.startsWith("image/")) {
-			return {
-				status: ToolResultStatus.Error,
-				message: "Model provided a file indx which is not an image",
-			};
+			throw Error("Model provided a file idex which is not an image");
 		}
 
 		// todo: should handle multiple images
@@ -96,7 +82,6 @@ const imageEditing: BackendTool = {
 		};
 
 		return {
-			status: ToolResultStatus.Success,
 			outputs: [
 				{
 					imageEditing: `An image has been generated for the following prompt: "${prompt}". Answer as if the user can already see the image. Do not try to insert the image or to add space for it. The user can already see the image. Do not try to describe the image as you the model cannot see it.`,

src/lib/server/tools/images/generation.ts

@@ -1,6 +1,5 @@
 import type { BackendTool } from "..";
 import { uploadFile } from "../../files/uploadFile";
-import { ToolResultStatus } from "$lib/types/Tool";
 import { MessageUpdateType } from "$lib/types/MessageUpdate";
 import { callSpace, getIpToken, type GradioImage } from "../utils";
 
@@ -81,7 +80,6 @@ const imageGeneration: BackendTool = {
 		}
 
 		return {
-			status: ToolResultStatus.Success,
 			outputs: [
 				{
 					imageGeneration: `An image has been generated for the following prompt: "${prompt}". Answer as if the user can already see the image. Do not try to insert the image or to add space for it. The user can already see the image. Do not try to describe the image as you the model cannot see it.`,

src/lib/server/tools/index.ts

@@ -1,5 +1,5 @@
 import type { MessageUpdate } from "$lib/types/MessageUpdate";
-import type { Tool, ToolResultError, ToolResultSuccess } from "$lib/types/Tool";
+import type { Tool, ToolResultSuccess } from "$lib/types/Tool";
 
 import calculator from "./calculator";
 import directlyAnswer from "./directlyAnswer";
@@ -15,17 +15,11 @@ export type BackendToolContext = Pick<
 	"conv" | "messages" | "assistant" | "ip" | "username"
 > & { preprompt?: string };
 
-// typescript can't narrow a discriminated union after applying a generic like Omit to it
-// so we have to define the omitted types and create a new union
-type ToolResultSuccessOmitted = Omit<ToolResultSuccess, "call">;
-type ToolResultErrorOmitted = Omit<ToolResultError, "call">;
-type ToolResultOmitted = ToolResultSuccessOmitted | ToolResultErrorOmitted;
-
 export interface BackendTool extends Tool {
 	call(
 		params: Record<string, string | number | boolean>,
 		context: BackendToolContext
-	): AsyncGenerator<MessageUpdate, ToolResultOmitted, undefined>;
+	): AsyncGenerator<MessageUpdate, Omit<ToolResultSuccess, "status" | "call" | "type">, undefined>;
 }
 
 export const allTools: BackendTool[] = [

src/lib/server/tools/web/search.ts

@@ -1,4 +1,3 @@
-import { ToolResultStatus } from "$lib/types/Tool";
 import type { BackendTool } from "..";
 import { runWebSearch } from "../../websearch/runWebSearch";
 
@@ -23,7 +22,6 @@ const websearch: BackendTool = {
 			.join("\n------------\n");
 
 		return {
-			status: ToolResultStatus.Success,
 			outputs: [{ websearch: chunks }],
 			display: false,
 		};

src/lib/server/tools/web/url.ts

@@ -1,6 +1,5 @@
 import { stringifyMarkdownElementTree } from "$lib/server/websearch/markdown/utils/stringify";
 import { scrapeUrl } from "$lib/server/websearch/scrape/scrape";
-import { ToolResultStatus } from "$lib/types/Tool";
 import type { BackendTool } from "..";
 
 const fetchUrl: BackendTool = {
@@ -22,7 +21,6 @@ const fetchUrl: BackendTool = {
 		const { title, markdownTree } = await scrapeUrl(url, Infinity);
 
 		return {
-			status: ToolResultStatus.Success,
 			outputs: [{ title, text: stringifyMarkdownElementTree(markdownTree) }],
 			display: false,
 		};

src/lib/utils/stringifyError.ts

@@ -0,0 +1,12 @@
+/** Takes an unknown error and attempts to convert it to a string */
+export function stringifyError(error: unknown): string {
+	if (error instanceof Error) return error.message;
+	if (typeof error === "string") return error;
+	if (typeof error === "object" && error !== null) {
+		// try a few common properties
+		if ("message" in error && typeof error.message === "string") return error.message;
+		if ("body" in error && typeof error.body === "string") return error.body;
+		if ("name" in error && typeof error.name === "string") return error.name;
+	}
+	return "Unknown error";
+}