orizu 0.3.1 → 0.3.3

package/vendor/orizu-gepa-python/src/orizu_gepa/reflection.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+import json
+import os
+import ssl
+import subprocess
+import tempfile
+import urllib.error
+import urllib.request
+from typing import Any
+
+from .optimizer import ReflectionResult, RowEvaluation, TextGepaConfig, build_reflection_prompt, extract_candidate_text
+
+
+def _read_anthropic_response_with_curl(api_key: str, payload: dict[str, Any]) -> dict[str, Any]:
+    with tempfile.NamedTemporaryFile("w", encoding="utf-8", suffix=".json") as body_file:
+        json.dump(payload, body_file)
+        body_file.flush()
+        curl_config = "\n".join([
+            'url = "https://api.anthropic.com/v1/messages"',
+            'request = "POST"',
+            'header = "anthropic-version: 2023-06-01"',
+            'header = "content-type: application/json"',
+            f'header = "x-api-key: {api_key}"',
+            f'data-binary = "@{body_file.name}"',
+            "fail-with-body",
+            "silent",
+            "show-error",
+        ])
+        result = subprocess.run(
+            ["curl", "--config", "-"],
+            input=curl_config,
+            text=True,
+            capture_output=True,
+            check=False,
+        )
+    if result.returncode != 0:
+        detail = "\n".join(part for part in [result.stderr.strip(), result.stdout.strip()] if part)
+        raise RuntimeError(f"Reflection LM failed via curl: {detail}")
+    return json.loads(result.stdout)
+
+
+def _merge_provider_settings(
+    payload: dict[str, Any],
+    settings: dict[str, Any],
+    *,
+    reserved_keys: set[str],
+) -> dict[str, Any]:
+    blocked = reserved_keys.intersection(settings)
+    if blocked:
+        raise RuntimeError(
+            f"reflection_provider_settings cannot override reserved request keys: {', '.join(sorted(blocked))}"
+        )
+    return {**payload, **settings}
+
+
+def _add_temperature(payload: dict[str, Any], config: TextGepaConfig) -> dict[str, Any]:
+    if config.reflection_temperature is None:
+        return payload
+    if "temperature" in payload and payload["temperature"] != config.reflection_temperature:
+        raise RuntimeError("reflection_temperature conflicts with reflection_provider_settings.temperature")
+    return {**payload, "temperature": config.reflection_temperature}
+
+
+def build_anthropic_reflection_payload(model: str, prompt: str, config: TextGepaConfig) -> dict[str, Any]:
+    payload = _merge_provider_settings({
+        "model": model,
+        "max_tokens": config.reflection_max_tokens,
+        "messages": [{"role": "user", "content": prompt}],
+    }, config.reflection_provider_settings, reserved_keys={"model", "max_tokens", "messages"})
+    payload = _add_temperature(payload, config)
+    if "temperature" in payload and "thinking" in payload:
+        raise RuntimeError("reflection_temperature cannot be combined with Anthropic thinking")
+    return payload
+
+
+def build_openai_reflection_payload(model: str, prompt: str, config: TextGepaConfig) -> dict[str, Any]:
+    payload = _merge_provider_settings({
+        "model": model,
+        "input": [{"role": "user", "content": prompt}],
+        "max_output_tokens": config.reflection_max_tokens,
+    }, config.reflection_provider_settings, reserved_keys={"model", "input", "max_output_tokens"})
+    return _add_temperature(payload, config)
+
+
+def _extract_openai_output_text(data: dict[str, Any]) -> str:
+    output_text = data.get("output_text")
+    if isinstance(output_text, str):
+        return output_text
+    text_parts: list[str] = []
+    for item in data.get("output", []):
+        if not isinstance(item, dict) or item.get("type") != "message":
+            continue
+        for content in item.get("content", []):
+            if isinstance(content, dict) and content.get("type") == "output_text":
+                text = content.get("text")
+                if isinstance(text, str):
+                    text_parts.append(text)
+    return "".join(text_parts)
+
+
+def reflect_with_anthropic(parent_text: str, parent_results: list[RowEvaluation], config: TextGepaConfig) -> ReflectionResult:
+    api_key = os.environ.get("ANTHROPIC_API_KEY")
+    if not api_key:
+        raise RuntimeError("ANTHROPIC_API_KEY is required for the default reflective LM")
+    model = config.reflection_model
+    if model.startswith("anthropic/"):
+        model = model.split("/", 1)[1]
+    prompt = build_reflection_prompt(parent_text, parent_results, config)
+    payload = build_anthropic_reflection_payload(model, prompt, config)
+    request = urllib.request.Request(
+        "https://api.anthropic.com/v1/messages",
+        data=json.dumps(payload).encode("utf-8"),
+        method="POST",
+        headers={
+            "x-api-key": api_key,
+            "anthropic-version": "2023-06-01",
+            "content-type": "application/json",
+        },
+    )
+    try:
+        with urllib.request.urlopen(request, timeout=180) as response:
+            data = json.loads(response.read().decode("utf-8"))
+    except urllib.error.HTTPError as error:
+        detail = error.read().decode("utf-8", errors="replace")
+        raise RuntimeError(f"Reflection LM failed: {error.code} {detail}") from error
+    except urllib.error.URLError as error:
+        if not isinstance(error.reason, ssl.SSLCertVerificationError):
+            raise
+        data = _read_anthropic_response_with_curl(api_key, payload)
+    parts = data.get("content") or []
+    text = "".join(part.get("text", "") for part in parts if isinstance(part, dict))
+    return ReflectionResult(
+        prompt=prompt,
+        response=text,
+        candidate_text=extract_candidate_text(text),
+    )
+
+
+def reflect_with_openai(parent_text: str, parent_results: list[RowEvaluation], config: TextGepaConfig) -> ReflectionResult:
+    api_key = os.environ.get("OPENAI_API_KEY")
+    if not api_key:
+        raise RuntimeError("OPENAI_API_KEY is required for OpenAI reflection models")
+    model = config.reflection_model
+    if model.startswith("openai/"):
+        model = model.split("/", 1)[1]
+    prompt = build_reflection_prompt(parent_text, parent_results, config)
+    payload = build_openai_reflection_payload(model, prompt, config)
+    request = urllib.request.Request(
+        "https://api.openai.com/v1/responses",
+        data=json.dumps(payload).encode("utf-8"),
+        method="POST",
+        headers={
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+        },
+    )
+    try:
+        with urllib.request.urlopen(request, timeout=180) as response:
+            data = json.loads(response.read().decode("utf-8"))
+    except urllib.error.HTTPError as error:
+        detail = error.read().decode("utf-8", errors="replace")
+        raise RuntimeError(f"OpenAI reflection LM failed: {error.code} {detail}") from error
+    text = _extract_openai_output_text(data)
+    if not text:
+        raise RuntimeError("OpenAI reflection LM returned no output text")
+    return ReflectionResult(
+        prompt=prompt,
+        response=text,
+        candidate_text=extract_candidate_text(text),
+    )
+
+
+def reflect_with_provider(parent_text: str, parent_results: list[RowEvaluation], config: TextGepaConfig) -> ReflectionResult:
+    if config.reflection_model.startswith("openai/"):
+        return reflect_with_openai(parent_text, parent_results, config)
+    return reflect_with_anthropic(parent_text, parent_results, config)

package/vendor/orizu-gepa-python/src/orizu_gepa/runner.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import tempfile
+import zipfile
+from pathlib import Path
+from typing import Any
+
+from .client import OrizuClient
+from .optimizer import DatasetRow, PromptContext, RunnerCallResult
+
+MAX_RUNNER_ZIP_BYTES = 25 * 1024 * 1024
+MAX_RUNNER_ZIP_ENTRIES = 1000
+MAX_RUNNER_UNCOMPRESSED_BYTES = 75 * 1024 * 1024
+MAX_RUNNER_OUTPUT_BYTES = 2 * 1024 * 1024
+DEFAULT_RUNNER_TIMEOUT_SECONDS = 120
+
+ALLOWED_RUNNER_ENV_KEYS = {
+    "PATH",
+    "SystemRoot",
+    "WINDIR",
+    "HOME",
+    "TMPDIR",
+    "TEMP",
+    "TMP",
+    "LANG",
+    "LC_ALL",
+    "PYTHONPATH",
+    "NODE_PATH",
+    "ANTHROPIC_API_KEY",
+    "OPENAI_API_KEY",
+    "GEMINI_API_KEY",
+    "GOOGLE_API_KEY",
+}
+
+
+def read_manifest(runner_dir: str | Path) -> dict[str, Any]:
+    path = Path(runner_dir) / "manifest.json"
+    return json.loads(path.read_text())
+
+
+def _safe_extract_zip(zip_path: Path, destination: Path) -> None:
+    if zip_path.stat().st_size > MAX_RUNNER_ZIP_BYTES:
+        raise RuntimeError("Runner artifact exceeds the maximum zip size")
+
+    with zipfile.ZipFile(zip_path) as archive:
+        infos = archive.infolist()
+        if len(infos) > MAX_RUNNER_ZIP_ENTRIES:
+            raise RuntimeError("Runner artifact contains too many files")
+
+        total_uncompressed = 0
+        destination_resolved = destination.resolve()
+        for info in infos:
+            total_uncompressed += info.file_size
+            if total_uncompressed > MAX_RUNNER_UNCOMPRESSED_BYTES:
+                raise RuntimeError("Runner artifact exceeds the maximum expanded size")
+
+            target = (destination / info.filename).resolve()
+            if os.path.commonpath([str(destination_resolved), str(target)]) != str(destination_resolved):
+                raise RuntimeError("Runner artifact contains an unsafe path")
+
+        archive.extractall(destination)
+
+
+def _runner_env(input_path: Path, output_path: Path) -> dict[str, str]:
+    env = {
+        key: value
+        for key, value in os.environ.items()
+        if key in ALLOWED_RUNNER_ENV_KEYS
+    }
+    env["ORIZU_RUNNER_INPUT_PATH"] = str(input_path)
+    env["ORIZU_RUNNER_OUTPUT_PATH"] = str(output_path)
+    return env
+
+
+def _bounded_stream(value: str) -> str:
+    encoded = value.encode("utf-8", errors="replace")
+    if len(encoded) <= MAX_RUNNER_OUTPUT_BYTES:
+        return value
+    return encoded[:MAX_RUNNER_OUTPUT_BYTES].decode("utf-8", errors="replace") + "\n[truncated]"
+
+
+def materialize_runner_version(client: OrizuClient, runner_version_id: str) -> tempfile.TemporaryDirectory[str]:
+    temp_dir = tempfile.TemporaryDirectory(prefix="orizu-gepa-runner-")
+    data = client._request_bytes(f"/api/cli/runner-versions/{runner_version_id}/download")
+    zip_path = Path(temp_dir.name) / "runner.zip"
+    runner_dir = Path(temp_dir.name) / "runner"
+    zip_path.write_bytes(data)
+    _safe_extract_zip(zip_path, runner_dir)
+    return temp_dir
+
+
+def run_file_contract_runner(
+    *,
+    runner_dir: str | Path,
+    row: dict[str, Any],
+    prompt_body: str | None,
+    body_kind: str,
+    provider_settings: dict[str, Any],
+    prompt_version_id: str,
+    runner_version_id: str,
+    run_id: str | None,
+    timeout_seconds: int = DEFAULT_RUNNER_TIMEOUT_SECONDS,
+) -> RunnerCallResult:
+    runner_path = Path(runner_dir)
+    manifest = read_manifest(runner_path)
+    command = manifest.get("command")
+    if not isinstance(command, list) or not all(isinstance(item, str) for item in command):
+        raise RuntimeError(f"Runner manifest at {runner_path / 'manifest.json'} must include command: string[]")
+
+    with tempfile.TemporaryDirectory(prefix="orizu-gepa-call-") as temp_dir:
+        input_path = Path(temp_dir) / "input.json"
+        output_path = Path(temp_dir) / "output.json"
+        input_path.write_text(json.dumps({
+            "row": row,
+            "prompt": {
+                "body": prompt_body,
+                "body_kind": body_kind,
+                "provider_settings": provider_settings,
+            },
+            "prompt_version_id": prompt_version_id,
+            "runner_version_id": runner_version_id,
+            "run_id": run_id,
+        }, ensure_ascii=False))
+
+        try:
+            result = subprocess.run(
+                command,
+                cwd=runner_path,
+                env=_runner_env(input_path, output_path),
+                text=True,
+                capture_output=True,
+                check=False,
+                timeout=timeout_seconds,
+            )
+        except subprocess.TimeoutExpired as exc:
+            stdout = _bounded_stream(exc.stdout or "")
+            stderr = _bounded_stream(exc.stderr or "")
+            raise RuntimeError(
+                f"Runner timed out after {timeout_seconds}s: {stderr or stdout}"
+            ) from exc
+
+        if result.returncode != 0:
+            stdout = _bounded_stream(result.stdout or "")
+            stderr = _bounded_stream(result.stderr or "")
+            raise RuntimeError(
+                f"Runner failed with exit code {result.returncode}: {stderr or stdout}"
+            )
+        data = json.loads(output_path.read_text())
+        known = {
+            "model_response",
+            "raw_api_response",
+            "token_in",
+            "token_out",
+            "latency_ms",
+            "cost_usd",
+            "error",
+        }
+        return RunnerCallResult(
+            model_response=data.get("model_response"),
+            raw_api_response=data.get("raw_api_response"),
+            token_in=data.get("token_in"),
+            token_out=data.get("token_out"),
+            latency_ms=data.get("latency_ms"),
+            cost_usd=data.get("cost_usd"),
+            error=data.get("error"),
+            extra={key: value for key, value in data.items() if key not in known},
+        )
+
+
+def make_candidate_runner(runner_dir: str | Path, run_id: str | None):
+    def _run(candidate_text: str, row: DatasetRow, prompt_context: PromptContext, candidate_id: str) -> RunnerCallResult:
+        return run_file_contract_runner(
+            runner_dir=runner_dir,
+            row=row.row,
+            prompt_body=candidate_text,
+            body_kind=prompt_context.body_kind,
+            provider_settings=prompt_context.provider_settings,
+            prompt_version_id=prompt_context.prompt_version_id,
+            runner_version_id=prompt_context.runner_version_id,
+            run_id=run_id,
+        )
+    return _run
+
+
+def make_scorer_runner(runner_dir: str | Path, run_id: str | None):
+    def _run(row: DatasetRow, candidate_result: RunnerCallResult, scorer_context: PromptContext, candidate_id: str) -> RunnerCallResult:
+        scorer_row = {
+            "source_row": row.row,
+            "candidate_id": candidate_id,
+            "candidate_output": candidate_result.model_response,
+            "candidate_raw_response": candidate_result.raw_api_response,
+            "candidate_error": candidate_result.error,
+        }
+        return run_file_contract_runner(
+            runner_dir=runner_dir,
+            row=scorer_row,
+            prompt_body=scorer_context.body,
+            body_kind=scorer_context.body_kind,
+            provider_settings=scorer_context.provider_settings,
+            prompt_version_id=scorer_context.prompt_version_id,
+            runner_version_id=scorer_context.runner_version_id,
+            run_id=run_id,
+        )
+    return _run