updat ethe agent to add plot

multiAgent.py

@@ -1,3 +1,5 @@
+import os
+
 from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool
 
 from cargo_tool import calculate_cargo_travel_time
@@ -28,23 +30,38 @@ from model import model
 def build_web_agent() -> CodeAgent:
     return CodeAgent(
         model=model,
-        tools=[DuckDuckGoSearchTool(), VisitWebpageTool(), calculate_cargo_travel_time],
+        tools=[DuckDuckGoSearchTool(), VisitWebpageTool()],
         name="web_agent",
-        description="Browses the web to find information",
+        description="Browses the web to find places. Always returns a list of: name, latitude, longitude. Does not calculate anything.",
         verbosity_level=0,
         max_steps=10,
     )
 
 
+MAP_FILE = "saved_map.png"
+
+
+def check_map_saved(final_answer, memory, agent=None) -> bool:
+    # Rejects the final answer until the map exists; the agent sees this message and keeps working
+    if not os.path.exists(MAP_FILE):
+        raise Exception(f'No map found. Plot the points with px.scatter_map and save with fig.write_image("{MAP_FILE}") before final_answer.')
+    # A real world map is tens of KB; a tiny file means a blank or fake placeholder
+    if os.path.getsize(MAP_FILE) < 10_000:
+        raise Exception(f"{MAP_FILE} is too small to be a real map. Do not write placeholder images; plot the points with px.scatter_map.")
+    return True
+
+
 def build_manager_agent() -> CodeAgent:
     return CodeAgent(
         model=model,
         tools=[calculate_cargo_travel_time],
         managed_agents=[build_web_agent()],
-        additional_authorized_imports=["geopandas", "plotly", "shapely", "json", "pandas", "numpy"],
+        # "plotly.*" (not "plotly") so the agent may import plotly.express
+        additional_authorized_imports=["geopandas.*", "plotly.*", "shapely.*", "json", "pandas", "numpy"],
         planning_interval=5,
         verbosity_level=2,
         max_steps=15,
+        final_answer_checks=[check_map_saved],
     )
 
 MAP_TASK = """
@@ -57,13 +74,23 @@ import plotly.express as px
 df = px.data.carshare()
 fig = px.scatter_map(df, lat="centroid_lat", lon="centroid_lon", text="name", color="peak_hour", size=100,
      color_continuous_scale=px.colors.sequential.Magma, size_max=15, zoom=1)
-fig.show()
-fig.write_image("saved_image.png")
+fig.write_image("saved_map.png")
 final_answer(fig)
 
 Never try to process strings using code: when you have a string to read, just print it and you'll see it.
+
+Follow this plan:
+1. Ask web_agent for Batman filming locations (name, latitude, longitude).
+2. Compute each travel time with calculate_cargo_travel_time. Never compute travel times yourself.
+3. Ask web_agent for supercar factories (name, latitude, longitude).
+4. Compute their travel times the same way, and keep only factories within 1 hour of any filming location's time. If fewer than 6 points remain in total, ask web_agent for more factories.
+5. Plot all points, save with fig.write_image("saved_map.png"), then call final_answer.
+
 """
 
 if __name__ == "__main__":
+    # Remove the map from an earlier run so the check only passes for a map made in this run
+    if os.path.exists(MAP_FILE):
+        os.remove(MAP_FILE)
     manager = build_manager_agent()
     manager.run(MAP_TASK)