tenant.openagents/omega

1from __future__ import annotations
2
3import json
4import os
5import shlex
6from collections.abc import Callable
7from pathlib import Path
8from typing import Any
9
10PROVIDER_ENV_MAP = {
11    "anthropic": "ANTHROPIC_API_KEY",
12    "openai": "OPENAI_API_KEY",
13    "google": "GEMINI_API_KEY",
14    "gemini": "GEMINI_API_KEY",
15    "deepseek": "DEEPSEEK_API_KEY",
16    "mistral": "MISTRAL_API_KEY",
17}
18
19
20def provider_api_env(
21    model_name: str | None,
22    get_env: Callable[[str], str | None] = os.environ.get,
23) -> dict[str, str]:
24    env: dict[str, str] = {}
25    if not model_name or "/" not in model_name:
26        return env
27
28    provider = model_name.split("/", 1)[0]
29    env_var = PROVIDER_ENV_MAP.get(provider) or (
30        f"{provider}_API_KEY".upper().replace("-", "_")
31    )
32    api_key = get_env(env_var)
33    if api_key:
34        env[env_var] = api_key
35    return env
36
37
38def add_openai_compatible_provider_env(
39    env: dict[str, str], providers_json: str | None
40) -> None:
41    if providers_json:
42        env["ZED_OPENAI_COMPATIBLE_PROVIDERS"] = providers_json
43
44
45def add_anthropic_available_models_env(
46    env: dict[str, str], models_json: str | None
47) -> None:
48    if models_json:
49        env["ZED_ANTHROPIC_AVAILABLE_MODELS"] = models_json
50
51
52def add_zed_eval_env(
53    env: dict[str, str], extra_env: dict[str, str], *, exclude: set[str] | None = None
54) -> None:
55    exclude = exclude or set()
56    for key, value in extra_env.items():
57        if key.startswith("ZED_EVAL_") and key not in exclude:
58            env[key] = value
59
60
61async def detect_workdir(
62    environment: Any,
63    exec_as_agent: Callable[..., Any],
64    get_env: Callable[[str], str | None],
65    error_message: str,
66) -> str:
67    override = get_env("EVAL_CLI_WORKDIR")
68    if override:
69        return override
70
71    result = await exec_as_agent(
72        environment,
73        command=(
74            "for d in /app /testbed /repo; do "
75            '  if [ -d "$d/.git" ]; then echo "$d"; exit 0; fi; '
76            "done; "
77            "find / -maxdepth 3 -name .git -type d 2>/dev/null "
78            '| head -1 | sed "s|/.git$||"'
79        ),
80    )
81    workdir = (result.stdout or "").strip()
82    if workdir:
83        return workdir
84
85    result = await exec_as_agent(
86        environment,
87        command=(
88            "for d in /app /testbed /repo /root /home; do "
89            '  if [ -d "$d" ]; then echo "$d"; exit 0; fi; '
90            "done; pwd"
91        ),
92    )
93    workdir = (result.stdout or "").strip()
94    if workdir:
95        return workdir
96    raise RuntimeError(error_message)
97
98
99def populate_context_from_result(logs_dir: Path, context: Any, logger: Any) -> None:
100    result_data = None
101    for json_file in logs_dir.rglob("result.json"):
102        try:
103            result_data = json.loads(json_file.read_text())
104            break
105        except (json.JSONDecodeError, OSError):
106            continue
107
108    if result_data is None:
109        logger.warning("Could not find or parse result.json from eval-cli")
110        return
111
112    if result_data.get("input_tokens") is not None:
113        context.n_input_tokens = result_data["input_tokens"]
114    if result_data.get("output_tokens") is not None:
115        context.n_output_tokens = result_data["output_tokens"]
116    if result_data.get("cache_read_input_tokens") is not None:
117        context.n_cache_tokens = result_data["cache_read_input_tokens"]
118    if isinstance(result_data.get("step_count"), int):
119        context.n_agent_steps = result_data["step_count"]
120
121    context.metadata = {
122        "status": result_data.get("status"),
123        "duration_secs": result_data.get("duration_secs"),
124        "model": result_data.get("model"),
125        "tool_call_count": result_data.get("tool_call_count"),
126    }
127
128
129def eval_cli_with_log_command(
130    parts: list[str],
131    log_path: str,
132    *,
133    timeout_message: str | None = None,
134    line_buffered: bool = False,
135) -> str:
136    """Run eval-cli, tee output, and preserve eval-cli's exit status.
137
138    POSIX shells return the last command's status for a pipeline, so `cmd | tee`
139    would otherwise hide eval-cli failures whenever `tee` succeeds.
140    """
141    status_file = "/tmp/zed-eval-eval-cli-status"
142    quoted_status_file = shlex.quote(status_file)
143    quoted_log_path = shlex.quote(log_path)
144    timeout_handler = 'if [ "$ec" -eq 2 ]; then ec=0; fi; '
145    if timeout_message:
146        timeout_handler = (
147            f'if [ "$ec" -eq 2 ]; then echo {shlex.quote(timeout_message)}; ec=0; fi; '
148        )
149    tee_command = f"tee {quoted_log_path}"
150    if line_buffered:
151        tee_command = (
152            "if command -v stdbuf >/dev/null 2>&1; "
153            f"then stdbuf -oL tee {quoted_log_path}; "
154            f"else tee {quoted_log_path}; fi"
155        )
156    return (
157        f'status_file={quoted_status_file}; rm -f "$status_file"; '
158        "( "
159        + " ".join(parts)
160        + "; ec=$?; "
161        + timeout_handler
162        + 'printf "%s\\n" "$ec" > "$status_file"; '
163        + 'exit "$ec" ) 2>&1 | '
164        + tee_command
165        + '; ec=1; if [ -s "$status_file" ]; then read ec < "$status_file"; fi; '
166        + 'rm -f "$status_file"; exit "$ec"'
167    )
168
169
170def patch_command(agent_dir: str) -> str:
171    patch_path = shlex.quote(f"{agent_dir}/patch.diff")
172    return (
173        "if git rev-parse --git-dir >/dev/null 2>&1; then "
174        "git add -A && "
175        "if git rev-parse --verify HEAD >/dev/null 2>&1; then "
176        f"git diff --cached HEAD -- > {patch_path} && "
177        f'echo "Patch size: $(wc -c < {patch_path}) bytes"; '
178        "else "
179        'echo "Git repo has no valid HEAD, skipping patch generation"; '
180        "fi; "
181        "else "
182        'echo "No git repo found, skipping patch generation"; '
183        "fi"
184    )
185