training/codex_client.py

"""Small text-generation adapter backed by `codex exec`."""

from pathlib import Path
from types import SimpleNamespace
import subprocess
import tempfile


class CodexModels:
    def generate_content(self, model, contents, config=None):
        del model, config
        with tempfile.TemporaryDirectory(prefix="needle-codex-") as temp_dir:
            output_path = Path(temp_dir) / "response.txt"
            process = subprocess.run(
                [
                    "codex",
                    "exec",
                    "--ephemeral",
                    "--sandbox",
                    "read-only",
                    "--skip-git-repo-check",
                    "--color",
                    "never",
                    "--output-last-message",
                    str(output_path),
                    "-",
                ],
                input=contents,
                text=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                timeout=600,
                check=False,
            )
            if process.returncode != 0:
                detail = (process.stderr or process.stdout).strip()[-4000:]
                raise RuntimeError(f"codex exec failed ({process.returncode}): {detail}")
            if not output_path.exists():
                raise RuntimeError("codex exec produced no final-response file")
            text = output_path.read_text(encoding="utf-8").strip()
            if not text:
                raise RuntimeError("codex exec produced an empty final response")
            return SimpleNamespace(text=text)


class CodexClient:
    def __init__(self):
        self.models = CodexModels()


class CodexClientPool:
    """Matches the tiny client-pool interface used by Needle's generator."""

    def __init__(self):
        self._client = CodexClient()

    def get(self):
        return self._client