diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/tiktok/system_prompt.md b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/tiktok/system_prompt.md index 4300e9313..cd0b05e0e 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/tiktok/system_prompt.md +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/tiktok/system_prompt.md @@ -14,11 +14,11 @@ Answer the delegated question from live TikTok data gathered with your verb, com -- 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. diff --git a/surfsense_backend/app/capabilities/instagram/details/schemas.py b/surfsense_backend/app/capabilities/instagram/details/schemas.py index 1372f881f..8ec21335b 100644 --- a/surfsense_backend/app/capabilities/instagram/details/schemas.py +++ b/surfsense_backend/app/capabilities/instagram/details/schemas.py @@ -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)." diff --git a/surfsense_backend/app/capabilities/instagram/scrape/schemas.py b/surfsense_backend/app/capabilities/instagram/scrape/schemas.py index f30078d9b..bfb95cfac 100644 --- a/surfsense_backend/app/capabilities/instagram/scrape/schemas.py +++ b/surfsense_backend/app/capabilities/instagram/scrape/schemas.py @@ -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)." diff --git a/surfsense_backend/app/capabilities/tiktok/scrape/definition.py b/surfsense_backend/app/capabilities/tiktok/scrape/definition.py index e2706a2ef..72bc16ace 100644 --- a/surfsense_backend/app/capabilities/tiktok/scrape/definition.py +++ b/surfsense_backend/app/capabilities/tiktok/scrape/definition.py @@ -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, diff --git a/surfsense_backend/app/capabilities/tiktok/scrape/executor.py b/surfsense_backend/app/capabilities/tiktok/scrape/executor.py index 7ef71486d..8ce1d5be5 100644 --- a/surfsense_backend/app/capabilities/tiktok/scrape/executor.py +++ b/surfsense_backend/app/capabilities/tiktok/scrape/executor.py @@ -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( diff --git a/surfsense_backend/app/capabilities/tiktok/scrape/schemas.py b/surfsense_backend/app/capabilities/tiktok/scrape/schemas.py index 853b70bd9..ac3792712 100644 --- a/surfsense_backend/app/capabilities/tiktok/scrape/schemas.py +++ b/surfsense_backend/app/capabilities/tiktok/scrape/schemas.py @@ -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 (/@), a hashtag " - "(/tag/), or a search URL. Provide these OR profiles/hashtags/" - "search_queries (at least one source is required)." + "(/tag/), 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 diff --git a/surfsense_backend/app/proprietary/platforms/instagram/fetch.py b/surfsense_backend/app/proprietary/platforms/instagram/fetch.py index 41f692bf1..ab94f81a0 100644 --- a/surfsense_backend/app/proprietary/platforms/instagram/fetch.py +++ b/surfsense_backend/app/proprietary/platforms/instagram/fetch.py @@ -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: diff --git a/surfsense_backend/app/proprietary/platforms/instagram/parsers.py b/surfsense_backend/app/proprietary/platforms/instagram/parsers.py index 5c4e8ea03..c76d32e5b 100644 --- a/surfsense_backend/app/proprietary/platforms/instagram/parsers.py +++ b/surfsense_backend/app/proprietary/platforms/instagram/parsers.py @@ -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'', re.DOTALL ) -_OG_RE = re.compile( - r' 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"), diff --git a/surfsense_backend/app/proprietary/platforms/instagram/scraper.py b/surfsense_backend/app/proprietary/platforms/instagram/scraper.py index 9f4f6e2eb..f248b5653 100644 --- a/surfsense_backend/app/proprietary/platforms/instagram/scraper.py +++ b/surfsense_backend/app/proprietary/platforms/instagram/scraper.py @@ -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() diff --git a/surfsense_backend/app/proprietary/platforms/instagram/url_resolver.py b/surfsense_backend/app/proprietary/platforms/instagram/url_resolver.py index 08f5c6553..197a90276 100644 --- a/surfsense_backend/app/proprietary/platforms/instagram/url_resolver.py +++ b/surfsense_backend/app/proprietary/platforms/instagram/url_resolver.py @@ -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 diff --git a/surfsense_backend/app/proprietary/platforms/tiktok/extraction/timestamps.py b/surfsense_backend/app/proprietary/platforms/tiktok/extraction/timestamps.py index 4b155c137..ed45baa8f 100644 --- a/surfsense_backend/app/proprietary/platforms/tiktok/extraction/timestamps.py +++ b/surfsense_backend/app/proprietary/platforms/tiktok/extraction/timestamps.py @@ -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") diff --git a/surfsense_backend/app/proprietary/platforms/tiktok/flows/video.py b/surfsense_backend/app/proprietary/platforms/tiktok/flows/video.py index 48e558677..0cdefe1c4 100644 --- a/surfsense_backend/app/proprietary/platforms/tiktok/flows/video.py +++ b/surfsense_backend/app/proprietary/platforms/tiktok/flows/video.py @@ -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 diff --git a/surfsense_backend/app/proprietary/platforms/tiktok/orchestrator.py b/surfsense_backend/app/proprietary/platforms/tiktok/orchestrator.py index 01fae6caf..a738233a6 100644 --- a/surfsense_backend/app/proprietary/platforms/tiktok/orchestrator.py +++ b/surfsense_backend/app/proprietary/platforms/tiktok/orchestrator.py @@ -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/``) — 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: diff --git a/surfsense_backend/scripts/e2e_instagram_scraper.py b/surfsense_backend/scripts/e2e_instagram_scraper.py index c81075932..a5e356f99 100644 --- a/surfsense_backend/scripts/e2e_instagram_scraper.py +++ b/surfsense_backend/scripts/e2e_instagram_scraper.py @@ -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"): diff --git a/surfsense_backend/scripts/e2e_tiktok_scrape.py b/surfsense_backend/scripts/e2e_tiktok_scrape.py index b399ca3fb..4a0063303 100644 --- a/surfsense_backend/scripts/e2e_tiktok_scrape.py +++ b/surfsense_backend/scripts/e2e_tiktok_scrape.py @@ -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 diff --git a/surfsense_backend/tests/unit/capabilities/instagram/test_schemas.py b/surfsense_backend/tests/unit/capabilities/instagram/test_schemas.py index afabbf833..fde5b07ab 100644 --- a/surfsense_backend/tests/unit/capabilities/instagram/test_schemas.py +++ b/surfsense_backend/tests/unit/capabilities/instagram/test_schemas.py @@ -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"]) diff --git a/surfsense_backend/tests/unit/capabilities/tiktok/scrape/test_executor.py b/surfsense_backend/tests/unit/capabilities/tiktok/scrape/test_executor.py index 1bf07b0c3..367831c76 100644 --- a/surfsense_backend/tests/unit/capabilities/tiktok/scrape/test_executor.py +++ b/surfsense_backend/tests/unit/capabilities/tiktok/scrape/test_executor.py @@ -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 diff --git a/surfsense_backend/tests/unit/platforms/instagram/test_budget.py b/surfsense_backend/tests/unit/platforms/instagram/test_budget.py index ea500be40..703002f0c 100644 --- a/surfsense_backend/tests/unit/platforms/instagram/test_budget.py +++ b/surfsense_backend/tests/unit/platforms/instagram/test_budget.py @@ -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) ], }, } diff --git a/surfsense_backend/tests/unit/platforms/instagram/test_discovery.py b/surfsense_backend/tests/unit/platforms/instagram/test_discovery.py index c95a80315..379e3ba28 100644 --- a/surfsense_backend/tests/unit/platforms/instagram/test_discovery.py +++ b/surfsense_backend/tests/unit/platforms/instagram/test_discovery.py @@ -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 diff --git a/surfsense_backend/tests/unit/platforms/instagram/test_fetch_resilience.py b/surfsense_backend/tests/unit/platforms/instagram/test_fetch_resilience.py index 859c70443..44fefa430 100644 --- a/surfsense_backend/tests/unit/platforms/instagram/test_fetch_resilience.py +++ b/surfsense_backend/tests/unit/platforms/instagram/test_fetch_resilience.py @@ -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 diff --git a/surfsense_backend/tests/unit/platforms/instagram/test_parsers.py b/surfsense_backend/tests/unit/platforms/instagram/test_parsers.py index 8cba74b3a..1cfa386a7 100644 --- a/surfsense_backend/tests/unit/platforms/instagram/test_parsers.py +++ b/surfsense_backend/tests/unit/platforms/instagram/test_parsers.py @@ -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 = ( diff --git a/surfsense_backend/tests/unit/platforms/tiktok/test_discovery.py b/surfsense_backend/tests/unit/platforms/tiktok/test_discovery.py deleted file mode 100644 index 4664f85e0..000000000 --- a/surfsense_backend/tests/unit/platforms/tiktok/test_discovery.py +++ /dev/null @@ -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/``) -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/`` 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 ( - '' - ) - - -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"] diff --git a/surfsense_backend/tests/unit/platforms/tiktok/test_fetch_resilience.py b/surfsense_backend/tests/unit/platforms/tiktok/test_fetch_resilience.py index 484157964..a2177233e 100644 --- a/surfsense_backend/tests/unit/platforms/tiktok/test_fetch_resilience.py +++ b/surfsense_backend/tests/unit/platforms/tiktok/test_fetch_resilience.py @@ -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 diff --git a/surfsense_backend/tests/unit/platforms/tiktok/test_target_resolver.py b/surfsense_backend/tests/unit/platforms/tiktok/test_target_resolver.py index c118f4f86..d246085ae 100644 --- a/surfsense_backend/tests/unit/platforms/tiktok/test_target_resolver.py +++ b/surfsense_backend/tests/unit/platforms/tiktok/test_target_resolver.py @@ -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" diff --git a/surfsense_backend/tests/unit/platforms/tiktok/test_user_search.py b/surfsense_backend/tests/unit/platforms/tiktok/test_user_search.py index 845b5e22d..c4b89e390 100644 --- a/surfsense_backend/tests/unit/platforms/tiktok/test_user_search.py +++ b/surfsense_backend/tests/unit/platforms/tiktok/test_user_search.py @@ -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" diff --git a/surfsense_backend/tests/unit/services/test_requesty_model_normalizer.py b/surfsense_backend/tests/unit/services/test_requesty_model_normalizer.py index 36b867923..2adae3ada 100644 --- a/surfsense_backend/tests/unit/services/test_requesty_model_normalizer.py +++ b/surfsense_backend/tests/unit/services/test_requesty_model_normalizer.py @@ -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), ] ) diff --git a/surfsense_evals/scripts/analyze_failure_timing.py b/surfsense_evals/scripts/analyze_failure_timing.py index f4f8aedba..14b76852f 100644 --- a/surfsense_evals/scripts/analyze_failure_timing.py +++ b/surfsense_evals/scripts/analyze_failure_timing.py @@ -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}) diff --git a/surfsense_evals/scripts/analyze_failures.py b/surfsense_evals/scripts/analyze_failures.py index ff5ec23f4..f60038c00 100644 --- a/surfsense_evals/scripts/analyze_failures.py +++ b/surfsense_evals/scripts/analyze_failures.py @@ -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() }, } diff --git a/surfsense_evals/scripts/check_extraction_sizes.py b/surfsense_evals/scripts/check_extraction_sizes.py index 712e693cb..1755e9b6c 100644 --- a/surfsense_evals/scripts/check_extraction_sizes.py +++ b/surfsense_evals/scripts/check_extraction_sizes.py @@ -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) diff --git a/surfsense_evals/scripts/compute_adjusted_accuracy.py b/surfsense_evals/scripts/compute_adjusted_accuracy.py index 13693c055..0cd4b3073 100644 --- a/surfsense_evals/scripts/compute_adjusted_accuracy.py +++ b/surfsense_evals/scripts/compute_adjusted_accuracy.py @@ -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") diff --git a/surfsense_evals/scripts/compute_blog_extras.py b/surfsense_evals/scripts/compute_blog_extras.py index abe88d08b..29922f54c 100644 --- a/surfsense_evals/scripts/compute_blog_extras.py +++ b/surfsense_evals/scripts/compute_blog_extras.py @@ -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}" + ) # --------------------------------------------------------------------------- diff --git a/surfsense_evals/scripts/compute_post_retry_accuracy.py b/surfsense_evals/scripts/compute_post_retry_accuracy.py index f3a716219..29007ed41 100644 --- a/surfsense_evals/scripts/compute_post_retry_accuracy.py +++ b/surfsense_evals/scripts/compute_post_retry_accuracy.py @@ -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 diff --git a/surfsense_evals/scripts/inspect_first30.py b/surfsense_evals/scripts/inspect_first30.py index e06c6c029..b3caedca6 100644 --- a/surfsense_evals/scripts/inspect_first30.py +++ b/surfsense_evals/scripts/inspect_first30.py @@ -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:") diff --git a/surfsense_evals/scripts/peek_crag_run.py b/surfsense_evals/scripts/peek_crag_run.py index 3a79d76c7..0720e24cf 100644 --- a/surfsense_evals/scripts/peek_crag_run.py +++ b/surfsense_evals/scripts/peek_crag_run.py @@ -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}") diff --git a/surfsense_evals/scripts/retry_failed_questions.py b/surfsense_evals/scripts/retry_failed_questions.py index 25facfe60..d65f9f0fb 100644 --- a/surfsense_evals/scripts/retry_failed_questions.py +++ b/surfsense_evals/scripts/retry_failed_questions.py @@ -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))) diff --git a/surfsense_evals/scripts/summarise_crag_run.py b/surfsense_evals/scripts/summarise_crag_run.py index 9722cd391..d15c4996f 100644 --- a/surfsense_evals/scripts/summarise_crag_run.py +++ b/surfsense_evals/scripts/summarise_crag_run.py @@ -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)) diff --git a/surfsense_evals/scripts/summarise_parser_compare_run.py b/surfsense_evals/scripts/summarise_parser_compare_run.py index c091043aa..7801a1318 100644 --- a/surfsense_evals/scripts/summarise_parser_compare_run.py +++ b/surfsense_evals/scripts/summarise_parser_compare_run.py @@ -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}" diff --git a/surfsense_evals/scripts/test_context_overflow_hypothesis.py b/surfsense_evals/scripts/test_context_overflow_hypothesis.py index 89bd6cb3d..8ccccba45 100644 --- a/surfsense_evals/scripts/test_context_overflow_hypothesis.py +++ b/surfsense_evals/scripts/test_context_overflow_hypothesis.py @@ -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__": diff --git a/surfsense_evals/src/surfsense_evals/core/arms/surfsense.py b/surfsense_evals/src/surfsense_evals/core/arms/surfsense.py index a84350dfd..f63b5cbe6 100644 --- a/surfsense_evals/src/surfsense_evals/core/arms/surfsense.py +++ b/surfsense_evals/src/surfsense_evals/core/arms/surfsense.py @@ -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( diff --git a/surfsense_evals/src/surfsense_evals/core/auth.py b/surfsense_evals/src/surfsense_evals/core/auth.py index a87e757c2..cf348e3ff 100644 --- a/surfsense_evals/src/surfsense_evals/core/auth.py +++ b/surfsense_evals/src/surfsense_evals/core/auth.py @@ -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, diff --git a/surfsense_evals/src/surfsense_evals/core/cli.py b/surfsense_evals/src/surfsense_evals/core/cli.py index 2dcbe7327..21e706c49 100644 --- a/surfsense_evals/src/surfsense_evals/core/cli.py +++ b/surfsense_evals/src/surfsense_evals/core/cli.py @@ -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) diff --git a/surfsense_evals/src/surfsense_evals/core/clients/documents.py b/surfsense_evals/src/surfsense_evals/core/clients/documents.py index b96bef6d8..2fd9b2766 100644 --- a/surfsense_evals/src/surfsense_evals/core/clients/documents.py +++ b/surfsense_evals/src/surfsense_evals/core/clients/documents.py @@ -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`` diff --git a/surfsense_evals/src/surfsense_evals/core/clients/new_chat.py b/surfsense_evals/src/surfsense_evals/core/clients/new_chat.py index a4c23d010..397193bba 100644 --- a/surfsense_evals/src/surfsense_evals/core/clients/new_chat.py +++ b/surfsense_evals/src/surfsense_evals/core/clients/new_chat.py @@ -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, diff --git a/surfsense_evals/src/surfsense_evals/core/clients/search_space.py b/surfsense_evals/src/surfsense_evals/core/clients/search_space.py index efd4a571d..19486aca1 100644 --- a/surfsense_evals/src/surfsense_evals/core/clients/search_space.py +++ b/surfsense_evals/src/surfsense_evals/core/clients/search_space.py @@ -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 diff --git a/surfsense_evals/src/surfsense_evals/core/config.py b/surfsense_evals/src/surfsense_evals/core/config.py index 9a5a71e89..80002157e 100644 --- a/surfsense_evals/src/surfsense_evals/core/config.py +++ b/surfsense_evals/src/surfsense_evals/core/config.py @@ -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 ), ) diff --git a/surfsense_evals/src/surfsense_evals/core/ingest_settings.py b/surfsense_evals/src/surfsense_evals/core/ingest_settings.py index 8328e0d46..216ae36a4 100644 --- a/surfsense_evals/src/surfsense_evals/core/ingest_settings.py +++ b/surfsense_evals/src/surfsense_evals/core/ingest_settings.py @@ -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__ = [ diff --git a/surfsense_evals/src/surfsense_evals/core/metrics/comparison.py b/surfsense_evals/src/surfsense_evals/core/metrics/comparison.py index 9e40dc5ca..332535871 100644 --- a/surfsense_evals/src/surfsense_evals/core/metrics/comparison.py +++ b/surfsense_evals/src/surfsense_evals/core/metrics/comparison.py @@ -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: diff --git a/surfsense_evals/src/surfsense_evals/core/metrics/mc_accuracy.py b/surfsense_evals/src/surfsense_evals/core/metrics/mc_accuracy.py index 8b0188ca4..958f62600 100644 --- a/surfsense_evals/src/surfsense_evals/core/metrics/mc_accuracy.py +++ b/surfsense_evals/src/surfsense_evals/core/metrics/mc_accuracy.py @@ -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: diff --git a/surfsense_evals/src/surfsense_evals/core/metrics/retrieval.py b/surfsense_evals/src/surfsense_evals/core/metrics/retrieval.py index d4cfe10ae..3fd25f634 100644 --- a/surfsense_evals/src/surfsense_evals/core/metrics/retrieval.py +++ b/surfsense_evals/src/surfsense_evals/core/metrics/retrieval.py @@ -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 diff --git a/surfsense_evals/src/surfsense_evals/core/parse/citations.py b/surfsense_evals/src/surfsense_evals/core/parse/citations.py index c57ffeb0b..38bd16d31 100644 --- a/surfsense_evals/src/surfsense_evals/core/parse/citations.py +++ b/surfsense_evals/src/surfsense_evals/core/parse/citations.py @@ -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+)*" diff --git a/surfsense_evals/src/surfsense_evals/core/parse/freeform_answer.py b/surfsense_evals/src/surfsense_evals/core/parse/freeform_answer.py index 959b045a5..104176d09 100644 --- a/surfsense_evals/src/surfsense_evals/core/parse/freeform_answer.py +++ b/surfsense_evals/src/surfsense_evals/core/parse/freeform_answer.py @@ -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 diff --git a/surfsense_evals/src/surfsense_evals/core/parsers/azure_di.py b/surfsense_evals/src/surfsense_evals/core/parsers/azure_di.py index 7d796ee99..16e618db4 100644 --- a/surfsense_evals/src/surfsense_evals/core/parsers/azure_di.py +++ b/surfsense_evals/src/surfsense_evals/core/parsers/azure_di.py @@ -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) diff --git a/surfsense_evals/src/surfsense_evals/core/parsers/llamacloud.py b/surfsense_evals/src/surfsense_evals/core/parsers/llamacloud.py index 300b4ec87..32fd97e47 100644 --- a/surfsense_evals/src/surfsense_evals/core/parsers/llamacloud.py +++ b/surfsense_evals/src/surfsense_evals/core/parsers/llamacloud.py @@ -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) diff --git a/surfsense_evals/src/surfsense_evals/core/pdf/render.py b/surfsense_evals/src/surfsense_evals/core/pdf/render.py index 624136d7c..21866f3e5 100644 --- a/surfsense_evals/src/surfsense_evals/core/pdf/render.py +++ b/surfsense_evals/src/surfsense_evals/core/pdf/render.py @@ -116,11 +116,7 @@ def _normalise_paragraphs(text: str) -> list[str]: def _escape_html(text: str) -> str: - return ( - text.replace("&", "&") - .replace("<", "<") - .replace(">", ">") - ) + return text.replace("&", "&").replace("<", "<").replace(">", ">") def render_pdf( diff --git a/surfsense_evals/src/surfsense_evals/core/providers/openrouter_pdf.py b/surfsense_evals/src/surfsense_evals/core/providers/openrouter_pdf.py index 985d88a68..5cd47b04e 100644 --- a/surfsense_evals/src/surfsense_evals/core/providers/openrouter_pdf.py +++ b/surfsense_evals/src/surfsense_evals/core/providers/openrouter_pdf.py @@ -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 diff --git a/surfsense_evals/src/surfsense_evals/core/registry.py b/surfsense_evals/src/surfsense_evals/core/registry.py index 65f64c39a..7fb64c36f 100644 --- a/surfsense_evals/src/surfsense_evals/core/registry.py +++ b/surfsense_evals/src/surfsense_evals/core/registry.py @@ -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 ``.""" - 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 "" - 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]: diff --git a/surfsense_evals/src/surfsense_evals/core/scenarios.py b/surfsense_evals/src/surfsense_evals/core/scenarios.py index 16874a069..fefdc6865 100644 --- a/surfsense_evals/src/surfsense_evals/core/scenarios.py +++ b/surfsense_evals/src/surfsense_evals/core/scenarios.py @@ -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}`." diff --git a/surfsense_evals/src/surfsense_evals/suites/__init__.py b/surfsense_evals/src/surfsense_evals/suites/__init__.py index a0a01223c..f3d26f865 100644 --- a/surfsense_evals/src/surfsense_evals/suites/__init__.py +++ b/surfsense_evals/src/surfsense_evals/suites/__init__.py @@ -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 diff --git a/surfsense_evals/src/surfsense_evals/suites/medical/cure/ingest.py b/surfsense_evals/src/surfsense_evals/suites/medical/cure/ingest.py index 0c32a38a1..84108b4df 100644 --- a/surfsense_evals/src/surfsense_evals/suites/medical/cure/ingest.py +++ b/surfsense_evals/src/surfsense_evals/suites/medical/cure/ingest.py @@ -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( diff --git a/surfsense_evals/src/surfsense_evals/suites/medical/cure/runner.py b/surfsense_evals/src/surfsense_evals/suites/medical/cure/runner.py index 4a85b6ba5..d2735c8d5 100644 --- a/surfsense_evals/src/surfsense_evals/suites/medical/cure/runner.py +++ b/surfsense_evals/src/surfsense_evals/suites/medical/cure/runner.py @@ -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: diff --git a/surfsense_evals/src/surfsense_evals/suites/medical/medxpertqa/ingest.py b/surfsense_evals/src/surfsense_evals/suites/medical/medxpertqa/ingest.py index ff43c7049..f50247acb 100644 --- a/surfsense_evals/src/surfsense_evals/suites/medical/medxpertqa/ingest.py +++ b/surfsense_evals/src/surfsense_evals/suites/medical/medxpertqa/ingest.py @@ -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 diff --git a/surfsense_evals/src/surfsense_evals/suites/medical/medxpertqa/runner.py b/surfsense_evals/src/surfsense_evals/suites/medical/medxpertqa/runner.py index ac0651996..f7a3331a9 100644 --- a/surfsense_evals/src/surfsense_evals/suites/medical/medxpertqa/runner.py +++ b/surfsense_evals/src/surfsense_evals/suites/medical/medxpertqa/runner.py @@ -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 diff --git a/surfsense_evals/src/surfsense_evals/suites/medical/mirage/ingest.py b/surfsense_evals/src/surfsense_evals/suites/medical/mirage/ingest.py index 8e891aabf..c4aa53fe2 100644 --- a/surfsense_evals/src/surfsense_evals/suites/medical/mirage/ingest.py +++ b/surfsense_evals/src/surfsense_evals/suites/medical/mirage/ingest.py @@ -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) diff --git a/surfsense_evals/src/surfsense_evals/suites/medical/mirage/runner.py b/surfsense_evals/src/surfsense_evals/suites/medical/mirage/runner.py index 8a6b04ae0..76e719f1f 100644 --- a/surfsense_evals/src/surfsense_evals/suites/medical/mirage/runner.py +++ b/surfsense_evals/src/surfsense_evals/suites/medical/mirage/runner.py @@ -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} diff --git a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/grader.py b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/grader.py index 7edad73eb..ecb5144e8 100644 --- a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/grader.py +++ b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/grader.py @@ -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), ) diff --git a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/ingest.py b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/ingest.py index 15cdbeb77..3c736756a 100644 --- a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/ingest.py +++ b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/ingest.py @@ -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 diff --git a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/prompt.py b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/prompt.py index 27d6a0d00..70229dc15 100644 --- a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/prompt.py +++ b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/prompt.py @@ -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: `." ), - "int": ( - "Respond with a single integer only. " - "Format your final line as `Answer: `." - ), + "int": ("Respond with a single integer only. Format your final line as `Answer: `."), "float": ( "Respond with a single decimal number only (no units). " "Format your final line as `Answer: `." diff --git a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/runner.py b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/runner.py index b7685766e..782ba5d9a 100644 --- a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/runner.py +++ b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/mmlongbench/runner.py @@ -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(), diff --git a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/parser_compare/ingest.py b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/parser_compare/ingest.py index 1360ec89e..a7fce60d1 100644 --- a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/parser_compare/ingest.py +++ b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/parser_compare/ingest.py @@ -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, ) diff --git a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/parser_compare/prompt.py b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/parser_compare/prompt.py index 7119bbd29..ccde69e71 100644 --- a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/parser_compare/prompt.py +++ b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/parser_compare/prompt.py @@ -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: `." ), - "int": ( - "Respond with a single integer only. " - "Format your final line as `Answer: `." - ), + "int": ("Respond with a single integer only. Format your final line as `Answer: `."), "float": ( "Respond with a single decimal number only (no units). " "Format your final line as `Answer: `." @@ -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"\n" + f'\n' f"{document_markdown.strip()}\n" f"\n\n" f"Question: {question.strip()}\n\n" diff --git a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/parser_compare/runner.py b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/parser_compare/runner.py index 297bd2fa0..6c009995f 100644 --- a/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/parser_compare/runner.py +++ b/surfsense_evals/src/surfsense_evals/suites/multimodal_doc/parser_compare/runner.py @@ -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, diff --git a/surfsense_evals/src/surfsense_evals/suites/research/crag/dataset.py b/surfsense_evals/src/surfsense_evals/suites/research/crag/dataset.py index 224dcae5c..7154e6d14 100644 --- a/surfsense_evals/src/surfsense_evals/suites/research/crag/dataset.py +++ b/surfsense_evals/src/surfsense_evals/suites/research/crag/dataset.py @@ -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 diff --git a/surfsense_evals/src/surfsense_evals/suites/research/crag/grader.py b/surfsense_evals/src/surfsense_evals/suites/research/crag/grader.py index 63f66702b..e49660a6f 100644 --- a/surfsense_evals/src/surfsense_evals/suites/research/crag/grader.py +++ b/surfsense_evals/src/surfsense_evals/suites/research/crag/grader.py @@ -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\": \"\"}' + '{"grade": "correct"|"missing"|"incorrect", "rationale": ""}' ) @@ -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( diff --git a/surfsense_evals/src/surfsense_evals/suites/research/crag/html_extract.py b/surfsense_evals/src/surfsense_evals/suites/research/crag/html_extract.py index dd618d7e3..271d43d56 100644 --- a/surfsense_evals/src/surfsense_evals/suites/research/crag/html_extract.py +++ b/surfsense_evals/src/surfsense_evals/suites/research/crag/html_extract.py @@ -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) diff --git a/surfsense_evals/src/surfsense_evals/suites/research/crag/ingest.py b/surfsense_evals/src/surfsense_evals/suites/research/crag/ingest.py index 4e0c2bdc5..1b66f45f9 100644 --- a/surfsense_evals/src/surfsense_evals/suites/research/crag/ingest.py +++ b/surfsense_evals/src/surfsense_evals/suites/research/crag/ingest.py @@ -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()) diff --git a/surfsense_evals/src/surfsense_evals/suites/research/crag/prompt.py b/surfsense_evals/src/surfsense_evals/suites/research/crag/prompt.py index 0d4327774..5b29fb90b 100644 --- a/surfsense_evals/src/surfsense_evals/suites/research/crag/prompt.py +++ b/surfsense_evals/src/surfsense_evals/suites/research/crag/prompt.py @@ -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, diff --git a/surfsense_evals/src/surfsense_evals/suites/research/crag/runner.py b/surfsense_evals/src/surfsense_evals/suites/research/crag/runner.py index efb7b4474..801e00220 100644 --- a/surfsense_evals/src/surfsense_evals/suites/research/crag/runner.py +++ b/surfsense_evals/src/surfsense_evals/suites/research/crag/runner.py @@ -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 "?" diff --git a/surfsense_evals/src/surfsense_evals/suites/research/frames/dataset.py b/surfsense_evals/src/surfsense_evals/suites/research/frames/dataset.py index 629874102..80c7075f8 100644 --- a/surfsense_evals/src/surfsense_evals/suites/research/frames/dataset.py +++ b/surfsense_evals/src/surfsense_evals/suites/research/frames/dataset.py @@ -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 diff --git a/surfsense_evals/src/surfsense_evals/suites/research/frames/grader.py b/surfsense_evals/src/surfsense_evals/suites/research/frames/grader.py index d280e3eaf..32343de42 100644 --- a/surfsense_evals/src/surfsense_evals/suites/research/frames/grader.py +++ b/surfsense_evals/src/surfsense_evals/suites/research/frames/grader.py @@ -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\": \"\"}' + '{"correct": true|false, "rationale": ""}' ) @@ -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)) diff --git a/surfsense_evals/src/surfsense_evals/suites/research/frames/ingest.py b/surfsense_evals/src/surfsense_evals/suites/research/frames/ingest.py index 0288e192e..3ea7246a6 100644 --- a/surfsense_evals/src/surfsense_evals/suites/research/frames/ingest.py +++ b/surfsense_evals/src/surfsense_evals/suites/research/frames/ingest.py @@ -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). diff --git a/surfsense_evals/src/surfsense_evals/suites/research/frames/runner.py b/surfsense_evals/src/surfsense_evals/suites/research/frames/runner.py index 450c7ddd6..c703ab23e 100644 --- a/surfsense_evals/src/surfsense_evals/suites/research/frames/runner.py +++ b/surfsense_evals/src/surfsense_evals/suites/research/frames/runner.py @@ -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(), diff --git a/surfsense_evals/src/surfsense_evals/suites/research/frames/wiki_fetch.py b/surfsense_evals/src/surfsense_evals/suites/research/frames/wiki_fetch.py index 7f6b63e50..2bc96ad3a 100644 --- a/surfsense_evals/src/surfsense_evals/suites/research/frames/wiki_fetch.py +++ b/surfsense_evals/src/surfsense_evals/suites/research/frames/wiki_fetch.py @@ -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: diff --git a/surfsense_evals/tests/core/test_auth.py b/surfsense_evals/tests/core/test_auth.py index 181d8e632..ad291b295 100644 --- a/surfsense_evals/tests/core/test_auth.py +++ b/surfsense_evals/tests/core/test_auth.py @@ -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" diff --git a/surfsense_evals/tests/core/test_clients.py b/surfsense_evals/tests/core/test_clients.py index aa98f0ad4..44b78e3ba 100644 --- a/surfsense_evals/tests/core/test_clients.py +++ b/surfsense_evals/tests/core/test_clients.py @@ -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) diff --git a/surfsense_evals/tests/core/test_ingest_settings.py b/surfsense_evals/tests/core/test_ingest_settings.py index fd7e7818a..cd1d1f827 100644 --- a/surfsense_evals/tests/core/test_ingest_settings.py +++ b/surfsense_evals/tests/core/test_ingest_settings.py @@ -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 diff --git a/surfsense_evals/tests/core/test_metrics.py b/surfsense_evals/tests/core/test_metrics.py index cde1bb957..73c85c371 100644 --- a/surfsense_evals/tests/core/test_metrics.py +++ b/surfsense_evals/tests/core/test_metrics.py @@ -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) diff --git a/surfsense_evals/tests/core/test_parse_answer_letter.py b/surfsense_evals/tests/core/test_parse_answer_letter.py index 5adbf4bc3..0662adba3 100644 --- a/surfsense_evals/tests/core/test_parse_answer_letter.py +++ b/surfsense_evals/tests/core/test_parse_answer_letter.py @@ -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"), diff --git a/surfsense_evals/tests/core/test_parse_citations.py b/surfsense_evals/tests/core/test_parse_citations.py index eb444dab2..488af590d 100644 --- a/surfsense_evals/tests/core/test_parse_citations.py +++ b/surfsense_evals/tests/core/test_parse_citations.py @@ -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 diff --git a/surfsense_evals/tests/core/test_parse_freeform_answer.py b/surfsense_evals/tests/core/test_parse_freeform_answer.py index bdc7d74fc..a39aad2e5 100644 --- a/surfsense_evals/tests/core/test_parse_freeform_answer.py +++ b/surfsense_evals/tests/core/test_parse_freeform_answer.py @@ -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 diff --git a/surfsense_evals/tests/core/test_parse_sse.py b/surfsense_evals/tests/core/test_parse_sse.py index 362717288..10998a881 100644 --- a/surfsense_evals/tests/core/test_parse_sse.py +++ b/surfsense_evals/tests/core/test_parse_sse.py @@ -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"] diff --git a/surfsense_evals/tests/core/test_provider_openrouter.py b/surfsense_evals/tests/core/test_provider_openrouter.py index eb78aa053..aeed6eae5 100644 --- a/surfsense_evals/tests/core/test_provider_openrouter.py +++ b/surfsense_evals/tests/core/test_provider_openrouter.py @@ -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" diff --git a/surfsense_evals/tests/suites/test_crag_dataset.py b/surfsense_evals/tests/suites/test_crag_dataset.py index 6221467fd..c52e0f56b 100644 --- a/surfsense_evals/tests/suites/test_crag_dataset.py +++ b/surfsense_evals/tests/suites/test_crag_dataset.py @@ -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": "full 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": "full 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="", + page_name="A", + page_url="https://x.test/p?q=1", + page_snippet="", + page_html="", ) b = CragPage( - page_name="B", page_url="https://x.test/p?q=1", - page_snippet="", page_html="", + page_name="B", + page_url="https://x.test/p?q=1", + page_snippet="", + page_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="", + page_name="A", + page_url="https://x.test/a", + page_snippet="", + page_html="", ) b = CragPage( - page_name="B", page_url="https://x.test/b", page_snippet="", page_html="", + page_name="B", + page_url="https://x.test/b", + page_snippet="", + page_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 diff --git a/surfsense_evals/tests/suites/test_crag_grader.py b/surfsense_evals/tests/suites/test_crag_grader.py index 74960afa6..e0599c8b6 100644 --- a/surfsense_evals/tests/suites/test_crag_grader.py +++ b/surfsense_evals/tests/suites/test_crag_grader.py @@ -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" diff --git a/surfsense_evals/tests/suites/test_crag_html_extract.py b/surfsense_evals/tests/suites/test_crag_html_extract.py index 3bf757dbd..368692177 100644 --- a/surfsense_evals/tests/suites/test_crag_html_extract.py +++ b/surfsense_evals/tests/suites/test_crag_html_extract.py @@ -112,7 +112,9 @@ class TestFallbackStripper: """ 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 diff --git a/surfsense_evals/tests/suites/test_frames_wiki_fetch.py b/surfsense_evals/tests/suites/test_frames_wiki_fetch.py index 4941756f4..483c7aa58 100644 --- a/surfsense_evals/tests/suites/test_frames_wiki_fetch.py +++ b/surfsense_evals/tests/suites/test_frames_wiki_fetch.py @@ -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 diff --git a/surfsense_evals/tests/suites/test_mmlongbench_grader.py b/surfsense_evals/tests/suites/test_mmlongbench_grader.py index 92cd5f0cb..89005b3cd 100644 --- a/surfsense_evals/tests/suites/test_mmlongbench_grader.py +++ b/surfsense_evals/tests/suites/test_mmlongbench_grader.py @@ -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) diff --git a/surfsense_mcp/mcp_server/features/knowledge_base/__init__.py b/surfsense_mcp/mcp_server/features/knowledge_base/__init__.py index 1a971bfe4..323fb4e48 100644 --- a/surfsense_mcp/mcp_server/features/knowledge_base/__init__.py +++ b/surfsense_mcp/mcp_server/features/knowledge_base/__init__.py @@ -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) diff --git a/surfsense_mcp/mcp_server/features/knowledge_base/document_tools.py b/surfsense_mcp/mcp_server/features/knowledge_base/document_tools.py index 497a2526c..f2cc20f9e 100644 --- a/surfsense_mcp/mcp_server/features/knowledge_base/document_tools.py +++ b/surfsense_mcp/mcp_server/features/knowledge_base/document_tools.py @@ -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: diff --git a/surfsense_mcp/mcp_server/features/knowledge_base/search_tools.py b/surfsense_mcp/mcp_server/features/knowledge_base/search_tools.py index a9e60810d..c0f7f83a9 100644 --- a/surfsense_mcp/mcp_server/features/knowledge_base/search_tools.py +++ b/surfsense_mcp/mcp_server/features/knowledge_base/search_tools.py @@ -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: diff --git a/surfsense_mcp/mcp_server/features/scrapers/__init__.py b/surfsense_mcp/mcp_server/features/scrapers/__init__.py index eb5f72ba5..9aabbe0e5 100644 --- a/surfsense_mcp/mcp_server/features/scrapers/__init__.py +++ b/surfsense_mcp/mcp_server/features/scrapers/__init__.py @@ -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) diff --git a/surfsense_mcp/mcp_server/features/scrapers/capability.py b/surfsense_mcp/mcp_server/features/scrapers/capability.py index 7245c9f84..82f24c6e7 100644 --- a/surfsense_mcp/mcp_server/features/scrapers/capability.py +++ b/surfsense_mcp/mcp_server/features/scrapers/capability.py @@ -38,9 +38,7 @@ async def run_scraper( return _render_markdown(platform, verb, resolved.name, result) -def _render_markdown( - platform: str, verb: str, workspace_name: str, result: Any -) -> str: +def _render_markdown(platform: str, verb: str, workspace_name: str, result: Any) -> str: """A readable header plus the structured payload, clipped to a safe size.""" header = f'# {platform}.{verb} — {_describe_size(result)} from "{workspace_name}"' body = clip(to_json(result)) diff --git a/surfsense_mcp/mcp_server/features/scrapers/platforms/google_maps.py b/surfsense_mcp/mcp_server/features/scrapers/platforms/google_maps.py index e1613ca4e..daa27cef2 100644 --- a/surfsense_mcp/mcp_server/features/scrapers/platforms/google_maps.py +++ b/surfsense_mcp/mcp_server/features/scrapers/platforms/google_maps.py @@ -16,9 +16,7 @@ from ..capability import run_scraper ReviewSort = Literal["newest", "mostRelevant", "highestRanking", "lowestRanking"] -def register( - mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext -) -> None: +def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None: """Register the Google Maps place and review tools.""" @mcp.tool( @@ -45,10 +43,7 @@ def register( ] = None, location: Annotated[ str | None, - Field( - description="Geographic scope for a search, e.g. " - "'Seattle, USA'." - ), + Field(description="Geographic scope for a search, e.g. 'Seattle, USA'."), ] = None, max_places: Annotated[ int, Field(ge=1, description="Maximum places to return.") @@ -56,8 +51,7 @@ def register( include_details: Annotated[ bool, Field( - description="True adds opening hours and extra contact info " - "(slower)." + description="True adds opening hours and extra contact info (slower)." ), ] = False, workspace: WorkspaceParam = None, @@ -96,16 +90,11 @@ def register( async def google_maps_reviews( urls: Annotated[ list[str] | None, - Field( - description="Google Maps URLs of places. Provide urls OR " - "place_ids." - ), + Field(description="Google Maps URLs of places. Provide urls OR place_ids."), ] = None, place_ids: Annotated[ list[str] | None, - Field( - description="Google place ids from surfsense_google_maps_scrape." - ), + Field(description="Google place ids from surfsense_google_maps_scrape."), ] = None, max_reviews: Annotated[ int, Field(ge=1, description="Maximum reviews per place.") diff --git a/surfsense_mcp/mcp_server/features/scrapers/platforms/google_search.py b/surfsense_mcp/mcp_server/features/scrapers/platforms/google_search.py index cc1a1f8ed..acabef1d1 100644 --- a/surfsense_mcp/mcp_server/features/scrapers/platforms/google_search.py +++ b/surfsense_mcp/mcp_server/features/scrapers/platforms/google_search.py @@ -14,9 +14,7 @@ from ..annotations import SCRAPE from ..capability import run_scraper -def register( - mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext -) -> None: +def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None: """Register the Google Search tool.""" @mcp.tool( @@ -46,9 +44,7 @@ def register( ] = "", site: Annotated[ str | None, - Field( - description="Restrict results to one domain, e.g. 'example.com'." - ), + Field(description="Restrict results to one domain, e.g. 'example.com'."), ] = None, workspace: WorkspaceParam = None, response_format: ResponseFormatParam = "markdown", diff --git a/surfsense_mcp/mcp_server/features/scrapers/platforms/instagram.py b/surfsense_mcp/mcp_server/features/scrapers/platforms/instagram.py index 1792bf7dd..62e7552d6 100644 --- a/surfsense_mcp/mcp_server/features/scrapers/platforms/instagram.py +++ b/surfsense_mcp/mcp_server/features/scrapers/platforms/instagram.py @@ -17,9 +17,7 @@ ResultType = Literal["posts", "reels"] SearchType = Literal["profile", "user"] -def register( - mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext -) -> None: +def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None: """Register the Instagram scrape and details tools (anonymous-only).""" @mcp.tool( diff --git a/surfsense_mcp/mcp_server/features/scrapers/platforms/reddit.py b/surfsense_mcp/mcp_server/features/scrapers/platforms/reddit.py index 035193ebc..1bc2d851e 100644 --- a/surfsense_mcp/mcp_server/features/scrapers/platforms/reddit.py +++ b/surfsense_mcp/mcp_server/features/scrapers/platforms/reddit.py @@ -17,9 +17,7 @@ RedditSort = Literal["relevance", "hot", "top", "new", "rising", "comments"] RedditTime = Literal["hour", "day", "week", "month", "year", "all"] -def register( - mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext -) -> None: +def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None: """Register the Reddit tool.""" @mcp.tool( diff --git a/surfsense_mcp/mcp_server/features/scrapers/platforms/tiktok.py b/surfsense_mcp/mcp_server/features/scrapers/platforms/tiktok.py index 7e2bff509..620936063 100644 --- a/surfsense_mcp/mcp_server/features/scrapers/platforms/tiktok.py +++ b/surfsense_mcp/mcp_server/features/scrapers/platforms/tiktok.py @@ -14,9 +14,7 @@ from ..annotations import SCRAPE from ..capability import run_scraper -def register( - mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext -) -> None: +def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None: """Register the TikTok tools.""" @mcp.tool( @@ -32,7 +30,7 @@ def register( description="TikTok URLs: a video, a profile " "('https://www.tiktok.com/@nasa'), a hashtag " "('https://www.tiktok.com/tag/food'), or a search URL. Provide " - "urls OR profiles/hashtags/search_queries." + "urls OR profiles/hashtags." ), ] = None, profiles: Annotated[ @@ -45,23 +43,12 @@ def register( hashtags: Annotated[ list[str] | None, Field( - description="Hashtag names to scrape, without the '#', e.g. " - "['food']." - ), - ] = None, - search_queries: Annotated[ - list[str] | None, - Field( - description="Keyword search terms, resolved via Google to public " - "TikTok videos (TikTok's own keyword search is login-walled). " - "Slower than hashtags/urls, so start with at most 3 queries and " - "expand only if nothing significant is found. For accounts by " - "keyword use surfsense_tiktok_user_search." + description="Hashtag names to scrape, without the '#', e.g. ['food']." ), ] = None, results_per_page: Annotated[ int, - Field(ge=1, description="Max videos per profile/hashtag/search target."), + Field(ge=1, description="Max videos per profile/hashtag target."), ] = 10, max_items: Annotated[ int, Field(ge=1, description="Maximum videos to return in total.") @@ -69,15 +56,13 @@ def register( workspace: WorkspaceParam = None, response_format: ResponseFormatParam = "markdown", ) -> str: - """Scrape public TikTok videos by hashtag, profile, URL, or keyword. + """Scrape public TikTok videos by hashtag, profile, or URL. Use for TikTok video research — a creator's videos, a hashtag feed, or a specific video/profile/hashtag URL — instead of a generic web search. - search_queries also finds videos on a topic (resolved via Google), but is - slower: start with at most 3 queries and expand only if nothing - significant is found. Returns videos with text, author, stats, music, and - the web URL. For accounts by keyword use surfsense_tiktok_user_search. - Example: hashtags=['food'], max_items=20. + Returns videos with text, author, stats, music, and the web URL. There is + no keyword-video search; for accounts by keyword use + surfsense_tiktok_user_search. Example: hashtags=['food'], max_items=20. """ return await run_scraper( client, @@ -88,7 +73,6 @@ def register( "urls": urls, "profiles": profiles, "hashtags": hashtags, - "search_queries": search_queries, "results_per_page": results_per_page, "max_items": max_items, }, diff --git a/surfsense_mcp/mcp_server/features/scrapers/platforms/web.py b/surfsense_mcp/mcp_server/features/scrapers/platforms/web.py index 9c24a4352..8c6faaaf3 100644 --- a/surfsense_mcp/mcp_server/features/scrapers/platforms/web.py +++ b/surfsense_mcp/mcp_server/features/scrapers/platforms/web.py @@ -14,9 +14,7 @@ from ..annotations import SCRAPE from ..capability import run_scraper -def register( - mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext -) -> None: +def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None: """Register the web crawl tool.""" @mcp.tool( diff --git a/surfsense_mcp/mcp_server/features/scrapers/platforms/youtube.py b/surfsense_mcp/mcp_server/features/scrapers/platforms/youtube.py index 5582c82bb..1871fe711 100644 --- a/surfsense_mcp/mcp_server/features/scrapers/platforms/youtube.py +++ b/surfsense_mcp/mcp_server/features/scrapers/platforms/youtube.py @@ -16,9 +16,7 @@ from ..capability import run_scraper CommentSort = Literal["TOP_COMMENTS", "NEWEST_FIRST"] -def register( - mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext -) -> None: +def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None: """Register the YouTube video and comment tools.""" @mcp.tool( diff --git a/surfsense_mcp/mcp_server/features/scrapers/run_history.py b/surfsense_mcp/mcp_server/features/scrapers/run_history.py index 9274a1a69..ed36cfca6 100644 --- a/surfsense_mcp/mcp_server/features/scrapers/run_history.py +++ b/surfsense_mcp/mcp_server/features/scrapers/run_history.py @@ -17,9 +17,7 @@ from ...core.workspace_context import WorkspaceContext, WorkspaceParam from .annotations import READ_RUNS -def register( - mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext -) -> None: +def register(mcp: FastMCP, client: SurfSenseClient, context: WorkspaceContext) -> None: """Register the run-history tools.""" @mcp.tool( @@ -29,9 +27,7 @@ def register( structured_output=False, ) async def list_scraper_runs( - limit: Annotated[ - int, Field(ge=1, description="Maximum runs to list.") - ] = 20, + limit: Annotated[int, Field(ge=1, description="Maximum runs to list.")] = 20, capability: Annotated[ str | None, Field( diff --git a/surfsense_web/content/docs/connectors/native/tiktok.mdx b/surfsense_web/content/docs/connectors/native/tiktok.mdx index fc4cf99c3..60506da47 100644 --- a/surfsense_web/content/docs/connectors/native/tiktok.mdx +++ b/surfsense_web/content/docs/connectors/native/tiktok.mdx @@ -11,15 +11,14 @@ The TikTok connector pulls structured public data from TikTok across four verbs: POST /api/v1/workspaces/{workspace_id}/scrapers/tiktok/scrape ``` -Give it URLs (a video, a profile, a hashtag, or a search page) and/or profiles, hashtags, or search terms; returns videos (caption, author, play/like/comment/share counts, music, hashtags, timestamps, and the web URL). At least one of `urls`, `profiles`, `hashtags`, or `search_queries` is required. +Give it URLs (a video, a profile, a hashtag, or a search page) and/or profiles or hashtags; returns videos (caption, author, play/like/comment/share counts, music, hashtags, timestamps, and the web URL). At least one of `urls`, `profiles`, or `hashtags` is required. | Field | Default | Description | |-------|---------|-------------| | `urls` | — | TikTok URLs: a video, a profile (`/@`), a hashtag (`/tag/`), or a search URL (max 20 sources per call) | | `profiles` | — | Profile usernames, with or without a leading `@` | | `hashtags` | — | Hashtag names, without the `#` | -| `search_queries` | — | Search terms to run on TikTok | -| `results_per_page` | `10` | Max videos per profile/hashtag/search target | +| `results_per_page` | `10` | Max videos per profile/hashtag target | | `max_items` | `10` | Max total videos returned across all sources (hard cap 100) | ```bash @@ -30,7 +29,7 @@ curl -X POST "$BASE_URL/api/v1/workspaces/1/scrapers/tiktok/scrape" \ ``` -Video and hashtag targets are the reliable video paths. A `profiles` target returns the account's **metadata** (name, followers, bio, verification) reliably, but TikTok often withholds its **video list** from automated clients — so a profile can return metadata with no videos. Keyword **video** search is login-walled and returns a surfaced error; to find accounts by keyword use **user search** below. +Video and hashtag targets are the reliable video paths. A `profiles` target returns the account's **metadata** (name, followers, bio, verification) reliably, but TikTok often withholds its **video list** from automated clients — so a profile can return metadata with no videos. There is no keyword **video** search (TikTok's own search is login-walled); to find accounts by keyword use **user search** below. ## Comments diff --git a/surfsense_web/lib/connectors-marketing/tiktok.tsx b/surfsense_web/lib/connectors-marketing/tiktok.tsx index 8c6bb8b01..6063fd1ae 100644 --- a/surfsense_web/lib/connectors-marketing/tiktok.tsx +++ b/surfsense_web/lib/connectors-marketing/tiktok.tsx @@ -159,7 +159,7 @@ export const tiktok: ConnectorPageContent = { schema: { requestNote: - "Provide at least one source: urls, profiles, hashtags, or search_queries. Up to 20 sources per call.", + "Provide at least one source: urls, profiles, or hashtags. Up to 20 sources per call. To find accounts by keyword, use the user search verb.", request: [ { name: "urls", @@ -180,13 +180,6 @@ export const tiktok: ConnectorPageContent = { defaultValue: "[]", description: "Hashtag names to scrape, without the # prefix. Max 20.", }, - { - name: "search_queries", - type: "string[]", - defaultValue: "[]", - description: - "Keyword search terms. Keyword video search is login-walled and returns no videos; use hashtags/profiles/urls for videos, or user_search for accounts. Max 20.", - }, { name: "results_per_page", type: "integer", @@ -270,7 +263,7 @@ export const tiktok: ConnectorPageContent = { { question: "Can I scrape TikTok comments and hashtags?", answer: - "Yes. Pass a video URL to the comments endpoint for the public comment thread. Pass hashtag names or /tag/ URLs to the TikTok hashtag scraper to pull videos under that tag. Keyword video search is login-walled, so hashtags and direct URLs are the reliable discovery paths.", + "Yes. Pass a video URL to the comments endpoint for the public comment thread. Pass hashtag names or /tag/ URLs to the TikTok hashtag scraper to pull videos under that tag. Keyword video search is login-walled on TikTok, so hashtags and direct URLs are the reliable discovery paths; to find accounts by keyword, use the user search verb.", }, { question: "What are the rate limits?",