| 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 |
|
|
