feat(tools): use gradio ETA for tools (#1411)

src/lib/components/chat/ToolUpdate.svelte

@@ -9,7 +9,7 @@
 	import CarbonTools from "~icons/carbon/tools";
 	import { ToolResultStatus, type ToolFront } from "$lib/types/Tool";
 	import { page } from "$app/stores";
-	import { onMount } from "svelte";
+	import { onDestroy } from "svelte";
 	import { browser } from "$app/environment";
 
 	export let tool: MessageToolUpdate[];
@@ -19,6 +19,8 @@
 	$: toolError = tool.some(isMessageToolErrorUpdate);
 	$: toolDone = tool.some(isMessageToolResultUpdate);
 
+	$: eta = tool.find((el) => el.subtype === MessageToolUpdateType.ETA)?.eta;
+
 	const availableTools: ToolFront[] = $page.data.tools;
 
 	let loadingBarEl: HTMLDivElement;
@@ -26,16 +28,24 @@
 
 	let isShowingLoadingBar = false;
 
-	onMount(() => {
-		if (!toolError && !toolDone && loading && loadingBarEl) {
+	$: !toolError &&
+		!toolDone &&
+		loading &&
+		loadingBarEl &&
+		eta &&
+		(() => {
 			loadingBarEl.classList.remove("hidden");
 			isShowingLoadingBar = true;
 			animation = loadingBarEl.animate([{ width: "0%" }, { width: "calc(100%+1rem)" }], {
-				duration: availableTools.find((tool) => tool.name === toolFnName)?.timeToUseMS,
+				duration: eta * 1000,
 				fill: "forwards",
 			});
+		})();
+
+	onDestroy(() => {
+		if (animation) {
+			animation.cancel();
 		}
-		return () => animation?.cancel();
 	});
 
 	// go to 100% quickly if loading is done

src/lib/server/textGeneration/tools.ts

@@ -70,7 +70,7 @@ async function* callTool(
 	};
 
 	try {
-		const toolResult = yield* tool.call(call.parameters, ctx);
+		const toolResult = yield* tool.call(call.parameters, ctx, uuid);
 
 		yield {
 			type: MessageUpdateType.Tool,

src/lib/server/tools/index.ts

@@ -119,7 +119,7 @@ export const configTools = z
 	.transform((val) => [...val, calculator, directlyAnswer, fetchUrl, websearch]);
 
 export function getCallMethod(tool: Omit<BaseTool, "call">): BackendCall {
-	return async function* (params, ctx) {
+	return async function* (params, ctx, uuid) {
 		if (
 			tool.endpoint === null ||
 			!tool.baseUrl ||
@@ -203,11 +203,12 @@ export function getCallMethod(tool: Omit<BaseTool, "call">): BackendCall {
 			}
 		});
 
-		const outputs = await callSpace(
+		const outputs = yield* callSpace(
 			tool.baseUrl,
 			tool.endpoint,
 			await Promise.all(inputs),
-			ipToken
+			ipToken,
+			uuid
 		);
 
 		if (!isValidOutputComponent(tool.outputComponent)) {

src/lib/server/tools/utils.ts

@@ -1,27 +1,20 @@
 import { env } from "$env/dynamic/private";
 import { Client } from "@gradio/client";
 import { SignJWT } from "jose";
-import { logger } from "../logger";
 import JSON5 from "json5";
+import {
+	MessageToolUpdateType,
+	MessageUpdateType,
+	type MessageToolUpdate,
+} from "$lib/types/MessageUpdate";
 
-export type GradioImage = {
-	path: string;
-	url: string;
-	orig_name: string;
-	is_stream: boolean;
-	meta: Record<string, unknown>;
-};
-
-type GradioResponse = {
-	data: unknown[];
-};
-
-export async function callSpace<TInput extends unknown[], TOutput extends unknown[]>(
+export async function* callSpace<TInput extends unknown[], TOutput extends unknown[]>(
 	name: string,
 	func: string,
 	parameters: TInput,
-	ipToken: string | undefined
-): Promise<TOutput> {
+	ipToken: string | undefined,
+	uuid: string
+): AsyncGenerator<MessageToolUpdate, TOutput, undefined> {
 	class CustomClient extends Client {
 		fetch(input: RequestInfo | URL, init?: RequestInit): Promise<Response> {
 			init = init || {};
@@ -34,15 +27,32 @@ export async function callSpace<TInput extends unknown[], TOutput extends unknow
 	}
 	const client = await CustomClient.connect(name, {
 		hf_token: (env.HF_TOKEN ?? env.HF_ACCESS_TOKEN) as unknown as `hf_${string}`,
+		events: ["status", "data"],
 	});
 
-	return await client
-		.predict(func, parameters)
-		.then((res) => (res as unknown as GradioResponse).data as TOutput)
-		.catch((e) => {
-			logger.error(e);
-			throw e;
-		});
+	const job = client.submit(func, parameters);
+
+	let data;
+	for await (const output of job) {
+		console.log({ output });
+		if (output.type === "data") {
+			data = output.data as TOutput;
+		}
+		if (output.type === "status" && output.eta) {
+			yield {
+				type: MessageUpdateType.Tool,
+				subtype: MessageToolUpdateType.ETA,
+				eta: output.eta,
+				uuid,
+			};
+		}
+	}
+
+	if (!data) {
+		throw new Error("No data found in tool call");
+	}
+
+	return data;
 }
 
 export async function getIpToken(ip: string, username?: string) {

src/lib/types/MessageUpdate.ts

@@ -75,7 +75,10 @@ export enum MessageToolUpdateType {
 	Result = "result",
 	/** Error while running tool */
 	Error = "error",
+	/** ETA update */
+	ETA = "eta",
 }
+
 interface MessageToolBaseUpdate<TSubType extends MessageToolUpdateType> {
 	type: MessageUpdateType.Tool;
 	subtype: TSubType;
@@ -91,10 +94,16 @@ export interface MessageToolResultUpdate
 export interface MessageToolErrorUpdate extends MessageToolBaseUpdate<MessageToolUpdateType.Error> {
 	message: string;
 }
+
+export interface MessageToolETAUpdate extends MessageToolBaseUpdate<MessageToolUpdateType.ETA> {
+	eta: number;
+}
+
 export type MessageToolUpdate =
 	| MessageToolCallUpdate
 	| MessageToolResultUpdate
-	| MessageToolErrorUpdate;
+	| MessageToolErrorUpdate
+	| MessageToolETAUpdate;
 
 // Everything else
 export interface MessageTitleUpdate {

src/lib/types/Tool.ts

@@ -177,5 +177,6 @@ export interface ToolCall {
 
 export type BackendCall = (
 	params: Record<string, string | number | boolean>,
-	context: BackendToolContext
+	context: BackendToolContext,
+	uuid: string
 ) => AsyncGenerator<MessageUpdate, Omit<ToolResultSuccess, "status" | "call" | "type">, undefined>;