refactor: streamline TikTok and Instagram scraping logic by removing search_queries and enhancing documentation for clarity

This commit is contained in:
DESKTOP-RTLN3BA\$punk 2026-07-13 17:11:25 -07:00
parent e8b3692b54
commit 2b018c4474
111 changed files with 1800 additions and 1580 deletions

View file

@ -14,11 +14,11 @@ Answer the delegated question from live TikTok data gathered with your verb, com
</available_tools>
<playbook>
- Finding videos on a topic: prefer `tiktok_scrape` with `hashtags` (no leading '#') or a direct TikTok URL in `urls` (fastest). `search_queries` also finds videos on a topic, but it is Google-backed and slow, so start with **at most 3** distinct queries and only add more if the first round returns nothing significant — never batch many phrasing variants of the same intent.
- Finding videos on a topic: call `tiktok_scrape` with `hashtags` (no leading '#'), or pass a TikTok URL in `urls`. There is no keyword-video search — use hashtags or a video URL.
- Scraping a specific video, profile, hashtag, or search page: pass its TikTok URL in `urls`.
- Profiles: a creator's `profiles` feed returns the account's metadata (name, followers, bio, verification) reliably, but its video list is often withheld by TikTok — treat an empty video list as a known limit, not a failure to retry endlessly. Prefer `hashtags` or a direct video URL for videos.
- Comments on a video: call `tiktok_comments` with the video URL(s) in `video_urls`.
- Finding accounts by keyword: call `tiktok_user_search` with `queries` — that is the path for accounts. Use `search_queries` on `tiktok_scrape` only when you want videos, not accounts.
- Finding accounts by keyword: call `tiktok_user_search` with `queries` — that is the path for accounts.
- "What's trending now": call `tiktok_trending` (no query needed); set `max_items` for how many.
- Controlling volume: use `max_items` for the total cap and `results_per_page` per target (per-verb equivalents: `comments_per_video`, `results_per_query`).
- Requested counts: `max_items` defaults low — when the task asks for N items, set `max_items` (and the per-target count) above N. A call that caps below the target can never satisfy it.

View file

@ -54,9 +54,7 @@ class DetailsInput(BaseModel):
@model_validator(mode="after")
def _exactly_one_source(self) -> DetailsInput:
if not self.urls and not self.search_queries:
raise ValueError(
"Provide at least one of 'urls' or 'search_queries'."
)
raise ValueError("Provide at least one of 'urls' or 'search_queries'.")
if self.urls and self.search_queries:
raise ValueError(
"Provide 'urls' OR 'search_queries', not both (they cannot be combined)."

View file

@ -77,9 +77,7 @@ class ScrapeInput(BaseModel):
@model_validator(mode="after")
def _exactly_one_source(self) -> ScrapeInput:
if not self.urls and not self.search_queries:
raise ValueError(
"Provide at least one of 'urls' or 'search_queries'."
)
raise ValueError("Provide at least one of 'urls' or 'search_queries'.")
if self.urls and self.search_queries:
raise ValueError(
"Provide 'urls' OR 'search_queries', not both (they cannot be combined)."

View file

@ -10,9 +10,8 @@ from app.capabilities.tiktok.scrape.schemas import ScrapeInput, ScrapeOutput
TIKTOK_SCRAPE = Capability(
name="tiktok.scrape",
description=(
"Scrape public TikTok videos. Use urls, profiles, hashtags, or "
"search_queries (search_queries are resolved via Google to public "
"videos; for accounts by keyword use tiktok.user_search)."
"Scrape public TikTok videos. Use urls, profiles, or hashtags. To find "
"accounts by keyword, use tiktok.user_search."
),
input_schema=ScrapeInput,
output_schema=ScrapeOutput,

View file

@ -26,7 +26,6 @@ def build_scrape_executor(scrape_fn: ScrapeFn | None = None) -> Executor:
startUrls=[{"url": url} for url in payload.urls],
profiles=payload.profiles,
hashtags=payload.hashtags,
searchQueries=payload.search_queries,
resultsPerPage=payload.results_per_page,
)
emit_progress(

View file

@ -4,7 +4,8 @@ A lean, agent-friendly surface over ``TikTokScrapeInput``
(``app/proprietary/platforms/tiktok``). The executor maps this to the full
scraper input; the scraper's ``TikTokVideoItem`` is reused verbatim as the
output element. Any TikTok URL kind (video, profile, hashtag, search) goes in
``urls``; ``profiles``/``hashtags``/``search_queries`` are typed shortcuts.
``urls``; ``profiles``/``hashtags`` are typed shortcuts. Keyword search is not a
video source here use ``tiktok.user_search`` to find accounts by keyword.
"""
from __future__ import annotations
@ -26,8 +27,8 @@ class ScrapeInput(BaseModel):
max_length=MAX_TIKTOK_SOURCES,
description=(
"TikTok URLs to scrape: a video, a profile (/@<user>), a hashtag "
"(/tag/<name>), or a search URL. Provide these OR profiles/hashtags/"
"search_queries (at least one source is required)."
"(/tag/<name>), or a search URL. Provide these OR profiles/hashtags "
"(at least one source is required)."
),
)
profiles: list[str] = Field(
@ -40,21 +41,11 @@ class ScrapeInput(BaseModel):
max_length=MAX_TIKTOK_SOURCES,
description="Hashtag names to scrape, without the leading '#'.",
)
search_queries: list[str] = Field(
default_factory=list,
max_length=MAX_TIKTOK_SOURCES,
description=(
"Search terms resolved via Google (site:tiktok.com) to public TikTok "
"videos, since TikTok's own keyword search is login-walled. Slower "
"than hashtags/urls. To find accounts by keyword, use "
"tiktok.user_search instead."
),
)
results_per_page: int = Field(
default=10,
ge=1,
le=MAX_TIKTOK_ITEMS,
description="Max videos to pull per profile/hashtag/search target.",
description="Max videos to pull per profile/hashtag target.",
)
max_items: int = Field(
default=10,
@ -65,10 +56,9 @@ class ScrapeInput(BaseModel):
@model_validator(mode="after")
def _require_a_source(self) -> ScrapeInput:
if not any((self.urls, self.profiles, self.hashtags, self.search_queries)):
if not any((self.urls, self.profiles, self.hashtags)):
raise ValueError(
"Provide at least one of 'urls', 'profiles', 'hashtags', or "
"'search_queries'."
"Provide at least one of 'urls', 'profiles', or 'hashtags'."
)
return self

View file

@ -222,9 +222,7 @@ class _RotatingSession:
await self.close()
self.rotations += 1
await self._open()
logger.info(
"[instagram] rotated proxy session (rotation #%d)", self.rotations
)
logger.info("[instagram] rotated proxy session (rotation #%d)", self.rotations)
return self.session
async def pace(self) -> None:
@ -378,9 +376,7 @@ async def _fetch(
if status == _BACKOFF_STATUS and backoffs < _MAX_BACKOFFS:
backoffs += 1
delay = _BACKOFF_BASE_S * (2 ** (backoffs - 1))
logger.warning(
"[instagram] 429 on %s; backing off %.1fs", path, delay
)
logger.warning("[instagram] 429 on %s; backing off %.1fs", path, delay)
await asyncio.sleep(delay + random.uniform(0, 1))
continue
if status in _ROTATE_STATUSES:

View file

@ -171,7 +171,9 @@ def _relay_child(node: dict[str, Any]) -> dict[str, Any]:
mt = node.get("media_type")
vv = node.get("video_versions")
video_url = (
vv[0].get("url") if isinstance(vv, list) and vv and isinstance(vv[0], dict) else None
vv[0].get("url")
if isinstance(vv, list) and vv and isinstance(vv[0], dict)
else None
)
is_video = mt == 2 or bool(video_url)
return {
@ -290,9 +292,7 @@ def parse_profile(user: dict[str, Any]) -> dict[str, Any]:
_APP_JSON_RE = re.compile(
r'<script type="application/json"[^>]*>(.*?)</script>', re.DOTALL
)
_OG_RE = re.compile(
r'<meta\s+property="og:([^"]+)"\s+content="([^"]*)"', re.IGNORECASE
)
_OG_RE = re.compile(r'<meta\s+property="og:([^"]+)"\s+content="([^"]*)"', re.IGNORECASE)
# og tags are the fallback source (used only when the relay blob is absent). They
# follow a fixed English shape because the fetch layer pins Accept-Language en-US:
# og:description = "{likes} likes, {comments} comments - {username} on {Month D, YYYY}: "{caption}""
@ -321,7 +321,9 @@ _MEDIA_ID_RE = re.compile(r"instagram://media\?id=(\d+)")
def _og_date_to_iso(value: str) -> str | None:
"""``"July 9, 2026"`` -> ``"2026-07-09"`` (date-only; og carries no time)."""
try:
return datetime.strptime(value, "%B %d, %Y").replace(tzinfo=UTC).date().isoformat()
return (
datetime.strptime(value, "%B %d, %Y").replace(tzinfo=UTC).date().isoformat()
)
except ValueError:
return None
@ -359,7 +361,7 @@ def _parse_og_meta(og: dict[str, str]) -> dict[str, Any]:
elif owner_date:
# No usable og:title: fall back to the caption after og:description's
# date prefix — still clean (the counts/username/date are stripped).
out["caption"] = _clean_caption(desc[owner_date.end():])
out["caption"] = _clean_caption(desc[owner_date.end() :])
return out
@ -438,13 +440,21 @@ def _media_from_relay(
mt = media.get("media_type")
cap = media.get("caption")
caption = (
cap.get("text") if isinstance(cap, dict) else (cap if isinstance(cap, str) else None)
cap.get("text")
if isinstance(cap, dict)
else (cap if isinstance(cap, str) else None)
)
carousel = media.get("carousel_media")
carousel = [c for c in carousel if isinstance(c, dict)] if isinstance(carousel, list) else []
carousel = (
[c for c in carousel if isinstance(c, dict)]
if isinstance(carousel, list)
else []
)
vv = media.get("video_versions")
video_url = (
vv[0].get("url") if isinstance(vv, list) and vv and isinstance(vv[0], dict) else None
vv[0].get("url")
if isinstance(vv, list) and vv and isinstance(vv[0], dict)
else None
)
is_video = mt == 2 or bool(video_url)
owner = media.get("user") if isinstance(media.get("user"), dict) else {}
@ -469,13 +479,18 @@ def _media_from_relay(
"type": _MEDIA_TYPE.get(mt) or ("Video" if is_video else "Image"),
"shortCode": media.get("code") or shortcode,
"caption": caption,
"hashtags": list(dict.fromkeys(_HASHTAG_RE.findall(caption))) if caption else [],
"mentions": list(dict.fromkeys(_MENTION_RE.findall(caption))) if caption else [],
"hashtags": list(dict.fromkeys(_HASHTAG_RE.findall(caption)))
if caption
else [],
"mentions": list(dict.fromkeys(_MENTION_RE.findall(caption)))
if caption
else [],
"url": url,
"commentsCount": _int(media.get("comment_count")),
"dimensionsHeight": _int(media.get("original_height")),
"dimensionsWidth": _int(media.get("original_width")),
"displayUrl": _iv2_url(media.get("image_versions2")) or media.get("display_uri"),
"displayUrl": _iv2_url(media.get("image_versions2"))
or media.get("display_uri"),
"images": [
u
for c in carousel
@ -535,8 +550,12 @@ def parse_post(
"type": "Video" if is_video else "Image",
"shortCode": shortcode,
"caption": caption,
"hashtags": list(dict.fromkeys(_HASHTAG_RE.findall(caption))) if caption else [],
"mentions": list(dict.fromkeys(_MENTION_RE.findall(caption))) if caption else [],
"hashtags": list(dict.fromkeys(_HASHTAG_RE.findall(caption)))
if caption
else [],
"mentions": list(dict.fromkeys(_MENTION_RE.findall(caption)))
if caption
else [],
"url": url,
"commentsCount": og_meta.get("comments"),
"displayUrl": og.get("image"),

View file

@ -328,9 +328,7 @@ async def _discover_via_google(
return resolved
async def _discover(
query: str, *, search_type: str, limit: int
) -> list[ResolvedUrl]:
async def _discover(query: str, *, search_type: str, limit: int) -> list[ResolvedUrl]:
"""Resolve a discovery query into profile targets - anonymously.
A query that is a valid handle resolves directly against the anonymous
@ -397,9 +395,7 @@ async def iter_instagram(
# posts / reels -> media feeds, de-duped by id across targets.
jobs = [
_media_flow(
r, input_model=input_model, cutoff=cutoff, per_target=per_target
)
_media_flow(r, input_model=input_model, cutoff=cutoff, per_target=per_target)
for r in targets
]
seen: set[str] = set()

View file

@ -25,9 +25,7 @@ from urllib.parse import urlparse
ResolvedKind = Literal["profile", "post", "reel"]
_INSTAGRAM_HOSTS = frozenset(
{"m.instagram.com", "www.instagram.com", "instagram.com"}
)
_INSTAGRAM_HOSTS = frozenset({"m.instagram.com", "www.instagram.com", "instagram.com"})
_STRIP_SEGMENTS = frozenset({"_u", "profilecard"})
_RESERVED = frozenset(
{"p", "s", "tv", "reel", "reels", "share", "explore", "stories", "accounts"}
@ -68,9 +66,7 @@ def resolve_url(url: str) -> ResolvedUrl | None:
if "instagram.com" not in url.lower():
token = url.strip().lstrip("@")
if token and "/" not in token and "." not in token:
return ResolvedUrl(
"profile", token, f"https://www.instagram.com/{token}/"
)
return ResolvedUrl("profile", token, f"https://www.instagram.com/{token}/")
segments = _segments(url)
if not segments:
return None
@ -83,9 +79,7 @@ def resolve_url(url: str) -> ResolvedUrl | None:
return ResolvedUrl("reel", code, url, numeric_post_id=code.isdigit())
if head == "stories" and len(segments) >= 2:
user = segments[1]
return ResolvedUrl(
"profile", user, f"https://www.instagram.com/{user}/"
)
return ResolvedUrl("profile", user, f"https://www.instagram.com/{user}/")
if head not in _RESERVED:
return ResolvedUrl("profile", head, url)
return None

View file

@ -5,11 +5,11 @@ from __future__ import annotations
from datetime import UTC, datetime
def epoch_to_iso(seconds: int | None) -> str | None:
def epoch_to_iso(seconds: int | str | None) -> str | None:
"""Convert a Unix-seconds timestamp to ``YYYY-MM-DDTHH:MM:SS.000Z``."""
if not seconds:
return None
stamp = datetime.fromtimestamp(seconds, tz=UTC)
stamp = datetime.fromtimestamp(int(seconds), tz=UTC)
return stamp.strftime("%Y-%m-%dT%H:%M:%S.000Z")

View file

@ -11,7 +11,9 @@ from ..targets.types import TikTokTarget
from . import FetchFn
async def iter_video(target: TikTokTarget, *, fetch: FetchFn) -> AsyncIterator[dict[str, Any]]:
async def iter_video(
target: TikTokTarget, *, fetch: FetchFn
) -> AsyncIterator[dict[str, Any]]:
html = await fetch(target.url)
if not html:
return

View file

@ -10,10 +10,6 @@ from __future__ import annotations
from collections.abc import AsyncIterator
from typing import Any
from urllib.parse import quote
from app.proprietary.platforms.google_search.schemas import GoogleSearchScrapeInput
from app.proprietary.platforms.google_search.scraper import scrape_serps
from .extraction.timestamps import now_iso
from .flows import FetchCommentsFn, FetchFn, FetchListingFn, FetchUsersFn
@ -35,26 +31,15 @@ from .targets.types import TikTokTarget
_PROFILE_URL = "https://www.tiktok.com/@{name}"
_HASHTAG_URL = "https://www.tiktok.com/tag/{tag}"
_SEARCH_URL = "https://www.tiktok.com/search?q={query}"
_EXPLORE_URL = "https://www.tiktok.com/explore"
# A ``searchQueries`` term whose Google discovery surfaced no scrapable video
# URLs degrades to one honest ErrorItem (mirrors the listing flow's contract:
# never vanish silently).
_EMPTY_DISCOVERY_MESSAGE = (
"No public TikTok videos found for this query via Google discovery. Try a "
"narrower phrasing, a hashtag, or a direct video URL."
)
def _resolve_targets(input_model: TikTokScrapeInput) -> list[TikTokTarget]:
"""Build the target list from the URL/profile/hashtag sources.
``searchQueries`` is deliberately excluded: TikTok's own keyword search is
login-walled for anonymous sessions, so it is routed through Google video
discovery in :func:`iter_tiktok` instead. A raw ``tiktok.com/search?...``
URL passed explicitly in ``startUrls``/``postURLs`` still resolves here and
keeps its native listing routing.
A raw ``tiktok.com/search?...`` URL passed explicitly in
``startUrls``/``postURLs`` still resolves here and keeps its native listing
routing; there is no keyword-search shortcut.
"""
targets: list[TikTokTarget] = []
for entry in input_model.startUrls:
@ -73,39 +58,6 @@ def _resolve_targets(input_model: TikTokScrapeInput) -> list[TikTokTarget]:
return targets
async def _discover_via_google(query: str, *, limit: int) -> list[TikTokTarget]:
"""Discover public TikTok video targets via Google ``site:tiktok.com``.
TikTok's anonymous keyword search is login-walled, so we reuse the existing
``google_search`` platform, classify each organic URL with ``resolve_target``,
and keep only video hits (``/@user/video/<id>``) the one kind that scrapes
reliably over plain HTTP. Profile/hashtag/search/photo/non-tiktok results are
dropped (accounts belong to the ``user_search`` verb). De-duped, capped at
``limit``.
"""
serps = await scrape_serps(
GoogleSearchScrapeInput(
queries=query, site="tiktok.com", maxPagesPerQuery=1
),
limit=1,
)
resolved: list[TikTokTarget] = []
seen: set[str] = set()
for serp in serps:
for org in serp.get("organicResults") or []:
url = org.get("url", "") if isinstance(org, dict) else ""
target = resolve_target(url)
if target is None or target.kind != "video":
continue
if target.value in seen:
continue
seen.add(target.value)
resolved.append(target)
if len(resolved) >= limit:
return resolved
return resolved
def _dispatch(
target: TikTokTarget,
*,
@ -128,11 +80,9 @@ async def iter_tiktok(
) -> AsyncIterator[dict[str, Any]]:
"""Yield normalized items for every resolved target, in order.
Direct sources (URLs, profiles, hashtags) resolve up front; ``searchQueries``
then run through Google video discovery. The video flow's ``fetch_html``
opens its own warmed proxy session per call when none is bound; the listing
flow drives its own browser. Neither binds a ContextVar across these
``yield``s, so the generator stays context-safe.
The video flow's ``fetch_html`` opens its own warmed proxy session per call
when none is bound; the listing flow drives its own browser. Neither binds a
ContextVar across these ``yield``s, so the generator stays context-safe.
"""
cap = input_model.resultsPerPage
for target in _resolve_targets(input_model):
@ -141,27 +91,6 @@ async def iter_tiktok(
):
yield item
# searchQueries -> Google-discovered public video URLs, de-duped across
# queries so the same video surfacing under two terms is scraped once.
seen_videos: set[str] = set()
for query in input_model.searchQueries:
discovered = await _discover_via_google(query, limit=cap)
if not discovered:
yield ErrorItem(
url=_SEARCH_URL.format(query=quote(query)),
input=query,
error=_EMPTY_DISCOVERY_MESSAGE,
errorCode="no_items",
scrapedAt=now_iso(),
).to_output()
continue
for target in discovered:
if target.value in seen_videos:
continue
seen_videos.add(target.value)
async for item in iter_video(target, fetch=fetch):
yield item
async def scrape_tiktok(
input_model: TikTokScrapeInput,
@ -174,7 +103,9 @@ async def scrape_tiktok(
from app.capabilities.core.progress import emit_progress
results: list[dict[str, Any]] = []
async for item in iter_tiktok(input_model, fetch=fetch, fetch_listing=fetch_listing):
async for item in iter_tiktok(
input_model, fetch=fetch, fetch_listing=fetch_listing
):
results.append(item)
emit_progress("scraping", current=len(results), total=limit, unit="item")
if limit is not None and len(results) >= limit:

View file

@ -54,9 +54,7 @@ from app.proprietary.platforms.instagram.url_resolver import resolve_url # noqa
_PROFILE = "natgeo"
_SEARCH_TERM = "national geographic"
_FIXTURE_DIR = (
_BACKEND_ROOT / "tests" / "unit" / "platforms" / "instagram" / "fixtures"
)
_FIXTURE_DIR = _BACKEND_ROOT / "tests" / "unit" / "platforms" / "instagram" / "fixtures"
# Fields to strip from dumped fixtures so we never commit PII / volatile tokens.
_PII_KEYS = frozenset(
@ -98,7 +96,9 @@ async def step0_probe() -> bool:
data = await fetch_json(
"api/v1/users/web_profile_info/", {"username": _PROFILE}
)
user = (data or {}).get("data", {}).get("user") if isinstance(data, dict) else None
user = (
(data or {}).get("data", {}).get("user") if isinstance(data, dict) else None
)
print(f" web_profile_info({_PROFILE}) -> user={'yes' if user else 'no'}")
return _check("sticky web_profile_info", minted and bool(user))
@ -179,9 +179,7 @@ async def step5_search() -> bool:
async def step6_dump_fixtures(post_url: str | None) -> bool:
_hr("STEP 6 — dump trimmed, anonymized fixtures for offline tests")
profile = await fetch_json(
"api/v1/users/web_profile_info/", {"username": _PROFILE}
)
profile = await fetch_json("api/v1/users/web_profile_info/", {"username": _PROFILE})
_FIXTURE_DIR.mkdir(parents=True, exist_ok=True)
wrote = []
if isinstance(profile, dict) and profile.get("data", {}).get("user"):

View file

@ -178,7 +178,9 @@ async def stage_pipeline() -> bool:
f"{len(items)} item(s)",
)
if items:
print(f" sample: {items[0].get('webVideoUrl')}{items[0].get('text', '')[:60]!r}")
print(
f" sample: {items[0].get('webVideoUrl')}{items[0].get('text', '')[:60]!r}"
)
return ok
@ -210,9 +212,7 @@ async def stage_comments(video_url: str) -> tuple[bool, list[dict[str, Any]]]:
# Comments load over a signed /api/comment/list XHR that TikTok serves to
# anonymous sessions once the panel opens. Pass if real comments come back
# OR a graceful ErrorItem (video has none / disabled / withheld).
items = await scrape_tiktok_comments(
[video_url], per_video=_COUNT, limit=_COUNT
)
items = await scrape_tiktok_comments([video_url], per_video=_COUNT, limit=_COUNT)
has_comment = any(it.get("id") and not it.get("errorCode") for it in items)
has_error = any(it.get("errorCode") == "no_comments" for it in items)
ok = _check(
@ -253,7 +253,9 @@ async def stage_trending() -> tuple[bool, list[dict[str, Any]]]:
f"{len(items)} item(s); videos={len(real)}",
)
if real:
print(f" sample: {real[0].get('webVideoUrl')}{real[0].get('text', '')[:60]!r}")
print(
f" sample: {real[0].get('webVideoUrl')}{real[0].get('text', '')[:60]!r}"
)
return ok, items

View file

@ -66,6 +66,4 @@ def test_details_wraps_profile_items():
def test_details_rejects_both_sources():
with pytest.raises(ValidationError):
DetailsInput(
urls=["https://www.instagram.com/natgeo/"], search_queries=["x"]
)
DetailsInput(urls=["https://www.instagram.com/natgeo/"], search_queries=["x"])

View file

@ -54,7 +54,6 @@ async def test_forwards_typed_sources_and_limit():
ScrapeInput(
profiles=["nasa"],
hashtags=["food"],
search_queries=["cats"],
results_per_page=7,
max_items=25,
)
@ -63,7 +62,6 @@ async def test_forwards_typed_sources_and_limit():
(actor_input, limit) = scraper.calls[0]
assert actor_input.profiles == ["nasa"]
assert actor_input.hashtags == ["food"]
assert actor_input.searchQueries == ["cats"]
assert actor_input.resultsPerPage == 7
# The outer collection limit is the caller's total-item cap.
assert limit == 25

View file

@ -42,8 +42,7 @@ def _profile_payload(n: int) -> dict:
"edge_owner_to_timeline_media": {
"count": n,
"edges": [
{"node": {"id": str(i), "shortcode": f"S{i}"}}
for i in range(n)
{"node": {"id": str(i), "shortcode": f"S{i}"}} for i in range(n)
],
},
}

View file

@ -33,9 +33,7 @@ async def test_google_discovery_keeps_only_profiles(monkeypatch):
"https://example.com/not-instagram",
),
)
targets = await scraper._discover(
"nat geo photos", search_type="profile", limit=10
)
targets = await scraper._discover("nat geo photos", search_type="profile", limit=10)
assert [(t.kind, t.value) for t in targets] == [("profile", "natgeo")]
@ -48,9 +46,7 @@ async def test_google_discovery_dedupes(monkeypatch):
"https://www.instagram.com/natgeo/",
),
)
targets = await scraper._discover(
"nat geo photos", search_type="profile", limit=10
)
targets = await scraper._discover("nat geo photos", search_type="profile", limit=10)
assert len(targets) == 1

View file

@ -102,7 +102,9 @@ async def test_warms_then_returns_json():
holder = _FakeHolder([_FakeSession(200, csrftoken=True)])
token = _current_session.set(holder)
try:
result = await fetch_json("api/v1/users/web_profile_info/", {"username": "natgeo"})
result = await fetch_json(
"api/v1/users/web_profile_info/", {"username": "natgeo"}
)
finally:
_current_session.reset(token)
assert result == _PAYLOAD

View file

@ -175,8 +175,14 @@ def test_parse_post_prefers_relay_json():
"image_versions2": {"candidates": [{"url": "https://cdn/c2.jpg"}]},
},
],
"usertags": {"in": [{"position": [0.5, 0.5], "user": {"username": "tagged1", "id": "77"}}]},
"coauthor_producers": [{"username": "coauthor1", "id": "88", "is_verified": True}],
"usertags": {
"in": [
{"position": [0.5, 0.5], "user": {"username": "tagged1", "id": "77"}}
]
},
"coauthor_producers": [
{"username": "coauthor1", "id": "88", "is_verified": True}
],
"location": {"id": "123", "name": "Bali"},
}
html = (

View file

@ -1,149 +0,0 @@
"""Offline tests for Google-backed TikTok video discovery.
``searchQueries`` are login-walled on TikTok's native search, so they route
through the ``google_search`` platform (``site:tiktok.com``): each organic URL
is classified with ``resolve_target`` and only video hits (``/@user/video/<id>``)
are kept profiles/hashtags/search/photo/non-tiktok are dropped (accounts
belong to the user-search verb). These tests inject a fake ``scrape_serps`` so
there is no network: they pin the classification, cross-query de-dup, the limit
cap, the barren-query ErrorItem, and that no ``/search?q=`` listing target is
ever built.
"""
from __future__ import annotations
import json
from app.proprietary.platforms.tiktok import (
TikTokScrapeInput,
orchestrator,
scrape_tiktok,
)
def _fake_serps(*organic_urls: str):
async def _scrape_serps(input_model, *, limit=None):
assert input_model.site == "tiktok.com"
assert input_model.maxPagesPerQuery == 1
return [{"organicResults": [{"url": u} for u in organic_urls]}]
return _scrape_serps
def _video_page(url: str) -> str:
"""Render a rehydration blob for a ``/@user/video/<id>`` URL."""
video_id = url.rsplit("/", 1)[1]
username = url.split("@")[1].split("/")[0]
blob = {
"__DEFAULT_SCOPE__": {
"webapp.video-detail": {
"itemInfo": {
"itemStruct": {
"id": video_id,
"desc": "hi",
"author": {"uniqueId": username},
"stats": {"diggCount": 1},
}
}
}
}
}
return (
'<script id="__UNIVERSAL_DATA_FOR_REHYDRATION__" '
f'type="application/json">{json.dumps(blob)}</script>'
)
async def _fetch_video(url: str) -> str:
return _video_page(url)
async def test_search_discovery_keeps_only_videos(monkeypatch):
# Only the video URL survives; profile / hashtag / search / photo /
# non-tiktok organic results are dropped.
monkeypatch.setattr(
orchestrator,
"scrape_serps",
_fake_serps(
"https://www.tiktok.com/@nasa/video/123",
"https://www.tiktok.com/@nasa",
"https://www.tiktok.com/tag/space",
"https://www.tiktok.com/search?q=space",
"https://www.tiktok.com/@nasa/photo/999",
"https://example.com/not-tiktok",
),
)
items = await scrape_tiktok(
TikTokScrapeInput(searchQueries=["space"], resultsPerPage=10),
fetch=_fetch_video,
)
assert [i["id"] for i in items] == ["123"]
async def test_search_discovery_dedupes_across_queries(monkeypatch):
# The same video surfacing under two queries is scraped once.
monkeypatch.setattr(
orchestrator,
"scrape_serps",
_fake_serps("https://www.tiktok.com/@nasa/video/123"),
)
items = await scrape_tiktok(
TikTokScrapeInput(searchQueries=["space", "rockets"], resultsPerPage=10),
fetch=_fetch_video,
)
assert [i["id"] for i in items] == ["123"]
async def test_search_discovery_respects_per_target_limit(monkeypatch):
monkeypatch.setattr(
orchestrator,
"scrape_serps",
_fake_serps(
"https://www.tiktok.com/@a/video/1",
"https://www.tiktok.com/@b/video/2",
"https://www.tiktok.com/@c/video/3",
),
)
items = await scrape_tiktok(
TikTokScrapeInput(searchQueries=["x"], resultsPerPage=2),
fetch=_fetch_video,
)
assert [i["id"] for i in items] == ["1", "2"]
async def test_search_barren_query_emits_error_item(monkeypatch):
# A query whose discovery finds no video URLs degrades to one ErrorItem.
monkeypatch.setattr(
orchestrator,
"scrape_serps",
_fake_serps(
"https://www.tiktok.com/@nasa",
"https://example.com/x",
),
)
items = await scrape_tiktok(
TikTokScrapeInput(searchQueries=["space"], resultsPerPage=10),
fetch=_fetch_video,
)
assert len(items) == 1
assert items[0]["errorCode"] == "no_items"
assert items[0]["input"] == "space"
async def test_search_never_builds_listing_target(monkeypatch):
# searchQueries must never hit the (login-walled) native search listing flow.
monkeypatch.setattr(
orchestrator,
"scrape_serps",
_fake_serps("https://www.tiktok.com/@nasa/video/123"),
)
async def _boom_listing(_url: str, _count: int) -> list[dict]:
raise AssertionError("searchQueries must not build a listing target")
items = await scrape_tiktok(
TikTokScrapeInput(searchQueries=["space"], resultsPerPage=10),
fetch=_fetch_video,
fetch_listing=_boom_listing,
)
assert [i["id"] for i in items] == ["123"]

View file

@ -84,7 +84,9 @@ async def test_warms_then_returns_html():
async def test_rotates_when_warm_fails_then_succeeds():
holder = _FakeHolder([_FakeSession(200, warms=False), _FakeSession(200, warms=True)])
holder = _FakeHolder(
[_FakeSession(200, warms=False), _FakeSession(200, warms=True)]
)
token = _current_session.set(holder)
try:
result = await client.fetch_html("https://www.tiktok.com/@scout2015")
@ -119,9 +121,7 @@ async def test_rotates_and_rewarms_on_403():
async def test_persistent_403_raises_blocked(monkeypatch):
_no_sleep(monkeypatch)
holder = _FakeHolder(
[_FakeSession(403) for _ in range(client._MAX_ROTATIONS + 1)]
)
holder = _FakeHolder([_FakeSession(403) for _ in range(client._MAX_ROTATIONS + 1)])
token = _current_session.set(holder)
try:
raised = False

View file

@ -6,7 +6,9 @@ from app.proprietary.platforms.tiktok.targets import resolve_target
def test_resolve_video_carries_username_and_id():
target = resolve_target("https://www.tiktok.com/@scout2015/video/6718335390845095173")
target = resolve_target(
"https://www.tiktok.com/@scout2015/video/6718335390845095173"
)
assert target is not None
assert target.kind == "video"
assert target.value == "6718335390845095173"

View file

@ -28,9 +28,7 @@ async def test_user_search_parses_dedupes_and_caps():
async def fake_fetch(_url: str, _cap: int) -> list[dict]:
return [_user("1", "nasa"), _user("1", "nasa"), _user("2", "nasa2")]
items = await search_tiktok_users(
["nasa"], per_query=2, fetch_users=fake_fetch
)
items = await search_tiktok_users(["nasa"], per_query=2, fetch_users=fake_fetch)
assert [i["id"] for i in items] == ["1", "2"]
first = items[0]
@ -49,9 +47,7 @@ async def test_user_search_empty_query_emits_error_item():
async def fake_fetch(_url: str, _cap: int) -> list[dict]:
return []
items = await search_tiktok_users(
["ghost"], per_query=5, fetch_users=fake_fetch
)
items = await search_tiktok_users(["ghost"], per_query=5, fetch_users=fake_fetch)
assert len(items) == 1
assert items[0]["errorCode"] == "no_users"

View file

@ -58,9 +58,7 @@ def test_chat_model_requires_slash_tools_and_context():
def test_excluded_provider_slug_is_filtered():
assert not is_requesty_chat_model(
_requesty_model(model_id="amazon/nova-pro-v1")
)
assert not is_requesty_chat_model(_requesty_model(model_id="amazon/nova-pro-v1"))
def test_image_generation_models_excluded_from_chat_and_flagged():
@ -89,9 +87,7 @@ def test_normalize_maps_context_window_and_capabilities():
name="GPT-4o mini",
),
_requesty_model(model_id="openai/gpt-4o-mini", tools=False),
_requesty_model(
model_id="black-forest-labs/flux", image_generation=True
),
_requesty_model(model_id="black-forest-labs/flux", image_generation=True),
]
)

View file

@ -21,8 +21,7 @@ PDFS = REPO / "data" / "multimodal_doc" / "mmlongbench" / "pdfs"
def main() -> None:
rows = [
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines()
if line.strip()
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines() if line.strip()
]
# 1) SSL clustering: failures by question index per arm
@ -35,11 +34,19 @@ def main() -> None:
arm_seen_count[arm] += 1
qid_order[f"{arm}::{row['qid']}"] = idx
err = row.get("error") or ""
cluster = "ssl" if "SSLError" in err else (
"empty" if not (row.get("raw_text") or "").strip() and not err else (
"5xx" if "502" in err or "503" in err else (
"size_limit" if "exceeds" in err.lower() and "limit" in err.lower() else (
"other_err" if err else "ok"
cluster = (
"ssl"
if "SSLError" in err
else (
"empty"
if not (row.get("raw_text") or "").strip() and not err
else (
"5xx"
if "502" in err or "503" in err
else (
"size_limit"
if "exceeds" in err.lower() and "limit" in err.lower()
else ("other_err" if err else "ok")
)
)
)
@ -100,19 +107,26 @@ def main() -> None:
err = row.get("error") or ""
empty = not (row.get("raw_text") or "").strip()
if err or empty:
by_pdf[row["doc_id"]].append({
"arm": row["arm"],
"qid": row["qid"],
"err_kind": (
"ssl" if "SSLError" in err
else "size_limit" if "exceeds" in err.lower() and "limit" in err.lower()
else "5xx" if "502" in err or "503" in err
else "json_decode" if "JSONDecodeError" in err
else "empty" if empty and not err
else "other"
),
"pages": row.get("pages"),
})
by_pdf[row["doc_id"]].append(
{
"arm": row["arm"],
"qid": row["qid"],
"err_kind": (
"ssl"
if "SSLError" in err
else "size_limit"
if "exceeds" in err.lower() and "limit" in err.lower()
else "5xx"
if "502" in err or "503" in err
else "json_decode"
if "JSONDecodeError" in err
else "empty"
if empty and not err
else "other"
),
"pages": row.get("pages"),
}
)
for doc, items in sorted(by_pdf.items(), key=lambda x: (-len(x[1]), x[0])):
kinds = Counter(i["err_kind"] for i in items)
arms = sorted({i["arm"] for i in items})

View file

@ -51,8 +51,7 @@ def _classify(error: str | None, raw_text: str) -> str:
def main() -> None:
rows = [
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines()
if line.strip()
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines() if line.strip()
]
by_arm_failures: dict[str, list[dict]] = defaultdict(list)
@ -121,7 +120,9 @@ def main() -> None:
print("=" * 90)
for entry in by_arm_failures.get("native_pdf", []):
err = (entry["error"] or "(no error string)")[:240].replace("\n", " ")
print(f" {entry['qid']} doc={entry['doc_id']} pages={entry['pages']} cluster={entry['cluster']}")
print(
f" {entry['qid']} doc={entry['doc_id']} pages={entry['pages']} cluster={entry['cluster']}"
)
print(f" err: {err}")
summary: dict[str, Any] = {
@ -130,18 +131,13 @@ def main() -> None:
"n": n_per_arm[arm],
"failures": len(by_arm_failures[arm]),
"rate": len(by_arm_failures[arm]) / n_per_arm[arm],
"clusters": {
cluster: len(items)
for cluster, items in error_clusters[arm].items()
},
"clusters": {cluster: len(items) for cluster, items in error_clusters[arm].items()},
"rows": by_arm_failures[arm],
}
for arm in sorted(n_per_arm)
},
"per_pdf": {
pdf: [
{**r, "arm": r["arm"]} for r in failures
]
pdf: [{**r, "arm": r["arm"]} for r in failures]
for pdf, failures in by_pdf_failures.items()
},
}

View file

@ -23,9 +23,7 @@ SAFE_CHARS = (CTX_TOKENS - PROMPT_OVERHEAD_TOKENS - MAX_OUTPUT_TOKENS) * CHARS_P
def main() -> None:
rows = [
json.loads(line)
for line in MAP.read_text(encoding="utf-8").splitlines()
if line.strip()
json.loads(line) for line in MAP.read_text(encoding="utf-8").splitlines() if line.strip()
]
total = len(rows)

View file

@ -67,14 +67,18 @@ def classify(error: str | None, raw_text: str) -> str:
def main() -> None:
rows = [
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines()
if line.strip()
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines() if line.strip()
]
by_arm: dict[str, dict] = defaultdict(lambda: {
"n": 0, "correct": 0,
"transient_ssl_or_5xx": 0, "transient_empty": 0,
"intrinsic_limit": 0, "other_error": 0,
})
by_arm: dict[str, dict] = defaultdict(
lambda: {
"n": 0,
"correct": 0,
"transient_ssl_or_5xx": 0,
"transient_empty": 0,
"intrinsic_limit": 0,
"other_error": 0,
}
)
for row in rows:
arm = row["arm"]
m = by_arm[arm]
@ -86,7 +90,9 @@ def main() -> None:
if kind != "ok":
m[kind] += 1
print(f"{'arm':<25} {'raw acc%':>8} {'transient':>10} {'intrinsic':>10} {'other':>6} {'adj acc% (no transient)':>22}")
print(
f"{'arm':<25} {'raw acc%':>8} {'transient':>10} {'intrinsic':>10} {'other':>6} {'adj acc% (no transient)':>22}"
)
print("-" * 88)
for arm in sorted(by_arm):
m = by_arm[arm]
@ -96,9 +102,7 @@ def main() -> None:
other = m["other_error"]
usable = m["n"] - transient
adj = m["correct"] / usable * 100 if usable else 0
print(
f"{arm:<25} {raw:>7.1f}% {transient:>10} {intrinsic:>10} {other:>6} {adj:>21.1f}%"
)
print(f"{arm:<25} {raw:>7.1f}% {transient:>10} {intrinsic:>10} {other:>6} {adj:>21.1f}%")
print()
print("transient = SSLError / 502 / 503 / empty stream / mid-stream JSON decode (would")

View file

@ -95,7 +95,7 @@ def _mcnemar_exact_pvalue(b: int, c: int) -> float:
k = min(b, c)
# Two-sided exact: 2 * P(X <= k) clipped at 1.0
cdf = sum(_binom_coef(n, i) for i in range(k + 1))
p = 2.0 * cdf / (2 ** n)
p = 2.0 * cdf / (2**n)
return min(1.0, p)
@ -116,7 +116,7 @@ def _mcnemar_table(rows: list[dict]) -> dict:
qids = sorted(by_qid)
out: dict[str, dict] = {"arms": arms, "n_qids": len(qids), "pairs": []}
for i, ai in enumerate(arms):
for aj in arms[i + 1:]:
for aj in arms[i + 1 :]:
b = c = both = neither = 0
for q in qids:
row = by_qid[q]
@ -132,12 +132,17 @@ def _mcnemar_table(rows: list[dict]) -> dict:
else:
neither += 1
p = _mcnemar_exact_pvalue(b, c)
out["pairs"].append({
"arm_i": ai, "arm_j": aj,
"b_i_only": b, "c_j_only": c,
"both_correct": both, "both_wrong": neither,
"p_value": p,
})
out["pairs"].append(
{
"arm_i": ai,
"arm_j": aj,
"b_i_only": b,
"c_j_only": c,
"both_correct": both,
"both_wrong": neither,
"p_value": p,
}
)
return out
@ -154,9 +159,7 @@ def _per_pdf_stats(rows: list[dict]) -> dict[str, dict]:
arm = r["arm"]
pdf = r["doc_id"]
graded = r.get("graded") or {}
bucket.setdefault(arm, {}).setdefault(pdf, []).append(
bool(graded.get("correct"))
)
bucket.setdefault(arm, {}).setdefault(pdf, []).append(bool(graded.get("correct")))
out: dict[str, dict] = {}
for arm, pdfs in bucket.items():
@ -207,7 +210,8 @@ def _per_arm_latency(rows: list[dict]) -> dict[str, dict]:
# Coefficient of variation: std / mean (unitless tail-fatness).
"cv": (
statistics.stdev(lats) / statistics.mean(lats)
if len(lats) > 1 and statistics.mean(lats) > 0 else 0.0
if len(lats) > 1 and statistics.mean(lats) > 0
else 0.0
),
}
return out
@ -259,24 +263,30 @@ def _print_latency(title: str, lat: dict[str, dict]) -> None:
print()
print(title)
print("-" * len(title))
header = (f"{'arm':<25} {'n':>4} {'mean':>7} {'std':>7} "
f"{'p50':>7} {'p90':>7} {'p95':>7} {'p99':>7} {'max':>7} {'CV':>5}")
header = (
f"{'arm':<25} {'n':>4} {'mean':>7} {'std':>7} "
f"{'p50':>7} {'p90':>7} {'p95':>7} {'p99':>7} {'max':>7} {'CV':>5}"
)
print(header)
print("-" * len(header))
for arm in sorted(lat, key=lambda a: lat[a]["mean_s"]):
s = lat[arm]
print(f"{arm:<25} {s['n']:>4} "
f"{s['mean_s']:>6.1f}s {s['std_s']:>6.1f}s "
f"{s['p50_s']:>6.1f}s {s['p90_s']:>6.1f}s {s['p95_s']:>6.1f}s "
f"{s['p99_s']:>6.1f}s {s['max_s']:>6.1f}s {s['cv']:>5.2f}")
print(
f"{arm:<25} {s['n']:>4} "
f"{s['mean_s']:>6.1f}s {s['std_s']:>6.1f}s "
f"{s['p50_s']:>6.1f}s {s['p90_s']:>6.1f}s {s['p95_s']:>6.1f}s "
f"{s['p99_s']:>6.1f}s {s['max_s']:>6.1f}s {s['cv']:>5.2f}"
)
def _print_tokens(title: str, toks: dict[str, dict]) -> None:
print()
print(title)
print("-" * len(title))
header = (f"{'arm':<25} {'in mean':>9} {'in p50':>9} {'in p95':>9} {'in max':>9}"
f" {'out mean':>9} {'out p95':>9}")
header = (
f"{'arm':<25} {'in mean':>9} {'in p50':>9} {'in p95':>9} {'in max':>9}"
f" {'out mean':>9} {'out p95':>9}"
)
print(header)
print("-" * len(header))
for arm in sorted(toks):
@ -285,25 +295,31 @@ def _print_tokens(title: str, toks: dict[str, dict]) -> None:
eout = e.get("output")
if not ein:
continue
print(f"{arm:<25} "
f"{ein['mean']:>9,.0f} {ein['p50']:>9,.0f} {ein['p95']:>9,.0f} {ein['max']:>9,.0f} "
f"{(eout or {}).get('mean', 0):>9,.0f} {(eout or {}).get('p95', 0):>9,.0f}")
print(
f"{arm:<25} "
f"{ein['mean']:>9,.0f} {ein['p50']:>9,.0f} {ein['p95']:>9,.0f} {ein['max']:>9,.0f} "
f"{(eout or {}).get('mean', 0):>9,.0f} {(eout or {}).get('p95', 0):>9,.0f}"
)
def _print_pdf_var(title: str, var: dict[str, dict]) -> None:
print()
print(title)
print("-" * len(title))
header = (f"{'arm':<25} {'n_pdfs':>7} {'mean':>7} {'std':>7} {'min':>7} "
f"{'p25':>7} {'p50':>7} {'p75':>7} {'max':>7} {'#0%':>5} {'#100%':>6}")
header = (
f"{'arm':<25} {'n_pdfs':>7} {'mean':>7} {'std':>7} {'min':>7} "
f"{'p25':>7} {'p50':>7} {'p75':>7} {'max':>7} {'#0%':>5} {'#100%':>6}"
)
print(header)
print("-" * len(header))
for arm in sorted(var, key=lambda a: -var[a]["mean"]):
s = var[arm]
print(f"{arm:<25} {s['n_pdfs']:>7} "
f"{s['mean']*100:>6.1f}% {s['std']*100:>6.1f}% {s['min']*100:>6.1f}% "
f"{s['p25']*100:>6.1f}% {s['p50']*100:>6.1f}% {s['p75']*100:>6.1f}% "
f"{s['max']*100:>6.1f}% {s['n_pdfs_zero']:>5} {s['n_pdfs_perfect']:>6}")
print(
f"{arm:<25} {s['n_pdfs']:>7} "
f"{s['mean'] * 100:>6.1f}% {s['std'] * 100:>6.1f}% {s['min'] * 100:>6.1f}% "
f"{s['p25'] * 100:>6.1f}% {s['p50'] * 100:>6.1f}% {s['p75'] * 100:>6.1f}% "
f"{s['max'] * 100:>6.1f}% {s['n_pdfs_zero']:>5} {s['n_pdfs_perfect']:>6}"
)
def _print_mcnemar(title: str, table: dict) -> None:
@ -311,8 +327,10 @@ def _print_mcnemar(title: str, table: dict) -> None:
print(title)
print("-" * len(title))
print(f"n_qids on which all arms have a graded row: {table['n_qids']}")
header = (f"{'arm_i':<25} {'arm_j':<25} {'b':>4} {'c':>4} "
f"{'both ok':>8} {'both wr':>8} {'p (2-sided)':>13} {'sig':>4}")
header = (
f"{'arm_i':<25} {'arm_j':<25} {'b':>4} {'c':>4} "
f"{'both ok':>8} {'both wr':>8} {'p (2-sided)':>13} {'sig':>4}"
)
print(header)
print("-" * len(header))
for pair in sorted(table["pairs"], key=lambda p: p["p_value"]):
@ -323,10 +341,12 @@ def _print_mcnemar(title: str, table: dict) -> None:
sig = "**"
elif pair["p_value"] < 0.05:
sig = "*"
print(f"{pair['arm_i']:<25} {pair['arm_j']:<25} "
f"{pair['b_i_only']:>4} {pair['c_j_only']:>4} "
f"{pair['both_correct']:>8} {pair['both_wrong']:>8} "
f"{pair['p_value']:>13.4f} {sig:>4}")
print(
f"{pair['arm_i']:<25} {pair['arm_j']:<25} "
f"{pair['b_i_only']:>4} {pair['c_j_only']:>4} "
f"{pair['both_correct']:>8} {pair['both_wrong']:>8} "
f"{pair['p_value']:>13.4f} {sig:>4}"
)
# ---------------------------------------------------------------------------

View file

@ -78,9 +78,11 @@ def _print_table(title: str, summary: dict[str, dict]) -> None:
# stable order: highest accuracy first
arms_sorted = sorted(summary.items(), key=lambda kv: -kv[1]["accuracy"])
for arm, s in arms_sorted:
print(f"{arm:<25} {s['n']:>4} {s['n_correct']:>7} "
f"{s['accuracy']*100:>6.1f}% {s['f1_mean']*100:>6.1f}% "
f"{s['n_failures']:>6} {s['failure_rate']*100:>6.1f}%")
print(
f"{arm:<25} {s['n']:>4} {s['n_correct']:>7} "
f"{s['accuracy'] * 100:>6.1f}% {s['f1_mean'] * 100:>6.1f}% "
f"{s['n_failures']:>6} {s['failure_rate'] * 100:>6.1f}%"
)
def main() -> int:
@ -103,9 +105,7 @@ def main() -> int:
raw_rows = _read_jsonl(raw_path)
retry_rows = _read_jsonl(retry_path)
retry_by_key: dict[tuple[str, str], dict] = {
_row_key(r): r for r in retry_rows
}
retry_by_key: dict[tuple[str, str], dict] = {_row_key(r): r for r in retry_rows}
merged_rows: list[dict] = []
n_replaced_recovered = 0

View file

@ -44,10 +44,7 @@ def main() -> None:
f"questions covering first 30 docs: total={len(qs_in_30)} "
f"answerable={answerable} unanswerable={unanswerable}"
)
print(
f"avg Qs/PDF: {len(qs_in_30) / 30:.1f} "
f"answerable/PDF: {answerable / 30:.1f}"
)
print(f"avg Qs/PDF: {len(qs_in_30) / 30:.1f} answerable/PDF: {answerable / 30:.1f}")
print(f"format mix in scope: {dict(fmts)}")
print()
print("25 new PDFs to ingest:")

View file

@ -27,10 +27,7 @@ def main() -> None:
grade = a.get("graded", {})
text = (a.get("raw_text") or "").strip()
tail = text[-200:] if text else ""
print(
f" [{arm_name}] grade={grade.get('grade')} "
f"method={grade.get('method')}"
)
print(f" [{arm_name}] grade={grade.get('grade')} method={grade.get('method')}")
print(f" -> {tail!r}")

View file

@ -132,17 +132,19 @@ def _load_failed_rows(raw_path: Path) -> list[FailedRow]:
row = json.loads(line)
if not _is_failure_row(row):
continue
out.append(FailedRow(
arm=str(row["arm"]),
qid=str(row["qid"]),
doc_id=str(row["doc_id"]),
answer_format=str(row.get("answer_format") or ""),
gold=str(row.get("gold") or ""),
pages=int(row.get("pages") or 0),
document_id=row.get("document_id"),
original_error=row.get("error"),
original_row=row,
))
out.append(
FailedRow(
arm=str(row["arm"]),
qid=str(row["qid"]),
doc_id=str(row["doc_id"]),
answer_format=str(row.get("answer_format") or ""),
gold=str(row.get("gold") or ""),
pages=int(row.get("pages") or 0),
document_id=row.get("document_id"),
original_error=row.get("error"),
original_row=row,
)
)
return out
@ -202,8 +204,12 @@ def _qid_index(qid: str) -> int:
def _build_native_request(
qid: str, question: str, answer_format: str, pdf_path: Path,
*, max_output_tokens: int,
qid: str,
question: str,
answer_format: str,
pdf_path: Path,
*,
max_output_tokens: int,
) -> ArmRequest:
return ArmRequest(
question_id=qid,
@ -214,12 +220,14 @@ def _build_native_request(
def _build_lc_request(
qid: str, question: str, answer_format: str, doc_id: str, md_path: Path,
qid: str,
question: str,
answer_format: str,
doc_id: str,
md_path: Path,
) -> ArmRequest:
if not md_path.exists():
raise FileNotFoundError(
f"Missing parser extraction at {md_path}; cannot retry LC arm."
)
raise FileNotFoundError(f"Missing parser extraction at {md_path}; cannot retry LC arm.")
markdown = md_path.read_text(encoding="utf-8")
return ArmRequest(
question_id=qid,
@ -256,7 +264,9 @@ class RetryOutcome:
async def _retry_one(
arm_obj: Any, request: ArmRequest, *,
arm_obj: Any,
request: ArmRequest,
*,
arm_name: str,
qid: str,
max_attempts: int,
@ -274,31 +284,44 @@ async def _retry_one(
attempt_error = result.error
if not attempt_error and not raw_text:
attempt_error = "EmptyResponse: stream ended with no text"
attempts.append(AttemptLog(
attempt=attempt,
started_iso=started_iso,
latency_ms=latency_ms,
error=attempt_error,
raw_text_chars=len(raw_text),
))
attempts.append(
AttemptLog(
attempt=attempt,
started_iso=started_iso,
latency_ms=latency_ms,
error=attempt_error,
raw_text_chars=len(raw_text),
)
)
final = result
if not attempt_error and raw_text:
return RetryOutcome(
arm=arm_name, qid=qid, attempts=attempts,
final_result=result, recovered=True,
arm=arm_name,
qid=qid,
attempts=attempts,
final_result=result,
recovered=True,
)
if attempt < max_attempts:
delay = min(max_delay, base_delay * (2 ** (attempt - 1)))
delay = delay * (0.5 + random.random())
logger.info(
"[%s::%s] attempt %d/%d failed (%s); sleeping %.1fs",
arm_name, qid, attempt, max_attempts, attempt_error, delay,
arm_name,
qid,
attempt,
max_attempts,
attempt_error,
delay,
)
await asyncio.sleep(delay)
assert final is not None
return RetryOutcome(
arm=arm_name, qid=qid, attempts=attempts,
final_result=final, recovered=False,
arm=arm_name,
qid=qid,
attempts=attempts,
final_result=final,
recovered=False,
)
@ -365,7 +388,8 @@ async def _run(args: argparse.Namespace) -> int:
by_arm_count[f.arm] = by_arm_count.get(f.arm, 0) + 1
logger.info(
"Loaded %d failed rows across %d arms: %s",
len(failed), len(by_arm_count),
len(failed),
len(by_arm_count),
", ".join(f"{a}={n}" for a, n in sorted(by_arm_count.items())),
)
@ -383,7 +407,8 @@ async def _run(args: argparse.Namespace) -> int:
engine=PdfEngine(args.pdf_engine),
)
native_arm = NativePdfArm(
provider=native_provider, max_output_tokens=args.max_output_tokens,
provider=native_provider,
max_output_tokens=args.max_output_tokens,
)
lc_arms: dict[str, BareLlmArm] = {}
@ -413,7 +438,8 @@ async def _run(args: argparse.Namespace) -> int:
if qrow is None:
logger.error(
"Could not find question text for %s (idx %d) — skipping",
f.doc_id, q_idx,
f.doc_id,
q_idx,
)
continue
question_text = str(qrow.get("question") or "").strip()
@ -430,7 +456,10 @@ async def _run(args: argparse.Namespace) -> int:
logger.error("PDF missing on disk: %s — skipping", pdf_path)
continue
request = _build_native_request(
f.qid, question_text, answer_format, pdf_path,
f.qid,
question_text,
answer_format,
pdf_path,
max_output_tokens=args.max_output_tokens,
)
arm_obj = native_arm
@ -440,11 +469,16 @@ async def _run(args: argparse.Namespace) -> int:
if not md_path_str or ext_blob.get("status") != "ok":
logger.error(
"Missing extraction for %s on %s — cannot retry; skipping",
f.arm, f.doc_id,
f.arm,
f.doc_id,
)
continue
request = _build_lc_request(
f.qid, question_text, answer_format, f.doc_id, Path(md_path_str),
f.qid,
question_text,
answer_format,
f.doc_id,
Path(md_path_str),
)
arm_obj = lc_arms[f.arm]
else:
@ -452,13 +486,17 @@ async def _run(args: argparse.Namespace) -> int:
continue
plan.append((f, request, arm_obj))
coros.append(_retry_one(
arm_obj, request,
arm_name=f.arm, qid=f.qid,
max_attempts=args.max_attempts,
base_delay=args.base_delay,
max_delay=args.max_delay,
))
coros.append(
_retry_one(
arm_obj,
request,
arm_name=f.arm,
qid=f.qid,
max_attempts=args.max_attempts,
base_delay=args.base_delay,
max_delay=args.max_delay,
)
)
if not coros:
logger.warning("Nothing to retry after request building.")
@ -467,13 +505,17 @@ async def _run(args: argparse.Namespace) -> int:
logger.info(
"Retrying %d failed rows with up to %d attempts each "
"(base_delay=%.1fs, max_delay=%.1fs, concurrency=%d).",
len(coros), args.max_attempts, args.base_delay, args.max_delay,
len(coros),
args.max_attempts,
args.base_delay,
args.max_delay,
args.concurrency,
)
started = time.monotonic()
outcomes: list[RetryOutcome] = await _gather_with_limit(
coros, concurrency=args.concurrency,
coros,
concurrency=args.concurrency,
)
elapsed = time.monotonic() - started
logger.info("Retry pass finished in %.1fs.", elapsed)
@ -489,12 +531,8 @@ async def _run(args: argparse.Namespace) -> int:
for (f, _req, _arm_obj), outcome in zip(plan, outcomes, strict=True):
per_arm_total[outcome.arm] = per_arm_total.get(outcome.arm, 0) + 1
if outcome.recovered:
per_arm_recovered[outcome.arm] = (
per_arm_recovered.get(outcome.arm, 0) + 1
)
per_arm_attempts_dist.setdefault(outcome.arm, []).append(
len(outcome.attempts)
)
per_arm_recovered[outcome.arm] = per_arm_recovered.get(outcome.arm, 0) + 1
per_arm_attempts_dist.setdefault(outcome.arm, []).append(len(outcome.attempts))
g = grade(
pred=extract_freeform_answer(outcome.final_result.raw_text or ""),
@ -555,12 +593,11 @@ async def _run(args: argparse.Namespace) -> int:
arm: {
"tried": per_arm_total.get(arm, 0),
"recovered": per_arm_recovered.get(arm, 0),
"still_failed": (
per_arm_total.get(arm, 0) - per_arm_recovered.get(arm, 0)
),
"still_failed": (per_arm_total.get(arm, 0) - per_arm_recovered.get(arm, 0)),
"recovery_rate": (
per_arm_recovered.get(arm, 0) / per_arm_total[arm]
if per_arm_total.get(arm) else 0.0
if per_arm_total.get(arm)
else 0.0
),
"attempts_distribution": sorted(per_arm_attempts_dist.get(arm, [])),
}
@ -593,8 +630,7 @@ async def _run(args: argparse.Namespace) -> int:
rec_total = sum(per_arm_recovered.values())
rate_total = (rec_total / total * 100) if total else 0.0
print("-" * len(header))
print(f"{'TOTAL':<25} {total:>6} {rec_total:>10} {total - rec_total:>11} "
f"{rate_total:>6.1f}%")
print(f"{'TOTAL':<25} {total:>6} {rec_total:>10} {total - rec_total:>11} {rate_total:>6.1f}%")
print()
print(f"Wrote {out_path.relative_to(REPO)}")
print(f"Wrote {summary_path.relative_to(REPO)}")
@ -604,27 +640,37 @@ async def _run(args: argparse.Namespace) -> int:
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--run-id", default="2026-05-14T00-53-19Z",
"--run-id",
default="2026-05-14T00-53-19Z",
help="Run timestamp under data/multimodal_doc/runs/. Default is the "
"n=171 production run we wrote up in the blog.",
"n=171 production run we wrote up in the blog.",
)
parser.add_argument("--max-attempts", type=int, default=5)
parser.add_argument("--base-delay", type=float, default=1.0,
help="Base seconds for exponential backoff (default 1s).")
parser.add_argument("--max-delay", type=float, default=30.0,
help="Cap on per-retry sleep (default 30s).")
parser.add_argument("--concurrency", type=int, default=2,
help="Parallel retries in flight (default 2 — keep low "
"to avoid the same transport stress that caused "
"the original failures).")
parser.add_argument(
"--base-delay",
type=float,
default=1.0,
help="Base seconds for exponential backoff (default 1s).",
)
parser.add_argument(
"--max-delay", type=float, default=30.0, help="Cap on per-retry sleep (default 30s)."
)
parser.add_argument(
"--concurrency",
type=int,
default=2,
help="Parallel retries in flight (default 2 — keep low "
"to avoid the same transport stress that caused "
"the original failures).",
)
parser.add_argument("--llm-model", default="anthropic/claude-sonnet-4.5")
parser.add_argument("--pdf-engine", default="native",
choices=[e.value for e in PdfEngine])
parser.add_argument("--pdf-engine", default="native", choices=[e.value for e in PdfEngine])
parser.add_argument("--max-output-tokens", type=int, default=512)
parser.add_argument(
"--include-surfsense", action="store_true",
"--include-surfsense",
action="store_true",
help="Also retry surfsense_agentic failures (requires backend + celery up). "
"Default is to skip them since the n=171 run had 0 SurfSense failures.",
"Default is to skip them since the n=171 run had 0 SurfSense failures.",
)
args = parser.parse_args()
raise SystemExit(asyncio.run(_run(args)))

View file

@ -23,12 +23,12 @@ def main() -> None:
d = metrics[arm]
print(
f"{arm:14s}: "
f"acc={d['accuracy']*100:5.1f}% (Wilson 95% CI "
f"{d['ci_low']*100:.1f}-{d['ci_high']*100:.1f}) | "
f"correct={d['correct_rate']*100:5.1f}% "
f"missing={d['missing_rate']*100:5.1f}% "
f"incorrect={d['incorrect_rate']*100:5.1f}% | "
f"truth={d['truthfulness_score']*100:+5.1f}%"
f"acc={d['accuracy'] * 100:5.1f}% (Wilson 95% CI "
f"{d['ci_low'] * 100:.1f}-{d['ci_high'] * 100:.1f}) | "
f"correct={d['correct_rate'] * 100:5.1f}% "
f"missing={d['missing_rate'] * 100:5.1f}% "
f"incorrect={d['incorrect_rate'] * 100:5.1f}% | "
f"truth={d['truthfulness_score'] * 100:+5.1f}%"
)
print()
@ -48,7 +48,7 @@ def main() -> None:
pieces = [f"{qt:20s} (n={n:3d}):"]
for arm in ("bare_llm", "long_context", "surfsense"):
if arm in row:
pieces.append(f"{arm}={row[arm]['truthfulness_score']*100:+7.1f}%")
pieces.append(f"{arm}={row[arm]['truthfulness_score'] * 100:+7.1f}%")
print(" ".join(pieces))
print()
@ -58,7 +58,7 @@ def main() -> None:
pieces = [f"{dom:10s} (n={n:3d}):"]
for arm in ("bare_llm", "long_context", "surfsense"):
if arm in row:
pieces.append(f"{arm}={row[arm]['truthfulness_score']*100:+7.1f}%")
pieces.append(f"{arm}={row[arm]['truthfulness_score'] * 100:+7.1f}%")
print(" ".join(pieces))

View file

@ -23,7 +23,9 @@ ARTIFACT = RUN_DIR / "run_artifact.json"
def main() -> None:
rows = [json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines() if line.strip()]
rows = [
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines() if line.strip()
]
print(f"raw rows: {len(rows)}")
by_qid: dict[str, list[dict]] = defaultdict(list)
@ -31,11 +33,19 @@ def main() -> None:
by_qid[row["qid"]].append(row)
print(f"unique questions: {len(by_qid)}")
arm_metrics: dict[str, dict] = defaultdict(lambda: {
"n": 0, "n_correct": 0, "n_failed": 0, "n_empty": 0,
"costs": [], "in_tokens": [], "out_tokens": [], "latency_ms": [],
"by_format": defaultdict(lambda: {"n": 0, "correct": 0}),
})
arm_metrics: dict[str, dict] = defaultdict(
lambda: {
"n": 0,
"n_correct": 0,
"n_failed": 0,
"n_empty": 0,
"costs": [],
"in_tokens": [],
"out_tokens": [],
"latency_ms": [],
"by_format": defaultdict(lambda: {"n": 0, "correct": 0}),
}
)
for row in rows:
arm = row["arm"]
@ -70,7 +80,9 @@ def main() -> None:
print()
print("=" * 100)
print(f"{'arm':<25} {'n':>4} {'acc%':>6} {'F1%':>6} {'fail':>5} {'$ mean':>10} {'$ median':>10} {'in tok mean':>12} {'out tok mean':>12} {'p50 ms':>8}")
print(
f"{'arm':<25} {'n':>4} {'acc%':>6} {'F1%':>6} {'fail':>5} {'$ mean':>10} {'$ median':>10} {'in tok mean':>12} {'out tok mean':>12} {'p50 ms':>8}"
)
print("=" * 100)
art = json.loads(ARTIFACT.read_text(encoding="utf-8"))
per_arm_art = art["metrics"]["per_arm"]
@ -110,7 +122,7 @@ def main() -> None:
print("Aggregated cost (from run_artifact.json):")
for arm, row in per_arm_art.items():
print(
f" {arm:<25} acc={row['accuracy']*100:5.1f}% "
f" {arm:<25} acc={row['accuracy'] * 100:5.1f}% "
f" $/Q LLM={row['llm_cost_per_q']:.4f} "
f" preprocess total=${row['preprocess_cost_total']:.2f} "
f" $/Q total={row['total_cost_per_q']:.4f}"

View file

@ -40,8 +40,7 @@ CONTEXT_HINTS = (
def main() -> None:
rows = [
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines()
if line.strip()
json.loads(line) for line in RAW.read_text(encoding="utf-8").splitlines() if line.strip()
]
extraction_size: dict[tuple[str, str], int] = {}
@ -73,12 +72,12 @@ def main() -> None:
print("=" * 80)
print("(b) Extraction size for OK vs FAILED rows per arm")
print("=" * 80)
arm_buckets: dict[str, dict[str, list[int]]] = defaultdict(
lambda: {"ok": [], "fail": []}
)
arm_buckets: dict[str, dict[str, list[int]]] = defaultdict(lambda: {"ok": [], "fail": []})
parser_arms = (
"azure_basic_lc", "azure_premium_lc",
"llamacloud_basic_lc", "llamacloud_premium_lc",
"azure_basic_lc",
"azure_premium_lc",
"llamacloud_basic_lc",
"llamacloud_premium_lc",
)
for row in rows:
arm = row["arm"]
@ -133,10 +132,13 @@ def main() -> None:
" 3M_2018_10K x llamacloud_premium = 908,733 chars (~227k tokens) "
"-- this is above Sonnet 4.5's 200k window."
)
print(" If transport hypothesis is correct, this should still fail with a "
"real overflow error.")
print(" If transport hypothesis is correct AND the model truncates silently, "
"it might 'succeed' but be wrong.")
print(
" If transport hypothesis is correct, this should still fail with a real overflow error."
)
print(
" If transport hypothesis is correct AND the model truncates silently, "
"it might 'succeed' but be wrong."
)
print()
for row in rows:
if row["doc_id"] != "3M_2018_10K.pdf":
@ -145,10 +147,7 @@ def main() -> None:
continue
err = row.get("error") or "(none)"
graded = row.get("graded") or {}
print(
f" {row['qid']:<40} correct={graded.get('correct')!s:<5} "
f"err={err[:100]}"
)
print(f" {row['qid']:<40} correct={graded.get('correct')!s:<5} err={err[:100]}")
if __name__ == "__main__":

View file

@ -72,9 +72,7 @@ class SurfSenseArm(Arm):
try:
await self._client.delete_thread(thread_id)
except Exception as exc: # noqa: BLE001
logger.debug(
"Failed to delete thread %s: %s", thread_id, exc
)
logger.debug("Failed to delete thread %s: %s", thread_id, exc)
letter = extract_answer_letter(answer.text)
return ArmResult(

View file

@ -83,6 +83,7 @@ async def acquire_token(config: Config, *, http: httpx.AsyncClient | None = None
)
if config.has_local_mode():
async def _login(client: httpx.AsyncClient) -> TokenBundle:
response = await client.post(
f"{config.surfsense_api_base}/auth/desktop/login",
@ -94,15 +95,12 @@ async def acquire_token(config: Config, *, http: httpx.AsyncClient | None = None
)
if response.status_code != 200:
raise CredentialError(
f"LOCAL login failed (HTTP {response.status_code}): "
f"{_safe_text(response)}"
f"LOCAL login failed (HTTP {response.status_code}): {_safe_text(response)}"
)
payload = response.json()
access = payload.get("access_token")
if not access:
raise CredentialError(
f"LOCAL login response missing access_token: {payload!r}"
)
raise CredentialError(f"LOCAL login response missing access_token: {payload!r}")
return TokenBundle(
access_token=access,
refresh_token=payload.get("refresh_token") or None,

View file

@ -204,8 +204,7 @@ async def _cmd_setup(args: argparse.Namespace) -> int:
if scenario not in SCENARIOS:
console.print(
f"[red]Unknown scenario {scenario!r}. Pick one of: "
f"{', '.join(SCENARIOS)}[/red]"
f"[red]Unknown scenario {scenario!r}. Pick one of: {', '.join(SCENARIOS)}[/red]"
)
return 2
@ -292,9 +291,7 @@ async def _cmd_setup(args: argparse.Namespace) -> int:
if not skip_vision_setup and (vision_required or vision_llm_slug is not None):
try:
vision_candidates = await ss_client.list_global_vision_models()
resolved = resolve_vision_llm(
vision_candidates, explicit_slug=vision_llm_slug
)
resolved = resolve_vision_llm(vision_candidates, explicit_slug=vision_llm_slug)
except VisionConfigError as exc:
console.print(f"[red]{exc}[/red]")
return 2
@ -524,10 +521,7 @@ async def _cmd_run(args: argparse.Namespace) -> int:
)
artifact = await benchmark.run(ctx, **extra_kwargs)
console.print(
f"[green]run OK[/green] {args.suite}/{args.benchmark}"
f"{artifact.raw_path}"
)
console.print(f"[green]run OK[/green] {args.suite}/{args.benchmark}{artifact.raw_path}")
return 0
@ -697,15 +691,21 @@ def _build_parser() -> argparse.ArgumentParser:
)
p_setup.set_defaults(_func=_cmd_setup, _async=True)
p_teardown = sub.add_parser("teardown", help="Soft-delete the suite SearchSpace + clear state slot.")
p_teardown = sub.add_parser(
"teardown", help="Soft-delete the suite SearchSpace + clear state slot."
)
p_teardown.add_argument("--suite", required=True)
p_teardown.set_defaults(_func=_cmd_teardown, _async=True)
p_models = sub.add_parser("models", help="LLM-config discovery helpers.")
models_sub = p_models.add_subparsers(dest="subcommand", required=True)
p_models_list = models_sub.add_parser("list", help="List global LLM configs.")
p_models_list.add_argument("--provider", default=None, help="Filter by provider, e.g. openrouter")
p_models_list.add_argument("--grep", default=None, help="Substring filter on name / model_name.")
p_models_list.add_argument(
"--provider", default=None, help="Filter by provider, e.g. openrouter"
)
p_models_list.add_argument(
"--grep", default=None, help="Substring filter on name / model_name."
)
p_models_list.set_defaults(_func=_cmd_models_list, _async=True)
p_suites = sub.add_parser("suites", help="List registered suites.")
@ -729,7 +729,9 @@ def _build_parser() -> argparse.ArgumentParser:
suite_parser = ingest_sub.add_parser(suite, help=f"Ingest a {suite} benchmark.")
suite_bench = suite_parser.add_subparsers(dest="benchmark", required=True)
for benchmark in registry.list_benchmarks(suite):
bp = suite_bench.add_parser(benchmark.name, help=getattr(benchmark, "description", benchmark.name))
bp = suite_bench.add_parser(
benchmark.name, help=getattr(benchmark, "description", benchmark.name)
)
if hasattr(benchmark, "add_run_args"):
benchmark.add_run_args(bp)
bp.set_defaults(_func=_cmd_ingest, _async=True)
@ -740,7 +742,9 @@ def _build_parser() -> argparse.ArgumentParser:
suite_parser = run_sub.add_parser(suite, help=f"Run a {suite} benchmark.")
suite_bench = suite_parser.add_subparsers(dest="benchmark", required=True)
for benchmark in registry.list_benchmarks(suite):
bp = suite_bench.add_parser(benchmark.name, help=getattr(benchmark, "description", benchmark.name))
bp = suite_bench.add_parser(
benchmark.name, help=getattr(benchmark, "description", benchmark.name)
)
if hasattr(benchmark, "add_run_args"):
benchmark.add_run_args(bp)
bp.set_defaults(_func=_cmd_run, _async=True)

View file

@ -84,8 +84,7 @@ class DocumentProcessingFailed(RuntimeError):
def __init__(self, statuses: Sequence[DocumentStatus]) -> None:
details = ", ".join(
f"id={s.document_id} ({s.title!r}): {s.reason or 'unknown'}"
for s in statuses
f"id={s.document_id} ({s.title!r}): {s.reason or 'unknown'}" for s in statuses
)
super().__init__(f"Document(s) failed to process: {details}")
self.statuses = list(statuses)
@ -240,9 +239,7 @@ class DocumentsClient:
# chunks (chunk_id -> document_id map)
# ------------------------------------------------------------------
async def list_chunks(
self, document_id: int, *, page_size: int = 100
) -> list[ChunkRow]:
async def list_chunks(self, document_id: int, *, page_size: int = 100) -> list[ChunkRow]:
"""Walk ``GET /documents/{id}/chunks`` until ``has_more=False``.
Used by ingestion to materialise the ``chunk_id -> document_id``

View file

@ -145,7 +145,7 @@ class NewChatClient:
if attempt > max_busy_retries:
raise
# Cap wait at 30s; backend retry hint is exponential anyway.
wait = min(30.0, 0.5 * (2 ** attempt))
wait = min(30.0, 0.5 * (2**attempt))
logger.info(
"thread_id=%s busy (%s); retry %d/%d after %.1fs",
thread_id,

View file

@ -177,16 +177,12 @@ class SearchSpaceClient:
response.raise_for_status()
payload = response.json()
if not isinstance(payload, list):
raise RuntimeError(
f"Unexpected /model-connections/global payload: {payload!r}"
)
raise RuntimeError(f"Unexpected /model-connections/global payload: {payload!r}")
entries: list[VisionModelEntry] = []
for connection in payload:
provider = str(connection.get("provider", ""))
for model in connection.get("models") or []:
if not model.get("enabled", True) or not model.get("supports_image_input"):
continue
entries.append(
VisionModelEntry.from_payload({**model, "provider": provider})
)
entries.append(VisionModelEntry.from_payload({**model, "provider": provider}))
return entries

View file

@ -104,7 +104,9 @@ def load_config() -> Config:
data_dir = Path(os.environ.get("EVAL_DATA_DIR") or (project_root / "data")).resolve()
reports_dir = Path(os.environ.get("EVAL_REPORTS_DIR") or (project_root / "reports")).resolve()
return Config(
surfsense_api_base=os.environ.get("SURFSENSE_API_BASE", "http://localhost:8000").rstrip("/"),
surfsense_api_base=os.environ.get("SURFSENSE_API_BASE", "http://localhost:8000").rstrip(
"/"
),
openrouter_api_key=os.environ.get("OPENROUTER_API_KEY") or None,
openrouter_base_url=os.environ.get(
"OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1"
@ -203,9 +205,7 @@ class SuiteState:
else None
),
native_arm_model=(
str(payload["native_arm_model"])
if payload.get("native_arm_model")
else None
str(payload["native_arm_model"]) if payload.get("native_arm_model") else None
),
)

View file

@ -95,10 +95,7 @@ class IngestSettings:
def render_label(self) -> str:
"""Human-readable single-line label for reports / log lines."""
return (
f"vision={'on' if self.use_vision_llm else 'off'}, "
f"mode={self.processing_mode}"
)
return f"vision={'on' if self.use_vision_llm else 'off'}, mode={self.processing_mode}"
def _coerce_bool(value: Any, default: bool) -> bool:
@ -122,9 +119,7 @@ def _coerce_mode(value: Any, default: str) -> str:
return default
val = str(value).strip().lower()
if val not in PROCESSING_MODES:
raise ValueError(
f"Invalid processing_mode {val!r}; must be one of {PROCESSING_MODES}"
)
raise ValueError(f"Invalid processing_mode {val!r}; must be one of {PROCESSING_MODES}")
return val
@ -274,10 +269,7 @@ def format_ingest_settings_md(settings: Any) -> str:
return "- SurfSense ingest settings: (not recorded — re-ingest to capture)"
vision = "on" if settings.get("use_vision_llm") else "off"
mode = settings.get("processing_mode") or "basic"
return (
f"- SurfSense ingest settings: vision_llm=`{vision}`, "
f"processing_mode=`{mode}`"
)
return f"- SurfSense ingest settings: vision_llm=`{vision}`, processing_mode=`{mode}`"
__all__ = [

View file

@ -67,17 +67,13 @@ def mcnemar_test(
"""
if len(arm_a_correct) != len(arm_b_correct):
raise ValueError(
f"Length mismatch: arm_a={len(arm_a_correct)}, arm_b={len(arm_b_correct)}"
)
raise ValueError(f"Length mismatch: arm_a={len(arm_a_correct)}, arm_b={len(arm_b_correct)}")
n = len(arm_a_correct)
b = sum(1 for a, c in zip(arm_a_correct, arm_b_correct, strict=False) if a and not c)
c = sum(1 for a, cc in zip(arm_a_correct, arm_b_correct, strict=False) if (not a) and cc)
discordant = b + c
if discordant == 0:
return McnemarResult(
n_total=n, b=b, c=c, statistic=0.0, p_value=1.0, method="degenerate"
)
return McnemarResult(n_total=n, b=b, c=c, statistic=0.0, p_value=1.0, method="degenerate")
if discordant < use_exact_below:
# Exact binomial: under H0 each discordant pair is a Bernoulli(0.5).
@ -92,13 +88,11 @@ def mcnemar_test(
# Chi-square with continuity correction (McNemar-Edwards).
chi = ((abs(b - c) - 1) ** 2) / discordant
p_value = _chi2_sf(chi, df=1)
return McnemarResult(
n_total=n, b=b, c=c, statistic=chi, p_value=p_value, method="chi2_cc"
)
return McnemarResult(n_total=n, b=b, c=c, statistic=chi, p_value=p_value, method="chi2_cc")
def _binom_pmf(n: int, k: int) -> float:
return math.comb(n, k) * (0.5 ** n)
return math.comb(n, k) * (0.5**n)
def _chi2_sf(x: float, *, df: int) -> float:

View file

@ -46,9 +46,7 @@ _Z_FOR_LEVEL: dict[float, float] = {
}
def wilson_ci(
n_correct: int, n_total: int, *, level: float = 0.95
) -> tuple[float, float]:
def wilson_ci(n_correct: int, n_total: int, *, level: float = 0.95) -> tuple[float, float]:
"""Two-sided Wilson score confidence interval for a proportion.
Returns ``(low, high)``. ``n_total == 0`` returns ``(0.0, 1.0)``
@ -70,9 +68,7 @@ def wilson_ci(
return low, high
def accuracy_with_wilson_ci(
n_correct: int, n_total: int, *, level: float = 0.95
) -> AccuracyResult:
def accuracy_with_wilson_ci(n_correct: int, n_total: int, *, level: float = 0.95) -> AccuracyResult:
if n_total < 0:
raise ValueError(f"n_total must be >= 0, got {n_total}")
if n_correct < 0 or n_correct > n_total:
@ -109,10 +105,7 @@ def per_task_accuracy(
bucket[1] += 1
if row.get(correct_key):
bucket[0] += 1
return {
task: accuracy_with_wilson_ci(c[0], c[1], level=level)
for task, c in counts.items()
}
return {task: accuracy_with_wilson_ci(c[0], c[1], level=level) for task, c in counts.items()}
def macro_accuracy(per_task: Mapping[str, AccuracyResult]) -> float:

View file

@ -61,7 +61,7 @@ def _dcg_at_k(grades: Sequence[float], k: int) -> float:
s = 0.0
for i, grade in enumerate(grades[:k], start=1):
# Standard log-base-2 discount; gain = 2^grade - 1 for graded relevance.
s += (2.0 ** grade - 1.0) / math.log2(i + 1)
s += (2.0**grade - 1.0) / math.log2(i + 1)
return s
@ -106,7 +106,9 @@ def score_run(
qids = set(per_query_qrels.keys()) & set(per_query_retrieved.keys())
if not qids:
return RetrievalScores(recall_at_k={k: 0.0 for k in ks}, mrr=0.0, ndcg_at_10=0.0, n_queries=0)
return RetrievalScores(
recall_at_k={k: 0.0 for k in ks}, mrr=0.0, ndcg_at_10=0.0, n_queries=0
)
recall_totals = {k: 0.0 for k in ks}
mrr_total = 0.0

View file

@ -35,7 +35,7 @@ from typing import Any
# the pattern source, so we splice the literal character in via an
# f-string. This keeps our pattern functionally identical to the TS
# reference and lets ``"\u200B" in CITATION_REGEX.pattern`` succeed.
_ZWSP = "\u200B"
_ZWSP = "\u200b"
CITATION_REGEX = re.compile(
rf"[\[【]{_ZWSP}?citation:\s*("
rf"https?://[^\]】{_ZWSP}]+|urlcite\d+|(?:doc-)?-?\d+(?:\s*,\s*(?:doc-)?-?\d+)*"

View file

@ -56,7 +56,7 @@ def extract_freeform_answer(text: str) -> str:
marker_matches = list(_ANSWER_MARKER.finditer(text))
if marker_matches:
last = marker_matches[-1]
tail = text[last.end():]
tail = text[last.end() :]
nl = tail.find("\n")
if nl >= 0:
tail = tail[:nl]
@ -77,7 +77,7 @@ def extract_freeform_answer(text: str) -> str:
# 2. Strip wrapping quotes / parens / trailing punctuation that
# confuse the grader without changing meaning.
candidate = candidate.strip().strip("`").strip()
if candidate.startswith(("\"", "'")) and candidate.endswith(("\"", "'")):
if candidate.startswith(('"', "'")) and candidate.endswith(('"', "'")):
candidate = candidate[1:-1].strip()
return candidate

View file

@ -64,8 +64,7 @@ async def parse_with_azure_di(
api_key = api_key or os.environ.get("AZURE_DI_KEY")
if not endpoint or not api_key:
raise ValueError(
"AZURE_DI_ENDPOINT and AZURE_DI_KEY must be set "
"(see surfsense_evals/.env)."
"AZURE_DI_ENDPOINT and AZURE_DI_KEY must be set (see surfsense_evals/.env)."
)
model_id = _AZURE_MODEL_BY_MODE.get(processing_mode, "prebuilt-read")
@ -86,7 +85,10 @@ async def parse_with_azure_di(
file_size_mb = await asyncio.to_thread(os.path.getsize, file_path) / (1024 * 1024)
logger.info(
"Azure DI parsing %s (mode=%s, model=%s, size=%.1fMB)",
file_path, processing_mode, model_id, file_size_mb,
file_path,
processing_mode,
model_id,
file_size_mb,
)
last_exc: Exception | None = None
@ -106,12 +108,12 @@ async def parse_with_azure_di(
result = await poller.result()
content = (result.content or "").strip()
if not content:
raise AzureDIError(
f"Azure DI returned empty content for {file_path}"
)
raise AzureDIError(f"Azure DI returned empty content for {file_path}")
logger.info(
"Azure DI OK: %s (%s) -> %d chars",
file_path, model_id, len(content),
file_path,
model_id,
len(content),
)
return content
@ -120,9 +122,7 @@ async def parse_with_azure_di(
except HttpResponseError as exc:
# 4xx that's not auth: don't retry, the request itself is broken.
if exc.status_code and 400 <= exc.status_code < 500:
raise AzureDIError(
f"Azure DI {exc.status_code} on {file_path}: {exc}"
) from exc
raise AzureDIError(f"Azure DI {exc.status_code} on {file_path}: {exc}") from exc
last_exc = exc
except (ServiceRequestError, ServiceResponseError) as exc:
last_exc = exc
@ -133,7 +133,10 @@ async def parse_with_azure_di(
sleep_for = delay + jitter
logger.warning(
"Azure DI attempt %d/%d failed (%s); retrying in %.1fs",
attempt, _MAX_RETRIES, type(last_exc).__name__, sleep_for,
attempt,
_MAX_RETRIES,
type(last_exc).__name__,
sleep_for,
)
await asyncio.sleep(sleep_for)

View file

@ -61,8 +61,7 @@ def _extract_markdown(result) -> str:
if result and hasattr(result[0], "text"):
return result[0].text
return "\n\n".join(
doc.page_content if hasattr(doc, "page_content") else str(doc)
for doc in result
doc.page_content if hasattr(doc, "page_content") else str(doc) for doc in result
)
return str(result)
@ -86,9 +85,7 @@ async def parse_with_llamacloud(
api_key = api_key or os.environ.get("LLAMA_CLOUD_API_KEY")
if not api_key:
raise ValueError(
"LLAMA_CLOUD_API_KEY must be set (see surfsense_evals/.env)."
)
raise ValueError("LLAMA_CLOUD_API_KEY must be set (see surfsense_evals/.env).")
parse_mode = _LLAMA_PARSE_MODE_MAP.get(processing_mode, "parse_page_with_llm")
@ -106,13 +103,19 @@ async def parse_with_llamacloud(
upload_timeout = max(120.0, 30.0 * file_size_mb)
logger.info(
"LlamaCloud parsing %s (mode=%s, parse_mode=%s, %.1fMB, "
"job_timeout=%.0fs)",
file_path, processing_mode, parse_mode, file_size_mb, job_timeout,
"LlamaCloud parsing %s (mode=%s, parse_mode=%s, %.1fMB, job_timeout=%.0fs)",
file_path,
processing_mode,
parse_mode,
file_size_mb,
job_timeout,
)
custom_timeout = httpx.Timeout(
connect=120.0, read=upload_timeout, write=upload_timeout, pool=120.0,
connect=120.0,
read=upload_timeout,
write=upload_timeout,
pool=120.0,
)
last_exc: Exception | None = None
@ -135,12 +138,12 @@ async def parse_with_llamacloud(
result = await parser.aparse(str(file_path))
content = _extract_markdown(result).strip()
if not content:
raise LlamaCloudError(
f"LlamaCloud returned empty content for {file_path}"
)
raise LlamaCloudError(f"LlamaCloud returned empty content for {file_path}")
logger.info(
"LlamaCloud OK: %s (%s) -> %d chars",
file_path, parse_mode, len(content),
file_path,
parse_mode,
len(content),
)
return content
@ -156,7 +159,10 @@ async def parse_with_llamacloud(
sleep_for = delay + jitter
logger.warning(
"LlamaCloud attempt %d/%d failed (%s); retrying in %.1fs",
attempt, _MAX_RETRIES, type(last_exc).__name__, sleep_for,
attempt,
_MAX_RETRIES,
type(last_exc).__name__,
sleep_for,
)
await asyncio.sleep(sleep_for)

View file

@ -116,11 +116,7 @@ def _normalise_paragraphs(text: str) -> list[str]:
def _escape_html(text: str) -> str:
return (
text.replace("&", "&amp;")
.replace("<", "&lt;")
.replace(">", "&gt;")
)
return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
def render_pdf(

View file

@ -121,9 +121,7 @@ class OpenRouterPdfProvider:
body: dict[str, Any] = {
"model": self._model,
"messages": messages,
"plugins": [
{"id": "file-parser", "pdf": {"engine": self._engine.value}}
],
"plugins": [{"id": "file-parser", "pdf": {"engine": self._engine.value}}],
}
if max_tokens:
body["max_tokens"] = max_tokens

View file

@ -177,7 +177,9 @@ class Benchmark(Protocol):
def add_run_args(self, parser: argparse.ArgumentParser) -> None: # pragma: no cover - protocol
"""Add benchmark-specific flags to ``run <suite> <benchmark>``."""
def report_section(self, artifacts: list[RunArtifact]) -> ReportSection: # pragma: no cover - protocol
def report_section(
self, artifacts: list[RunArtifact]
) -> ReportSection: # pragma: no cover - protocol
...
@ -224,9 +226,7 @@ def get(suite: str, name: str) -> Benchmark:
return _REGISTRY[(suite, name)]
except KeyError as exc:
available = ", ".join(f"{s}/{n}" for s, n in sorted(_REGISTRY)) or "<none>"
raise KeyError(
f"Unknown benchmark '{suite}/{name}'. Registered: {available}"
) from exc
raise KeyError(f"Unknown benchmark '{suite}/{name}'. Registered: {available}") from exc
def list_suites() -> list[str]:

View file

@ -45,10 +45,7 @@ def format_scenario_md(extra: Mapping[str, Any] | None) -> str:
"(text-only model can't see images) — that's the point."
)
else:
body = (
f"- Scenario: head-to-head — both arms answer with `{surf_slug}` "
"via OpenRouter."
)
body = f"- Scenario: head-to-head — both arms answer with `{surf_slug}` via OpenRouter."
if vision_slug:
body += f" SurfSense ingest VLM: `{vision_slug}`."

View file

@ -60,7 +60,5 @@ def discover_suites() -> list[str]:
importlib.import_module(benchmark_name)
imported.append(benchmark_name)
except Exception as exc: # noqa: BLE001
logger.warning(
"Failed to import benchmark %s: %s", benchmark_name, exc
)
logger.warning("Failed to import benchmark %s: %s", benchmark_name, exc)
return imported

View file

@ -154,13 +154,11 @@ async def run_ingest(
if not batches:
logger.warning("Discipline %s produced 0 batches; skipping upload", discipline)
continue
logger.info(
"Uploading %d batches for discipline %s", len(batches), discipline
)
logger.info("Uploading %d batches for discipline %s", len(batches), discipline)
upload_result = await docs_client.upload(
files=[b.path for b in batches],
search_space_id=ctx.search_space_id,
use_vision_llm=settings.use_vision_llm,
use_vision_llm=settings.use_vision_llm,
processing_mode=settings.processing_mode,
)
new_doc_ids = list(upload_result.document_ids)
@ -177,9 +175,7 @@ async def run_ingest(
)
title_to_doc = {s.title: s.document_id for s in statuses}
per_discipline_path = (
ctx.maps_dir() / f"cure_corpus_map_{discipline}.jsonl"
)
per_discipline_path = ctx.maps_dir() / f"cure_corpus_map_{discipline}.jsonl"
with per_discipline_path.open("w", encoding="utf-8") as fh:
fh.write(settings_header_line(settings) + "\n")
for batch in batches:
@ -202,9 +198,7 @@ async def run_ingest(
try:
chunks = await docs_client.list_chunks(int(doc_id))
except Exception as exc: # noqa: BLE001
logger.warning(
"Failed to list chunks for doc_id=%s: %s", doc_id, exc
)
logger.warning("Failed to list chunks for doc_id=%s: %s", doc_id, exc)
continue
for chunk in chunks:
fh.write(

View file

@ -191,12 +191,15 @@ class CureBenchmark:
def add_run_args(self, parser: argparse.ArgumentParser) -> None:
parser.add_argument("--lang", default="en", choices=("en", "es", "fr"))
parser.add_argument("--discipline", default=None,
help="Restrict to one discipline (default: all ingested).")
parser.add_argument(
"--discipline", default=None, help="Restrict to one discipline (default: all ingested)."
)
parser.add_argument("--n", dest="sample_n", type=int, default=None)
parser.add_argument("--concurrency", type=int, default=4)
parser.add_argument(
"--max-passages-per-discipline", type=int, default=None,
"--max-passages-per-discipline",
type=int,
default=None,
help="(ingest only) cap corpus rows per discipline for smoke testing.",
)
# Per-upload knobs forwarded to /documents/fileupload at ingest;
@ -233,11 +236,13 @@ class CureBenchmark:
# Disciplines to query are determined by the per-discipline maps
# actually present (either user-filtered or whatever was ingested).
ingested_disciplines = sorted({
row_disc
for path in maps_dir.glob("cure_corpus_map_*.jsonl")
for row_disc in [path.stem[len("cure_corpus_map_"):]]
})
ingested_disciplines = sorted(
{
row_disc
for path in maps_dir.glob("cure_corpus_map_*.jsonl")
for row_disc in [path.stem[len("cure_corpus_map_") :]]
}
)
if discipline_filter:
disciplines = [discipline_filter]
else:

View file

@ -55,15 +55,15 @@ def _hf_hub_download(*args, **kwargs):
@dataclass
class MedXpertQuestion:
qid: str # e.g. "MM-26"
question: str # full question text (case + ask)
options: dict[str, str] # A-E
label: str # "A".."E"
image_files: list[str] # filenames inside images.zip
qid: str # e.g. "MM-26"
question: str # full question text (case + ask)
options: dict[str, str] # A-E
label: str # "A".."E"
image_files: list[str] # filenames inside images.zip
medical_task: str
body_system: str
question_type: str
split: str # "test" or "dev"
split: str # "test" or "dev"
def _load_jsonl(path: Path, *, split: str) -> list[MedXpertQuestion]:
@ -84,17 +84,19 @@ def _load_jsonl(path: Path, *, split: str) -> list[MedXpertQuestion]:
images = row.get("images") or []
if not isinstance(images, list):
images = []
out.append(MedXpertQuestion(
qid=qid,
question=question,
options=opts,
label=label,
image_files=[str(x).strip() for x in images if str(x).strip()],
medical_task=str(row.get("medical_task") or "").strip(),
body_system=str(row.get("body_system") or "").strip(),
question_type=str(row.get("question_type") or "").strip(),
split=split,
))
out.append(
MedXpertQuestion(
qid=qid,
question=question,
options=opts,
label=label,
image_files=[str(x).strip() for x in images if str(x).strip()],
medical_task=str(row.get("medical_task") or "").strip(),
body_system=str(row.get("body_system") or "").strip(),
question_type=str(row.get("question_type") or "").strip(),
split=split,
)
)
return out
@ -204,7 +206,7 @@ async def _upload_pdfs(
name_to_id: dict[str, int] = {}
pdf_list = list(pdf_paths)
for batch_start in range(0, len(pdf_list), batch_size):
batch = pdf_list[batch_start:batch_start + batch_size]
batch = pdf_list[batch_start : batch_start + batch_size]
result = await docs_client.upload(
files=batch,
search_space_id=ctx.search_space_id,
@ -226,8 +228,10 @@ async def _upload_pdfs(
name_to_id[s.title] = s.document_id
logger.info(
"Uploaded MedXpertQA batch %d-%d: %d new, %d duplicate",
batch_start, batch_start + len(batch),
len(result.document_ids), len(result.duplicate_document_ids),
batch_start,
batch_start + len(batch),
len(result.document_ids),
len(result.duplicate_document_ids),
)
return name_to_id
@ -310,9 +314,11 @@ async def run_ingest(
# Materialise into bench_dir so the path is stable.
try:
from os import link as _link
_link(local_zip, images_zip_local)
except OSError:
from shutil import copy2
copy2(local_zip, images_zip_local)
_ensure_images_extracted(images_zip_local, images_dir)
@ -354,17 +360,22 @@ async def run_ingest(
questions_jsonl = bench_dir / "questions.jsonl"
with questions_jsonl.open("w", encoding="utf-8") as fh:
for q in questions:
fh.write(json.dumps({
"qid": q.qid,
"question": q.question,
"options": q.options,
"label": q.label,
"image_files": q.image_files,
"medical_task": q.medical_task,
"body_system": q.body_system,
"question_type": q.question_type,
"split": q.split,
}) + "\n")
fh.write(
json.dumps(
{
"qid": q.qid,
"question": q.question,
"options": q.options,
"label": q.label,
"image_files": q.image_files,
"medical_task": q.medical_task,
"body_system": q.body_system,
"question_type": q.question_type,
"split": q.split,
}
)
+ "\n"
)
logger.info("Wrote %d MedXpertQA questions to %s", len(questions), questions_jsonl)
map_path = ctx.maps_dir() / "medxpertqa_doc_map.jsonl"
@ -376,13 +387,18 @@ async def run_ingest(
local = pdf_paths.get(q.qid)
if local is None:
continue
fh.write(json.dumps({
"qid": q.qid,
"document_id": name_to_id.get(local.name),
"pdf_path": str(local),
"n_images": len(q.image_files),
"split": q.split,
}) + "\n")
fh.write(
json.dumps(
{
"qid": q.qid,
"document_id": name_to_id.get(local.name),
"pdf_path": str(local),
"n_images": len(q.image_files),
"split": q.split,
}
)
+ "\n"
)
logger.info("Wrote MedXpertQA doc map to %s", map_path)
new_state = ctx.suite_state

View file

@ -129,19 +129,21 @@ def _load_questions(
n_images = int(map_row.get("n_images", 0))
if require_images and n_images <= 0:
continue
out.append(MXQuestion(
qid=qid,
question=str(row.get("question") or ""),
options={str(k).upper(): str(v) for k, v in (row.get("options") or {}).items()},
label=str(row.get("label") or "").strip().upper(),
medical_task=str(row.get("medical_task") or "").strip(),
body_system=str(row.get("body_system") or "").strip(),
question_type=str(row.get("question_type") or "").strip(),
split=str(row.get("split") or ""),
n_images=n_images,
pdf_path=Path(map_row["pdf_path"]),
document_id=map_row.get("document_id"),
))
out.append(
MXQuestion(
qid=qid,
question=str(row.get("question") or ""),
options={str(k).upper(): str(v) for k, v in (row.get("options") or {}).items()},
label=str(row.get("label") or "").strip().upper(),
medical_task=str(row.get("medical_task") or "").strip(),
body_system=str(row.get("body_system") or "").strip(),
question_type=str(row.get("question_type") or "").strip(),
split=str(row.get("split") or ""),
n_images=n_images,
pdf_path=Path(map_row["pdf_path"]),
document_id=map_row.get("document_id"),
)
)
out.sort(key=lambda q: (q.split, q.qid))
if sample_n is not None and sample_n > 0:
out = out[:sample_n]
@ -182,51 +184,81 @@ class MedXpertQAMMBenchmark:
def add_run_args(self, parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--split", default="test", choices=["test", "dev", "all"],
"--split",
default="test",
choices=["test", "dev", "all"],
help="Which MedXpertQA-MM split to run (default: test).",
)
parser.add_argument(
"--task", default="all",
"--task",
default="all",
help="Filter by medical_task value (e.g. Diagnosis, Treatment, Basic Medicine).",
)
parser.add_argument(
"--body-system", dest="body_filter", default="all",
"--body-system",
dest="body_filter",
default="all",
help="Filter by body_system value (e.g. Cardiovascular, Lymphatic).",
)
parser.add_argument(
"--require-images", dest="require_images", action="store_true",
"--require-images",
dest="require_images",
action="store_true",
help="Skip rare MM rows that ended up with zero resolvable images.",
)
parser.add_argument("--n", dest="sample_n", type=int, default=None,
help="Run only the first N questions after filters apply.")
parser.add_argument("--concurrency", type=int, default=4,
help="Parallel question workers per arm.")
parser.add_argument("--no-mentions", dest="no_mentions", action="store_true",
help="SurfSense arm: skip mentioned_document_ids (unscoped retrieval).")
parser.add_argument(
"--pdf-engine", default="native",
"--n",
dest="sample_n",
type=int,
default=None,
help="Run only the first N questions after filters apply.",
)
parser.add_argument(
"--concurrency", type=int, default=4, help="Parallel question workers per arm."
)
parser.add_argument(
"--no-mentions",
dest="no_mentions",
action="store_true",
help="SurfSense arm: skip mentioned_document_ids (unscoped retrieval).",
)
parser.add_argument(
"--pdf-engine",
default="native",
choices=[e.value for e in PdfEngine],
help="OpenRouter file-parser engine for the native arm.",
)
parser.add_argument(
"--max-output-tokens", type=int, default=512,
"--max-output-tokens",
type=int,
default=512,
help="Cap on completion length for both arms.",
)
# Ingest-only knobs (forwarded by the CLI to ingest.run_ingest).
parser.add_argument(
"--max-questions", dest="max_questions", type=int, default=None,
"--max-questions",
dest="max_questions",
type=int,
default=None,
help="(ingest only) cap on number of MM questions to render + upload.",
)
parser.add_argument(
"--upload-batch-size", dest="upload_batch_size", type=int, default=8,
"--upload-batch-size",
dest="upload_batch_size",
type=int,
default=8,
help="(ingest only) PDFs per fileupload call.",
)
parser.add_argument(
"--skip-upload", dest="skip_upload", action="store_true",
"--skip-upload",
dest="skip_upload",
action="store_true",
help="(ingest only) render PDFs locally but don't push to SurfSense.",
)
parser.add_argument(
"--include-dev", dest="include_dev", action="store_true",
"--include-dev",
dest="include_dev",
action="store_true",
help="(ingest only) shorthand for --split all.",
)
# Per-upload knobs forwarded to /documents/fileupload at ingest;
@ -270,7 +302,8 @@ class MedXpertQAMMBenchmark:
doc_map, ingest_settings = _load_doc_map(map_path)
questions = _load_questions(
questions_jsonl, doc_map,
questions_jsonl,
doc_map,
split_filter=split_filter,
task_filter=task_filter if task_filter != "all" else None,
body_filter=body_filter if body_filter != "all" else None,
@ -378,13 +411,18 @@ class MedXpertQAMMBenchmark:
manifest_path = run_dir / "run_artifact.json"
manifest_path.write_text(
json.dumps({
"suite": self.suite,
"benchmark": self.name,
"raw_path": "raw.jsonl",
"metrics": metrics,
"extra": artifact.extra,
}, indent=2, sort_keys=True) + "\n",
json.dumps(
{
"suite": self.suite,
"benchmark": self.name,
"raw_path": "raw.jsonl",
"metrics": metrics,
"extra": artifact.extra,
},
indent=2,
sort_keys=True,
)
+ "\n",
encoding="utf-8",
)
return artifact
@ -536,8 +574,12 @@ def _compute_metrics(
cost_pct = _safe_pct(surf_cost_agg.mean, native_cost_agg.mean)
lat_pct = _safe_pct(surf_lat_agg.median, native_lat_agg.median)
per_task = _per_field(questions, native_correct, surf_correct, key=lambda q: q.medical_task or "unknown")
per_body = _per_field(questions, native_correct, surf_correct, key=lambda q: q.body_system or "unknown")
per_task = _per_field(
questions, native_correct, surf_correct, key=lambda q: q.medical_task or "unknown"
)
per_body = _per_field(
questions, native_correct, surf_correct, key=lambda q: q.body_system or "unknown"
)
return {
"native": {
@ -593,8 +635,7 @@ def _per_field(
"native_accuracy": (sum(n_correct) / len(pairs)) if pairs else 0.0,
"surfsense_accuracy": (sum(s_correct) / len(pairs)) if pairs else 0.0,
"delta_accuracy_pp": (
100.0 * (sum(s_correct) - sum(n_correct)) / len(pairs)
if pairs else 0.0
100.0 * (sum(s_correct) - sum(n_correct)) / len(pairs) if pairs else 0.0
),
}
return out

View file

@ -48,9 +48,7 @@ from ....core.registry import RunContext
logger = logging.getLogger(__name__)
MIRAGE_BENCHMARK_URL = (
"https://raw.githubusercontent.com/Teddy-XiongGZ/MIRAGE/main/benchmark.json"
)
MIRAGE_BENCHMARK_URL = "https://raw.githubusercontent.com/Teddy-XiongGZ/MIRAGE/main/benchmark.json"
# Upstream only ships ONE zip — top-10k retrievals across 5 retrievers,
# ~16 GB. We default to skipping it (see `--skip-snippet-filter`) and
# ingesting the chosen corpus in full; this URL is only fetched when
@ -100,8 +98,7 @@ def _reuse_cached_dest(dest: Path, *, expect_zip: bool, label: str) -> Path | No
return None
if expect_zip and not _is_valid_zip(dest):
logger.warning(
"Cached %s at %s failed ZIP validation (size=%d B); deleting "
"and re-downloading.",
"Cached %s at %s failed ZIP validation (size=%d B); deleting and re-downloading.",
label,
dest,
dest.stat().st_size,
@ -176,10 +173,13 @@ async def _fetch_to_path(
)
try:
async with httpx.AsyncClient(
timeout=httpx.Timeout(timeout_s, connect=20.0),
follow_redirects=True,
) as client, client.stream("GET", url, headers=headers) as response:
async with (
httpx.AsyncClient(
timeout=httpx.Timeout(timeout_s, connect=20.0),
follow_redirects=True,
) as client,
client.stream("GET", url, headers=headers) as response,
):
if existing_bytes and response.status_code == 200:
logger.warning(
"Server ignored Range header for %s; restarting from 0.",
@ -223,7 +223,7 @@ async def _fetch_to_path(
raise
except _RETRYABLE_NET_EXC as exc:
last_exc = exc
wait = min(60.0, 2.0 ** attempt)
wait = min(60.0, 2.0**attempt)
logger.warning(
"Network error fetching %s (%s: %s); retrying in %.0fs.",
label,
@ -236,7 +236,7 @@ async def _fetch_to_path(
last_exc = exc
# Truncated body — drop the partial and retry from scratch.
partial.unlink(missing_ok=True)
wait = min(60.0, 2.0 ** attempt)
wait = min(60.0, 2.0**attempt)
logger.warning(
"Truncated ZIP for %s; restarting from byte 0 in %.0fs.",
label,
@ -278,9 +278,9 @@ class _LargeDownloadAbort(RuntimeError):
"""Raised when a download exceeds the safety threshold without opt-in."""
def __init__(self, label: str, size_bytes: int) -> None:
gb = size_bytes / (1024 ** 3)
gb = size_bytes / (1024**3)
super().__init__(
f"{label} would download ~{gb:.1f} GB, above the {_LARGE_DOWNLOAD_BYTES / (1024 ** 3):.0f} GB safety cap. "
f"{label} would download ~{gb:.1f} GB, above the {_LARGE_DOWNLOAD_BYTES / (1024**3):.0f} GB safety cap. "
"Re-run with `--allow-large-download` to acknowledge, or use "
"`--skip-snippet-filter` to bypass this download entirely and "
"ingest the full corpus instead."
@ -320,9 +320,7 @@ def _read_snippet_ids(zip_path: Path, *, tasks: list[str]) -> dict[str, set[str]
return out
def _load_corpus(
corpus_name: str, snippet_ids: set[str] | None
) -> Iterable[SnippetRow]:
def _load_corpus(corpus_name: str, snippet_ids: set[str] | None) -> Iterable[SnippetRow]:
"""Stream rows from a MedRAG HF corpus.
* ``snippet_ids=None`` yield every row (full-corpus ingestion path).
@ -541,10 +539,7 @@ async def run_ingest(
logger.warning("Failed to list chunks for doc_id=%s: %s", doc_id, exc)
continue
for chunk in chunks:
fh.write(
json.dumps({"chunk_id": chunk.id, "document_id": doc_id})
+ "\n"
)
fh.write(json.dumps({"chunk_id": chunk.id, "document_id": doc_id}) + "\n")
new_state = ctx.suite_state
new_state.ingestion_maps["mirage"] = str(snippet_map_path)

View file

@ -134,15 +134,23 @@ class MirageBenchmark:
choices=("all", *_TASKS),
help="Run a single task or all (default: all).",
)
parser.add_argument("--n", dest="sample_n", type=int, default=None,
help="Stratified sample size across tasks.")
parser.add_argument(
"--n",
dest="sample_n",
type=int,
default=None,
help="Stratified sample size across tasks.",
)
parser.add_argument("--concurrency", type=int, default=4)
parser.add_argument(
"--corpus", default="MedRAG/textbooks",
"--corpus",
default="MedRAG/textbooks",
help="HF MedRAG corpus to ingest from (default: MedRAG/textbooks).",
)
parser.add_argument(
"--max-snippets-per-task", type=int, default=None,
"--max-snippets-per-task",
type=int,
default=None,
help="Cap the per-task ingestion to N snippets (smoke).",
)
# Mutually exclusive: by default we skip the upstream 16 GB
@ -152,18 +160,24 @@ class MirageBenchmark:
# --allow-large-download).
snippet_group = parser.add_mutually_exclusive_group()
snippet_group.add_argument(
"--use-snippet-filter", dest="use_snippet_filter", action="store_true",
"--use-snippet-filter",
dest="use_snippet_filter",
action="store_true",
default=False,
help="Download retrieved_snippets_10k.zip (~16 GB) and "
"filter the corpus to those ids before ingest. "
"Default: skip and ingest entire corpus.",
"filter the corpus to those ids before ingest. "
"Default: skip and ingest entire corpus.",
)
snippet_group.add_argument(
"--skip-snippet-filter", dest="use_snippet_filter", action="store_false",
"--skip-snippet-filter",
dest="use_snippet_filter",
action="store_false",
help="(Default) Skip the 16 GB upstream zip; ingest entire corpus.",
)
parser.add_argument(
"--allow-large-download", action="store_true", default=False,
"--allow-large-download",
action="store_true",
default=False,
help="Permit downloads larger than 2 GB (e.g. retrieved_snippets_10k.zip).",
)
# Per-upload knobs; ignored at run-time (runner reads the
@ -196,16 +210,13 @@ class MirageBenchmark:
"`python -m surfsense_evals ingest medical mirage` first."
)
benchmark = json.loads(bench_path.read_text(encoding="utf-8"))
ingest_settings = read_settings_header(
ctx.maps_dir() / "mirage_snippet_map.jsonl"
)
ingest_settings = read_settings_header(ctx.maps_dir() / "mirage_snippet_map.jsonl")
questions = _load_questions(benchmark, tasks=tasks, sample_n=sample_n)
if not questions:
raise RuntimeError(
f"No MIRAGE questions matched task={task_filter!r} sample_n={sample_n!r}."
)
logger.info("MIRAGE: scheduled %d questions across tasks %s",
len(questions), tasks)
logger.info("MIRAGE: scheduled %d questions across tasks %s", len(questions), tasks)
arm = SurfSenseArm(
client=ctx.new_chat_client(),
@ -255,7 +266,10 @@ class MirageBenchmark:
per_task_acc[task] = acc.to_dict()
macro = macro_accuracy(
{t: accuracy_with_wilson_ci(d["n_correct"], d["n_total"]) for t, d in per_task_acc.items()}
{
t: accuracy_with_wilson_ci(d["n_correct"], d["n_total"])
for t, d in per_task_acc.items()
}
)
metrics = {"per_task": per_task_acc, "macro_accuracy": macro}

View file

@ -112,8 +112,9 @@ def _grade_int(pred: str, gold: str) -> GradeResult:
if p_match is None:
return GradeResult(False, 0.0, "int_eq", str(p_match), str(g_val))
p_val = int(p_match.group(0).replace(",", ""))
return GradeResult(p_val == g_val, 1.0 if p_val == g_val else 0.0,
"int_eq", str(p_val), str(g_val))
return GradeResult(
p_val == g_val, 1.0 if p_val == g_val else 0.0, "int_eq", str(p_val), str(g_val)
)
_FLOAT_RE = re.compile(r"-?\d+(?:[.,]\d+)?")
@ -145,15 +146,15 @@ def _grade_list(pred: str, gold: str) -> GradeResult:
return _grade_str(pred, gold)
inter = g_items & p_items
if not inter:
return GradeResult(False, 0.0, "list_set",
", ".join(sorted(p_items)),
", ".join(sorted(g_items)))
return GradeResult(
False, 0.0, "list_set", ", ".join(sorted(p_items)), ", ".join(sorted(g_items))
)
precision = len(inter) / len(p_items) if p_items else 0.0
recall = len(inter) / len(g_items)
f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0
return GradeResult(f1 >= 0.999, f1, "list_set",
", ".join(sorted(p_items)),
", ".join(sorted(g_items)))
return GradeResult(
f1 >= 0.999, f1, "list_set", ", ".join(sorted(p_items)), ", ".join(sorted(g_items))
)
def _grade_none(pred: str, gold: str) -> GradeResult:
@ -188,8 +189,11 @@ def _grade_none(pred: str, gold: str) -> GradeResult:
expressed_unknown = True
break
return GradeResult(
expressed_unknown, 1.0 if expressed_unknown else 0.0,
"none_match", p, _normalise_text(gold),
expressed_unknown,
1.0 if expressed_unknown else 0.0,
"none_match",
p,
_normalise_text(gold),
)

View file

@ -41,6 +41,7 @@ logger = logging.getLogger(__name__)
HF_REPO_ID = "yubo2333/MMLongBench-Doc"
HF_REPO_TYPE = "dataset"
# Lazy import: huggingface_hub + pyarrow are heavyweight; keep the
# benchmark module importable on machines that have only the core
# install (e.g. CI lint jobs).
@ -63,11 +64,11 @@ def _list_repo_files() -> list[str]:
@dataclass
class MMLongBenchQuestion:
doc_id: str # filename inside the documents/ folder
doc_id: str # filename inside the documents/ folder
doc_type: str
question: str
answer: str
answer_format: str # Str / Int / Float / List / None
answer_format: str # Str / Int / Float / List / None
evidence_pages: list[int]
evidence_sources: list[str]
@ -161,7 +162,9 @@ def _download_questions_parquet(cache_dir: Path) -> Path:
)
parquet_paths.append(Path(local))
logger.info("Cached MMLongBench parquet shard %s -> %s", rel, local)
return parquet_paths[0] if len(parquet_paths) == 1 else _merge_parquets(parquet_paths, cache_dir)
return (
parquet_paths[0] if len(parquet_paths) == 1 else _merge_parquets(parquet_paths, cache_dir)
)
def _merge_parquets(paths: list[Path], cache_dir: Path) -> Path:
@ -221,7 +224,7 @@ async def _upload_pdfs(
name_to_id: dict[str, int] = {}
pdf_list = list(pdf_paths)
for batch_start in range(0, len(pdf_list), batch_size):
batch = pdf_list[batch_start:batch_start + batch_size]
batch = pdf_list[batch_start : batch_start + batch_size]
result = await docs_client.upload(
files=batch,
search_space_id=ctx.search_space_id,
@ -243,8 +246,10 @@ async def _upload_pdfs(
name_to_id[s.title] = s.document_id
logger.info(
"Uploaded MMLongBench batch %d-%d: %d new, %d duplicate",
batch_start, batch_start + len(batch),
len(result.document_ids), len(result.duplicate_document_ids),
batch_start,
batch_start + len(batch),
len(result.document_ids),
len(result.duplicate_document_ids),
)
return name_to_id
@ -299,15 +304,20 @@ async def run_ingest(
questions_jsonl = bench_dir / "questions.jsonl"
with questions_jsonl.open("w", encoding="utf-8") as fh:
for q in questions:
fh.write(json.dumps({
"doc_id": q.doc_id,
"doc_type": q.doc_type,
"question": q.question,
"answer": q.answer,
"answer_format": q.answer_format,
"evidence_pages": q.evidence_pages,
"evidence_sources": q.evidence_sources,
}) + "\n")
fh.write(
json.dumps(
{
"doc_id": q.doc_id,
"doc_type": q.doc_type,
"question": q.question,
"answer": q.answer,
"answer_format": q.answer_format,
"evidence_pages": q.evidence_pages,
"evidence_sources": q.evidence_sources,
}
)
+ "\n"
)
logger.info("Wrote %d MMLongBench questions to %s", len(questions), questions_jsonl)
# Step 2: download unique PDFs
@ -348,12 +358,17 @@ async def run_ingest(
local = pdf_paths.get(doc_id)
if local is None:
continue
fh.write(json.dumps({
"doc_id": doc_id,
"document_id": name_to_id.get(local.name),
"pdf_path": str(local),
"n_questions": sum(1 for q in questions if q.doc_id == doc_id),
}) + "\n")
fh.write(
json.dumps(
{
"doc_id": doc_id,
"document_id": name_to_id.get(local.name),
"pdf_path": str(local),
"n_questions": sum(1 for q in questions if q.doc_id == doc_id),
}
)
+ "\n"
)
logger.info("Wrote MMLongBench doc map to %s", map_path)
new_state = ctx.suite_state

View file

@ -18,10 +18,7 @@ _FORMAT_HINTS: dict[str, str] = {
"Respond with the answer as a short phrase, no full sentence. "
"Format your final line as `Answer: <text>`."
),
"int": (
"Respond with a single integer only. "
"Format your final line as `Answer: <integer>`."
),
"int": ("Respond with a single integer only. Format your final line as `Answer: <integer>`."),
"float": (
"Respond with a single decimal number only (no units). "
"Format your final line as `Answer: <number>`."

View file

@ -58,8 +58,8 @@ logger = logging.getLogger(__name__)
@dataclass
class MMLBQuestion:
qid: str # synthesised from doc_id + index
doc_id: str # filename inside the documents/ folder
qid: str # synthesised from doc_id + index
doc_id: str # filename inside the documents/ folder
doc_type: str
question: str
gold_answer: str
@ -126,18 +126,20 @@ def _load_questions(
continue
idx = per_doc_counter.get(doc_id, 0)
per_doc_counter[doc_id] = idx + 1
out.append(MMLBQuestion(
qid=f"{doc_id}::Q{idx:03d}",
doc_id=doc_id,
doc_type=str(row.get("doc_type") or "").strip(),
question=str(row.get("question") or "").strip(),
gold_answer=gold,
answer_format=answer_format,
evidence_pages=list(row.get("evidence_pages") or []),
evidence_sources=list(row.get("evidence_sources") or []),
pdf_path=Path(map_row["pdf_path"]),
document_id=map_row.get("document_id"),
))
out.append(
MMLBQuestion(
qid=f"{doc_id}::Q{idx:03d}",
doc_id=doc_id,
doc_type=str(row.get("doc_type") or "").strip(),
question=str(row.get("question") or "").strip(),
gold_answer=gold,
answer_format=answer_format,
evidence_pages=list(row.get("evidence_pages") or []),
evidence_sources=list(row.get("evidence_sources") or []),
pdf_path=Path(map_row["pdf_path"]),
document_id=map_row.get("document_id"),
)
)
out.sort(key=lambda q: (q.doc_id, q.qid))
if sample_n is not None and sample_n > 0:
out = out[:sample_n]
@ -202,41 +204,61 @@ class MMLongBenchDocBenchmark:
help="Filter to one answer format. 'none' = unanswerable probes only.",
)
parser.add_argument(
"--n", dest="sample_n", type=int, default=None,
"--n",
dest="sample_n",
type=int,
default=None,
help="Run only the first N questions after filters apply.",
)
parser.add_argument(
"--skip-unanswerable", dest="skip_unanswerable", action="store_true",
"--skip-unanswerable",
dest="skip_unanswerable",
action="store_true",
help="Drop ~22%% unanswerable questions (use to compare against baselines that don't include them).",
)
parser.add_argument(
"--concurrency", type=int, default=4,
"--concurrency",
type=int,
default=4,
help="Parallel question workers per arm.",
)
parser.add_argument(
"--no-mentions", dest="no_mentions", action="store_true",
"--no-mentions",
dest="no_mentions",
action="store_true",
help="SurfSense arm: skip mentioned_document_ids (unscoped retrieval).",
)
parser.add_argument(
"--pdf-engine", default="native",
"--pdf-engine",
default="native",
choices=[e.value for e in PdfEngine],
help="OpenRouter file-parser engine for the native arm.",
)
parser.add_argument(
"--max-output-tokens", type=int, default=512,
"--max-output-tokens",
type=int,
default=512,
help="Cap on completion length for both arms.",
)
# Ingest-only knobs (forwarded by the CLI to ingest.run_ingest).
parser.add_argument(
"--max-docs", dest="max_docs", type=int, default=None,
"--max-docs",
dest="max_docs",
type=int,
default=None,
help="(ingest only) cap on number of unique PDFs to download + upload.",
)
parser.add_argument(
"--upload-batch-size", dest="upload_batch_size", type=int, default=8,
"--upload-batch-size",
dest="upload_batch_size",
type=int,
default=8,
help="(ingest only) PDFs per fileupload call.",
)
parser.add_argument(
"--skip-upload", dest="skip_upload", action="store_true",
"--skip-upload",
dest="skip_upload",
action="store_true",
help="(ingest only) cache PDFs locally but don't push to SurfSense.",
)
# Per-upload knobs forwarded to /documents/fileupload at ingest;
@ -278,7 +300,8 @@ class MMLongBenchDocBenchmark:
doc_map, ingest_settings = _load_doc_map(map_path)
questions = _load_questions(
questions_jsonl, doc_map,
questions_jsonl,
doc_map,
doc_filter=doc_filter,
format_filter=None if format_filter == "all" else format_filter,
sample_n=sample_n,
@ -292,9 +315,7 @@ class MMLongBenchDocBenchmark:
api_key = os.environ.get("OPENROUTER_API_KEY")
if not api_key:
raise RuntimeError(
"OPENROUTER_API_KEY env var is required for the native arm."
)
raise RuntimeError("OPENROUTER_API_KEY env var is required for the native arm.")
# Native arm slug differs from SurfSense slug only in cost-arbitrage
# scenario; otherwise both arms answer with provider_model.
@ -362,18 +383,30 @@ class MMLongBenchDocBenchmark:
"evidence_sources": q.evidence_sources,
"document_id": q.document_id,
}
fh.write(json.dumps({
**meta,
**n_res.to_jsonl(),
"graded": _grade_to_jsonl(n_g),
}) + "\n")
fh.write(json.dumps({
**meta,
**s_res.to_jsonl(),
"graded": _grade_to_jsonl(s_g),
}) + "\n")
fh.write(
json.dumps(
{
**meta,
**n_res.to_jsonl(),
"graded": _grade_to_jsonl(n_g),
}
)
+ "\n"
)
fh.write(
json.dumps(
{
**meta,
**s_res.to_jsonl(),
"graded": _grade_to_jsonl(s_g),
}
)
+ "\n"
)
metrics = _compute_metrics(questions, native_results, surf_results, native_grades, surf_grades)
metrics = _compute_metrics(
questions, native_results, surf_results, native_grades, surf_grades
)
artifact = RunArtifact(
suite=self.suite,
benchmark=self.name,
@ -398,13 +431,18 @@ class MMLongBenchDocBenchmark:
manifest_path = run_dir / "run_artifact.json"
manifest_path.write_text(
json.dumps({
"suite": self.suite,
"benchmark": self.name,
"raw_path": "raw.jsonl",
"metrics": metrics,
"extra": artifact.extra,
}, indent=2, sort_keys=True) + "\n",
json.dumps(
{
"suite": self.suite,
"benchmark": self.name,
"raw_path": "raw.jsonl",
"metrics": metrics,
"extra": artifact.extra,
},
indent=2,
sort_keys=True,
)
+ "\n",
encoding="utf-8",
)
return artifact
@ -450,9 +488,7 @@ class MMLongBenchDocBenchmark:
f"(McNemar p={_fmt(delta.get('mcnemar_p_value'), 4)}, "
f"method={delta.get('mcnemar_method')})"
)
body_lines.append(
f" - F1 (mean): SurfSense {_pp(delta.get('f1_pp'))} pp"
)
body_lines.append(f" - F1 (mean): SurfSense {_pp(delta.get('f1_pp'))} pp")
body_lines.append(
f" - Bootstrap 95% CI on accuracy delta: "
f"[{_pp(delta.get('bootstrap_ci_low'))}pp, {_pp(delta.get('bootstrap_ci_high'))}pp]"
@ -472,8 +508,8 @@ class MMLongBenchDocBenchmark:
for fmt, vals in sorted(per_format.items()):
body_lines.append(
f" - {fmt}: SurfSense {_pp(vals.get('delta_accuracy_pp'))} pp "
f"(n={vals.get('n')}, native acc={vals.get('native_accuracy', 0)*100:.1f}%, "
f"surf acc={vals.get('surfsense_accuracy', 0)*100:.1f}%)"
f"(n={vals.get('n')}, native acc={vals.get('native_accuracy', 0) * 100:.1f}%, "
f"surf acc={vals.get('surfsense_accuracy', 0) * 100:.1f}%)"
)
return ReportSection(
@ -576,8 +612,7 @@ def _compute_metrics(
"native_accuracy": (sum(n_correct) / len(pairs)) if pairs else 0.0,
"surfsense_accuracy": (sum(s_correct) / len(pairs)) if pairs else 0.0,
"delta_accuracy_pp": (
100.0 * (sum(s_correct) - sum(n_correct)) / len(pairs)
if pairs else 0.0
100.0 * (sum(s_correct) - sum(n_correct)) / len(pairs) if pairs else 0.0
),
}
@ -593,8 +628,12 @@ def _compute_metrics(
"latency_ms_mean": native_latency_agg.mean,
"latency_ms_median": native_latency_agg.median,
"latency_ms_p95": native_latency_agg.p95,
"input_tokens_mean": (sum(native_in_tokens) / len(native_in_tokens)) if native_in_tokens else 0.0,
"output_tokens_mean": (sum(native_out_tokens) / len(native_out_tokens)) if native_out_tokens else 0.0,
"input_tokens_mean": (sum(native_in_tokens) / len(native_in_tokens))
if native_in_tokens
else 0.0,
"output_tokens_mean": (sum(native_out_tokens) / len(native_out_tokens))
if native_out_tokens
else 0.0,
},
"surfsense": {
**surf_acc.to_dict(),

View file

@ -53,9 +53,9 @@ logger = logging.getLogger(__name__)
# Order matters for the manifest only (deterministic JSONL diffs);
# the runner doesn't rely on it.
PARSER_ARMS: tuple[tuple[str, str, str], ...] = (
("azure_basic_lc", "azure", "basic"),
("azure_premium_lc", "azure", "premium"),
("llamacloud_basic_lc", "llamacloud", "basic"),
("azure_basic_lc", "azure", "basic"),
("azure_premium_lc", "azure", "premium"),
("llamacloud_basic_lc", "llamacloud", "basic"),
("llamacloud_premium_lc", "llamacloud", "premium"),
)
@ -98,9 +98,7 @@ class PdfManifestRow:
"pdf_path": str(self.pdf_path),
"document_id": self.document_id,
"pages": self.pages,
"extractions": {
arm: ext.to_jsonl() for arm, ext in self.extractions.items()
},
"extractions": {arm: ext.to_jsonl() for arm, ext in self.extractions.items()},
}
@ -124,7 +122,9 @@ async def _run_one_extraction(
markdown = await parse_with_azure_di(pdf_path, processing_mode=mode)
elif parser == "llamacloud":
markdown = await parse_with_llamacloud(
pdf_path, processing_mode=mode, estimated_pages=estimated_pages,
pdf_path,
processing_mode=mode,
estimated_pages=estimated_pages,
)
else:
raise ValueError(f"Unknown parser {parser!r}")
@ -168,14 +168,17 @@ async def _extract_one_pdf(
error="(cached)",
)
logger.info(
"Cached extraction reused: %s (%d chars)", out_path.name, len(cached),
"Cached extraction reused: %s (%d chars)",
out_path.name,
len(cached),
)
coros.append(_noop())
else:
coros.append(
_run_one_extraction(
pdf_path,
parser=parser, mode=mode,
parser=parser,
mode=mode,
out_path=out_path,
estimated_pages=estimated_pages,
)
@ -190,16 +193,24 @@ async def _extract_one_pdf(
err_msg = f"{type(err).__name__}: {err}"
logger.warning(
"Extraction FAILED for %s [%s/%s]: %s",
pdf_path.name, parser, mode, err_msg,
pdf_path.name,
parser,
mode,
err_msg,
)
out[arm_name] = ExtractionResult(
arm=arm_name, parser=parser, mode=mode,
status="failed", error=err_msg,
arm=arm_name,
parser=parser,
mode=mode,
status="failed",
error=err_msg,
)
else:
markdown, elapsed = result
out[arm_name] = ExtractionResult(
arm=arm_name, parser=parser, mode=mode,
arm=arm_name,
parser=parser,
mode=mode,
markdown_path=out_path,
chars=len(markdown),
elapsed_s=elapsed,
@ -288,9 +299,7 @@ async def run_ingest(
rows_in_scope = rows_in_scope[:max_docs]
if not rows_in_scope:
raise RuntimeError(
"No PDFs in scope for parser_compare. Check --docs / --max-docs."
)
raise RuntimeError("No PDFs in scope for parser_compare. Check --docs / --max-docs.")
bench_dir = ctx.benchmark_data_dir()
extractions_dir = bench_dir / "extractions"
@ -317,7 +326,8 @@ async def run_ingest(
logger.info(
"parser_compare: extracting %d PDFs x 4 parsers (concurrency=%d)",
len(rows_in_scope), pdf_concurrency,
len(rows_in_scope),
pdf_concurrency,
)
manifest_rows = await asyncio.gather(*(_process(r) for r in rows_in_scope))
@ -337,12 +347,13 @@ async def run_ingest(
# Quick summary log
total_extractions = sum(len(mr.extractions) for mr in manifest_rows)
failures = sum(
1 for mr in manifest_rows for ext in mr.extractions.values()
if ext.status != "ok"
1 for mr in manifest_rows for ext in mr.extractions.values() if ext.status != "ok"
)
logger.info(
"parser_compare ingest done: %d PDFs, %d extractions, %d failures",
len(manifest_rows), total_extractions, failures,
len(manifest_rows),
total_extractions,
failures,
)

View file

@ -34,10 +34,7 @@ _FORMAT_HINTS: dict[str, str] = {
"Respond with the answer as a short phrase, no full sentence. "
"Format your final line as `Answer: <text>`."
),
"int": (
"Respond with a single integer only. "
"Format your final line as `Answer: <integer>`."
),
"int": ("Respond with a single integer only. Format your final line as `Answer: <integer>`."),
"float": (
"Respond with a single decimal number only (no units). "
"Format your final line as `Answer: <number>`."
@ -69,11 +66,7 @@ _BASE_INSTRUCTION = (
def build_native_pdf_prompt(question: str, *, answer_format: str) -> str:
"""Prompt for ``NativePdfArm`` — PDF attached separately as a file part."""
return (
f"{_BASE_INSTRUCTION}\n\n"
f"Question: {question.strip()}\n\n"
f"{_format_hint(answer_format)}\n"
)
return f"{_BASE_INSTRUCTION}\n\nQuestion: {question.strip()}\n\n{_format_hint(answer_format)}\n"
def build_surfsense_prompt(question: str, *, answer_format: str) -> str:
@ -82,11 +75,7 @@ def build_surfsense_prompt(question: str, *, answer_format: str) -> str:
# SurfSense's agent already injects retrieved chunks via its tool
# loop; the prompt only carries the user-visible question + format
# hint, mirroring how a human asks the SurfSense UI.
return (
f"{_BASE_INSTRUCTION}\n\n"
f"Question: {question.strip()}\n\n"
f"{_format_hint(answer_format)}\n"
)
return f"{_BASE_INSTRUCTION}\n\nQuestion: {question.strip()}\n\n{_format_hint(answer_format)}\n"
def build_long_context_prompt(
@ -105,7 +94,7 @@ def build_long_context_prompt(
return (
f"{_BASE_INSTRUCTION}\n\n"
f"<document name=\"{document_label}\">\n"
f'<document name="{document_label}">\n'
f"{document_markdown.strip()}\n"
f"</document>\n\n"
f"Question: {question.strip()}\n\n"

View file

@ -72,7 +72,7 @@ logger = logging.getLogger(__name__)
# Cost tariff (per the user's spec: $1 / 1k pages basic, $10 / 1k pages premium).
# Held as dollars-per-page so per-PDF math is a pure multiply.
PREPROCESS_USD_PER_PAGE = {
"basic": 1.0 / 1000.0,
"basic": 1.0 / 1000.0,
"premium": 10.0 / 1000.0,
}
@ -183,17 +183,19 @@ def _select_questions(
if ext_blob.get("status") == "ok" and ext_blob.get("markdown_path"):
extractions[arm_name] = Path(ext_blob["markdown_path"])
out.append(PCQuestion(
qid=f"{doc_id}::Q{idx:03d}",
doc_id=doc_id,
question=str(row.get("question") or "").strip(),
gold_answer=str(row.get("answer") or "").strip(),
answer_format=answer_format,
pdf_path=Path(map_row["pdf_path"]),
document_id=map_row.get("document_id"),
pages=int(map_row.get("pages", 0)),
extractions=extractions,
))
out.append(
PCQuestion(
qid=f"{doc_id}::Q{idx:03d}",
doc_id=doc_id,
question=str(row.get("question") or "").strip(),
gold_answer=str(row.get("answer") or "").strip(),
answer_format=answer_format,
pdf_path=Path(map_row["pdf_path"]),
document_id=map_row.get("document_id"),
pages=int(map_row.get("pages", 0)),
extractions=extractions,
)
)
per_doc_taken[doc_id] = per_doc_taken.get(doc_id, 0) + 1
out.sort(key=lambda q: (q.doc_id, q.qid))
@ -242,65 +244,86 @@ class ParserCompareBenchmark:
def add_run_args(self, parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--docs", default=None,
"--docs",
default=None,
help="Comma-separated doc_ids to include (default: all in manifest).",
)
parser.add_argument(
"--sample-per-doc", type=int, default=1,
"--sample-per-doc",
type=int,
default=1,
help="Take the first N answerable questions per PDF (default 1).",
)
parser.add_argument(
"--skip-unanswerable", dest="skip_unanswerable",
action="store_true", default=True,
"--skip-unanswerable",
dest="skip_unanswerable",
action="store_true",
default=True,
help="Drop 'None' format probes (default true; we want signal not "
"hallucination probes for n=5).",
"hallucination probes for n=5).",
)
parser.add_argument(
"--include-unanswerable", dest="skip_unanswerable",
"--include-unanswerable",
dest="skip_unanswerable",
action="store_false",
help="Override --skip-unanswerable; include unanswerable probes too.",
)
parser.add_argument(
"--skip-format", default=None,
"--skip-format",
default=None,
help="Comma-separated answer_format values to skip (e.g. 'none,float').",
)
parser.add_argument(
"--concurrency", type=int, default=2,
"--concurrency",
type=int,
default=2,
help="Parallel question workers per arm (default 2).",
)
parser.add_argument(
"--no-mentions", dest="no_mentions", action="store_true",
"--no-mentions",
dest="no_mentions",
action="store_true",
help="SurfSense arm: skip mentioned_document_ids (full-corpus retrieval).",
)
parser.add_argument(
"--pdf-engine", default="native",
"--pdf-engine",
default="native",
choices=[e.value for e in PdfEngine],
help="OpenRouter file-parser engine for native_pdf arm.",
)
parser.add_argument(
"--max-output-tokens", type=int, default=512,
"--max-output-tokens",
type=int,
default=512,
help="Cap on completion length for every arm.",
)
parser.add_argument(
"--llm-model", default="anthropic/claude-sonnet-4.5",
"--llm-model",
default="anthropic/claude-sonnet-4.5",
help="OpenRouter slug used by the 5 OpenRouter-driven arms. "
"SurfSense arm uses whatever provider_model is pinned on the suite.",
"SurfSense arm uses whatever provider_model is pinned on the suite.",
)
parser.add_argument(
"--skip-arms", default=None,
"--skip-arms",
default=None,
help="Comma-separated arm names to skip (e.g. 'llamacloud_premium_lc').",
)
# Ingest-only flags (forwarded by the CLI to ingest.run_ingest).
parser.add_argument(
"--max-docs", type=int, default=None,
"--max-docs",
type=int,
default=None,
help="(ingest only) cap number of unique PDFs to process.",
)
parser.add_argument(
"--force-reextract", action="store_true",
"--force-reextract",
action="store_true",
help="(ingest only) re-call parsers even if cached .md exists.",
)
parser.add_argument(
"--pdf-concurrency", type=int, default=2,
"--pdf-concurrency",
type=int,
default=2,
help="(ingest only) parallel PDFs (each fans out to 4 parsers).",
)
@ -312,9 +335,7 @@ class ParserCompareBenchmark:
from .ingest import run_ingest
docs_raw: str | None = opts.get("docs")
docs_filter = (
[d.strip() for d in docs_raw.split(",") if d.strip()] if docs_raw else None
)
docs_filter = [d.strip() for d in docs_raw.split(",") if d.strip()] if docs_raw else None
await run_ingest(
ctx,
docs_filter=docs_filter,
@ -329,15 +350,14 @@ class ParserCompareBenchmark:
async def run(self, ctx: RunContext, **opts: Any) -> RunArtifact:
docs_raw: str | None = opts.get("docs")
docs_filter = (
[d.strip() for d in docs_raw.split(",") if d.strip()] if docs_raw else None
)
docs_filter = [d.strip() for d in docs_raw.split(",") if d.strip()] if docs_raw else None
sample_per_doc = int(opts.get("sample_per_doc") or 1)
skip_unanswerable = bool(opts.get("skip_unanswerable", True))
skip_format_raw: str | None = opts.get("skip_format")
skip_format = (
[f.strip() for f in skip_format_raw.split(",") if f.strip()]
if skip_format_raw else None
if skip_format_raw
else None
)
concurrency = int(opts.get("concurrency") or 2)
no_mentions = bool(opts.get("no_mentions"))
@ -346,8 +366,7 @@ class ParserCompareBenchmark:
llm_model = str(opts.get("llm_model") or "anthropic/claude-sonnet-4.5")
skip_arms_raw: str | None = opts.get("skip_arms")
skip_arms = (
{a.strip() for a in skip_arms_raw.split(",") if a.strip()}
if skip_arms_raw else set()
{a.strip() for a in skip_arms_raw.split(",") if a.strip()} if skip_arms_raw else set()
)
active_arms = [a for a in ARM_NAMES if a not in skip_arms]
@ -373,19 +392,20 @@ class ParserCompareBenchmark:
doc_map = _read_doc_map(map_path)
questions = _select_questions(
questions_jsonl, doc_map,
questions_jsonl,
doc_map,
docs_filter=docs_filter,
sample_per_doc=sample_per_doc,
skip_unanswerable=skip_unanswerable,
skip_format=skip_format,
)
if not questions:
raise RuntimeError(
"No questions matched filters; broaden --docs / --skip-format."
)
raise RuntimeError("No questions matched filters; broaden --docs / --skip-format.")
logger.info(
"parser_compare: scheduled %d questions across %d arms (%s)",
len(questions), len(active_arms), ",".join(active_arms),
len(questions),
len(active_arms),
",".join(active_arms),
)
api_key = os.environ.get("OPENROUTER_API_KEY")
@ -396,16 +416,20 @@ class ParserCompareBenchmark:
arms: dict[str, Any] = {}
if "native_pdf" in active_arms:
native_provider = OpenRouterPdfProvider(
api_key=api_key, base_url=ctx.config.openrouter_base_url,
model=llm_model, engine=PdfEngine(pdf_engine_name),
api_key=api_key,
base_url=ctx.config.openrouter_base_url,
model=llm_model,
engine=PdfEngine(pdf_engine_name),
)
arms["native_pdf"] = NativePdfArm(
provider=native_provider, max_output_tokens=max_output_tokens,
provider=native_provider,
max_output_tokens=max_output_tokens,
)
for arm_name, _, _ in PARSER_ARMS:
if arm_name in active_arms:
lc_provider = OpenRouterChatProvider(
api_key=api_key, base_url=ctx.config.openrouter_base_url,
api_key=api_key,
base_url=ctx.config.openrouter_base_url,
model=llm_model,
)
arms[arm_name] = BareLlmArm(
@ -441,9 +465,7 @@ class ParserCompareBenchmark:
def _lc_req(q: PCQuestion, arm_name: str) -> ArmRequest:
md_path = q.extractions.get(arm_name)
if md_path is None or not md_path.exists():
raise FileNotFoundError(
f"Missing extraction for {arm_name} on {q.doc_id}"
)
raise FileNotFoundError(f"Missing extraction for {arm_name} on {q.doc_id}")
markdown = md_path.read_text(encoding="utf-8")
return ArmRequest(
question_id=q.qid,
@ -483,14 +505,15 @@ class ParserCompareBenchmark:
# Run all arms in parallel (each arm bounded by `concurrency`).
per_arm_tasks: dict[str, list] = {
arm_name: [_answer_one(arm_name, q) for q in questions]
for arm_name in active_arms
arm_name: [_answer_one(arm_name, q) for q in questions] for arm_name in active_arms
}
per_arm_results: dict[str, list[ArmResult]] = {}
gathered = await asyncio.gather(*[
_gather_with_limit(per_arm_tasks[arm_name], concurrency=concurrency)
for arm_name in active_arms
])
gathered = await asyncio.gather(
*[
_gather_with_limit(per_arm_tasks[arm_name], concurrency=concurrency)
for arm_name in active_arms
]
)
for arm_name, results in zip(active_arms, gathered, strict=True):
per_arm_results[arm_name] = results
@ -520,21 +543,29 @@ class ParserCompareBenchmark:
for arm_name in active_arms:
res = per_arm_results[arm_name][i]
g = per_arm_grades[arm_name][i]
fh.write(json.dumps({
**base,
**res.to_jsonl(),
"graded": {
"correct": g.correct,
"f1": g.f1,
"method": g.method,
"normalised_pred": g.normalised_pred,
"normalised_gold": g.normalised_gold,
},
}) + "\n")
fh.write(
json.dumps(
{
**base,
**res.to_jsonl(),
"graded": {
"correct": g.correct,
"f1": g.f1,
"method": g.method,
"normalised_pred": g.normalised_pred,
"normalised_gold": g.normalised_gold,
},
}
)
+ "\n"
)
# Aggregate per-arm metrics + cost
metrics = _compute_metrics(
questions, per_arm_results, per_arm_grades, active_arms,
questions,
per_arm_results,
per_arm_grades,
active_arms,
)
artifact = RunArtifact(
@ -564,13 +595,18 @@ class ParserCompareBenchmark:
manifest_path = run_dir / "run_artifact.json"
manifest_path.write_text(
json.dumps({
"suite": self.suite,
"benchmark": self.name,
"raw_path": "raw.jsonl",
"metrics": metrics,
"extra": artifact.extra,
}, indent=2, sort_keys=True) + "\n",
json.dumps(
{
"suite": self.suite,
"benchmark": self.name,
"raw_path": "raw.jsonl",
"metrics": metrics,
"extra": artifact.extra,
},
indent=2,
sort_keys=True,
)
+ "\n",
encoding="utf-8",
)
return artifact
@ -602,10 +638,7 @@ class ParserCompareBenchmark:
f"(LLM: `{extra.get('llm_model', '?')}`, "
f"engine: `{extra.get('pdf_engine', 'native')}`)."
)
body.append(
"- Preprocess tariff: basic = $1 / 1k pages, "
"premium = $10 / 1k pages."
)
body.append("- Preprocess tariff: basic = $1 / 1k pages, premium = $10 / 1k pages.")
body.append("")
body.append("### Per-arm summary")
body.append("")
@ -620,13 +653,13 @@ class ParserCompareBenchmark:
continue
body.append(
f"| `{arm_name}` "
f"| {row['accuracy']*100:.1f}% "
f"| {row['accuracy'] * 100:.1f}% "
f"({row['n_correct']}/{row['n']}) "
f"| {row['f1_mean']*100:.1f}% "
f"| {row['f1_mean'] * 100:.1f}% "
f"| ${row['llm_cost_per_q']:.4f} "
f"| ${row['preprocess_cost_total']:.4f} "
f"| ${row['total_cost_per_q']:.4f} "
f"| {row['latency_ms_median']/1000:.1f}s |"
f"| {row['latency_ms_median'] / 1000:.1f}s |"
)
body.append("")
@ -679,8 +712,7 @@ class ParserCompareBenchmark:
else:
row_cells.append("" if g.get("correct") else "")
body.append(
f"| `{doc_id}` | {info.get('pages', '?')} | "
+ " | ".join(row_cells) + " |"
f"| `{doc_id}` | {info.get('pages', '?')} | " + " | ".join(row_cells) + " |"
)
return ReportSection(
@ -740,16 +772,16 @@ def _compute_metrics(
preprocess_per_page = 0.0
preprocess_label = "unknown"
preprocess_cost_total = sum(
pages * preprocess_per_page for pages in pdf_pages.values()
)
preprocess_cost_total = sum(pages * preprocess_per_page for pages in pdf_pages.values())
preprocess_cost_per_q = preprocess_cost_total / n if n else 0.0
total_cost_per_q = llm_cost_per_q + preprocess_cost_per_q
latencies = sorted(int(r.latency_ms or 0) for r in results)
latency_median = latencies[len(latencies) // 2] if latencies else 0
latency_p95 = latencies[int(len(latencies) * 0.95)] if len(latencies) >= 20 else (
latencies[-1] if latencies else 0
latency_p95 = (
latencies[int(len(latencies) * 0.95)]
if len(latencies) >= 20
else (latencies[-1] if latencies else 0)
)
in_tokens = [int(r.input_tokens or 0) for r in results]
@ -775,15 +807,21 @@ def _compute_metrics(
# Per-PDF breakdown (correct / not for each arm)
per_pdf: dict[str, dict[str, Any]] = {}
for i, q in enumerate(questions):
slot = per_pdf.setdefault(q.doc_id, {
"pages": q.pages,
"arms": {},
})
slot = per_pdf.setdefault(
q.doc_id,
{
"pages": q.pages,
"arms": {},
},
)
for arm_name in active_arms:
slot["arms"].setdefault(arm_name, {
"correct": per_arm_grades[arm_name][i].correct,
"f1": per_arm_grades[arm_name][i].f1,
})
slot["arms"].setdefault(
arm_name,
{
"correct": per_arm_grades[arm_name][i].correct,
"f1": per_arm_grades[arm_name][i].f1,
},
)
return {
"per_arm": per_arm,

View file

@ -80,7 +80,7 @@ class CragPage:
class CragQuestion:
"""One row of CRAG (Tasks 1 & 2)."""
qid: str # synthesised "C00000".."C02705"
qid: str # synthesised "C00000".."C02705"
interaction_id: str
query_time: str
query: str
@ -89,9 +89,9 @@ class CragQuestion:
domain: str
question_type: str
static_or_dynamic: str
popularity: str # may be "" for web-sourced questions
split: int # 0=validation, 1=public_test
raw_index: int # row index in the source JSONL
popularity: str # may be "" for web-sourced questions
split: int # 0=validation, 1=public_test
raw_index: int # row index in the source JSONL
pages: list[CragPage] = field(default_factory=list)
def to_dict(self) -> dict[str, Any]:
@ -166,16 +166,19 @@ def _parse_pages(raw_search_results: Any) -> list[CragPage]:
if not url or not html.strip():
# No URL or empty HTML => useless for retrieval.
continue
pages.append(CragPage(
page_name=str(entry.get("page_name") or "").strip(),
page_url=url,
page_snippet=str(entry.get("page_snippet") or "").strip(),
page_html=html,
page_last_modified=(
str(entry.get("page_last_modified")).strip()
if entry.get("page_last_modified") else None
),
))
pages.append(
CragPage(
page_name=str(entry.get("page_name") or "").strip(),
page_url=url,
page_snippet=str(entry.get("page_snippet") or "").strip(),
page_html=html,
page_last_modified=(
str(entry.get("page_last_modified")).strip()
if entry.get("page_last_modified")
else None
),
)
)
return pages
@ -217,21 +220,23 @@ def iter_questions(jsonl_bz2_path: Path) -> list[CragQuestion]:
continue
interaction_id = str(row.get("interaction_id") or "").strip()
pages = _parse_pages(row.get("search_results"))
out.append(CragQuestion(
qid=f"C{raw_idx:05d}",
interaction_id=interaction_id,
query_time=str(row.get("query_time") or "").strip(),
query=query,
gold_answer=answer,
alt_answers=_parse_alt_answers(row.get("alt_ans")),
domain=str(row.get("domain") or "").strip().lower(),
question_type=str(row.get("question_type") or "").strip().lower(),
static_or_dynamic=str(row.get("static_or_dynamic") or "").strip().lower(),
popularity=str(row.get("popularity") or "").strip().lower(),
split=int(row.get("split") or 0),
raw_index=raw_idx,
pages=pages,
))
out.append(
CragQuestion(
qid=f"C{raw_idx:05d}",
interaction_id=interaction_id,
query_time=str(row.get("query_time") or "").strip(),
query=query,
gold_answer=answer,
alt_answers=_parse_alt_answers(row.get("alt_ans")),
domain=str(row.get("domain") or "").strip().lower(),
question_type=str(row.get("question_type") or "").strip().lower(),
static_or_dynamic=str(row.get("static_or_dynamic") or "").strip().lower(),
popularity=str(row.get("popularity") or "").strip().lower(),
split=int(row.get("split") or 0),
raw_index=raw_idx,
pages=pages,
)
)
return out

View file

@ -58,10 +58,10 @@ class CragGradeResult:
"""One graded (pred, gold) pair under CRAG's 3-class rubric."""
grade: GradeClass
score: int # +1 / 0 / -1
method: str # exact, numeric, substring, refusal,
# false_premise_correct, false_premise_miss,
# llm_judge, lexical_miss, ...
score: int # +1 / 0 / -1
method: str # exact, numeric, substring, refusal,
# false_premise_correct, false_premise_miss,
# llm_judge, lexical_miss, ...
normalised_pred: str = ""
normalised_gold: str = ""
judge_rationale: str = ""
@ -112,10 +112,27 @@ def _normalise(s: str) -> str:
_WORD_NUMBERS = {
"zero": 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5,
"six": 6, "seven": 7, "eight": 8, "nine": 9, "ten": 10, "eleven": 11,
"twelve": 12, "thirteen": 13, "fourteen": 14, "fifteen": 15, "sixteen": 16,
"seventeen": 17, "eighteen": 18, "nineteen": 19, "twenty": 20,
"zero": 0,
"one": 1,
"two": 2,
"three": 3,
"four": 4,
"five": 5,
"six": 6,
"seven": 7,
"eight": 8,
"nine": 9,
"ten": 10,
"eleven": 11,
"twelve": 12,
"thirteen": 13,
"fourteen": 14,
"fifteen": 15,
"sixteen": 16,
"seventeen": 17,
"eighteen": 18,
"nineteen": 19,
"twenty": 20,
}
_NUMERIC_RE = re.compile(r"-?\d+(?:[.,]\d+)?")
@ -274,8 +291,11 @@ def grade_deterministic(
continue
if n_pred == cand_norm:
return CragGradeResult(
grade="correct", score=1, method="exact",
normalised_pred=n_pred, normalised_gold=cand_norm,
grade="correct",
score=1,
method="exact",
normalised_pred=n_pred,
normalised_gold=cand_norm,
)
p_num = _maybe_number(pred)
c_num = _maybe_number(candidate)
@ -289,21 +309,30 @@ def grade_deterministic(
tol = abs(c_num) * 0.01
if abs(p_num - c_num) <= tol:
return CragGradeResult(
grade="correct", score=1, method="numeric",
normalised_pred=n_pred, normalised_gold=cand_norm,
grade="correct",
score=1,
method="numeric",
normalised_pred=n_pred,
normalised_gold=cand_norm,
)
# Numeric question with different numbers — keep looking
# at other candidates rather than declaring miss now;
# alt answers may include word forms that pass.
if _whole_word_substring(n_pred, cand_norm):
return CragGradeResult(
grade="correct", score=1, method="substring",
normalised_pred=n_pred, normalised_gold=cand_norm,
grade="correct",
score=1,
method="substring",
normalised_pred=n_pred,
normalised_gold=cand_norm,
)
if _whole_word_substring(cand_norm, n_pred) and len(n_pred) >= 3:
return CragGradeResult(
grade="correct", score=1, method="substring_reverse",
normalised_pred=n_pred, normalised_gold=cand_norm,
grade="correct",
score=1,
method="substring_reverse",
normalised_pred=n_pred,
normalised_gold=cand_norm,
)
return CragGradeResult(
@ -326,21 +355,21 @@ _JUDGE_SYSTEM = (
"answer (and any alternative valid answers), and a model's "
"prediction, classify the prediction into exactly one of three "
"categories:\n\n"
"* \"correct\" — the prediction expresses the same factual "
'* "correct" — the prediction expresses the same factual '
"content as the gold answer (paraphrasing OK; numbers as words "
"OK; partial-but-correct names OK; non-contradictory extra "
"detail OK).\n"
"* \"missing\" — the prediction explicitly refuses, says \"I "
'* "missing" — the prediction explicitly refuses, says "I '
"don't know\", says there is insufficient information, or hedges "
"without committing.\n"
"* \"incorrect\" — the prediction commits to a fact that is "
'* "incorrect" — the prediction commits to a fact that is '
"different from the gold answer, or fails to flag a false "
"premise when the question contains one.\n\n"
"Special case: if the question contains a false premise and the "
"gold answer says so, then a prediction that flags the false "
"premise is \"correct\".\n\n"
'premise is "correct".\n\n'
"Respond with ONLY a JSON object on a single line:\n"
'{\"grade\": \"correct\"|\"missing\"|\"incorrect\", \"rationale\": \"<one short sentence>\"}'
'{"grade": "correct"|"missing"|"incorrect", "rationale": "<one short sentence>"}'
)
@ -444,15 +473,17 @@ def _parse_judge_response(text: str) -> tuple[GradeClass, str]:
# Methods that should *not* trigger the LLM judge — the deterministic
# verdict is conclusive (refusal, exact match, numeric mismatch, etc.).
_TERMINAL_METHODS = frozenset({
"refusal",
"exact",
"numeric",
"substring",
"substring_reverse",
"false_premise_flagged",
"empty_gold",
})
_TERMINAL_METHODS = frozenset(
{
"refusal",
"exact",
"numeric",
"substring",
"substring_reverse",
"false_premise_flagged",
"empty_gold",
}
)
async def grade_with_judge(

View file

@ -42,7 +42,7 @@ class ExtractionResult:
"""Outcome of converting one HTML blob to plain markdown."""
text: str
method: str # "trafilatura" | "fallback_strip" | "empty"
method: str # "trafilatura" | "fallback_strip" | "empty"
n_chars: int
@property
@ -94,11 +94,30 @@ class _StripHTMLParser(HTMLParser):
"""
_SKIP_TAGS = frozenset({"script", "style", "nav", "header", "footer", "aside", "svg"})
_BLOCK_TAGS = frozenset({
"p", "div", "section", "article", "li", "ul", "ol",
"h1", "h2", "h3", "h4", "h5", "h6", "br", "tr",
"td", "th", "table", "blockquote", "pre",
})
_BLOCK_TAGS = frozenset(
{
"p",
"div",
"section",
"article",
"li",
"ul",
"ol",
"h1",
"h2",
"h3",
"h4",
"h5",
"h6",
"br",
"tr",
"td",
"th",
"table",
"blockquote",
"pre",
}
)
def __init__(self) -> None:
super().__init__(convert_charrefs=True)

View file

@ -158,7 +158,10 @@ def _materialise_pages(
logger.info(
"CRAG page extraction: %s; empty=%d, total_files=%d across %d questions",
method_counts, n_empty, len(file_to_url), len(qid_to_files),
method_counts,
n_empty,
len(file_to_url),
len(qid_to_files),
)
return qid_to_files, file_to_url
@ -215,8 +218,10 @@ async def _upload_pages(
name_to_id[f"{s.title}.md"] = s.document_id
logger.info(
"CRAG upload batch %d-%d: %d new, %d duplicate",
batch_start, batch_start + len(batch),
len(result.document_ids), len(result.duplicate_document_ids),
batch_start,
batch_start + len(batch),
len(result.document_ids),
len(result.duplicate_document_ids),
)
return name_to_id
@ -243,24 +248,26 @@ def _resolve_question_doc_ids(
doc_ids.append(doc_id)
else:
missing.append(fn)
rows.append({
"qid": q.qid,
"interaction_id": q.interaction_id,
"raw_index": q.raw_index,
"question": q.query,
"gold_answer": q.gold_answer,
"alt_answers": list(q.alt_answers),
"domain": q.domain,
"question_type": q.question_type,
"static_or_dynamic": q.static_or_dynamic,
"popularity": q.popularity,
"query_time": q.query_time,
"split": q.split,
"page_filenames": filenames,
"document_ids": doc_ids,
"missing_pages": missing,
"n_pages": len(filenames),
})
rows.append(
{
"qid": q.qid,
"interaction_id": q.interaction_id,
"raw_index": q.raw_index,
"question": q.query,
"gold_answer": q.gold_answer,
"alt_answers": list(q.alt_answers),
"domain": q.domain,
"question_type": q.question_type,
"static_or_dynamic": q.static_or_dynamic,
"popularity": q.popularity,
"query_time": q.query_time,
"split": q.split,
"page_filenames": filenames,
"document_ids": doc_ids,
"missing_pages": missing,
"n_pages": len(filenames),
}
)
return rows
@ -305,7 +312,7 @@ async def run_ingest(
settings = settings or IngestSettings(
use_vision_llm=False,
processing_mode="basic",
)
)
bench_dir = ctx.benchmark_data_dir()
pages_dir = bench_dir / "pages"
raw_cache = bench_dir / ".raw_cache"
@ -336,10 +343,13 @@ async def run_ingest(
n_pages_total = sum(len(q.pages) for q in questions)
logger.info(
"CRAG: extracting up to %d pages across %d questions ...",
n_pages_total, len(questions),
n_pages_total,
len(questions),
)
qid_to_files, file_to_url = _materialise_pages(
questions, pages_dir=pages_dir, overwrite=overwrite_extract,
questions,
pages_dir=pages_dir,
overwrite=overwrite_extract,
)
n_pages_extracted = sum(len(v) for v in qid_to_files.values())

View file

@ -37,7 +37,7 @@ _BASE_INSTRUCTIONS = (
"is factually wrong), say so explicitly in your final answer "
"rather than answering as if the premise were true.\n"
"2. If you are not confident in an answer, prefer saying \"I don't "
"know\" over guessing. A wrong commit is penalised more than a "
'know" over guessing. A wrong commit is penalised more than a '
"refusal.\n"
"3. Keep the final answer short — a name, a number, a date, a "
"phrase. Do not repeat the question.\n\n"
@ -125,9 +125,7 @@ def build_long_context_prompt(
if len(body) > per_page_char_cap:
body = body[:per_page_char_cap].rstrip() + "\n[...truncated...]"
title_clean = (title or f"page_{idx}").strip().replace("\n", " ")
blocks.append(
f"--- PAGE {idx}: {title_clean} ---\n{body}\n"
)
blocks.append(f"--- PAGE {idx}: {title_clean} ---\n{body}\n")
contexts_block = "\n".join(blocks) if blocks else "(no pages retrieved)"
return _LONG_CONTEXT_TEMPLATE.format(
instructions=_BASE_INSTRUCTIONS,

View file

@ -125,21 +125,23 @@ def _filter_questions(
continue
if qtype_filter and qtype_filter not in qtype:
continue
out.append(CragRunnerQuestion(
qid=str(row.get("qid") or "").strip(),
raw_index=int(row.get("raw_index") or 0),
question=str(row.get("question") or "").strip(),
gold_answer=str(row.get("gold_answer") or "").strip(),
alt_answers=list(row.get("alt_answers") or []),
domain=domain,
question_type=qtype,
static_or_dynamic=str(row.get("static_or_dynamic") or "").lower(),
popularity=str(row.get("popularity") or "").lower(),
query_time=str(row.get("query_time") or "").strip(),
page_filenames=list(row.get("page_filenames") or []),
document_ids=list(row.get("document_ids") or []),
missing_pages=list(row.get("missing_pages") or []),
))
out.append(
CragRunnerQuestion(
qid=str(row.get("qid") or "").strip(),
raw_index=int(row.get("raw_index") or 0),
question=str(row.get("question") or "").strip(),
gold_answer=str(row.get("gold_answer") or "").strip(),
alt_answers=list(row.get("alt_answers") or []),
domain=domain,
question_type=qtype,
static_or_dynamic=str(row.get("static_or_dynamic") or "").lower(),
popularity=str(row.get("popularity") or "").lower(),
query_time=str(row.get("query_time") or "").strip(),
page_filenames=list(row.get("page_filenames") or []),
document_ids=list(row.get("document_ids") or []),
missing_pages=list(row.get("missing_pages") or []),
)
)
out.sort(key=lambda q: q.raw_index)
if sample_n is not None and sample_n > 0:
out = out[:sample_n]
@ -190,15 +192,22 @@ class CragBenchmark:
def add_run_args(self, parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--n", dest="sample_n", type=int, default=None,
"--n",
dest="sample_n",
type=int,
default=None,
help="Run only the first N questions after filters.",
)
parser.add_argument(
"--domain", dest="domain_filter", default=None,
"--domain",
dest="domain_filter",
default=None,
help="Filter to a single CRAG domain (finance|music|movie|sports|open).",
)
parser.add_argument(
"--qtype", dest="qtype_filter", default=None,
"--qtype",
dest="qtype_filter",
default=None,
help=(
"Filter to questions whose question_type contains this "
"substring (case-insensitive). Examples: 'multi-hop', "
@ -206,31 +215,46 @@ class CragBenchmark:
),
)
parser.add_argument(
"--concurrency", type=int, default=4,
"--concurrency",
type=int,
default=4,
help="Parallel question workers per arm.",
)
parser.add_argument(
"--max-output-tokens", type=int, default=512,
"--max-output-tokens",
type=int,
default=512,
help="Cap on completion length for the chat-completion arms.",
)
parser.add_argument(
"--per-page-char-cap", dest="per_page_char_cap", type=int, default=12_000,
"--per-page-char-cap",
dest="per_page_char_cap",
type=int,
default=12_000,
help="Long-context arm: max chars per page before truncation (default 12k).",
)
parser.add_argument(
"--skip-bare", dest="skip_bare", action="store_true",
"--skip-bare",
dest="skip_bare",
action="store_true",
help="Skip the bare-LLM arm (saves cost on re-runs).",
)
parser.add_argument(
"--skip-long-context", dest="skip_long_context", action="store_true",
"--skip-long-context",
dest="skip_long_context",
action="store_true",
help="Skip the long-context arm.",
)
parser.add_argument(
"--skip-surfsense", dest="skip_surfsense", action="store_true",
"--skip-surfsense",
dest="skip_surfsense",
action="store_true",
help="Skip the SurfSense arm (useful when iterating on the LLM arms only).",
)
parser.add_argument(
"--no-mention-scope", dest="no_mention_scope", action="store_true",
"--no-mention-scope",
dest="no_mention_scope",
action="store_true",
help=(
"SurfSense arm: don't pass mentioned_document_ids; let "
"the agent retrieve over the entire SearchSpace. Default "
@ -239,37 +263,56 @@ class CragBenchmark:
),
)
parser.add_argument(
"--no-judge", dest="no_judge", action="store_true",
"--no-judge",
dest="no_judge",
action="store_true",
help="Disable the LLM-as-judge fallback grader.",
)
parser.add_argument(
"--judge-model", dest="judge_model",
"--judge-model",
dest="judge_model",
default="anthropic/claude-sonnet-4.5",
help="OpenRouter slug for the LLM judge.",
)
parser.add_argument(
"--judge-concurrency", dest="judge_concurrency", type=int, default=4,
"--judge-concurrency",
dest="judge_concurrency",
type=int,
default=4,
help="Parallel judge calls.",
)
# Ingest knobs
parser.add_argument(
"--n-questions", dest="n_questions", type=int, default=None,
"--n-questions",
dest="n_questions",
type=int,
default=None,
help="(ingest only) cap on number of questions to materialise + ingest.",
)
parser.add_argument(
"--upload-batch-size", dest="upload_batch_size", type=int, default=16,
"--upload-batch-size",
dest="upload_batch_size",
type=int,
default=16,
help="(ingest only) markdown files per fileupload call.",
)
parser.add_argument(
"--skip-upload", dest="skip_upload", action="store_true",
"--skip-upload",
dest="skip_upload",
action="store_true",
help="(ingest only) extract pages locally but don't push to SurfSense.",
)
parser.add_argument(
"--overwrite-extract", dest="overwrite_extract", action="store_true",
"--overwrite-extract",
dest="overwrite_extract",
action="store_true",
help="(ingest only) re-run trafilatura even when cached markdown exists.",
)
parser.add_argument(
"--sample-seed", dest="sample_seed", type=int, default=17,
"--sample-seed",
dest="sample_seed",
type=int,
default=17,
help="(ingest only) RNG seed for the stratified sample.",
)
add_ingest_settings_args(parser, defaults=_DEFAULT_INGEST_SETTINGS)
@ -362,12 +405,14 @@ class CragBenchmark:
if not api_key:
logger.warning("CRAG: --no-judge implied (no OPENROUTER_API_KEY for judge)")
else:
judge = CragLlmJudge(config=CragJudgeConfig(
api_key=api_key,
model=judge_model,
base_url=ctx.config.openrouter_base_url,
concurrency=judge_concurrency,
))
judge = CragLlmJudge(
config=CragJudgeConfig(
api_key=api_key,
model=judge_model,
base_url=ctx.config.openrouter_base_url,
concurrency=judge_concurrency,
)
)
run_timestamp = utc_iso_timestamp()
run_dir = ctx.runs_dir(run_timestamp=run_timestamp)
@ -393,29 +438,53 @@ class CragBenchmark:
# internally concurrency-bounded.
tasks: list[Any] = []
if bare_arm is not None:
tasks.append(_gather_with_limit((_bare_one(q) for q in questions), concurrency=concurrency))
tasks.append(
_gather_with_limit((_bare_one(q) for q in questions), concurrency=concurrency)
)
else:
tasks.append(_make_skipped_results(questions, "bare_llm"))
if long_context_arm is not None:
tasks.append(_gather_with_limit((_long_context_one(q) for q in questions), concurrency=concurrency))
tasks.append(
_gather_with_limit(
(_long_context_one(q) for q in questions), concurrency=concurrency
)
)
else:
tasks.append(_make_skipped_results(questions, "long_context"))
if surf_arm is not None:
tasks.append(_gather_with_limit((_surf_one(q) for q in questions), concurrency=concurrency))
tasks.append(
_gather_with_limit((_surf_one(q) for q in questions), concurrency=concurrency)
)
else:
tasks.append(_make_skipped_results(questions, "surfsense"))
bare_results, long_context_results, surf_results = await asyncio.gather(*tasks)
bare_grades = await _grade_results(questions, bare_results, judge=judge) if bare_arm else _empty_grades(questions)
lc_grades = await _grade_results(questions, long_context_results, judge=judge) if long_context_arm else _empty_grades(questions)
surf_grades = await _grade_results(questions, surf_results, judge=judge) if surf_arm else _empty_grades(questions)
bare_grades = (
await _grade_results(questions, bare_results, judge=judge)
if bare_arm
else _empty_grades(questions)
)
lc_grades = (
await _grade_results(questions, long_context_results, judge=judge)
if long_context_arm
else _empty_grades(questions)
)
surf_grades = (
await _grade_results(questions, surf_results, judge=judge)
if surf_arm
else _empty_grades(questions)
)
with raw_path.open("w", encoding="utf-8") as fh:
for q, b_res, l_res, s_res, b_g, l_g, s_g in zip(
questions,
bare_results, long_context_results, surf_results,
bare_grades, lc_grades, surf_grades,
bare_results,
long_context_results,
surf_results,
bare_grades,
lc_grades,
surf_grades,
strict=False,
):
meta = {
@ -431,18 +500,29 @@ class CragBenchmark:
"alt_answers": q.alt_answers,
}
for res, grade in (
(b_res, b_g), (l_res, l_g), (s_res, s_g),
(b_res, b_g),
(l_res, l_g),
(s_res, s_g),
):
fh.write(json.dumps({
**meta,
**res.to_jsonl(),
"graded": grade.to_dict(),
}) + "\n")
fh.write(
json.dumps(
{
**meta,
**res.to_jsonl(),
"graded": grade.to_dict(),
}
)
+ "\n"
)
metrics = _compute_metrics(
questions=questions,
bare_results=bare_results, long_context_results=long_context_results, surf_results=surf_results,
bare_grades=bare_grades, lc_grades=lc_grades, surf_grades=surf_grades,
bare_results=bare_results,
long_context_results=long_context_results,
surf_results=surf_results,
bare_grades=bare_grades,
lc_grades=lc_grades,
surf_grades=surf_grades,
arms_active={
"bare_llm": bare_arm is not None,
"long_context": long_context_arm is not None,
@ -481,13 +561,18 @@ class CragBenchmark:
manifest_path = run_dir / "run_artifact.json"
manifest_path.write_text(
json.dumps({
"suite": self.suite,
"benchmark": self.name,
"raw_path": "raw.jsonl",
"metrics": metrics,
"extra": artifact.extra,
}, indent=2, sort_keys=True) + "\n",
json.dumps(
{
"suite": self.suite,
"benchmark": self.name,
"raw_path": "raw.jsonl",
"metrics": metrics,
"extra": artifact.extra,
},
indent=2,
sort_keys=True,
)
+ "\n",
encoding="utf-8",
)
return artifact
@ -547,7 +632,9 @@ class CragBenchmark:
body_lines.append("- Headline truthfulness scores (CRAG paper rubric):")
for label, key in (
("Bare LLM", "bare_llm"), ("Long-Context", "long_context"), ("SurfSense", "surfsense"),
("Bare LLM", "bare_llm"),
("Long-Context", "long_context"),
("SurfSense", "surfsense"),
):
d = m.get(key, {})
body_lines.append(
@ -583,9 +670,7 @@ class CragBenchmark:
for arm in ("bare_llm", "long_context", "surfsense"):
if arm not in row:
continue
pieces.append(
f"{arm}={_signed_pct(row[arm].get('truthfulness_score'))}"
)
pieces.append(f"{arm}={_signed_pct(row[arm].get('truthfulness_score'))}")
body_lines.append(" ".join(pieces))
if per_qtype:
@ -596,9 +681,7 @@ class CragBenchmark:
for arm in ("bare_llm", "long_context", "surfsense"):
if arm not in row:
continue
pieces.append(
f"{arm}={_signed_pct(row[arm].get('truthfulness_score'))}"
)
pieces.append(f"{arm}={_signed_pct(row[arm].get('truthfulness_score'))}")
body_lines.append(" ".join(pieces))
return ReportSection(
@ -669,32 +752,31 @@ async def _grade_results(
rows: list[CragGradeRow] = []
for q, r in zip(questions, results, strict=False):
pred = extract_freeform_answer(r.raw_text or "")
rows.append(CragGradeRow(
qid=q.qid,
question=q.question,
gold=q.gold_answer,
alt_answers=q.alt_answers,
pred=pred,
question_type=q.question_type,
))
rows.append(
CragGradeRow(
qid=q.qid,
question=q.question,
gold=q.gold_answer,
alt_answers=q.alt_answers,
pred=pred,
question_type=q.question_type,
)
)
return await grade_many(rows=rows, judge=judge)
def _empty_grades(questions: list[CragRunnerQuestion]) -> list[CragGradeResult]:
return [
CragGradeResult(grade="missing", score=0, method="skipped_arm")
for _ in questions
]
return [CragGradeResult(grade="missing", score=0, method="skipped_arm") for _ in questions]
async def _make_skipped_results(
questions: list[CragRunnerQuestion], arm_name: str,
questions: list[CragRunnerQuestion],
arm_name: str,
) -> list[ArmResult]:
"""Stand-in results so downstream code can assume parallel lists."""
return [
ArmResult(arm=arm_name, question_id=q.qid, raw_text="", error="skipped")
for q in questions
ArmResult(arm=arm_name, question_id=q.qid, raw_text="", error="skipped") for q in questions
]
@ -776,20 +858,41 @@ def _compute_metrics(
deltas: dict[str, Any] = {}
for label, ref_correct, ref_t, chal_correct, chal_t, both_active in (
("surfsense_vs_bare", bare_correct, bare_t, surf_correct, surf_t,
arms_active.get("bare_llm") and arms_active.get("surfsense")),
("surfsense_vs_long_context", lc_correct, lc_t, surf_correct, surf_t,
arms_active.get("long_context") and arms_active.get("surfsense")),
("long_context_vs_bare", bare_correct, bare_t, lc_correct, lc_t,
arms_active.get("bare_llm") and arms_active.get("long_context")),
(
"surfsense_vs_bare",
bare_correct,
bare_t,
surf_correct,
surf_t,
arms_active.get("bare_llm") and arms_active.get("surfsense"),
),
(
"surfsense_vs_long_context",
lc_correct,
lc_t,
surf_correct,
surf_t,
arms_active.get("long_context") and arms_active.get("surfsense"),
),
(
"long_context_vs_bare",
bare_correct,
bare_t,
lc_correct,
lc_t,
arms_active.get("bare_llm") and arms_active.get("long_context"),
),
):
if not both_active:
continue
mc = mcnemar_test(ref_correct, chal_correct)
boot = bootstrap_delta_ci(ref_correct, chal_correct, n_resamples=2000)
deltas[label] = {
"accuracy_pp": 100.0 * (sum(chal_correct) - sum(ref_correct)) / max(1, len(chal_correct)),
"truthfulness_score_pp": 100.0 * (chal_t["truthfulness_score"] - ref_t["truthfulness_score"]),
"accuracy_pp": 100.0
* (sum(chal_correct) - sum(ref_correct))
/ max(1, len(chal_correct)),
"truthfulness_score_pp": 100.0
* (chal_t["truthfulness_score"] - ref_t["truthfulness_score"]),
"mcnemar_p_value": mc.p_value,
"mcnemar_method": mc.method,
"mcnemar_b_ref_only": mc.b,
@ -800,12 +903,18 @@ def _compute_metrics(
out["deltas"] = deltas
out["per_domain"] = _per_facet_truthfulness(
questions, bare_grades, lc_grades, surf_grades,
questions,
bare_grades,
lc_grades,
surf_grades,
arms_active=arms_active,
key_fn=lambda q: q.domain or "(unspecified)",
)
out["per_question_type"] = _per_facet_truthfulness(
questions, bare_grades, lc_grades, surf_grades,
questions,
bare_grades,
lc_grades,
surf_grades,
arms_active=arms_active,
key_fn=lambda q: q.question_type or "(unspecified)",
)
@ -867,11 +976,11 @@ def _arm_summary_lines(d: dict[str, Any], *, indent: str) -> str:
high = d.get("ci_high", 0.0)
lines = [
f"{indent}- Accuracy: {acc * 100:.1f}% (Wilson 95% CI: {low * 100:.1f}% {high * 100:.1f}%)",
f"{indent}- 3-class: correct={d.get('correct_rate', 0)*100:.1f}%, "
f"missing={d.get('missing_rate', 0)*100:.1f}%, "
f"incorrect={d.get('incorrect_rate', 0)*100:.1f}%",
f"{indent}- 3-class: correct={d.get('correct_rate', 0) * 100:.1f}%, "
f"missing={d.get('missing_rate', 0) * 100:.1f}%, "
f"incorrect={d.get('incorrect_rate', 0) * 100:.1f}%",
f"{indent}- Truthfulness score (correct - incorrect)/total: "
f"{d.get('truthfulness_score', 0)*100:+.1f}%",
f"{d.get('truthfulness_score', 0) * 100:+.1f}%",
f"{indent}- Cost / question: ${_dollars(d.get('cost_micros_mean'))} (mean), "
f"${_dollars(d.get('cost_micros_median'))} (median)",
f"{indent}- Latency: p50 {_ms_to_s(d.get('latency_ms_median'))}, "
@ -916,7 +1025,7 @@ def _pct(value: Any) -> str:
if value is None:
return "?"
try:
return f"{float(value)*100:.1f}%"
return f"{float(value) * 100:.1f}%"
except (TypeError, ValueError):
return "?"
@ -925,7 +1034,7 @@ def _signed_pct(value: Any) -> str:
if value is None:
return "?"
try:
return f"{float(value)*100:+.1f}%"
return f"{float(value) * 100:+.1f}%"
except (TypeError, ValueError):
return "?"

View file

@ -51,12 +51,12 @@ def _hf_hub_download(*args: Any, **kwargs: Any) -> str:
class FramesQuestion:
"""One row of FRAMES (post-parse)."""
qid: str # synthesised "Q000" .. "Q823"
qid: str # synthesised "Q000" .. "Q823"
question: str
gold_answer: str
wiki_urls: list[str] # deduped, in original order
reasoning_types: list[str] # split on "|"
raw_index: int # row index from the TSV (for debugging)
wiki_urls: list[str] # deduped, in original order
reasoning_types: list[str] # split on "|"
raw_index: int # row index from the TSV (for debugging)
def to_dict(self) -> dict[str, Any]:
return {
@ -146,14 +146,16 @@ def load_questions(tsv_path: Path) -> list[FramesQuestion]:
if val and val not in urls:
urls.append(val)
reasoning = _parse_reasoning_types(row.get("reasoning_types"))
out.append(FramesQuestion(
qid=f"Q{int(raw_idx):03d}",
question=prompt,
gold_answer=answer,
wiki_urls=urls,
reasoning_types=reasoning,
raw_index=int(raw_idx),
))
out.append(
FramesQuestion(
qid=f"Q{int(raw_idx):03d}",
question=prompt,
gold_answer=answer,
wiki_urls=urls,
reasoning_types=reasoning,
raw_index=int(raw_idx),
)
)
return out

View file

@ -90,10 +90,27 @@ def _normalise(s: str) -> str:
_WORD_NUMBERS = {
"zero": 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5,
"six": 6, "seven": 7, "eight": 8, "nine": 9, "ten": 10, "eleven": 11,
"twelve": 12, "thirteen": 13, "fourteen": 14, "fifteen": 15, "sixteen": 16,
"seventeen": 17, "eighteen": 18, "nineteen": 19, "twenty": 20,
"zero": 0,
"one": 1,
"two": 2,
"three": 3,
"four": 4,
"five": 5,
"six": 6,
"seven": 7,
"eight": 8,
"nine": 9,
"ten": 10,
"eleven": 11,
"twelve": 12,
"thirteen": 13,
"fourteen": 14,
"fifteen": 15,
"sixteen": 16,
"seventeen": 17,
"eighteen": 18,
"nineteen": 19,
"twenty": 20,
}
_NUMERIC_RE = re.compile(r"-?\d+(?:[.,]\d+)?")
@ -194,7 +211,7 @@ _JUDGE_SYSTEM = (
"expresses a different fact, omits the central answer, or hedges "
"without committing.\n\n"
"Respond with ONLY a JSON object on a single line:\n"
'{\"correct\": true|false, \"rationale\": \"<one short sentence>\"}'
'{"correct": true|false, "rationale": "<one short sentence>"}'
)
@ -324,10 +341,7 @@ async def grade_many(
if not rows:
return []
coros = [
grade_with_judge(pred=p, gold=g, question=q, judge=judge)
for _qid, q, g, p in rows
]
coros = [grade_with_judge(pred=p, gold=g, question=q, judge=judge) for _qid, q, g, p in rows]
return list(await asyncio.gather(*coros))

View file

@ -160,8 +160,10 @@ async def _upload_markdowns(
name_to_id[s.title] = s.document_id
logger.info(
"FRAMES upload batch %d-%d: %d new, %d duplicate",
batch_start, batch_start + len(batch),
len(result.document_ids), len(result.duplicate_document_ids),
batch_start,
batch_start + len(batch),
len(result.document_ids),
len(result.duplicate_document_ids),
)
return name_to_id
@ -188,14 +190,16 @@ def _resolve_question_doc_ids(
doc_id = name_to_id.get(stem) or name_to_id.get(article.markdown_path.name)
if doc_id is not None and doc_id not in doc_ids:
doc_ids.append(doc_id)
rows.append({
"qid": q.qid,
"raw_index": q.raw_index,
"n_wiki_urls": len(q.wiki_urls),
"wiki_titles": titles,
"document_ids": doc_ids,
"missing_urls": missing,
})
rows.append(
{
"qid": q.qid,
"raw_index": q.raw_index,
"n_wiki_urls": len(q.wiki_urls),
"wiki_titles": titles,
"document_ids": doc_ids,
"missing_urls": missing,
}
)
return rows
@ -238,7 +242,7 @@ async def run_ingest(
settings = settings or IngestSettings(
use_vision_llm=False,
processing_mode="basic",
)
)
bench_dir = ctx.benchmark_data_dir()
wiki_cache = bench_dir / "wiki"
wiki_cache.mkdir(parents=True, exist_ok=True)
@ -250,8 +254,7 @@ async def run_ingest(
questions = load_questions(tsv_path)
if not questions:
raise RuntimeError(
"FRAMES test.tsv contained no parseable rows; upstream may "
"have changed schema."
"FRAMES test.tsv contained no parseable rows; upstream may have changed schema."
)
logger.info("FRAMES: parsed %d questions from %s", len(questions), tsv_path.name)
if max_questions is not None and max_questions > 0:
@ -269,19 +272,23 @@ async def run_ingest(
unique_urls = list(seen_urls.keys())
logger.info(
"FRAMES: %d unique Wikipedia URLs across %d questions",
len(unique_urls), len(questions),
len(unique_urls),
len(questions),
)
# 3. Fetch (with cache).
fetcher = WikiFetcher(cache_dir=wiki_cache, rate_limit_rps=fetch_rate_limit_rps)
n_cached = sum(
1 for url in unique_urls
1
for url in unique_urls
if (wiki_cache / cache_filename_for_title(_safe_title(url))).exists()
)
fetched, missing_urls = await _fetch_articles(fetcher, unique_urls)
logger.info(
"FRAMES: fetched=%d, cache_hits=%d, missing=%d",
len(fetched), n_cached, len(missing_urls),
len(fetched),
n_cached,
len(missing_urls),
)
# 4. Upload to SurfSense (deduped by filename).

View file

@ -65,7 +65,7 @@ class FramesRunnerQuestion:
question: str
gold_answer: str
reasoning_types: list[str]
document_ids: list[int] # subset of corpus relevant to this Q (may be empty)
document_ids: list[int] # subset of corpus relevant to this Q (may be empty)
n_wiki_urls: int
missing_urls: list[str]
@ -107,16 +107,18 @@ def _load_questions(
reasoning = list(row.get("reasoning_types") or [])
if reasoning_filter and reasoning_filter not in [r.lower() for r in reasoning]:
continue
out.append(FramesRunnerQuestion(
qid=qid,
raw_index=int(row.get("raw_index") or 0),
question=str(row.get("question") or "").strip(),
gold_answer=str(row.get("gold_answer") or "").strip(),
reasoning_types=reasoning,
document_ids=list(map_row.get("document_ids") or []),
n_wiki_urls=int(map_row.get("n_wiki_urls") or 0),
missing_urls=list(map_row.get("missing_urls") or []),
))
out.append(
FramesRunnerQuestion(
qid=qid,
raw_index=int(row.get("raw_index") or 0),
question=str(row.get("question") or "").strip(),
gold_answer=str(row.get("gold_answer") or "").strip(),
reasoning_types=reasoning,
document_ids=list(map_row.get("document_ids") or []),
n_wiki_urls=int(map_row.get("n_wiki_urls") or 0),
missing_urls=list(map_row.get("missing_urls") or []),
)
)
out.sort(key=lambda q: q.raw_index)
if sample_n is not None and sample_n > 0:
out = out[:sample_n]
@ -166,7 +168,10 @@ class FramesBenchmark:
def add_run_args(self, parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--n", dest="sample_n", type=int, default=None,
"--n",
dest="sample_n",
type=int,
default=None,
help="Run only the first N questions after filters (default: all 824).",
)
parser.add_argument(
@ -180,11 +185,15 @@ class FramesBenchmark:
),
)
parser.add_argument(
"--concurrency", type=int, default=4,
"--concurrency",
type=int,
default=4,
help="Parallel question workers per arm.",
)
parser.add_argument(
"--scope-mentions", dest="scope_mentions", action="store_true",
"--scope-mentions",
dest="scope_mentions",
action="store_true",
help=(
"SurfSense arm: scope retrieval to the per-question "
"document_ids (oracle-retrieval upper bound). Default "
@ -192,11 +201,15 @@ class FramesBenchmark:
),
)
parser.add_argument(
"--max-output-tokens", type=int, default=512,
"--max-output-tokens",
type=int,
default=512,
help="Cap on completion length for both arms.",
)
parser.add_argument(
"--no-judge", dest="no_judge", action="store_true",
"--no-judge",
dest="no_judge",
action="store_true",
help=(
"Disable LLM-as-judge fallback grading; use only the "
"deterministic grader (faster but more pessimistic)."
@ -217,19 +230,30 @@ class FramesBenchmark:
)
# Ingest-only knobs.
parser.add_argument(
"--max-questions", dest="max_questions", type=int, default=None,
"--max-questions",
dest="max_questions",
type=int,
default=None,
help="(ingest only) cap on number of questions to materialise + ingest.",
)
parser.add_argument(
"--upload-batch-size", dest="upload_batch_size", type=int, default=16,
"--upload-batch-size",
dest="upload_batch_size",
type=int,
default=16,
help="(ingest only) markdown files per fileupload call.",
)
parser.add_argument(
"--skip-upload", dest="skip_upload", action="store_true",
"--skip-upload",
dest="skip_upload",
action="store_true",
help="(ingest only) cache wiki articles locally but don't push to SurfSense.",
)
parser.add_argument(
"--fetch-rps", dest="fetch_rate_limit_rps", type=float, default=2.0,
"--fetch-rps",
dest="fetch_rate_limit_rps",
type=float,
default=2.0,
help="(ingest only) max requests/second to the Wikipedia API.",
)
add_ingest_settings_args(parser, defaults=_DEFAULT_INGEST_SETTINGS)
@ -270,21 +294,18 @@ class FramesBenchmark:
doc_map, ingest_settings = _load_doc_map(map_path)
questions = _load_questions(
questions_jsonl, doc_map,
questions_jsonl,
doc_map,
sample_n=sample_n,
reasoning_filter=reasoning_filter,
)
if not questions:
raise RuntimeError(
"No FRAMES questions matched the filters; broaden --reasoning/--n."
)
raise RuntimeError("No FRAMES questions matched the filters; broaden --reasoning/--n.")
logger.info("FRAMES: scheduled %d questions", len(questions))
api_key = os.environ.get("OPENROUTER_API_KEY")
if not api_key:
raise RuntimeError(
"OPENROUTER_API_KEY env var is required for the bare-LLM arm."
)
raise RuntimeError("OPENROUTER_API_KEY env var is required for the bare-LLM arm.")
bare_provider = OpenRouterChatProvider(
api_key=api_key,
@ -303,12 +324,14 @@ class FramesBenchmark:
judge: LlmJudge | None = None
if not no_judge:
judge = LlmJudge(config=JudgeConfig(
api_key=api_key,
model=judge_model,
base_url=ctx.config.openrouter_base_url,
concurrency=judge_concurrency,
))
judge = LlmJudge(
config=JudgeConfig(
api_key=api_key,
model=judge_model,
base_url=ctx.config.openrouter_base_url,
concurrency=judge_concurrency,
)
)
run_timestamp = utc_iso_timestamp()
run_dir = ctx.runs_dir(run_timestamp=run_timestamp)
@ -318,9 +341,7 @@ class FramesBenchmark:
return await bare_arm.answer(_make_bare_request(q, max_output_tokens))
async def _surf_one(q: FramesRunnerQuestion) -> ArmResult:
return await surf_arm.answer(
_make_surfsense_request(q, scope_mentions=scope_mentions)
)
return await surf_arm.answer(_make_surfsense_request(q, scope_mentions=scope_mentions))
bare_results, surf_results = await asyncio.gather(
_gather_with_limit((_bare_one(q) for q in questions), concurrency=concurrency),
@ -343,16 +364,26 @@ class FramesBenchmark:
"n_missing_urls": len(q.missing_urls),
"gold": q.gold_answer,
}
fh.write(json.dumps({
**meta,
**b_res.to_jsonl(),
"graded": b_g.to_dict(),
}) + "\n")
fh.write(json.dumps({
**meta,
**s_res.to_jsonl(),
"graded": s_g.to_dict(),
}) + "\n")
fh.write(
json.dumps(
{
**meta,
**b_res.to_jsonl(),
"graded": b_g.to_dict(),
}
)
+ "\n"
)
fh.write(
json.dumps(
{
**meta,
**s_res.to_jsonl(),
"graded": s_g.to_dict(),
}
)
+ "\n"
)
metrics = _compute_metrics(questions, bare_results, surf_results, bare_grades, surf_grades)
artifact = RunArtifact(
@ -380,13 +411,18 @@ class FramesBenchmark:
manifest_path = run_dir / "run_artifact.json"
manifest_path.write_text(
json.dumps({
"suite": self.suite,
"benchmark": self.name,
"raw_path": "raw.jsonl",
"metrics": metrics,
"extra": artifact.extra,
}, indent=2, sort_keys=True) + "\n",
json.dumps(
{
"suite": self.suite,
"benchmark": self.name,
"raw_path": "raw.jsonl",
"metrics": metrics,
"extra": artifact.extra,
},
indent=2,
sort_keys=True,
)
+ "\n",
encoding="utf-8",
)
return artifact
@ -451,8 +487,8 @@ class FramesBenchmark:
for tag, vals in sorted(per_reasoning.items()):
body_lines.append(
f" - {tag}: SurfSense {_pp(vals.get('delta_accuracy_pp'))} pp "
f"(n={vals.get('n')}, bare acc={vals.get('bare_accuracy', 0)*100:.1f}%, "
f"surf acc={vals.get('surfsense_accuracy', 0)*100:.1f}%)"
f"(n={vals.get('n')}, bare acc={vals.get('bare_accuracy', 0) * 100:.1f}%, "
f"surf acc={vals.get('surfsense_accuracy', 0) * 100:.1f}%)"
)
return ReportSection(
@ -553,8 +589,7 @@ def _compute_metrics(
"bare_accuracy": (sum(b_correct) / len(pairs)) if pairs else 0.0,
"surfsense_accuracy": (sum(s_correct) / len(pairs)) if pairs else 0.0,
"delta_accuracy_pp": (
100.0 * (sum(s_correct) - sum(b_correct)) / len(pairs)
if pairs else 0.0
100.0 * (sum(s_correct) - sum(b_correct)) / len(pairs) if pairs else 0.0
),
}
@ -571,8 +606,12 @@ def _compute_metrics(
"latency_ms_mean": bare_latency_agg.mean,
"latency_ms_median": bare_latency_agg.median,
"latency_ms_p95": bare_latency_agg.p95,
"input_tokens_mean": (sum(bare_in_tokens) / len(bare_in_tokens)) if bare_in_tokens else 0.0,
"output_tokens_mean": (sum(bare_out_tokens) / len(bare_out_tokens)) if bare_out_tokens else 0.0,
"input_tokens_mean": (sum(bare_in_tokens) / len(bare_in_tokens))
if bare_in_tokens
else 0.0,
"output_tokens_mean": (sum(bare_out_tokens) / len(bare_out_tokens))
if bare_out_tokens
else 0.0,
},
"surfsense": {
**surf_acc.to_dict(),

View file

@ -49,10 +49,10 @@ USER_AGENT = (
class WikiArticle:
"""One fetched article + metadata."""
title: str # canonical title returned by MW (post-redirect)
source_url: str # the URL we were asked to fetch
markdown_path: Path # where the cached body lives on disk
n_chars: int # length of the body (post-prepend H1)
title: str # canonical title returned by MW (post-redirect)
source_url: str # the URL we were asked to fetch
markdown_path: Path # where the cached body lives on disk
n_chars: int # length of the body (post-prepend H1)
redirected_from: str | None = None
@ -168,10 +168,13 @@ class WikiFetcher:
break
except (httpx.HTTPError, RuntimeError) as exc:
last_exc = exc
wait = 1.0 * (2 ** attempt)
wait = 1.0 * (2**attempt)
logger.warning(
"wiki fetch %r attempt %d failed: %s; retry in %.1fs",
title, attempt + 1, exc, wait,
title,
attempt + 1,
exc,
wait,
)
await asyncio.sleep(wait)
else:
@ -217,10 +220,14 @@ class WikiFetcher:
}
headers = {"User-Agent": USER_AGENT, "Accept": "application/json"}
if http is not None:
response = await http.get(WIKI_API, params=params, headers=headers, timeout=self._timeout)
response = await http.get(
WIKI_API, params=params, headers=headers, timeout=self._timeout
)
else:
async with httpx.AsyncClient(timeout=self._timeout) as client:
response = await client.get(WIKI_API, params=params, headers=headers, timeout=self._timeout)
response = await client.get(
WIKI_API, params=params, headers=headers, timeout=self._timeout
)
response.raise_for_status()
data = response.json()
if "error" in data:

View file

@ -52,9 +52,7 @@ async def test_acquire_token_local_mode_posts_desktop_login_json():
200, json={"access_token": "T", "refresh_token": "R", "token_type": "bearer"}
)
)
config = _make_config(
surfsense_user_email="u@example.com", surfsense_user_password="pw"
)
config = _make_config(surfsense_user_email="u@example.com", surfsense_user_password="pw")
bundle = await acquire_token(config)
assert bundle.access_token == "T"
assert bundle.refresh_token == "R"

View file

@ -94,10 +94,18 @@ async def test_documents_status_parses_state(respx_mock, http):
200,
json={
"items": [
{"id": 1, "title": "a.pdf", "document_type": "FILE",
"status": {"state": "ready", "reason": None}},
{"id": 2, "title": "b.pdf", "document_type": "FILE",
"status": {"state": "failed", "reason": "ETL boom"}},
{
"id": 1,
"title": "a.pdf",
"document_type": "FILE",
"status": {"state": "ready", "reason": None},
},
{
"id": 2,
"title": "b.pdf",
"document_type": "FILE",
"status": {"state": "failed", "reason": "ETL boom"},
},
]
},
)
@ -137,14 +145,26 @@ async def test_documents_upload_returns_payload(respx_mock, http, tmp_path: Path
async def test_documents_list_chunks_paginated(respx_mock, http):
respx_mock.get("/api/v1/documents/5/chunks").mock(
side_effect=[
httpx.Response(200, json={
"items": [{"id": 1, "content": "a"}, {"id": 2, "content": "b"}],
"total": 3, "page": 0, "page_size": 2, "has_more": True,
}),
httpx.Response(200, json={
"items": [{"id": 3, "content": "c"}],
"total": 3, "page": 1, "page_size": 2, "has_more": False,
}),
httpx.Response(
200,
json={
"items": [{"id": 1, "content": "a"}, {"id": 2, "content": "b"}],
"total": 3,
"page": 0,
"page_size": 2,
"has_more": True,
},
),
httpx.Response(
200,
json={
"items": [{"id": 3, "content": "c"}],
"total": 3,
"page": 1,
"page_size": 2,
"has_more": False,
},
),
]
)
client = DocumentsClient(http, _BASE)
@ -191,15 +211,17 @@ def _sse_body(events: list[dict]) -> bytes:
@pytest.mark.asyncio
@respx.mock(base_url=_BASE)
async def test_ask_accumulates_text_deltas(respx_mock, http):
body = _sse_body([
{"type": "start", "messageId": "m1"},
{"type": "text-start", "id": "t1"},
{"type": "text-delta", "id": "t1", "delta": "Answer "},
{"type": "text-delta", "id": "t1", "delta": "is "},
{"type": "text-delta", "id": "t1", "delta": "B [citation:42]."},
{"type": "text-end", "id": "t1"},
{"type": "finish"},
])
body = _sse_body(
[
{"type": "start", "messageId": "m1"},
{"type": "text-start", "id": "t1"},
{"type": "text-delta", "id": "t1", "delta": "Answer "},
{"type": "text-delta", "id": "t1", "delta": "is "},
{"type": "text-delta", "id": "t1", "delta": "B [citation:42]."},
{"type": "text-end", "id": "t1"},
{"type": "finish"},
]
)
respx_mock.post("/api/v1/new_chat").mock(
return_value=httpx.Response(
200,
@ -208,9 +230,7 @@ async def test_ask_accumulates_text_deltas(respx_mock, http):
)
)
client = NewChatClient(http, _BASE)
answer = await client.ask(
thread_id=1, search_space_id=2, user_query="What is the answer?"
)
answer = await client.ask(thread_id=1, search_space_id=2, user_query="What is the answer?")
assert answer.text == "Answer is B [citation:42]."
assert answer.finished_normally is True
assert any(c["chunk_id"] == 42 for c in answer.citations)
@ -219,23 +239,21 @@ async def test_ask_accumulates_text_deltas(respx_mock, http):
@pytest.mark.asyncio
@respx.mock(base_url=_BASE)
async def test_ask_409_thread_busy_retries(respx_mock, http):
body = _sse_body([
{"type": "text-delta", "id": "t1", "delta": "ok"},
{"type": "finish"},
])
body = _sse_body(
[
{"type": "text-delta", "id": "t1", "delta": "ok"},
{"type": "finish"},
]
)
busy = httpx.Response(
409,
json={"detail": {"errorCode": "THREAD_BUSY", "message": "busy"}},
headers={"Retry-After": "1"},
)
success = httpx.Response(
200, content=body, headers={"Content-Type": "text/event-stream"}
)
success = httpx.Response(200, content=body, headers={"Content-Type": "text/event-stream"})
respx_mock.post("/api/v1/new_chat").mock(side_effect=[busy, success])
client = NewChatClient(http, _BASE)
answer = await client.ask(
thread_id=1, search_space_id=2, user_query="hi", max_busy_retries=2
)
answer = await client.ask(thread_id=1, search_space_id=2, user_query="hi", max_busy_retries=2)
assert answer.text == "ok"
@ -250,6 +268,4 @@ async def test_ask_409_exhausts_retries(respx_mock, http):
respx_mock.post("/api/v1/new_chat").mock(return_value=busy)
client = NewChatClient(http, _BASE)
with pytest.raises(ThreadBusyError):
await client.ask(
thread_id=1, search_space_id=2, user_query="hi", max_busy_retries=1
)
await client.ask(thread_id=1, search_space_id=2, user_query="hi", max_busy_retries=1)

View file

@ -46,32 +46,24 @@ class TestMerge:
def test_explicit_false_overrides_default_true(self) -> None:
defaults = IngestSettings(use_vision_llm=True)
merged = IngestSettings.merge(
defaults, {"use_vision_llm": False}
)
merged = IngestSettings.merge(defaults, {"use_vision_llm": False})
assert merged.use_vision_llm is False
def test_explicit_true_overrides_default_false(self) -> None:
defaults = IngestSettings(use_vision_llm=False)
merged = IngestSettings.merge(
defaults, {"use_vision_llm": True}
)
merged = IngestSettings.merge(defaults, {"use_vision_llm": True})
assert merged.use_vision_llm is True
def test_none_means_silent(self) -> None:
# Argparse with BooleanOptionalAction yields None when the
# operator passed neither --use-vision-llm nor --no-vision-llm.
defaults = IngestSettings(use_vision_llm=True)
merged = IngestSettings.merge(
defaults, {"use_vision_llm": None}
)
merged = IngestSettings.merge(defaults, {"use_vision_llm": None})
assert merged.use_vision_llm is True
def test_processing_mode_override(self) -> None:
defaults = IngestSettings(processing_mode="basic")
merged = IngestSettings.merge(
defaults, {"processing_mode": "premium"}
)
merged = IngestSettings.merge(defaults, {"processing_mode": "premium"})
assert merged.processing_mode == "premium"
def test_processing_mode_invalid_raises(self) -> None:
@ -134,9 +126,7 @@ class TestAddArgs:
p = argparse.ArgumentParser()
add_ingest_settings_args(
p,
defaults=IngestSettings(
use_vision_llm=False, processing_mode="basic"
),
defaults=IngestSettings(use_vision_llm=False, processing_mode="basic"),
)
return p
@ -158,31 +148,21 @@ class TestAddArgs:
args = parser.parse_args(["--processing-mode", mode])
assert args.processing_mode == mode
def test_processing_mode_rejects_unknown(
self, parser: argparse.ArgumentParser
) -> None:
def test_processing_mode_rejects_unknown(self, parser: argparse.ArgumentParser) -> None:
with pytest.raises(SystemExit):
parser.parse_args(["--processing-mode", "exotic"])
def test_vision_flags_mutually_exclusive(
self, parser: argparse.ArgumentParser
) -> None:
def test_vision_flags_mutually_exclusive(self, parser: argparse.ArgumentParser) -> None:
with pytest.raises(SystemExit):
parser.parse_args(["--use-vision-llm", "--no-vision-llm"])
def test_full_pipeline(self, parser: argparse.ArgumentParser) -> None:
# Operator passes flags + defaults are reasonable. Merge
# should yield exactly what they asked for.
args = parser.parse_args(
["--use-vision-llm", "--processing-mode", "premium"]
)
defaults = IngestSettings(
use_vision_llm=False, processing_mode="basic"
)
args = parser.parse_args(["--use-vision-llm", "--processing-mode", "premium"])
defaults = IngestSettings(use_vision_llm=False, processing_mode="basic")
merged = IngestSettings.merge(defaults, vars(args))
assert merged == IngestSettings(
use_vision_llm=True, processing_mode="premium"
)
assert merged == IngestSettings(use_vision_llm=True, processing_mode="premium")
# ---------------------------------------------------------------------------
@ -240,16 +220,12 @@ class TestHeader:
class TestFormatMd:
def test_full_settings(self) -> None:
out = format_ingest_settings_md(
{"use_vision_llm": True, "processing_mode": "premium"}
)
out = format_ingest_settings_md({"use_vision_llm": True, "processing_mode": "premium"})
assert "vision_llm=`on`" in out
assert "processing_mode=`premium`" in out
def test_default_off(self) -> None:
out = format_ingest_settings_md(
{"use_vision_llm": False, "processing_mode": "basic"}
)
out = format_ingest_settings_md({"use_vision_llm": False, "processing_mode": "basic"})
assert "vision_llm=`off`" in out
assert "processing_mode=`basic`" in out

View file

@ -25,7 +25,12 @@ from surfsense_evals.core.metrics import (
@pytest.mark.parametrize(
"k,n,low,high",
[
(80, 100, 0.7111, 0.8666), # cross-checked vs statsmodels.proportion_confint(method='wilson')
(
80,
100,
0.7111,
0.8666,
), # cross-checked vs statsmodels.proportion_confint(method='wilson')
(50, 100, 0.4038, 0.5962),
(0, 0, 0.0, 1.0),
(0, 10, 0.0, 0.2775),
@ -74,7 +79,7 @@ def test_mcnemar_exact_branch_strong_signal():
assert res.b == 0
assert res.c == 10
assert res.method == "exact"
expected = 2 * (0.5 ** 10)
expected = 2 * (0.5**10)
assert math.isclose(res.p_value, expected, rel_tol=1e-9)

View file

@ -11,7 +11,11 @@ from surfsense_evals.core.parse.answer_letter import AnswerLetterResult
@pytest.mark.parametrize(
"text,expected_letter,expected_strategy",
[
('```json\n{"step_by_step_thinking": "...", "answer_choice": "B"}\n```', "B", "json_envelope"),
(
'```json\n{"step_by_step_thinking": "...", "answer_choice": "B"}\n```',
"B",
"json_envelope",
),
('Reasoning... {"step_by_step_thinking": "x", "answer_choice": "C"}', "C", "json_envelope"),
("Long reasoning.\nAnswer: D", "D", "answer_line"),
("The correct answer is (A).", "A", "answer_line"),

View file

@ -91,7 +91,7 @@ def test_regex_pattern_matches_ts_source():
assert "https?://" in pattern
assert "urlcite" in pattern
assert "doc-" in pattern
assert "\u200B" in pattern
assert "\u200b" in pattern
assert "" in pattern and "" in pattern

View file

@ -44,11 +44,14 @@ class TestExtractFreeformAnswer:
assert extract_freeform_answer("ANSWER: yes") == "yes"
assert extract_freeform_answer("answer: no") == "no"
@pytest.mark.parametrize("text,expected", [
("Answer: 1, 2, 3", "1, 2, 3"),
("Answer: 3.14", "3.14"),
("Answer: spaced ", "spaced"),
])
@pytest.mark.parametrize(
"text,expected",
[
("Answer: 1, 2, 3", "1, 2, 3"),
("Answer: 3.14", "3.14"),
("Answer: spaced ", "spaced"),
],
)
def test_various_payloads(self, text: str, expected: str) -> None:
assert extract_freeform_answer(text) == expected

View file

@ -22,12 +22,16 @@ async def _astream(lines):
@pytest.mark.asyncio
async def test_basic_data_frame():
events = await _alist(
iter_sse_events(_astream([
'data: {"type": "text-delta", "delta": "hi"}',
"",
'data: {"type": "finish"}',
"",
]))
iter_sse_events(
_astream(
[
'data: {"type": "text-delta", "delta": "hi"}',
"",
'data: {"type": "finish"}',
"",
]
)
)
)
assert [e.data for e in events] == [
'{"type": "text-delta", "delta": "hi"}',
@ -38,10 +42,14 @@ async def test_basic_data_frame():
@pytest.mark.asyncio
async def test_done_sentinel_passes_through():
events = await _alist(
iter_sse_events(_astream([
"data: [DONE]",
"",
]))
iter_sse_events(
_astream(
[
"data: [DONE]",
"",
]
)
)
)
assert [e.data for e in events] == ["[DONE]"]
@ -49,11 +57,15 @@ async def test_done_sentinel_passes_through():
@pytest.mark.asyncio
async def test_multiline_data_joins_with_newline():
events = await _alist(
iter_sse_events(_astream([
"data: line1",
"data: line2",
"",
]))
iter_sse_events(
_astream(
[
"data: line1",
"data: line2",
"",
]
)
)
)
assert events[0].data == "line1\nline2"
@ -61,13 +73,17 @@ async def test_multiline_data_joins_with_newline():
@pytest.mark.asyncio
async def test_comments_and_other_fields_ignored():
events = await _alist(
iter_sse_events(_astream([
": heartbeat",
"event: foo",
"id: 123",
"data: payload",
"",
]))
iter_sse_events(
_astream(
[
": heartbeat",
"event: foo",
"id: 123",
"data: payload",
"",
]
)
)
)
assert [e.data for e in events] == ["payload"]
@ -77,8 +93,12 @@ async def test_handles_missing_trailing_blank():
"""Some servers omit the final blank line; the consumer should still emit."""
events = await _alist(
iter_sse_events(_astream([
"data: only-one",
]))
iter_sse_events(
_astream(
[
"data: only-one",
]
)
)
)
assert [e.data for e in events] == ["only-one"]

View file

@ -36,11 +36,18 @@ async def test_payload_shape_matches_openrouter_docs(respx_mock, tiny_pdf: Path)
return httpx.Response(
200,
json={
"choices": [{
"message": {"content": "Answer: B"},
"finish_reason": "stop",
}],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "cost": 0.0001},
"choices": [
{
"message": {"content": "Answer: B"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15,
"cost": 0.0001,
},
},
)
@ -63,8 +70,7 @@ async def test_payload_shape_matches_openrouter_docs(respx_mock, tiny_pdf: Path)
assert file_part["file"]["filename"] == tiny_pdf.name
assert file_part["file"]["file_data"].startswith("data:application/pdf;base64,")
assert (
base64.b64decode(file_part["file"]["file_data"].split(",", 1)[1])
== tiny_pdf.read_bytes() # noqa: ASYNC240 — test fixture, sync read is fine
base64.b64decode(file_part["file"]["file_data"].split(",", 1)[1]) == tiny_pdf.read_bytes() # noqa: ASYNC240 — test fixture, sync read is fine
)
assert user["content"][1] == {"type": "text", "text": "What is the diagnosis?"}
assert captured["headers"]["authorization"] == "Bearer sk-or-test"
@ -85,22 +91,22 @@ async def test_chat_array_content_concatenates(respx_mock, tiny_pdf: Path):
return_value=httpx.Response(
200,
json={
"choices": [{
"message": {
"content": [
{"type": "text", "text": "Hello "},
{"type": "text", "text": "world"},
{"type": "image_url", "image_url": "ignored"},
]
"choices": [
{
"message": {
"content": [
{"type": "text", "text": "Hello "},
{"type": "text", "text": "world"},
{"type": "image_url", "image_url": "ignored"},
]
}
}
}],
],
"usage": {"prompt_tokens": 1, "completion_tokens": 1},
},
)
)
provider = OpenRouterPdfProvider(
api_key="sk-or-test", base_url=_BASE, model="x/y"
)
provider = OpenRouterPdfProvider(api_key="sk-or-test", base_url=_BASE, model="x/y")
response = await provider.complete(prompt="hi", pdf_path=tiny_pdf)
assert response.text == "Hello world"

View file

@ -65,13 +65,15 @@ class TestParser:
interaction_id="abc",
query="Who directed Inception?",
answer="Christopher Nolan",
pages=[{
"page_name": "Inception (film)",
"page_url": "https://en.wikipedia.org/wiki/Inception",
"page_snippet": "snippet",
"page_result": "<html>full html</html>",
"page_last_modified": "2024-01-01",
}],
pages=[
{
"page_name": "Inception (film)",
"page_url": "https://en.wikipedia.org/wiki/Inception",
"page_snippet": "snippet",
"page_result": "<html>full html</html>",
"page_last_modified": "2024-01-01",
}
],
),
]
path = _make_jsonl_bz2(rows, tmp_path)
@ -120,8 +122,7 @@ class TestParser:
def test_alt_answers_parsed(self, tmp_path: Path) -> None:
rows = [
_row(interaction_id="z", query="q?", answer="42",
alt_ans=["forty-two", "42.0"]),
_row(interaction_id="z", query="q?", answer="42", alt_ans=["forty-two", "42.0"]),
]
path = _make_jsonl_bz2(rows, tmp_path)
parsed = iter_questions(path)
@ -143,22 +144,32 @@ class TestParser:
class TestPageHash:
def test_url_hash_stable(self) -> None:
a = CragPage(
page_name="A", page_url="https://x.test/p?q=1",
page_snippet="", page_html="<html/>",
page_name="A",
page_url="https://x.test/p?q=1",
page_snippet="",
page_html="<html/>",
)
b = CragPage(
page_name="B", page_url="https://x.test/p?q=1",
page_snippet="", page_html="<html/>",
page_name="B",
page_url="https://x.test/p?q=1",
page_snippet="",
page_html="<html/>",
)
assert a.url_hash == b.url_hash
assert len(a.url_hash) == 12
def test_url_hash_unique(self) -> None:
a = CragPage(
page_name="A", page_url="https://x.test/a", page_snippet="", page_html="<html/>",
page_name="A",
page_url="https://x.test/a",
page_snippet="",
page_html="<html/>",
)
b = CragPage(
page_name="B", page_url="https://x.test/b", page_snippet="", page_html="<html/>",
page_name="B",
page_url="https://x.test/b",
page_snippet="",
page_html="<html/>",
)
assert a.url_hash != b.url_hash
@ -174,21 +185,23 @@ class TestStratifiedSample:
(5, "sports", "multi-hop"),
):
for _ in range(n):
out.append(CragQuestion(
qid=f"C{idx:05d}",
interaction_id=f"i{idx}",
query_time="2024-01-01",
query=f"q{idx}?",
gold_answer="a",
alt_answers=[],
domain=domain,
question_type=qtype,
static_or_dynamic="static",
popularity="head",
split=0,
raw_index=idx,
pages=[],
))
out.append(
CragQuestion(
qid=f"C{idx:05d}",
interaction_id=f"i{idx}",
query_time="2024-01-01",
query=f"q{idx}?",
gold_answer="a",
alt_answers=[],
domain=domain,
question_type=qtype,
static_or_dynamic="static",
popularity="head",
split=0,
raw_index=idx,
pages=[],
)
)
idx += 1
return out

View file

@ -152,7 +152,9 @@ class TestGradeDeterministicHappyPath:
class TestGradeDeterministicRefusal:
def test_idk_maps_to_missing(self) -> None:
result = grade_deterministic(
pred="I don't know.", gold="Tim Cook", question_type="simple",
pred="I don't know.",
gold="Tim Cook",
question_type="simple",
)
assert result.grade == "missing"
assert result.score == 0
@ -225,8 +227,11 @@ class TestGradeDeterministicLexicalMiss:
class TestGradeResultShape:
def test_to_dict_round_trip(self) -> None:
result = CragGradeResult(
grade="correct", score=1, method="exact",
normalised_pred="x", normalised_gold="x",
grade="correct",
score=1,
method="exact",
normalised_pred="x",
normalised_gold="x",
)
d = result.to_dict()
assert d["grade"] == "correct"

View file

@ -112,7 +112,9 @@ class TestFallbackStripper:
</body></html>
"""
result = extract_main_content(
html, url="https://x.test/", page_name="Title",
html,
url="https://x.test/",
page_name="Title",
)
assert result.ok
assert "content one" in result.text

View file

@ -63,14 +63,22 @@ class TestCacheFilename:
@pytest.mark.asyncio
@respx.mock
async def test_fetch_success_writes_markdown(tmp_path: Path) -> None:
respx.get(WIKI_API).mock(return_value=httpx.Response(
200,
json={"query": {"pages": [{
"pageid": 1,
"title": "James Buchanan",
"extract": "James Buchanan was the 15th president of the United States.",
}]}},
))
respx.get(WIKI_API).mock(
return_value=httpx.Response(
200,
json={
"query": {
"pages": [
{
"pageid": 1,
"title": "James Buchanan",
"extract": "James Buchanan was the 15th president of the United States.",
}
]
}
},
)
)
fetcher = WikiFetcher(cache_dir=tmp_path, rate_limit_rps=100) # disable throttle
article = await fetcher.fetch("https://en.wikipedia.org/wiki/James_Buchanan")
assert article is not None
@ -83,13 +91,21 @@ async def test_fetch_success_writes_markdown(tmp_path: Path) -> None:
@pytest.mark.asyncio
@respx.mock
async def test_fetch_missing_page_returns_none(tmp_path: Path) -> None:
respx.get(WIKI_API).mock(return_value=httpx.Response(
200,
json={"query": {"pages": [{
"title": "DoesNotExist",
"missing": True,
}]}},
))
respx.get(WIKI_API).mock(
return_value=httpx.Response(
200,
json={
"query": {
"pages": [
{
"title": "DoesNotExist",
"missing": True,
}
]
}
},
)
)
fetcher = WikiFetcher(cache_dir=tmp_path, rate_limit_rps=100)
article = await fetcher.fetch("https://en.wikipedia.org/wiki/DoesNotExist")
assert article is None

View file

@ -99,7 +99,9 @@ class TestListFormat:
assert 0.0 < r.f1 < 1.0
def test_extra_items_lower_precision(self) -> None:
r = grade(pred="apple, banana, cherry, date", gold="apple, banana, cherry", answer_format="List")
r = grade(
pred="apple, banana, cherry, date", gold="apple, banana, cherry", answer_format="List"
)
assert 0.0 < r.f1 < 1.0
# Recall=1, precision=3/4 → F1 ~= 0.857
assert r.f1 == pytest.approx(2 * (3 / 4) * 1 / (3 / 4 + 1), rel=1e-3)

View file

@ -14,9 +14,7 @@ from ...core.workspace_context import WorkspaceContext
from . import document_tools, search_tools
def register(
mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext
) -> None:
def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None:
"""Register every knowledge-base tool on the server."""
search_tools.register(mcp, client, context)
document_tools.register(mcp, client, context)

View file

@ -20,9 +20,7 @@ from .annotations import DELETE, WRITE, DocumentId
from .note_ingestion import build_note_document
def register(
mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext
) -> None:
def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None:
"""Register the knowledge-base write and delete tools."""
@mcp.tool(
@ -136,8 +134,7 @@ def register(
str,
Field(
min_length=1,
description="New full text; replaces the existing content "
"entirely.",
description="New full text; replaces the existing content entirely.",
),
],
) -> str:

View file

@ -17,9 +17,7 @@ from ...core.workspace_context import WorkspaceContext, WorkspaceParam
from .annotations import READ, DocumentId, DocumentTypes
def register(
mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext
) -> None:
def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None:
"""Register the knowledge-base read tools."""
@mcp.tool(
@ -81,12 +79,8 @@ def register(
int | None,
Field(description="Only documents in this folder. Omit for all."),
] = None,
page: Annotated[
int, Field(ge=0, description="Zero-based page number.")
] = 0,
page_size: Annotated[
int, Field(ge=1, description="Documents per page.")
] = 20,
page: Annotated[int, Field(ge=0, description="Zero-based page number.")] = 0,
page_size: Annotated[int, Field(ge=1, description="Documents per page.")] = 20,
workspace: WorkspaceParam = None,
response_format: ResponseFormatParam = "markdown",
) -> str:

View file

@ -35,9 +35,7 @@ _REGISTRARS = (
)
def register(
mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext
) -> None:
def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None:
"""Register every scraper and run-history tool on the server."""
for module in _REGISTRARS:
module.register(mcp, client, context)

Some files were not shown because too many files have changed in this diff Show more