feat : added xperimental advanced search mode.

agent.py

@@ -43,16 +43,16 @@ class SmolAlbert(CodeAgent):
         Initialize the SmolAlbert agent with Tavily tools and a model.
         """
         # Set up the agent with the Tavily tool and a model
-        search_tool = TavilySearchTool()
-        image_search_tool = TavilyImageURLSearchTool()
-        extract_tool = TavilyExtractTool()
-        image_query_tool = ImageQueryTool()
+        self.search_tool = TavilySearchTool()
+        self.image_search_tool = TavilyImageURLSearchTool()
+        self.extract_tool = TavilyExtractTool()
+        self.image_query_tool = ImageQueryTool()
         model = InferenceClientModel(
             model_id=self.model_id,
             provider=self.provider,
             token=os.getenv("HF_API_KEY"))
         self.agent = CodeAgent(
-            tools=[search_tool, image_search_tool, extract_tool, image_query_tool], 
+            tools=[self.search_tool, self.image_search_tool, self.extract_tool, self.image_query_tool], 
             model=model,
             stream_outputs=True,
             instructions=(
@@ -61,9 +61,21 @@ class SmolAlbert(CodeAgent):
                 "Example format: ... (see [1](https://example1.com)) ... (see [2](https://example2.com)) ... "
                 "Do not invent URL(s) — only use the ones you were provided."
                 "If the answer includes an image URL, include it as an inline Markdown image: ![image](<image_url>)"
-            )
+            ),
+            #additional_authorized_imports=[""]
         )
 
+        self.advanced_mode = False
+
+    def enable_advanced_mode(self, enable: bool):
+        """
+        Enable or disable advanced mode for the search tool.
+        Advanced mode uses more credits but yields better results.
+        """
+        self.advanced_mode = enable
+        self.search_tool.enable_advanced_mode(enable)
+        self.extract_tool.enable_advanced_mode(enable)
+
     def run(self, task: str, additional_args: dict | None = None) -> str:
         """
         Run the agent with a given query and return the final answer.

agent_ui.py

@@ -299,7 +299,13 @@ class AgentUI:
         self.agent = agent
         self.description = getattr(agent, "description", None)
 
-    def interact_with_agent(self, prompt, verbose_messages, quiet_messages):
+    def set_advanced_mode(self, enabled: bool):
+        """
+        Configure the agent to enable/disable advanced mode.
+        """
+        self.agent.enable_advanced_mode(enabled)
+
+    def interact_with_agent(self, prompt: str, verbose_messages: list, quiet_messages: list):
         """
         Interacts with the agent and streams results into two separate histories:
             - verbose_messages: full reasoning stream (Chatterbox)
@@ -437,6 +443,12 @@ class AgentUI:
         """
         return self.agent.get_search_credits()
 
+    def get_advanced_mode(self) -> bool:
+        """
+        Return the agent's current advanced_mode flag for initializing the checkbox on page load.
+        """
+        return getattr(self.agent, "advanced_mode", False)
+
     def create_app(self):
         import gradio as gr
 
@@ -463,6 +475,19 @@ class AgentUI:
                     )
                     submit_btn = gr.Button("Submit", variant="primary")
 
+                # Advanced search mode checkbox
+                advanced_checkbox = gr.Checkbox(
+                    label="Advanced search mode",
+                    value=getattr(self.agent, "advanced_mode", False),
+                    info="Toggle advanced search behavior for the agent (x2 search credits).",
+                    container=True,
+                )
+
+                # call agent configuration when checkbox changes
+                advanced_checkbox.change(self.set_advanced_mode, advanced_checkbox, None)
+                # ensure the checkbox reflects the current agent state each time a page/session loads
+                agent.load(self.get_advanced_mode, None, advanced_checkbox)
+
                 tavily_credits = gr.Textbox(
                     label="Tavily Credits",
                     value=self.get_tavily_credits(),
@@ -554,5 +579,3 @@ class AgentUI:
             verbose_chatbot.clear(self.clear_history, inputs=None, outputs=[stored_messages_verbose, stored_messages_quiet])
 
         return agent
-
-__all__ = ["stream_to_gradio", "AgentUI"]

web_tools.py

@@ -46,6 +46,10 @@ class TavilyBaseClient:
 
         return f"{plan_usage}/{plan_limit}"
 
+# ---------------------------------------------------------------------
+# Tavily Tools
+# ---------------------------------------------------------------------
+
 class TavilySearchTool(TavilyBaseClient, Tool):
     """
     A tool to perform web searches using the Tavily API.
@@ -60,16 +64,63 @@ class TavilySearchTool(TavilyBaseClient, Tool):
     }
     output_type = "string"
 
+    # Consumes 1 Tavily credit per query
+    __basic_params = {
+            "search_depth": "basic",
+            "max_results": 10,              # fetch up to 10 sources
+            "auto_parameters": False,       # keep manual control
+            "include_raw_content": False,
+        }
+
+    # Consumes 2 Tavily credits per query
+    __advanced_params = {
+            "search_depth": "advanced",     # 'advanced' yields better relevance
+            "max_results": 10,              # fetch up to 10 sources
+            "chunks_per_source": 3,         # number of content snippets to return per source
+            "auto_parameters": False,       # keep manual control
+            "include_raw_content": False,
+        }
+
+    def __init__(self):
+        """
+        Construct the TavilySearchTool.
+        """
+        # Call superclass constructor
+        super().__init__()
+
+
+        self.params = TavilySearchTool.__basic_params
+
+    def enable_advanced_mode(self, enable: bool = True):
+        """
+        Enable or disable advanced mode for the search tool.
+        Advanced mode uses more credits but yields better results.
+        """
+        if enable:
+            self.params = TavilySearchTool.__advanced_params
+        else:
+            self.params = TavilySearchTool.__basic_params
+        
+        print(f"TavilySearchTool advanced mode has been {'enabled' if enable else 'disabled'}.")
+
     def forward(self, query: str):
-        response = self._tavily_client.search(query)
+
+        params = self.params
+        params["query"] = query
+
+        try:
+            response = self._tavily_client.search(**params)
+        except Exception as e:
+            return f"Error calling Tavily API: {e}"
+        
         return response
 
 class TavilyExtractTool(TavilyBaseClient, Tool):
     """
-    A tool to extract information from web pages using the Tavily API.
+    A tool to extract raw information from web pages using the Tavily API.
     """
     name = "tavily_extract"
-    description = "Extract information from web pages using Tavily."
+    description = "Extract raw information from web pages using Tavily."
     inputs = {
         "url": {
             "type": "string",
@@ -78,8 +129,37 @@ class TavilyExtractTool(TavilyBaseClient, Tool):
     }
     output_type = "string"
 
+    def __init__(self):
+        """
+        Construct the TavilyExtractTool.
+        """
+        # Call superclass constructor
+        super().__init__()
+
+        self.extract_depth = "basic"
+
+    def enable_advanced_mode(self, enable: bool = True):
+        """
+        Enable or disable advanced mode for the extract tool.
+        Advanced mode uses more credits but yields better results (retrieves more data, including tables and embedded content).
+        """
+        if enable:
+            self.extract_depth = "advanced"
+        else:
+            self.extract_depth = "basic"
+
+        print(f"TavilyExtractTool advanced mode has been {'enabled' if enable else 'disabled'}.")
+
     def forward(self, url: str):
-        response = self._tavily_client.extract(url)
+        try:
+            response = self._tavily_client.extract(
+                urls=url,
+                extract_depth=self.extract_depth)
+        except Exception as e:
+            return f"Error calling Tavily extract API: {e}"
+
+        # Tavily's Extract API can return raw HTML + text.
+        # you may trim or sanitize here if needed.
         return response
 
 class TavilyImageURLSearchTool(TavilyBaseClient, Tool):
@@ -97,14 +177,22 @@ class TavilyImageURLSearchTool(TavilyBaseClient, Tool):
     output_type = "string"
 
     def forward(self, query: str):
-        response = self._tavily_client.search(query, include_images=True)
-        
+        try:
+            response = self._tavily_client.search(
+                query,
+                include_images=True,
+                include_image_descriptions=True,
+                max_results=5
+            )
+        except Exception as e:
+            return f"Error calling Tavily API: {e}"
+
         images = response.get("images", [])
-        
-        if images:
-            # Return the URL of the first image
-            first_image = images[0]
-            if isinstance(first_image, dict):
-                return first_image.get("url")
-            return first_image
-        return "none"
+        if not images:
+            return "none"
+
+        # Return the URL of the first image
+        first_image = images[0]
+        if isinstance(first_image, dict):
+            return first_image.get("url", "none")
+        return first_image or "none"