You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
swarms/aop/test_aop.py

67 lines
1.7 KiB

from swarms.structs.aop import AOP
# Initialize the AOP instance
aop = AOP(
name="example_system",
description="A simple example of tools, agents, and swarms",
)
# Define a simple tool
@aop.tool(name="calculator", description="A simple calculator tool")
async def calculator(operation: str, x: float, y: float):
"""
Performs basic arithmetic operations
"""
if operation == "add":
return x + y
elif operation == "multiply":
return x * y
else:
raise ValueError("Unsupported operation")
# Define an agent that uses the calculator tool
@aop.agent(
name="math_agent",
description="Agent that performs mathematical operations",
)
async def math_agent(operation: str, numbers: list[float]):
"""
Agent that chains multiple calculations together
"""
result = numbers[0]
for num in numbers[1:]:
# Using the calculator tool within the agent
result = await aop.call_tool_or_agent(
"calculator",
{"operation": operation, "x": result, "y": num},
)
return result
# Define a swarm that coordinates multiple agents
@aop.swarm(
name="math_swarm",
description="Swarm that coordinates mathematical operations",
)
async def math_swarm(numbers: list[float]):
"""
Swarm that performs multiple operations on a set of numbers
"""
# Perform addition and multiplication in parallel
results = await aop.call_tool_or_agent_concurrently(
names=["math_agent", "math_agent"],
arguments=[
{"operation": "add", "numbers": numbers},
{"operation": "multiply", "numbers": numbers},
],
)
return {"sum": results[0], "product": results[1]}
# Example usage
if __name__ == "__main__":
aop.run_sse()