diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..84b8b76263673f1d889c0a9398adeea821b74b92 --- /dev/null +++ b/.gitignore @@ -0,0 +1,33 @@ +# Python +__pycache__/ +*.py[cod] +*.egg-info/ +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +.coverage +htmlcov/ + +# Environments +.venv/ +venv/ +.env +.env.* + +# Build +build/ +dist/ + +# Editors and OS +.vscode/ +.idea/ +.DS_Store + +# Runtime artifacts +logs/ +runs/ +artifacts/ +*.log + +__pycache__/ +*.py[cod] diff --git a/.mypy_cache/.gitignore b/.mypy_cache/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..f7208da358e6ca5f313f75fa765c7a677bf9d8da --- /dev/null +++ b/.mypy_cache/.gitignore @@ -0,0 +1,2 @@ +# Automatically created by mypy +* diff --git a/.mypy_cache/CACHEDIR.TAG b/.mypy_cache/CACHEDIR.TAG new file mode 100644 index 0000000000000000000000000000000000000000..a001cd238467894c46ce678ea349f0df86f4d27e --- /dev/null +++ b/.mypy_cache/CACHEDIR.TAG @@ -0,0 +1,3 @@ +Signature: 8a477f597d28d172789f06886806bc55 +# This file is a cache directory tag automatically created by mypy. +# For information about cache directory tags see https://bford.info/cachedir/ diff --git a/.pytest_cache/.gitignore b/.pytest_cache/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..bc1a1f6167d09c909aad37280b760bb715d0f1da --- /dev/null +++ b/.pytest_cache/.gitignore @@ -0,0 +1,2 @@ +# Created by pytest automatically. +* diff --git a/.pytest_cache/CACHEDIR.TAG b/.pytest_cache/CACHEDIR.TAG new file mode 100644 index 0000000000000000000000000000000000000000..fce15ad7eaa74e5682b644c84efb75334c112f95 --- /dev/null +++ b/.pytest_cache/CACHEDIR.TAG @@ -0,0 +1,4 @@ +Signature: 8a477f597d28d172789f06886806bc55 +# This file is a cache directory tag created by pytest. +# For information about cache directory tags, see: +# https://bford.info/cachedir/spec.html diff --git a/.pytest_cache/README.md b/.pytest_cache/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b89018ced91c0a8af7f3f23ce8901870da89f3a0 --- /dev/null +++ b/.pytest_cache/README.md @@ -0,0 +1,8 @@ +# pytest cache directory # + +This directory contains data from the pytest's cache plugin, +which provides the `--lf` and `--ff` options, as well as the `cache` fixture. + +**Do not** commit this to version control. + +See [the docs](https://docs.pytest.org/en/stable/how-to/cache.html) for more information. diff --git a/.pytest_cache/v/cache/nodeids b/.pytest_cache/v/cache/nodeids new file mode 100644 index 0000000000000000000000000000000000000000..b2d79c3f2aab168f18adde48d2f90ead5ceaaeca --- /dev/null +++ b/.pytest_cache/v/cache/nodeids @@ -0,0 +1,58 @@ +[ + "tests/unit/test_cli_run.py::test_main_supports_help_and_execution[argv0-0]", + "tests/unit/test_cli_run.py::test_main_supports_help_and_execution[argv1-0]", + "tests/unit/test_cli_run.py::test_run_command_evaluator_failure", + "tests/unit/test_cli_run.py::test_run_command_json_output_failure", + "tests/unit/test_cli_run.py::test_run_command_json_output_success", + "tests/unit/test_cli_run.py::test_run_command_max_steps_forwarded", + "tests/unit/test_cli_run.py::test_run_command_rejects_empty_task", + "tests/unit/test_cli_run.py::test_run_command_runtime_exception", + "tests/unit/test_cli_run.py::test_run_command_success_with_text_output", + "tests/unit/test_cli_run.py::test_run_command_unknown_agent", + "tests/unit/test_cli_run.py::test_run_command_unknown_workflow", + "tests/unit/test_cli_run.py::test_run_command_verbose_failure_has_no_stack_trace_for_known_error", + "tests/unit/test_conductor.py::test_conductor_executes_registered_agent", + "tests/unit/test_conductor.py::test_conductor_rejects_mismatched_result_name", + "tests/unit/test_conductor.py::test_conductor_rejects_missing_agent", + "tests/unit/test_evaluator.py::test_evaluator_fails_unsuccessful_result", + "tests/unit/test_evaluator.py::test_evaluator_passes_valid_result", + "tests/unit/test_evaluator.py::test_evaluator_rejects_invalid_result", + "tests/unit/test_evaluator.py::test_evaluator_rejects_missing_result", + "tests/unit/test_memory.py::test_memory_store_adds_and_gets_entry", + "tests/unit/test_memory.py::test_memory_store_clear_removes_all_entries", + "tests/unit/test_memory.py::test_memory_store_deletes_entry", + "tests/unit/test_memory.py::test_memory_store_initializes_with_entries", + "tests/unit/test_memory.py::test_memory_store_normalizes_keys", + "tests/unit/test_memory.py::test_memory_store_raises_for_missing_key", + "tests/unit/test_memory.py::test_memory_store_rejects_delete_for_missing_key", + "tests/unit/test_memory.py::test_memory_store_rejects_duplicate_key", + "tests/unit/test_memory.py::test_memory_store_rejects_empty_key_lookup", + "tests/unit/test_memory.py::test_memory_store_set_replaces_existing_entry", + "tests/unit/test_models.py::test_agent_context_contains_task", + "tests/unit/test_models.py::test_agent_result_success", + "tests/unit/test_models.py::test_task_defaults_to_pending", + "tests/unit/test_planner.py::test_planner_output_is_bounded", + "tests/unit/test_planner.py::test_planner_returns_structured_plan", + "tests/unit/test_planner.py::test_planner_steps_have_sequential_ids", + "tests/unit/test_policies.py::test_agent_policy_combines_response_and_execution_rules", + "tests/unit/test_policies.py::test_execution_policy_defaults_to_restricted_permissions", + "tests/unit/test_policies.py::test_invalid_policy_limit_is_rejected", + "tests/unit/test_policies.py::test_response_policy_defaults", + "tests/unit/test_registry.py::test_registry_raises_when_agent_is_missing", + "tests/unit/test_registry.py::test_registry_registers_and_resolves_agent", + "tests/unit/test_registry.py::test_registry_rejects_duplicate_agent_names", + "tests/unit/test_registry.py::test_registry_returns_sorted_names", + "tests/unit/test_tool_registry.py::test_registry_accepts_initial_tools", + "tests/unit/test_tool_registry.py::test_registry_normalizes_tool_names", + "tests/unit/test_tool_registry.py::test_registry_raises_when_tool_is_missing", + "tests/unit/test_tool_registry.py::test_registry_registers_and_resolves_tool", + "tests/unit/test_tool_registry.py::test_registry_rejects_duplicate_tool_names", + "tests/unit/test_tool_registry.py::test_registry_rejects_empty_tool_name", + "tests/unit/test_tool_registry.py::test_registry_returns_sorted_names", + "tests/unit/test_tool_runner.py::test_tool_runner_executes_registered_tool", + "tests/unit/test_tool_runner.py::test_tool_runner_preserves_arguments", + "tests/unit/test_workflow.py::test_workflow_evaluates_failed_execution", + "tests/unit/test_workflow.py::test_workflow_runs_plan_execution_and_evaluation", + "tests/unit/test_workflow.py::test_workflow_stops_when_planner_fails", + "tests/unit/test_workflow.py::test_workflow_stores_all_results_in_memory" +] \ No newline at end of file diff --git a/.ruff_cache/.gitignore b/.ruff_cache/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..b100327770759f0efb7859fe78aca4801d69ed61 --- /dev/null +++ b/.ruff_cache/.gitignore @@ -0,0 +1,2 @@ +# Automatically created by ruff. +* diff --git a/.ruff_cache/0.15.2/12944381863948167845 b/.ruff_cache/0.15.2/12944381863948167845 new file mode 100644 index 0000000000000000000000000000000000000000..7208871bb6071aef8ba7ecc460c15fe1406b67a5 Binary files /dev/null and b/.ruff_cache/0.15.2/12944381863948167845 differ diff --git a/.ruff_cache/0.15.2/9042146699542604200 b/.ruff_cache/0.15.2/9042146699542604200 new file mode 100644 index 0000000000000000000000000000000000000000..3ab34f87c02da06000e5e492fc8d746c305cfcf1 Binary files /dev/null and b/.ruff_cache/0.15.2/9042146699542604200 differ diff --git a/.ruff_cache/0.16.0/3506816079193222284 b/.ruff_cache/0.16.0/3506816079193222284 new file mode 100644 index 0000000000000000000000000000000000000000..6c40e2ed24d3b2bb6eadcd7ee0011ebb3c039da3 Binary files /dev/null and b/.ruff_cache/0.16.0/3506816079193222284 differ diff --git a/.ruff_cache/0.16.0/6668019937987631587 b/.ruff_cache/0.16.0/6668019937987631587 new file mode 100644 index 0000000000000000000000000000000000000000..b65c8147c2072790daa230c229664926d80d9443 Binary files /dev/null and b/.ruff_cache/0.16.0/6668019937987631587 differ diff --git a/.ruff_cache/0.16.0/9602145004845211386 b/.ruff_cache/0.16.0/9602145004845211386 new file mode 100644 index 0000000000000000000000000000000000000000..9347d27987cf58b0640b4e9d3f731b6f8da77cb5 Binary files /dev/null and b/.ruff_cache/0.16.0/9602145004845211386 differ diff --git a/.ruff_cache/CACHEDIR.TAG b/.ruff_cache/CACHEDIR.TAG new file mode 100644 index 0000000000000000000000000000000000000000..bc1ecb967a482524e7736038de0df6e08f9ee452 --- /dev/null +++ b/.ruff_cache/CACHEDIR.TAG @@ -0,0 +1 @@ +Signature: 8a477f597d28d172789f06886806bc55 \ No newline at end of file diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/API.md b/API.md new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..00db2dcfbca78a6361331804f1a16d17f7f237b0 --- /dev/null +++ b/README.md @@ -0,0 +1,101 @@ +# Aegis + +Aegis is a modular Python framework for composing agent workflows, registries, memory, tool execution, and evaluation. The repository now ships a runnable CLI that executes tasks through the built-in orchestration layer. + +## What it does now + +- Registers agents and tools through typed registries +- Executes workflows through a conductor and workflow engine +- Stores run-scoped results in an in-memory memory store +- Evaluates run output with a deterministic evaluator +- Exposes a user-facing CLI: `aegis run` + +## Installation + +```bash +cd /Users/giancarlovizhnay/aegis +python -m venv .venv +source .venv/bin/activate +pip install -e ".[dev]" +aegis run --help +aegis run "Run the default Aegis workflow" +``` + +## CLI usage + +Default text output: + +```bash +aegis run "Review this repository" +``` + +Run a specific agent: + +```bash +aegis run --agent default-task "Review this repository" +``` + +Run a named workflow: + +```bash +aegis run --workflow review "Review this repository" +``` + +JSON output: + +```bash +aegis run --output json "Review this repository" +``` + +Limit execution steps: + +```bash +aegis run --max-steps 10 "Review this repository" +``` + +## Available options + +- `--agent `: run a specific registered agent +- `--workflow `: run a named workflow +- `--output text|json`: choose human-readable or structured output +- `--max-steps `: pass a step limit into the runtime context +- `--verbose`: show debug details for unexpected failures only + +## Example output + +```text +Run ID: run_... +Status: completed +Agent: default-task +Workflow: default +Result: +{ + "summary": "Received task: Review this repository" +} +``` + +## Exit codes + +- `0`: success +- `1`: unexpected internal error +- `2`: invalid CLI input +- `3`: configuration or registry error +- `4`: policy violation +- `5`: execution failure +- `6`: evaluation failure + +## Current limitations + +- Persistence is still in-memory and resets on restart. +- The default workflow is intentionally small and deterministic. +- No external model provider is wired in yet. +- No HTTP API is exposed yet. +- Tool integrations are still skeletal. + +## End-to-end example + +```bash +aegis run --workflow review --output json "Review this repository" +``` + +This returns a structured JSON report with the run ID, selected workflow, execution result, evaluation payload, and any error information. diff --git a/ROADMAP.md b/ROADMAP.md new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/TESTING.md b/TESTING.md new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..ff0138fee2ea4de3b5ad0506d37e852cc7dd66b2 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,51 @@ +[build-system] +requires = ["hatchling>=1.27"] +build-backend = "hatchling.build" + +[project] +name = "aegis-agent" +version = "0.1.0" +description = "A modular, reliable AI agent engineering framework." +readme = "README.md" +requires-python = ">=3.11" +license = { file = "LICENSE" } +authors = [ + { name = "Giancarlo Vizhnay" } +] +dependencies = [ + "pydantic>=2.10,<3.0", +] + +[project.optional-dependencies] +dev = [ + "mypy>=1.14", + "pytest>=8.3", + "pytest-cov>=6.0", + "pytest-asyncio>=0.25", + "ruff>=0.9", +] + +[project.scripts] +aegis = "aegis.cli.app:main" + +[tool.hatch.build.targets.wheel] +packages = ["src/aegis"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["src"] +addopts = "-q --strict-markers --disable-warnings" + +[tool.ruff] +target-version = "py311" +line-length = 100 +src = ["src", "tests"] + +[tool.ruff.lint] +select = ["E", "F", "I", "B", "UP", "SIM", "RUF"] + +[tool.mypy] +python_version = "3.11" +strict = true +mypy_path = "src" +packages = ["aegis"] diff --git a/src/aegis/__init__.py b/src/aegis/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b2d59a80b5f6a3456abfff49a859cbf277c1a5d7 --- /dev/null +++ b/src/aegis/__init__.py @@ -0,0 +1,11 @@ +from aegis.core import Agent, AgentContext, AgentResult, Task, TaskStatus + +__version__ = "0.1.0" + +__all__ = [ + "Agent", + "AgentContext", + "AgentResult", + "Task", + "TaskStatus", +] diff --git a/src/aegis/agents/__init__.py b/src/aegis/agents/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f727a5c266c2d9a7c0ed2a9e959f4f1a48c5bf22 --- /dev/null +++ b/src/aegis/agents/__init__.py @@ -0,0 +1,10 @@ +from aegis.agents.evaluator import EvaluationReport, EvaluatorAgent +from aegis.agents.planner import ExecutionPlan, PlannerAgent, PlanStep + +__all__ = [ + "EvaluationReport", + "EvaluatorAgent", + "ExecutionPlan", + "PlanStep", + "PlannerAgent", +] diff --git a/src/aegis/agents/evaluator.py b/src/aegis/agents/evaluator.py new file mode 100644 index 0000000000000000000000000000000000000000..502f74b3ee2e0fc94decaa792c87e5ca4bcf9aa8 --- /dev/null +++ b/src/aegis/agents/evaluator.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + +from aegis.core import Agent, AgentContext, AgentResult + + +class EvaluationReport(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + passed: bool + score: int = Field(ge=0, le=100) + issues: tuple[str, ...] = () + evaluated_agent: str = Field(min_length=1) + + +class EvaluatorAgent(Agent): + name = "evaluator" + + async def run(self, context: AgentContext) -> AgentResult: + raw_result: Any = context.inputs.get("result") + + if raw_result is None: + return AgentResult( + agent_name=self.name, + success=False, + errors=["Missing required input: result"], + ) + + try: + result = AgentResult.model_validate(raw_result) + except ValueError as exc: + return AgentResult( + agent_name=self.name, + success=False, + errors=[f"Invalid agent result: {exc}"], + ) + + issues: list[str] = [] + + if not result.success: + issues.append("Agent execution reported failure.") + + if not result.output: + issues.append("Agent result contains no output.") + + if result.errors: + issues.append("Agent result contains errors.") + + score = max(0, 100 - (25 * len(issues))) + passed = not issues + + report = EvaluationReport( + passed=passed, + score=score, + issues=tuple(issues), + evaluated_agent=result.agent_name, + ) + + return AgentResult( + agent_name=self.name, + success=True, + output=report.model_dump(), + ) diff --git a/src/aegis/agents/planner.py b/src/aegis/agents/planner.py new file mode 100644 index 0000000000000000000000000000000000000000..c934c687e550ecb353b24e20b6f83b428257d943 --- /dev/null +++ b/src/aegis/agents/planner.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from pydantic import BaseModel, Field + +from aegis.core import Agent, AgentContext, AgentResult + + +class PlanStep(BaseModel): + id: int = Field(ge=1) + action: str = Field(min_length=1) + expected_output: str = Field(min_length=1) + + +class ExecutionPlan(BaseModel): + task_name: str = Field(min_length=1) + steps: tuple[PlanStep, ...] = Field(min_length=1) + + +class PlannerAgent(Agent): + name = "planner" + + async def run(self, context: AgentContext) -> AgentResult: + task = context.task + + plan = ExecutionPlan( + task_name=task.name, + steps=( + PlanStep( + id=1, + action="Analyze the task requirements.", + expected_output="A clear statement of scope and constraints.", + ), + PlanStep( + id=2, + action="Select the required tools and agents.", + expected_output="A bounded execution strategy.", + ), + PlanStep( + id=3, + action="Execute and validate the work.", + expected_output="A verified result or explicit failure state.", + ), + ), + ) + + return AgentResult( + agent_name=self.name, + success=True, + output=plan.model_dump(), + ) diff --git a/src/aegis/cli/__init__.py b/src/aegis/cli/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4e3a9eed363471047332a618ab9b213abc98d831 --- /dev/null +++ b/src/aegis/cli/__init__.py @@ -0,0 +1,3 @@ +from aegis.cli.app import main + +__all__ = ["main"] diff --git a/src/aegis/cli/app.py b/src/aegis/cli/app.py new file mode 100644 index 0000000000000000000000000000000000000000..6942b3a97003f2481f19cd7e9f651e70bd344296 --- /dev/null +++ b/src/aegis/cli/app.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +import argparse +from collections.abc import Sequence + +from aegis.cli.commands.run import add_arguments, run_command + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="aegis", description="Aegis CLI") + subparsers = parser.add_subparsers(dest="command") + + run_parser = subparsers.add_parser("run", help="Run a task through Aegis") + add_arguments(run_parser) + run_parser.set_defaults(command_handler=run_command) + + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + parser = build_parser() + + try: + args = parser.parse_args(list(argv) if argv is not None else None) + except SystemExit as exc: + return int(exc.code) if isinstance(exc.code, int) else 1 + + command_handler = getattr(args, "command_handler", None) + if command_handler is None: + parser.print_help() + return 0 + + return int(command_handler(args)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/aegis/cli/commands/__init__.py b/src/aegis/cli/commands/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/aegis/cli/commands/run.py b/src/aegis/cli/commands/run.py new file mode 100644 index 0000000000000000000000000000000000000000..766eeedd45f9cf48c3233a6944be0d2fb40a81f9 --- /dev/null +++ b/src/aegis/cli/commands/run.py @@ -0,0 +1,193 @@ +from __future__ import annotations + +import argparse +import sys +import traceback +from typing import Any +from uuid import uuid4 + +from aegis.cli.output import format_run_json, format_run_text +from aegis.runtime import ( + AegisRunError, + AegisRuntime, + EvaluationFailureError, + ExecutionFailureError, + InternalFailureError, + InvalidConfigurationError, + InvalidTaskError, + PolicyViolationError, + RunOutcome, + UnknownAgentError, + UnknownWorkflowError, +) + + +def add_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "task", + help="Task to run through Aegis.", + ) + group = parser.add_mutually_exclusive_group() + group.add_argument( + "--agent", + help="Run a specific registered agent.", + ) + group.add_argument( + "--workflow", + help="Run a named workflow.", + ) + parser.add_argument( + "--output", + choices=("text", "json"), + default="text", + help="Output format.", + ) + parser.add_argument( + "--max-steps", + type=_positive_int, + default=None, + help="Maximum execution steps for the run.", + ) + parser.add_argument( + "--verbose", + action="store_true", + help="Show debug details for unexpected failures.", + ) + + +def run_command(args: argparse.Namespace, runtime: Any | None = None) -> int: + runner: Any = runtime or AegisRuntime() + run_uuid = uuid4() + run_id = f"run_{run_uuid.hex}" + requested_agent = _normalize_optional(str(args.agent)) if getattr(args, "agent", None) else None + requested_workflow = ( + _normalize_optional(str(args.workflow)) if getattr(args, "workflow", None) else None + ) + max_steps = int(args.max_steps) if getattr(args, "max_steps", None) is not None else None + + try: + outcome = runner.run( + task=str(args.task), + agent_name=requested_agent, + workflow_name=requested_workflow, + max_steps=max_steps, + run_id=run_uuid, + ) + except ( + InvalidTaskError, + InvalidConfigurationError, + UnknownAgentError, + UnknownWorkflowError, + ) as exc: + return _emit_failure( + run_id=run_id, + status="failed", + agent=requested_agent, + workflow=requested_workflow, + message=str(exc), + exit_code=exc.exit_code, + output_format=args.output, + verbose=args.verbose, + ) + except PolicyViolationError as exc: + return _emit_failure( + run_id=run_id, + status="failed", + agent=requested_agent, + workflow=requested_workflow, + message=str(exc), + exit_code=exc.exit_code, + output_format=args.output, + verbose=args.verbose, + ) + except (ExecutionFailureError, EvaluationFailureError) as exc: + return _emit_failure( + run_id=run_id, + status="failed", + agent=requested_agent, + workflow=requested_workflow, + message=str(exc), + exit_code=exc.exit_code, + output_format=args.output, + verbose=args.verbose, + ) + except AegisRunError as exc: + return _emit_failure( + run_id=run_id, + status="failed", + agent=requested_agent, + workflow=requested_workflow, + message=str(exc), + exit_code=exc.exit_code, + output_format=args.output, + verbose=args.verbose, + ) + except Exception as exc: # pragma: no cover - defensive boundary + error = InternalFailureError("Unexpected internal error.") + return _emit_failure( + run_id=run_id, + status="failed", + agent=requested_agent, + workflow=requested_workflow, + message=str(error), + exit_code=error.exit_code, + output_format=args.output, + verbose=args.verbose, + exc=exc, + ) + + if args.output == "json": + print(format_run_json(outcome)) + else: + print(format_run_text(outcome)) + + return 0 + + + + +def _emit_failure( + *, + run_id: str, + status: str, + agent: str | None, + workflow: str | None, + message: str, + exit_code: int, + output_format: str, + verbose: bool, + exc: BaseException | None = None, +) -> int: + outcome = RunOutcome( + run_id=run_id, + status=status, + agent=agent, + workflow=workflow, + result=None, + evaluation=None, + error=message, + ) + + if output_format == "json": + print(format_run_json(outcome)) + else: + _emit_error(message, verbose, exc) + + return exit_code + + +def _emit_error(message: str, verbose: bool, exc: BaseException | None = None) -> None: + print(f"Error: {message}", file=sys.stderr) + if verbose and exc is not None: + traceback.print_exception(exc, file=sys.stderr) + + +def _normalize_optional(value: str) -> str: + return value.strip() + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("must be a positive integer") + return parsed diff --git a/src/aegis/cli/output.py b/src/aegis/cli/output.py new file mode 100644 index 0000000000000000000000000000000000000000..c8804c4daf2d46bf0f58d0e48a9d4d8b2c29192c --- /dev/null +++ b/src/aegis/cli/output.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import json +from collections.abc import Mapping +from typing import Any + +from aegis.runtime.runner import RunOutcome + + +def format_run_text(outcome: RunOutcome) -> str: + lines: list[str] = [ + f"Run ID: {outcome.run_id}", + f"Status: {outcome.status}", + ] + + if outcome.agent is not None: + lines.append(f"Agent: {outcome.agent}") + + if outcome.workflow is not None: + lines.append(f"Workflow: {outcome.workflow}") + + if outcome.result is not None: + lines.append("Result:") + lines.append(_format_payload(outcome.result)) + else: + lines.append("Result:") + lines.append("") + + if outcome.evaluation is not None: + lines.append("Evaluation:") + lines.append(_format_payload(outcome.evaluation)) + + if outcome.error is not None: + lines.append("Error:") + lines.append(outcome.error) + + return "\n".join(lines) + + +def format_run_json(outcome: RunOutcome) -> str: + return json.dumps(outcome.model_dump(mode="json"), indent=2, sort_keys=True) + + +def _format_payload(payload: Mapping[str, Any]) -> str: + return json.dumps(payload, indent=2, sort_keys=True) diff --git a/src/aegis/core/__init__.py b/src/aegis/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..250cceb000065efc7325168e971059c9f38ff570 --- /dev/null +++ b/src/aegis/core/__init__.py @@ -0,0 +1,41 @@ +from aegis.core.conductor import Conductor +from aegis.core.contracts import Agent, ModelProvider, Tool +from aegis.core.memory import ( + DuplicateMemoryKeyError, + MemoryEntry, + MemoryKeyNotFoundError, + MemoryStore, + MemoryStoreError, +) +from aegis.core.models import AgentContext, AgentResult, Task, TaskStatus +from aegis.core.registry import ( + AgentNotFoundError, + AgentRegistry, + AgentRegistryError, + DuplicateAgentError, +) +from aegis.core.tool_runner import ToolRunner +from aegis.core.workflow import WorkflowEngine, WorkflowResult + +__all__ = [ + "Agent", + "AgentContext", + "AgentNotFoundError", + "AgentRegistry", + "AgentRegistryError", + "AgentResult", + "Conductor", + "DuplicateAgentError", + "DuplicateMemoryKeyError", + "MemoryEntry", + "MemoryKeyNotFoundError", + "MemoryStore", + "MemoryStoreError", + "ModelProvider", + "Task", + "TaskStatus", + "Tool", + "ToolRunner", + "WorkflowEngine", + "WorkflowResult", +] diff --git a/src/aegis/core/conductor.py b/src/aegis/core/conductor.py new file mode 100644 index 0000000000000000000000000000000000000000..e31359849b8c15133710cdb76231ac1407d3d1f3 --- /dev/null +++ b/src/aegis/core/conductor.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from aegis.core.models import AgentContext, AgentResult +from aegis.core.registry import AgentRegistry + + +class Conductor: + def __init__(self, registry: AgentRegistry) -> None: + self._registry = registry + + async def execute( + self, + agent_name: str, + context: AgentContext, + ) -> AgentResult: + agent = self._registry.get(agent_name) + result = await agent.run(context) + + if result.agent_name != agent.name: + raise ValueError("Agent result name does not match the registered agent name.") + + return result diff --git a/src/aegis/core/contracts.py b/src/aegis/core/contracts.py new file mode 100644 index 0000000000000000000000000000000000000000..f554dfb05dd13d29eb7706de8907ac418fc8b4fd --- /dev/null +++ b/src/aegis/core/contracts.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Protocol, runtime_checkable + +from aegis.core.models import AgentContext, AgentResult + + +class Agent(ABC): + name: str + + @abstractmethod + async def run(self, context: AgentContext) -> AgentResult: + """Execute the agent against the supplied context.""" + raise NotImplementedError + + +@runtime_checkable +class ModelProvider(Protocol): + async def generate(self, prompt: str) -> str: + """Generate model output for a prompt.""" + ... + + +@runtime_checkable +class Tool(Protocol): + name: str + + async def execute(self, **kwargs: object) -> object: + """Execute the tool with validated arguments.""" + ... diff --git a/src/aegis/core/memory.py b/src/aegis/core/memory.py new file mode 100644 index 0000000000000000000000000000000000000000..531ee39888957df06014d91bc3235c02499c7eb2 --- /dev/null +++ b/src/aegis/core/memory.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + + +class MemoryStoreError(Exception): + """Base error for memory store operations.""" + + +class DuplicateMemoryKeyError(MemoryStoreError): + """Raised when adding a key that already exists.""" + + +class MemoryKeyNotFoundError(MemoryStoreError): + """Raised when a requested key does not exist.""" + + +class MemoryEntry(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + key: str = Field(min_length=1) + value: Any + metadata: dict[str, Any] = Field(default_factory=dict) + + +class MemoryStore: + def __init__(self, entries: Iterable[MemoryEntry] = ()) -> None: + self._entries: dict[str, MemoryEntry] = {} + + for entry in entries: + self.add(entry) + + def add(self, entry: MemoryEntry) -> None: + key = self._normalize_key(entry.key) + + if key in self._entries: + raise DuplicateMemoryKeyError(f"Memory key already exists: {key}") + + normalized_entry = entry.model_copy(update={"key": key}) + self._entries[key] = normalized_entry + + def set(self, entry: MemoryEntry) -> None: + key = self._normalize_key(entry.key) + normalized_entry = entry.model_copy(update={"key": key}) + self._entries[key] = normalized_entry + + def get(self, key: str) -> MemoryEntry: + normalized_key = self._normalize_key(key) + + try: + return self._entries[normalized_key] + except KeyError as exc: + raise MemoryKeyNotFoundError(f"Memory key not found: {normalized_key}") from exc + + def delete(self, key: str) -> None: + normalized_key = self._normalize_key(key) + + if normalized_key not in self._entries: + raise MemoryKeyNotFoundError(f"Memory key not found: {normalized_key}") + + del self._entries[normalized_key] + + def contains(self, key: str) -> bool: + return self._normalize_key(key) in self._entries + + def keys(self) -> tuple[str, ...]: + return tuple(sorted(self._entries)) + + def clear(self) -> None: + self._entries.clear() + + def __len__(self) -> int: + return len(self._entries) + + @staticmethod + def _normalize_key(key: str) -> str: + normalized = key.strip() + + if not normalized: + raise ValueError("Memory key must not be empty.") + + return normalized diff --git a/src/aegis/core/models.py b/src/aegis/core/models.py new file mode 100644 index 0000000000000000000000000000000000000000..75452d6bca30100d395e9f4f17a28b7098ac2864 --- /dev/null +++ b/src/aegis/core/models.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +from enum import StrEnum +from typing import Any +from uuid import UUID, uuid4 + +from pydantic import BaseModel, ConfigDict, Field + + +class TaskStatus(StrEnum): + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + BLOCKED = "blocked" + + +class Task(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + id: UUID = Field(default_factory=uuid4) + name: str = Field(min_length=1, max_length=120) + description: str = Field(min_length=1) + status: TaskStatus = TaskStatus.PENDING + metadata: dict[str, Any] = Field(default_factory=dict) + + +class AgentContext(BaseModel): + model_config = ConfigDict(extra="forbid") + + run_id: UUID = Field(default_factory=uuid4) + task: Task + inputs: dict[str, Any] = Field(default_factory=dict) + state: dict[str, Any] = Field(default_factory=dict) + + +class AgentResult(BaseModel): + model_config = ConfigDict(extra="forbid") + + agent_name: str = Field(min_length=1) + success: bool + output: dict[str, Any] = Field(default_factory=dict) + errors: list[str] = Field(default_factory=list) + metadata: dict[str, Any] = Field(default_factory=dict) diff --git a/src/aegis/core/policies.py b/src/aegis/core/policies.py new file mode 100644 index 0000000000000000000000000000000000000000..9a60ae58dac0f5d574b01d98b6ba32d52576ee84 --- /dev/null +++ b/src/aegis/core/policies.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from enum import StrEnum + +from pydantic import BaseModel, ConfigDict, Field + + +class ApprovalMode(StrEnum): + NEVER = "never" + HIGH_RISK_ONLY = "high_risk_only" + ALWAYS = "always" + + +class ResponsePolicy(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + max_input_chars: int = Field(default=20_000, ge=1) + max_output_tokens: int = Field(default=600, ge=1) + + answer_first: bool = True + include_reasoning: bool = False + include_sources: bool = True + include_warnings: bool = True + allow_repetition: bool = False + + +class ExecutionPolicy(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + max_steps: int = Field(default=8, ge=1) + timeout_seconds: float = Field(default=60.0, gt=0) + retry_limit: int = Field(default=1, ge=0) + max_tool_calls: int = Field(default=6, ge=0) + max_parallel_tasks: int = Field(default=3, ge=1) + + approval_mode: ApprovalMode = ApprovalMode.HIGH_RISK_ONLY + allow_network: bool = False + allow_file_writes: bool = False + allow_shell: bool = False + + +class AgentPolicy(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + response: ResponsePolicy = Field(default_factory=ResponsePolicy) + execution: ExecutionPolicy = Field(default_factory=ExecutionPolicy) diff --git a/src/aegis/core/registry.py b/src/aegis/core/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..b8f38107cb206fe9b3b4d2dd40feced1adcd13fb --- /dev/null +++ b/src/aegis/core/registry.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from collections.abc import Iterable + +from aegis.core.contracts import Agent + + +class AgentRegistryError(Exception): + """Base registry error.""" + + +class DuplicateAgentError(AgentRegistryError): + """Raised when an agent name is already registered.""" + + +class AgentNotFoundError(AgentRegistryError): + """Raised when an agent cannot be found.""" + + +class AgentRegistry: + def __init__(self, agents: Iterable[Agent] | None = None) -> None: + self._agents: dict[str, Agent] = {} + + if agents is not None: + for agent in agents: + self.register(agent) + + def register(self, agent: Agent) -> None: + name = agent.name.strip() + + if not name: + raise ValueError("Agent name cannot be empty.") + + if name in self._agents: + raise DuplicateAgentError(f"Agent already registered: {name}") + + self._agents[name] = agent + + def get(self, name: str) -> Agent: + normalized_name = name.strip() + + try: + return self._agents[normalized_name] + except KeyError as exc: + raise AgentNotFoundError(f"Agent not found: {normalized_name}") from exc + + def contains(self, name: str) -> bool: + return name.strip() in self._agents + + def names(self) -> tuple[str, ...]: + return tuple(sorted(self._agents)) + + def __len__(self) -> int: + return len(self._agents) diff --git a/src/aegis/core/tool_registry.py b/src/aegis/core/tool_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..788a4de15cb8535ea4d2210fe29467805f5880c7 --- /dev/null +++ b/src/aegis/core/tool_registry.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from collections.abc import Iterable + +from aegis.core.contracts import Tool + + +class ToolRegistryError(Exception): + """Base tool registry error.""" + + +class DuplicateToolError(ToolRegistryError): + """Raised when a tool name is already registered.""" + + +class ToolNotFoundError(ToolRegistryError): + """Raised when a tool cannot be found.""" + + +class ToolRegistry: + def __init__(self, tools: Iterable[Tool] | None = None) -> None: + self._tools: dict[str, Tool] = {} + + if tools is not None: + for tool in tools: + self.register(tool) + + def register(self, tool: Tool) -> None: + name = tool.name.strip() + + if not name: + raise ValueError("Tool name cannot be empty.") + + if name in self._tools: + raise DuplicateToolError(f"Tool already registered: {name}") + + self._tools[name] = tool + + def get(self, name: str) -> Tool: + normalized_name = name.strip() + + try: + return self._tools[normalized_name] + except KeyError as exc: + raise ToolNotFoundError(f"Tool not found: {normalized_name}") from exc + + def contains(self, name: str) -> bool: + return name.strip() in self._tools + + def names(self) -> tuple[str, ...]: + return tuple(sorted(self._tools)) + + def __len__(self) -> int: + return len(self._tools) diff --git a/src/aegis/core/tool_runner.py b/src/aegis/core/tool_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..5fa16a5e664315d38333f32bffff7e39c4cffa51 --- /dev/null +++ b/src/aegis/core/tool_runner.py @@ -0,0 +1,16 @@ +from __future__ import annotations + +from aegis.core.tool_registry import ToolRegistry + + +class ToolRunner: + def __init__(self, registry: ToolRegistry) -> None: + self._registry = registry + + async def execute( + self, + tool_name: str, + **kwargs: object, + ) -> object: + tool = self._registry.get(tool_name) + return await tool.execute(**kwargs) diff --git a/src/aegis/core/workflow.py b/src/aegis/core/workflow.py new file mode 100644 index 0000000000000000000000000000000000000000..ea24781f5ffef493c92b5ecb784c3faba2931e5b --- /dev/null +++ b/src/aegis/core/workflow.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict, Field + +from aegis.core.conductor import Conductor +from aegis.core.memory import MemoryEntry, MemoryStore +from aegis.core.models import AgentContext, AgentResult + + +class WorkflowResult(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + plan: AgentResult + execution: AgentResult + evaluation: AgentResult + memory_keys: tuple[str, ...] = Field(default_factory=tuple) + + +class WorkflowEngine: + def __init__( + self, + conductor: Conductor, + memory: MemoryStore, + ) -> None: + self._conductor = conductor + self._memory = memory + + async def run( + self, + context: AgentContext, + target_agent: str, + ) -> WorkflowResult: + plan_result = await self._conductor.execute("planner", context) + self._store_result(context, "plan", plan_result) + + if not plan_result.success: + raise RuntimeError("Planner failed to produce a valid plan.") + + execution_result = await self._conductor.execute( + target_agent, + context, + ) + self._store_result(context, "execution", execution_result) + + evaluation_context = context.model_copy( + deep=True, + update={ + "inputs": { + **context.inputs, + "result": execution_result.model_dump(), + } + }, + ) + + evaluation_result = await self._conductor.execute( + "evaluator", + evaluation_context, + ) + self._store_result(context, "evaluation", evaluation_result) + + return WorkflowResult( + plan=plan_result, + execution=execution_result, + evaluation=evaluation_result, + memory_keys=self._workflow_memory_keys(context), + ) + + def _store_result( + self, + context: AgentContext, + stage: str, + result: AgentResult, + ) -> None: + key = self._memory_key(context, stage) + + self._memory.set( + MemoryEntry( + key=key, + value=result.model_dump(), + metadata={ + "run_id": str(context.run_id), + "stage": stage, + "agent_name": result.agent_name, + }, + ) + ) + + def _workflow_memory_keys( + self, + context: AgentContext, + ) -> tuple[str, ...]: + prefix = f"workflow.{context.run_id}." + + return tuple(key for key in self._memory.keys() if key.startswith(prefix)) # noqa: SIM118 + + @staticmethod + def _memory_key( + context: AgentContext, + stage: str, + ) -> str: + return f"workflow.{context.run_id}.{stage}" diff --git a/src/aegis/evaluation/__init__.py b/src/aegis/evaluation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/aegis/memory/__init__.py b/src/aegis/memory/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/aegis/observability/__init__.py b/src/aegis/observability/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/aegis/providers/__init__.py b/src/aegis/providers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/aegis/runtime/__init__.py b/src/aegis/runtime/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7744c9362e1ef6f83818f6778555702a431e4f12 --- /dev/null +++ b/src/aegis/runtime/__init__.py @@ -0,0 +1,27 @@ +from aegis.runtime.errors import ( + AegisRunError, + EvaluationFailureError, + ExecutionFailureError, + InternalFailureError, + InvalidConfigurationError, + InvalidTaskError, + PolicyViolationError, + UnknownAgentError, + UnknownWorkflowError, +) +from aegis.runtime.runner import AegisRuntime, RunOutcome, WorkflowDefinition + +__all__ = [ + "AegisRunError", + "AegisRuntime", + "EvaluationFailureError", + "ExecutionFailureError", + "InternalFailureError", + "InvalidConfigurationError", + "InvalidTaskError", + "PolicyViolationError", + "RunOutcome", + "UnknownAgentError", + "UnknownWorkflowError", + "WorkflowDefinition", +] diff --git a/src/aegis/runtime/errors.py b/src/aegis/runtime/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..12e3cef0f00b00d5800fe52e4e01efcb44ba9165 --- /dev/null +++ b/src/aegis/runtime/errors.py @@ -0,0 +1,37 @@ +from __future__ import annotations + + +class AegisRunError(Exception): + exit_code = 1 + + +class InvalidTaskError(AegisRunError): + exit_code = 2 + + +class InvalidConfigurationError(AegisRunError): + exit_code = 3 + + +class UnknownAgentError(InvalidConfigurationError): + pass + + +class UnknownWorkflowError(InvalidConfigurationError): + pass + + +class PolicyViolationError(AegisRunError): + exit_code = 4 + + +class ExecutionFailureError(AegisRunError): + exit_code = 5 + + +class EvaluationFailureError(AegisRunError): + exit_code = 6 + + +class InternalFailureError(AegisRunError): + exit_code = 1 diff --git a/src/aegis/runtime/runner.py b/src/aegis/runtime/runner.py new file mode 100644 index 0000000000000000000000000000000000000000..45811313be7ec48965baabbe8726904615bba319 --- /dev/null +++ b/src/aegis/runtime/runner.py @@ -0,0 +1,264 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any +from uuid import UUID, uuid4 + +from pydantic import BaseModel, ConfigDict + +from aegis.agents import EvaluatorAgent, PlannerAgent +from aegis.core import ( + Agent, + AgentContext, + AgentRegistry, + AgentResult, + Conductor, + MemoryStore, + Task, + WorkflowEngine, +) +from aegis.core.models import TaskStatus +from aegis.runtime.errors import ( + EvaluationFailureError, + ExecutionFailureError, + InvalidConfigurationError, + InvalidTaskError, + PolicyViolationError, + UnknownAgentError, + UnknownWorkflowError, +) + + +class RunOutcome(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + run_id: str + status: str + agent: str | None = None + workflow: str | None = None + result: dict[str, Any] | None = None + evaluation: dict[str, Any] | None = None + error: str | None = None + + +@dataclass(frozen=True, slots=True) +class WorkflowDefinition: + name: str + target_agent: str + description: str + minimum_steps: int = 3 + + +class DefaultTaskAgent(Agent): + name = "default-task" + + async def run(self, context: AgentContext) -> AgentResult: + task_text = context.task.description.strip() + if not task_text: + return AgentResult( + agent_name=self.name, + success=False, + errors=["Task description is empty."], + ) + + max_steps = context.state.get("max_steps") + constraints: dict[str, Any] = {} + if isinstance(max_steps, int): + constraints["max_steps"] = max_steps + + return AgentResult( + agent_name=self.name, + success=True, + output={ + "task": task_text, + "summary": f"Received task: {task_text}", + "next_step": ( + "Break the task into concrete subtasks, execute them, and verify the result." + ), + "constraints": constraints, + }, + ) + + +class AegisRuntime: + def __init__( + self, + *, + registry: AgentRegistry | None = None, + memory: MemoryStore | None = None, + workflows: dict[str, WorkflowDefinition] | None = None, + conductor: Conductor | None = None, + workflow_engine: WorkflowEngine | None = None, + ) -> None: + self._registry = registry or self._build_default_registry() + self._memory = memory or MemoryStore() + self._workflows = workflows or self._build_default_workflows() + self._conductor = conductor or Conductor(self._registry) + self._workflow_engine = workflow_engine or WorkflowEngine( + conductor=self._conductor, + memory=self._memory, + ) + + def run( + self, + *, + task: str, + agent_name: str | None = None, + workflow_name: str | None = None, + max_steps: int | None = None, + run_id: UUID | None = None, + ) -> RunOutcome: + normalized_task = task.strip() + if not normalized_task: + raise InvalidTaskError("Task cannot be empty.") + + requested_workflow = self._normalize_name(workflow_name) + requested_agent = self._normalize_name(agent_name) + + if requested_agent is not None and requested_workflow is not None: + raise InvalidConfigurationError("Specify either --agent or --workflow, not both.") + + if max_steps is not None and max_steps < 3: + raise PolicyViolationError("max-steps must be at least 3 for the default workflow.") + + selected_workflow = requested_workflow or "default" + workflow = self._workflows.get(selected_workflow) + + if requested_workflow is not None and workflow is None: + raise UnknownWorkflowError(f"Unknown workflow: {requested_workflow}") + + if requested_agent is not None: + if not self._registry.contains(requested_agent): + raise UnknownAgentError(f"Unknown agent: {requested_agent}") + target_agent = requested_agent + workflow_label: str | None = None + else: + workflow = workflow or self._workflows["default"] + target_agent = workflow.target_agent + workflow_label = workflow.name + + if not self._registry.contains(target_agent): + raise UnknownAgentError(f"Unknown agent: {target_agent}") + + task_model = Task( + id=run_id or uuid4(), + name=self._task_name(normalized_task), + description=normalized_task, + status=TaskStatus.PENDING, + metadata={ + "requested_agent": requested_agent, + "requested_workflow": requested_workflow, + "max_steps": max_steps, + }, + ) + context = AgentContext( + run_id=task_model.id, + task=task_model, + inputs={ + "task": normalized_task, + "requested_agent": requested_agent, + "requested_workflow": requested_workflow, + "max_steps": max_steps, + }, + state={ + "max_steps": max_steps, + }, + ) + + try: + workflow_result = self._run_async( + self._workflow_engine.run( + context, + target_agent=target_agent, + ) + ) + except UnknownAgentError: + raise + except UnknownWorkflowError: + raise + except ExecutionFailureError: + raise + except EvaluationFailureError: + raise + except Exception as exc: # pragma: no cover - defensive boundary + raise ExecutionFailureError(str(exc)) from exc + + execution = workflow_result.execution + evaluation = workflow_result.evaluation + + if not execution.success: + message = execution.errors[0] if execution.errors else "Execution failed." + raise ExecutionFailureError(message) + + if not evaluation.success: + message = evaluation.errors[0] if evaluation.errors else "Evaluation failed." + raise EvaluationFailureError(message) + + evaluation_output = evaluation.output + passed = evaluation_output.get("passed") if isinstance(evaluation_output, dict) else None + if passed is False: + issues = ( + evaluation_output.get("issues", []) if isinstance(evaluation_output, dict) else [] + ) + if isinstance(issues, list) and issues: + message = str(issues[0]) + else: + message = "Evaluation failed." + raise EvaluationFailureError(message) + + return RunOutcome( + run_id=f"run_{task_model.id.hex}", + status="completed", + agent=target_agent, + workflow=workflow_label, + result=execution.model_dump(mode="json"), + evaluation=evaluation.model_dump(mode="json"), + error=None, + ) + + def _build_default_registry(self) -> AgentRegistry: + return AgentRegistry( + [ + PlannerAgent(), + EvaluatorAgent(), + DefaultTaskAgent(), + ] + ) + + def _build_default_workflows(self) -> dict[str, WorkflowDefinition]: + workflow = WorkflowDefinition( + name="default", + target_agent="default-task", + description="Plan, execute, and evaluate the default Aegis workflow.", + ) + review = WorkflowDefinition( + name="review", + target_agent="default-task", + description="Review-oriented alias for the default Aegis workflow.", + ) + return { + workflow.name: workflow, + review.name: review, + } + + @staticmethod + def _normalize_name(name: str | None) -> str | None: + if name is None: + return None + normalized = name.strip() + return normalized or None + + @staticmethod + def _task_name(task: str) -> str: + if len(task) <= 120: + return task + return "CLI task" + + @staticmethod + def _run_async(result: object) -> Any: + if not hasattr(result, "__await__"): + return result + + import asyncio + + return asyncio.run(result) # type: ignore[arg-type] diff --git a/src/aegis/safety/__init__.py b/src/aegis/safety/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/aegis/tools/__init__.py b/src/aegis/tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/unit/test_cli_run.py b/tests/unit/test_cli_run.py new file mode 100644 index 0000000000000000000000000000000000000000..f84c5468dc702da28f96a748e0e253d534445300 --- /dev/null +++ b/tests/unit/test_cli_run.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +import argparse +import json + +import pytest + +from aegis.cli.app import main +from aegis.cli.commands.run import run_command +from aegis.runtime import AegisRuntime, EvaluationFailureError, RunOutcome + + +def _args(**kwargs: object) -> argparse.Namespace: + defaults = { + "task": "Review this repository", + "agent": None, + "workflow": None, + "output": "text", + "max_steps": None, + "verbose": False, + } + defaults.update(kwargs) + return argparse.Namespace(**defaults) + + +class SuccessfulRuntime: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def run( + self, + *, + task: str, + agent_name: str | None = None, + workflow_name: str | None = None, + max_steps: int | None = None, + run_id=None, + ) -> RunOutcome: + self.calls.append( + { + "task": task, + "agent_name": agent_name, + "workflow_name": workflow_name, + "max_steps": max_steps, + "run_id": run_id, + } + ) + return RunOutcome( + run_id="run_123", + status="completed", + agent=agent_name or "default-task", + workflow=workflow_name or "default", + result={"summary": task}, + evaluation={"passed": True, "score": 100}, + error=None, + ) + + +class RaisingRuntime: + def __init__(self, exc: BaseException) -> None: + self._exc = exc + + def run(self, **kwargs: object) -> RunOutcome: + raise self._exc + + +@pytest.mark.parametrize( + "argv, expected", + [ + (["run", "--help"], 0), + (["run", "Review this repository"], 0), + ], +) +def test_main_supports_help_and_execution( + argv: list[str], expected: int, capsys: pytest.CaptureFixture[str] +) -> None: + code = main(argv) + captured = capsys.readouterr() + assert code == expected + if argv[-1] == "--help": + assert "usage:" in captured.out.lower() + else: + assert "Run ID:" in captured.out + + +def test_run_command_success_with_text_output(capsys: pytest.CaptureFixture[str]) -> None: + runtime = SuccessfulRuntime() + + code = run_command(_args(), runtime=runtime) + + captured = capsys.readouterr() + assert code == 0 + assert "Run ID: run_123" in captured.out + assert "Status: completed" in captured.out + assert "Result:" in captured.out + assert captured.err == "" + assert runtime.calls[0]["run_id"] is not None + assert str(runtime.calls[0]["run_id"]).startswith("run_") is False + + +def test_run_command_rejects_empty_task(capsys: pytest.CaptureFixture[str]) -> None: + code = run_command(_args(task=" "), runtime=AegisRuntime()) + + captured = capsys.readouterr() + assert code == 2 + assert "Task cannot be empty" in captured.err + assert "Traceback" not in captured.err + + +def test_run_command_unknown_agent(capsys: pytest.CaptureFixture[str]) -> None: + code = run_command(_args(agent="missing"), runtime=AegisRuntime()) + + captured = capsys.readouterr() + assert code == 3 + assert "Unknown agent: missing" in captured.err + assert "Traceback" not in captured.err + + +def test_run_command_unknown_workflow(capsys: pytest.CaptureFixture[str]) -> None: + code = run_command(_args(workflow="missing"), runtime=AegisRuntime()) + + captured = capsys.readouterr() + assert code == 3 + assert "Unknown workflow: missing" in captured.err + assert "Traceback" not in captured.err + + +def test_run_command_runtime_exception(capsys: pytest.CaptureFixture[str]) -> None: + code = run_command(_args(), runtime=RaisingRuntime(RuntimeError("boom"))) + + captured = capsys.readouterr() + assert code == 1 + assert "Unexpected internal error." in captured.err + assert "Traceback" not in captured.err + + +def test_run_command_evaluator_failure(capsys: pytest.CaptureFixture[str]) -> None: + code = run_command(_args(), runtime=RaisingRuntime(EvaluationFailureError("evaluation failed"))) + + captured = capsys.readouterr() + assert code == 6 + assert "evaluation failed" in captured.err + + +def test_run_command_json_output_success(capsys: pytest.CaptureFixture[str]) -> None: + runtime = SuccessfulRuntime() + code = run_command(_args(output="json"), runtime=runtime) + + captured = capsys.readouterr() + payload = json.loads(captured.out) + assert code == 0 + assert payload["status"] == "completed" + assert payload["agent"] == "default-task" + assert payload["workflow"] == "default" + assert payload["error"] is None + + +def test_run_command_json_output_failure(capsys: pytest.CaptureFixture[str]) -> None: + code = run_command(_args(output="json", agent="missing"), runtime=AegisRuntime()) + + captured = capsys.readouterr() + payload = json.loads(captured.out) + assert code == 3 + assert payload["status"] == "failed" + assert payload["error"] == "Unknown agent: missing" + assert payload["result"] is None + + +def test_run_command_max_steps_forwarded(capsys: pytest.CaptureFixture[str]) -> None: + runtime = SuccessfulRuntime() + code = run_command(_args(max_steps=10), runtime=runtime) + + capsys.readouterr() + assert code == 0 + assert runtime.calls[0]["max_steps"] == 10 + + +def test_run_command_verbose_failure_has_no_stack_trace_for_known_error( + capsys: pytest.CaptureFixture[str] +) -> None: + code = run_command(_args(verbose=True, agent="missing"), runtime=AegisRuntime()) + + captured = capsys.readouterr() + assert code == 3 + assert "Traceback" not in captured.err + assert "Unknown agent: missing" in captured.err diff --git a/tests/unit/test_conductor.py b/tests/unit/test_conductor.py new file mode 100644 index 0000000000000000000000000000000000000000..eb42b9254904661bed0c40b558a2aaa76d40e5ff --- /dev/null +++ b/tests/unit/test_conductor.py @@ -0,0 +1,72 @@ +import pytest + +from aegis.core import Agent, AgentContext, AgentResult, Task +from aegis.core.conductor import Conductor +from aegis.core.registry import AgentNotFoundError, AgentRegistry + + +class StubAgent(Agent): + def __init__( + self, + name: str, + *, + result_name: str | None = None, + ) -> None: + self.name = name + self._result_name = result_name or name + + async def run(self, context: AgentContext) -> AgentResult: + return AgentResult( + agent_name=self._result_name, + success=True, + output={"task": context.task.name}, + ) + + +@pytest.mark.asyncio +async def test_conductor_executes_registered_agent() -> None: + registry = AgentRegistry([StubAgent("planner")]) + conductor = Conductor(registry) + context = AgentContext( + task=Task( + name="Create plan", + description="Build a deterministic execution plan.", + ) + ) + + result = await conductor.execute("planner", context) + + assert result.success is True + assert result.output == {"task": "Create plan"} + + +@pytest.mark.asyncio +async def test_conductor_rejects_missing_agent() -> None: + conductor = Conductor(AgentRegistry()) + context = AgentContext( + task=Task( + name="Missing agent", + description="Verify missing-agent behavior.", + ) + ) + + with pytest.raises(AgentNotFoundError): + await conductor.execute("missing", context) + + +@pytest.mark.asyncio +async def test_conductor_rejects_mismatched_result_name() -> None: + registry = AgentRegistry([StubAgent("planner", result_name="research")]) + conductor = Conductor(registry) + context = AgentContext( + task=Task( + name="Validate result", + description="Verify agent identity enforcement.", + ) + ) + + with pytest.raises( + ValueError, + match="Agent result name does not match", + ): + await conductor.execute("planner", context) diff --git a/tests/unit/test_evaluator.py b/tests/unit/test_evaluator.py new file mode 100644 index 0000000000000000000000000000000000000000..3225fd573987a8929e2be02c381147cc20e7874b --- /dev/null +++ b/tests/unit/test_evaluator.py @@ -0,0 +1,86 @@ +import pytest + +from aegis.agents import EvaluatorAgent +from aegis.core import AgentContext, AgentResult, Task + + +def make_context(result: AgentResult | dict[str, object] | None) -> AgentContext: + inputs: dict[str, object] = {} + + if result is not None: + inputs["result"] = result.model_dump() if isinstance(result, AgentResult) else result + + return AgentContext( + task=Task( + name="Evaluate result", + description="Evaluate an agent execution result.", + ), + inputs=inputs, + ) + + +@pytest.mark.asyncio +async def test_evaluator_passes_valid_result() -> None: + evaluator = EvaluatorAgent() + context = make_context( + AgentResult( + agent_name="planner", + success=True, + output={"steps": ["analyze", "execute"]}, + ) + ) + + result = await evaluator.run(context) + + assert result.success is True + assert result.output["passed"] is True + assert result.output["score"] == 100 + assert result.output["issues"] == () + assert result.output["evaluated_agent"] == "planner" + + +@pytest.mark.asyncio +async def test_evaluator_fails_unsuccessful_result() -> None: + evaluator = EvaluatorAgent() + context = make_context( + AgentResult( + agent_name="tool", + success=False, + errors=["Execution failed."], + ) + ) + + result = await evaluator.run(context) + + assert result.success is True + assert result.output["passed"] is False + assert result.output["score"] == 25 + assert len(result.output["issues"]) == 3 + + +@pytest.mark.asyncio +async def test_evaluator_rejects_missing_result() -> None: + evaluator = EvaluatorAgent() + context = make_context(None) + + result = await evaluator.run(context) + + assert result.success is False + assert result.errors == ["Missing required input: result"] + + +@pytest.mark.asyncio +async def test_evaluator_rejects_invalid_result() -> None: + evaluator = EvaluatorAgent() + context = make_context( + { + "agent_name": "", + "success": True, + } + ) + + result = await evaluator.run(context) + + assert result.success is False + assert result.errors + assert result.errors[0].startswith("Invalid agent result:") diff --git a/tests/unit/test_memory.py b/tests/unit/test_memory.py new file mode 100644 index 0000000000000000000000000000000000000000..4ab4eb8de3459ea450755138bf14d20f4f96a3ce --- /dev/null +++ b/tests/unit/test_memory.py @@ -0,0 +1,116 @@ +import pytest + +from aegis.core import ( + DuplicateMemoryKeyError, + MemoryEntry, + MemoryKeyNotFoundError, + MemoryStore, +) + + +def test_memory_store_adds_and_gets_entry() -> None: + store = MemoryStore() + entry = MemoryEntry( + key="planner.output", + value={"steps": ["analyze", "execute"]}, + metadata={"agent": "planner"}, + ) + + store.add(entry) + + stored = store.get("planner.output") + + assert stored.key == "planner.output" + assert stored.value == {"steps": ["analyze", "execute"]} + assert stored.metadata == {"agent": "planner"} + assert len(store) == 1 + + +def test_memory_store_initializes_with_entries() -> None: + store = MemoryStore( + [ + MemoryEntry(key="beta", value=2), + MemoryEntry(key="alpha", value=1), + ] + ) + + assert store.keys() == ("alpha", "beta") + assert len(store) == 2 + + +def test_memory_store_rejects_duplicate_key() -> None: + store = MemoryStore([MemoryEntry(key="result", value=1)]) + + with pytest.raises( + DuplicateMemoryKeyError, + match="Memory key already exists: result", + ): + store.add(MemoryEntry(key="result", value=2)) + + +def test_memory_store_set_replaces_existing_entry() -> None: + store = MemoryStore([MemoryEntry(key="result", value=1)]) + + store.set(MemoryEntry(key="result", value=2)) + + assert store.get("result").value == 2 + assert len(store) == 1 + + +def test_memory_store_raises_for_missing_key() -> None: + store = MemoryStore() + + with pytest.raises( + MemoryKeyNotFoundError, + match="Memory key not found: missing", + ): + store.get("missing") + + +def test_memory_store_deletes_entry() -> None: + store = MemoryStore([MemoryEntry(key="result", value=1)]) + + store.delete("result") + + assert store.contains("result") is False + assert len(store) == 0 + + +def test_memory_store_rejects_delete_for_missing_key() -> None: + store = MemoryStore() + + with pytest.raises( + MemoryKeyNotFoundError, + match="Memory key not found: missing", + ): + store.delete("missing") + + +def test_memory_store_normalizes_keys() -> None: + store = MemoryStore() + + store.add(MemoryEntry(key=" planner.output ", value=1)) + + assert store.contains("planner.output") is True + assert store.get(" planner.output ").key == "planner.output" + + +def test_memory_store_rejects_empty_key_lookup() -> None: + store = MemoryStore() + + with pytest.raises(ValueError, match="Memory key must not be empty"): + store.get(" ") + + +def test_memory_store_clear_removes_all_entries() -> None: + store = MemoryStore( + [ + MemoryEntry(key="one", value=1), + MemoryEntry(key="two", value=2), + ] + ) + + store.clear() + + assert store.keys() == () + assert len(store) == 0 diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py new file mode 100644 index 0000000000000000000000000000000000000000..257e4a41896ef2073637bb848ae5a1ec75956a13 --- /dev/null +++ b/tests/unit/test_models.py @@ -0,0 +1,27 @@ +from aegis.core import AgentContext, AgentResult, Task, TaskStatus + + +def test_task_defaults_to_pending() -> None: + task = Task(name="Plan", description="Create an execution plan.") + + assert task.status is TaskStatus.PENDING + assert task.name == "Plan" + + +def test_agent_context_contains_task() -> None: + task = Task(name="Research", description="Retrieve relevant evidence.") + context = AgentContext(task=task, inputs={"query": "AI agents"}) + + assert context.task == task + assert context.inputs["query"] == "AI agents" + + +def test_agent_result_success() -> None: + result = AgentResult( + agent_name="planner", + success=True, + output={"steps": ["inspect", "plan", "execute"]}, + ) + + assert result.success is True + assert result.errors == [] diff --git a/tests/unit/test_planner.py b/tests/unit/test_planner.py new file mode 100644 index 0000000000000000000000000000000000000000..8affc597a163a8d55b2ad562672a1a0a3cfac831 --- /dev/null +++ b/tests/unit/test_planner.py @@ -0,0 +1,53 @@ +import pytest + +from aegis.agents import PlannerAgent +from aegis.core import AgentContext, Task + + +@pytest.mark.asyncio +async def test_planner_returns_structured_plan() -> None: + planner = PlannerAgent() + context = AgentContext( + task=Task( + name="Build API", + description="Create and validate a small API endpoint.", + ) + ) + + result = await planner.run(context) + + assert result.success is True + assert result.agent_name == "planner" + assert result.output["task_name"] == "Build API" + assert len(result.output["steps"]) == 3 + + +@pytest.mark.asyncio +async def test_planner_steps_have_sequential_ids() -> None: + planner = PlannerAgent() + context = AgentContext( + task=Task( + name="Validate repository", + description="Run repository validation checks.", + ) + ) + + result = await planner.run(context) + step_ids = [step["id"] for step in result.output["steps"]] + + assert step_ids == [1, 2, 3] + + +@pytest.mark.asyncio +async def test_planner_output_is_bounded() -> None: + planner = PlannerAgent() + context = AgentContext( + task=Task( + name="Small task", + description="Produce a concise execution plan.", + ) + ) + + result = await planner.run(context) + + assert len(result.output["steps"]) <= 3 diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py new file mode 100644 index 0000000000000000000000000000000000000000..dccd3e6ce96b2f86974fda26f2323515bd568a29 --- /dev/null +++ b/tests/unit/test_policies.py @@ -0,0 +1,33 @@ +import pytest +from pydantic import ValidationError + +from aegis.core.policies import AgentPolicy, ApprovalMode, ExecutionPolicy, ResponsePolicy + + +def test_response_policy_defaults() -> None: + policy = ResponsePolicy() + + assert policy.max_output_tokens == 600 + assert policy.answer_first is True + assert policy.include_reasoning is False + + +def test_execution_policy_defaults_to_restricted_permissions() -> None: + policy = ExecutionPolicy() + + assert policy.approval_mode is ApprovalMode.HIGH_RISK_ONLY + assert policy.allow_network is False + assert policy.allow_file_writes is False + assert policy.allow_shell is False + + +def test_agent_policy_combines_response_and_execution_rules() -> None: + policy = AgentPolicy() + + assert policy.response.max_input_chars == 20_000 + assert policy.execution.max_steps == 8 + + +def test_invalid_policy_limit_is_rejected() -> None: + with pytest.raises(ValidationError): + ResponsePolicy(max_output_tokens=0) diff --git a/tests/unit/test_registry.py b/tests/unit/test_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..486df510d853fb02bfc1c40eb77bfa187d828795 --- /dev/null +++ b/tests/unit/test_registry.py @@ -0,0 +1,56 @@ +import pytest + +from aegis.core import Agent, AgentContext, AgentResult +from aegis.core.registry import ( + AgentNotFoundError, + AgentRegistry, + DuplicateAgentError, +) + + +class StubAgent(Agent): + def __init__(self, name: str) -> None: + self.name = name + + async def run(self, context: AgentContext) -> AgentResult: + return AgentResult( + agent_name=self.name, + success=True, + output={"task": context.task.name}, + ) + + +def test_registry_registers_and_resolves_agent() -> None: + agent = StubAgent("planner") + registry = AgentRegistry() + + registry.register(agent) + + assert registry.get("planner") is agent + assert registry.contains("planner") is True + assert len(registry) == 1 + + +def test_registry_rejects_duplicate_agent_names() -> None: + registry = AgentRegistry([StubAgent("planner")]) + + with pytest.raises(DuplicateAgentError): + registry.register(StubAgent("planner")) + + +def test_registry_raises_when_agent_is_missing() -> None: + registry = AgentRegistry() + + with pytest.raises(AgentNotFoundError): + registry.get("missing") + + +def test_registry_returns_sorted_names() -> None: + registry = AgentRegistry( + [ + StubAgent("research"), + StubAgent("planner"), + ] + ) + + assert registry.names() == ("planner", "research") diff --git a/tests/unit/test_tool_registry.py b/tests/unit/test_tool_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..d86d373f43c20d32a43308d873adf0447d9efc61 --- /dev/null +++ b/tests/unit/test_tool_registry.py @@ -0,0 +1,74 @@ +import pytest + +from aegis.core.tool_registry import ( + DuplicateToolError, + ToolNotFoundError, + ToolRegistry, +) + + +class StubTool: + def __init__(self, name: str) -> None: + self.name = name + + +def test_registry_registers_and_resolves_tool() -> None: + tool = StubTool("calculator") + registry = ToolRegistry() + + registry.register(tool) + + assert registry.get("calculator") is tool + assert registry.contains("calculator") is True + assert len(registry) == 1 + + +def test_registry_accepts_initial_tools() -> None: + calculator = StubTool("calculator") + search = StubTool("search") + + registry = ToolRegistry([calculator, search]) + + assert registry.get("calculator") is calculator + assert registry.get("search") is search + assert len(registry) == 2 + + +def test_registry_rejects_duplicate_tool_names() -> None: + registry = ToolRegistry([StubTool("calculator")]) + + with pytest.raises(DuplicateToolError): + registry.register(StubTool("calculator")) + + +def test_registry_raises_when_tool_is_missing() -> None: + registry = ToolRegistry() + + with pytest.raises(ToolNotFoundError): + registry.get("missing") + + +def test_registry_returns_sorted_names() -> None: + registry = ToolRegistry( + [ + StubTool("search"), + StubTool("calculator"), + ] + ) + + assert registry.names() == ("calculator", "search") + + +def test_registry_normalizes_tool_names() -> None: + tool = StubTool(" calculator ") + registry = ToolRegistry([tool]) + + assert registry.contains("calculator") is True + assert registry.get(" calculator ") is tool + + +def test_registry_rejects_empty_tool_name() -> None: + registry = ToolRegistry() + + with pytest.raises(ValueError, match="Tool name cannot be empty"): + registry.register(StubTool(" ")) diff --git a/tests/unit/test_tool_runner.py b/tests/unit/test_tool_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..a19f8f0d37e4ed5e0e3fcf4e711281f63986703b --- /dev/null +++ b/tests/unit/test_tool_runner.py @@ -0,0 +1,45 @@ +import pytest + +from aegis.core.tool_registry import ToolRegistry +from aegis.core.tool_runner import ToolRunner + + +class EchoTool: + name = "echo" + + async def execute(self, **kwargs: object) -> object: + return kwargs + + +@pytest.mark.asyncio +async def test_tool_runner_executes_registered_tool() -> None: + registry = ToolRegistry([EchoTool()]) + runner = ToolRunner(registry) + + result = await runner.execute( + "echo", + message="hello", + value=42, + ) + + assert result == { + "message": "hello", + "value": 42, + } + + +@pytest.mark.asyncio +async def test_tool_runner_preserves_arguments() -> None: + registry = ToolRegistry([EchoTool()]) + runner = ToolRunner(registry) + + result = await runner.execute( + "echo", + alpha=1, + beta=2, + ) + + assert result == { + "alpha": 1, + "beta": 2, + } diff --git a/tests/unit/test_workflow.py b/tests/unit/test_workflow.py new file mode 100644 index 0000000000000000000000000000000000000000..98164156f1f6293c816fff2fc83f732cde3f7ad5 --- /dev/null +++ b/tests/unit/test_workflow.py @@ -0,0 +1,179 @@ +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