diff --git a/evaluations/evaluations/datasets/frames.py b/evaluations/evaluations/datasets/frames.py index f1860bf7..64371b26 100644 --- a/evaluations/evaluations/datasets/frames.py +++ b/evaluations/evaluations/datasets/frames.py @@ -10,6 +10,7 @@ import ast import json import logging import re +import time from collections.abc import Mapping from datetime import UTC, datetime from pathlib import Path @@ -27,6 +28,9 @@ from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator logger = logging.getLogger(__name__) USER_AGENT = "haiku.rag-evaluations (https://github.com/ggozad/haiku.rag)" +FETCH_ATTEMPTS = 3 +THROTTLE_SECONDS = 1.0 +RATE_LIMIT_BACKOFF_SECONDS = 60.0 def load_frames_test() -> Dataset: @@ -153,6 +157,13 @@ def _fetch_article_page( return response.text, "html", revid, title +def _backoff_seconds(error: Exception, attempt: int) -> float: + if isinstance(error, httpx.HTTPStatusError) and error.response.status_code == 429: + retry_after = error.response.headers.get("retry-after") + return float(retry_after) if retry_after else RATE_LIMIT_BACKOFF_SECONDS + return 5.0 * attempt + + def fetch_article( uri: str, cache_dir: Path, client: httpx.Client | None ) -> dict[str, Any] | None: @@ -170,17 +181,24 @@ def fetch_article( return row assert client is not None - try: - title = unquote(urlsplit(uri).path[len("/wiki/") :]) - if title.startswith("Category:"): - content, format, revid = _fetch_category_page( - urlsplit(uri).netloc, title, client - ) - else: - content, format, revid, title = _fetch_article_page(uri, client) - except Exception as e: - logger.warning(f"Failed to fetch {uri}: {e}") - return None + title = unquote(urlsplit(uri).path[len("/wiki/") :]) + # Wikimedia throttles sustained bot traffic; pace uncached fetches. + time.sleep(THROTTLE_SECONDS) + for attempt in range(1, FETCH_ATTEMPTS + 1): + try: + if title.startswith("Category:"): + content, format, revid = _fetch_category_page( + urlsplit(uri).netloc, title, client + ) + else: + content, format, revid, title = _fetch_article_page(uri, client) + break + except Exception as e: + if attempt == FETCH_ATTEMPTS: + logger.warning(f"Failed to fetch {uri}: {e}") + return None + logger.info(f"Retrying {uri} after error: {e}") + time.sleep(_backoff_seconds(e, attempt)) row: dict[str, Any] = { "uri": uri, @@ -228,6 +246,11 @@ def load_frames_corpus() -> list[dict[str, Any]]: if index % 100 == 0: logger.info(f"Fetched {index}/{len(uris)} articles") logger.info(f"Fetched {len(rows)}/{len(uris)} articles") + if len(rows) < len(uris): + raise RuntimeError( + f"Fetched only {len(rows)}/{len(uris)} FRAMES articles; " + "refusing to build a partial corpus. Re-run to resume from cache." + ) _cached_corpus = rows return _cached_corpus diff --git a/evaluations/tests/test_datasets.py b/evaluations/tests/test_datasets.py index 40805377..61f96bc3 100644 --- a/evaluations/tests/test_datasets.py +++ b/evaluations/tests/test_datasets.py @@ -1,6 +1,9 @@ from pathlib import Path +import pytest + from evaluations.datasets.frames import ( + FETCH_ATTEMPTS, build_frames_case, fetch_article, map_frames_document, @@ -564,3 +567,124 @@ class TestFrames: content = Path(row["path"]).read_text() assert "1908 Summer Olympics" in content assert "2012 Summer Olympics" in content + + def test_fetch_article_retries_transient_failures( + self, tmp_path: Path, monkeypatch + ) -> None: + sleeps: list[float] = [] + monkeypatch.setattr( + "evaluations.datasets.frames.time.sleep", lambda s: sleeps.append(s) + ) + + class FlakyResponse: + text = "ok" + headers = {"etag": 'W/"42/uuid"'} + + def raise_for_status(self) -> None: + pass + + class FlakyClient: + def __init__(self) -> None: + self.calls = 0 + + def get(self, url: str, params: dict | None = None) -> FlakyResponse: + self.calls += 1 + if self.calls < 3: + raise OSError("connection reset") + return FlakyResponse() + + client = FlakyClient() + row = fetch_article( + "https://en.wikipedia.org/wiki/Capybara", + tmp_path, + client=client, # ty: ignore[invalid-argument-type] + ) + assert row is not None + assert row["revid"] == "42" + assert client.calls == 3 + # One throttle sleep before fetching plus one backoff per failure. + assert len(sleeps) == 3 + + def test_fetch_article_honors_retry_after_on_rate_limit( + self, tmp_path: Path, monkeypatch + ) -> None: + import httpx + + sleeps: list[float] = [] + monkeypatch.setattr( + "evaluations.datasets.frames.time.sleep", lambda s: sleeps.append(s) + ) + request = httpx.Request("GET", "https://en.wikipedia.org/x") + + class OkResponse: + text = "ok" + headers = {"etag": 'W/"42/uuid"'} + + def raise_for_status(self) -> None: + pass + + class RateLimitedClient: + def __init__(self) -> None: + self.calls = 0 + + def get(self, url: str, params: dict | None = None) -> OkResponse: + self.calls += 1 + if self.calls == 1: + raise httpx.HTTPStatusError( + "429 too many requests", + request=request, + response=httpx.Response( + 429, headers={"retry-after": "13"}, request=request + ), + ) + return OkResponse() + + row = fetch_article( + "https://en.wikipedia.org/wiki/Capybara", + tmp_path, + client=RateLimitedClient(), # ty: ignore[invalid-argument-type] + ) + assert row is not None + assert 13.0 in sleeps + + def test_fetch_article_gives_up_after_max_attempts( + self, tmp_path: Path, monkeypatch + ) -> None: + monkeypatch.setattr("evaluations.datasets.frames.time.sleep", lambda s: None) + + class DeadClient: + def __init__(self) -> None: + self.calls = 0 + + def get(self, url: str, params: dict | None = None): + self.calls += 1 + raise OSError("connection reset") + + client = DeadClient() + row = fetch_article( + "https://en.wikipedia.org/wiki/Capybara", + tmp_path, + client=client, # ty: ignore[invalid-argument-type] + ) + assert row is None + assert client.calls == FETCH_ATTEMPTS + + def test_load_corpus_raises_on_partial_fetch(self, monkeypatch) -> None: + import evaluations.datasets.frames as frames + + monkeypatch.setattr(frames, "_cached_corpus", None) + monkeypatch.setattr( + frames, + "load_frames_test", + lambda: [ + { + "wiki_links": "['https://en.wikipedia.org/wiki/A', " + "'https://en.wikipedia.org/wiki/B']" + } + ], + ) + monkeypatch.setattr( + frames, "fetch_article", lambda uri, cache_dir, client: None + ) + with pytest.raises(RuntimeError, match="0/2"): + frames.load_frames_corpus()