Enrich a CSV of Companies Overnight with the Batch API

httpx==0.28.1
pytest==8.4.2

name,website
Apple,https://www.apple.com
Microsoft,https://www.microsoft.com
Nvidia,https://www.nvidia.com
Shopify,https://www.shopify.com
Toyota,https://global.toyota
Nestlé,https://www.nestle.com
Siemens,https://www.siemens.com
Spotify,https://www.spotify.com
Atlassian,https://www.atlassian.com
Salesforce,https://www.salesforce.com
Netflix,https://www.netflix.com
Airbnb,https://www.airbnb.com

"""Build batch requests, validate results, join CSV rows, and write retries without network calls."""

from __future__ import annotations

import csv
import json
import re
from dataclasses import dataclass, field
from typing import Any

MODEL = "openai/gpt-5.6-sol"  # one model per file; every line must match
BATCH_URL = "/v1/responses"  # the only url the Batch API accepts
CUSTOM_ID_MAX_BYTES = 64

ENRICHMENT_FIELDS = (
    "headquarters_city",
    "founded_year",
    "description",
    "latest_news_headline",
    "latest_news_url",
)
RESPONSE_SCHEMA: dict[str, Any] = {
    "type": "object",
    "properties": {
        "headquarters_city": {"type": "string"},
        "founded_year": {"type": "integer"},
        "description": {"type": "string"},
        "latest_news_headline": {"type": "string"},
        "latest_news_url": {"type": "string"},
    },
    "required": list(ENRICHMENT_FIELDS),
    "additionalProperties": False,
}
INSTRUCTIONS = (
    "You are a research assistant filling in a company database. Use web search "
    "to confirm every field against a current source. For the news fields, use the "
    "most recent news article about the company and return that article's URL. "
    "Return only the JSON object; no markdown, no commentary."
)


def make_custom_id(index: int, name: str) -> str:
    """Build a unique, readable id of at most 64 bytes: `row-003-nvidia`."""
    slug = re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-") or "row"
    prefix = f"row-{index:03d}-"
    budget = CUSTOM_ID_MAX_BYTES - len(prefix)
    return prefix + slug.encode("utf-8")[:budget].decode("utf-8", "ignore")


def read_companies(path: str) -> list[dict[str, str]]:
    """Read the CSV and attach a custom_id to every row."""
    with open(path, newline="", encoding="utf-8") as handle:
        rows = [dict(row) for row in csv.DictReader(handle)]
    for index, row in enumerate(rows, start=1):
        row["custom_id"] = make_custom_id(index, row["name"])
    return rows


def build_request_body(row: dict[str, str]) -> dict[str, Any]:
    """One Agent API request body. Shared by the batch file and a live spot check."""
    prompt = (
        f"Company: {row['name']}\nWebsite: {row['website']}\n\n"
        "Find the headquarters city, the founding year, a one-sentence description "
        "of what the company does, and the most recent news headline about the "
        "company with the URL of that article."
    )
    return {
        "model": MODEL,
        "instructions": INSTRUCTIONS,
        "input": prompt,
        "tools": [{"type": "web_search"}],
        "response_format": {
            "type": "json_schema",
            "json_schema": {"name": "company_enrichment", "schema": RESPONSE_SCHEMA},
        },
    }


def build_batch_lines(rows: list[dict[str, str]]) -> list[dict[str, Any]]:
    """One batch request per row: custom_id, method, url, body."""
    return [
        {
            "custom_id": row["custom_id"],
            "method": "POST",
            "url": BATCH_URL,
            "body": build_request_body(row),
        }
        for row in rows
    ]


def write_jsonl(lines: list[dict[str, Any]], path: str) -> None:
    """One compact JSON object per line, LF endings, no blank lines."""
    with open(path, "w", encoding="utf-8", newline="\n") as handle:
        handle.writelines(json.dumps(line, ensure_ascii=False) + "\n" for line in lines)


def read_jsonl(path: str) -> list[dict[str, Any]]:
    with open(path, encoding="utf-8") as handle:
        return [json.loads(line) for line in handle if line.strip()]


@dataclass
class Result:
    """What the recipe keeps for one row after the batch finishes."""

    custom_id: str
    status: str  # completed | incomplete | failed | error | parse_error | missing
    fields: dict[str, Any] = field(default_factory=dict)
    error: str = ""
    response_id: str = ""
    cost_usd: float = 0.0
    tool_calls_cost_usd: float = 0.0

    @property
    def ok(self) -> bool:
        return self.status == "completed"


def message_text(body: dict[str, Any]) -> str:
    """Concatenate the output_text parts of every message in a Response."""
    parts: list[str] = []
    for item in body.get("output") or []:
        if item.get("type") != "message":
            continue
        for part in item.get("content") or []:
            if part.get("type") == "output_text":
                parts.append(part.get("text", ""))
    return "".join(parts)


def parse_enrichment(text: str) -> dict[str, Any] | None:
    """Parse the model's JSON. Return None when it is not the object you asked for."""
    cleaned = text.strip()
    if cleaned.startswith("`"):  # strip a markdown fence if the model added one anyway
        cleaned = re.sub(r"^`+(?:json)?\s*|\s*`+$", "", cleaned)
    try:
        data = json.loads(cleaned)
    except json.JSONDecodeError:
        return None
    if not isinstance(data, dict) or set(data) != set(ENRICHMENT_FIELDS):
        return None
    if type(data["founded_year"]) is not int:
        return None
    if any(not isinstance(data[k], str) for k in ENRICHMENT_FIELDS if k != "founded_year"):
        return None
    return {k: data[k] for k in ENRICHMENT_FIELDS}


def parse_output_line(line: dict[str, Any]) -> Result:
    """One line of the output file: a full Response under response.body."""
    body = (line.get("response") or {}).get("body") or {}
    status = body.get("status", "failed")
    cost = (body.get("usage") or {}).get("cost") or {}
    result = Result(
        custom_id=line["custom_id"],
        status=status,
        response_id=body.get("id", ""),
        cost_usd=float(cost.get("total_cost") or 0.0),
        tool_calls_cost_usd=float(cost.get("tool_calls_cost") or 0.0),
    )
    if status != "completed":
        error = body.get("error") or {}
        result.error = error.get("message", f"response status {status}")
        return result
    fields = parse_enrichment(message_text(body))
    if fields is None:
        result.status = "parse_error"
        result.error = "response text was not the expected JSON object"
        return result
    result.fields = fields
    return result


def parse_error_line(line: dict[str, Any]) -> Result:
    """One line of the error file. Two shapes: an API rejection, or a stopped run."""
    response = line.get("response")
    if response:  # the API rejected the request: status_code plus an error body
        err = (response.get("body") or {}).get("error") or {}
        message = err.get("message", "request rejected")
        return Result(
            custom_id=line["custom_id"],
            status="error",
            error=f"HTTP {response.get('status_code')}: {message}",
        )
    err = line.get("error") or {}  # the run stopped: batch_expired, poll_timeout, ...
    return Result(
        custom_id=line["custom_id"],
        status="error",
        error=f"{err.get('code', 'unknown')}: {err.get('message', '')}".strip(": "),
    )


def collect_results(
    rows: list[dict[str, str]],
    output_lines: list[dict[str, Any]],
    error_lines: list[dict[str, Any]],
) -> dict[str, Result]:
    """Index every result by custom_id. Rows with no line at all become `missing`."""
    results: dict[str, Result] = {}
    for line in output_lines:
        results[line["custom_id"]] = parse_output_line(line)
    for line in error_lines:
        results[line["custom_id"]] = parse_error_line(line)
    for row in rows:
        results.setdefault(
            row["custom_id"],
            Result(row["custom_id"], "missing", error="no output or error line"),
        )
    return results


def write_enriched_csv(rows: list[dict[str, str]], results: dict[str, Result], path: str) -> None:
    """Join by custom_id and keep every input row, enriched or not."""
    base_columns = [c for c in rows[0] if c != "custom_id"] if rows else []
    columns = list(dict.fromkeys(base_columns + list(ENRICHMENT_FIELDS) + ["status", "error", "response_id"]))
    with open(path, "w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=columns)
        writer.writeheader()
        for row in rows:
            result = results[row["custom_id"]]
            out = {c: row.get(c, "") for c in base_columns}
            out.update({k: result.fields.get(k, "") for k in ENRICHMENT_FIELDS})
            out["status"], out["error"] = result.status, result.error
            out["response_id"] = result.response_id
            writer.writerow(out)


def retry_lines(batch_lines: list[dict], results: dict[str, Result]) -> list[dict]:
    """The original request lines for every row that did not complete."""
    failed = {cid for cid, result in results.items() if not result.ok}
    return [line for line in batch_lines if line["custom_id"] in failed]


def summarize(results: dict[str, Result]) -> dict[str, Any]:
    """Counts by status plus the summed per-request cost from the output file."""
    counts: dict[str, int] = {}
    for result in results.values():
        counts[result.status] = counts.get(result.status, 0) + 1
    return {
        "total": len(results),
        "completed": counts.get("completed", 0),
        "needs_retry": sum(1 for r in results.values() if not r.ok),
        "by_status": dict(sorted(counts.items())),
        "cost_usd": round(sum(r.cost_usd for r in results.values()), 4),
        "tool_calls_cost_usd": round(sum(r.tool_calls_cost_usd for r in results.values()), 4),
    }

"""Enrich companies.csv overnight with the Batch API.

Steps: read the CSV, write batch_input.jsonl, upload it, create a batch, poll
until a final status, download the output and error files, join by custom_id
into companies_enriched.csv, write a retry file, and print a summary.

State (file id, batch id, idempotency key) is saved to batch_state.json after
each step, so `--resume` picks up where a crash or Ctrl-C left off.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
import time
from typing import Any

import enrichment as en
import httpx

API = "https://api.perplexity.ai"
FINAL_STATUSES = {"completed", "expired", "cancelled", "failed"}
POLL_START_SECONDS = 30
POLL_MAX_SECONDS = 300


class BatchError(Exception):
    """A batch or file endpoint answered with an error."""


def check(response: httpx.Response, step: str) -> httpx.Response:
    """Raise a readable BatchError on any non-2xx answer."""
    if response.status_code == 404:
        raise BatchError(
            f"{step}: 404. The batch or file was not found, belongs to another "
            "organization, or has expired. See the Batch API guide."
        )
    if response.status_code >= 400:
        raise BatchError(f"{step}: HTTP {response.status_code} {response.text.strip()}")
    return response


def upload_input(client: httpx.Client, path: str) -> str:
    """POST /v1/files with purpose=batch. The API validates every line first."""
    with open(path, "rb") as handle:
        response = client.post(
            f"{API}/v1/files",
            data={"purpose": "batch"},
            files={"file": (os.path.basename(path), handle, "application/jsonl")},
        )
    return check(response, "upload").json()["id"]


def create_batch(
    client: httpx.Client, file_id: str, idempotency_key: str, job: str
) -> dict[str, Any]:
    """POST /v1/batches. A retried create with the same Idempotency-Key returns the same batch."""
    response = client.post(
        f"{API}/v1/batches",
        headers={"Idempotency-Key": idempotency_key},
        json={
            "input_file_id": file_id,
            "endpoint": en.BATCH_URL,
            "completion_window": "24h",
            "metadata": {"job": job},
        },
    )
    return check(response, "create batch").json()


def poll_batch(client: httpx.Client, batch_id: str) -> dict[str, Any]:
    """GET /v1/batches/{id} with backoff until the status is final."""
    delay = POLL_START_SECONDS
    while True:
        try:
            response = client.get(f"{API}/v1/batches/{batch_id}")
            if response.status_code == 429:
                wait = float(response.headers.get("Retry-After", delay))
                print(f"  rate limited while polling; waiting {wait:.0f}s")
                time.sleep(wait)
                continue
            batch = check(response, "poll").json()
        except httpx.HTTPError as error:  # network blip; keep the overnight job alive
            print(f"  poll failed ({error}); retrying in {delay}s")
            time.sleep(delay)
            continue
        counts = batch["request_counts"]
        print(
            f"  {time.strftime('%H:%M:%S')} {batch['status']}: "
            f"{counts['completed']} done, {counts['failed']} failed, {counts['total']} total"
        )
        if batch["status"] in FINAL_STATUSES:
            return batch
        time.sleep(delay)
        delay = min(int(delay * 1.5), POLL_MAX_SECONDS)


def download_file(client: httpx.Client, file_id: str | None, path: str) -> list[dict]:
    """GET /v1/files/{id}/content to disk, then parse it. No id means no lines."""
    if not file_id:
        with open(path, "wb") as handle:
            handle.write(b"")
        return []
    response = client.get(f"{API}/v1/files/{file_id}/content")
    with open(path, "wb") as handle:
        handle.write(check(response, f"download {file_id}").content)
    return en.read_jsonl(path)


def load_state(path: str) -> dict[str, Any]:
    if os.path.exists(path):
        with open(path, encoding="utf-8") as handle:
            return json.load(handle)
    return {}


def save_state(path: str, state: dict[str, Any]) -> None:
    with open(path, "w", encoding="utf-8") as handle:
        json.dump(state, handle, indent=2)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--csv", default="companies.csv")
    parser.add_argument("--out", default="companies_enriched.csv")
    parser.add_argument("--requests", help="upload this JSONL (for example batch_retry.jsonl)")
    parser.add_argument("--job", default=f"company-enrichment-{time.strftime('%Y-%m-%d')}")
    parser.add_argument("--resume", action="store_true", help="reuse batch_state.json")
    args = parser.parse_args()

    api_key = os.environ.get("PERPLEXITY_API_KEY")
    if not api_key:
        print("Set PERPLEXITY_API_KEY first.", file=sys.stderr)
        return 2

    started = time.monotonic()
    rows = en.read_companies(args.csv)
    if args.requests:  # a retry run: upload an existing file, keep only its rows
        input_path = args.requests
        batch_lines = en.read_jsonl(input_path)
        wanted = {line["custom_id"] for line in batch_lines}
        rows = [row for row in rows if row["custom_id"] in wanted]
    else:
        input_path = "batch_input.jsonl"
        batch_lines = en.build_batch_lines(rows)
        en.write_jsonl(batch_lines, input_path)
    if not batch_lines:
        print("Nothing to send.")
        return 0
    with open(input_path, "rb") as handle:
        content_hash = hashlib.sha256(handle.read()).hexdigest()[:16]
    print(f"{input_path}: {len(batch_lines)} requests, model {batch_lines[0]['body']['model']}")

    state_path = "batch_state.json"
    state = load_state(state_path) if args.resume else {}
    if state and state.get("content_hash") != content_hash:
        print("batch_state.json is for a different input file; drop --resume.")
        return 2
    state["content_hash"] = content_hash

    headers = {"Authorization": f"Bearer {api_key}"}
    try:
        with httpx.Client(headers=headers, timeout=120.0) as client:
            if "file_id" not in state:
                state["file_id"] = upload_input(client, input_path)
                save_state(state_path, state)
            print(f"Input file: {state['file_id']}")

            if "batch_id" not in state:
                state.setdefault("job", args.job)
                state.setdefault("idempotency_key", f"{state['job']}-{state['file_id']}")
                save_state(state_path, state)  # persist before a create response can be lost
                batch = create_batch(client, state["file_id"], state["idempotency_key"], state["job"])
                state["batch_id"] = batch["id"]
                save_state(state_path, state)
            print(f"Batch: {state['batch_id']} (polling; safe to Ctrl-C and --resume)")

            batch = poll_batch(client, state["batch_id"])
            output_lines = download_file(client, batch.get("output_file_id"), "batch_output.jsonl")
            error_lines = download_file(client, batch.get("error_file_id"), "batch_errors.jsonl")
    except (BatchError, httpx.HTTPError) as error:
        print(f"Batch API failed: {error}", file=sys.stderr)
        return 1

    results = en.collect_results(rows, output_lines, error_lines)
    en.write_enriched_csv(rows, results, args.out)
    retries = en.retry_lines(batch_lines, results)
    en.write_jsonl(retries, "batch_retry.jsonl")

    summary = en.summarize(results)
    usage = batch.get("usage") or {}
    final_at = [batch.get(k) for k in ("completed_at", "expired_at", "cancelled_at", "failed_at")]
    ended_at = next((t for t in final_at if t), None)
    batch_minutes = (ended_at - batch["created_at"]) / 60 if ended_at else 0.0
    print(
        f"\nBatch {batch['id']} ended {batch['status']} after {batch_minutes:.1f} min "
        f"(this process waited {(time.monotonic() - started) / 60:.1f} min)"
    )
    if batch["status"] == "failed":
        print(f"  batch errors: {json.dumps(batch.get('errors'))}")
    print(
        f"  rows: {summary['total']}  completed: {summary['completed']}  "
        f"needs retry: {summary['needs_retry']}  by status: {summary['by_status']}"
    )
    print(
        f"  tokens (output file only): {usage.get('input_tokens', 0)} in, "
        f"{usage.get('output_tokens', 0)} out, {usage.get('total_tokens', 0)} total"
    )
    print(f"  reported cost: ${summary['cost_usd']:.4f} (includes ${summary['tool_calls_cost_usd']:.4f} for tools)")
    print(f"  files: {args.out}, batch_output.jsonl, batch_errors.jsonl, batch_retry.jsonl")
    for result in results.values():
        if not result.ok:
            print(f"  retry {result.custom_id}: {result.status}: {result.error}")
    return 1 if summary["needs_retry"] or batch["status"] != "completed" else 0


if __name__ == "__main__":
    sys.exit(main())

"""Offline tests for enrichment.py. No API key, no network."""

from __future__ import annotations

import csv
import json

import enrichment as en

ROWS = [
    {"name": "Apple", "website": "https://www.apple.com"},
    {"name": "Nestlé", "website": "https://www.nestle.com"},
    {"name": "Toyota", "website": "https://global.toyota"},
]
GOOD_JSON = {
    "headquarters_city": "Cupertino",
    "founded_year": 1976,
    "description": "Apple designs consumer electronics, software, and services.",
    "latest_news_headline": "Apple reports quarterly results",
    "latest_news_url": "https://example.com/apple-results",
}


def write_rows(tmp_path, rows=ROWS):
    path = tmp_path / "companies.csv"
    with open(path, "w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=["name", "website"])
        writer.writeheader()
        writer.writerows(rows)
    return str(path)


def output_line(custom_id, text, status="completed", cost=0.01):
    """Build an output-file line in the shape the Batch API guide documents."""
    return {
        "id": f"batch_req_{custom_id}",
        "custom_id": custom_id,
        "response": {
            "status_code": 200,
            "request_id": "01M4RFGZMBKRRZ4QYPDRCZ0CA0",
            "body": {
                "id": f"resp_{custom_id}",
                "object": "response",
                "status": status,
                "model": en.MODEL,
                "service_tier": "flex",
                "output": [
                    {"type": "search_results", "results": []},
                    {
                        "type": "message",
                        "role": "assistant",
                        "status": status,
                        "content": [{"type": "output_text", "text": text}],
                    },
                ],
                "usage": {"input_tokens": 10, "output_tokens": 5, "cost": {"total_cost": cost}},
            },
        },
        "error": None,
    }


# --- JSONL writer ---------------------------------------------------------


def test_batch_file_is_valid_jsonl(tmp_path):
    rows = en.read_companies(write_rows(tmp_path))
    lines = en.build_batch_lines(rows)
    path = tmp_path / "batch_input.jsonl"
    en.write_jsonl(lines, str(path))

    raw = path.read_bytes()
    assert raw.endswith(b"\n")
    text_lines = raw.decode("utf-8").split("\n")
    assert text_lines[-1] == ""  # exactly one trailing newline
    assert all(line.strip() for line in text_lines[:-1]), "no blank lines"
    assert "\r" not in raw.decode("utf-8")

    parsed = [json.loads(line) for line in text_lines[:-1]]
    assert len(parsed) == len(ROWS)
    ids = [line["custom_id"] for line in parsed]
    assert len(set(ids)) == len(ids), "custom_id must be unique"
    assert all(1 <= len(cid.encode("utf-8")) <= 64 for cid in ids)
    assert {line["method"] for line in parsed} == {"POST"}
    assert {line["url"] for line in parsed} == {"/v1/responses"}
    assert {line["body"]["model"] for line in parsed} == {en.MODEL}, "one model per file"
    assert parsed[0]["body"]["tools"] == [{"type": "web_search"}]
    assert parsed[0]["body"]["response_format"]["type"] == "json_schema"


def test_custom_id_is_safe_for_unicode_and_long_names():
    long_name = "Ü" * 200
    cid = en.make_custom_id(7, long_name)
    assert cid.startswith("row-007-")
    assert len(cid.encode("utf-8")) <= 64
    assert en.make_custom_id(2, "Nestlé") == "row-002-nestl"
    assert en.make_custom_id(3, "!!!") == "row-003-row"


# --- output and error line parsing -----------------------------------------


def test_parse_completed_output_line():
    line = output_line("row-001-apple", json.dumps(GOOD_JSON), cost=0.03204)
    line["response"]["body"]["usage"]["cost"]["tool_calls_cost"] = 0.005
    result = en.parse_output_line(line)
    assert result.ok
    assert result.fields == GOOD_JSON
    assert result.response_id == "resp_row-001-apple"
    assert result.cost_usd == 0.03204  # the captured total already includes tools
    assert result.tool_calls_cost_usd == 0.005
    summary = en.summarize({result.custom_id: result})
    assert summary["cost_usd"] == 0.032 and summary["tool_calls_cost_usd"] == 0.005


def test_parse_tolerates_markdown_fences_and_rejects_prose():
    fence = "`" * 3
    fenced = fence + "json\n" + json.dumps(GOOD_JSON) + "\n" + fence
    assert en.parse_output_line(output_line("a", fenced)).ok
    prose = en.parse_output_line(output_line("b", "Apple is based in Cupertino."))
    assert prose.status == "parse_error"
    partial = en.parse_output_line(output_line("c", json.dumps({"founded_year": 1976})))
    assert partial.status == "parse_error"


def test_incomplete_response_is_not_ok():
    result = en.parse_output_line(output_line("a", "", status="incomplete"))
    assert not result.ok
    assert result.status == "incomplete"


def test_schema_rejects_wrong_types_and_extra_fields():
    for changes in ({"founded_year": True}, {"founded_year": "1976"}, {"description": []}, {"extra": "x"}):
        assert en.parse_enrichment(json.dumps({**GOOD_JSON, **changes})) is None


def test_existing_enrichment_columns_are_not_duplicated(tmp_path):
    rows = [{"name": "Apple", "custom_id": "a", "headquarters_city": ""}]
    results = {"a": en.Result("a", "completed", fields=GOOD_JSON)}
    path = tmp_path / "enriched.csv"
    en.write_enriched_csv(rows, results, str(path))
    with path.open() as handle:
        columns = next(csv.reader(handle))
    assert len(columns) == len(set(columns))


def test_parse_error_line_api_rejection():
    # Shape 1 from the guide: error is null, response holds the status and body.
    line = {
        "id": "batch_req_01M4RFRF95XTJN5HBN1ASJKVMX",
        "custom_id": "request-7",
        "response": {
            "status_code": 400,
            "request_id": "01M4RFRF95XTJN5HBN1ASJKVMX",
            "body": {
                "error": {
                    "message": 'validation failed: model "openai/gpt-nope-9" is not supported',
                    "type": "invalid_request",
                    "code": 400,
                }
            },
        },
        "error": None,
    }
    result = en.parse_error_line(line)
    assert result.status == "error"
    assert result.error.startswith("HTTP 400: validation failed")


def test_parse_error_line_stopped_run():
    # Shape 2 from the guide: response is null, error has a code and message.
    line = {
        "id": "batch_req_01M4RFJC22KH0N346H7E8XH1DY",
        "custom_id": "cancel-1",
        "response": None,
        "error": {
            "code": "batch_cancelled",
            "message": "The batch was cancelled before this request finished.",
        },
    }
    result = en.parse_error_line(line)
    assert result.status == "error"
    assert result.error == (
        "batch_cancelled: The batch was cancelled before this request finished."
    )


# --- join, retry file, summary ----------------------------------------------


def test_join_keeps_every_row_and_marks_missing(tmp_path):
    rows = en.read_companies(write_rows(tmp_path))
    apple, nestle, toyota = (row["custom_id"] for row in rows)
    outputs = [
        output_line(toyota, json.dumps(GOOD_JSON)),  # order differs from the input
        output_line(apple, json.dumps(GOOD_JSON)),
    ]
    errors = [
        {"custom_id": nestle, "response": None, "error": {"code": "batch_expired", "message": "x"}}
    ]
    results = en.collect_results(rows, outputs, errors)
    assert results[apple].ok and results[toyota].ok
    assert results[nestle].status == "error"

    # Drop Nestlé's error line: the row must still appear, marked missing.
    results = en.collect_results(rows, outputs, [])
    assert results[nestle].status == "missing"

    path = tmp_path / "companies_enriched.csv"
    en.write_enriched_csv(rows, results, str(path))
    with open(path, newline="", encoding="utf-8") as handle:
        written = list(csv.DictReader(handle))
    assert [row["name"] for row in written] == ["Apple", "Nestlé", "Toyota"]
    assert written[0]["headquarters_city"] == "Cupertino"
    assert written[0]["status"] == "completed"
    assert written[1]["status"] == "missing"
    assert written[1]["headquarters_city"] == ""
    assert written[2]["founded_year"] == "1976"


def test_retry_file_contains_only_failed_rows(tmp_path):
    rows = en.read_companies(write_rows(tmp_path))
    batch_lines = en.build_batch_lines(rows)
    apple, nestle, toyota = (row["custom_id"] for row in rows)
    outputs = [
        output_line(apple, json.dumps(GOOD_JSON)),
        output_line(toyota, "not json"),
    ]
    results = en.collect_results(rows, outputs, [])
    retries = en.retry_lines(batch_lines, results)
    assert [line["custom_id"] for line in retries] == [nestle, toyota]
    assert retries[0] == batch_lines[1]  # the exact original request, ready to upload

    path = tmp_path / "batch_retry.jsonl"
    en.write_jsonl(retries, str(path))
    assert en.read_jsonl(str(path)) == retries

    summary = en.summarize(results)
    assert summary == {
        "total": 3,
        "completed": 1,
        "needs_retry": 2,
        "by_status": {"completed": 1, "missing": 1, "parse_error": 1},
        "cost_usd": 0.02,
        "tool_calls_cost_usd": 0.0,
    }

"""Offline lifecycle checks: create, resume, download, join, and partial failure."""

import json

import httpx
import pytest
import run_batch as rb
from test_enrichment import GOOD_JSON, output_line


@pytest.mark.parametrize("complete", [True, False])
def test_runner_preserves_idempotency_after_lost_create_response(
    tmp_path, monkeypatch, complete
):
    monkeypatch.chdir(tmp_path)
    (tmp_path / "companies.csv").write_text(
        "name,website\nApple,https://www.apple.com\n"
    )
    monkeypatch.setenv("PERPLEXITY_API_KEY", "unused-test-key")
    monkeypatch.setattr("sys.argv", ["run_batch.py", "--job", "first-day"])
    calls = []
    create_bodies = []

    def handler(request):
        calls.append((request.method, request.url.path))
        if request.url.path == "/v1/files":
            return httpx.Response(200, json={"id": "file-1"})
        if request.url.path == "/v1/batches":
            create_bodies.append(
                (request.headers["Idempotency-Key"], json.loads(request.content))
            )
            if len(create_bodies) == 1:
                raise httpx.ReadTimeout("create response lost", request=request)
            return httpx.Response(200, json={"id": "batch-1"})
        if request.url.path == "/v1/batches/batch-1":
            return httpx.Response(
                200,
                json={
                    "id": "batch-1",
                    "status": "completed",
                    "created_at": 100,
                    "completed_at": 160,
                    "request_counts": {"total": 1, "completed": 1, "failed": 0},
                    "output_file_id": "output-1",
                    "error_file_id": None,
                },
            )
        if request.url.path == "/v1/files/output-1/content":
            text = json.dumps(GOOD_JSON) if complete else "not JSON"
            return httpx.Response(
                200, text=json.dumps(output_line("row-001-apple", text)) + "\n"
            )
        raise AssertionError(f"unexpected request: {request.method} {request.url.path}")

    real_client = httpx.Client
    monkeypatch.setattr(
        rb.httpx,
        "Client",
        lambda **kwargs: real_client(transport=httpx.MockTransport(handler), **kwargs),
    )
    assert rb.main() == 1
    saved = json.loads((tmp_path / "batch_state.json").read_text())
    assert saved["idempotency_key"] == "first-day-file-1"
    monkeypatch.setattr("sys.argv", ["run_batch.py", "--resume", "--job", "next-day"])
    assert rb.main() == (0 if complete else 1)
    assert create_bodies[0] == create_bodies[1]
    assert calls.count(("POST", "/v1/files")) == 1
    assert (tmp_path / "companies_enriched.csv").exists()
    assert (tmp_path / "batch_errors.jsonl").read_text() == ""
    retry = (tmp_path / "batch_retry.jsonl").read_text()
    assert bool(retry.strip()) is not complete