orchestrator/run.py

#!/usr/bin/env python3
"""CLI entry point for the orchestration graph.

Usage:
    python3 run.py "Write a function that validates emails"
    python3 run.py --route-only "Send a Slack message to #general"
"""

import os
import sys
import time

from dotenv import load_dotenv

os.chdir(os.path.dirname(os.path.abspath(__file__)))
load_dotenv(".env")

from graph import get_app, retriever_node, router_node  # noqa: E402  (env must be loaded before graph imports composio)
from telemetry import build_record, log_run  # noqa: E402


def _format_retrieved(retrieved: list[dict] | None) -> str:
    if not retrieved:
        return "[retrieved: none]"
    names = ", ".join(m["name"] for m in retrieved)
    return f"[retrieved: {names}]"


def main():
    if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help"):
        print("Usage: python3 run.py [--route-only] <task description>")
        sys.exit(1)

    route_only = "--route-only" in sys.argv
    args = [a for a in sys.argv[1:] if a != "--route-only"]
    task = " ".join(args)

    if route_only:
        retrieval = retriever_node({"task": task, "messages": []})
        retrieved = retrieval.get("retrieved", [])
        out = router_node({"task": task, "messages": [], "retrieved": retrieved})
        print(_format_retrieved(retrieved))
        print(out["route"])
        return

    started = time.monotonic()
    result: dict | None = None
    error: str | None = None
    try:
        result = get_app().invoke({"task": task, "messages": []})
    except Exception as exc:
        error = f"{type(exc).__name__}: {exc}"
    finished = time.monotonic()

    log_run(
        build_record(
            task=task,
            started=started,
            finished=finished,
            result=result,
            success=error is None,
            error=error,
        )
    )

    if error is not None:
        print(f"[error] {error}", file=sys.stderr)
        sys.exit(1)

    print(_format_retrieved(result.get("retrieved")))
    print(f"[{result['route']}]")
    print()
    output = result.get("result", "")
    if isinstance(output, list):
        output = "\n".join(str(c) for c in output)
    print(output)


if __name__ == "__main__":
    main()