feat: add tavily search citations (#1778)

* feat: add citations for website knowledge * chore: remove console logs Signed-off-by: Ryan Hopper-Lowe <[email protected]> * feat: add citations for tavily search --------- Signed-off-by: Ryan Hopper-Lowe <[email protected]>
obot-platform · Feb 19, 2025 · 64f76ae · 64f76ae
1 parent 02b9fe8
commit 64f76ae
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 3 deletions.
diff --git a/ui/admin/app/lib/model/chatEvents.ts b/ui/admin/app/lib/model/chatEvents.ts
@@ -3,14 +3,31 @@ export type ToolInput = {
 	content: string;
 };
 
-export type KnowledgeToolOutput = { url?: string; content: string }[];
+export type KnowledgeToolOutput = {
+	url?: string;
+	content: string;
+}[];
 
 export type GoogleSearchOutput = {
 	duration: { search: number; refine: number; response: number };
 	query: string;
 	results: { url: string; content: string }[];
 };
 
+export type TavilySearchOutput = {
+	query: string;
+	follow_up_questions: string;
+	answer: string;
+	images: string[];
+	results: {
+		title: string;
+		url: string;
+		content: string;
+		score: number;
+		raw_content: string;
+	}[];
+};
+
 export type ToolCall = {
 	name: string;
 	description: string;
@@ -19,6 +36,7 @@ export type ToolCall = {
 	metadata?: {
 		category?: string;
 		icon?: string;
+		toolBundle?: string;
 	};
 };
 

diff --git a/ui/admin/app/lib/store/chat/message-store.ts b/ui/admin/app/lib/store/chat/message-store.ts
@@ -4,6 +4,7 @@ import {
 	ChatEvent,
 	GoogleSearchOutput,
 	KnowledgeToolOutput,
+	TavilySearchOutput,
 	ToolCall,
 } from "~/lib/model/chatEvents";
 import { Message, promptMessage, toolCallMessage } from "~/lib/model/messages";
@@ -196,7 +197,6 @@ export const createMessageStore = () => {
 			const { toolCall } = event;
 
 			const sources = pullSources(toolCall);
-
 			if (sources) parsedSources.push(...sources);
 
 			// if the toolCall is an output event
@@ -223,14 +223,23 @@ function pullSources(toolCall: ToolCall) {
 
 	const [err, output] = handleTry(() => JSON.parse(toolCall.output));
 
-	if (err) return [];
+	if (err) {
+		console.error(err);
+		return [];
+	}
 
 	if (toolCall.name === "Knowledge") {
 		const o = output as KnowledgeToolOutput;
 		return o;
 	}
 
 	if (toolCall.name === "Search") {
+		if (toolCall.metadata?.toolBundle === "Tavily Search") {
+			const o = output as TavilySearchOutput;
+
+			return o.results;
+		}
+
 		const o = output as GoogleSearchOutput;
 		return o.results;
 	}