Spaces:
Sleeping
Sleeping
File size: 1,839 Bytes
1eb9c9d 1ea134d dd7fa38 1ea134d 395c3d4 1ea134d 1eb9c9d 395c3d4 1eb9c9d dd7fa38 1eb9c9d 4fbfc2b 395c3d4 1eb9c9d 1ea134d 1eb9c9d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import pytest
from deepengineer.deepsearch.draw_agent import (
draw_matplotlib_image_from_prompt,
multiple_steps_draw_image_agent,
SaveMatplotlibFigTool,
)
from deepengineer.common_path import DATA_DIR
from pathlib import Path
from smolagents import CodeAgent, LiteLLMModel
@pytest.mark.expensive
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_matplotlib_image_from_prompt(
prompt, output_path, multiple_steps=False
)
assert output_path.exists()
@pytest.mark.expensive
def test_save_matplotlib_fig_tool():
model = LiteLLMModel(model_id="mistral/mistral-medium-latest")
agent = CodeAgent(
tools=[SaveMatplotlibFigTool(output_dir=Path(DATA_DIR))],
model=model,
additional_authorized_imports=[
"matplotlib.*",
"numpy.*",
"pandas.*",
"seaborn.*",
],
max_steps=20,
verbosity_level=2,
)
agent.run(
"""
Propose moi un schéma très détaillé d'un réacteur nucléaire hélium graphite.
Save the figure as figure.png. Return a markdown string where you explain what you did and then include the image.
"""
)
assert (Path(DATA_DIR) / "figure.png").exists()
@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."""
output_path = Path(DATA_DIR) / "figure.png"
output_path.unlink(missing_ok=True)
output_path = multiple_steps_draw_image_agent(prompt, output_path)
assert output_path.exists()
|