responsible-request 0.4.0__py3-none-any.whl

responsible_request/helpers/batch.py

@@ -0,0 +1,127 @@
+"""Running many requests and calibrating the latency baseline."""
+
+from __future__ import annotations
+
+import asyncio
+import statistics
+import uuid
+from collections.abc import Awaitable, Callable, Iterable, Sequence
+from typing import Any
+
+import openai
+
+from ..client import get_throttle
+from ..config import ThrottleConfig
+from ..cost import BudgetExceeded
+from ..estimator import LatencyEstimator, metric_value
+from ..records import tags
+
+Call = Callable[..., Awaitable[Any]]
+
+
+async def run_batch(
+    client: openai.AsyncOpenAI,
+    requests: Iterable[dict[str, Any]],
+    *,
+    call: Call | None = None,
+    progress: bool = False,
+    return_exceptions: bool = True,
+    batch_id: str | None = None,
+) -> list[Any]:
+    """Run many requests concurrently and return their results in input order.
+
+    The throttle paces the requests, so all of them can be submitted at once. Each request
+    record is tagged with ``batch_id`` and ``item_index`` for joining logs with results.
+    Once a cost budget is exhausted, the remaining requests fail fast with
+    :class:`responsible_request.BudgetExceeded`.
+
+    :param requests: Keyword arguments for ``call``, one dict per request.
+    :param call: Coroutine function to call, defaults to ``client.chat.completions.create``.
+    :param progress: Show a tqdm progress bar (requires the ``progress`` extra).
+    :param return_exceptions: Put exceptions into the result list instead of raising the first.
+    :param batch_id: Tag value, a random id by default.
+    """
+    call = call or client.chat.completions.create
+    items: Sequence[dict[str, Any]] = list(requests)
+    batch_id = batch_id or uuid.uuid4().hex[:12]
+    bar = _progress_bar(len(items)) if progress else None
+
+    async def one(index: int, kwargs: dict[str, Any]) -> Any:
+        try:
+            with tags(batch_id=batch_id, item_index=index):
+                return await call(**kwargs)
+        except openai.APIConnectionError as exc:
+            # older SDKs wrap exceptions raised by the transport
+            if isinstance(exc.__cause__, BudgetExceeded):
+                raise exc.__cause__ from None
+            raise
+        finally:
+            if bar is not None:
+                bar.update(1)
+
+    try:
+        return await asyncio.gather(
+            *(one(i, kw) for i, kw in enumerate(items)), return_exceptions=return_exceptions
+        )
+    finally:
+        if bar is not None:
+            bar.close()
+
+
+def run_batch_sync(
+    client: openai.AsyncOpenAI, requests: Iterable[dict[str, Any]], **kwargs: Any
+) -> list[Any]:
+    """Blocking version of :func:`run_batch` for scripts (runs its own event loop)."""
+
+    async def main() -> list[Any]:
+        try:
+            return await run_batch(client, requests, **kwargs)
+        finally:
+            await client.close()
+
+    return asyncio.run(main())
+
+
+def _progress_bar(total: int) -> Any:
+    try:
+        from tqdm.auto import tqdm
+    except ImportError as exc:  # pragma: no cover
+        raise ImportError(
+            "progress=True requires tqdm: pip install responsible-request[progress]"
+        ) from exc
+    return tqdm(total=total)
+
+
+async def calibrate(
+    client: openai.AsyncOpenAI,
+    model: str,
+    *,
+    n: int = 5,
+    messages: list[Any] | None = None,
+    **kwargs: Any,
+) -> float:
+    """Measure the latency baseline of ``model`` with ``n`` sequential requests and pin it.
+
+    Run this while the endpoint is idle (e.g. at night) or use a representative request of your
+    experiment via ``messages``/``kwargs`` so that the baseline matches your workload. Returns
+    the pinned baseline (in units of the configured metric).
+    """
+    throttle = get_throttle(client)
+    config: ThrottleConfig = throttle.config
+    messages = messages or [{"role": "user", "content": "Reply with the single word: OK"}]
+    values: list[float] = []
+    for _ in range(n):
+        with tags(calibration=True):
+            await client.chat.completions.create(model=model, messages=messages, **kwargs)
+        # Use the transport's own measurement, which excludes time spent in the rate limiter.
+        records = [lane.last_record for lane in throttle.find(model) if lane.last_record]
+        value = metric_value(records[-1], config) if records else None
+        if value is not None:
+            values.append(value)
+    if not values:
+        raise RuntimeError("calibration produced no usable measurements")
+    baseline = statistics.median(values)
+    for lane in throttle.find(model):
+        if isinstance(lane.estimator, LatencyEstimator):
+            lane.estimator.pin_baseline(baseline)
+    return baseline

responsible_request/helpers/messages.py

@@ -0,0 +1,18 @@
+"""Message builders."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def cached_system_message(content: str) -> dict[str, Any]:
+    """A system message marked for OpenRouter's prompt cache (``cache_control: ephemeral``).
+
+    The provider caches the prompt prefix up to this message, so repeated requests with the same
+    system prompt pay less for input tokens; the model still runs. This is unrelated to the
+    response cache (``cache=True``), which replays stored responses without sending a request.
+    """
+    return {
+        "role": "system",
+        "content": [{"type": "text", "text": content, "cache_control": {"type": "ephemeral"}}],
+    }

responsible_request/helpers/params.py

@@ -0,0 +1,18 @@
+"""Sampling-parameter presets."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def reproducible(seed: int = 0, **overrides: Any) -> dict[str, Any]:
+    """Parameters that make chat completions as repeatable as the backend allows.
+
+    Greedy decoding (``temperature=0``, ``top_p=1``) plus a fixed ``seed``. Spread it into a call
+    (``create(..., **rr.reproducible())``) or set it for every request with
+    ``rr.AsyncOpenAI(default_params=rr.reproducible())``.
+
+    Note: outputs can still differ between runs, because batched GPU inference is not bit-exact
+    (the batch composition depends on the other requests the server is processing).
+    """
+    return {"temperature": 0.0, "top_p": 1.0, "seed": seed, **overrides}

responsible_request/helpers/structured.py

@@ -0,0 +1,150 @@
+"""Structured output: request JSON matching a Pydantic model and validate the answer."""
+
+from __future__ import annotations
+
+import copy
+import json
+import re
+from collections.abc import Iterable
+from typing import Any, TypeVar
+
+import openai
+from pydantic import BaseModel, ValidationError
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class StructuredOutputError(Exception):
+    """The model did not produce valid output within the allowed number of attempts."""
+
+    def __init__(self, message: str, *, last_content: str | None, errors: list[Exception]) -> None:
+        super().__init__(message)
+        self.last_content = last_content
+        self.errors = errors
+
+
+def _strictify(node: Any) -> None:
+    if isinstance(node, dict):
+        if node.get("type") == "object" and "properties" in node:
+            node["additionalProperties"] = False
+            node["required"] = list(node["properties"])
+        node.pop("default", None)
+        for value in node.values():
+            _strictify(value)
+    elif isinstance(node, list):
+        for item in node:
+            _strictify(item)
+
+
+def response_format_from_model(model: type[BaseModel], *, strict: bool = True) -> dict[str, Any]:
+    """Build a ``json_schema`` ``response_format`` from a Pydantic model.
+
+    With ``strict=True`` every object forbids additional properties and marks all properties as
+    required, as strict structured-output backends demand.
+    """
+    schema = copy.deepcopy(model.model_json_schema())
+    if strict:
+        _strictify(schema)
+    return {
+        "type": "json_schema",
+        "json_schema": {"name": model.__name__, "schema": schema, "strict": strict},
+    }
+
+
+_FENCE = re.compile(r"^\s*```(?:json)?\s*(.*?)\s*```\s*$", re.DOTALL)
+
+
+def extract_json(text: str) -> str:
+    """Strip Markdown code fences and surrounding prose around a JSON object."""
+    match = _FENCE.match(text)
+    if match:
+        return match.group(1)
+    start, end = text.find("{"), text.rfind("}")
+    if start != -1 and end > start and not text.lstrip().startswith("{"):
+        return text[start : end + 1]
+    return text
+
+
+def _answer_text(message: Any) -> str:
+    """The answer text: ``content``, else ``reasoning``, else ``reasoning_content`` (some
+    reasoning models put the JSON there and leave ``content`` empty)."""
+    for key in ("content", "reasoning", "reasoning_content"):
+        value = getattr(message, key, None)  # the SDK's models keep unknown fields as attributes
+        if isinstance(value, str) and value.strip():
+            return value
+    return ""
+
+
+async def structured(
+    client: openai.AsyncOpenAI,
+    *,
+    model: str,
+    messages: Iterable[Any],
+    schema: type[T],
+    retries: int = 2,
+    strict: bool = True,
+    fallback: bool = True,
+    **kwargs: Any,
+) -> T:
+    """Ask for JSON matching ``schema`` and return it as a validated ``schema`` instance.
+
+    If the answer does not validate, the model is shown the validation error and asked again
+    (up to ``retries`` times). If the backend rejects ``json_schema`` response formats and
+    ``fallback`` is set, it switches to ``json_object`` mode with the schema in the prompt.
+    If the message ``content`` is empty, the answer is read from ``reasoning``, then from
+    ``reasoning_content`` (some reasoning models put the JSON there).
+
+    Extra keyword arguments are passed to ``client.chat.completions.create``.
+    """
+    msgs: list[Any] = list(messages)
+    response_format: dict[str, Any] = response_format_from_model(schema, strict=strict)
+    errors: list[Exception] = []
+    content: str | None = None
+
+    attempt = 0
+    while attempt <= retries:
+        try:
+            params: dict[str, Any] = {
+                "model": model,
+                "messages": msgs,
+                "response_format": response_format,
+                **kwargs,
+            }
+            response = await client.chat.completions.create(**params)
+        except openai.BadRequestError as exc:
+            if not fallback or response_format["type"] != "json_schema":
+                raise
+            errors.append(exc)
+            schema_text = json.dumps(response_format["json_schema"]["schema"])
+            msgs = [
+                {
+                    "role": "system",
+                    "content": f"Respond only with a JSON object matching this schema:\n"
+                    f"{schema_text}",
+                },
+                *msgs,
+            ]
+            response_format = {"type": "json_object"}
+            continue  # the fallback does not count as a retry
+
+        content = _answer_text(response.choices[0].message)
+        try:
+            return schema.model_validate_json(extract_json(content))
+        except ValidationError as exc:
+            errors.append(exc)
+            msgs = [
+                *msgs,
+                {"role": "assistant", "content": content},
+                {
+                    "role": "user",
+                    "content": "Your answer did not match the required JSON schema:\n"
+                    f"{exc}\nRespond again with only the corrected JSON object.",
+                },
+            ]
+        attempt += 1
+
+    raise StructuredOutputError(
+        f"no valid {schema.__name__} after {retries + 1} attempts",
+        last_content=content,
+        errors=errors,
+    )

responsible_request/limiter.py

@@ -0,0 +1,165 @@
+"""Evenly spaced rate limiting whose rate can change at runtime, plus a concurrency cap.
+
+The limiter is safe to share between threads that each run their own event loop (e.g. several
+threads calling ``asyncio.run``): its state is guarded by a ``threading.Lock``, and every waiter
+sleeps in its own loop and is woken through ``call_soon_threadsafe``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+import time
+from collections import deque
+
+
+class _Waiter:
+    """A future in the caller's event loop that another thread can resolve."""
+
+    __slots__ = ("future", "loop")
+
+    def __init__(self) -> None:
+        self.loop = asyncio.get_running_loop()
+        self.future: asyncio.Future[None] = self.loop.create_future()
+
+    def renew(self) -> None:
+        self.future = self.loop.create_future()
+
+    def wake(self) -> bool:
+        """Resolve the current future. Returns False if the waiter's loop is closed."""
+        try:
+            self.loop.call_soon_threadsafe(_resolve, self.future)
+        except RuntimeError:
+            return False
+        return True
+
+
+def _resolve(future: asyncio.Future[None]) -> None:
+    if not future.done():
+        future.set_result(None)
+
+
+class AdaptiveLimiter:
+    """Paces acquisitions at ``rpm`` and caps the number of concurrent holders.
+
+    Use as ``async with limiter: ...``. The concurrency slot is held for the duration of the
+    block. Requests leave the limiter at least ``60 / rpm`` seconds apart (no bursts), in the
+    order in which they got a concurrency slot. ``set_rpm`` applies to requests that are already
+    waiting, and ``pause`` holds back all requests for a while.
+    """
+
+    def __init__(self, rpm: float, max_concurrency: int) -> None:
+        self._lock = threading.Lock()
+        self._rpm = rpm
+        self._interval = 60 / rpm
+        self.max_concurrency = max_concurrency
+        self.in_flight = 0
+        self._slots_held = 0
+        self._last_sent = -float("inf")
+        self._paused_until = -float("inf")
+        self._slot_waiters: deque[_Waiter] = deque()
+        self._rate_waiters: deque[_Waiter] = deque()
+
+    @property
+    def rpm(self) -> float:
+        return self._rpm
+
+    @property
+    def paused_s(self) -> float:
+        """Remaining time of the current pause (0 if not paused)."""
+        return max(0.0, self._paused_until - time.monotonic())
+
+    def set_rpm(self, rpm: float) -> None:
+        with self._lock:
+            self._rpm = rpm
+            self._interval = 60 / rpm
+            self._wake_rate_head()  # it re-computes its wake-up time for the new rate
+
+    def pause(self, seconds: float) -> None:
+        """Send nothing for ``seconds`` (extends, never shortens, a running pause)."""
+        with self._lock:
+            self._paused_until = max(self._paused_until, time.monotonic() + seconds)
+
+    async def acquire(self) -> None:
+        # Concurrency first, then rate: requests leave the limiter evenly spaced.
+        await self._acquire_slot()
+        try:
+            await self._acquire_rate()
+        except BaseException:
+            self._release_slot()
+            raise
+
+    def release(self) -> None:
+        with self._lock:
+            self.in_flight -= 1
+        self._release_slot()
+
+    async def __aenter__(self) -> None:
+        await self.acquire()
+
+    async def __aexit__(self, *exc: object) -> None:
+        self.release()
+
+    # ------------------------------------------------------------------ internals
+
+    async def _acquire_slot(self) -> None:
+        with self._lock:
+            if self._slots_held < self.max_concurrency and not self._slot_waiters:
+                self._slots_held += 1
+                return
+            waiter = _Waiter()
+            self._slot_waiters.append(waiter)
+        try:
+            await waiter.future  # resolved by _release_slot, which hands over its slot
+        except BaseException:
+            with self._lock:
+                granted = waiter not in self._slot_waiters
+                if not granted:
+                    self._slot_waiters.remove(waiter)
+            if granted:  # the slot was already handed to us: pass it on
+                self._release_slot()
+            raise
+
+    def _release_slot(self) -> None:
+        with self._lock:
+            while self._slot_waiters:
+                if self._slot_waiters.popleft().wake():
+                    return  # the slot moves to the woken waiter
+            self._slots_held -= 1
+
+    async def _acquire_rate(self) -> None:
+        waiter = _Waiter()
+        with self._lock:
+            self._rate_waiters.append(waiter)
+        try:
+            while True:
+                with self._lock:
+                    waiter.renew()  # wake-ups from now on resolve the new future
+                    timeout: float | None = None
+                    if self._rate_waiters[0] is waiter:
+                        now = time.monotonic()
+                        ready = max(self._last_sent + self._interval, self._paused_until)
+                        if now >= ready:
+                            self._last_sent = now
+                            self._rate_waiters.popleft()
+                            self._wake_rate_head()
+                            self.in_flight += 1
+                            return
+                        timeout = ready - now
+                # Sleep in this thread's own loop until the slot is due or we are woken.
+                await asyncio.wait({waiter.future}, timeout=timeout)
+        except BaseException:
+            with self._lock:
+                if waiter in self._rate_waiters:
+                    was_head = self._rate_waiters[0] is waiter
+                    self._rate_waiters.remove(waiter)
+                    if was_head:
+                        self._wake_rate_head()
+            raise
+
+    def _wake_rate_head(self) -> None:
+        """Wake the first waiter for the rate slot (call with the lock held)."""
+        while self._rate_waiters:
+            if self._rate_waiters[0].wake():
+                return
+            self._rate_waiters.popleft()  # its loop is gone

responsible_request/logfiles.py

@@ -0,0 +1,120 @@
+"""The files of a JSONL log: the active file plus the segments loguru rotated away from it,
+optionally compressed."""
+
+from __future__ import annotations
+
+import bz2
+import gzip
+import io
+import lzma
+import os
+import re
+import shutil
+from collections.abc import Callable
+from datetime import datetime
+from pathlib import Path
+from typing import IO, Any
+
+COMPRESSIONS = ("gz", "bz2", "xz", "zst")
+"""Values of ``LogConfig.compression``: formats that are read back as a single JSONL stream."""
+
+# loguru names a rotated segment ``{stem}.{%Y-%m-%d_%H-%M-%S_%f}[.{n}]{suffix}``
+_DATE = r"\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2}_\d{6}(?:\.\d+)?"
+
+
+def _zstd() -> Any:
+    try:
+        from compression import zstd  # Python 3.14+
+
+        return zstd
+    except ImportError:
+        pass
+    try:
+        import zstandard
+
+        return zstandard
+    except ImportError:
+        raise ImportError(
+            "zstd needs Python 3.14+ or: pip install responsible-request[zstd]"
+        ) from None
+
+
+def compression_of(path: Path) -> str | None:
+    """The compression format of ``path`` from its extension, or None for a plain file."""
+    ext = path.suffix.lstrip(".")
+    return ext if ext in COMPRESSIONS else None
+
+
+def strip_compression(path: Path) -> Path:
+    """``requests.jsonl.gz`` -> ``requests.jsonl``."""
+    return path.with_suffix("") if compression_of(path) else path
+
+
+def open_binary(path: Path) -> IO[bytes]:
+    """Open a plain or compressed log file for reading (decompressed bytes)."""
+    fmt = compression_of(path)
+    if fmt == "zst":  # zstandard's reader can't iterate over lines, a buffer on top of it can
+        return io.BufferedReader(_zstd().open(path, "rb"))
+    opener: Any = {"gz": gzip.open, "bz2": bz2.open, "xz": lzma.open}.get(fmt or "")
+    stream: IO[bytes] = opener(path, "rb") if opener else path.open("rb")
+    return stream
+
+
+def rotated_segments(path: Path) -> list[Path]:
+    """Segments rotated away from the log at ``path`` (and compressed copies of ``path`` itself,
+    which loguru writes when a sink without rotation is removed), oldest first."""
+    pattern = re.compile(
+        rf"{re.escape(path.stem)}(?:\.{_DATE})?{re.escape(path.suffix)}"
+        rf"(?:\.(?:{'|'.join(COMPRESSIONS)}))?"
+    )
+    try:
+        names = {p.name for p in path.parent.iterdir()}
+    except FileNotFoundError:
+        return []
+    found = []
+    for name in names:
+        if name == path.name or not pattern.fullmatch(name):
+            continue
+        plain = name.rsplit(".", 1)[0]
+        if compression_of(Path(name)) and plain != path.name and plain in names:
+            continue  # still being compressed: loguru removes the plain file once it is done
+        found.append(path.parent / name)
+    return sorted(found, key=lambda p: (_mtime(p), p.name))
+
+
+def _mtime(path: Path) -> float:
+    try:
+        return path.stat().st_mtime
+    except FileNotFoundError:
+        return 0.0
+
+
+def log_segments(path: Path) -> list[Path]:
+    """All files of the log at ``path``, oldest first. A compressed ``path`` is read on its own."""
+    if compression_of(path):
+        return [path]
+    return rotated_segments(path) + ([path] if path.exists() else [])
+
+
+def _compress_zstd(path_in: str) -> None:
+    """Compress a finished segment to ``{path_in}.zst``, the way loguru does for other formats."""
+    path_out = Path(f"{path_in}.zst")
+    if path_out.exists():  # keep the earlier archive under a dated name
+        root, ext = os.path.splitext(path_in)
+        date = datetime.fromtimestamp(path_out.stat().st_mtime).strftime("%Y-%m-%d_%H-%M-%S_%f")
+        renamed, n = Path(f"{root}.{date}{ext}.zst"), 1
+        while renamed.exists():
+            n += 1
+            renamed = Path(f"{root}.{date}.{n}{ext}.zst")
+        path_out.rename(renamed)
+    with open(path_in, "rb") as src, _zstd().open(path_out, "wb") as dst:
+        shutil.copyfileobj(src, dst)
+    os.remove(path_in)
+
+
+def loguru_compression(fmt: str | None) -> str | Callable[[str], None] | None:
+    """The ``compression`` argument for loguru's ``logger.add``."""
+    if fmt == "zst":
+        _zstd()  # fail now rather than in the writer thread at the first rotation
+        return _compress_zstd
+    return fmt