From 0f320dd362940454de63238ce50861f4be93df89 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Wed, 23 Oct 2024 23:38:17 +0300
Subject: [PATCH] WebSearch: Add endpoint for Tavily

---
 public/scripts/secrets.js |  1 +
 src/endpoints/search.js   | 48 +++++++++++++++++++++++++++++++++++++++
 src/endpoints/secrets.js  |  1 +
 3 files changed, 50 insertions(+)

diff --git a/public/scripts/secrets.js b/public/scripts/secrets.js
index 9a2c23d03..24de91514 100644
--- a/public/scripts/secrets.js
+++ b/public/scripts/secrets.js
@@ -34,6 +34,7 @@ export const SECRET_KEYS = {
     STABILITY: 'api_key_stability',
     BLOCKENTROPY: 'api_key_blockentropy',
     CUSTOM_OPENAI_TTS: 'api_key_custom_openai_tts',
+    TAVILY: 'api_key_tavily',
 };
 
 const INPUT_MAP = {
diff --git a/src/endpoints/search.js b/src/endpoints/search.js
index e15c2b7a6..6e21dfa5c 100644
--- a/src/endpoints/search.js
+++ b/src/endpoints/search.js
@@ -206,6 +206,54 @@ router.post('/searxng', jsonParser, async (request, response) => {
     }
 });
 
+router.post('/tavily', jsonParser, async (request, response) => {
+    try {
+        const apiKey = readSecret(request.user.directories, SECRET_KEYS.TAVILY);
+
+        if (!apiKey) {
+            console.log('No Tavily key found');
+            return response.sendStatus(400);
+        }
+
+        const { query } = request.body;
+
+        const body = {
+            query: query,
+            api_key: apiKey,
+            search_depth: 'basic',
+            topic: 'general',
+            include_answer: true,
+            include_raw_content: false,
+            include_images: false,
+            include_image_descriptions: false,
+            include_domains: [],
+            max_results: 10,
+        };
+
+        const result = await fetch('https://api.tavily.com/search', {
+            method: 'POST',
+            headers: {
+                'Content-Type': 'application/json',
+            },
+            body: JSON.stringify(body),
+        });
+
+        console.log('Tavily query', query);
+
+        if (!result.ok) {
+            const text = await result.text();
+            console.log('Tavily request failed', result.statusText, text);
+            return response.status(500).send(text);
+        }
+
+        const data = await result.json();
+        return response.json(data);
+    } catch (error) {
+        console.log(error);
+        return response.sendStatus(500);
+    }
+});
+
 router.post('/visit', jsonParser, async (request, response) => {
     try {
         const url = request.body.url;
diff --git a/src/endpoints/secrets.js b/src/endpoints/secrets.js
index 078792cb6..e450c994c 100644
--- a/src/endpoints/secrets.js
+++ b/src/endpoints/secrets.js
@@ -47,6 +47,7 @@ export const SECRET_KEYS = {
     STABILITY: 'api_key_stability',
     BLOCKENTROPY: 'api_key_blockentropy',
     CUSTOM_OPENAI_TTS: 'api_key_custom_openai_tts',
+    TAVILY: 'api_key_tavily',
 };
 
 // These are the keys that are safe to expose, even if allowKeysExposure is false