Update app.py

app.py

@@ -9,17 +9,13 @@ from bs4 import BeautifulSoup
 import regex
 import pandas as pd
 import torch
-from functools import lru_cache
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
-from smolagents import PromptTemplates, CodeAgent, WebSearchTool, WikipediaSearchTool, VisitWebpageTool, PythonInterpreterTool
-import smolagents.tools as _tools
+from smolagents import CodeAgent, VisitWebpageTool
 from smolagents.models import ChatMessage
-# from huggingface_hub import InferenceClient, hf_hub_download
+from custom_tools import WebpageStructureAnalyzerTool, CachedWebSearchTool, CachedWikiTool, PreloadedPythonTool
 
 subprocess.run(["playwright", "install"], check=True)
 
-print(dir(_tools))
-
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
@@ -46,30 +42,6 @@ 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(WikipediaSearchTool):
-    @lru_cache(maxsize=128)
-    def run(self, page: str):
-        return super().run(page)
-
-class PreloadedPythonTool(PythonInterpreterTool):
-    """
-    A PythonInterpreterTool that automatically prepends the necessary imports
-    (bs4, BeautifulSoup, regex) so you never hit NameError inside your code blocks.
-    """
-    def run(self, code: str) -> str:
-        preamble = (
-            "import bs4\n"
-            "from bs4 import BeautifulSoup\n"
-            "import regex\n"
-        )
-        return super().run(preamble + code)
-
 # --- Basic Model Definition ---
 # ----- THIS IS WHERE YOU CAN BUILD WHAT YOU WANT ------
 class BasicModel:
@@ -194,13 +166,14 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
         # (1) Create the LLM wrapper
         llm_model = BasicModel(model_id=MODEL_ID, hf_token=hf_token)
         # (2) Create the tools
+        web_structure_analyzer_tool = WebpageStructureAnalyzerTool()
         wiki_tool = CachedWikiTool()
         search_tool = CachedWebSearchTool()
         python_tool = PythonInterpreterTool()
         html_parse_tool = VisitWebpageTool()
         #(3) Create the CodeAgent, passing the LLM wrapper and tools
         agent = CodeAgent(model=llm_model, 
-                          tools=[wiki_tool, search_tool, python_tool, html_parse_tool],
+                          tools=[web_structure_analyzer_tool, wiki_tool, search_tool, python_tool, html_parse_tool],
                           max_steps=5,
                           add_base_tools=True,
                           additional_authorized_imports=["dateparser", "bs4", "regex"])