From 64f76ae4bcb9c7764250f427e8200afb0ebd979e Mon Sep 17 00:00:00 2001
From: Ryan Hopper-Lowe <46546486+ryanhopperlowe@users.noreply.github.com>
Date: Wed, 19 Feb 2025 14:25:24 -0600
Subject: [PATCH] feat: add tavily search citations (#1778)

* feat: add citations for website knowledge

* chore: remove console logs

Signed-off-by: Ryan Hopper-Lowe <ryan@acorn.io>

* feat: add citations for tavily search

---------

Signed-off-by: Ryan Hopper-Lowe <ryan@acorn.io>
---
 ui/admin/app/lib/model/chatEvents.ts         | 20 +++++++++++++++++++-
 ui/admin/app/lib/store/chat/message-store.ts | 13 +++++++++++--
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/ui/admin/app/lib/model/chatEvents.ts b/ui/admin/app/lib/model/chatEvents.ts
index cce1599ac..8c24aec8e 100644
--- a/ui/admin/app/lib/model/chatEvents.ts
+++ b/ui/admin/app/lib/model/chatEvents.ts
@@ -3,7 +3,10 @@ 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 };
@@ -11,6 +14,20 @@ export type GoogleSearchOutput = {
 	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
index d1ef9c8bc..80d0d6aa5 100644
--- 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,7 +223,10 @@ 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;
@@ -231,6 +234,12 @@ function pullSources(toolCall: ToolCall) {
 	}
 
 	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;
 	}