feat: add support for images

src/deepengineer/backend/fast_api_app.py

@@ -17,7 +17,7 @@ class AgentRequest(BaseModel):
 async def run_agent(request: AgentRequest):
     task = request.task
     # Run the agent synchronously in a background thread
-    result = await anyio.to_thread.run_sync(main_search, task)
+    result, output_image_path = await anyio.to_thread.run_sync(main_search, task)
     return JSONResponse(content={"result": result})
 
 

src/deepengineer/backend/gradio_tools.py

@@ -25,7 +25,7 @@ def run_agent_stream(user_input: str):
     done = threading.Event()
 
     def _worker():
-        answer_container["text"] = main_search(user_input, log_queue)
+        answer_container["text"], output_image_path = main_search(user_input, log_queue)
         done.set()
 
     threading.Thread(target=_worker, daemon=True).start()

src/deepengineer/deepsearch/main_agent.py

@@ -18,87 +18,16 @@ from pathlib import Path
 import queue
 
 
-def _create_output_image_path():
+def create_output_image_path():
     random_name_images = random.randint(1000000, 9999999)
     output_image_path = Path(DATA_DIR) / f"images_{random_name_images}"
     output_image_path.mkdir(parents=True, exist_ok=True)
     return output_image_path
 
 
-def create_main_deep_search_agent(
-    main_model_id="deepseek/deepseek-reasoner",
-    web_search_model_id="deepseek/deepseek-reasoner",
-) -> tuple[CodeAgent, Path]:
-    """
-    Deepsearch agent use multiple agents to answer the question.
-    """
-    main_model = LiteLLMModel(model_id=main_model_id)
-
-    database = DataBase()
-    web_search_agent = create_web_search_agent(
-        model_id=web_search_model_id, database=database
-    )
-
-    output_image_path = _create_output_image_path()
-
-    manager_agent = CodeAgent(
-        model=main_model,
-        tools=[SaveMatplotlibFigTool(output_dir=output_image_path)],
-        max_steps=12,
-        verbosity_level=2,
-        additional_authorized_imports=[
-            "matplotlib.*",
-            "numpy.*",
-            "pandas.*",
-            "seaborn.*",
-            "scipy.*",
-            "sympy.*",
-        ],
-        planning_interval=4,
-        managed_agents=[web_search_agent],
-    )
-
-    return manager_agent, output_image_path
-
-
-def main_deep_search(task: str):
-
-    MAIN_PROMPT = """
-You are DeepDraft, an advanced research and analysis agent specialized in deep technical research, data visualization, and comprehensive information synthesis. You have access to powerful tools for web search, document analysis, and data visualization.
-
-You will be given a task to complete. This task is related to engineering, science, and technology. I want you to answer the question with a very detailed answer. The output should be written in markdown and include sources and images.
-
-## Your Capabilities
-
-### **Python and scientific reasoning**
-You will have to answer the question as an engineer. Make hypothesis, test them, and draw conclusions. When in doubt, go back to the basic equations and laws of physics. Make simple models that you can test using python.
-You are a coding agent so you can use python to test your hypothesis.
-You have access to these libraries: matplotlib, numpy, pandas, seaborn, scipy, sympy.
-You do not have access to any other library.
-
-### Managed Agent: **web_search_agent**
-This agent can search the web using Linkup API for comprehensive research with sourced answers. It can also search arXiv, PubMed, and ScienceDirect, download the documents and extract the relevant information.
-When calling this agent, you should provide a detailed task, he is extremely powerful.
-
-
-### **Data Visualization Tools**
-- You can always use the tool `SaveMatplotlibFigTool` to save a figure at the end of a matplotlib code block. You can then include the figure in your final answer.
-
-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:
-{task}
-"""
-    agent, output_image_path = create_main_deep_search_agent()
-
-    answer = agent.run(MAIN_PROMPT.format(task=task))
-
-    print(answer)
-
 
 def create_main_search_agent(
-    model_id="deepseek/deepseek-reasoner", database: DataBase | None = None, log_queue: queue.Queue | None = None
+    model_id="deepseek/deepseek-reasoner", database: DataBase | None = None, log_queue: queue.Queue | None = None, output_image_path: Path | None = None,
 ):
     """
     Simple agent that can search the web and answer the question. This is much faster and better for simple questions that do not require deep research.
@@ -107,8 +36,9 @@ def create_main_search_agent(
     model = LiteLLMModel(model_id=model_id)
     if database is None:
         database = DataBase()
-
-    output_image_path = _create_output_image_path()
+        
+    output_image_path = output_image_path or DATA_DIR / "images"
+    output_image_path.mkdir(parents=True, exist_ok=True)
 
     # Web search and crawling tools
     WEB_SEARCH_TOOLS = [
@@ -140,6 +70,7 @@ def create_main_search_agent(
 
 
 def main_search(task: str, log_queue: queue.Queue | None = None):
+    output_image_path = create_output_image_path()
     MAIN_PROMPT = """
 You are DeepDraft, an advanced research and analysis agent specialized in deep technical research, data visualization, and comprehensive information synthesis. You have access to powerful tools for web search, document analysis, and data visualization.
 
@@ -165,9 +96,10 @@ Failure or 'I cannot answer' or 'None found' will not be tolerated, success will
 Run verification steps if that's needed, you must make sure you find the correct answer! Here is the task:
 {task}
 """
-    agent = create_main_search_agent(model_id="mistral/mistral-medium-latest", log_queue=log_queue)
+    agent = create_main_search_agent(model_id="mistral/mistral-medium-latest", log_queue=log_queue, output_image_path=output_image_path)
     answer = agent.run(MAIN_PROMPT.format(task=task))
-    return answer
+    
+    return answer, output_image_path
 
 
 if __name__ == "__main__":

src/deepengineer/deepsearch/main_deep_search.py

@@ -0,0 +1,86 @@
+"""work in progress"""
+from pathlib import Path
+import random
+from smolagents import CodeAgent, LiteLLMModel
+from deepengineer.webcrawler.crawl_database import DataBase
+from deepengineer.deepsearch.scawl_web_agent import create_web_search_agent
+from deepengineer.deepsearch.draw_agent import SaveMatplotlibFigTool
+from deepengineer.common_path import DATA_DIR
+
+def _create_output_image_path(image_folder_suffix: int | None = None):
+    output_image_path = Path(DATA_DIR) / f"images_{image_folder_suffix}"
+    output_image_path.mkdir(parents=True, exist_ok=True)
+    return output_image_path
+
+def create_main_deep_search_agent(
+    main_model_id="deepseek/deepseek-reasoner",
+    web_search_model_id="deepseek/deepseek-reasoner",
+) -> tuple[CodeAgent, Path]:
+    """
+    Deepsearch agent use multiple agents to answer the question.
+    """
+    main_model = LiteLLMModel(model_id=main_model_id)
+
+    database = DataBase()
+    web_search_agent = create_web_search_agent(
+        model_id=web_search_model_id, database=database
+    )
+    
+    
+    image_folder_suffix = random.randint(1000000, 9999999)
+    output_image_path = _create_output_image_path(image_folder_suffix=image_folder_suffix)
+
+    manager_agent = CodeAgent(
+        model=main_model,
+        tools=[SaveMatplotlibFigTool(output_dir=output_image_path)],
+        max_steps=12,
+        verbosity_level=2,
+        additional_authorized_imports=[
+            "matplotlib.*",
+            "numpy.*",
+            "pandas.*",
+            "seaborn.*",
+            "scipy.*",
+            "sympy.*",
+        ],
+        planning_interval=4,
+        managed_agents=[web_search_agent],
+    )
+
+    return manager_agent, output_image_path
+
+
+def main_deep_search(task: str):
+
+    MAIN_PROMPT = """
+You are DeepDraft, an advanced research and analysis agent specialized in deep technical research, data visualization, and comprehensive information synthesis. You have access to powerful tools for web search, document analysis, and data visualization.
+
+You will be given a task to complete. This task is related to engineering, science, and technology. I want you to answer the question with a very detailed answer. The output should be written in markdown and include sources and images.
+
+## Your Capabilities
+
+### **Python and scientific reasoning**
+You will have to answer the question as an engineer. Make hypothesis, test them, and draw conclusions. When in doubt, go back to the basic equations and laws of physics. Make simple models that you can test using python.
+You are a coding agent so you can use python to test your hypothesis.
+You have access to these libraries: matplotlib, numpy, pandas, seaborn, scipy, sympy.
+You do not have access to any other library.
+
+### Managed Agent: **web_search_agent**
+This agent can search the web using Linkup API for comprehensive research with sourced answers. It can also search arXiv, PubMed, and ScienceDirect, download the documents and extract the relevant information.
+When calling this agent, you should provide a detailed task, he is extremely powerful.
+
+
+### **Data Visualization Tools**
+- You can always use the tool `SaveMatplotlibFigTool` to save a figure at the end of a matplotlib code block. You can then include the figure in your final answer.
+
+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:
+{task}
+"""
+    agent, output_image_path = create_main_deep_search_agent()
+
+    answer = agent.run(MAIN_PROMPT.format(task=task))
+
+    print(answer)