Initial commit

src/yc_bench/cli/__init__.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import json
+import sys
+from contextlib import contextmanager
+from decimal import Decimal
+from uuid import UUID
+
+import typer
+
+from ..db.session import build_engine, build_session_factory, session_scope
+
+app = typer.Typer(name="yc-bench", add_completion=False)
+
+
+# ---------------------------------------------------------------------------
+# Helpers shared across command modules
+# ---------------------------------------------------------------------------
+
+@contextmanager
+def get_db():
+    """Yield a transactional SQLAlchemy session, commit on success."""
+    engine = build_engine()
+    factory = build_session_factory(engine)
+    with session_scope(factory) as session:
+        yield session
+
+
+class _JSONEncoder(json.JSONEncoder):
+    """Handle UUID, Decimal, datetime serialisation."""
+
+    def default(self, o):
+        if isinstance(o, UUID):
+            return str(o)
+        if isinstance(o, Decimal):
+            return float(o)
+        from datetime import datetime, date
+        if isinstance(o, datetime):
+            return o.isoformat()
+        if isinstance(o, date):
+            return o.isoformat()
+        return super().default(o)
+
+
+def json_output(data: dict | list) -> None:
+    """Print JSON to stdout (captured by run_command executor)."""
+    typer.echo(json.dumps(data, cls=_JSONEncoder, indent=2))
+
+
+def error_output(message: str, code: int = 1) -> None:
+    """Print JSON error and exit."""
+    typer.echo(json.dumps({"error": message}), err=False)
+    raise typer.Exit(code=code)
+
+
+# ---------------------------------------------------------------------------
+# Register sub-command groups
+# ---------------------------------------------------------------------------
+
+from .sim_commands import sim_app        # noqa: E402
+from .company_commands import company_app  # noqa: E402
+from .market_commands import market_app    # noqa: E402
+from .task_commands import task_app        # noqa: E402
+from .finance_commands import finance_app  # noqa: E402
+from .report_commands import report_app    # noqa: E402
+from .employee_commands import employee_app  # noqa: E402
+from .scratchpad_commands import scratchpad_app  # noqa: E402
+
+app.add_typer(sim_app, name="sim")
+app.add_typer(company_app, name="company")
+app.add_typer(employee_app, name="employee")
+app.add_typer(market_app, name="market")
+app.add_typer(task_app, name="task")
+app.add_typer(finance_app, name="finance")
+app.add_typer(report_app, name="report")
+app.add_typer(scratchpad_app, name="scratchpad")
+
+
+@app.command("run")
+def run_command_cli(
+    model: str = typer.Option(..., help="LiteLLM model string (e.g. openrouter/z-ai/glm-5)"),
+    seed: int = typer.Option(..., help="Random seed for deterministic world generation"),
+    horizon_years: int = typer.Option(3, help="Simulation horizon in years"),
+    company_name: str = typer.Option("BenchCo", help="Name of the simulated company"),
+    start_date: str = typer.Option("2025-01-01", help="Simulation start date (YYYY-MM-DD)"),
+    config_name: str = typer.Option(
+        "default", "--config",
+        help="Preset name ('default', 'fast_test', 'high_reward') or path to a .toml file",
+    ),
+):
+    """Run a full benchmark: migrate DB, seed world, run agent loop to completion."""
+    from dotenv import find_dotenv, load_dotenv
+    load_dotenv(find_dotenv(usecwd=True), override=False)
+
+    from ..runner.main import run_benchmark
+    from ..runner.args import RunArgs
+    args = RunArgs(
+        model=model,
+        seed=seed,
+        horizon_years=horizon_years,
+        company_name=company_name,
+        start_date=start_date,
+        config_name=config_name,
+    )
+    raise SystemExit(run_benchmark(args))
+
+
+def app_main():
+    """Entry point for `yc-bench` console_script."""
+    app()