tenant.openagents/omega

1from __future__ import annotations
2
3import argparse
4import datetime
5import json
6import os
7import pathlib
8import shutil
9import subprocess
10import tarfile
11from typing import Any
12
13from . import config, source
14
15ALL_SWE_ATLAS_PARTS = ["qna", "rf", "tw"]
16
17
18def utc_timestamp() -> str:
19    return datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%dT%H%M%SZ")
20
21
22def utc_now() -> str:
23    return datetime.datetime.now(datetime.timezone.utc).isoformat()
24
25
26def print_json(data: Any) -> None:
27    print(json.dumps(data, indent=2, sort_keys=True))
28
29
30def write_json(path: pathlib.Path, data: dict[str, Any]) -> None:
31    path.parent.mkdir(parents=True, exist_ok=True)
32    path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n")
33
34
35def load_json(path: pathlib.Path) -> dict[str, Any] | None:
36    try:
37        data = json.loads(path.read_text())
38    except (OSError, json.JSONDecodeError):
39        return None
40    return data if isinstance(data, dict) else None
41
42
43def safe_extract_archive(archive: tarfile.TarFile, destination: pathlib.Path) -> None:
44    destination = destination.resolve()
45    members = archive.getmembers()
46    for member in members:
47        if member.issym() or member.islnk():
48            raise ValueError(f"archive links are not supported: {member.name}")
49        target = (destination / member.name).resolve()
50        if destination != target and destination not in target.parents:
51            raise ValueError(f"archive member escapes destination: {member.name}")
52    archive.extractall(destination, members=members)
53
54
55def command_exists(name: str) -> bool:
56    return shutil.which(name) is not None
57
58
59def run_command(
60    command: list[str], *, capture: bool = False
61) -> subprocess.CompletedProcess[str]:
62    return subprocess.run(command, check=True, text=True, capture_output=capture)
63
64
65def dedupe_preserving_order(values: list[str]) -> list[str]:
66    seen = set()
67    result = []
68    for value in values:
69        if value in seen:
70            continue
71        seen.add(value)
72        result.append(value)
73    return result
74
75
76def parse_comma_values(values: list[str] | None) -> list[str]:
77    result = []
78    for value in values or []:
79        result.extend(part.strip() for part in value.split(",") if part.strip())
80    return result
81
82
83def parse_parts(values: list[str] | None) -> list[str]:
84    requested = parse_comma_values(values)
85    if any(value.lower() == "all" for value in requested):
86        return list(ALL_SWE_ATLAS_PARTS)
87    return dedupe_preserving_order(
88        [config.canonical_part(value) for value in requested]
89    )
90
91
92def configure_modal_environment(args: argparse.Namespace) -> None:
93    os.environ["AGENT_EVALS_APP_NAME"] = args.app_name
94    os.environ["AGENT_EVALS_VOLUME"] = args.volume
95    os.environ["AGENT_EVALS_MODAL_TOKEN_SECRET"] = args.modal_token_secret
96    os.environ["AGENT_EVALS_LLM_PROVIDERS_SECRET"] = args.api_secret
97
98
99def import_modal_app(args: argparse.Namespace):
100    configure_modal_environment(args)
101    from . import modal_app
102
103    return modal_app
104
105
106class AppNotDeployedError(RuntimeError):
107    """Raised when a Modal function lookup fails because the app (or that
108    specific function) has not been deployed. Surfaced as a clean CLI error
109    pointing the user at `zed-eval deploy`.
110    """
111
112
113def deploy_app(args: argparse.Namespace) -> None:
114    """Publish the Modal app. This is the ONLY code path that deploys.
115
116    Deploying replaces the live app and cancels any in-flight eval runs, so it
117    must happen exclusively via the `deploy` subcommand. Every other command
118    just looks up already-deployed functions with `deployed_function`. After
119    changing harness code, run `zed-eval deploy` once.
120    """
121    modal_app = import_modal_app(args)
122    print(f"Deploying Modal app '{args.app_name}'...")
123    modal_app.app.deploy(name=args.app_name)
124
125
126def deployed_function(args: argparse.Namespace, function_name: str):
127    """Look up an already-deployed Modal function. Never deploys.
128
129    `modal.Function.from_name` resolves a published function without deploying,
130    which is exactly what we want: deploying would cancel in-flight runs. We
131    hydrate eagerly so a missing app/function fails here with an actionable
132    message instead of deep inside a later `.spawn()`/`.remote()` call.
133    """
134    import modal
135    from modal.exception import NotFoundError
136
137    configure_modal_environment(args)
138    function = modal.Function.from_name(args.app_name, function_name)
139    try:
140        function.hydrate()
141    except NotFoundError as error:
142        raise AppNotDeployedError(
143            f"app '{args.app_name}' not deployed (or function "
144            f"'{function_name}' missing) — run 'zed-eval deploy' first"
145        ) from error
146    return function
147
148
149def modal_call_id(call: Any) -> str:
150    for attribute in ("object_id", "id", "function_call_id"):
151        value = getattr(call, attribute, None)
152        if value:
153            return str(value)
154    return str(call)
155
156
157def print_table(rows: list[dict[str, Any]], columns: list[tuple[str, str]]) -> None:
158    if not rows:
159        print("No rows")
160        return
161    widths = []
162    for key, title in columns:
163        width = len(title)
164        for row in rows:
165            width = max(width, len(str(row.get(key) or "")))
166        widths.append(min(width, 80))
167    print(
168        "  ".join(
169            title.ljust(widths[index]) for index, (_key, title) in enumerate(columns)
170        )
171    )
172    print("  ".join("-" * width for width in widths))
173    for row in rows:
174        cells = []
175        for index, (key, _title) in enumerate(columns):
176            text = str(row.get(key) or "")
177            if len(text) > widths[index]:
178                text = text[: widths[index] - 1] + "…"
179            cells.append(text.ljust(widths[index]))
180        print("  ".join(cells))
181
182
183def default_namespace(args: argparse.Namespace) -> str:
184    return source.sanitize_namespace(args.namespace or source.default_namespace())
185
186
187def resolve_run_location(args: argparse.Namespace, run_id: str) -> tuple[str, str]:
188    """Resolve (namespace, experiment_name) for a run-lookup command.
189
190    When `--experiment-name` is given, behaves exactly as before: namespace comes
191    from `--namespace` or the local git default. When it is omitted, the run is
192    located via the launch-time local run index (`run_index`) so a bare run id is
193    enough. Raises a friendly error pointing at the explicit flags when the run
194    can't be found.
195    """
196    from . import run_index
197
198    explicit_experiment = getattr(args, "experiment_name", None)
199    explicit_namespace = getattr(args, "namespace", None)
200    if explicit_experiment:
201        namespace = explicit_namespace or source.default_namespace()
202        return (
203            source.sanitize_namespace(namespace),
204            source.sanitize_namespace(explicit_experiment),
205        )
206    entry = run_index.lookup(run_id)
207    if not entry:
208        raise ValueError(
209            f"could not locate run '{run_id}' in the local run index "
210            f"({run_index.index_path()}). Pass --experiment-name (and "
211            "--namespace if it isn't yours). The index is written automatically "
212            "when you launch a run from this machine."
213        )
214    namespace = (
215        explicit_namespace or entry.get("namespace") or source.default_namespace()
216    )
217    return (
218        source.sanitize_namespace(namespace),
219        source.sanitize_namespace(entry["experiment_name"]),
220    )
221