cli.py

#!/usr/bin/env python3
"""Command-line interface for the multi-agent orchestrator.

Provides two modes:
1. **Goal execution**: Decomposes a goal into subtasks, builds a DAG,
   runs all agents, and prints results with trajectory analysis.
2. **Trajectory analysis**: Reads a trajectory JSON file and prints
   a comprehensive analysis report.

Usage::

    # Execute a goal
    python cli.py --goal "Build a REST API for user management"

    # Execute with verbose output
    python cli.py --goal "Optimize database queries" --verbose

    # Analyze an existing trajectory file
    python cli.py --analyze trajectory.json

    # Combine both
    python cli.py --goal "Build a CLI tool" --verbose --analyze saved.json
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import sys
from pathlib import Path
from typing import Any

# ---------------------------------------------------------------------------
# Logging setup
# ---------------------------------------------------------------------------

logger = logging.getLogger("orchestrator.cli")


def _configure_logging(verbose: bool) -> None:
    """Set up logging with appropriate verbosity level."""
    level = logging.DEBUG if verbose else logging.INFO
    fmt = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
    logging.basicConfig(level=level, format=fmt, stream=sys.stderr)


# ---------------------------------------------------------------------------
# Goal execution
# ---------------------------------------------------------------------------

async def _run_goal(goal: str, verbose: bool) -> None:
    """Decompose a goal, build a DAG, run orchestration, and print results.

    Args:
        goal: The high-level goal string to accomplish.
        verbose: Whether to print detailed output.
    """
    from agents.base import AgentSpec
    from agents.code import CodeAgent
    from agents.critic import CriticAgent
    from agents.planner import PlannerAgent
    from agents.research import ResearchAgent
    from agents.synthesis import SynthesisAgent
    from core.dag import TaskDAG
    from core.orchestrator import (
        AgentRegistry,
        Orchestrator,
        OrchestratorConfig,
        RetryPolicy,
    )
    from core.task import Task
    from core.trajectory import TrajectoryStore
    from eval.trajectory_analyzer import TrajectoryAnalyzer

    store = TrajectoryStore()

    # -- Step 1: Plan -------------------------------------------------------
    print(f"[1/4] Planning: decomposing goal...")
    if verbose:
        print(f"       Goal: {goal}")

    planner_spec = AgentSpec(
        name="planner",
        model="gpt-4",
        system_prompt="You are a planning agent. Decompose goals into subtasks.",
    )
    planner = PlannerAgent(spec=planner_spec, trajectory_store=store)

    plan_task = Task(
        id="root-plan",
        name="goal_decomposition",
        agent_name="planner",
        input_payload={"goal": goal},
    )

    try:
        dag = await planner.run(plan_task)
    except Exception as exc:
        print(f"Error during planning: {exc}", file=sys.stderr)
        sys.exit(1)

    print(f"       Created DAG with {len(dag)} subtasks")
    if verbose:
        for tid, task in dag.tasks.items():
            deps_str = ", ".join(task.dependencies) if task.dependencies else "none"
            print(f"       - {task.name} (agent={task.agent_name}, deps=[{deps_str}])")

    # -- Step 2: Register agents --------------------------------------------
    print("[2/4] Registering agents...")
    registry = AgentRegistry()

    def _make_research() -> ResearchAgent:
        spec = AgentSpec(name="research", model="gpt-4")
        return ResearchAgent(spec=spec, trajectory_store=store)

    def _make_code() -> CodeAgent:
        spec = AgentSpec(name="code", model="gpt-4")
        return CodeAgent(spec=spec, trajectory_store=store)

    def _make_synthesis() -> SynthesisAgent:
        spec = AgentSpec(name="synthesis", model="gpt-4")
        return SynthesisAgent(spec=spec, trajectory_store=store)

    def _make_critic() -> CriticAgent:
        spec = AgentSpec(name="critic", model="gpt-4")
        return CriticAgent(spec=spec, trajectory_store=store)

    registry.register("research", _make_research)
    registry.register("code", _make_code)
    registry.register("synthesis", _make_synthesis)
    registry.register("critic", _make_critic)

    # -- Step 3: Run orchestration ------------------------------------------
    print("[3/4] Executing DAG...")
    config = OrchestratorConfig(
        retry_policy=RetryPolicy(
            max_retries=2,
            base_delay=0.5,
            max_delay=10.0,
            exponential_base=2.0,
        ),
        enable_critic=False,
        max_concurrent_tasks=5,
        timeout_per_task=60.0,
    )

    orchestrator = Orchestrator(
        config=config,
        agent_registry=registry,
        trajectory_store=store,
    )

    try:
        results = await orchestrator.run(dag)
    except RuntimeError as exc:
        print(f"Error during orchestration: {exc}", file=sys.stderr)
        sys.exit(1)

    print(f"       Completed {len(results)} tasks")

    # -- Step 4: Print results and analysis ---------------------------------
    print("[4/4] Results and analysis:")
    print()

    if verbose:
        print("=== Task Results ===")
        for task_id, result in results.items():
            task = dag.get_task(task_id)
            print(f"\n--- {task.name} ({task_id}) ---")
            if isinstance(result, dict):
                for key, value in result.items():
                    val_str = str(value)[:200]
                    print(f"  {key}: {val_str}")
            else:
                print(f"  {str(result)[:500]}")
        print()

    # Trajectory analysis
    analyzer = TrajectoryAnalyzer(store)
    report = analyzer.to_report(dag=dag)
    _print_report(report, verbose)

    # Save trajectory to file
    traj_path = Path("trajectory_output.json")
    traj_path.write_text(store.to_json())
    print(f"\nTrajectory saved to: {traj_path}")


# ---------------------------------------------------------------------------
# Trajectory analysis from file
# ---------------------------------------------------------------------------

def _analyze_file(filepath: str, verbose: bool) -> None:
    """Load a trajectory JSON file and print analysis.

    Args:
        filepath: Path to the trajectory JSON file.
        verbose: Whether to print detailed output.
    """
    from core.trajectory import TrajectoryStore
    from eval.trajectory_analyzer import TrajectoryAnalyzer

    path = Path(filepath)
    if not path.exists():
        print(f"Error: File not found: {filepath}", file=sys.stderr)
        sys.exit(1)

    try:
        raw = path.read_text()
        store = TrajectoryStore.from_json(raw)
    except (json.JSONDecodeError, KeyError) as exc:
        print(f"Error: Could not parse trajectory file: {exc}", file=sys.stderr)
        sys.exit(1)

    print(f"Loaded trajectory with {len(store)} task trajectories from: {filepath}")
    print()

    analyzer = TrajectoryAnalyzer(store)
    report = analyzer.to_report()
    _print_report(report, verbose)


# ---------------------------------------------------------------------------
# Report formatting
# ---------------------------------------------------------------------------

def _print_report(report: dict[str, Any], verbose: bool) -> None:
    """Pretty-print a trajectory analysis report.

    Args:
        report: The report dict from TrajectoryAnalyzer.to_report().
        verbose: Whether to include per-agent details.
    """
    print("=== Trajectory Analysis ===")
    print(f"  Total tasks:        {report['total_tasks']}")
    print(f"  Total events:       {report['total_events']}")
    print(f"  Total tokens:       {report['total_tokens']:,}")
    print(f"  Retry rate:         {report['retry_rate']:.2%}")
    print(f"  Critic pass rate:   {report['critic_pass_rate']:.2%}")

    if report["critical_path_duration_ms"] is not None:
        print(f"  Critical path:      {report['critical_path_duration_ms']:.1f}ms")

    latencies = report.get("avg_latency_per_agent", {})
    if latencies:
        print("\n  Avg latency by agent:")
        for agent, lat in sorted(latencies.items()):
            print(f"    {agent:20s} {lat:.1f}ms")

    if verbose:
        perf = report.get("agent_performance", {})
        if perf:
            print("\n  Per-agent performance:")
            for agent, stats in sorted(perf.items()):
                print(f"\n    [{agent}]")
                print(f"      Total calls:      {stats['total_calls']}")
                print(f"      Avg latency:      {stats['avg_latency']:.1f}ms")
                print(f"      Errors:           {stats['error_count']}")
                print(f"      Retries:          {stats['retry_count']}")
                if stats['avg_critic_score'] is not None:
                    print(f"      Avg critic score: {stats['avg_critic_score']:.2f}")


# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------

def _build_parser() -> argparse.ArgumentParser:
    """Build the CLI argument parser."""
    parser = argparse.ArgumentParser(
        prog="orchestrator",
        description="Multi-agent orchestrator CLI — decompose goals, run agents, analyze results.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Examples:\n"
            "  python cli.py --goal 'Build a REST API'\n"
            "  python cli.py --goal 'Optimize queries' --verbose\n"
            "  python cli.py --analyze trajectory.json\n"
            "  python cli.py --analyze trajectory.json --verbose\n"
        ),
    )

    parser.add_argument(
        "--goal",
        type=str,
        help="High-level goal string to decompose and execute.",
    )
    parser.add_argument(
        "--analyze",
        type=str,
        metavar="FILE",
        help="Path to a trajectory JSON file to analyze.",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        default=False,
        help="Enable detailed output including per-agent stats and task results.",
    )

    return parser


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    """Entry point for the CLI."""
    parser = _build_parser()
    args = parser.parse_args()

    if not args.goal and not args.analyze:
        parser.print_help()
        print("\nError: At least one of --goal or --analyze is required.", file=sys.stderr)
        sys.exit(1)

    _configure_logging(args.verbose)

    if args.analyze:
        _analyze_file(args.analyze, args.verbose)

    if args.goal:
        try:
            asyncio.run(_run_goal(args.goal, args.verbose))
        except KeyboardInterrupt:
            print("\nInterrupted.", file=sys.stderr)
            sys.exit(130)


if __name__ == "__main__":
    main()