hibam1

updat ethe agent to add plot - update

15c82ee 17 days ago

3.49 kB

import os

from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool

from cargo_tool import calculate_cargo_travel_time
from model import model

# TASK for single agent:

# TASK = """
# Find all Batman filming locations in the world, 
# calculate the time to transfer via cargo plane to here (we're in Gotham, 40.7128° N, 74.0060° W), and return them to me as a pandas dataframe.
# Also give me some supercar factories with the same cargo plane transfer time."""

# def run_single_agent() -> None:
#     agent = CodeAgent(
#         model=model,
#         tools=[DuckDuckGoSearchTool(), VisitWebpageTool(), calculate_cargo_travel_time],
#         additional_authorized_imports=["pandas"],
#         max_steps=20,
#     )
#     print(agent.run(TASK))


# if __name__ == "__main__":
#     run_single_agent()

# TASK for multi-agent:

def build_web_agent() -> CodeAgent:
    return CodeAgent(
        model=model,
        tools=[DuckDuckGoSearchTool(), VisitWebpageTool()],
        name="web_agent",
        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()],
        # "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 = """
Find all Batman filming locations in the world, calculate the time to transfer via cargo plane to here (we're in Gotham, 40.7128° N, 74.0060° W).
Also give me some supercar factories with the same cargo plane transfer time. You need at least 6 points in total.
Represent this as spatial map of the world, with the locations represented as scatter points with a color that depends on the travel time, and save it to saved_map.png!
Here's an example of how to plot and return a map:
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.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.
"""

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)

	import os

	from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool

	from cargo_tool import calculate_cargo_travel_time
	from model import model

	# TASK for single agent:

	# TASK = """
	# Find all Batman filming locations in the world,
	# calculate the time to transfer via cargo plane to here (we're in Gotham, 40.7128° N, 74.0060° W), and return them to me as a pandas dataframe.
	# Also give me some supercar factories with the same cargo plane transfer time."""

	# def run_single_agent() -> None:
	# agent = CodeAgent(
	# model=model,
	# tools=[DuckDuckGoSearchTool(), VisitWebpageTool(), calculate_cargo_travel_time],
	# additional_authorized_imports=["pandas"],
	# max_steps=20,
	# )
	# print(agent.run(TASK))


	# if __name__ == "__main__":
	# run_single_agent()

	# TASK for multi-agent:

	def build_web_agent() -> CodeAgent:
	return CodeAgent(
	model=model,
	tools=[DuckDuckGoSearchTool(), VisitWebpageTool()],
	name="web_agent",
	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()],
	# "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 = """
	Find all Batman filming locations in the world, calculate the time to transfer via cargo plane to here (we're in Gotham, 40.7128° N, 74.0060° W).
	Also give me some supercar factories with the same cargo plane transfer time. You need at least 6 points in total.
	Represent this as spatial map of the world, with the locations represented as scatter points with a color that depends on the travel time, and save it to saved_map.png!

	Here's an example of how to plot and return a map:
	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.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.

	"""

	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)