parser_client.py

from __future__ import annotations

import base64
import os
import time
from dataclasses import dataclass

import requests

from parser_output import shape_parse_response

DEFAULT_ENVIRONMENT = "staging"
ENVIRONMENT_CONFIG = {
    "production": ("COHERE_API_KEY", "https://api.cohere.com/v2/parse"),
    "staging": ("COHERE_STG_API_KEY", "https://stg.api.cohere.ai/v2/parse"),
}
DEFAULT_MODEL = "parse-v5.0"
RETRYABLE_STATUSES = {429, 500, 502, 503, 504}


class ParserServiceError(RuntimeError):
    """A safe parser-service error that may be displayed to a user."""


def selected_environment() -> str:
    environment = os.environ.get("COHERE_ENV", DEFAULT_ENVIRONMENT).strip().lower()
    if environment not in ENVIRONMENT_CONFIG:
        raise ParserServiceError(
            "COHERE_ENV must be either 'staging' or 'production'."
        )
    return environment


def api_key_name_for_environment(environment: str | None = None) -> str:
    selected = environment or selected_environment()
    try:
        return ENVIRONMENT_CONFIG[selected][0]
    except KeyError as error:
        raise ParserServiceError(
            "COHERE_ENV must be either 'staging' or 'production'."
        ) from error


@dataclass(frozen=True)
class ParserConfig:
    api_url: str
    api_key: str
    model: str
    timeout: float
    max_retries: int = 4

    @classmethod
    def from_environment(cls) -> "ParserConfig":
        environment = selected_environment()
        api_key_name, default_api_url = ENVIRONMENT_CONFIG[environment]
        api_key = os.environ.get(api_key_name, "").strip()
        if not api_key:
            raise ParserServiceError(
                f"Cohere Parse is not configured for {environment}."
            )
        return cls(
            api_url=os.environ.get("COHERE_API_URL", "").strip()
            or default_api_url,
            api_key=api_key,
            model=os.environ.get("COHERE_MODEL", DEFAULT_MODEL).strip(),
            timeout=float(os.environ.get("COHERE_TIMEOUT", "300")),
        )


def build_payload(model: str, png: bytes) -> dict:
    encoded = base64.b64encode(png).decode("ascii")
    return {
        "model": model,
        "document": {
            "type": "image_url",
            "image_url": f"data:image/png;base64,{encoded}",
        },
        "output_format": "blocks",
    }


class ParserClient:
    def __init__(
        self,
        config: ParserConfig | None = None,
        *,
        session: requests.Session | None = None,
        sleeper=time.sleep,
    ) -> None:
        self.config = config or ParserConfig.from_environment()
        self.session = session or requests.Session()
        self.sleeper = sleeper

    def _retry_delay(self, response: requests.Response, attempt: int) -> float:
        header = response.headers.get("Retry-After", "")
        try:
            return min(60.0, max(0.0, float(header)))
        except ValueError:
            return min(60.0, 2.0 ** (attempt + 1))

    def parse_page(self, png: bytes) -> dict:
        started = time.perf_counter()
        headers = {
            "accept": "application/json",
            "content-type": "application/json",
            "Authorization": f"Bearer {self.config.api_key}",
        }
        payload = build_payload(self.config.model, png)

        for attempt in range(self.config.max_retries):
            try:
                response = self.session.post(
                    self.config.api_url,
                    headers=headers,
                    json=payload,
                    timeout=self.config.timeout,
                )
            except (requests.Timeout, requests.ConnectionError) as exc:
                if attempt + 1 < self.config.max_retries:
                    self.sleeper(min(60.0, 2.0 ** (attempt + 1)))
                    continue
                raise ParserServiceError(
                    "Cohere Parse could not be reached. Try again shortly."
                ) from exc

            request_id = (
                response.headers.get("x-request-id")
                or response.headers.get("request-id")
            )
            if response.status_code in RETRYABLE_STATUSES:
                if attempt + 1 < self.config.max_retries:
                    self.sleeper(self._retry_delay(response, attempt))
                    continue
                raise ParserServiceError(
                    f"Cohere Parse is temporarily unavailable ({response.status_code})."
                )
            if not response.ok:
                raise ParserServiceError(
                    f"Cohere Parse rejected the request ({response.status_code})."
                )

            try:
                body = response.json()
            except requests.JSONDecodeError as exc:
                raise ParserServiceError(
                    "Cohere Parse returned an invalid response."
                ) from exc

            try:
                parsed = shape_parse_response(body)
            except (TypeError, ValueError) as exc:
                raise ParserServiceError(
                    "Cohere Parse returned an invalid response."
                ) from exc

            return {
                **parsed,
                "elapsed_ms": round((time.perf_counter() - started) * 1000),
                "model": self.config.model,
                "usage": body.get("meta"),
                "request_id": request_id or body.get("id"),
            }

        raise ParserServiceError("Cohere Parse is temporarily unavailable.")