Merge remote-tracking branch 'upstream/dev' into fix/onboarding

This commit is contained in:
Anish Sarkar 2026-07-14 10:18:03 +05:30
commit 2d837ea18c
185 changed files with 4729 additions and 3630 deletions

View file

@ -75,6 +75,21 @@ jobs:
echo "Windows signing: skipped"
fi
- name: Verify Apple notarization credentials
if: matrix.os == 'macos-latest'
shell: bash
# Fails in seconds with the same 401 notarytool would throw after the
# full build, instead of wasting ~15 min of build time on bad secrets.
run: |
xcrun notarytool history \
--apple-id "$APPLE_ID" \
--password "$APPLE_APP_SPECIFIC_PASSWORD" \
--team-id "$APPLE_TEAM_ID" > /dev/null
env:
APPLE_ID: ${{ secrets.APPLE_ID }}
APPLE_APP_SPECIFIC_PASSWORD: ${{ secrets.APPLE_APP_SPECIFIC_PASSWORD }}
APPLE_TEAM_ID: ${{ secrets.APPLE_TEAM_ID }}
- name: Setup pnpm
uses: pnpm/action-setup@v5

View file

@ -79,7 +79,7 @@ We follow a **branch protection model** to keep `main` stable:
```
3. **Choose your setup method**:
- **Docker Setup**: Follow the [Docker Setup Guide](./DOCKER_SETUP.md)
- **Docker Setup**: Follow the [Building from Source (Contributors)](https://www.surfsense.com/docs/docker-installation#building-from-source-contributors) section of the Docker Installation guide
- **Manual Setup**: Follow the [Installation Guide](https://www.surfsense.com/docs/)
4. **Configure services**:

View file

@ -20,9 +20,9 @@
<a href="https://trendshift.io/repositories/13606" target="_blank"><img src="https://trendshift.io/api/badge/repositories/13606" alt="MODSetter%2FSurfSense | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
</div>
# SurfSense: Dale inteligencia competitiva a tus agentes de IA
# SurfSense: NotebookLM para investigación de inteligencia competitiva
SurfSense es la **plataforma de inteligencia competitiva de código abierto para agentes de IA**. Tus agentes monitorean a la competencia, siguen los rankings y escuchan a tu mercado con datos en vivo de **Reddit, YouTube, Google Maps, Google Search y la web abierta**, a través de una única **API REST** o un **servidor MCP**. Agentes programados y activados por eventos convierten lo que encuentran en informes y alertas, y una base de conocimiento integrada mantiene cada hallazgo disponible para búsqueda con citas.
SurfSense es la **plataforma de inteligencia competitiva de código abierto para agentes de IA**, como NotebookLM pero con conectores de scraping en vivo. Tus agentes monitorean a la competencia, siguen los rankings y escuchan a tu mercado con datos en vivo de **Reddit, YouTube, Instagram, TikTok, Google Maps, Google Search y la web abierta**, a través de una única **API REST** o un **servidor MCP**. Agentes programados y activados por eventos convierten lo que encuentran en informes y alertas, y una base de conocimiento integrada mantiene cada hallazgo disponible para búsqueda con citas.
> [!NOTE]
> **📢 Una nota para nuestros usuarios de la alternativa a NotebookLM**
@ -103,6 +103,8 @@ Las automatizaciones ejecutan turnos completos de agente según un horario o en
|---|---|---|
| **Reddit** | Publicaciones, comentarios y flujos de subreddits sin los límites de tasa de la API oficial | [Reddit Scraper API](https://www.surfsense.com/reddit) |
| **YouTube** | Videos, transcripciones e hilos de comentarios para escuchar a tu marca y productos | [YouTube Scraper API](https://www.surfsense.com/youtube) |
| **Instagram** | Perfiles, publicaciones y reels públicos sin la Graph API | [Instagram Scraper API](https://www.surfsense.com/instagram) |
| **TikTok** | Videos, comentarios, hashtags y perfiles sin aprobación de la Research API | [TikTok Scraper API](https://www.surfsense.com/tiktok) |
| **Google Maps** | Lugares, calificaciones y reseñas para investigar competidores locales y prospectos | [Google Maps Scraper API](https://www.surfsense.com/google-maps) |
| **Google Search** | SERPs en vivo para seguimiento de posiciones y monitoreo de mercado | [Google Search API](https://www.surfsense.com/google-search) |
| **Web Crawl** (rastreo web) | Cualquier página de la web abierta como contenido limpio y estructurado | [Web Crawling API](https://www.surfsense.com/web-crawl) |
@ -244,7 +246,7 @@ https://github.com/user-attachments/assets/a0a16566-6967-4374-ac51-9b3e07fbecd7
| Característica | Google NotebookLM | SurfSense |
|---------|-------------------|-----------|
| **Datos de mercado en vivo para agentes** | No | Conectores de Reddit, YouTube, Google Maps, Google Search y rastreo web vía API REST y MCP |
| **Datos de mercado en vivo para agentes** | No | Conectores de Reddit, YouTube, Instagram, TikTok, Google Maps, Google Search y rastreo web vía API REST y MCP |
| **Servidor MCP** | No | Cada conector expuesto como herramienta nativa de agente, más servidores MCP propios con aplicaciones OAuth de un clic |
| **Fuentes por Notebook** | 50 (gratis) a 600 (Ultra, $249.99/mes) | Ilimitadas |
| **Número de Notebooks** | 100 (gratis) a 500 (niveles de pago) | Ilimitado |

View file

@ -20,9 +20,9 @@
<a href="https://trendshift.io/repositories/13606" target="_blank"><img src="https://trendshift.io/api/badge/repositories/13606" alt="MODSetter%2FSurfSense | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
</div>
# SurfSense: अपने AI एजेंट्स को दें कॉम्पिटिटिव इंटेलिजेंस
# SurfSense: कॉम्पिटिटिव इंटेलिजेंस रिसर्च के लिए NotebookLM
SurfSense **AI एजेंट्स के लिए ओपन सोर्स कॉम्पिटिटिव इंटेलिजेंस प्लेटफ़ॉर्म** है। आपके एजेंट प्रतिस्पर्धियों पर नज़र रखते हैं, रैंकिंग ट्रैक करते हैं, और **Reddit, YouTube, Google Maps, Google Search और ओपन वेब** से लाइव डेटा के साथ आपके बाज़ार की बात सुनते हैं, वह भी एक ही **REST API** या **MCP सर्वर** के ज़रिए। शेड्यूल्ड और इवेंट-ट्रिगर्ड एजेंट अपनी खोजों को ब्रीफ़ और अलर्ट में बदलते हैं, और एक बिल्ट-इन नॉलेज बेस हर खोज को साइटेशन के साथ खोजने योग्य बनाए रखता है।
SurfSense **AI एजेंट्स के लिए ओपन सोर्स कॉम्पिटिटिव इंटेलिजेंस प्लेटफ़ॉर्म** है, बिलकुल NotebookLM जैसा, पर लाइव स्क्रैपिंग कनेक्टर्स के साथ। आपके एजेंट प्रतिस्पर्धियों पर नज़र रखते हैं, रैंकिंग ट्रैक करते हैं, और **Reddit, YouTube, Instagram, TikTok, Google Maps, Google Search और ओपन वेब** से लाइव डेटा के साथ आपके बाज़ार की बात सुनते हैं, वह भी एक ही **REST API** या **MCP सर्वर** के ज़रिए। शेड्यूल्ड और इवेंट-ट्रिगर्ड एजेंट अपनी खोजों को ब्रीफ़ और अलर्ट में बदलते हैं, और एक बिल्ट-इन नॉलेज बेस हर खोज को साइटेशन के साथ खोजने योग्य बनाए रखता है।
> [!NOTE]
> **📢 हमारे NotebookLM-विकल्प उपयोगकर्ताओं के लिए एक सूचना**
@ -103,6 +103,8 @@ SurfSense **AI एजेंट्स के लिए ओपन सोर्स
|---|---|---|
| **Reddit** | आधिकारिक API की रेट लिमिट के बिना पोस्ट, कमेंट और सबरेडिट स्ट्रीम | [Reddit Scraper API](https://www.surfsense.com/reddit) |
| **YouTube** | ब्रांड और प्रोडक्ट लिसनिंग के लिए वीडियो, ट्रांसक्रिप्ट और कमेंट थ्रेड | [YouTube Scraper API](https://www.surfsense.com/youtube) |
| **Instagram** | Graph API के बिना सार्वजनिक प्रोफ़ाइल, पोस्ट और रील्स | [Instagram Scraper API](https://www.surfsense.com/instagram) |
| **TikTok** | Research API अप्रूवल के बिना वीडियो, कमेंट, हैशटैग और प्रोफ़ाइल | [TikTok Scraper API](https://www.surfsense.com/tiktok) |
| **Google Maps** | स्थानीय प्रतिस्पर्धी और लीड रिसर्च के लिए स्थान, रेटिंग और रिव्यू | [Google Maps Scraper API](https://www.surfsense.com/google-maps) |
| **Google Search** | रैंक ट्रैकिंग और मार्केट मॉनिटरिंग के लिए लाइव SERP | [Google Search API](https://www.surfsense.com/google-search) |
| **Web Crawl** | ओपन वेब का कोई भी पेज साफ़-सुथरे, स्ट्रक्चर्ड कंटेंट के रूप में | [Web Crawling API](https://www.surfsense.com/web-crawl) |
@ -244,7 +246,7 @@ https://github.com/user-attachments/assets/a0a16566-6967-4374-ac51-9b3e07fbecd7
| फ़ीचर | Google NotebookLM | SurfSense |
|---------|-------------------|-----------|
| **एजेंट्स के लिए लाइव मार्केट डेटा** | नहीं | REST API और MCP के ज़रिए Reddit, YouTube, Google Maps, Google Search और वेब क्रॉल कनेक्टर |
| **एजेंट्स के लिए लाइव मार्केट डेटा** | नहीं | REST API और MCP के ज़रिए Reddit, YouTube, Instagram, TikTok, Google Maps, Google Search और वेब क्रॉल कनेक्टर |
| **MCP सर्वर** | नहीं | हर कनेक्टर नेटिव एजेंट टूल के रूप में उपलब्ध, साथ ही वन-क्लिक OAuth ऐप्स के साथ अपने MCP सर्वर लाने की सुविधा |
| **प्रति नोटबुक स्रोत** | 50 (Free) से 600 (Ultra, $249.99/माह) | असीमित |
| **नोटबुक की संख्या** | 100 (Free) से 500 (सशुल्क टियर) | असीमित |

View file

@ -20,9 +20,9 @@
<a href="https://trendshift.io/repositories/13606" target="_blank"><img src="https://trendshift.io/api/badge/repositories/13606" alt="MODSetter%2FSurfSense | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
</div>
# SurfSense: Give Your AI Agents Competitive Intelligence
# SurfSense: NotebookLM for Competitive Intelligence Research
SurfSense is the **open-source competitive intelligence platform for AI agents**. Your agents monitor competitors, track rankings, and listen to your market with live data from **Reddit, YouTube, Google Maps, Google Search, and the open web**, through one **REST API** or **MCP server**. Scheduled and event-triggered agents turn what they find into briefs and alerts, and a built-in knowledge base keeps every finding searchable with citations.
SurfSense is the **open-source competitive intelligence platform for AI agents**, like NotebookLM but with live scraping connectors. Your agents monitor competitors, track rankings, and listen to your market with live data from **Reddit, YouTube, Instagram, TikTok, Google Maps, Google Search, and the open web**, through one **REST API** or **MCP server**. Scheduled and event-triggered agents turn what they find into briefs and alerts, and a built-in knowledge base keeps every finding searchable with citations.
> [!NOTE]
> **📢 A note for our NotebookLM-alternative users**
@ -103,6 +103,8 @@ Automations run full agent turns on a schedule or in response to events, then wr
|---|---|---|
| **Reddit** | Posts, comments, and subreddit streams without the official API's rate limits | [Reddit Scraper API](https://www.surfsense.com/reddit) |
| **YouTube** | Videos, transcripts, and comment threads for brand and product listening | [YouTube Scraper API](https://www.surfsense.com/youtube) |
| **Instagram** | Public profiles, posts, and reels without the Graph API | [Instagram Scraper API](https://www.surfsense.com/instagram) |
| **TikTok** | Videos, comments, hashtags, and profiles without Research API approval | [TikTok Scraper API](https://www.surfsense.com/tiktok) |
| **Google Maps** | Places, ratings, and reviews for local competitor and lead research | [Google Maps Scraper API](https://www.surfsense.com/google-maps) |
| **Google Search** | Live SERPs for rank tracking and market monitoring | [Google Search API](https://www.surfsense.com/google-search) |
| **Web Crawl** | Any page on the open web as clean, structured content | [Web Crawling API](https://www.surfsense.com/web-crawl) |
@ -243,7 +245,7 @@ Still comparing us as a NotebookLM alternative? Here is the honest breakdown.
| Feature | Google NotebookLM | SurfSense |
|---------|-------------------|-----------|
| **Live market data for agents** | No | Reddit, YouTube, Google Maps, Google Search, and web crawl connectors via REST API and MCP |
| **Live market data for agents** | No | Reddit, YouTube, Instagram, TikTok, Google Maps, Google Search, and web crawl connectors via REST API and MCP |
| **MCP server** | No | Every connector exposed as a native agent tool, plus bring-your-own MCP servers with one-click OAuth apps |
| **Sources per Notebook** | 50 (Free) to 600 (Ultra, $249.99/mo) | Unlimited |
| **Number of Notebooks** | 100 (Free) to 500 (paid tiers) | Unlimited |

View file

@ -20,9 +20,9 @@
<a href="https://trendshift.io/repositories/13606" target="_blank"><img src="https://trendshift.io/api/badge/repositories/13606" alt="MODSetter%2FSurfSense | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
</div>
# SurfSense: Dê Inteligência Competitiva aos Seus Agentes de IA
# SurfSense: NotebookLM para Pesquisa de Inteligência Competitiva
O SurfSense é a **plataforma open source de inteligência competitiva para agentes de IA**. Seus agentes monitoram concorrentes, acompanham rankings e escutam o seu mercado com dados ao vivo do **Reddit, YouTube, Google Maps, Google Search e da web aberta**, por meio de uma única **API REST** ou de um **servidor MCP**. Agentes agendados ou acionados por eventos transformam o que encontram em relatórios e alertas, e uma base de conhecimento integrada mantém cada descoberta pesquisável, com citações.
O SurfSense é a **plataforma open source de inteligência competitiva para agentes de IA**, como o NotebookLM, mas com conectores de scraping ao vivo. Seus agentes monitoram concorrentes, acompanham rankings e escutam o seu mercado com dados ao vivo do **Reddit, YouTube, Instagram, TikTok, Google Maps, Google Search e da web aberta**, por meio de uma única **API REST** ou de um **servidor MCP**. Agentes agendados ou acionados por eventos transformam o que encontram em relatórios e alertas, e uma base de conhecimento integrada mantém cada descoberta pesquisável, com citações.
> [!NOTE]
> **📢 Um recado para nossos usuários que buscavam uma alternativa ao NotebookLM**
@ -103,6 +103,8 @@ As automações executam turnos completos de agente de forma agendada ou em resp
|---|---|---|
| **Reddit** | Posts, comentários e fluxos de subreddits sem os limites de requisição da API oficial | [Reddit Scraper API](https://www.surfsense.com/reddit) |
| **YouTube** | Vídeos, transcrições e threads de comentários para monitoramento de marca e produto | [YouTube Scraper API](https://www.surfsense.com/youtube) |
| **Instagram** | Perfis, posts e reels públicos sem a Graph API | [Instagram Scraper API](https://www.surfsense.com/instagram) |
| **TikTok** | Vídeos, comentários, hashtags e perfis sem aprovação da Research API | [TikTok Scraper API](https://www.surfsense.com/tiktok) |
| **Google Maps** | Estabelecimentos, avaliações e reviews para pesquisa local de concorrentes e leads | [Google Maps Scraper API](https://www.surfsense.com/google-maps) |
| **Google Search** | SERPs ao vivo para acompanhamento de rankings e monitoramento de mercado | [Google Search API](https://www.surfsense.com/google-search) |
| **Web Crawl** (rastreamento web) | Qualquer página da web aberta como conteúdo limpo e estruturado | [Web Crawling API](https://www.surfsense.com/web-crawl) |
@ -244,7 +246,7 @@ Ainda nos comparando como alternativa ao NotebookLM? Aqui está o comparativo ho
| Recurso | Google NotebookLM | SurfSense |
|---------|-------------------|-----------|
| **Dados de mercado ao vivo para agentes** | Não | Conectores de Reddit, YouTube, Google Maps, Google Search e rastreamento web via API REST e MCP |
| **Dados de mercado ao vivo para agentes** | Não | Conectores de Reddit, YouTube, Instagram, TikTok, Google Maps, Google Search e rastreamento web via API REST e MCP |
| **Servidor MCP** | Não | Cada conector exposto como ferramenta nativa de agente, além de servidores MCP próprios com apps OAuth em um clique |
| **Fontes por Notebook** | 50 (gratuito) a 600 (Ultra, US$ 249,99/mês) | Ilimitadas |
| **Número de Notebooks** | 100 (gratuito) a 500 (planos pagos) | Ilimitado |

View file

@ -20,9 +20,9 @@
<a href="https://trendshift.io/repositories/13606" target="_blank"><img src="https://trendshift.io/api/badge/repositories/13606" alt="MODSetter%2FSurfSense | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
</div>
# SurfSense为你的 AI 智能体注入竞争情报能力
# SurfSense面向竞争情报研究的 NotebookLM
SurfSense 是**面向 AI 智能体的开源竞争情报平台**。你的智能体可以通过一个 **REST API****MCP 服务器**,利用来自 **Reddit、YouTube、Google Maps、Google Search 和开放网络**的实时数据,监控竞争对手、追踪排名、倾听市场动态。定时和事件触发的智能体会把发现的内容转化为简报和预警,内置的知识库则让每一条发现都可搜索、可引用。
SurfSense 是**面向 AI 智能体的开源竞争情报平台**,就像 NotebookLM但配备了实时抓取连接器。你的智能体可以通过一个 **REST API****MCP 服务器**,利用来自 **Reddit、YouTube、Instagram、TikTok、Google Maps、Google Search 和开放网络**的实时数据,监控竞争对手、追踪排名、倾听市场动态。定时和事件触发的智能体会把发现的内容转化为简报和预警,内置的知识库则让每一条发现都可搜索、可引用。
> [!NOTE]
> **📢 致我们的 NotebookLM 替代品用户**
@ -103,6 +103,8 @@ SurfSense 是**面向 AI 智能体的开源竞争情报平台**。你的智能
|---|---|---|
| **Reddit** | 帖子、评论和子版块信息流,不受官方 API 速率限制 | [Reddit Scraper API](https://www.surfsense.com/reddit) |
| **YouTube** | 视频、字幕转录和评论串,用于品牌和产品舆情监听 | [YouTube Scraper API](https://www.surfsense.com/youtube) |
| **Instagram** | 公开主页、帖子和 Reels无需 Graph API | [Instagram Scraper API](https://www.surfsense.com/instagram) |
| **TikTok** | 视频、评论、话题标签和主页,无需 Research API 审批 | [TikTok Scraper API](https://www.surfsense.com/tiktok) |
| **Google Maps** | 地点、评分和评论,用于本地竞争对手和潜在客户调研 | [Google Maps Scraper API](https://www.surfsense.com/google-maps) |
| **Google Search** | 实时搜索结果页,用于排名追踪和市场监控 | [Google Search API](https://www.surfsense.com/google-search) |
| **Web Crawl** | 把开放网络上的任意页面转为干净、结构化的内容 | [Web Crawling API](https://www.surfsense.com/web-crawl) |
@ -244,7 +246,7 @@ https://github.com/user-attachments/assets/a0a16566-6967-4374-ac51-9b3e07fbecd7
| 功能 | Google NotebookLM | SurfSense |
|---------|-------------------|-----------|
| **面向智能体的实时市场数据** | 无 | 通过 REST API 和 MCP 提供 Reddit、YouTube、Google Maps、Google Search 和网页爬取连接器 |
| **面向智能体的实时市场数据** | 无 | 通过 REST API 和 MCP 提供 Reddit、YouTube、Instagram、TikTok、Google Maps、Google Search 和网页爬取连接器 |
| **MCP 服务器** | 无 | 每个连接器都作为原生智能体工具暴露,还可自带 MCP 服务器并使用一键 OAuth 应用 |
| **每个笔记本的来源数** | 50 个(免费版)至 600 个Ultra 版249.99 美元/月) | 无限制 |
| **笔记本数量** | 100 个(免费版)至 500 个(付费档位) | 无限制 |

View file

@ -1 +1 @@
0.0.31
0.0.32

View file

@ -120,8 +120,10 @@ STRIPE_RECONCILIATION_BATCH_SIZE=100
# BACKEND_URL=https://api.yourdomain.com
# Auth
AUTH_TYPE=GOOGLE or LOCAL
REGISTRATION_ENABLED=TRUE or FALSE
# AUTH_TYPE: GOOGLE or LOCAL
AUTH_TYPE=LOCAL
# REGISTRATION_ENABLED: TRUE or FALSE
REGISTRATION_ENABLED=TRUE
# For Google Auth Only
GOOGLE_OAUTH_CLIENT_ID=924507538m
GOOGLE_OAUTH_CLIENT_SECRET=GOCSV

View file

@ -13,7 +13,7 @@ Answer the delegated question from live Instagram data gathered with your verbs,
<playbook>
- Known profile/post/reel links: call `instagram_scrape` with the links in `urls` (use `result_type` to pick posts or reels). Hashtag/place URLs are unsupported (login-walled).
- Finding a profile on a topic: call `instagram_scrape` with `search_queries` (resolved to public profiles via Google; `search_type` is profile-only).
- Finding a profile on a topic: call `instagram_scrape` with `search_queries` (resolved to public profiles via Google; `search_type` is profile-only). Google-backed discovery is slow (~30-60s per query), so start with **at most 3** distinct queries per task and only add more if the first round returns nothing significant — never batch many phrasing variants of the same intent.
- Profile metadata (follower counts, bio, post count): call `instagram_details`.
- Batch multiple URLs (or queries) into one call rather than many single-item calls.
<include snippet="run_reader"/>

View file

@ -14,11 +14,11 @@ Answer the delegated question from live TikTok data gathered with your verb, com
</available_tools>
<playbook>
- Finding videos on a topic: call `tiktok_scrape` with `hashtags` (no leading '#'), or pass a TikTok URL in `urls`.
- 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`. Keyword search returns no videos, so do not use `search_queries` for videos — use `hashtags` or a video URL.
- Finding accounts by keyword: call `tiktok_user_search` with `queries` — that is the path for accounts.
- "What's trending now": call `tiktok_trending` (no query needed); set `max_items` for how many.
- Controlling volume: use `max_items` for the total cap and `results_per_page` per target (per-verb equivalents: `comments_per_video`, `results_per_query`).
- Requested counts: `max_items` defaults low — when the task asks for N items, set `max_items` (and the per-target count) above N. A call that caps below the target can never satisfy it.

View file

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

View file

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

View file

@ -10,8 +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."
"Scrape public TikTok videos. Use urls, profiles, or hashtags. To find "
"accounts by keyword, use tiktok.user_search."
),
input_schema=ScrapeInput,
output_schema=ScrapeOutput,

View file

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

View file

@ -4,7 +4,8 @@ A lean, agent-friendly surface over ``TikTokScrapeInput``
(``app/proprietary/platforms/tiktok``). The executor maps this to the full
scraper input; the scraper's ``TikTokVideoItem`` is reused verbatim as the
output element. Any TikTok URL kind (video, profile, hashtag, search) goes in
``urls``; ``profiles``/``hashtags``/``search_queries`` are typed shortcuts.
``urls``; ``profiles``/``hashtags`` are typed shortcuts. Keyword search is not a
video source here use ``tiktok.user_search`` to find accounts by keyword.
"""
from __future__ import annotations
@ -26,8 +27,8 @@ class ScrapeInput(BaseModel):
max_length=MAX_TIKTOK_SOURCES,
description=(
"TikTok URLs to scrape: a video, a profile (/@<user>), a hashtag "
"(/tag/<name>), or a search URL. Provide these OR profiles/hashtags/"
"search_queries (at least one source is required)."
"(/tag/<name>), or a search URL. Provide these OR profiles/hashtags "
"(at least one source is required)."
),
)
profiles: list[str] = Field(
@ -40,16 +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 to run on TikTok.",
)
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,
@ -60,10 +56,9 @@ class ScrapeInput(BaseModel):
@model_validator(mode="after")
def _require_a_source(self) -> ScrapeInput:
if not any((self.urls, self.profiles, self.hashtags, self.search_queries)):
if not any((self.urls, self.profiles, self.hashtags)):
raise ValueError(
"Provide at least one of 'urls', 'profiles', 'hashtags', or "
"'search_queries'."
"Provide at least one of 'urls', 'profiles', or 'hashtags'."
)
return self

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -10,7 +10,6 @@ from __future__ import annotations
from collections.abc import AsyncIterator
from typing import Any
from urllib.parse import quote
from .extraction.timestamps import now_iso
from .flows import FetchCommentsFn, FetchFn, FetchListingFn, FetchUsersFn
@ -32,12 +31,16 @@ 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"
def _resolve_targets(input_model: TikTokScrapeInput) -> list[TikTokTarget]:
"""Build the target list from every input source, dropping unresolved URLs."""
"""Build the target list from the URL/profile/hashtag sources.
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:
resolved = resolve_target(entry.url)
@ -52,10 +55,6 @@ def _resolve_targets(input_model: TikTokScrapeInput) -> list[TikTokTarget]:
targets.append(TikTokTarget("profile", name, _PROFILE_URL.format(name=name)))
for tag in input_model.hashtags:
targets.append(TikTokTarget("hashtag", tag, _HASHTAG_URL.format(tag=tag)))
for query in input_model.searchQueries:
targets.append(
TikTokTarget("search", query, _SEARCH_URL.format(query=quote(query)))
)
return targets
@ -104,7 +103,9 @@ async def scrape_tiktok(
from app.capabilities.core.progress import emit_progress
results: list[dict[str, Any]] = []
async for item in iter_tiktok(input_model, fetch=fetch, fetch_listing=fetch_listing):
async for item in iter_tiktok(
input_model, fetch=fetch, fetch_listing=fetch_listing
):
results.append(item)
emit_progress("scraping", current=len(results), total=limit, unit="item")
if limit is not None and len(results) >= limit:

View file

@ -15,6 +15,7 @@ from app.db import Connection, Model, ModelSource
from app.services.model_resolver import ensure_v1, to_litellm
from app.services.openrouter_model_normalizer import normalize_openrouter_models
from app.services.provider_registry import Transport, provider_label, spec_for
from app.services.requesty_model_normalizer import normalize_requesty_models
logger = logging.getLogger(__name__)
@ -148,7 +149,7 @@ async def verify_connection(conn: Connection) -> VerifyResult:
if spec.transport == Transport.OLLAMA and base_url:
url = f"{base_url.rstrip('/')}/api/version"
elif spec.discovery in {"openai_models", "openrouter"} and base_url:
elif spec.discovery in {"openai_models", "openrouter", "requesty"} and base_url:
url = f"{ensure_v1(base_url)}/models"
elif spec.discovery == "anthropic_models" and base_url:
url = f"{base_url.rstrip('/')}/models"
@ -363,6 +364,16 @@ async def _openrouter_models(conn: Connection) -> list[dict[str, Any]]:
return normalize_openrouter_models(response.json().get("data", []))
async def _requesty_models(conn: Connection) -> list[dict[str, Any]]:
base_url = _base_url_or_default(conn) or "https://router.requesty.ai/v1"
async with httpx.AsyncClient(timeout=DISCOVERY_TIMEOUT_SECONDS) as client:
response = await client.get(
f"{ensure_v1(base_url)}/models", headers=_auth_headers(conn)
)
response.raise_for_status()
return normalize_requesty_models(response.json().get("data", []))
def _litellm_static_models(conn: Connection) -> list[dict[str, Any]]:
provider = conn.provider
prefix = spec_for(provider).litellm_prefix or provider
@ -446,6 +457,8 @@ async def discover_models(conn: Connection) -> list[dict[str, Any]]:
results = await _ollama_tags_then_show(conn)
elif spec.discovery == "openrouter":
results = await _openrouter_models(conn)
elif spec.discovery == "requesty":
results = await _requesty_models(conn)
elif spec.discovery == "anthropic_models":
results = await _discover_anthropic_models(conn)
elif spec.discovery == "openai_models":

View file

@ -23,6 +23,7 @@ DiscoveryKind = Literal[
"anthropic_models",
"bedrock_models",
"openrouter",
"requesty",
"static",
"none",
]
@ -78,6 +79,15 @@ REGISTRY: dict[str, ProviderSpec] = {
"bearer",
"OpenRouter",
),
"requesty": ProviderSpec(
Transport.OPENAI_COMPATIBLE,
"openai",
"requesty",
"https://router.requesty.ai/v1",
False,
"bearer",
"Requesty",
),
"openai_compatible": ProviderSpec(
Transport.OPENAI_COMPATIBLE,
"openai",
@ -87,6 +97,15 @@ REGISTRY: dict[str, ProviderSpec] = {
"bearer",
"OpenAI-compatible provider",
),
"openai_compatible_raw": ProviderSpec(
Transport.NATIVE,
"openai",
"none",
None,
True,
"bearer",
"OpenAI-compatible raw endpoint",
),
"lm_studio": ProviderSpec(
Transport.OPENAI_COMPATIBLE,
"openai",

View file

@ -123,6 +123,7 @@ PROVIDER_PRESTIGE_YAML: dict[str, int] = {
"perplexity": 28,
"bedrock": 28,
"openrouter": 25,
"requesty": 25,
"ollama_chat": 12,
"custom": 12,
}

View file

@ -0,0 +1,123 @@
"""Shared Requesty model normalization.
Requesty (https://router.requesty.ai) is an OpenAI-compatible LLM router.
Its ``/v1/models`` catalogue carries richer, Requesty-specific capability
metadata than a generic OpenAI-compatible ``/models`` response, so keep all
Requesty filtering and capability extraction here -- mirroring
``openrouter_model_normalizer`` -- so GLOBAL catalogue generation and BYOK
discovery agree.
Unlike OpenRouter, Requesty exposes capabilities as flat booleans
(``supports_tool_calling`` / ``supports_reasoning`` / ``supports_vision`` /
``supports_image_generation``) rather than an ``architecture`` block plus a
``supported_parameters`` array, and it reports context size as
``context_window`` rather than ``context_length``. This module maps those
fields onto the same normalized shape the rest of the backend consumes.
"""
from __future__ import annotations
from typing import Any
from app.db import ModelSource
MIN_CONTEXT_LENGTH = 100_000
EXCLUDED_PROVIDER_SLUGS: set[str] = {"amazon"}
EXCLUDED_MODEL_IDS: set[str] = set()
EXCLUDED_MODEL_SUFFIXES: tuple[str, ...] = ("-deep-research",)
def is_image_output_model(model: dict[str, Any]) -> bool:
return bool(model.get("supports_image_generation"))
def is_text_output_model(model: dict[str, Any]) -> bool:
# Requesty entries are chat-completion models (``api == "chat"``). Treat a
# model as text output whenever it is not an image-generation model.
return not is_image_output_model(model)
def supports_image_input(model: dict[str, Any]) -> bool:
return bool(model.get("supports_vision"))
def supports_tool_calling(model: dict[str, Any]) -> bool:
return bool(model.get("supports_tool_calling"))
def has_sufficient_context(model: dict[str, Any]) -> bool:
return int(model.get("context_window") or 0) >= MIN_CONTEXT_LENGTH
def is_compatible_provider(model: dict[str, Any]) -> bool:
model_id = str(model.get("id") or "")
slug = model_id.split("/", 1)[0] if "/" in model_id else ""
return slug not in EXCLUDED_PROVIDER_SLUGS
def is_allowed_model(model: dict[str, Any]) -> bool:
model_id = str(model.get("id") or "")
if model_id in EXCLUDED_MODEL_IDS:
return False
base_id = model_id.split(":")[0]
return not base_id.endswith(EXCLUDED_MODEL_SUFFIXES)
def is_requesty_chat_model(model: dict[str, Any]) -> bool:
return (
"/" in str(model.get("id") or "")
and is_text_output_model(model)
and supports_tool_calling(model)
and has_sufficient_context(model)
and is_compatible_provider(model)
and is_allowed_model(model)
)
def is_requesty_image_model(model: dict[str, Any]) -> bool:
return (
"/" in str(model.get("id") or "")
and is_image_output_model(model)
and is_compatible_provider(model)
and is_allowed_model(model)
)
def normalize_requesty_models(
raw_models: list[dict[str, Any]],
) -> list[dict[str, Any]]:
normalized: list[dict[str, Any]] = []
for model in raw_models:
if not is_requesty_chat_model(model):
continue
model_id = str(model.get("id") or "")
normalized.append(
{
"model_id": model_id,
"display_name": model.get("name") or model_id,
"source": ModelSource.DISCOVERED,
"supports_chat": True,
"max_input_tokens": model.get("context_window"),
"supports_image_input": supports_image_input(model),
"supports_tools": supports_tool_calling(model),
"supports_image_generation": False,
"metadata": model,
}
)
return normalized
__all__ = [
"MIN_CONTEXT_LENGTH",
"has_sufficient_context",
"is_allowed_model",
"is_compatible_provider",
"is_image_output_model",
"is_requesty_chat_model",
"is_requesty_image_model",
"is_text_output_model",
"normalize_requesty_models",
"supports_image_input",
"supports_tool_calling",
]

View file

@ -1,6 +1,6 @@
[project]
name = "surf-new-backend"
version = "0.0.31"
version = "0.0.32"
description = "SurfSense Backend"
requires-python = ">=3.12"
dependencies = [

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -64,6 +64,22 @@ def test_openai_compatible_resolver_uses_explicit_api_base() -> None:
assert ensure_v1("http://example.com/v1") == "http://example.com/v1"
def test_openai_compatible_raw_resolver_does_not_append_v1() -> None:
model, kwargs = to_litellm(
{
"provider": "openai_compatible_raw",
"base_url": "https://ark.cn-beijing.volces.com/api/v3",
"api_key": "ark-key",
"extra": {},
},
"ep-20260101000000-test",
)
assert model == "openai/ep-20260101000000-test"
assert kwargs["api_base"] == "https://ark.cn-beijing.volces.com/api/v3"
assert kwargs["api_key"] == "ark-key"
def test_ollama_resolver_uses_native_api_base() -> None:
model, kwargs = to_litellm(
{

View file

@ -0,0 +1,103 @@
"""Unit tests for Requesty model normalization.
Mirrors the OpenRouter normalizer coverage but exercises Requesty's flat
boolean capability fields (``supports_tool_calling`` / ``supports_vision``)
and ``context_window`` sizing.
"""
from __future__ import annotations
import pytest
from app.services.requesty_model_normalizer import (
is_requesty_chat_model,
is_requesty_image_model,
normalize_requesty_models,
supports_image_input,
supports_tool_calling,
)
pytestmark = pytest.mark.unit
def _requesty_model(
*,
model_id: str,
context_window: int = 128_000,
tools: bool = True,
vision: bool = False,
image_generation: bool = False,
name: str | None = None,
) -> dict:
"""Return a synthetic Requesty ``/v1/models`` entry.
Only the fields the normalizer inspects are populated; the live payload
carries many more (pricing, ``supports_caching``, ``description``, ...).
"""
return {
"id": model_id,
"name": name or model_id,
"api": "chat",
"object": "model",
"context_window": context_window,
"supports_tool_calling": tools,
"supports_vision": vision,
"supports_image_generation": image_generation,
}
def test_chat_model_requires_slash_tools_and_context():
assert is_requesty_chat_model(_requesty_model(model_id="openai/gpt-4o-mini"))
assert not is_requesty_chat_model(
_requesty_model(model_id="openai/gpt-4o-mini", tools=False)
)
assert not is_requesty_chat_model(
_requesty_model(model_id="openai/gpt-4o-mini", context_window=8_000)
)
assert not is_requesty_chat_model(_requesty_model(model_id="bare-model"))
def test_excluded_provider_slug_is_filtered():
assert not is_requesty_chat_model(_requesty_model(model_id="amazon/nova-pro-v1"))
def test_image_generation_models_excluded_from_chat_and_flagged():
image_model = _requesty_model(
model_id="google/gemini-2.5-flash-image", image_generation=True
)
assert not is_requesty_chat_model(image_model)
assert is_requesty_image_model(image_model)
def test_capability_helpers_read_flat_booleans():
model = _requesty_model(
model_id="anthropic/claude-sonnet-4-5", vision=True, tools=True
)
assert supports_image_input(model) is True
assert supports_tool_calling(model) is True
def test_normalize_maps_context_window_and_capabilities():
normalized = normalize_requesty_models(
[
_requesty_model(
model_id="openai/gpt-4o-mini",
context_window=128_000,
vision=True,
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),
]
)
assert len(normalized) == 1
entry = normalized[0]
assert entry["model_id"] == "openai/gpt-4o-mini"
assert entry["display_name"] == "GPT-4o mini"
assert entry["supports_chat"] is True
assert entry["max_input_tokens"] == 128_000
assert entry["supports_image_input"] is True
assert entry["supports_tools"] is True
assert entry["supports_image_generation"] is False
assert entry["metadata"]["id"] == "openai/gpt-4o-mini"

View file

@ -36,6 +36,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra == 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -70,6 +73,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version < '3.13' and sys_platform == 'linux' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra == 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -104,6 +110,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra == 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -138,6 +147,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform == 'emscripten' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra == 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -172,6 +184,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra == 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -206,6 +221,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform == 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra == 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -240,6 +258,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform == 'emscripten' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra == 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -274,6 +295,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra == 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -308,6 +332,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version < '3.13' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra == 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -342,6 +369,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version < '3.13' and sys_platform == 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra == 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -387,6 +417,10 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform == 'linux' and extra != 'extra-16-surf-new-backend-cpu' and extra == 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -421,6 +455,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform == 'linux' and extra != 'extra-16-surf-new-backend-cpu' and extra == 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -455,6 +492,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version < '3.13' and sys_platform == 'linux' and extra != 'extra-16-surf-new-backend-cpu' and extra == 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -489,6 +529,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra == 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -523,6 +566,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform == 'emscripten' and extra != 'extra-16-surf-new-backend-cpu' and extra == 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -557,6 +603,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra == 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -591,6 +640,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform == 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra == 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -625,6 +677,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform == 'emscripten' and extra != 'extra-16-surf-new-backend-cpu' and extra == 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -659,6 +714,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra == 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -693,6 +751,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version < '3.13' and sys_platform != 'linux' and sys_platform != 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra == 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -727,6 +788,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version < '3.13' and sys_platform == 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra == 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -772,6 +836,10 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform == 'linux' and extra == 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -806,6 +874,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform == 'linux' and extra == 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -840,6 +911,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version < '3.13' and sys_platform == 'linux' and extra == 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -874,6 +948,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform == 'win32' and extra == 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -908,6 +985,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform == 'emscripten' and extra == 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -942,6 +1022,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -976,6 +1059,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform == 'win32' and extra == 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -1010,6 +1096,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform == 'emscripten' and extra == 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -1044,6 +1133,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -1078,6 +1170,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version < '3.13' and sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -1112,6 +1207,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version < '3.13' and sys_platform == 'win32' and extra == 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -1157,6 +1255,10 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -1191,6 +1293,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform == 'emscripten' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -1258,6 +1363,12 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -1292,6 +1403,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform == 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -1326,6 +1440,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform == 'emscripten' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -1393,6 +1510,12 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -1460,6 +1583,12 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version < '3.13' and sys_platform != 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
"python_version < '0'",
"python_version < '0'",
@ -1494,6 +1623,9 @@ resolution-markers = [
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_version < '0'",
"python_full_version < '3.13' and sys_platform == 'win32' and extra != 'extra-16-surf-new-backend-cpu' and extra != 'extra-16-surf-new-backend-cu126' and extra != 'extra-16-surf-new-backend-cu128'",
]
conflicts = [[
@ -10353,7 +10485,7 @@ wheels = [
[[package]]
name = "surf-new-backend"
version = "0.0.31"
version = "0.0.32"
source = { editable = "." }
dependencies = [
{ name = "alembic" },

View file

@ -1,7 +1,7 @@
{
"name": "surfsense_browser_extension",
"displayName": "Surfsense Browser Extension",
"version": "0.0.31",
"version": "0.0.32",
"description": "Extension to collect Browsing History for SurfSense.",
"author": "https://github.com/MODSetter",
"engines": {

View file

@ -1,7 +1,7 @@
{
"name": "surfsense-desktop",
"productName": "SurfSense",
"version": "0.0.31",
"version": "0.0.32",
"description": "SurfSense Desktop App",
"main": "dist/main.js",
"scripts": {

View file

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

View file

@ -12,12 +12,10 @@ Outputs (printed to stdout + written to `failures_n171.json`):
from __future__ import annotations
import json
import re
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any
REPO = Path(__file__).resolve().parents[1]
RUN = REPO / "data" / "multimodal_doc" / "runs" / "2026-05-14T00-53-19Z" / "parser_compare"
RAW = RUN / "raw.jsonl"
@ -53,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)
@ -123,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] = {
@ -132,18 +131,13 @@ def main() -> None:
"n": n_per_arm[arm],
"failures": len(by_arm_failures[arm]),
"rate": len(by_arm_failures[arm]) / n_per_arm[arm],
"clusters": {
cluster: len(items)
for cluster, items in error_clusters[arm].items()
},
"clusters": {cluster: len(items) for cluster, items in error_clusters[arm].items()},
"rows": by_arm_failures[arm],
}
for arm in sorted(n_per_arm)
},
"per_pdf": {
pdf: [
{**r, "arm": r["arm"]} for r in failures
]
pdf: [{**r, "arm": r["arm"]} for r in failures]
for pdf, failures in by_pdf_failures.items()
},
}

View file

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

View file

@ -15,7 +15,6 @@ from pathlib import Path
import httpx
from dotenv import load_dotenv
REPO = Path(__file__).resolve().parents[1]
PDF_DIR = REPO / "data" / "multimodal_doc" / "mmlongbench" / "pdfs"

View file

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

View file

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

View file

@ -47,9 +47,7 @@ def _row_key(row: dict) -> tuple[str, str]:
def _is_failure(row: dict) -> bool:
if row.get("error"):
return True
if not (row.get("raw_text") or "").strip():
return True
return False
return bool(not (row.get("raw_text") or "").strip())
def _summarise(rows_by_arm: dict[str, list[dict]]) -> dict[str, dict]:
@ -80,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:
@ -105,9 +105,7 @@ def main() -> int:
raw_rows = _read_jsonl(raw_path)
retry_rows = _read_jsonl(retry_path)
retry_by_key: dict[tuple[str, str], dict] = {
_row_key(r): r for r in retry_rows
}
retry_by_key: dict[tuple[str, str], dict] = {_row_key(r): r for r in retry_rows}
merged_rows: list[dict] = []
n_replaced_recovered = 0

View file

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

View file

@ -27,7 +27,6 @@ from __future__ import annotations
import json
from pathlib import Path
REPO = Path(__file__).resolve().parents[1]
MAP_PATH = REPO / "data" / "multimodal_doc" / "maps" / "mmlongbench_doc_map.jsonl"
PDF_DIR = REPO / "data" / "multimodal_doc" / "mmlongbench" / "pdfs"

View file

@ -10,26 +10,24 @@ from collections import defaultdict
def main() -> None:
raw_path = sorted(glob.glob("data/research/runs/*/crag/raw.jsonl"))[-1]
print(f"Reading: {raw_path}")
rows = [json.loads(line) for line in open(raw_path, encoding="utf-8") if line.strip()]
with open(raw_path, encoding="utf-8") as fh:
rows = [json.loads(line) for line in fh if line.strip()]
by_q: dict[str, dict[str, dict]] = defaultdict(dict)
for r in rows:
by_q[r["qid"]][r["arm"]] = r
for qid, arms in list(by_q.items()):
b = arms.get("bare_llm", {})
l = arms.get("long_context", {})
lc = arms.get("long_context", {})
s = arms.get("surfsense", {})
print(f"\n=== {qid} ({b.get('domain')}/{b.get('question_type')}) ===")
print(f" question: {b.get('extra', {}).get('question', '?')!r}")
print(f" gold: {b.get('gold')!r}")
for arm_name, a in (("bare_llm", b), ("long_context", l), ("surfsense", s)):
for arm_name, a in (("bare_llm", b), ("long_context", lc), ("surfsense", s)):
grade = a.get("graded", {})
text = (a.get("raw_text") or "").strip()
tail = text[-200:] if text else ""
print(
f" [{arm_name}] grade={grade.get('grade')} "
f"method={grade.get('method')}"
)
print(f" [{arm_name}] grade={grade.get('grade')} method={grade.get('method')}")
print(f" -> {tail!r}")

View file

@ -10,7 +10,8 @@ from collections import defaultdict
def main() -> None:
raw_path = sorted(glob.glob("data/research/runs/*/crag/raw.jsonl"))[-1]
print(f"Reading: {raw_path}")
rows = [json.loads(line) for line in open(raw_path, encoding="utf-8") if line.strip()]
with open(raw_path, encoding="utf-8") as fh:
rows = [json.loads(line) for line in fh if line.strip()]
by_q: dict[str, dict[str, dict]] = defaultdict(dict)
for r in rows:
by_q[r["qid"]][r["arm"]] = r

View file

@ -106,9 +106,7 @@ def _is_failure_row(row: dict[str, Any]) -> bool:
if row.get("error"):
return True
if not (row.get("raw_text") or "").strip():
return True
return False
return bool(not (row.get("raw_text") or "").strip())
@dataclass
@ -134,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
@ -204,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,
@ -216,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,
@ -258,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,
@ -276,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,
)
@ -367,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())),
)
@ -385,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] = {}
@ -415,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()
@ -428,11 +452,14 @@ async def _run(args: argparse.Namespace) -> int:
if f.arm == "native_pdf":
pdf_path = Path(map_row["pdf_path"])
if not pdf_path.exists():
if not await asyncio.to_thread(pdf_path.exists):
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
@ -442,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:
@ -454,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.")
@ -469,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)
@ -491,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 ""),
@ -557,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, [])),
}
@ -595,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)}")
@ -606,27 +640,37 @@ async def _run(args: argparse.Namespace) -> int:
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--run-id", default="2026-05-14T00-53-19Z",
"--run-id",
default="2026-05-14T00-53-19Z",
help="Run timestamp under data/multimodal_doc/runs/. Default is the "
"n=171 production run we wrote up in the blog.",
"n=171 production run we wrote up in the blog.",
)
parser.add_argument("--max-attempts", type=int, default=5)
parser.add_argument("--base-delay", type=float, default=1.0,
help="Base seconds for exponential backoff (default 1s).")
parser.add_argument("--max-delay", type=float, default=30.0,
help="Cap on per-retry sleep (default 30s).")
parser.add_argument("--concurrency", type=int, default=2,
help="Parallel retries in flight (default 2 — keep low "
"to avoid the same transport stress that caused "
"the original failures).")
parser.add_argument(
"--base-delay",
type=float,
default=1.0,
help="Base seconds for exponential backoff (default 1s).",
)
parser.add_argument(
"--max-delay", type=float, default=30.0, help="Cap on per-retry sleep (default 30s)."
)
parser.add_argument(
"--concurrency",
type=int,
default=2,
help="Parallel retries in flight (default 2 — keep low "
"to avoid the same transport stress that caused "
"the original failures).",
)
parser.add_argument("--llm-model", default="anthropic/claude-sonnet-4.5")
parser.add_argument("--pdf-engine", default="native",
choices=[e.value for e in PdfEngine])
parser.add_argument("--pdf-engine", default="native", choices=[e.value for e in PdfEngine])
parser.add_argument("--max-output-tokens", type=int, default=512)
parser.add_argument(
"--include-surfsense", action="store_true",
"--include-surfsense",
action="store_true",
help="Also retry surfsense_agentic failures (requires backend + celery up). "
"Default is to skip them since the n=171 run had 0 SurfSense failures.",
"Default is to skip them since the n=171 run had 0 SurfSense failures.",
)
args = parser.parse_args()
raise SystemExit(asyncio.run(_run(args)))

View file

@ -11,7 +11,8 @@ def main() -> None:
if not runs:
print("(no CRAG runs found)")
return
m = json.load(open(runs[-1], encoding="utf-8"))
with open(runs[-1], encoding="utf-8") as fh:
m = json.load(fh)
metrics = m["metrics"]
print(f"Reading: {runs[-1]}")
@ -22,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()
@ -47,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()
@ -57,7 +58,7 @@ def main() -> None:
pieces = [f"{dom:10s} (n={n:3d}):"]
for arm in ("bare_llm", "long_context", "surfsense"):
if arm in row:
pieces.append(f"{arm}={row[arm]['truthfulness_score']*100:+7.1f}%")
pieces.append(f"{arm}={row[arm]['truthfulness_score'] * 100:+7.1f}%")
print(" ".join(pieces))

View file

@ -16,7 +16,6 @@ import statistics
from collections import defaultdict
from pathlib import Path
REPO = Path(__file__).resolve().parents[1]
RUN_DIR = REPO / "data" / "multimodal_doc" / "runs" / "2026-05-14T00-53-19Z" / "parser_compare"
RAW = RUN_DIR / "raw.jsonl"
@ -24,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)
@ -32,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"]
@ -71,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"]
@ -91,7 +102,7 @@ def main() -> None:
print()
print("by answer_format (accuracy):")
formats = sorted({f for m in arm_metrics.values() for f in m["by_format"].keys()})
formats = sorted({f for m in arm_metrics.values() for f in m["by_format"]})
header = f"{'arm':<25} " + " ".join(f"{f:>10}" for f in formats)
print(header)
print("-" * len(header))
@ -111,7 +122,7 @@ def main() -> None:
print("Aggregated cost (from run_artifact.json):")
for arm, row in per_arm_art.items():
print(
f" {arm:<25} acc={row['accuracy']*100:5.1f}% "
f" {arm:<25} acc={row['accuracy'] * 100:5.1f}% "
f" $/Q LLM={row['llm_cost_per_q']:.4f} "
f" preprocess total=${row['preprocess_cost_total']:.2f} "
f" $/Q total={row['total_cost_per_q']:.4f}"

View file

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

View file

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

View file

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

View file

@ -32,14 +32,13 @@ from __future__ import annotations
import argparse
import asyncio
import contextlib
import json
import logging
import sys
from dataclasses import dataclass
from typing import Any
import sys
import httpx
from rich.console import Console
from rich.table import Table
@ -51,10 +50,8 @@ from rich.table import Table
# Terminal, PowerShell, cmd) all interpret ANSI escapes natively.
if sys.platform == "win32":
for _stream in (sys.stdout, sys.stderr):
try:
with contextlib.suppress(AttributeError, ValueError):
_stream.reconfigure(encoding="utf-8", errors="replace")
except (AttributeError, ValueError):
pass
from . import registry
from .auth import CredentialError, acquire_token, client_with_auth
@ -207,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
@ -295,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
@ -527,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
@ -700,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.")
@ -732,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)
@ -743,7 +742,9 @@ def _build_parser() -> argparse.ArgumentParser:
suite_parser = run_sub.add_parser(suite, help=f"Run a {suite} benchmark.")
suite_bench = suite_parser.add_subparsers(dest="benchmark", required=True)
for benchmark in registry.list_benchmarks(suite):
bp = suite_bench.add_parser(benchmark.name, help=getattr(benchmark, "description", benchmark.name))
bp = suite_bench.add_parser(
benchmark.name, help=getattr(benchmark, "description", benchmark.name)
)
if hasattr(benchmark, "add_run_args"):
benchmark.add_run_args(bp)
bp.set_defaults(_func=_cmd_run, _async=True)

View file

@ -18,6 +18,7 @@ Document processing is asynchronous:
from __future__ import annotations
import asyncio
import contextlib
import logging
import mimetypes
from collections.abc import Iterable, Sequence
@ -83,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)
@ -157,10 +157,8 @@ class DocumentsClient:
)
finally:
for _, (_, file_obj, _) in opened:
try:
with contextlib.suppress(Exception):
file_obj.close()
except Exception: # noqa: BLE001
pass
response.raise_for_status()
return FileUploadResult.from_payload(response.json())
@ -241,9 +239,7 @@ class DocumentsClient:
# chunks (chunk_id -> document_id map)
# ------------------------------------------------------------------
async def list_chunks(
self, document_id: int, *, page_size: int = 100
) -> list[ChunkRow]:
async def list_chunks(self, document_id: int, *, page_size: int = 100) -> list[ChunkRow]:
"""Walk ``GET /documents/{id}/chunks`` until ``has_more=False``.
Used by ingestion to materialise the ``chunk_id -> document_id``

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -3,7 +3,7 @@
from __future__ import annotations
from .answer_letter import AnswerLetterResult, extract_answer_letter
from .citations import CITATION_REGEX, CitationToken, ChunkCitation, UrlCitation, parse_citations
from .citations import CITATION_REGEX, ChunkCitation, CitationToken, UrlCitation, parse_citations
from .freeform_answer import extract_freeform_answer
from .sse import SseEvent, iter_sse_events

View file

@ -15,7 +15,7 @@ from __future__ import annotations
import re
from dataclasses import dataclass
from typing import Any, Union
from typing import Any
# Pattern preserves the TS source verbatim:
# /[\[【]\u200B?citation:\s*(https?:\/\/[^\]】\u200B]+|urlcite\d+|(?:doc-)?-?\d+(?:\s*,\s*(?:doc-)?-?\d+)*)\s*\u200B?[\]】]/g
@ -35,7 +35,7 @@ from typing import Any, Union
# 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+)*"
@ -64,7 +64,7 @@ class UrlCitation:
return {"kind": "url", "url": self.url}
CitationToken = Union[ChunkCitation, UrlCitation]
CitationToken = ChunkCitation | UrlCitation
def parse_citations(text: str, *, url_map: dict[str, str] | None = None) -> list[CitationToken]:

View file

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

View file

@ -25,6 +25,7 @@ import asyncio
import logging
import os
import random
from pathlib import Path
logger = logging.getLogger(__name__)
@ -63,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")
@ -82,10 +82,13 @@ async def parse_with_azure_di(
ServiceResponseError,
)
file_size_mb = os.path.getsize(file_path) / (1024 * 1024)
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
@ -96,21 +99,21 @@ async def parse_with_azure_di(
credential=AzureKeyCredential(api_key),
)
async with client:
with open(file_path, "rb") as fh:
poller = await client.begin_analyze_document(
model_id,
body=fh,
output_content_format=DocumentContentFormat.MARKDOWN,
)
body = await asyncio.to_thread(Path(file_path).read_bytes)
poller = await client.begin_analyze_document(
model_id,
body=body,
output_content_format=DocumentContentFormat.MARKDOWN,
)
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
@ -119,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
@ -132,7 +133,10 @@ async def parse_with_azure_di(
sleep_for = delay + jitter
logger.warning(
"Azure DI attempt %d/%d failed (%s); retrying in %.1fs",
attempt, _MAX_RETRIES, type(last_exc).__name__, sleep_for,
attempt,
_MAX_RETRIES,
type(last_exc).__name__,
sleep_for,
)
await asyncio.sleep(sleep_for)

View file

@ -61,8 +61,7 @@ def _extract_markdown(result) -> str:
if result and hasattr(result[0], "text"):
return result[0].text
return "\n\n".join(
doc.page_content if hasattr(doc, "page_content") else str(doc)
for doc in result
doc.page_content if hasattr(doc, "page_content") else str(doc) for doc in result
)
return str(result)
@ -86,9 +85,7 @@ async def parse_with_llamacloud(
api_key = api_key or os.environ.get("LLAMA_CLOUD_API_KEY")
if not api_key:
raise ValueError(
"LLAMA_CLOUD_API_KEY must be set (see surfsense_evals/.env)."
)
raise ValueError("LLAMA_CLOUD_API_KEY must be set (see surfsense_evals/.env).")
parse_mode = _LLAMA_PARSE_MODE_MAP.get(processing_mode, "parse_page_with_llm")
@ -98,7 +95,7 @@ async def parse_with_llamacloud(
from llama_cloud_services.parse.base import JobFailedException
from llama_cloud_services.parse.utils import ResultType
file_size_mb = os.path.getsize(file_path) / (1024 * 1024)
file_size_mb = await asyncio.to_thread(os.path.getsize, file_path) / (1024 * 1024)
# Match backend's per-page timeout heuristic so big PDFs don't drop
# mid-job: 60s baseline + 30s/page (premium agent runs longer than
# basic; both fit comfortably here).
@ -106,13 +103,19 @@ async def parse_with_llamacloud(
upload_timeout = max(120.0, 30.0 * file_size_mb)
logger.info(
"LlamaCloud parsing %s (mode=%s, parse_mode=%s, %.1fMB, "
"job_timeout=%.0fs)",
file_path, processing_mode, parse_mode, file_size_mb, job_timeout,
"LlamaCloud parsing %s (mode=%s, parse_mode=%s, %.1fMB, job_timeout=%.0fs)",
file_path,
processing_mode,
parse_mode,
file_size_mb,
job_timeout,
)
custom_timeout = httpx.Timeout(
connect=120.0, read=upload_timeout, write=upload_timeout, pool=120.0,
connect=120.0,
read=upload_timeout,
write=upload_timeout,
pool=120.0,
)
last_exc: Exception | None = None
@ -135,12 +138,12 @@ async def parse_with_llamacloud(
result = await parser.aparse(str(file_path))
content = _extract_markdown(result).strip()
if not content:
raise LlamaCloudError(
f"LlamaCloud returned empty content for {file_path}"
)
raise LlamaCloudError(f"LlamaCloud returned empty content for {file_path}")
logger.info(
"LlamaCloud OK: %s (%s) -> %d chars",
file_path, parse_mode, len(content),
file_path,
parse_mode,
len(content),
)
return content
@ -156,7 +159,10 @@ async def parse_with_llamacloud(
sleep_for = delay + jitter
logger.warning(
"LlamaCloud attempt %d/%d failed (%s); retrying in %.1fs",
attempt, _MAX_RETRIES, type(last_exc).__name__, sleep_for,
attempt,
_MAX_RETRIES,
type(last_exc).__name__,
sleep_for,
)
await asyncio.sleep(sleep_for)

View file

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

View file

@ -29,8 +29,8 @@ from typing import Any
import httpx
from .openrouter_pdf import (
OpenRouterResponse,
_DEFAULT_HEADERS,
OpenRouterResponse,
_parse_chat_completion,
)

View file

@ -34,7 +34,7 @@ import base64
import logging
import time
from dataclasses import dataclass
from enum import Enum
from enum import StrEnum
from pathlib import Path
from typing import Any
@ -43,7 +43,7 @@ import httpx
logger = logging.getLogger(__name__)
class PdfEngine(str, Enum):
class PdfEngine(StrEnum):
NATIVE = "native"
MISTRAL_OCR = "mistral-ocr"
CLOUDFLARE_AI = "cloudflare-ai"
@ -121,9 +121,7 @@ class OpenRouterPdfProvider:
body: dict[str, Any] = {
"model": self._model,
"messages": messages,
"plugins": [
{"id": "file-parser", "pdf": {"engine": self._engine.value}}
],
"plugins": [{"id": "file-parser", "pdf": {"engine": self._engine.value}}],
}
if max_tokens:
body["max_tokens"] = max_tokens

View file

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

View file

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

View file

@ -20,7 +20,7 @@ from __future__ import annotations
import importlib
import logging
import pkgutil
from typing import Iterable
from collections.abc import Iterable
logger = logging.getLogger(__name__)
@ -60,7 +60,5 @@ def discover_suites() -> list[str]:
importlib.import_module(benchmark_name)
imported.append(benchmark_name)
except Exception as exc: # noqa: BLE001
logger.warning(
"Failed to import benchmark %s: %s", benchmark_name, exc
)
logger.warning("Failed to import benchmark %s: %s", benchmark_name, exc)
return imported

View file

@ -6,7 +6,6 @@ import argparse
from typing import Any
from ....core.registry import (
Benchmark,
ReportSection,
RunArtifact,
RunContext,

View file

@ -12,7 +12,7 @@ Recall@k / MRR / nDCG@10 against qrels.
from __future__ import annotations
from .runner import CureBenchmark
from ....core import registry as _registry
from .runner import CureBenchmark
_registry.register(CureBenchmark())

View file

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

View file

@ -34,7 +34,6 @@ from ....core.ingest_settings import (
)
from ....core.metrics.retrieval import score_run
from ....core.registry import (
Benchmark,
ReportSection,
RunArtifact,
RunContext,
@ -192,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;
@ -234,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:
@ -276,7 +280,7 @@ class CureBenchmark:
)
per_query_retrieved: dict[str, list[str]] = {}
for q, res in zip(queries, results):
for q, res in zip(queries, results, strict=False):
chunk_ids: list[int] = []
seen: set[int] = set()
for citation in res.citations:
@ -311,7 +315,7 @@ class CureBenchmark:
run_dir = ctx.runs_dir(run_timestamp=run_timestamp)
raw_path = run_dir / "raw.jsonl"
with raw_path.open("w", encoding="utf-8") as fh:
for q, res in zip(queries, results):
for q, res in zip(queries, results, strict=False):
fh.write(
json.dumps(
{

View file

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

View file

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

View file

@ -11,7 +11,7 @@ document — the corpus is millions of biomedical snippets.
from __future__ import annotations
from .runner import MirageBenchmark
from ....core import registry as _registry
from .runner import MirageBenchmark
_registry.register(MirageBenchmark())

View file

@ -48,9 +48,7 @@ from ....core.registry import RunContext
logger = logging.getLogger(__name__)
MIRAGE_BENCHMARK_URL = (
"https://raw.githubusercontent.com/Teddy-XiongGZ/MIRAGE/main/benchmark.json"
)
MIRAGE_BENCHMARK_URL = "https://raw.githubusercontent.com/Teddy-XiongGZ/MIRAGE/main/benchmark.json"
# Upstream only ships ONE zip — top-10k retrievals across 5 retrievers,
# ~16 GB. We default to skipping it (see `--skip-snippet-filter`) and
# ingesting the chosen corpus in full; this URL is only fetched when
@ -93,6 +91,24 @@ class SnippetRow:
# ---------------------------------------------------------------------------
def _reuse_cached_dest(dest: Path, *, expect_zip: bool, label: str) -> Path | None:
"""Return ``dest`` if a usable cache hit, else ``None`` (and delete corrupt zips)."""
if not dest.exists():
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.",
label,
dest,
dest.stat().st_size,
)
dest.unlink(missing_ok=True)
return None
logger.info("Using cached %s at %s", label, dest)
return dest
async def _fetch_to_path(
url: str,
*,
@ -127,19 +143,9 @@ async def _fetch_to_path(
surprise-grabbing 16 GB on a slow link.
"""
if dest.exists():
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.",
label,
dest,
dest.stat().st_size,
)
dest.unlink(missing_ok=True)
else:
logger.info("Using cached %s at %s", label, dest)
return dest
cached = _reuse_cached_dest(dest, expect_zip=expect_zip, label=label)
if cached is not None:
return cached
dest.parent.mkdir(parents=True, exist_ok=True)
partial = dest.with_suffix(dest.suffix + ".partial")
@ -167,42 +173,44 @@ async def _fetch_to_path(
)
try:
async with httpx.AsyncClient(
timeout=httpx.Timeout(timeout_s, connect=20.0),
follow_redirects=True,
) as client:
async with 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.",
label,
)
partial.unlink(missing_ok=True)
existing_bytes = 0
elif response.status_code == 416:
# Range not satisfiable — the .partial is at or
# past the end. Treat as "already downloaded";
# validate by closing and re-opening for atomic
# rename below.
logger.info(
"Server reports %s already complete (HTTP 416).",
label,
)
elif response.status_code not in (200, 206):
response.raise_for_status()
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.",
label,
)
partial.unlink(missing_ok=True)
existing_bytes = 0
elif response.status_code == 416:
# Range not satisfiable — the .partial is at or
# past the end. Treat as "already downloaded";
# validate by closing and re-opening for atomic
# rename below.
logger.info(
"Server reports %s already complete (HTTP 416).",
label,
)
elif response.status_code not in (200, 206):
response.raise_for_status()
total_size = _planned_total_size(response, existing_bytes)
if (
total_size is not None
and total_size > _LARGE_DOWNLOAD_BYTES
and not allow_large_download
):
raise _LargeDownloadAbort(label, total_size)
total_size = _planned_total_size(response, existing_bytes)
if (
total_size is not None
and total_size > _LARGE_DOWNLOAD_BYTES
and not allow_large_download
):
raise _LargeDownloadAbort(label, total_size)
mode = "ab" if existing_bytes else "wb"
with partial.open(mode) as fh:
async for chunk in response.aiter_bytes(chunk_size=1 << 18):
fh.write(chunk)
mode = "ab" if existing_bytes else "wb"
with partial.open(mode) as fh:
async for chunk in response.aiter_bytes(chunk_size=1 << 18):
fh.write(chunk)
# Optional content sanity check before promoting to dest.
if expect_zip and not _is_valid_zip(partial):
raise zipfile.BadZipFile(
@ -215,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,
@ -228,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,
@ -270,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."
@ -312,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).
@ -533,10 +539,7 @@ async def run_ingest(
logger.warning("Failed to list chunks for doc_id=%s: %s", doc_id, exc)
continue
for chunk in chunks:
fh.write(
json.dumps({"chunk_id": chunk.id, "document_id": doc_id})
+ "\n"
)
fh.write(json.dumps({"chunk_id": chunk.id, "document_id": doc_id}) + "\n")
new_state = ctx.suite_state
new_state.ingestion_maps["mirage"] = str(snippet_map_path)

View file

@ -8,7 +8,6 @@ from __future__ import annotations
from collections.abc import Mapping
_PROMPT_TEMPLATE = """\
You are a helpful medical expert. Answer the following multiple-choice
question using the relevant medical knowledge available to you (and any

View file

@ -29,7 +29,6 @@ from ....core.ingest_settings import (
)
from ....core.metrics.mc_accuracy import accuracy_with_wilson_ci, macro_accuracy
from ....core.registry import (
Benchmark,
ReportSection,
RunArtifact,
RunContext,
@ -135,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
@ -153,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
@ -197,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(),
@ -229,7 +239,7 @@ class MirageBenchmark:
run_dir = ctx.runs_dir(run_timestamp=run_timestamp)
raw_path = run_dir / "raw.jsonl"
with raw_path.open("w", encoding="utf-8") as fh:
for q, res in zip(questions, results):
for q, res in zip(questions, results, strict=False):
fh.write(
json.dumps(
{
@ -246,7 +256,7 @@ class MirageBenchmark:
for task in tasks:
n_correct = 0
n_total = 0
for q, res in zip(questions, results):
for q, res in zip(questions, results, strict=False):
if q.task != task:
continue
n_total += 1
@ -256,7 +266,10 @@ class MirageBenchmark:
per_task_acc[task] = acc.to_dict()
macro = macro_accuracy(
{t: accuracy_with_wilson_ci(d["n_correct"], d["n_total"]) for t, d in per_task_acc.items()}
{
t: accuracy_with_wilson_ci(d["n_correct"], d["n_total"])
for t, d in per_task_acc.items()
}
)
metrics = {"per_task": per_task_acc, "macro_accuracy": macro}

View file

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

View file

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

View file

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

View file

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

View file

@ -41,8 +41,6 @@ from typing import Any
from ....core.config import set_suite_state
from ....core.parsers import (
AzureDIError,
LlamaCloudError,
count_pdf_pages,
parse_with_azure_di,
parse_with_llamacloud,
@ -55,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"),
)
@ -100,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()},
}
@ -126,12 +122,14 @@ 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}")
out_path.parent.mkdir(parents=True, exist_ok=True)
out_path.write_text(markdown, encoding="utf-8")
await asyncio.to_thread(out_path.write_text, markdown, encoding="utf-8")
return markdown, time.monotonic() - started
@ -170,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,
)
@ -192,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,
@ -290,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"
@ -319,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))
@ -339,12 +347,13 @@ async def run_ingest(
# Quick summary log
total_extractions = sum(len(mr.extractions) for mr in manifest_rows)
failures = sum(
1 for mr in manifest_rows for ext in mr.extractions.values()
if ext.status != "ok"
1 for mr in manifest_rows for ext in mr.extractions.values() if ext.status != "ok"
)
logger.info(
"parser_compare ingest done: %d PDFs, %d extractions, %d failures",
len(manifest_rows), total_extractions, failures,
len(manifest_rows),
total_extractions,
failures,
)

View file

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

View file

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

View file

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

View file

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

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