Update app.py

app.py

@@ -5,10 +5,13 @@ import requests
 import inspect
 import subprocess
 import dateparser
+from bs4 import BeautifulSoup
+import regex
 import pandas as pd
 import torch
+from functools import lru_cache
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from smolagents import CodeAgent, WebSearchTool, WikipediaSearchTool
+from smolagents import CodeAgent, WebSearchTool, WikipediaSearchTool, PythonTool, HTMLParseTool
 from smolagents.models import ChatMessage
 # from huggingface_hub import InferenceClient, hf_hub_download
 
@@ -60,6 +63,17 @@ def check_token_access():
     except Exception as e:
         print("❌ Token check failed:", e)
 
+class CachedWebSearchTool(WebSearchTool):
+    @lru_cache(maxsize=128)
+    def run(self, query: str):
+        # identical queries return instantly
+        return super().run(query)
+
+class CachedWikiTool(WikipediaTool):
+    @lru_cache(maxsize=128)
+    def run(self, page: str):
+        return super().run(page)
+
 # --- Basic Agent Definition ---
 # ----- THIS IS WHERE YOU CAN BUILD WHAT YOU WANT ------
 class BasicAgent:
@@ -103,13 +117,15 @@ class BasicAgent:
             # temperature=1.0,
         )
         # Introduce tools
-        wiki_tool = WikipediaSearchTool()
-        search_tool = WebSearchTool()
+        wiki_tool = CachedWikiTool()
+        search_tool = CachedWebSearchTool()
+        python_tool = PythonTool()
+        html_parse_tool = HTMLParseTool()
         # Initialize the agent
         self.agent = CodeAgent(model=self, 
-                               tools=[wiki_tool, search_tool], 
+                               tools=[wiki_tool, search_tool, python_tool, html_parse_tool], 
                                add_base_tools=True,
-                               additional_authorized_imports=["dateparser"])
+                               additional_authorized_imports=["dateparser", "bs4", "regex"])
         
     def _serialize_messages(self, messages):
         prompt = []
@@ -307,7 +323,8 @@ with gr.Blocks() as demo:
 
     run_button.click(
         fn=run_and_submit_all,
-        outputs=[status_output, results_table]
+        outputs=[status_output, results_table],
+        stream=True
     )
 
 if __name__ == "__main__":