From 0732d9b6c82dcfacae2109395390d3b02366c0c2 Mon Sep 17 00:00:00 2001
From: Siddhant Rai <siddhant.rai.5686@gmail.com>
Date: Thu, 20 Mar 2025 08:27:00 +0530
Subject: [PATCH] fix: brave and duckduckgo retrievers

---
 application/retriever/brave_search.py    | 13 ++++---
 application/retriever/duckduck_search.py | 47 ++++++++----------------
 frontend/src/upload/Upload.tsx           |  1 +
 3 files changed, 25 insertions(+), 36 deletions(-)

diff --git a/application/retriever/brave_search.py b/application/retriever/brave_search.py
index ed490734..33bdb894 100644
--- a/application/retriever/brave_search.py
+++ b/application/retriever/brave_search.py
@@ -1,15 +1,16 @@
 import json
-from application.retriever.base import BaseRetriever
+
+from langchain_community.tools import BraveSearch
+
 from application.core.settings import settings
 from application.llm.llm_creator import LLMCreator
-from langchain_community.tools import BraveSearch
+from application.retriever.base import BaseRetriever
 
 
 class BraveRetSearch(BaseRetriever):
 
     def __init__(
         self,
-        question,
         source,
         chat_history,
         prompt,
@@ -19,7 +20,7 @@ class BraveRetSearch(BaseRetriever):
         user_api_key=None,
         decoded_token=None,
     ):
-        self.question = question
+        self.question = ""
         self.source = source
         self.chat_history = chat_history
         self.prompt = prompt
@@ -93,7 +94,9 @@ class BraveRetSearch(BaseRetriever):
         for line in completion:
             yield {"answer": str(line)}
 
-    def search(self):
+    def search(self, query: str = ""):
+        if query:
+            self.question = query
         return self._get_data()
 
     def get_params(self):
diff --git a/application/retriever/duckduck_search.py b/application/retriever/duckduck_search.py
index 9ce73995..29bbba18 100644
--- a/application/retriever/duckduck_search.py
+++ b/application/retriever/duckduck_search.py
@@ -1,15 +1,15 @@
-from application.retriever.base import BaseRetriever
-from application.core.settings import settings
-from application.llm.llm_creator import LLMCreator
 from langchain_community.tools import DuckDuckGoSearchResults
 from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
 
+from application.core.settings import settings
+from application.llm.llm_creator import LLMCreator
+from application.retriever.base import BaseRetriever
+
 
 class DuckDuckSearch(BaseRetriever):
 
     def __init__(
         self,
-        question,
         source,
         chat_history,
         prompt,
@@ -19,7 +19,7 @@ class DuckDuckSearch(BaseRetriever):
         user_api_key=None,
         decoded_token=None,
     ):
-        self.question = question
+        self.question = ""
         self.source = source
         self.chat_history = chat_history
         self.prompt = prompt
@@ -38,41 +38,24 @@ class DuckDuckSearch(BaseRetriever):
         self.user_api_key = user_api_key
         self.decoded_token = decoded_token
 
-    def _parse_lang_string(self, input_string):
-        result = []
-        current_item = ""
-        inside_brackets = False
-        for char in input_string:
-            if char == "[":
-                inside_brackets = True
-            elif char == "]":
-                inside_brackets = False
-                result.append(current_item)
-                current_item = ""
-            elif inside_brackets:
-                current_item += char
-
-        if inside_brackets:
-            result.append(current_item)
-
-        return result
-
     def _get_data(self):
         if self.chunks == 0:
             docs = []
         else:
             wrapper = DuckDuckGoSearchAPIWrapper(max_results=self.chunks)
-            search = DuckDuckGoSearchResults(api_wrapper=wrapper)
+            search = DuckDuckGoSearchResults(api_wrapper=wrapper, output_format="list")
             results = search.run(self.question)
-            results = self._parse_lang_string(results)
 
             docs = []
             for i in results:
                 try:
-                    text = i.split("title:")[0]
-                    title = i.split("title:")[1].split("link:")[0]
-                    link = i.split("link:")[1]
-                    docs.append({"text": text, "title": title, "link": link})
+                    docs.append(
+                        {
+                            "text": i.get("snippet", "").strip(),
+                            "title": i.get("title", "").strip(),
+                            "link": i.get("link", "").strip(),
+                        }
+                    )
                 except IndexError:
                     pass
         if settings.LLM_NAME == "llama.cpp":
@@ -110,7 +93,9 @@ class DuckDuckSearch(BaseRetriever):
         for line in completion:
             yield {"answer": str(line)}
 
-    def search(self):
+    def search(self, query: str = ""):
+        if query:
+            self.question = query
         return self._get_data()
 
     def get_params(self):
diff --git a/frontend/src/upload/Upload.tsx b/frontend/src/upload/Upload.tsx
index e70c930f..8c016b29 100644
--- a/frontend/src/upload/Upload.tsx
+++ b/frontend/src/upload/Upload.tsx
@@ -471,6 +471,7 @@ function Upload({
       }, 3000);
     };
     xhr.open('POST', `${apiHost}/api/remote`);
+    xhr.setRequestHeader('Authorization', `Bearer ${token}`);
     xhr.send(formData);
   };
   const { getRootProps, getInputProps, isDragActive } = useDropzone({