Update search.py

Replaced serper by ddg
BaLaurent · Aug 20, 2024 · 7b574fb · 7b574fb
1 parent 755d48b
commit 7b574fb
Showing 1 changed file with 22 additions and 39 deletions.
diff --git a/llm_osint/tools/search.py b/llm_osint/tools/search.py
@@ -1,54 +1,37 @@
-from langchain.utilities import GoogleSerperAPIWrapper
+from langchain_community.tools import DuckDuckGoSearchRun, DuckDuckGoSearchResults
 from langchain.agents import Tool
 
 from llm_osint import cache_utils
 
 
-class GoogleSerperSearchWrapper(GoogleSerperAPIWrapper):
+class CachedDuckDuckGoSearchRun(DuckDuckGoSearchRun):
     @cache_utils.cache_func
     def run(self, query: str) -> str:
         return super().run(query)
 
-    def _parse_results(self, results: dict) -> str:
-        snippets = []
-
-        if results.get("answerBox"):
-            answer_box = results.get("answerBox", {})
-            if answer_box.get("answer"):
-                return answer_box.get("answer")
-            elif answer_box.get("snippet"):
-                return answer_box.get("snippet").replace("\n", " ")
-            elif answer_box.get("snippetHighlighted"):
-                return ", ".join(answer_box.get("snippetHighlighted"))
-
-        if results.get("knowledgeGraph"):
-            kg = results.get("knowledgeGraph", {})
-            title = kg.get("title")
-            entity_type = kg.get("type")
-            if entity_type:
-                snippets.append(f"{title}: {entity_type}.")
-            description = kg.get("description")
-            if description:
-                snippets.append(description)
-            for attribute, value in kg.get("attributes", {}).items():
-                snippets.append(f"{title} {attribute}: {value}.")
-
-        for result in results["organic"][: self.k]:
-            if "snippet" in result:
-                snippets.append(f'{result["title"]}: {result["snippet"]} (link {result["link"]})')
-            for attribute, value in result.get("attributes", {}).items():
-                snippets.append(f'{result["title"]}: {attribute} = {value}.')
-
-        if len(snippets) == 0:
-            return "No good results found"
-
-        return "\n\n".join(snippets)
-
 
 def get_search_tool(**kwargs) -> Tool:
-    search = GoogleSerperSearchWrapper(**kwargs)
+    search = CachedDuckDuckGoSearchRun(**kwargs)
     return Tool(
         name="Search Term",
         func=search.run,
-        description="useful for when you need to find information about general things, names, usernames, places, etc. the input should be a search term",
+        description="Useful for finding information on general topics, names, usernames, places, etc. The input should be a search term.",
+    )
+
+
+def get_detailed_search_tool(**kwargs) -> Tool:
+    search = DuckDuckGoSearchResults(**kwargs)
+    return Tool(
+        name="Detailed Search",
+        func=search.run,
+        description="Useful for getting detailed search results, including links and sources. The input should be a search term.",
+    )
+
+
+def get_news_search_tool(**kwargs) -> Tool:
+    search = DuckDuckGoSearchResults(backend="news", **kwargs)
+    return Tool(
+        name="News Search",
+        func=search.run,
+        description="Useful for searching recent news articles. The input should be a search term.",
     )