File size: 4,511 Bytes
8a87cbb | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | import pytest
from aegis.agents import EvaluatorAgent, PlannerAgent
from aegis.core import (
Agent,
AgentContext,
AgentRegistry,
AgentResult,
Conductor,
MemoryStore,
Task,
WorkflowEngine,
)
class WorkerAgent(Agent):
name = "worker"
async def run(self, context: AgentContext) -> AgentResult:
return AgentResult(
agent_name=self.name,
success=True,
output={
"task": context.task.name,
"completed": True,
},
)
class FailingWorkerAgent(Agent):
name = "failing-worker"
async def run(self, context: AgentContext) -> AgentResult:
return AgentResult(
agent_name=self.name,
success=False,
errors=["Worker execution failed."],
)
class FailingPlannerAgent(Agent):
name = "planner"
async def run(self, context: AgentContext) -> AgentResult:
return AgentResult(
agent_name=self.name,
success=False,
errors=["Planning failed."],
)
def make_context() -> AgentContext:
return AgentContext(
task=Task(
name="Execute workflow",
description="Run and evaluate a target agent.",
)
)
@pytest.mark.asyncio
async def test_workflow_runs_plan_execution_and_evaluation() -> None:
registry = AgentRegistry(
[
PlannerAgent(),
WorkerAgent(),
EvaluatorAgent(),
]
)
memory = MemoryStore()
engine = WorkflowEngine(
conductor=Conductor(registry),
memory=memory,
)
context = make_context()
result = await engine.run(context, "worker")
assert result.plan.agent_name == "planner"
assert result.plan.success is True
assert result.execution.agent_name == "worker"
assert result.execution.success is True
assert result.execution.output["completed"] is True
assert result.evaluation.agent_name == "evaluator"
assert result.evaluation.success is True
assert result.evaluation.output["passed"] is True
assert result.evaluation.output["score"] == 100
assert len(result.memory_keys) == 3
@pytest.mark.asyncio
async def test_workflow_stores_all_results_in_memory() -> None:
registry = AgentRegistry(
[
PlannerAgent(),
WorkerAgent(),
EvaluatorAgent(),
]
)
memory = MemoryStore()
engine = WorkflowEngine(
conductor=Conductor(registry),
memory=memory,
)
context = make_context()
await engine.run(context, "worker")
prefix = f"workflow.{context.run_id}"
plan = memory.get(f"{prefix}.plan")
execution = memory.get(f"{prefix}.execution")
evaluation = memory.get(f"{prefix}.evaluation")
assert plan.metadata["stage"] == "plan"
assert plan.metadata["agent_name"] == "planner"
assert execution.metadata["stage"] == "execution"
assert execution.metadata["agent_name"] == "worker"
assert evaluation.metadata["stage"] == "evaluation"
assert evaluation.metadata["agent_name"] == "evaluator"
@pytest.mark.asyncio
async def test_workflow_evaluates_failed_execution() -> None:
registry = AgentRegistry(
[
PlannerAgent(),
FailingWorkerAgent(),
EvaluatorAgent(),
]
)
engine = WorkflowEngine(
conductor=Conductor(registry),
memory=MemoryStore(),
)
result = await engine.run(
make_context(),
"failing-worker",
)
assert result.execution.success is False
assert result.evaluation.success is True
assert result.evaluation.output["passed"] is False
assert result.evaluation.output["score"] == 25
@pytest.mark.asyncio
async def test_workflow_stops_when_planner_fails() -> None:
registry = AgentRegistry(
[
FailingPlannerAgent(),
WorkerAgent(),
EvaluatorAgent(),
]
)
memory = MemoryStore()
engine = WorkflowEngine(
conductor=Conductor(registry),
memory=memory,
)
context = make_context()
with pytest.raises(
RuntimeError,
match="Planner failed to produce a valid plan",
):
await engine.run(context, "worker")
prefix = f"workflow.{context.run_id}"
assert memory.contains(f"{prefix}.plan") is True
assert memory.contains(f"{prefix}.execution") is False
assert memory.contains(f"{prefix}.evaluation") is False
|