From 62b10f6b906b90871b898a7f57b9db6ee775d062 Mon Sep 17 00:00:00 2001
From: Mishig Davaadorj <dmishig@gmail.com>
Date: Mon, 1 Jul 2024 23:31:08 +0100
Subject: [PATCH] [tool] hugging face docs tool

---
 src/lib/server/tools/hfDocs.ts | 47 ++++++++++++++++++++++++++++++++++
 src/lib/server/tools/index.ts  |  2 ++
 2 files changed, 49 insertions(+)
 create mode 100644 src/lib/server/tools/hfDocs.ts

diff --git a/src/lib/server/tools/hfDocs.ts b/src/lib/server/tools/hfDocs.ts
new file mode 100644
index 00000000000..70c040bbc21
--- /dev/null
+++ b/src/lib/server/tools/hfDocs.ts
@@ -0,0 +1,47 @@
+import type { BackendTool } from ".";
+import { callSpace, getIpToken } from "./utils";
+
+const hfDocs: BackendTool = {
+	name: "hf_docs",
+	displayName: "Hugging Face Docs",
+	description:
+		"Use this tool to get relevant docs snippets regarding Hugging Face open source libraries (transformers, diffusers, accelerate, huggingface_hub) and Hugging Face hf.co platform.",
+	isOnByDefault: true,
+	parameterDefinitions: {
+		query: {
+			required: true,
+			type: "string",
+			description:
+				"A search query which will be used to fetch the most relevant docs snippets regarding the user's query",
+		},
+	},
+	async *call({ query }, { messages, ip, username }) {
+		const ipToken = await getIpToken(ip, username);
+
+		const userMessages = messages.filter(({ from }) => from === "user");
+		const previousUserMessages = userMessages.slice(0, -1);
+
+		const queryWithPreviousMsgs =
+			(previousUserMessages.length
+				? `Previous questions: \n${previousUserMessages
+						.map(({ content }) => `- ${content}`)
+						.join("\n")}`
+				: "") +
+			"\n\nCurrent Question: " +
+			String(query);
+
+		const outputs = await callSpace<string[], string[]>(
+			"huggingchat/hf-docs",
+			"/predict",
+			[queryWithPreviousMsgs, "RAG-friendly"],
+			ipToken
+		);
+
+		return {
+			outputs: [{ hfDocs: outputs[0] }],
+			display: false,
+		};
+	},
+};
+
+export default hfDocs;
diff --git a/src/lib/server/tools/index.ts b/src/lib/server/tools/index.ts
index 8b963642e4e..66a5a186a1a 100644
--- a/src/lib/server/tools/index.ts
+++ b/src/lib/server/tools/index.ts
@@ -2,6 +2,7 @@ import type { MessageUpdate } from "$lib/types/MessageUpdate";
 import type { Tool, ToolResultSuccess } from "$lib/types/Tool";
 
 import calculator from "./calculator";
+import hfDocs from "./hfDocs";
 import directlyAnswer from "./directlyAnswer";
 import imageEditing from "./images/editing";
 import imageGeneration from "./images/generation";
@@ -30,4 +31,5 @@ export const allTools: BackendTool[] = [
 	imageEditing,
 	documentParser,
 	calculator,
+	hfDocs,
 ];