app.py

#!/usr/bin/env python3
"""
ROMA Example App
================
Demonstrates the Recursive Open Meta-Agent framework via four progressively
richer patterns:

  1. one_shot    — high-level solve() shortcut
  2. pipeline    — manual Atomizer → Planner → Executor → Aggregator → Verifier
  3. async_demo  — async event-driven solving with live status output
  4. custom      — bring your own task from the command line

Usage:
    python app.py                        # interactive menu
    python app.py --demo one_shot
    python app.py --demo pipeline
    python app.py --demo async_demo
    python app.py --demo custom --task "Summarise the history of Rome in 3 bullets"
"""

import argparse
import asyncio
import os
import sys
import textwrap
from pathlib import Path

# ---------------------------------------------------------------------------
# Env loading (optional: python-dotenv)
# ---------------------------------------------------------------------------
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

# ---------------------------------------------------------------------------
# Validate that at least one LLM key is present before importing ROMA
# ---------------------------------------------------------------------------
LLM_KEYS = [
    "OPENROUTER_API_KEY",
    "OPENAI_API_KEY",
    "ANTHROPIC_API_KEY",
    "GOOGLE_API_KEY",
    "FIREWORKS_API_KEY",
    "XAI_API_KEY",
]

if not any(len(os.getenv(k, "")) >= 20 for k in LLM_KEYS):
    sys.exit(
        "\n[ERROR] No LLM API key found.\n"
        "Set at least one of: " + ", ".join(LLM_KEYS) + "\n"
        "Copy .env.example → .env and fill in your key.\n"
    )

# ---------------------------------------------------------------------------
# ROMA imports
# ---------------------------------------------------------------------------
import dspy
from roma_dspy import (
    Atomizer,
    Planner,
    Executor,
    Aggregator,
    Verifier,
    RecursiveSolver,
    SubTask,
    solve,
    async_event_solve,
)
from roma_dspy.config.schemas.agents import AgentConfig, LLMConfig
from roma_dspy.config.schemas.root import AgentsConfig, ROMAConfig, RuntimeConfig

# ---------------------------------------------------------------------------
# Pretty-print helpers
# ---------------------------------------------------------------------------
WIDTH = 72

def banner(title: str) -> None:
    print("\n" + "=" * WIDTH)
    print(f"  {title}")
    print("=" * WIDTH)

def section(title: str) -> None:
    print(f"\n{'─' * WIDTH}")
    print(f"  {title}")
    print(f"{'─' * WIDTH}")

def wrap(text: str, indent: int = 2) -> str:
    prefix = " " * indent
    return textwrap.fill(
        str(text), width=WIDTH - indent,
        initial_indent=prefix, subsequent_indent=prefix,
    )

def show_result(label: str, value) -> None:
    section(label)
    print(wrap(str(value)))

# ---------------------------------------------------------------------------
# Model / config helpers
# ---------------------------------------------------------------------------

# Provider preference order: OpenRouter → OpenAI → Anthropic → xAI → Google → Fireworks
def _key_is_real(env_var: str, min_len: int = 20) -> bool:
    """Return True only if the key looks like a real credential (not a placeholder)."""
    val = os.getenv(env_var, "")
    return len(val) >= min_len


def _pick_real_model(candidates: dict[str, str]) -> str:
    """Return first model whose provider key passes the reality check."""
    for model, env_key in candidates.items():
        if _key_is_real(env_key):
            return model
    raise RuntimeError(
        "No valid LLM key found. Make sure your .env key is the full secret, "
        "not a placeholder like 'sk-...'."
    )


FAST_MODEL = _pick_real_model({
    "openrouter/google/gemini-2.5-flash":             "OPENROUTER_API_KEY",
    "openai/gpt-4o-mini":                             "OPENAI_API_KEY",
    "anthropic/claude-haiku-4-5-20251001":            "ANTHROPIC_API_KEY",
    "xai/grok-3-mini-latest":                         "XAI_API_KEY",
    "google/gemini-2.0-flash":                        "GOOGLE_API_KEY",
    "fireworks_ai/accounts/fireworks/models/kimi-k2": "FIREWORKS_API_KEY",
})
REASON_MODEL = _pick_real_model({
    "openrouter/openai/gpt-4o-mini":                  "OPENROUTER_API_KEY",
    "openai/gpt-4o-mini":                             "OPENAI_API_KEY",
    "anthropic/claude-sonnet-4-5":                    "ANTHROPIC_API_KEY",
    "xai/grok-3-latest":                              "XAI_API_KEY",
    "google/gemini-2.0-flash":                        "GOOGLE_API_KEY",
    "fireworks_ai/accounts/fireworks/models/kimi-k2": "FIREWORKS_API_KEY",
})
EXEC_MODEL = _pick_real_model({
    "openrouter/anthropic/claude-sonnet-4-5":         "OPENROUTER_API_KEY",
    "anthropic/claude-sonnet-4-5":                    "ANTHROPIC_API_KEY",
    "openai/gpt-4o-mini":                             "OPENAI_API_KEY",
    "xai/grok-3-latest":                              "XAI_API_KEY",
    "google/gemini-2.0-flash":                        "GOOGLE_API_KEY",
    "fireworks_ai/accounts/fireworks/models/kimi-k2": "FIREWORKS_API_KEY",
})


def make_config(
    fast: str = FAST_MODEL,
    reason: str = REASON_MODEL,
    executor: str = EXEC_MODEL,
    max_depth: int = 3,
) -> ROMAConfig:
    """
    Build a ROMAConfig that assigns specific models to each pipeline stage.
    Runtime timeout is set to 600 s so it doesn't conflict with the default
    LLM timeout (also 600 s).
    """
    agents = AgentsConfig(
        atomizer=AgentConfig(
            llm=LLMConfig(model=fast,   temperature=0.1, max_tokens=1_000),
        ),
        planner=AgentConfig(
            llm=LLMConfig(model=reason, temperature=0.3, max_tokens=3_000),
        ),
        executor=AgentConfig(
            llm=LLMConfig(model=executor, temperature=0.5, max_tokens=4_000),
        ),
        aggregator=AgentConfig(
            llm=LLMConfig(model=reason, temperature=0.2, max_tokens=4_000),
        ),
        verifier=AgentConfig(
            llm=LLMConfig(model=fast,   temperature=0.0, max_tokens=500),
        ),
    )
    return ROMAConfig(
        agents=agents,
        runtime=RuntimeConfig(max_depth=max_depth, timeout=600, verbose=False),
    )


# ---------------------------------------------------------------------------
# Demo 1 — High-level one-shot solve
# ---------------------------------------------------------------------------

def demo_one_shot() -> None:
    """
    The simplest possible ROMA usage: pass a task string to solve().
    ROMA internally runs Atomizer → Planner → Executor → Aggregator → Verifier
    and returns a completed TaskNode.
    """
    banner("Demo 1 · One-Shot Solve")
    print(wrap(
        "Uses the top-level solve() helper. ROMA orchestrates everything "
        "internally — it checks whether the task needs decomposition, plans "
        "subtasks if required, executes them, aggregates, and verifies."
    ))

    task = (
        "List 5 key milestones in the history of artificial intelligence, "
        "one sentence each."
    )

    section("Task")
    print(wrap(task))
    print()

    config = make_config(max_depth=2)
    print(f"  fast model    : {FAST_MODEL}")
    print(f"  executor model: {EXEC_MODEL}")

    print("\n  Solving…", flush=True)
    node = solve(task, config=config)

    show_result(f"Status: {node.status}", node.result or "(no result)")


# ---------------------------------------------------------------------------
# Demo 2 — Manual pipeline (Atomizer → Planner → Executor → Aggregator → Verifier)
# ---------------------------------------------------------------------------

def demo_pipeline() -> None:
    """
    Walks through each ROMA module individually so you can inspect and
    customise every step.
    """
    banner("Demo 2 · Manual Pipeline")
    print(wrap(
        "Instantiates each module separately with a specific model. This gives "
        "fine-grained control — different models per stage, custom tools on "
        "the Executor, or your own aggregation logic."
    ))

    task = (
        "Compare Python and Rust for systems programming: "
        "cover performance, safety, ecosystem, and learning curve."
    )

    section("Task")
    print(wrap(task))

    # --- Step 1: Atomizer ---------------------------------------------------
    section("Step 1 · Atomizer")
    print(wrap(
        "Decides whether the task is atomic (directly executable) or needs "
        "planning. Returns is_atomic (bool) and node_type."
    ))

    atomizer = Atomizer(model=FAST_MODEL)
    atomized = atomizer.forward(task)

    print(f"\n  is_atomic : {atomized.is_atomic}")
    print(f"  node_type : {atomized.node_type}")

    if atomized.is_atomic:
        print(wrap("\nTask is atomic — skipping Planner, going straight to Executor."))
        executor = Executor(model=EXEC_MODEL)
        res = executor.forward(task)
        show_result("Executor Result", res.output)
        return

    # --- Step 2: Planner ----------------------------------------------------
    section("Step 2 · Planner")
    print(wrap(
        "Decomposes the task into ordered SubTask objects, each with a typed "
        "goal (RETRIEVE / WRITE / THINK / CODE_INTERPRET / IMAGE_GENERATION)."
    ))

    planner = Planner(model=REASON_MODEL)
    plan = planner.forward(task)

    print(f"\n  Subtasks ({len(plan.subtasks)}):")
    for i, sub in enumerate(plan.subtasks, 1):
        print(f"\n  [{i}] {sub.goal}")
        print(f"       type        : {sub.task_type}")
        if sub.dependencies:
            print(f"       depends_on  : {sub.dependencies}")

    # --- Step 3: Executor ---------------------------------------------------
    section("Step 3 · Executor")
    print(wrap(
        "Resolves each subtask independently. Results are stored back on the "
        "SubTask objects so the Aggregator can reference them."
    ))

    executor = Executor(model=EXEC_MODEL)
    subtasks_with_results: list[SubTask] = []

    for i, sub in enumerate(plan.subtasks, 1):
        print(f"\n  Executing [{i}/{len(plan.subtasks)}]: {sub.goal[:60]}…")
        res = executor.forward(sub.goal)
        # Store result on a new SubTask (result field is Optional[str])
        filled = sub.model_copy(update={"result": res.output})
        subtasks_with_results.append(filled)
        preview = str(res.output)[:120]
        print(wrap(f"→ {preview}{'…' if len(str(res.output)) > 120 else ''}", indent=4))

    # --- Step 4: Aggregator -------------------------------------------------
    section("Step 4 · Aggregator")
    print(wrap(
        "Synthesises all SubTask results into a single, coherent final answer."
    ))

    aggregator = Aggregator(model=REASON_MODEL)
    aggregated = aggregator.forward(task, subtasks_with_results)

    show_result("Aggregated Answer", aggregated.synthesized_result)

    # --- Step 5: Verifier ---------------------------------------------------
    section("Step 5 · Verifier")
    print(wrap(
        "Validates the aggregated answer against the original goal. Returns a "
        "boolean verdict and optional corrective feedback."
    ))

    verifier = Verifier(model=FAST_MODEL)
    verdict = verifier.forward(task, aggregated.synthesized_result)

    print(f"\n  verdict  : {verdict.verdict}")
    if verdict.feedback:
        print(wrap(f"feedback : {verdict.feedback}"))


# ---------------------------------------------------------------------------
# Demo 3 — Async event-driven solving
# ---------------------------------------------------------------------------

async def _async_demo_inner() -> None:
    """
    async_event_solve() uses an event-loop controller that runs independent
    subtasks concurrently.  We poll the DAG while waiting to stream progress.
    """
    banner("Demo 3 · Async Event-Driven Solve")
    print(wrap(
        "async_event_solve() uses an internal event-loop scheduler that runs "
        "independent subtasks concurrently. Here we show the final TaskNode "
        "after completion, with its child nodes listed."
    ))

    task = (
        "Create a concise study plan for learning machine learning from scratch: "
        "include foundational math, programming prerequisites, key algorithms, "
        "and recommended resources."
    )

    section("Task")
    print(wrap(task))
    print()

    config = make_config(max_depth=2)
    print(f"  model : {FAST_MODEL}")
    print("\n  Solving asynchronously…", flush=True)

    node = await async_event_solve(task, config=config, concurrency=3)

    section("Result")
    print(f"  status     : {node.status}")
    print(f"  task_id    : {node.task_id}")
    print(f"  children   : {len(node.children)}")

    if node.children:
        print("\n  Child nodes:")
        for child_id in node.children:
            print(f"    • {child_id}")

    show_result("Final Answer", node.result or "(no result)")


def demo_async() -> None:
    asyncio.run(_async_demo_inner())


# ---------------------------------------------------------------------------
# Demo 4 — Custom task via CLI
# ---------------------------------------------------------------------------

def demo_custom(task: str) -> None:
    """Run any task you provide via --task, using RecursiveSolver directly."""
    banner("Demo 4 · Custom Task")
    section("Task")
    print(wrap(task))
    print()

    config = make_config(max_depth=3)
    solver = RecursiveSolver(config=config)

    print("  Solving with RecursiveSolver…\n", flush=True)
    node = solver.solve(task=task)

    show_result(f"Status: {node.status}", node.result or "(no result)")

    # Show execution DAG if available
    dag = getattr(solver, "last_dag", None)
    if dag:
        section("Execution DAG")
        try:
            from roma_dspy.core.utils.trace_formatter import format_dag_summary
            print(format_dag_summary(dag))
        except Exception:
            # Graceful fallback
            nodes = [dag.get_node(nid) for nid in (dag.get_all_node_ids() or [])]
            for n in nodes:
                if n:
                    print(f"  [{n.status:10s}] {(n.goal or '')[:65]}")

        try:
            exec_id = dag.get_execution_id()
            print(f"\n  Execution ID: {exec_id}")
        except Exception:
            pass


# ---------------------------------------------------------------------------
# Interactive menu
# ---------------------------------------------------------------------------

DEMOS: dict[str, tuple[str, str, object]] = {
    "1": ("one_shot",   "High-level one-shot solve()",              demo_one_shot),
    "2": ("pipeline",   "Manual 5-stage pipeline (step-by-step)",   demo_pipeline),
    "3": ("async_demo", "Async event-driven solving",               demo_async),
}

def interactive_menu() -> None:
    banner("ROMA — Recursive Open Meta-Agent · Example App")
    print(wrap(
        "ROMA decomposes complex goals into a DAG of subtasks and solves them "
        "recursively using a configurable pipeline of LLM-powered modules: "
        "Atomizer → Planner → Executor → Aggregator → Verifier."
    ))
    print()
    print("  Available demos:")
    for key, (name, desc, _) in DEMOS.items():
        print(f"    [{key}] {desc}")
    print("    [q] Quit")
    print()
    print(f"  Using models:")
    print(f"    fast     : {FAST_MODEL}")
    print(f"    reasoning: {REASON_MODEL}")
    print(f"    executor : {EXEC_MODEL}")
    print()

    while True:
        choice = input("  Choose a demo [1-3 / q]: ").strip().lower()
        if choice == "q":
            print("  Bye!\n")
            sys.exit(0)
        if choice in DEMOS:
            _, _, fn = DEMOS[choice]
            fn()
            print()
            cont = input("  Run another demo? [y/n]: ").strip().lower()
            if cont != "y":
                print("  Bye!\n")
                sys.exit(0)
        else:
            print("  Invalid choice. Enter 1, 2, 3, or q.")


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="ROMA example app — demonstrates recursive multi-agent solving.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        "--demo",
        choices=["one_shot", "pipeline", "async_demo", "custom"],
        help="Run a specific demo directly.",
    )
    parser.add_argument(
        "--task",
        type=str,
        default=None,
        help="Custom task string (used with --demo custom).",
    )
    args = parser.parse_args()

    if args.demo is None:
        interactive_menu()
    elif args.demo == "one_shot":
        demo_one_shot()
    elif args.demo == "pipeline":
        demo_pipeline()
    elif args.demo == "async_demo":
        demo_async()
    elif args.demo == "custom":
        if not args.task:
            parser.error("--demo custom requires --task <your task>")
        demo_custom(args.task)


if __name__ == "__main__":
    main()