Fix results

app.py

@@ -25,7 +25,6 @@ from scripts.text_web_browser import (
     FindNextTool,
     PageDownTool,
     PageUpTool,
-    SearchInformationTool,
     SimpleTextBrowser,
     VisitTool,
 )
@@ -42,7 +41,95 @@ from smolagents import (
 from smolagents.agent_types import AgentText, AgentImage, AgentAudio
 from smolagents.gradio_ui import pull_messages_from_step, handle_agent_output_types
 
+from smolagents import Tool
 
+
+class GoogleSearchTool(Tool):
+    name = "web_search"
+    description = """Performs a google web search for your query then returns a string of the top search results."""
+    inputs = {
+        "query": {"type": "string", "description": "The search query to perform."},
+        "filter_year": {
+            "type": "integer",
+            "description": "Optionally restrict results to a certain year",
+            "nullable": True,
+        },
+    }
+    output_type = "string"
+
+    def __init__(self):
+        super().__init__(self)
+        import os
+
+        self.serpapi_key = os.getenv("SERPER_API_KEY")
+
+    def forward(self, query: str, filter_year: Optional[int] = None) -> str:
+        import requests
+
+        if self.serpapi_key is None:
+            raise ValueError("Missing SerpAPI key. Make sure you have 'SERPER_API_KEY' in your env variables.")
+
+        params = {
+            "engine": "google",
+            "q": query,
+            "api_key": self.serpapi_key,
+            "google_domain": "google.com",
+        }
+
+        headers = {
+        'X-API-KEY': self.serpapi_key,
+        'Content-Type': 'application/json'
+        }
+
+        if filter_year is not None:
+            params["tbs"] = f"cdr:1,cd_min:01/01/{filter_year},cd_max:12/31/{filter_year}"
+
+        response = requests.request("POST", "https://google.serper.dev/search", headers=headers, data=json.dumps(params))
+
+
+        if response.status_code == 200:
+            results = response.json()
+        else:
+            raise ValueError(response.json())
+
+        if "organic" not in results.keys():
+            print("REZZZ", results.keys())
+            if filter_year is not None:
+                raise Exception(
+                    f"No results found for query: '{query}' with filtering on year={filter_year}. Use a less restrictive query or do not filter on year."
+                )
+            else:
+                raise Exception(f"No results found for query: '{query}'. Use a less restrictive query.")
+        if len(results["organic"]) == 0:
+            year_filter_message = f" with filter year={filter_year}" if filter_year is not None else ""
+            return f"No results found for '{query}'{year_filter_message}. Try with a more general query, or remove the year filter."
+
+        web_snippets = []
+        if "organic" in results:
+            for idx, page in enumerate(results["organic"]):
+                date_published = ""
+                if "date" in page:
+                    date_published = "\nDate published: " + page["date"]
+
+                source = ""
+                if "source" in page:
+                    source = "\nSource: " + page["source"]
+
+                snippet = ""
+                if "snippet" in page:
+                    snippet = "\n" + page["snippet"]
+
+                redacted_version = f"{idx}. [{page['title']}]({page['link']}){date_published}{source}\n{snippet}"
+
+                redacted_version = redacted_version.replace("Your browser can't play this video.", "")
+                web_snippets.append(redacted_version)
+
+        return "## Search Results\n" + "\n\n".join(web_snippets)
+
+# web_search = GoogleSearchTool()
+
+# print(web_search(query="Donald Trump news"))
+# quit()
 AUTHORIZED_IMPORTS = [
     "requests",
     "zipfile",
@@ -74,8 +161,6 @@ login(os.getenv("HF_TOKEN"))
 
 append_answer_lock = threading.Lock()
 
-SET = "validation"
-
 custom_role_conversions = {"tool-call": "assistant", "tool-response": "user"}
 
 user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
@@ -104,7 +189,7 @@ ti_tool = TextInspectorTool(model, text_limit)
 browser = SimpleTextBrowser(**BROWSER_CONFIG)
 
 WEB_TOOLS = [
-    SearchInformationTool(browser),
+    GoogleSearchTool(),
     VisitTool(browser),
     PageUpTool(browser),
     PageDownTool(browser),
@@ -128,24 +213,6 @@ def create_agent():
 
 document_inspection_tool = TextInspectorTool(model, 20000)
 
-
-# augmented_question = """You have one question to answer. It is paramount that you provide a correct answer.
-# Give it all you can: I know for a fact that you have access to all the relevant tools to solve it and find the correct answer (the answer does exist). Failure or 'I cannot answer' or 'None found' will not be tolerated, success will be rewarded.
-# Run verification steps if that's needed, you must make sure you find the correct answer!
-# Here is the task:
-# """ + example["question"]
-
-# if example["file_name"]:
-#     prompt_use_files = "\n\nTo solve the task above, you will have to use this attached file:"
-#     prompt_use_files += get_single_file_description(
-#         example["file_name"], example["question"], visual_inspection_tool, document_inspection_tool
-#     )
-#     augmented_question += prompt_use_files
-
-
-# final_result = agent.run(augmented_question)
-
-
 def stream_to_gradio(
     agent,
     task: str,