refactor: rename functions in draw agent

src/deepengineer/deepsearch/draw_agent.py

@@ -13,9 +13,10 @@ from PIL import Image
 from smolagents import CodeAgent, LiteLLMModel
 from smolagents.agents import ActionStep
 from smolagents import tool, Tool
+from deepengineer.webcrawler.crawl_database import DataBase
 
 
-def save_fig(image_path: Path = Path("figure.png")) -> str:
+def _find_and_save_matplotlib_figure(image_path: Path = Path("figure.png")) -> str:
     """Save the current matplotlib figure to *path*.
     Save fig takes no arguments. The output path is hardcoded to "figure.png".
     """
@@ -48,7 +49,7 @@ class SaveMatplotlibFigTool(Tool):
             image_name = image_name + ".png"
         output_path = self.output_dir / image_name
         output_path.unlink(missing_ok=True)
-        save_fig(output_path)
+        _find_and_save_matplotlib_figure(output_path)
         if output_path.exists():
             return f"![]({image_name})"
         else:
@@ -58,7 +59,7 @@ class SaveMatplotlibFigTool(Tool):
 def _capture_snapshot(
     memory_step: ActionStep, agent: CodeAgent, image_path: Path = Path("figure.png")
 ) -> None:
-    save_fig(image_path)
+    _find_and_save_matplotlib_figure(image_path)
     if not plt.get_fignums():
         return
 
@@ -116,7 +117,7 @@ User instructions:
 """
 
 
-def draw_image_agent(
+def draw_matplotlib_image_from_prompt(
     prompt: str,
     image_path: str = Path("figure.png"),
     model_id: str = "mistral/mistral-medium-latest",
@@ -174,7 +175,9 @@ class DrawImageTool(Tool):
     def forward(self, prompt: str, image_name: str) -> str:
         if not image_name.endswith(".png"):
             image_name = image_name + ".png"
-        output_path = draw_image_agent(prompt, self.output_dir / image_name)
+        output_path = draw_matplotlib_image_from_prompt(
+            prompt, self.output_dir / image_name
+        )
         if output_path.exists():
             return f"![]({image_name})"
         else:

src/deepengineer/deepsearch/scawl_web_agent.py

@@ -203,11 +203,14 @@ class FindInMarkdownTool(Tool):
         return find_in_markdown(markdown, search_queries)
 
 
-def create_web_search_agent(model_id="deepseek/deepseek-chat"):
+def create_web_search_agent(
+    model_id="deepseek/deepseek-reasoner", database: DataBase | None = None
+):
     """Create a web search agent with search, crawling, and PDF analysis capabilities."""
 
     model = LiteLLMModel(model_id=model_id)
-    database = DataBase()
+    if database is None:
+        database = DataBase()
 
     # Web search and crawling tools
     WEB_SEARCH_TOOLS = [

tests/webcrawler/test_draw_agent.py

@@ -1,6 +1,6 @@
 import pytest
 from deepengineer.deepsearch.draw_agent import (
-    draw_image_agent,
+    draw_matplotlib_image_from_prompt,
     multiple_steps_draw_image_agent,
     SaveMatplotlibFigTool,
 )
@@ -14,7 +14,9 @@ def test_draw_image_agent():
     prompt = """Propose moi un schéma très détaillé d'un réacteur nucléaire hélium graphite."""
     output_path = Path(DATA_DIR) / "figure.png"
     output_path.unlink(missing_ok=True)
-    output_path = draw_image_agent(prompt, output_path, multiple_steps=False)
+    output_path = draw_matplotlib_image_from_prompt(
+        prompt, output_path, multiple_steps=False
+    )
     assert output_path.exists()
 
 
@@ -41,9 +43,6 @@ def test_save_matplotlib_fig_tool():
     assert (Path(DATA_DIR) / "figure.png").exists()
 
 
-test_save_matplotlib_fig_tool()
-
-
 @pytest.mark.skip(reason="This function is not working yet")
 def test_run_agent_step_by_step():
     prompt = """Propose moi un schéma très détaillé d'un réacteur nucléaire hélium graphite."""