From e547654518f84febf5e1c41d5a524c37e25334a3 Mon Sep 17 00:00:00 2001 From: pkuwkl Date: Sat, 18 Jul 2026 15:23:44 +0800 Subject: [PATCH 1/4] refactor(paper): add source-first Paper Flow V1 --- .../references/develop-components.md | 10 +- .../references/develop-components.md | 10 +- .github/workflows/e2e.yml | 22 +- README.md | 39 +- contexts/design/README.md | 3 +- contexts/design/flow/paper.md | 424 ++------ contexts/design/knowledge/paper.md | 93 ++ contexts/design/library/local.md | 203 ++-- contexts/design/operations/naming.md | 2 +- contexts/design/preprocess/pdf.md | 10 +- contexts/design/rag/document.md | 60 +- contexts/usage/README.md | 2 +- docs/README.md | 13 +- docs/library.md | 126 ++- examples/flows/paper.py | 146 +++ examples/library/README.md | 4 +- examples/library/data/ai_infrastructure.db | Bin 98304 -> 176128 bytes quantmind/configs/paper.py | 46 +- quantmind/flows/__init__.py | 11 +- quantmind/flows/_paper_summary.py | 243 +++++ quantmind/flows/paper.py | 603 ++++++++--- quantmind/knowledge/__init__.py | 51 +- quantmind/knowledge/_base.py | 30 +- quantmind/knowledge/_flatten.py | 10 +- quantmind/knowledge/_tree.py | 14 +- quantmind/knowledge/earnings.py | 4 - quantmind/knowledge/factor.py | 4 - quantmind/knowledge/news.py | 4 - quantmind/knowledge/paper.py | 604 ++++++++++- quantmind/knowledge/thesis.py | 3 - quantmind/library/__init__.py | 13 +- .../library/_internal/llamaindex_retriever.py | 15 + .../library/_internal/retrieval_targets.py | 110 +- quantmind/library/_internal/sqlite_store.py | 975 +++++++++++++++++- quantmind/library/_types.py | 18 +- quantmind/library/local.py | 279 ++++- quantmind/preprocess/fetch/_types.py | 1 + quantmind/preprocess/fetch/arxiv.py | 10 +- quantmind/rag/document.py | 22 +- .../build_ai_infrastructure_bundle.py | 4 +- scripts/verify_pdf_rag_e2e.py | 261 ++++- tests/configs/test_paper.py | 16 +- tests/flows/test_paper.py | 573 +++++----- tests/knowledge/test_base.py | 33 +- tests/knowledge/test_earnings.py | 7 +- tests/knowledge/test_factor.py | 8 +- tests/knowledge/test_news.py | 7 +- tests/knowledge/test_paper.py | 245 +++-- tests/knowledge/test_roundtrip.py | 42 +- tests/knowledge/test_thesis.py | 4 +- tests/knowledge/test_tree.py | 8 +- tests/library/test_example_bundle.py | 10 +- tests/library/test_local.py | 49 +- tests/library/test_paper.py | 408 ++++++++ tests/library/test_types.py | 7 +- tests/paper_helpers.py | 200 ++++ tests/preprocess/fetch/test_arxiv.py | 5 + tests/rag/test_document.py | 11 + tests/test_verify_pdf_rag_e2e.py | 162 ++- 59 files changed, 4859 insertions(+), 1438 deletions(-) create mode 100644 contexts/design/knowledge/paper.md create mode 100644 examples/flows/paper.py create mode 100644 quantmind/flows/_paper_summary.py create mode 100644 tests/library/test_paper.py create mode 100644 tests/paper_helpers.py diff --git a/.agents/skills/quantmind-dev/references/develop-components.md b/.agents/skills/quantmind-dev/references/develop-components.md index 7170e47..239e3be 100644 --- a/.agents/skills/quantmind-dev/references/develop-components.md +++ b/.agents/skills/quantmind-dev/references/develop-components.md @@ -43,11 +43,13 @@ apply throughout. - Pydantic models, `frozen=True`, `extra="forbid"`. - Every `BaseKnowledge` subclass **must** require `as_of: datetime` (financial time-sensitivity is mandatory) and a typed `source: SourceRef` - (no bare strings), and **must** override `embedding_text()`. + (no bare strings). +- Canonical models do not select retrieval text or store vectors. Put + rebuildable text projections in `quantmind.library`. - Pick one shape: `FlattenKnowledge` (atomic card), `TreeKnowledge` - (hierarchical artifact), or `GraphKnowledge` (placeholder). Whole-document - objects are `TreeKnowledge` even when a flatten card exists alongside - (e.g. `Paper` vs `PaperKnowledgeCard`). + (hierarchical artifact), or `GraphKnowledge` (placeholder) for conventional + knowledge. Source-first paper revisions and independently versioned paper + artifacts use their dedicated frozen models instead of `BaseKnowledge`. ### `quantmind/configs/` — operation cfg + typed inputs diff --git a/.claude/skills/quantmind-dev/references/develop-components.md b/.claude/skills/quantmind-dev/references/develop-components.md index 7170e47..239e3be 100644 --- a/.claude/skills/quantmind-dev/references/develop-components.md +++ b/.claude/skills/quantmind-dev/references/develop-components.md @@ -43,11 +43,13 @@ apply throughout. - Pydantic models, `frozen=True`, `extra="forbid"`. - Every `BaseKnowledge` subclass **must** require `as_of: datetime` (financial time-sensitivity is mandatory) and a typed `source: SourceRef` - (no bare strings), and **must** override `embedding_text()`. + (no bare strings). +- Canonical models do not select retrieval text or store vectors. Put + rebuildable text projections in `quantmind.library`. - Pick one shape: `FlattenKnowledge` (atomic card), `TreeKnowledge` - (hierarchical artifact), or `GraphKnowledge` (placeholder). Whole-document - objects are `TreeKnowledge` even when a flatten card exists alongside - (e.g. `Paper` vs `PaperKnowledgeCard`). + (hierarchical artifact), or `GraphKnowledge` (placeholder) for conventional + knowledge. Source-first paper revisions and independently versioned paper + artifacts use their dedicated frozen models instead of `BaseKnowledge`. ### `quantmind/configs/` — operation cfg + typed inputs diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index 0b01d7b..7376606 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -19,6 +19,11 @@ on: - 'quantmind/preprocess/fetch/arxiv.py' - 'quantmind/preprocess/format/pdf.py' - 'quantmind/rag/**' + - 'quantmind/configs/paper.py' + - 'quantmind/flows/_paper_summary.py' + - 'quantmind/flows/paper.py' + - 'quantmind/knowledge/paper.py' + - 'quantmind/library/**' - 'scripts/verify_pdf_rag_e2e.py' - 'pyproject.toml' schedule: @@ -38,7 +43,7 @@ jobs: runs-on: ubuntu-latest outputs: news: ${{ steps.filter.outputs.news }} - pdf_rag: ${{ steps.filter.outputs.pdf_rag }} + paper_flow: ${{ steps.filter.outputs.paper_flow }} steps: - name: Checkout @@ -60,9 +65,14 @@ jobs: - 'quantmind/preprocess/format/html.py' - 'quantmind/preprocess/fetch/**' - 'pyproject.toml' - pdf_rag: + paper_flow: - '.github/workflows/e2e.yml' - 'scripts/verify_pdf_rag_e2e.py' + - 'quantmind/configs/paper.py' + - 'quantmind/flows/_paper_summary.py' + - 'quantmind/flows/paper.py' + - 'quantmind/knowledge/paper.py' + - 'quantmind/library/**' - 'quantmind/preprocess/fetch/arxiv.py' - 'quantmind/preprocess/format/pdf.py' - 'quantmind/rag/**' @@ -98,11 +108,13 @@ jobs: - name: Run live news E2E run: .venv/bin/python scripts/verify_news_e2e.py - pdf-rag: + paper-flow: needs: changes - if: github.event_name != 'pull_request' || needs.changes.outputs.pdf_rag == 'true' + if: github.event_name != 'pull_request' || needs.changes.outputs.paper_flow == 'true' runs-on: ubuntu-latest timeout-minutes: 10 + env: + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} steps: - name: Checkout @@ -125,5 +137,5 @@ jobs: - name: Install project runtime dependencies run: uv pip install --python .venv/bin/python -e . - - name: Run live PDF RAG E2E + - name: Run live Paper Flow V1 E2E run: .venv/bin/python scripts/verify_pdf_rag_e2e.py diff --git a/README.md b/README.md index 25e3e51..3734012 100644 --- a/README.md +++ b/README.md @@ -156,7 +156,7 @@ We use [uv](https://github.com/astral-sh/uv) for fast and reliable Python packag Component-specific guides and architecture notes live under [`docs/`](docs/). -#### Run a single paper through `paper_flow` +#### Build an exact paper source, chunks, and cited summary ```python import asyncio @@ -167,11 +167,12 @@ from quantmind.flows import paper_flow async def main() -> None: - paper = await paper_flow( - ArxivIdentifier(id="2401.12345"), + result = await paper_flow( + ArxivIdentifier(id="1706.03762v7"), cfg=PaperFlowCfg(model="gpt-4o-mini"), ) - print(paper.model_dump_json(indent=2)) + print(result.global_summary.summary) + print(result.source_revision.id, result.chunk_set.id) asyncio.run(main()) @@ -219,17 +220,17 @@ async def main() -> None: "Pull arXiv 2401.12345 about cross-sectional momentum; use gpt-4o-mini.", target_flow=paper_flow, ) - paper = await paper_flow(inp, cfg=cfg) - print(paper.model_dump_json(indent=2)) + result = await paper_flow(inp, cfg=cfg) + print(result.global_summary.summary) asyncio.run(main()) ``` -> **Note**: QuantMind is mid-migration to OpenAI Agents SDK -> (see [#71](https://github.com/LLMQuant/quant-mind/issues/71)). PR5 lands the -> apex layer (`flows/` + `magic.py`); the remaining work is the `mind/` -> memory + store layer scheduled for PR6 and PR7. +Paper Flow V1 is source-first: code preserves the exact PDF revision and +page-aware chunks before accepting a bounded, cited model summary. See the +[complete persist/reopen/search example](examples/flows/paper.py) and +[design contract](contexts/design/flow/paper.md). --- @@ -238,7 +239,7 @@ asyncio.run(main()) - [x] Better `flow` design for user-friendly usage - [x] First production level example (Quant Paper Agent) - [ ] Migrate Agent layer to OpenAI Agents SDK -- [ ] Standardize knowledge format with `knowledge/` (Pydantic-based) +- [x] Standardize knowledge format with `knowledge/` (Pydantic-based) - [ ] Additional content sources (financial news, blogs, reports) - [ ] Cross-step working memory (`mind/memory`) for batch document processing @@ -254,18 +255,10 @@ QuantMind is designed with a larger vision: to become a comprehensive intelligen The foundation we're building today—starting with papers—will expand to encompass the entire financial information ecosystem. > [!NOTE] -> **Future Conceptual Example (PR6 brings `FilesystemMemory`):** -> -> ```python -> from quantmind.configs.paper import ArxivIdentifier -> from quantmind.flows import paper_flow -> from quantmind.knowledge import Paper -> from quantmind.mind.memory import FilesystemMemory # PR6 -> -> memory = FilesystemMemory("./mem/factor-research/") -> for arxiv_id in arxiv_ids: -> paper: Paper = await paper_flow(ArxivIdentifier(id=arxiv_id), memory=memory) -> ``` +> The current source-first paper path produces independently versioned source, +> chunk-set, and cited-summary artifacts. Future agent memory and cross-document +> reasoning can build on `LocalKnowledgeLibrary.search()` without changing +> those canonical artifacts. This future state represents our commitment to moving beyond simple data aggregation and toward genuine machine intelligence in the financial domain. diff --git a/contexts/design/README.md b/contexts/design/README.md index ef3e394..188079a 100644 --- a/contexts/design/README.md +++ b/contexts/design/README.md @@ -24,7 +24,8 @@ implementation must preserve. | Domain | Design | |---|---| -| Flow | [Paper extraction from input to validated result](flow/paper.md) | +| Flow | [Source-first paper flow](flow/paper.md) | +| Knowledge | [Paper sources, artifacts, citations, and locators](knowledge/paper.md) | | Flow | [News collection](flow/news.md) | | Preprocess | [Page-aware multimodal PDF parsing](preprocess/pdf.md) | | RAG | [Page-aware document chunking and retrieval](rag/document.md) | diff --git a/contexts/design/flow/paper.md b/contexts/design/flow/paper.md index de1b509..4894810 100644 --- a/contexts/design/flow/paper.md +++ b/contexts/design/flow/paper.md @@ -1,360 +1,150 @@ -# Paper Extraction: End-to-End Design +# Build a source-first paper result ## Quick Summary -- **Purpose**: Define how a paper input becomes a validated `Paper`. -- **Read when**: Changing paper inputs, parsing, section trees, source tracking, page ranges, or future PageIndex support. -- **Status**: Mixed. Page-aware PDF parsing and document RAG are implemented; [Current Gaps](#current-gaps) lists the remaining paper assembly work. -- **Core rule**: A model or PageIndex may suggest a section tree. Code creates the final IDs, links, order, page ranges, citations, and source-backed text. -- **Page numbering**: PDF page ranges start at 1 and include both the first and last page. +- **Purpose**: Define how one exact PDF revision becomes a durable page-aware chunk set and a cited global summary. +- **Read when**: Changing `paper_flow`, paper inputs, summarization limits, citation validation, or the end-to-end paper verifier. +- **Status**: Implemented by `quantmind.flows.paper_flow` for PDF-backed arXiv, HTTP, and local inputs. +- **Core rule**: Preserve and validate the source revision and chunk set before any model-generated summary is accepted. +- **Canonical models**: [Paper source and artifact design](../knowledge/paper.md). ## Contents -- [Overview](#overview) -- [Who Owns Each Step](#who-owns-each-step) -- [Supported Inputs and Fetching](#supported-inputs-and-fetching) -- [Keep PDF Pages Separate](#keep-pdf-pages-separate) -- [Which Source Provides Each Field](#which-source-provides-each-field) -- [Model Output Is a Draft](#model-output-is-a-draft) -- [Build and Validate the Final Paper](#build-and-validate-the-final-paper) -- [Rules for a Valid Paper](#rules-for-a-valid-paper) -- [Get Node Content from Source Pages](#get-node-content-from-source-pages) -- [Retries and Failures by Step](#retries-and-failures-by-step) -- [What Paper Extraction Does Not Do](#what-paper-extraction-does-not-do) -- [How PageIndex Can Fit Later](#how-pageindex-can-fit-later) -- [Fixed Paper Test Data](#fixed-paper-test-data) -- [Current Gaps](#current-gaps) - -## Overview - -Paper extraction has six steps and returns one validated result: - -```mermaid -flowchart LR - input["Paper input"] --> resolve["Resolve and fetch"] - resolve --> source["Keep pages separate
and record source facts"] - source --> draft["Suggest a section tree"] - draft --> assembly["Build the final Paper in code"] - assembly --> validate["Validate the tree and source links"] - validate --> paper["Validated Paper"] -``` +- [Contract](#contract) +- [Execution Order](#execution-order) +- [Source Revision](#source-revision) +- [Chunk Set](#chunk-set) +- [Cited Global Summary](#cited-global-summary) +- [Bounded Model Calls](#bounded-model-calls) +- [Failure Semantics](#failure-semantics) +- [Persistence and Retrieval](#persistence-and-retrieval) +- [Verification Slice](#verification-slice) +- [Out of Scope](#out-of-scope) + +## Contract -Code, not the model, records source facts, creates IDs and tree links, copies -text from source pages, creates citations, and validates the result. A model or -future PageIndex integration may suggest sections, titles, summaries, and page -ranges, but that output remains a draft. +`paper_flow(input, *, cfg)` returns one validated `PaperFlowResult`: -Paper extraction does not require PageIndex. It first preserves ordered source -pages and passes them through a small `PaperSourceDocument -> draft` interface. -PageIndex can later implement that interface without changing the `Paper` type -or creating the final IDs. +```text +PaperFlowResult +├── source_revision: PaperSourceRevision +├── chunk_set: PaperChunkSet +└── global_summary: PaperGlobalSummary +``` -## Who Owns Each Step +The result is source-first. It does not return a model-authored paper tree, and V1 does not define `PaperTree`. The source revision is an immutable anchor for independently versioned artifacts. A different splitter configuration may produce another chunk set for the same source; a different model, prompt, input chunk set, or output bound may produce another summary. These versions coexist instead of overwriting each other. -This operation creates one extraction result. It does not store papers, search -across them, or write answers. +V1 accepts: -| Work | Owner | +| Input | Behavior | |---|---| -| Resolve identifiers, fetch bytes, parse pages, and hash source content | `quantmind.preprocess` | -| Chunk or retrieve page-aware document evidence when requested | `quantmind.rag` | -| Configure the operation and select the input variant | `quantmind.configs` | -| Suggest a section tree and build the final `Paper` | `quantmind.flows` | -| Define the `Paper`, `TreeKnowledge`, `TreeNode`, source, citation, and extraction models | `quantmind.knowledge` | -| Store knowledge and make it searchable by meaning | `quantmind.library` | -| Navigate a tree and write an answer | A future `quantmind.mind` consumer or agent application | +| `ArxivIdentifier` | Resolve and fetch the exact PDF revision. The canonical source records a versioned arXiv ID such as `1706.03762v7`. | +| `HttpUrl` | Fetch the URL and require PDF content. | +| `LocalFilePath` | Read the local file and require PDF content. | +| `RawText` | Reject because it has no physical page evidence. | +| `DoiIdentifier` | Reject until an exact open-PDF resolver exists. | -The `flow/` directory groups this work because it spans several packages. It -does not require a new `*_flow` public API and does not override the -[public operation naming rules](../operations/naming.md). +## Execution Order -## Supported Inputs and Fetching +The operation has a strict order: -### Supported inputs +1. Resolve the input to exact bytes and source metadata. +2. Parse the PDF into ordered physical pages, blocks, screenshots, and extracted images. +3. Build and validate `PaperSourceRevision`, including content-addressed asset references and blobs. +4. Chunk each page with LlamaIndex while retaining page and character spans. +5. Build and validate `PaperChunkSet` with code-owned IDs and membership. +6. Invoke the bounded summarizer using only the validated chunk set. +7. Resolve model-returned chunk/page coordinates into canonical citations in code. +8. Build and validate `PaperGlobalSummary` and the cross-artifact `PaperFlowResult`. -The planned pipeline accepts the existing `PaperInput` types with the following -behavior: +A summarization failure occurs after source and chunks exist in memory, but `paper_flow` returns no partial success value. Persistence is a separate explicit operation. -| Input | What happens | -|---|---| -| `ArxivIdentifier` | Resolve the identifier to an exact arXiv version, provider metadata, and PDF bytes. Preserve the resolved URL and the time that version became available. | -| `HttpUrl` | Follow a limited number of redirects and accept supported PDF, HTML, Markdown, or plain-text content. Record the final URL and exact fetched content. | -| `LocalFilePath` | Read a supported PDF, HTML, Markdown, or plain-text file. The caller owns file retention; extraction records a local source reference and content hash. | -| `RawText` | Treat the supplied text as one source document without pages. It can produce a `Paper`, but it cannot claim PDF page ranges. | +## Source Revision -`DoiIdentifier` remains unsupported until an open-access resolver can fetch the -exact paper content. A DOI landing page alone contains metadata and may not -contain the paper itself. +`PaperSourceRevision` owns the exact fetched bytes and deterministic parser result. Its ID is derived from the source SHA-256 hash. It records: -### Explicitly unsupported inputs +- typed source metadata, financial `as_of`, `available_at`, publication time, exact arXiv revision, title, and authors; +- parser name, parser version, cleanup version, and every 1-based physical page; +- page text, source blocks, bounding boxes, screenshot references, and extracted-image references; +- content-addressed references for the raw PDF and every retained visual asset; +- exact asset bytes while crossing from the flow to persistence. -V1 does not claim support for: +The raw asset hash must equal the source hash. Every asset ID is derived from source revision, asset kind, page, and content hash. Page references must resolve to assets on that page. Canonical JSON excludes blob bytes; [`LocalKnowledgeLibrary.put_paper()`](../library/local.md) stores them in linked SQLite rows. -- password-protected or corrupt PDFs; -- image-only scans that require OCR; -- authenticated, paywalled, or dynamically rendered sources without a - separately authorized fetcher; -- unsupported binary or media content types; -- a DOI that cannot be resolved to accessible paper content; -- a source whose exact fetched content cannot be identified. +For arXiv, the source must record an exact version suffix. Fetch time is never used as publication time. `available_at` uses exact revision metadata when supplied and otherwise falls back to the earliest reliable observation supported by the input. -Unsupported input fails before a model or tree builder runs. The pipeline must -not send an error page, login page, or unresolved metadata page to a model and -call the result a successfully extracted paper. +## Chunk Set -### Fetching rules +`PaperChunkSet` is built before summarization. Its producer identity contains the LlamaIndex sentence-splitter version, chunk size, and overlap. The producer configuration hash and source revision deterministically identify the artifact. -Fetching produces one exact source version. Resolve redirects, arXiv revisions, -and the content selected by the server before parsing. The final URI and -SHA-256 hash identify the exact bytes used. A retry may receive changed -content; when the hash changes, treat it as a new source version rather than -silently merging it with the earlier attempt. +Each `PaperChunk` has a stable code-owned ID, contiguous position, exact text hash, and one or more `PaperSourceSpan` values. A span retains its 1-based page number, page-local character offsets, block boxes, and visual asset IDs. The chunk-set content hash covers ordered membership, content hashes, and source spans. -## Keep PDF Pages Separate +Chunk IDs, artifact IDs, hashes, positions, and membership are not model outputs. Repeating the same source and configuration produces the same source, chunk-set, and chunk IDs. -PDF parsing must preserve ordered pages before any tree builder runs. The next -step receives data shaped like this: +## Cited Global Summary -```text -PaperSourceDocument -├── source details and exact content hash -├── page or character range format -└── ordered pages - ├── page number - ├── extracted text, including an empty string when the page has no text - └── optional parser-provided layout or outline hints -``` +The model returns only a draft containing summary prose and citation coordinates: chunk index, page number, and an optional quote. It never chooses canonical UUIDs, artifact lineage, source facts, or timestamps. -Required properties: - -- PDF page ranges in the fixed test PDF, model draft, and citations start at 1 - and include both the first and last page. -- Every physical page remains represented and in order. Empty pages are not - dropped because doing so would renumber later page references. -- Text cleanup may remove parser noise, but it must not erase the page - boundary or change which page owns an anchor. -- Record the exact source hash, parser name and version, and cleanup version - with the extraction run. -- HTML, Markdown, plain text, and `RawText` have no pages. They retain source - text and character positions but do not invent PDF page numbers. -- A page-based tree builder, including a future PageIndex integration, accepts - only a source document whose range unit is `pdf_page`. - -For PDF inputs, [`ParsedDocument`](../preprocess/pdf.md) is the implemented -deterministic source value. It preserves blocks, coordinates, and optional page -screenshots as well as page text. `PaperSourceDocument` is the flow's temporary -view over that value, not another public knowledge model. The caller keeps raw -PDF, HTML, and screenshot files; this design does not embed them inside -`Paper`. - -## Which Source Provides Each Field - -Facts from the fetched source take priority over model suggestions. A model may -fill a missing summary or section title, but it must not overwrite a URL, -content hash, publication time, or other known source fact. - -| Field or fact | Who sets it and where it is stored | -|---|---| -| Resolved source URI and source kind | Resolver/fetcher; `SourceRef.kind` and `SourceRef.uri` | -| Exact fetched bytes and content hash | Fetcher; `SourceRef.content_hash`, while the caller keeps the raw file | -| Fetch time | Fetcher clock; `SourceRef.fetched_at` | -| Publication/version time | Provider metadata for the exact version; do not replace it with the first-version date | -| Source availability time | Provider metadata for the exact version; `Paper.available_at`. Use fetch time as the latest possible availability time only when the publication time is unknown. | -| Latest date covered by the paper | A study or data cutoff stated by the source; otherwise use the exact version's publication time for `Paper.as_of`, never fetch time | -| Authors and title | Provider or document metadata; `Paper.authors` and root title unless that metadata is missing or clearly invalid | -| Extraction model, operation, run, and time | Runtime; `ExtractionRef` | -| Parser and cleanup versions | Runtime; keep them with the extraction run until `Paper` has a dedicated field | -| Summaries, section titles, methodology, findings, and limitations | Model or tree builder; validate them before building the final `Paper` | -| Final item/node IDs and tree links | Code only | - -For an arXiv revision, `available_at` refers to the exact revision used, not the -first submission date. `as_of` is the latest date covered by the paper, not the -time it was fetched. Unknown publication or availability data remains unknown -rather than being guessed by a model. - -## Model Output Is a Draft - -A model or PageIndex integration returns a draft with only the information -needed to suggest a section tree. It may contain: - -- temporary node keys used only by that integration; -- candidate titles and summaries; -- an ordered parent/child outline using those temporary keys; -- suggested page ranges using the source document's range format; -- confidence values or errors needed to accept, repair, or reject the draft. - -A draft must not set the final `Paper.id`, `TreeNode.node_id`, -`parent_id`/`children_ids`, copied source text, known source facts, or a -object tied to one provider and exposed through the public result. - -The interface is intentionally small: one `PaperSourceDocument` produces one -draft. The default implementation can use an LLM; a future PageIndex -integration can return the same draft shape. Neither defines the public -`Paper` type. - -## Build and Validate the Final Paper - -Code treats the accepted draft as untrusted input and performs these steps: - -1. Generate the final item and node IDs. Temporary integration IDs never - appear in the returned result. -2. Set `parent_id` and ordered `children_ids` together for each parent-child - link so they cannot disagree. -3. Assign sibling positions from the accepted order. -4. Check page ranges against the preserved source pages. -5. Copy source content from the stated page range and build citations from the - same pages. Model-written paraphrases never become source text or quotes. -6. Apply known source facts and extraction run details. -7. Construct the final `Paper` and run every validation rule below. - -Code owns the final IDs, but separate extraction runs do not need to create the -same IDs. Merging results across runs is a separate decision. - -## Rules for a Valid Paper - -A successful `Paper` satisfies all of the following before it is returned: - -### IDs and root - -- `root_node_id` identifies exactly one entry in `nodes`. -- Every dictionary key equals that node's `node_id`. -- The root has `parent_id=None`; every other node has exactly one parent. -- All item and node IDs are created by code and unique within the result. - -### Parent and child links - -- Every referenced parent and child exists. -- A parent lists each child, and each child points back to that parent. -- A child appears at most once in a parent's `children_ids`. -- The same accepted draft produces the same sibling order, and sibling - positions are unique within one parent. - -### Reachability and safe traversal - -- Every node is reachable from the root. -- The graph has no cycles and no node points to itself. -- No node is shared by multiple parents. -- Depth-first walk and root-to-node path lookup finish for every node. Looking - up an unknown node returns the documented safe result. - -### Source page ranges and citations - -- A PDF page range uses `pdf_page`, starts at 1 or later, has - `start_page <= end_page`, and ends within the preserved PDF page count. -- Citation page ranges also start at 1, include both ends, and stay within the - PDF page count. -- Citation quotes and node content come from the identified source range. -- Inputs without pages do not carry made-up PDF page ranges. -- Sibling page ranges may overlap. A child range is not required to fit - completely inside its parent range unless a future rule adds that - requirement. - -## Get Node Content from Source Pages - -Branch nodes may retain source content. `content=None` remains useful for a -navigation-only node, but being a branch is not itself a reason to discard its -source text. - -When content is present, code copies it from the node's 1-based inclusive page -range. Page-level ranges can include a heading or paragraph also used by an -adjacent node; this is expected when ranges overlap. Do not ask the model to -reproduce source text or create a parent node by joining model summaries. - -Choosing a section and loading its source text remain separate operations. An -agent can read titles and summaries to select a node, then fetch that node's -page range. A future PageIndex integration may help build the outline, but it -does not fetch the source text. +Code accepts the draft only when: -## Retries and Failures by Step +- every chunk index exists; +- every cited page is present in that chunk's source spans; +- every supplied quote occurs verbatim in the cited chunk; +- citation count meets `summary_min_citations`; +- distinct cited-page count meets `summary_min_pages`. -| Step | Same input, same result? | Retry and failure behavior | -|---|---|---| -| Input validation and fetching rules | Yes, for one configuration | Invalid or unsupported input fails immediately. | -| Network fetch | No; the remote source may change | Retry only a limited number of temporary failures. Record and hash the exact successful response. Permanent HTTP or content failures stop the operation. | -| PDF or text parsing | Yes, for the same bytes and parser version | Parser failure stops before the model or tree builder runs. Do not skip corrupt pages or renumber later pages. | -| Model or PageIndex draft | No when a model is used | Limited retries may repair network or invalid-draft failures against the same source version. Record every attempt. | -| Build the final `Paper` | Yes for one accepted draft and source version, except for new UUID values | Invalid links, page ranges, or conflicts with source facts reject the draft. | -| Final validation | Yes | Any failure rejects the entire result; never return a partial `Paper` as success. | +Code then creates `PaperCitation` values and a `PaperGlobalSummary`. Its producer identity includes model, prompt version, input chunk-set ID, instructions hash, and maximum output tokens. Its lineage contains the exact input chunk-set locator. Summary citations must resolve through that chunk set to source pages. -A successful result means one `Paper` passed its source, page-range, and tree -checks. Fetching bytes, receiving a model response, or constructing a Pydantic -object alone is not success. A failure records enough about the failed step and -source to diagnose or retry it, but it does not store a partial `Paper`. +## Bounded Model Calls -## What Paper Extraction Does Not Do +`PaperFlowCfg` makes runtime and usage bounds explicit: -Paper extraction returns one validated `Paper`. Other components handle: +- `max_summary_tool_calls` caps adaptive chunk reads; +- `max_summary_concurrency` bounds simultaneous chunk-read tools; +- `max_summary_input_tokens` caps estimated manifest and chunk-read input; +- `max_summary_output_tokens` caps the structured response and validated summary size; +- `max_turns` caps Agents SDK turns and defaults to 16 for the paper summarizer; +- `timeout_seconds` bounds the complete summarization operation; +- `summary_prompt_version` and `summary_instructions` version the semantic producer. -- storing and safely updating it in `LocalKnowledgeLibrary`; -- generating embeddings and building search records; -- searching across a collection; -- PageIndex-style tree navigation within one selected document; -- writing answers, managing conversation state, and citing a final response; -- deciding whether to keep raw source bytes. +The summarizer receives a bounded manifest and may adaptively read up to eight consecutive chunks per tool call. A shared concurrency-safe budget is reserved before every read. The operation rejects work that would exceed a call or token limit rather than silently running without a bound. -Keeping this work separate allows Paper extraction to be implemented before -PageIndex. PageIndex can then be added without replacing collection-wide search -or stored knowledge. +## Failure Semantics -## How PageIndex Can Fit Later - -Future integration must preserve these decisions: +- Non-PDF input raises `UnsupportedContentTypeError`. +- Unresolved DOI input raises `NotImplementedError`. +- Fetching, parsing, or missing parser assets raise their source error and produce no result. +- Empty chunk output is invalid. +- Invalid or insufficient summary citations raise `PaperCitationValidationError`. +- Call, token, output, concurrency, and timeout violations fail the summary operation. +- Any canonical identity, content hash, membership, lineage, or cross-artifact mismatch fails Pydantic validation. -1. PageIndex receives ordered source pages before they are joined into one text. -2. Its node IDs remain temporary. Code creates the final IDs and tree links. -3. It may propose titles, summaries, ordering, and 1-based inclusive page - ranges through the draft described above. -4. It does not become the `Paper`, `TreeKnowledge`, or `TreeNode` type. -5. Navigation uses titles and summaries; loading source text separately fetches - the selected page range. -6. Sibling page ranges may overlap, and a child range does not need to fit - completely inside its parent range. +No failure is converted into a partially valid `PaperFlowResult`. Callers may retry with the same source and producer settings; stable IDs make successful repeated runs idempotent. -A PageIndex adapter belongs with other opinionated document retrieval in -[`quantmind.rag`](../rag/document.md). It still returns the limited draft above; -it does not become the canonical tree or a generic retrieval backend. +## Persistence and Retrieval -## Fixed Paper Test Data +The flow performs no library write. Callers explicitly pass its result to `LocalKnowledgeLibrary.put_paper()`. That method prepares all required summary and chunk embeddings before one SQLite transaction, so provider failure cannot leave a partial paper. -The fixed test files live at: +Collection search uses library-owned projections. A summary hit resolves to `PaperGlobalSummary`; a chunk hit resolves to `PaperChunk`. `SemanticHit.locator` carries source revision, artifact, artifact kind, and optional member ID. `SemanticHit.projection` explains the rebuildable embedding projection used for ranking. -```text -tests/fixtures/paper/golden/ -├── paper.pdf -└── expected.json +## Verification Slice + +Offline tests use fixed generated PDFs and fake summary/embedding providers. They verify exact source bytes, page evidence, stable IDs, citation rejection, bounds, atomic writes, reopen behavior, projection reuse, multi-version coexistence, search, and locator resolution. + +The bounded live slice is: + +```bash +python scripts/verify_pdf_rag_e2e.py ``` -It is a small four-page test paper with a nested subsection, multi-page -sections, and intentionally overlapping sibling page ranges. `expected.json` -contains only stable facts: page count, titles and paths, 1-based inclusive -page ranges, distinctive text anchors, tree shape, and validation rules. It -does not pin summaries or any other wording produced by a model. - -The offline test parses the fixed PDF, checks every text anchor against its -declared page, validates the tree and all rules, and confirms that overlapping -ranges work. Paper extraction and future PageIndex work must reuse these files -rather than create a competing test paper. - -## Current Gaps - -The repository does not yet guarantee the target pipeline above: - -- `pdf_to_markdown()` remains a compatibility view, while the primary - `parse_pdf()` path now preserves pages, blocks, coordinates, and artifacts. -- `paper_flow()` has not yet adopted `ParsedDocument` or the document RAG - boundary; it still consumes the compatibility Markdown view. -- `paper_flow()` sends the flattened document to one extraction agent and asks - it to return the final `Paper` directly. -- The model currently controls IDs, edges, citations, source fields, and - content instead of returning the limited draft described above. -- Resolved source URL, content hash, publication/version time, fetch time, and - exact-version availability are not consistently collected from the source. -- `DoiIdentifier` raises `NotImplementedError` because there is no accessible - content resolver. -- `TreeKnowledge` provides traversal helpers but does not yet check every root, - link, reachability, cycle, or page-range rule when it is created. -- The structured-output failure tracked by issue #91 remains separate from - this design issue. - -Those gaps are future implementation work. Adding PageIndex first would not fix -them and is not required to implement the Paper steps above. +It fetches exact arXiv revision `1706.03762v7`, parses at least 15 physical pages, creates multiple page-aware chunks, generates a cited global summary with bounded `gpt-4o-mini` calls, persists with `text-embedding-3-small`, closes and reopens the database, runs summary and chunk searches, resolves every returned locator, and prints useful summary, page, citation, and score diagnostics. The dedicated `paper-flow` job in `.github/workflows/e2e.yml` owns this non-required public-network check. + +## Out of Scope + +- model-authored canonical IDs, hashes, lineage, or source metadata; +- a V1 paper tree or PageIndex structure; +- HTML, Markdown, raw-text, or scanned-document OCR paper ingestion; +- DOI-to-open-PDF resolution; +- question answering over search results; +- hidden or unbounded model calls; +- implicit persistence from `paper_flow`. diff --git a/contexts/design/knowledge/paper.md b/contexts/design/knowledge/paper.md new file mode 100644 index 0000000..1139078 --- /dev/null +++ b/contexts/design/knowledge/paper.md @@ -0,0 +1,93 @@ +# Model paper sources and artifacts independently + +## Quick Summary + +- **Purpose**: Define canonical Paper Flow V1 source, chunk, summary, citation, and locator models. +- **Read when**: Changing `quantmind.knowledge.paper`, stable paper identities, artifact lineage, or paper search resolution. +- **Status**: Implemented by `quantmind.knowledge.paper` and persisted by `quantmind.library`. +- **Core rule**: Source revisions are immutable anchors; derived artifacts are independently versioned and never own retrieval vectors. + +## Contents + +- [Model Layers](#model-layers) +- [Stable Identity](#stable-identity) +- [Source and Asset Integrity](#source-and-asset-integrity) +- [Artifact Versioning](#artifact-versioning) +- [Citation and Lineage Integrity](#citation-and-lineage-integrity) +- [Retrieval Boundary](#retrieval-boundary) +- [Compatibility Boundary](#compatibility-boundary) + +## Model Layers + +Paper Flow V1 separates three layers: + +| Layer | Canonical aggregate | Addressable members | Purpose | +|---|---|---|---| +| Exact source | `PaperSourceRevision` | `PaperAssetRef`, `PaperParsedPage`, `PaperParsedBlock` | Preserve fetched bytes, page-aware parser output, metadata, and visual evidence. | +| Deterministic artifact | `PaperChunkSet` | `PaperChunk`, `PaperSourceSpan` | Record one exact chunking of the source before any summary call. | +| Semantic artifact | `PaperGlobalSummary` | `PaperCitation` | Store one independently versioned model summary with resolvable chunk/page evidence. | + +`PaperFlowResult` validates one compatible source, chunk set, and summary combination. It is a transfer result, not a fourth stored artifact. + +All models are frozen Pydantic values with `extra="forbid"`. Canonical values contain no embedding vectors, provider node objects, or storage handles. + +## Stable Identity + +IDs are generated and checked in code: + +- source revision ID: UUIDv5 over the exact source SHA-256 hash; +- asset ID: UUIDv5 over source revision, asset kind, page, and asset content hash; +- artifact ID: UUIDv5 over source revision, artifact kind, and producer configuration hash; +- chunk ID: UUIDv5 over chunk-set ID, position, content hash, and source-span hash. + +Producer configuration is canonical JSON with sorted keys before SHA-256 hashing. Chunk-set content hashes cover ordered chunk membership and spans. Summary content hashes cover summary prose and ordered citations. + +These identities make an identical run idempotent. They also keep a changed splitter or summary producer from overwriting an older artifact. + +## Source and Asset Integrity + +`PaperSourceRevision` requires a typed `SourceRef` whose `content_hash` equals the parsed manifest hash and raw-asset hash. arXiv sources require an exact revision suffix. Pages are contiguous and 1-based. Every page visual reference must name a known asset from that page. + +`PaperAssetRef` records media type, content hash, byte length, kind, and optional page. Exact blobs are keyed by content hash while the result crosses into persistence. When blobs are loaded, every reference must have bytes with matching length and SHA-256 hash. + +The source's canonical JSON excludes blobs. This keeps canonical hashes stable and reviewable while allowing SQLite to store exact bytes in a normalized linked table. Rehydration checks both directions: stored blob bytes must match their table hashes, and table asset metadata must match the canonical source manifest. + +Every chunk span is also checked against that manifest: its page must exist, its character range must fit the page text, and every visual asset ID must resolve to a screenshot or image from the same page. This check runs for a complete `PaperFlowResult` and when a stored chunk set is rehydrated independently. + +## Artifact Versioning + +`PaperChunkSet.producer` records splitter identity, installed splitter version, chunk size, and overlap. Its members have contiguous positions and must all point back to the artifact and source revision. + +`PaperGlobalSummary.producer` records: + +- model identity; +- prompt version; +- exact input chunk-set ID; +- summary-instruction hash; +- maximum output tokens. + +Changing any producer field creates a distinct artifact ID. Multiple chunk sets and summaries may coexist for one source revision. Loading a complete `PaperFlowResult` without explicit artifact IDs is allowed only when one unambiguous linked pair exists. + +## Citation and Lineage Integrity + +A `PaperCitation` identifies the exact chunk set, chunk, page, and optional verbatim quote. `PaperFlowResult` rejects citations to missing chunks, pages outside the cited chunk spans, or quotes absent from chunk text. + +`PaperGlobalSummary.derived_from` contains `ArtifactLocator` values. At least one locator must point to its producer's exact input chunk set, with the same source revision and no member ID. The library stores this relationship explicitly so lineage can be checked independently from the summary JSON. + +## Retrieval Boundary + +Canonical paper models do not implement `embedding_text()` and do not select retrieval text. `quantmind.library` projects: + +- one text-embedding target for the global summary; +- one text-embedding target per paper chunk; +- no aggregate target for a chunk set. + +`ArtifactLocator` addresses a source revision, artifact, artifact kind, and optional member. The optional source revision keeps the locator usable for legacy `BaseKnowledge` results; V1 paper locators always set it. `LocalKnowledgeLibrary.resolve()` returns the canonical summary, chunk set, chunk, knowledge item, or tree node selected by a locator. + +`SearchProjection` is separate from the locator. It records the rebuildable projection kind, version, modality, model, dimensions, and content hash that produced a ranked `SemanticHit`. + +## Compatibility Boundary + +`LegacyPaper` retains the pre-V1 `TreeKnowledge` shape only so existing version-2 databases and the bundled legacy example can be opened. It is not exported as `Paper`, is not produced by `paper_flow`, and is not part of the V1 paper contract. + +There is no `PaperTree` in V1. A future tree or PageIndex artifact must be independently versioned, linked to an exact source and inputs, and added only with a real retrieval use case and migration plan. diff --git a/contexts/design/library/local.md b/contexts/design/library/local.md index cdec382..9ee8ba5 100644 --- a/contexts/design/library/local.md +++ b/contexts/design/library/local.md @@ -1,140 +1,141 @@ -# Local Knowledge Library Design +# Store and search canonical knowledge locally ## Quick Summary -- **Purpose**: Define how the local library stores validated knowledge and searches it by meaning. -- **Read when**: Changing `LocalKnowledgeLibrary`, search records, filters, query results, or file ownership. -- **Status**: Current design owned at runtime by `quantmind.library`. -- **Core rule**: Stored `BaseKnowledge` is the data to keep. Search text and vectors can be rebuilt. The caller stores raw PDF, HTML, and media files. +- **Purpose**: Define how the local library stores validated knowledge and source-first papers, then searches rebuildable projections. +- **Read when**: Changing `LocalKnowledgeLibrary`, SQLite schema, semantic projections, filters, hit locators, or file ownership. +- **Status**: Implemented by `quantmind.library` with SQLite schema version 3 and private LlamaIndex ranking. +- **Core rule**: Canonical sources and artifacts are durable; projection text, vectors, and private indexes are rebuildable. - **User guide**: [`docs/library.md`](../../../docs/library.md) ## Contents -- [Key Decisions](#key-decisions) -- [Who Owns What](#who-owns-what) -- [Stored Knowledge and Rebuildable Search Data](#stored-knowledge-and-rebuildable-search-data) -- [What Can Match a Query](#what-can-match-a-query) -- [When to Rebuild Search Data](#when-to-rebuild-search-data) -- [Time Fields and Look-Ahead](#time-fields-and-look-ahead) -- [Why SQLite and LlamaIndex Ranking](#why-sqlite-and-llamaindex-ranking) -- [Independent tree navigation](#independent-tree-navigation) +- [Public Contract](#public-contract) +- [Ownership](#ownership) +- [Canonical Storage](#canonical-storage) +- [Paper Transaction](#paper-transaction) +- [Search Projections](#search-projections) +- [Search and Resolution](#search-and-resolution) +- [Selective Rebuild](#selective-rebuild) +- [Financial Time](#financial-time) +- [Integrity and Migration](#integrity-and-migration) +- [PageIndex Boundary](#pageindex-boundary) - [Out of Scope](#out-of-scope) -## Key Decisions +## Public Contract -`quantmind.library` stores validated QuantMind knowledge and searches it by -meaning. It replaces an older design for one generic storage layer that would -hide raw files, knowledge JSON, embeddings, and indexes behind one backend API. +`LocalKnowledgeLibrary` is the only public backend class. Its common operations are: -The public API is intentionally small: - -- `LocalKnowledgeLibrary` -- `SemanticQuery` -- `SemanticHit` +| Operation | Contract | +|---|---| +| `open()` | Open or migrate a SQLite library without network I/O. | +| `put()` | Store one conventional `BaseKnowledge` item and its required projections. | +| `put_paper()` | Store one `PaperFlowResult`, including exact source assets, two artifacts, lineage, and required projections. | +| `get()` | Rehydrate one conventional knowledge item. | +| `get_paper()` | Rehydrate one unambiguous source/chunk-set/summary result, or use explicit artifact IDs when versions coexist. | +| `get_artifact()` | Rehydrate a paper chunk set or global summary by artifact ID. | +| `search()` | Filter and rank rebuildable projections, returning `SemanticHit` evidence. | +| `resolve()` | Resolve a hit locator to its canonical aggregate or member. | -There is no public `Storage`, `VectorStore`, `Retriever`, provider registry, or -backend class hierarchy. Add a shared backend API only after a second working -implementation proves which behavior is truly shared. +The public types are `LocalKnowledgeLibrary`, `SemanticQuery`, `SemanticHit`, and `SearchProjection`. SQLite rows, embedding providers, LlamaIndex nodes, retrievers, and indexes remain private. -## Who Owns What +## Ownership -| Package or caller | Responsibility | +| Owner | Responsibility | |---|---| -| `quantmind.knowledge` | Define immutable knowledge models and the text used for embeddings; perform no I/O | -| `quantmind.library` | Store validated knowledge, maintain rebuildable search records, and return `SemanticHit` results | -| [`quantmind.rag`](../rag/document.md) | Chunk and retrieve evidence within one parsed document without storing canonical knowledge | -| `quantmind.flows` | Produce validated knowledge and optionally pass it to a library | -| `quantmind.mind` or an agent application | Search the library and use matches to write answers | -| Caller or source-specific pipeline | Retain raw PDF, HTML, media, and operational files | +| `quantmind.knowledge` | Define immutable canonical models; perform no I/O and choose no retrieval text. | +| `quantmind.library` | Persist canonical values, own retrieval projections, maintain vectors, and resolve search evidence. | +| [`quantmind.rag`](../rag/document.md) | Build page-aware chunks or transient document-local evidence without owning canonical persistence. | +| `quantmind.flows` | Produce validated results and leave persistence explicit. | +| Caller | Choose database and temporary parser-artifact locations and manage application lifecycle. | + +For conventional `BaseKnowledge`, source references remain pointers and the caller retains external raw files. For `PaperFlowResult`, the exact source PDF, screenshots, and extracted images are part of the durable source revision and are copied into the library transaction. + +## Canonical Storage + +SQLite schema version 3 keeps conventional and source-first paper storage explicit. + +Conventional knowledge uses: + +- `knowledge_items` for canonical typed aggregate payloads; +- `knowledge_nodes` for normalized tree members; +- `semantic_records` for rebuildable conventional projections and vectors. + +Source-first papers use: + +- `paper_sources` for immutable source-revision metadata and canonical manifests; +- `paper_source_assets` for linked content-addressed bytes and asset metadata; +- `paper_artifacts` for independently versioned chunk sets and summaries; +- `paper_artifact_members` for directly addressable chunks; +- `paper_artifact_lineage` for summary-to-chunk-set derivation; +- `paper_projections` for rebuildable summary and chunk text embeddings. + +There is no single opaque paper JSON blob and no canonical vector field. Aggregate JSON and normalized relationship rows are cross-checked during rehydration. + +## Paper Transaction + +`put_paper()` first validates the complete `PaperFlowResult`, computes its source and artifact canonical forms, determines affected projections, and obtains every required embedding. Only then does it begin one `BEGIN IMMEDIATE` transaction. + +The transaction writes or reuses the source, asset blobs, artifacts, members, lineage, and all required projections. Any constraint, integrity, or write failure rolls back the transaction. An embedding-provider failure occurs before the transaction and therefore leaves no partial source or artifact rows. + +Putting an unchanged result is idempotent and reuses valid vectors. The same source may own multiple chunk-set and summary versions. Artifact identity includes producer configuration, so one version never silently replaces another. + +## Search Projections + +Retrieval text selection lives in `quantmind.library._internal.retrieval_targets`, not canonical models. -The V1 library does not store raw source files. `SourceRef` and citations point -back to the source, but they do not contain the source file itself. +- Each supported flat knowledge item produces one whole-item projection. +- Each `TreeKnowledge` produces one aggregate projection and one projection for every non-root node. +- A paper global summary produces one aggregate projection. +- Every paper chunk produces one member projection. +- A paper chunk-set aggregate is not searchable by itself. -## Stored Knowledge and Rebuildable Search Data +Every stored projection records target identity, exact projection text and hash, projection schema version, canonical schema version, source content hash, embedding model, vector dimension, and bytes. Search returns the exact projected text as `SemanticHit.matched_text`. -Validated `BaseKnowledge` is the record that must be preserved. Embeddings, -the text sent to the embedder, filter columns, and vector indexes can all be -rebuilt, even when they share one SQLite database with the knowledge records. +## Search and Resolution -The local implementation separates three concerns: +`SemanticQuery.artifact_kinds` filters artifact granularity, including `paper_summary` and `paper_chunk_set`. Existing `item_types`, source, confidence, tag, tree, `as_of`, and `available_at` filters continue to apply where meaningful. Filtering happens before ranking. -- `knowledge_items` stores one complete validated knowledge item. -- `knowledge_nodes` stores each `TreeNode` in a separate row with its parent, - position, content hash, and owning item. -- `semantic_records` stores searchable text and vector details for whole items, - roots, and nodes. +Each `SemanticHit` carries two complementary values: -Concrete types do not get a table per Pydantic subtype. Storing each complete -item plus separate node rows lets the library rebuild the original typed item -and later navigate a tree one node at a time. +- `locator`: source revision, artifact ID, artifact kind, and optional member ID needed to resolve canonical evidence; +- `projection`: projection version, modality, embedding model, dimensions, and projection content hash used for ranking. -## What Can Match a Query +For paper hits, `search()` rehydrates and validates the owning source and artifact before returning evidence. Summary citations come from the canonical summary. Chunk citations are reconstructed from page and character spans. Callers pass the locator to `resolve()` to obtain a `PaperGlobalSummary` or exact `PaperChunk`. -- A `FlattenKnowledge` item produces one searchable record from its exact - `embedding_text()` value. -- A `TreeKnowledge` item produces one record for the whole item with - `node_id=None`, plus one record for every non-root node. -- Each record says whether it represents a whole item or one node. -- `SemanticHit.matched_text` is the exact text used to rank the match. -- Callers use `get(item_id)` to load the full knowledge item and node paths. +Compatibility fields `item_id`, `node_id`, and `item_type` remain on `SemanticHit` for conventional knowledge callers. For paper hits they mirror artifact ID, member ID, and artifact kind. -Putting unchanged knowledge with the same item ID is safe to repeat and reuses -its embeddings. The library does not merge separate extraction runs that -created different item IDs. +## Selective Rebuild -## When to Rebuild Search Data +A projection is reusable only when all recorded identities still match: projection hash and schema version, canonical schema version, source content hash, embedding model, and requested vector dimension. Invalid vector bytes or dimensions are rejected. -Every stored vector records the information needed to decide whether it is -reusable: embedding model, dimension, embedding-text hash, source content hash, -knowledge model version, and embedding-text format version. +`put()` and `put_paper()` send only affected target text to the embedding provider. Changing summary prose or producer configuration does not re-embed unchanged chunks. Closing and reopening the library loads persisted vectors without document re-embedding; a search embeds only the query. -When any of those values changes, rebuild only the affected search records. -Deleting an item removes the item, its nodes, and its search records in one -transaction. Invalid vector sizes or bytes, unreadable knowledge data, and -search rows without an item raise an error instead of returning an incomplete -match that looks valid. +The private LlamaIndex retrieval state is rebuilt lazily after open or mutation. It is never canonical state and may be replaced without migrating source or artifact payloads. -## Time Fields and Look-Ahead +## Financial Time -`as_of` and `available_at` answer different questions: +`as_of` is the latest date covered by knowledge. `available_at` is when the exact source version became observable. `available_at_before` excludes records with a later or unknown availability time; filtering only on `as_of_before` may still leak future information. -- `as_of` is the latest date covered by the knowledge. -- `available_at` is when that source version became observable. +Paper projections inherit both times from their exact source revision. This prevents a later summary run from changing the historical availability of the underlying paper. -`available_at_before` excludes records that became available after the cutoff -or have no known availability time. Filtering only by `as_of_before` can still -leak future information. Apply source kind, item type, confidence, tags, tree -ID, and both time cutoffs before ranking results. +## Integrity and Migration -## Why SQLite and LlamaIndex Ranking +Open performs an explicit SQLite user-version migration from schema 2 to schema 3 by adding paper tables and indexes without rewriting conventional knowledge. Older canonical class references for the pre-V1 paper tree load as `LegacyPaper` solely for database compatibility. -SQLite provides transactions, foreign keys, and reliable reconstruction of -typed knowledge. LlamaIndex owns the private collection-wide vector retrieval -and ranking mechanics. On the first search after open or a write, private -retrieval state is rebuilt from the filtered semantic records stored in SQLite; -unchanged records reuse their persisted embeddings and are not sent to the -embedding provider again. +Reads fail closed when canonical hashes, counts, IDs, membership, lineage, source relationships, vector bytes, or asset metadata disagree. Blob SHA-256 hashes and byte lengths are checked, and stored asset table fields must match the canonical source manifest. Missing IDs raise `KeyError`; stale or corrupt linked state raises `RuntimeError` with context. -LlamaIndex nodes and retrievers remain implementation details. They do not -enter `SemanticQuery`, `SemanticHit`, canonical Pydantic payloads, or public -signatures. A future approximate or remote index may replace the private -search implementation without replacing stored knowledge or changing the -public result type. +## PageIndex Boundary -## Independent tree navigation +The library is canonical storage plus collection-wide semantic search, not a vector database abstraction. A future PageIndex path may first select a paper through `search()` and then navigate a separate, independently versioned tree artifact through an opinionated operation under `quantmind.rag`. -`LocalKnowledgeLibrary` is canonical knowledge storage with rebuildable -retrieval capabilities; it is not defined as a vector database. A future -PageIndex path can select a paper through collection-wide semantic retrieval, -then navigate that selected document's tree through a separate operation and -separately rebuildable state. PageIndex does not have to be served through -`search()` or LlamaIndex ranking. Opinionated document retrieval, including a -future PageIndex adapter, belongs under [`quantmind.rag`](../rag/document.md). +PageIndex is not required to use `LocalKnowledgeLibrary.search()` or private LlamaIndex vector ranking. Paper Flow V1 deliberately stores chunks and a cited summary without defining a paper tree. ## Out of Scope -- raw source file storage; -- a generic framework for retrieving context and writing answers; -- a public embedder, vector-store, retriever, or backend registry; -- PageIndex tree construction or navigation; -- merging knowledge across runs without a separate stable-ID design. +- a public storage, vector-store, retriever, or provider hierarchy; +- implicit persistence inside flows; +- answer synthesis or agent memory; +- merging distinct canonical identities; +- a V1 paper tree or PageIndex navigation API; +- treating rebuildable projections as canonical knowledge. diff --git a/contexts/design/operations/naming.md b/contexts/design/operations/naming.md index 69d4cf8..be9e5fd 100644 --- a/contexts/design/operations/naming.md +++ b/contexts/design/operations/naming.md @@ -48,7 +48,7 @@ use `*_pipeline` when the function combines multiple public operations. `PaperInput`. - Config types name the domain and stage, such as `NewsCollectionCfg` or a future `PaperExtractionCfg`. -- Result types describe returned data, such as `NewsBatch`, `Paper`, or +- Result types describe returned data, such as `NewsBatch`, `PaperFlowResult`, or `PageIndex`. - Keep provider names out of public function names unless callers are choosing provider-specific behavior. diff --git a/contexts/design/preprocess/pdf.md b/contexts/design/preprocess/pdf.md index 3931da7..4e29959 100644 --- a/contexts/design/preprocess/pdf.md +++ b/contexts/design/preprocess/pdf.md @@ -5,7 +5,7 @@ - **Purpose**: Define the deterministic PDF value shared by paper extraction, document RAG, and a future PageIndex adapter. - **Read when**: Changing PDF parsing, page artifacts, or multimodal source evidence. - **Status**: Implemented by `quantmind.preprocess.format.parse_pdf`. -- **Core rule**: Parsing preserves every physical page and its source coordinates before any chunker or tree builder runs. +- **Core rule**: Parsing preserves every physical page and its source coordinates before any chunker or semantic artifact producer runs. ## Contents @@ -21,14 +21,16 @@ Each `TextBlock` keeps its page ownership, text, bounding box, and parser-provid ## Artifacts and ownership -The caller chooses an artifact directory. When supplied, parsing renders one PNG screenshot per page and stores a stable path reference on the page. The library does not copy screenshots or PDF bytes into canonical SQLite knowledge. Without an artifact directory, pages remain valid and `screenshot_path` is absent. +The caller chooses an artifact directory. When supplied, parsing renders one PNG screenshot per page and stores a path reference on the page. Without an artifact directory, pages remain valid and `screenshot_path` is absent. + +`ParsedDocument` paths and bytes are preprocessing values, not canonical storage. A source-first paper flow reads the referenced bytes, converts them to content-addressed `PaperAssetRef` values, and carries the blobs in `PaperSourceRevision` until `LocalKnowledgeLibrary.put_paper()` copies them into linked SQLite rows. Other parser callers continue to own their files. ## Downstream consumers Preprocessing ends after producing `ParsedDocument`. It does not chunk, index, rank, or answer a query. - [`quantmind.rag`](../rag/document.md) converts the parsed value into LlamaIndex-backed chunks and page-aware retrieval evidence. -- [`paper_flow`](../flow/paper.md) uses the preserved source pages when assembling a canonical `Paper`. -- A future PageIndex adapter may consume the same ordered pages to propose and navigate a document tree. +- [`paper_flow`](../flow/paper.md) uses the preserved source pages to build an exact source revision and a canonical chunk set before generating a cited summary. +- A future PageIndex adapter may consume the same ordered pages to propose independently versioned navigation evidence. Flattened Markdown remains a compatibility view produced from the preserved pages. It is not the primary parsing result. diff --git a/contexts/design/rag/document.md b/contexts/design/rag/document.md index 6d2a929..5d4903e 100644 --- a/contexts/design/rag/document.md +++ b/contexts/design/rag/document.md @@ -2,36 +2,62 @@ ## Quick Summary -- **Purpose**: Define how an ordered `ParsedDocument` becomes chunks and ranked evidence without leaking LlamaIndex types. -- **Read when**: Changing document chunking, document-local retrieval, RAG evidence, or a future PageIndex adapter. +- **Purpose**: Define how an ordered `ParsedDocument` becomes deterministic page-aware chunks and ranked document-local evidence. +- **Read when**: Changing document chunking, source spans, document-local retrieval, Paper Flow chunk construction, or a future PageIndex adapter. - **Status**: Implemented by `quantmind.rag.document`. -- **Core rule**: `quantmind.rag` is an opinionated LlamaIndex data-plane package, not a generic retriever or backend framework. +- **Core rule**: LlamaIndex owns splitting and ranking mechanics; QuantMind owns stable identity and source/page provenance. ## Contents -- [Package boundary](#package-boundary) -- [Chunk and retrieval contract](#chunk-and-retrieval-contract) -- [What this package does not abstract](#what-this-package-does-not-abstract) -- [Collection search and PageIndex](#collection-search-and-pageindex) +- [Package Boundary](#package-boundary) +- [Chunk Contract](#chunk-contract) +- [Document-Local Retrieval](#document-local-retrieval) +- [Paper Flow Boundary](#paper-flow-boundary) +- [Collection Search and PageIndex](#collection-search-and-pageindex) +- [What This Package Does Not Abstract](#what-this-package-does-not-abstract) -## Package boundary +## Package Boundary `quantmind.preprocess` owns deterministic source parsing and returns a page-aware [`ParsedDocument`](../preprocess/pdf.md). `quantmind.rag` may import that value and apply LlamaIndex transformations and retrieval. Preprocessing never imports RAG, so parsing remains usable without a query or index. -LlamaIndex is a required dependency and owns the chunker, nodes, indexes, retrievers, ranking mechanics, and their supported parameters. QuantMind adds only the work that LlamaIndex cannot own: stable source hashes, page ownership, block coordinates, screenshot/image references, and conversion back to typed QuantMind evidence. +LlamaIndex is a required dependency and owns `SentenceSplitter`, BM25, nodes, indexes, retrievers, ranking mechanics, and upstream parameters. QuantMind adds stable source hashes, page ownership, page-local character spans, block coordinates, screenshot/image references, and conversion back to typed evidence. -## Chunk and retrieval contract +## Chunk Contract -`chunk_parsed_document()` applies LlamaIndex `SentenceSplitter` to each non-empty page without erasing physical page boundaries. `SentenceSplitterConfig` exposes the selected upstream parameters by their upstream meaning rather than reimplementing the algorithm. +`chunk_parsed_document()` applies LlamaIndex `SentenceSplitter` independently to each non-empty physical page. Splitting by page prevents a chunk from erasing page ownership or spanning an implicit page boundary. -Each `ParsedChunk` retains the exact source hash, 1-based page number, available block bounding boxes, and screenshot/image references. `retrieve_parsed_document()` uses LlamaIndex BM25 and returns ranked `ParsedDocumentHit` values. LlamaIndex `Document`, node, retriever, index, and score wrapper types remain private implementation details. +`SentenceSplitterConfig` exposes chunk size, overlap, separator, paragraph separator, and tokenizer behavior by their upstream meanings. It does not reimplement the splitter. -## What this package does not abstract +Each `ParsedChunk` records: -The package does not define a public `Retriever`, `VectorStore`, backend registry, provider protocol, query engine, or answer-synthesis framework. Add another direct, opinionated operation only when a real pipeline needs it. Do not build an abstraction solely to hide an upstream LlamaIndex call. +- deterministic SHA-256-derived chunk ID; +- exact document source hash; +- 1-based physical page number; +- page-local `start_char` and `end_char` offsets; +- exact chunk text; +- overlapping parser block bounding boxes; +- page screenshot and extracted-image paths. -## Collection search and PageIndex +Repeating chunking for the same parsed document and configuration produces the same ordered texts, spans, and IDs. Empty pages remain in `ParsedDocument` but produce no chunks. -Document-local RAG and collection search have different responsibilities. [`LocalKnowledgeLibrary`](../library/local.md) stores canonical knowledge in SQLite and privately uses LlamaIndex for collection-wide semantic ranking. It does not own transient PDF parsing or document-local BM25 chunks. +## Document-Local Retrieval -A future PageIndex implementation may live under `quantmind.rag` as another opinionated document operation. It can consume `ParsedDocument` and return a limited tree draft or navigation evidence, while canonical IDs, links, citations, and source-backed text remain owned by QuantMind code. PageIndex does not have to run through `LocalKnowledgeLibrary.search()` or LlamaIndex vector ranking. +`retrieve_parsed_document()` chunks one document, uses LlamaIndex BM25, and returns ranked `ParsedDocumentHit` values. Each hit wraps one `ParsedChunk` and a score. The operation is transient and requires no canonical library write. + +LlamaIndex `Document`, node, retriever, index, and score-wrapper types remain private. Public callers receive frozen QuantMind dataclasses with enough evidence to trace a hit back to physical pages and parser artifacts. + +## Paper Flow Boundary + +[`paper_flow`](../flow/paper.md) uses `chunk_parsed_document()` as the deterministic split stage. It converts `ParsedChunk` values into canonical `PaperChunk` members only after an exact `PaperSourceRevision` exists. + +The conversion replaces parser paths with canonical source asset IDs and validates character spans against page evidence. The resulting `PaperChunkSet` is a durable, independently versioned artifact. `quantmind.rag` itself does not import or construct canonical paper models. + +## Collection Search and PageIndex + +Document-local RAG and collection search have different responsibilities. [`LocalKnowledgeLibrary`](../library/local.md) stores canonical sources and artifacts in SQLite and privately uses LlamaIndex for collection-wide embedding ranking. It does not persist transient `ParsedDocumentHit` values. + +Paper Flow V1 defines no paper tree. A future PageIndex implementation may live under `quantmind.rag` as another opinionated document operation. It may consume `ParsedDocument` or an exact paper source and return a bounded draft or navigation evidence, while canonical IDs, links, citations, and source-backed text remain code-owned. + +## What This Package Does Not Abstract + +The package does not define a public `Retriever`, `VectorStore`, backend registry, provider protocol, generic query engine, answer-synthesis framework, or canonical paper tree. Add another direct opinionated operation only when a real pipeline needs it; do not add a wrapper solely to hide an upstream LlamaIndex call. diff --git a/contexts/usage/README.md b/contexts/usage/README.md index 3f973b1..4d5193b 100644 --- a/contexts/usage/README.md +++ b/contexts/usage/README.md @@ -26,7 +26,7 @@ current public API, focused examples, and component-specific guidance. |---|---| | Current operations, inputs, results, and sources | [Public component catalog](../../docs/README.md) | | Installation and common usage | [Root README usage](../../README.md#-usage-examples) | -| Paper extraction | [Paper extraction design](../design/flow/paper.md) | +| Source-first paper flow | [Paper flow design](../design/flow/paper.md) | | News collection | [News design and behavior](../design/flow/news.md) | | Search local knowledge by meaning | [Library guide](../../docs/library.md) and [focused example](../../examples/library/README.md) | | Runnable operation examples | [`examples/flows/`](../../examples/flows/) | diff --git a/docs/README.md b/docs/README.md index eff681e..5758058 100644 --- a/docs/README.md +++ b/docs/README.md @@ -14,10 +14,10 @@ harness. | Operation | Import | Input and config | Result | Example | Design or guide | |---|---|---|---|---|---| -| Paper extraction | `quantmind.flows.paper_flow` | `PaperInput`, `PaperFlowCfg` | `Paper` | [README usage](../README.md#-usage-examples) | [Paper E2E design](../contexts/design/flow/paper.md) | +| Source-first paper flow | `quantmind.flows.paper_flow` | `PaperInput`, `PaperFlowCfg` | `PaperFlowResult` | [Persist and search a paper](../examples/flows/paper.py) | [Paper flow design](../contexts/design/flow/paper.md) | | News collection | `quantmind.flows.collect_news` | `NewsWindow`, `NewsCollectionCfg` | `NewsBatch` from `quantmind.preprocess` | [Collect news](../examples/flows/collect_news.py) | [News collection design](../contexts/design/flow/news.md) | | Bounded fan-out | `quantmind.flows.batch_run` | Operation inputs and shared config | `BatchResult` | [README usage](../README.md#-usage-examples) | API docstrings | -| Local semantic search | `quantmind.library.LocalKnowledgeLibrary` | `BaseKnowledge`, `SemanticQuery` | `list[SemanticHit]` | [Library example](../examples/library/README.md) | [Library guide](library.md) | +| Local semantic search | `quantmind.library.LocalKnowledgeLibrary` | `BaseKnowledge` or `PaperFlowResult`, `SemanticQuery` | `list[SemanticHit]` | [Library example](../examples/library/README.md) | [Library guide](library.md) | | Page-aware document RAG | `quantmind.rag.chunk_parsed_document`, `quantmind.rag.retrieve_parsed_document` | `ParsedDocument`, splitter config, and query | `tuple[ParsedDocumentHit, ...]` | [Paper RAG](../examples/rag/paper.py) | [Document RAG design](../contexts/design/rag/document.md) | Import public inputs and configs from `quantmind.configs` and current public @@ -29,7 +29,7 @@ layer shown in the catalog. | Source | Source selection | Operation | Live-network component smoke test | |---|---|---|---| | PR Newswire | `NewsWindow(source="pr-newswire", ...)` | `collect_news` | `python scripts/verify_news_e2e.py` | -| arXiv Transformer PDF | `fetch_arxiv("1706.03762v7")` | `parse_pdf` and `quantmind.rag` retrieval | `python scripts/verify_pdf_rag_e2e.py` | +| arXiv Transformer PDF | `ArxivIdentifier(id="1706.03762v7")` | `paper_flow`, persistence, reopen, search, and resolution | `python scripts/verify_pdf_rag_e2e.py` | The PR Newswire smoke test checks the public RSS feed, a complete preceding 24-hour listing window, and ticker-hint recall on a bounded sample of up to 25 @@ -38,6 +38,13 @@ manual dispatch, and only on pull requests that change its dependency paths. It is not a required merge check, so external PR Newswire availability cannot block unrelated changes. +The `paper-flow` job fetches exact arXiv revision `1706.03762v7`, preserves at +least 15 pages, runs bounded `gpt-4o-mini` summarization, persists summary and +chunk projections with `text-embedding-3-small`, reopens the database, searches +both artifact kinds, and resolves every hit. It runs daily, manually, and on +pull requests that change its dependency paths, and it remains non-required +because arXiv and model providers are public-network dependencies. + ## Verification Run the deterministic required verification for every change: diff --git a/docs/library.md b/docs/library.md index 3323d34..c427cd0 100644 --- a/docs/library.md +++ b/docs/library.md @@ -1,64 +1,96 @@ # Local Semantic Knowledge Library -This page only explains how to run the bundled example. The canonical storage, -retrieval, financial-time, and PageIndex boundaries live in the -[local library design](../contexts/design/library/local.md). +`LocalKnowledgeLibrary` persists canonical QuantMind values in SQLite and ranks rebuildable text-embedding projections with a private LlamaIndex retriever. The canonical storage, transaction, financial-time, migration, and PageIndex boundaries live in the [local library design](../contexts/design/library/local.md). -## Usage +## Store a Paper Flow Result -The bundled AI-infrastructure scenario contains primary-source-backed `News`, -`Earnings`, and a real research `Paper` tree. Its canonical source JSON is -precompiled into a SQLite database with six `text-embedding-3-small` targets, -so the user-facing example does not perform ingestion or document embedding. -Put `OPENAI_API_KEY` in `.env` and run: +Run `paper_flow()` first, then explicitly store its complete result: -```bash -python examples/library/semantic_search.py +```python +result = await paper_flow( + ArxivIdentifier(id="1706.03762v7"), + cfg=PaperFlowCfg(model="gpt-4o-mini"), +) + +library = await LocalKnowledgeLibrary.open( + ".quantmind/library.db", + embedding_model="text-embedding-3-small", +) +try: + await library.put_paper(result) +finally: + await library.close() ``` -The example: +`put_paper()` persists the exact source PDF and retained parser assets, one page-aware chunk-set artifact, one cited global-summary artifact, explicit lineage, and required summary/chunk projections. It obtains all affected embeddings before opening the SQLite transaction, so an embedding failure leaves no partial paper. -1. Opens the ready-to-search SQLite bundle through `LocalKnowledgeLibrary`. -2. Embeds only the user's query with `text-embedding-3-small`. -3. Searches a concrete AI-infrastructure question with a no-look-ahead - availability cutoff. -4. Resolves every hit with `get()` and prints source, citation, financial time, - and the root-to-node path and content for tree evidence. +Putting the same result again is safe and reuses valid vectors. A changed splitter or summary producer creates another independently addressable artifact version for the same source. -Maintainers can regenerate the model-specific database from the auditable JSON -after changing the source data or storage schema: +## Reopen, Search, and Resolve -```bash -python scripts/examples/build_ai_infrastructure_bundle.py +Opening a library performs no embedding or network request. Search embeds only the query when stored projections are reusable: + +```python +library = await LocalKnowledgeLibrary.open( + ".quantmind/library.db", + embedding_model="text-embedding-3-small", +) +try: + summary_hits = await library.search( + SemanticQuery( + text="What is the paper's central contribution?", + artifact_kinds=["paper_summary"], + top_k=3, + ) + ) + chunk_hits = await library.search( + SemanticQuery( + text="How does multi-head attention work?", + artifact_kinds=["paper_chunk_set"], + top_k=5, + ) + ) + evidence = [ + await library.resolve(hit.locator) + for hit in (*summary_hits, *chunk_hits) + ] +finally: + await library.close() ``` -The bundle's facts and short citations come directly from the -[Compute Trends Across Three Eras of Machine Learning paper](https://arxiv.org/abs/2202.05924), -[Microsoft's FY2025 AI-datacenter investment announcement](https://blogs.microsoft.com/on-the-issues/2025/01/03/the-golden-opportunity-for-american-ai/), -and [NVIDIA's FY2026 Q1 results](https://investor.nvidia.com/news/press-release-details/2025/NVIDIA-Announces-Financial-Results-for-First-Quarter-Fiscal-2026/default.aspx). +A `paper_summary` hit resolves to `PaperGlobalSummary`. A `paper_chunk_set` hit has a member ID and resolves to the exact `PaperChunk`, including source-page spans. Every `SemanticHit` also includes: + +- `matched_text`, the exact library-owned projection used for ranking; +- `projection`, the projection version, model, dimensions, and content hash; +- source metadata, financial time, and canonical citations; +- compatibility fields `item_id`, `node_id`, and `item_type`. + +Use `get_artifact(artifact_id)` when the aggregate ID is already known. Use `get_paper(source_revision_id, chunk_set_id=..., summary_id=...)` to reconstruct a compatible result. Artifact IDs may be omitted only when one unambiguous linked chunk-set/summary pair exists. + +The complete runnable path is [examples/flows/paper.py](../examples/flows/paper.py). + +## Conventional Knowledge + +`put(item)` and `get(item_id)` remain available for supported `BaseKnowledge` values such as `News`, `Earnings`, `Factor`, `Thesis`, and generic trees. Canonical models do not implement `embedding_text()`; library-owned projection rules select searchable text. -Representative output has this shape; scores depend on the selected embedding -model: +`SemanticQuery` supports item type, source kind, confidence, tag, tree, `as_of`, and `available_at` filters. Use `available_at_before` to prevent look-ahead: an `as_of` cutoff alone does not prove that the source was observable at that time. -```text -Bundle: .../examples/library/data/ai_infrastructure.db -Query: What evidence shows demand for AI infrastructure is expanding? +## Bundled Compatibility Example -1. score=0.812 type=paper - path: ... > ... - matched: ... - source: https://... - citation: ... +The bundled AI-infrastructure scenario contains primary-source-backed `News`, `Earnings`, and one pre-V1 `LegacyPaper` tree. Its canonical JSON is precompiled into a SQLite database with six `text-embedding-3-small` targets, so the example embeds only the query: + +```bash +python examples/library/semantic_search.py +``` + +`LegacyPaper` exists only so older databases and this auditable example remain readable. New paper ingestion uses `PaperFlowResult` and `put_paper()`. + +Maintainers can regenerate the bundle after changing source data, projection rules, or the storage schema: + +```bash +python scripts/examples/build_ai_infrastructure_bundle.py ``` -`SemanticHit.matched_text` is the exact projection used in ranking. Use -`get(item_id)` to resolve the full canonical item and, for node hits, the node -identified by `node_id`. - -Missing IDs raise `KeyError`. Stored canonical payloads that no longer validate -and orphaned or mismatched derived records raise `RuntimeError` with stale-data -context. Invalid vector bytes and inconsistent stored dimensions raise a clear -corrupt-index `RuntimeError`; provider or query dimension mismatches raise -`ValueError`. `delete()` does not deserialize canonical JSON, so a stale -canonical payload can still be removed together with all derived targets in one -transaction. +The bundle's facts and short citations come directly from the [Compute Trends Across Three Eras of Machine Learning paper](https://arxiv.org/abs/2202.05924), [Microsoft's FY2025 AI-datacenter investment announcement](https://blogs.microsoft.com/on-the-issues/2025/01/03/the-golden-opportunity-for-american-ai/), and [NVIDIA's FY2026 Q1 results](https://investor.nvidia.com/news/press-release-details/2025/NVIDIA-Announces-Financial-Results-for-First-Quarter-Fiscal-2026/default.aspx). + +Missing IDs raise `KeyError`. Canonical payloads, linked rows, or asset metadata that no longer agree raise `RuntimeError` with stale-data context. Invalid vector bytes and inconsistent dimensions raise a corrupt-index `RuntimeError`; provider or query dimension mismatches raise `ValueError`. diff --git a/examples/flows/paper.py b/examples/flows/paper.py new file mode 100644 index 0000000..e4dad7f --- /dev/null +++ b/examples/flows/paper.py @@ -0,0 +1,146 @@ +"""Build, persist, reopen, and search the Transformer paper artifacts.""" + +import asyncio +import os +from collections.abc import Sequence +from pathlib import Path + +from dotenv import load_dotenv + +from quantmind.configs import PaperFlowCfg +from quantmind.configs.paper import ArxivIdentifier +from quantmind.flows import paper_flow +from quantmind.knowledge import PaperChunk, PaperGlobalSummary +from quantmind.library import ( + LocalKnowledgeLibrary, + SemanticHit, + SemanticQuery, +) + +_ARXIV_ID = "1706.03762v7" +_EMBEDDING_MODEL = "text-embedding-3-small" + + +async def _search_and_resolve( + library: LocalKnowledgeLibrary, +) -> tuple[ + list[SemanticHit], + list[SemanticHit], + Sequence[object], +]: + """Run both V1 retrieval grains and resolve every returned locator.""" + summary_hits = await library.search( + SemanticQuery( + text="What is the paper's central contribution?", + artifact_kinds=["paper_summary"], + top_k=3, + ) + ) + chunk_hits = await library.search( + SemanticQuery( + text="How does multi-head attention work?", + artifact_kinds=["paper_chunk_set"], + top_k=5, + ) + ) + resolved = [ + await library.resolve(hit.locator) + for hit in (*summary_hits, *chunk_hits) + ] + return summary_hits, chunk_hits, resolved + + +async def main() -> None: + """Run the common Paper Flow V1 path and print auditable evidence.""" + load_dotenv() + if not os.getenv("OPENAI_API_KEY"): + raise SystemExit("Set OPENAI_API_KEY before running this example.") + + workspace = Path(".quantmind") + workspace.mkdir(exist_ok=True) + result = await paper_flow( + ArxivIdentifier(id=_ARXIV_ID), + cfg=PaperFlowCfg( + model="gpt-4o-mini", + output_dir=str(workspace / "attention-assets"), + ), + ) + + database = workspace / "library.db" + library = await LocalKnowledgeLibrary.open( + database, + embedding_model=_EMBEDDING_MODEL, + ) + try: + await library.put_paper(result) + ( + first_summary_hits, + first_chunk_hits, + first_resolved, + ) = await _search_and_resolve(library) + finally: + await library.close() + + library = await LocalKnowledgeLibrary.open( + database, + embedding_model=_EMBEDDING_MODEL, + ) + try: + restored = await library.get_paper( + result.source_revision.id, + chunk_set_id=result.chunk_set.id, + summary_id=result.global_summary.id, + ) + ( + second_summary_hits, + second_chunk_hits, + second_resolved, + ) = await _search_and_resolve(library) + + print(restored.global_summary.summary) + print( + f"chunks={len(restored.chunk_set.chunks)} " + f"source_pages={len(restored.source_revision.parsed.pages)}" + ) + print( + "citations=" + + ", ".join( + f"page {citation.page_number} / chunk {citation.chunk_id}" + for citation in restored.global_summary.citations + ) + ) + print( + "scores_before_reopen=" + f"summary={[hit.score for hit in first_summary_hits]} " + f"chunks={[hit.score for hit in first_chunk_hits]}" + ) + print( + "scores_after_reopen=" + f"summary={[hit.score for hit in second_summary_hits]} " + f"chunks={[hit.score for hit in second_chunk_hits]}" + ) + for hit, resolved in zip( + (*second_summary_hits, *second_chunk_hits), + second_resolved, + strict=True, + ): + if isinstance(resolved, PaperGlobalSummary): + detail = "global summary" + elif isinstance(resolved, PaperChunk): + pages = sorted( + {span.page_number for span in resolved.source_spans} + ) + detail = f"chunk pages={pages} text={resolved.text[:120]!r}" + else: + detail = type(resolved).__name__ + print(f"score={hit.score:.3f} {detail}") + print( + f"resolved_before_reopen={len(first_resolved)} " + f"resolved_after_reopen={len(second_resolved)}" + ) + finally: + await library.close() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/library/README.md b/examples/library/README.md index 99c9a7c..f92599e 100644 --- a/examples/library/README.md +++ b/examples/library/README.md @@ -1,7 +1,9 @@ # Local Knowledge Library Example This example searches an auditable AI-infrastructure knowledge bundle containing -`News`, `Earnings`, and a research `Paper` tree. +`News`, `Earnings`, and a pre-V1 `LegacyPaper` tree. New paper ingestion uses +`PaperFlowResult` and `LocalKnowledgeLibrary.put_paper()`; the legacy type is +retained only so this existing bundle remains readable. ## Run diff --git a/examples/library/data/ai_infrastructure.db b/examples/library/data/ai_infrastructure.db index 10f4eedc1727baf2eb29e2226b3f921cca2e7450..6b8340b15c0e70eec69bce9fe43048d1301101af 100644 GIT binary patch delta 4418 zcmd^C?`s=n9M3gP?s7?bP1o+4u3ew4-P)|tI;V^X+vwWs7TTs=mX1}F^P1eY+tMWF zlHr`1W}8D{;>Pnrd?Ug@d?PXjgK3B0TR{Z>0bl7~APk1+^V~g8`do6Uc1#dFD9t_3 z_xru)`TqKRzU%A$^yY}tycsvv6k3L=KM<;zj|2g#My;C2a+~Ilty%ir`k)GiF zz0#Ag#y<+a&L8%FwDz8*71_}Pwsc_0r%iueVpq!5`o`JgCixaq?c8#TrtjR z*@~&Y*t7tSbYr<%)3SQ7X<_*K$uKV;I>g;PT+vqZy1i`{tL3b&YNn}IOqZp>>9jJH zQ9x$u_`CvK1i-im(4%#jCL9t5dx*iZe%UaMVnH=>Pz&^SUNQ>SkKHBhysj3iYpZ%0 zBvTpXgp!V17S{Bfp}|yF6n+^3HPd)oS68o8bdySf)M5stmgnce^hsrU4vaG*M~{FP zVz!ypd~p>VpI1eL<-C)D@fJ8HFL-qNGr2ST1ibSr_nV5_PH%(x=D-h zXn25+DP4YfL>1ZxRZv!u{d-sA@cIr7#jtQ$_=qOVLGpD4%~n2^NKj{_k^!?f6>jA{9&B_-S;D)D&4vLzF*w( zb!*|Rcb%3Ggr&!xaK}|o_;cxD@H0p8ZFJTM^Z1J_J;p8(a;{Rg}Z$IT7 zwce8Q|4x0AUJ5eG*$f;LdbUyKJ*{Pq$}%9vgaE3Dh)4WG_*hsHBEGNq4Zi68)b~d4 zuJ=sGH)zXxyl3v!f8%QP{R6ywY?7@ z!>ECx)1ny%mKw?t0?`gAW^|h{*g<>jB)~af`}8iO8`f-}mz zf;dlFIh{!-r!!(~Zl_;X6P(A&C-cp`Q9+am{wM{SSQqs?g~7(1 znvQ7^si1}~SsihU<{lb0BIF*;fW1qic(7D74P3HqnN>r8@G6TqU{=;X>Wwdu7~7l! zuwjb>YWHdLAcFWWYcWP>*jz6T53Z?~n<|00h0>R1pYEI6K{XR?G@EF&N-T~BXFLLK zkI$N=-?$|4^27x9Vb~=@no`5u?UoaCQfSQukb$E@UmJ=pQ0Kr%Ad2BdLTxBl%_1V4 zEf&rh=i%3(C|nN#OdWu?j@loG}30`uqwSm>XT*yX}e6LZfcDvKb#*_FKFfk zw_Is8QK^k<@8wN=*jhw|a><&w1ti-r00QiFt4QZD(Bai1u6(Pb;LN5#khC^AS zqS<8Nq6trlJH^d|U}a&gnA1@6yn;tU_(@;{ejga "PaperFlowCfg": + if self.chunk_overlap >= self.chunk_size: + raise ValueError("chunk_overlap must be smaller than chunk_size") + if self.min_summary_pages > self.min_summary_citations: + raise ValueError( + "min_summary_pages cannot exceed min_summary_citations" + ) + return self diff --git a/quantmind/flows/__init__.py b/quantmind/flows/__init__.py index 6290275..8889955 100644 --- a/quantmind/flows/__init__.py +++ b/quantmind/flows/__init__.py @@ -7,16 +7,21 @@ - ``batch_run`` runs any flow over a list of inputs with bounded concurrency and aggregated results. - ``BatchResult`` is the shape returned by ``batch_run``. -- ``UnsupportedContentTypeError`` is raised when ``paper_flow`` cannot - route fetched bytes through the format layer. +- ``UnsupportedContentTypeError`` is raised when ``paper_flow`` does not + resolve a page-aware PDF. """ from quantmind.flows.batch import BatchResult, batch_run from quantmind.flows.news import collect_news -from quantmind.flows.paper import UnsupportedContentTypeError, paper_flow +from quantmind.flows.paper import ( + PaperCitationValidationError, + UnsupportedContentTypeError, + paper_flow, +) __all__ = [ "BatchResult", + "PaperCitationValidationError", "UnsupportedContentTypeError", "batch_run", "collect_news", diff --git a/quantmind/flows/_paper_summary.py b/quantmind/flows/_paper_summary.py new file mode 100644 index 0000000..179b8b7 --- /dev/null +++ b/quantmind/flows/_paper_summary.py @@ -0,0 +1,243 @@ +"""Bounded Agents SDK synthesis for one paper chunk-set artifact.""" + +import asyncio +import hashlib +import json +from dataclasses import replace +from typing import Any, Protocol + +from agents import Agent, ModelSettings, function_tool +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from quantmind.configs import PaperFlowCfg +from quantmind.flows._runner import run_with_observability +from quantmind.knowledge import PaperChunkSet, PaperSourceRevision + +_SUMMARY_INSTRUCTIONS = """\ +Write one accurate global summary of the supplied paper. Use +`read_chunk_group` to inspect source chunks before writing. Cover the central +contribution, architecture or methodology, principal results, and important +limitations. Return only summary prose plus citations. Each citation uses the +zero-based chunk index and a physical page owned by that chunk. Never invent +IDs, source metadata, storage links, pages, or quotations. Read useful +consecutive chunks in groups of up to eight; do not spend one call per chunk. +""" + + +class PaperSummaryBudgetExceeded(RuntimeError): + """The summary exceeded a configured runtime or usage boundary.""" + + +class PaperSummaryCitationDraft(BaseModel): + """Model-owned citation coordinates before code resolves canonical IDs.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + chunk_index: int = Field(ge=0) + page_number: int = Field(ge=1) + quote: str | None = Field(default=None, max_length=500) + + +class PaperSummaryDraft(BaseModel): + """Limited model output containing prose and non-canonical coordinates.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + summary: str = Field(min_length=1) + citations: tuple[PaperSummaryCitationDraft, ...] = Field(min_length=1) + + @field_validator("summary") + @classmethod + def _summary_is_not_blank(cls, value: str) -> str: + if not value.strip(): + raise ValueError("paper summary draft must not be blank") + return value + + +class _PaperSummaryProvider(Protocol): + """Deterministic test seam and production summary-provider boundary.""" + + async def summarize( + self, + source: PaperSourceRevision, + chunk_set: PaperChunkSet, + *, + cfg: PaperFlowCfg, + ) -> PaperSummaryDraft: + """Create one bounded draft from the selected chunk set.""" + ... + + +def _estimated_tokens(text: str) -> int: + return max(1, (len(text) + 3) // 4) + + +class _SummaryBudget: + """Concurrency-safe accounting for adaptive chunk reads.""" + + def __init__(self, cfg: PaperFlowCfg, initial_text: str) -> None: + initial_tokens = _estimated_tokens(initial_text) + if initial_tokens > cfg.max_summary_input_tokens: + raise PaperSummaryBudgetExceeded( + "paper summary manifest exceeds max_summary_input_tokens" + ) + self._max_calls = cfg.max_summary_tool_calls + self._max_input_tokens = cfg.max_summary_input_tokens + self._calls = 0 + self._input_tokens = initial_tokens + self._lock = asyncio.Lock() + self._semaphore = asyncio.Semaphore(cfg.max_summary_concurrency) + + async def read( + self, + chunk_set: PaperChunkSet, + *, + start: int, + count: int, + ) -> str: + """Return one bounded group while reserving calls and input tokens.""" + if start < 0 or count < 1 or count > 8: + raise ValueError( + "start must be non-negative and count must be 1..8" + ) + if start >= len(chunk_set.chunks): + raise ValueError("start is outside the chunk-set manifest") + async with self._semaphore: + selected = chunk_set.chunks[start : start + count] + payload = json.dumps( + [ + { + "chunk_index": chunk.position, + "pages": sorted( + {span.page_number for span in chunk.source_spans} + ), + "text": chunk.text, + } + for chunk in selected + ], + ensure_ascii=False, + ) + tokens = _estimated_tokens(payload) + async with self._lock: + if self._calls >= self._max_calls: + raise PaperSummaryBudgetExceeded( + "paper summary exceeded max_summary_tool_calls" + ) + if self._input_tokens + tokens > self._max_input_tokens: + raise PaperSummaryBudgetExceeded( + "paper summary exceeded max_summary_input_tokens" + ) + self._calls += 1 + self._input_tokens += tokens + return payload + + +def _summary_manifest( + source: PaperSourceRevision, + chunk_set: PaperChunkSet, + cfg: PaperFlowCfg, +) -> str: + manifest = [ + { + "chunk_index": chunk.position, + "pages": sorted({span.page_number for span in chunk.source_spans}), + "characters": len(chunk.text), + "preview": chunk.text[:160], + } + for chunk in chunk_set.chunks + ] + return json.dumps( + { + "title": source.title, + "authors": source.authors, + "page_count": len(source.parsed.pages), + "required_citations": cfg.min_summary_citations, + "required_source_pages": cfg.min_summary_pages, + "chunks": manifest, + }, + ensure_ascii=False, + ) + + +def _summary_instructions(cfg: PaperFlowCfg) -> str: + instructions = _SUMMARY_INSTRUCTIONS + if cfg.summary_instructions: + instructions = ( + f"{instructions}\n\nAdditional summary requirements:\n" + f"{cfg.summary_instructions}" + ) + return instructions + + +def _summary_instructions_hash(cfg: PaperFlowCfg) -> str: + return hashlib.sha256( + _summary_instructions(cfg).encode("utf-8") + ).hexdigest() + + +def _bounded_model_settings(cfg: PaperFlowCfg) -> ModelSettings: + settings = cfg.model_settings or ModelSettings() + configured = settings.max_tokens or cfg.max_summary_output_tokens + return replace( + settings, + max_tokens=min(configured, cfg.max_summary_output_tokens), + parallel_tool_calls=True, + ) + + +class _AgentsPaperSummaryProvider: + """Use one SDK agent with bounded adaptive access to chunk groups.""" + + async def summarize( + self, + source: PaperSourceRevision, + chunk_set: PaperChunkSet, + *, + cfg: PaperFlowCfg, + ) -> PaperSummaryDraft: + manifest = _summary_manifest(source, chunk_set, cfg) + budget = _SummaryBudget(cfg, manifest) + + @function_tool + async def read_chunk_group(start: int, count: int) -> str: + """Read up to eight consecutive paper chunks. + + Args: + start: Zero-based first chunk index. + count: Number of chunks to read, from one through eight. + """ + return await budget.read( + chunk_set, + start=start, + count=count, + ) + + agent: Agent[Any] = Agent( + name="paper_global_summarizer", + instructions=_summary_instructions(cfg), + model=cfg.model, + model_settings=_bounded_model_settings(cfg), + tools=[read_chunk_group], + output_type=PaperSummaryDraft, + ) + try: + output = await asyncio.wait_for( + run_with_observability( + agent, + manifest, + cfg=cfg, + memory=None, + extra_run_hooks=[], + ), + timeout=cfg.timeout_seconds, + ) + except asyncio.TimeoutError as exc: + raise PaperSummaryBudgetExceeded( + "paper summary exceeded timeout_seconds" + ) from exc + draft = PaperSummaryDraft.model_validate(output) + if _estimated_tokens(draft.summary) > cfg.max_summary_output_tokens: + raise PaperSummaryBudgetExceeded( + "paper summary exceeded max_summary_output_tokens" + ) + return draft diff --git a/quantmind/flows/paper.py b/quantmind/flows/paper.py index 5b8373b..6638dee 100644 --- a/quantmind/flows/paper.py +++ b/quantmind/flows/paper.py @@ -1,19 +1,12 @@ -"""Paper extraction flow. +"""Source-first paper flow that returns chunks before a cited summary.""" -`paper_flow` ingests one of the ``PaperInput`` discriminated-union -variants, fetches and converts the raw payload to markdown via -``preprocess.fetch`` + ``preprocess.format``, then runs an -``Agent(output_type=Paper)`` to produce a typed ``Paper`` -``TreeKnowledge`` object. - -Customization happens through the configured ``PaperFlowCfg`` (Layer 1) -or the keyword arguments on this function (Layer 2). To swap the whole -flow, fork this file (Layer 3 — design doc §9). -""" - -from typing import Any, TypeVar - -from agents import Agent, RunHooks, Tool +import hashlib +import mimetypes +from dataclasses import dataclass +from datetime import datetime, timezone +from importlib.metadata import version +from pathlib import Path +from typing import Literal from quantmind.configs import PaperFlowCfg from quantmind.configs.paper import ( @@ -24,180 +17,496 @@ PaperInput, RawText, ) -from quantmind.flows._runner import run_with_observability -from quantmind.knowledge import Paper +from quantmind.flows._paper_summary import ( + PaperSummaryDraft, + _AgentsPaperSummaryProvider, + _PaperSummaryProvider, + _summary_instructions_hash, +) +from quantmind.knowledge import ( + ArtifactLocator, + PaperAssetRef, + PaperBoundingBox, + PaperChunk, + PaperChunkingConfig, + PaperChunkSet, + PaperCitation, + PaperFlowResult, + PaperGlobalSummary, + PaperParsedBlock, + PaperParsedManifest, + PaperParsedPage, + PaperSourceRevision, + PaperSourceSpan, + PaperSummaryProducer, + SourceRef, +) +from quantmind.knowledge.paper import ( + _paper_artifact_id, + _paper_asset_id, + _paper_chunk_id, + _paper_chunk_set_content_hash, + _paper_source_id, + _paper_summary_content_hash, + _stable_hash, + _text_hash, +) from quantmind.preprocess.fetch import ( Fetched, + RawPaper, fetch_arxiv, fetch_url, read_local_file, ) -from quantmind.preprocess.format import html_to_markdown, pdf_to_markdown +from quantmind.preprocess.format import ParsedDocument, parse_pdf +from quantmind.rag import SentenceSplitterConfig, chunk_parsed_document + -P = TypeVar("P", bound=Paper) +class UnsupportedContentTypeError(ValueError): + """The source is not a page-aware PDF supported by Paper Flow V1.""" -_DEFAULT_INSTRUCTIONS = """\ -You are extracting a research paper into a structured QuantMind ``Paper`` -TreeKnowledge object. Build the section tree top-down: every node has a -title and a short summary; leaf nodes additionally carry the section -markdown content. Cite supporting passages on each node. -Honour these flags from the run config: -- extract_methodology={extract_methodology}: when true, every methodology - section becomes its own subtree with a per-step summary. -- extract_limitations={extract_limitations}: when true, surface - limitations as a dedicated top-level child rather than inlining them. -- asset_class_hint={asset_class_hint!r}: when set, prefer this asset - class for ``Paper.asset_classes`` if the paper does not state one - explicitly. +class PaperCitationValidationError(ValueError): + """A generated summary did not provide valid source coverage.""" -Set ``as_of`` to the publication date when given; otherwise use today's -date. Set the ``source`` provenance ref using the metadata supplied in -the prompt. -""" +@dataclass(frozen=True) +class _FetchedPaperSource: + """Exact fetched bytes and code-owned source facts before parsing.""" -class UnsupportedContentTypeError(ValueError): - """Fetched bytes have a content type paper_flow cannot route to a parser.""" + bytes: bytes + media_type: str + kind: Literal["arxiv", "http", "local"] + uri: str + fetched_at: datetime + available_at: datetime + published_at: datetime | None = None + arxiv_id: str | None = None + title: str | None = None + authors: tuple[str, ...] = () async def paper_flow( input: PaperInput, *, cfg: PaperFlowCfg | None = None, - extra_tools: list[Tool] | None = None, - extra_instructions: str | None = None, - output_type: type[P] | None = None, - memory: object | None = None, - extra_run_hooks: list[RunHooks[Any]] | None = None, - extra_input_guardrails: list[Any] | None = None, - extra_output_guardrails: list[Any] | None = None, -) -> P | Paper: - """Extract a ``Paper`` from a typed ``PaperInput``. - - See design doc §4.1 for the rationale on each kwarg. ``memory`` is a - PR6 placeholder — non-None values are accepted but unused in PR5. + _summary_provider: _PaperSummaryProvider | None = None, +) -> PaperFlowResult: + """Build a page-aware chunk set and one cited global summary. + + IDs, source metadata, artifact membership, lineage, and citation links are + created and validated by code. The model returns only summary prose and + chunk/page coordinates through a bounded summarization seam. + + Args: + input: Typed paper source. V1 requires a PDF-backed input. + cfg: Splitter, summary model, and explicit usage/runtime limits. + + Returns: + The exact source revision, one chunk-set artifact, and one cited + global-summary artifact. Raises: - UnsupportedContentTypeError: When fetched bytes are not PDF / - HTML / markdown / plain-text. - NotImplementedError: When ``input`` is a ``DoiIdentifier`` (the - unpaywall fallback is its own follow-up issue). + UnsupportedContentTypeError: If the resolved content is not a PDF. + PaperCitationValidationError: If generated citations are invalid or + do not meet configured source-coverage requirements. + NotImplementedError: If a DOI input has no exact open PDF resolver. """ cfg = cfg or PaperFlowCfg() - out_type: type[Paper] = output_type or Paper # type: ignore[assignment] - - raw_md, source_meta = await _fetch_and_format(input) - - # Agent's `model_settings` parameter is non-optional (defaults to a - # fresh ``ModelSettings()``); only forward when cfg has one set. - agent_kwargs: dict[str, Any] = { - "name": "paper_extractor", - "instructions": _compose_instructions( - _DEFAULT_INSTRUCTIONS, extra_instructions, cfg + fetched = await _fetch_paper_source(input) + parsed = await parse_pdf( + fetched.bytes, + artifact_dir=cfg.output_dir, + ) + source = _build_source_revision(fetched, parsed) + parsed_chunks = chunk_parsed_document( + parsed, + config=SentenceSplitterConfig( + chunk_size=cfg.chunk_size, + chunk_overlap=cfg.chunk_overlap, ), - "model": cfg.model, - "tools": list(extra_tools or []), - "output_type": out_type, - "input_guardrails": list(extra_input_guardrails or []), - "output_guardrails": list(extra_output_guardrails or []), - } - if cfg.model_settings is not None: - agent_kwargs["model_settings"] = cfg.model_settings - agent: Agent[Any] = Agent(**agent_kwargs) - return await run_with_observability( - agent, - _format_input(raw_md, source_meta), - cfg=cfg, - memory=memory, - extra_run_hooks=list(extra_run_hooks or []), + ) + chunk_set = _build_chunk_set(source, parsed_chunks, cfg) + provider = _summary_provider or _AgentsPaperSummaryProvider() + draft = await provider.summarize(source, chunk_set, cfg=cfg) + summary = _build_global_summary(source, chunk_set, draft, cfg) + return PaperFlowResult( + source_revision=source, + chunk_set=chunk_set, + global_summary=summary, ) -async def _fetch_and_format( - input: PaperInput, -) -> tuple[str, dict[str, Any]]: - """Dispatch on the input variant; return (markdown, source metadata).""" +def _aware_or_now(value: datetime | None) -> datetime: + if value is None: + return datetime.now(timezone.utc) + if value.tzinfo is None or value.utcoffset() is None: + return value.replace(tzinfo=timezone.utc) + return value.astimezone(timezone.utc) + + +def _require_pdf(raw: Fetched) -> None: + media_type = (raw.content_type or "").lower() + if not media_type.startswith("application/pdf"): + raise UnsupportedContentTypeError( + "Paper Flow V1 requires a page-aware PDF; resolved content type " + f"was {media_type!r}" + ) + + +async def _fetch_paper_source(input: PaperInput) -> _FetchedPaperSource: if isinstance(input, ArxivIdentifier): - raw = await fetch_arxiv(input.id) - md = await pdf_to_markdown(raw.bytes) - return md, { - "source": "arxiv", - "arxiv_id": raw.arxiv_id, - "title": raw.title, - "authors": list(raw.authors), - } + raw_paper: RawPaper = await fetch_arxiv(input.id) + _require_pdf(raw_paper) + fetched_at = _aware_or_now(raw_paper.fetched_at) + available_at = _aware_or_now( + raw_paper.updated_at or raw_paper.published_at or fetched_at + ) + uri = raw_paper.resolved_url or raw_paper.source_url + if uri is None: + raise ValueError("resolved arXiv paper is missing its source URL") + return _FetchedPaperSource( + bytes=raw_paper.bytes, + media_type="application/pdf", + kind="arxiv", + uri=uri, + fetched_at=fetched_at, + available_at=available_at, + published_at=raw_paper.published_at, + arxiv_id=raw_paper.arxiv_id, + title=raw_paper.title, + authors=raw_paper.authors, + ) if isinstance(input, HttpUrl): - raw = await fetch_url(input.url) - md = await _format_by_content_type(raw) - return md, { - "source": "web", - "url": input.url, - "content_type": raw.content_type, - } + raw_http = await fetch_url(input.url) + _require_pdf(raw_http) + fetched_at = _aware_or_now(raw_http.fetched_at) + return _FetchedPaperSource( + bytes=raw_http.bytes, + media_type="application/pdf", + kind="http", + uri=raw_http.resolved_url or raw_http.source_url or input.url, + fetched_at=fetched_at, + available_at=fetched_at, + ) if isinstance(input, LocalFilePath): - raw = await read_local_file(input.path) - md = await _format_by_content_type(raw) - return md, { - "source": "local", - "path": str(input.path), - "content_type": raw.content_type, - } + raw_local = await read_local_file(input.path) + _require_pdf(raw_local) + observed_at = _aware_or_now(raw_local.fetched_at) + path = Path(input.path).expanduser().resolve() + return _FetchedPaperSource( + bytes=raw_local.bytes, + media_type="application/pdf", + kind="local", + uri=raw_local.source_url or path.as_uri(), + fetched_at=observed_at, + available_at=observed_at, + ) if isinstance(input, RawText): - return input.text, {"source": "inline"} + raise UnsupportedContentTypeError( + "Paper Flow V1 requires a page-aware PDF; RawText has no physical " + "page evidence" + ) if isinstance(input, DoiIdentifier): - # PR4's CrossrefMetadata exposes only `primary_url` (publisher - # landing page), not a direct PDF link. Adding the unpaywall - # fallback that turns a DOI into an OA PDF URL is its own - # follow-up issue. raise NotImplementedError( - "DOI inputs require an OA PDF resolver (unpaywall fallback) " - "which is tracked as a PR4 follow-up. Use ArxivIdentifier or " - "HttpUrl for now." + "DOI inputs require an exact open PDF resolver before they can " + "produce a paper source revision" ) raise TypeError(f"Unsupported PaperInput variant: {type(input)!r}") -async def _format_by_content_type(raw: Fetched) -> str: - """Route a ``Fetched`` payload through the right format helper.""" - ct = (raw.content_type or "").lower() - if ct.startswith("application/pdf"): - return await pdf_to_markdown(raw.bytes) - if ct.startswith("text/html"): - return await html_to_markdown( - raw.bytes.decode("utf-8", errors="replace") +def _asset_from_path( + path_value: str, + *, + source_revision_id, + kind: Literal["screenshot", "image"], + page_number: int, + assets: dict, + blobs: dict[str, bytes], +) -> PaperAssetRef: + path = Path(path_value) + try: + content = path.read_bytes() + except OSError as exc: + raise RuntimeError( + f"Parser asset for page {page_number} is missing: {path}" + ) from exc + content_hash = hashlib.sha256(content).hexdigest() + media_type = ( + mimetypes.guess_type(path.name)[0] or "application/octet-stream" + ) + asset = PaperAssetRef( + asset_id=_paper_asset_id( + source_revision_id, + kind=kind, + page_number=page_number, + content_hash=content_hash, + ), + kind=kind, + media_type=media_type, + content_hash=content_hash, + size_bytes=len(content), + page_number=page_number, + ) + assets[asset.asset_id] = asset + blobs[content_hash] = content + return asset + + +def _build_source_revision( + fetched: _FetchedPaperSource, + parsed: ParsedDocument, +) -> PaperSourceRevision: + source_id = _paper_source_id(parsed.source_hash) + blobs = {parsed.source_hash: fetched.bytes} + raw = PaperAssetRef( + asset_id=_paper_asset_id( + source_id, + kind="raw", + page_number=None, + content_hash=parsed.source_hash, + ), + kind="raw", + media_type=fetched.media_type, + content_hash=parsed.source_hash, + size_bytes=len(fetched.bytes), + ) + assets = {raw.asset_id: raw} + pages: list[PaperParsedPage] = [] + for page in parsed.pages: + screenshot = ( + _asset_from_path( + page.screenshot_path, + source_revision_id=source_id, + kind="screenshot", + page_number=page.page_number, + assets=assets, + blobs=blobs, + ) + if page.screenshot_path is not None + else None + ) + images = tuple( + _asset_from_path( + image_path, + source_revision_id=source_id, + kind="image", + page_number=page.page_number, + assets=assets, + blobs=blobs, + ) + for image_path in page.image_paths + ) + pages.append( + PaperParsedPage( + page_number=page.page_number, + width=page.width, + height=page.height, + text=page.text, + blocks=tuple( + PaperParsedBlock( + text=block.text, + bbox=PaperBoundingBox( + x0=block.bbox.x0, + y0=block.bbox.y0, + x1=block.bbox.x1, + y1=block.bbox.y1, + ), + font_name=block.font_name, + font_size=block.font_size, + confidence=block.confidence, + ) + for block in page.blocks + ), + screenshot_asset_id=( + screenshot.asset_id if screenshot is not None else None + ), + image_asset_ids=tuple(image.asset_id for image in images), + ) ) - if ct.startswith("text/markdown") or ct.startswith("text/plain"): - return raw.bytes.decode("utf-8", errors="replace") - raise UnsupportedContentTypeError( - f"Unsupported content-type for paper input: {ct!r}" + source_ref = SourceRef( + kind=fetched.kind, + uri=fetched.uri, + fetched_at=fetched.fetched_at, + content_hash=parsed.source_hash, + ) + return PaperSourceRevision( + id=source_id, + source=source_ref, + as_of=fetched.published_at or fetched.available_at, + available_at=fetched.available_at, + published_at=fetched.published_at, + arxiv_id=fetched.arxiv_id, + title=fetched.title, + authors=fetched.authors, + parsed=PaperParsedManifest( + source_hash=parsed.source_hash, + parser_name=parsed.parser_name, + parser_version=parsed.parser_version, + cleanup_version=parsed.cleanup_version, + pages=tuple(pages), + ), + raw_asset_id=raw.asset_id, + assets=tuple(assets.values()), + blobs=blobs, ) -def _compose_instructions( - base: str, extra: str | None, cfg: PaperFlowCfg -) -> str: - """Render the system instructions, appending ``extra`` if provided.""" - instructions = base.format( - extract_methodology=cfg.extract_methodology, - extract_limitations=cfg.extract_limitations, - asset_class_hint=cfg.asset_class_hint, +def _build_chunk_set( + source: PaperSourceRevision, + parsed_chunks, + cfg: PaperFlowCfg, +) -> PaperChunkSet: + producer = PaperChunkingConfig( + splitter_version=version("llama-index-core"), + chunk_size=cfg.chunk_size, + chunk_overlap=cfg.chunk_overlap, + ) + producer_hash = _stable_hash(producer.model_dump(mode="json")) + artifact_id = _paper_artifact_id( + source.id, + "paper_chunk_set", + producer_hash, + ) + page_assets = { + page.page_number: tuple( + asset_id + for asset_id in ( + (page.screenshot_asset_id,) + if page.screenshot_asset_id is not None + else () + ) + + page.image_asset_ids + ) + for page in source.parsed.pages + } + chunks: list[PaperChunk] = [] + for position, parsed_chunk in enumerate(parsed_chunks): + if parsed_chunk.source_hash != source.source.content_hash: + raise ValueError("parsed chunk belongs to another source revision") + span = PaperSourceSpan( + page_number=parsed_chunk.page_number, + start_char=parsed_chunk.start_char, + end_char=parsed_chunk.end_char, + block_boxes=tuple( + PaperBoundingBox( + x0=box.x0, + y0=box.y0, + x1=box.x1, + y1=box.y1, + ) + for box in parsed_chunk.block_boxes + ), + asset_ids=page_assets.get(parsed_chunk.page_number, ()), + ) + content_hash = _text_hash(parsed_chunk.text) + chunks.append( + PaperChunk( + chunk_id=_paper_chunk_id( + artifact_id, + position=position, + content_hash=content_hash, + spans=(span,), + ), + chunk_set_id=artifact_id, + source_revision_id=source.id, + position=position, + text=parsed_chunk.text, + content_hash=content_hash, + source_spans=(span,), + ) + ) + if not chunks: + raise ValueError("paper source produced no non-empty chunks") + chunk_tuple = tuple(chunks) + return PaperChunkSet( + id=artifact_id, + source_revision_id=source.id, + producer=producer, + producer_config_hash=producer_hash, + content_hash=_paper_chunk_set_content_hash(chunk_tuple), + chunks=chunk_tuple, ) - if extra: - instructions = f"{instructions}\n\nAdditional instructions:\n{extra}" - return instructions -def _format_input(raw_md: str, source_meta: dict[str, Any]) -> str: - """Concatenate metadata + content into the prompt the agent sees.""" - lines: list[str] = [] - for key, value in source_meta.items(): - if value is None: +def _build_global_summary( + source: PaperSourceRevision, + chunk_set: PaperChunkSet, + draft: PaperSummaryDraft, + cfg: PaperFlowCfg, +) -> PaperGlobalSummary: + citations: list[PaperCitation] = [] + seen: set[tuple[int, int, str | None]] = set() + for draft_citation in draft.citations: + try: + chunk = chunk_set.chunks[draft_citation.chunk_index] + except IndexError as exc: + raise PaperCitationValidationError( + "paper summary cites an unknown chunk index" + ) from exc + pages = {span.page_number for span in chunk.source_spans} + if draft_citation.page_number not in pages: + raise PaperCitationValidationError( + "paper summary citation page is not owned by its chunk" + ) + quote = draft_citation.quote + if quote is not None and quote not in chunk.text: + raise PaperCitationValidationError( + "paper summary citation quote is not present in its chunk" + ) + key = (draft_citation.chunk_index, draft_citation.page_number, quote) + if key in seen: continue - if isinstance(value, (list, tuple)): - value = ", ".join(map(str, value)) - lines.append(f"{key}: {value}") - header = "\n".join(lines) - return ( - f"--- Source metadata ---\n{header}\n\n--- Paper content ---\n{raw_md}" + seen.add(key) + citations.append( + PaperCitation( + chunk_set_id=chunk_set.id, + chunk_id=chunk.chunk_id, + page_number=draft_citation.page_number, + quote=quote, + ) + ) + cited_pages = {citation.page_number for citation in citations} + if len(citations) < cfg.min_summary_citations: + raise PaperCitationValidationError( + "paper summary has fewer citations than min_summary_citations" + ) + if len(cited_pages) < cfg.min_summary_pages: + raise PaperCitationValidationError( + "paper summary has fewer source pages than min_summary_pages" + ) + + producer = PaperSummaryProducer( + model=cfg.model, + prompt_version=cfg.summary_prompt_version, + input_chunk_set_id=chunk_set.id, + instructions_hash=_summary_instructions_hash(cfg), + max_output_tokens=cfg.max_summary_output_tokens, + ) + producer_hash = _stable_hash(producer.model_dump(mode="json")) + artifact_id = _paper_artifact_id( + source.id, + "paper_summary", + producer_hash, + ) + citation_tuple = tuple(citations) + summary_text = draft.summary.strip() + return PaperGlobalSummary( + id=artifact_id, + source_revision_id=source.id, + producer=producer, + producer_config_hash=producer_hash, + content_hash=_paper_summary_content_hash( + summary_text, + citation_tuple, + ), + summary=summary_text, + citations=citation_tuple, + derived_from=( + ArtifactLocator( + source_revision_id=source.id, + artifact_id=chunk_set.id, + artifact_kind="paper_chunk_set", + ), + ), ) diff --git a/quantmind/knowledge/__init__.py b/quantmind/knowledge/__init__.py index 972a535..9e904c5 100644 --- a/quantmind/knowledge/__init__.py +++ b/quantmind/knowledge/__init__.py @@ -2,14 +2,14 @@ The standard defines three shapes that share `BaseKnowledge`: -- `FlattenKnowledge` — atomic cards (`News`, `Earnings`, `PaperKnowledgeCard`). -- `TreeKnowledge` — hierarchical artifacts (`Paper`). +- `FlattenKnowledge` — atomic cards (`News`, `Earnings`, `Factor`, `Thesis`). +- `TreeKnowledge` — hierarchical artifacts such as future navigation trees. - `GraphKnowledge` — cross-item edges (placeholder, not implemented). Every concrete subclass is frozen Pydantic v2 with ``extra="forbid"``, suitable for ``Agent(output_type=...)`` and round-tripping through JSON. -Subclasses MUST override ``embedding_text()`` so the store layer knows -what to embed. +Paper sources and artifacts are separate immutable models. Search projection +text and vectors remain rebuildable library-owned data. """ from quantmind.knowledge._base import ( @@ -24,7 +24,27 @@ from quantmind.knowledge.earnings import Earnings from quantmind.knowledge.factor import Factor from quantmind.knowledge.news import News -from quantmind.knowledge.paper import Paper, PaperKnowledgeCard +from quantmind.knowledge.paper import ( + ArtifactLocator, + LegacyPaper, + PaperArtifact, + PaperArtifactKind, + PaperAssetRef, + PaperBoundingBox, + PaperChunk, + PaperChunkingConfig, + PaperChunkSet, + PaperCitation, + PaperFlowResult, + PaperGlobalSummary, + PaperParsedBlock, + PaperParsedManifest, + PaperParsedPage, + PaperSourceRevision, + PaperSourceSpan, + PaperSummaryProducer, + ResolvedPaperArtifact, +) from quantmind.knowledge.thesis import Thesis __all__ = [ @@ -42,7 +62,24 @@ "Earnings", "Factor", "News", - "Paper", - "PaperKnowledgeCard", + "ArtifactLocator", + "LegacyPaper", + "PaperArtifact", + "PaperArtifactKind", + "PaperAssetRef", + "PaperBoundingBox", + "PaperChunk", + "PaperChunkingConfig", + "PaperChunkSet", + "PaperCitation", + "PaperFlowResult", + "PaperGlobalSummary", + "PaperParsedBlock", + "PaperParsedManifest", + "PaperParsedPage", + "PaperSourceRevision", + "PaperSourceSpan", + "PaperSummaryProducer", + "ResolvedPaperArtifact", "Thesis", ] diff --git a/quantmind/knowledge/_base.py b/quantmind/knowledge/_base.py index b6ae31c..c449c08 100644 --- a/quantmind/knowledge/_base.py +++ b/quantmind/knowledge/_base.py @@ -1,21 +1,22 @@ """Base types for the quantmind.knowledge data standard. -Three shapes share `BaseKnowledge`: +Three conventional shapes share `BaseKnowledge`: -- `FlattenKnowledge` — atomic cards (`News`, `Earnings`, `PaperKnowledgeCard`, ...) -- `TreeKnowledge` — hierarchical artifacts (`Paper`, `EarningsCallTranscript`, ...) +- `FlattenKnowledge` — atomic cards (`News`, `Earnings`, `Factor`, ...) +- `TreeKnowledge` — hierarchical conventional artifacts - `GraphKnowledge` — cross-item edges (placeholder, future PR) +Source-first paper revisions and artifacts use separate immutable models rather +than inheriting `BaseKnowledge`. + The `as_of` field is mandatory by design: financial knowledge is only useful when its information cutoff is explicit. Optional `available_at` records when the source became observable so research can prevent look-ahead independently from information time. `source` and `extraction` are typed references (not bare strings) so dedup, audit, and re-runs all have stable keys. -Subclasses MUST override `embedding_text()` to declare what string the store -layer should embed for them. The contract is enforced at runtime, not at -type-check time, because `BaseKnowledge` itself can be referenced as a -generic return type without forcing every consumer to know about embeddings. +Canonical knowledge does not choose text or store vectors for a retrieval +method. Search projections are rebuildable library-owned derived data. """ from datetime import datetime, timedelta, timezone @@ -68,8 +69,8 @@ class BaseKnowledge(BaseModel): """Root of every quantmind knowledge type. All three shapes (Flatten / Tree / Graph) share this field set; subclasses - add domain-specific payload. Subclasses MUST override `embedding_text()` - to declare what string the store layer should embed. + add domain-specific payload. Retrieval-specific text selection belongs to + the library indexing boundary. """ model_config = ConfigDict(extra="forbid", frozen=True) @@ -101,17 +102,6 @@ class BaseKnowledge(BaseModel): tags: list[str] = Field(default_factory=list) disclaimers: list[str] = Field(default_factory=list) - def embedding_text(self) -> str: - """Return the canonical string the store should embed for this item. - - Flatten subclasses: typically ``summary + key attrs``. - Tree subclasses: typically ``root.title + root.summary``. - Subclasses MUST override. - """ - raise NotImplementedError( - f"{type(self).__name__}.embedding_text() must be overridden" - ) - # ── Convenience helpers ──────────────────────────────────────────── # These are shared by every shape (Flatten / Tree / Graph) because they # operate on `BaseKnowledge` metadata, not domain payload. diff --git a/quantmind/knowledge/_flatten.py b/quantmind/knowledge/_flatten.py index cd1129d..c266b9c 100644 --- a/quantmind/knowledge/_flatten.py +++ b/quantmind/knowledge/_flatten.py @@ -2,9 +2,8 @@ A flatten card is semantically indivisible: one source artifact maps to one card whose body is the answer (no structure to navigate). Flatten is the -right shape for `News`, `Earnings`, `Factor`, `Thesis`, and the summary card -of a paper (`PaperKnowledgeCard`). Whole research papers are NOT flatten — -they are `TreeKnowledge`. +right shape for `News`, `Earnings`, `Factor`, and `Thesis`. Paper source +revisions and artifacts use their dedicated source-first models. This file only declares the marker base; concrete subclasses live in `paper.py` / `news.py` / `earnings.py` / etc. @@ -16,7 +15,6 @@ class FlattenKnowledge(BaseKnowledge): """Marker base for flat domain cards. - Subclasses add a typed payload (e.g. ``summary``, ``methodology``, - ``revenue``). They MUST override `embedding_text()` to produce a stable - string suitable for the store's vector index. + Subclasses add a typed payload (e.g. ``headline``, ``guidance``, or + ``revenue``). Search text is selected by the library projection boundary. """ diff --git a/quantmind/knowledge/_tree.py b/quantmind/knowledge/_tree.py index 6725528..74720d8 100644 --- a/quantmind/knowledge/_tree.py +++ b/quantmind/knowledge/_tree.py @@ -1,9 +1,9 @@ """TreeKnowledge — hierarchical-artifact shape. A tree's structure carries information: nodes derive meaning from their -position under a parent. `TreeKnowledge` is the right shape for whole -research papers (sections / subsections), regulatory filings (10-K parts), -and earnings-call transcripts (intro / Q&A / per-question). +position under a parent. `TreeKnowledge` is the right shape for regulatory +filings (10-K parts), earnings-call transcripts (intro / Q&A / per-question), +and a future paper-navigation artifact. Retrieval over a tree is reasoning-driven (PageIndex-style): an agent reads the root summary plus children summaries, picks the most likely branch, @@ -39,10 +39,6 @@ class TreeNode(BaseModel): citations: list[Citation] = Field(default_factory=list) children_ids: list[UUID] = Field(default_factory=list) - def embedding_text(self) -> str: - """Default: title + summary. Override per domain if needed.""" - return f"{self.title}\n{self.summary}" - class TreeKnowledge(BaseKnowledge): """Hierarchical knowledge artifact. @@ -85,7 +81,3 @@ def find_path(self, node_id: UUID) -> list[TreeNode]: cursor = node.parent_id path.reverse() return path - - def embedding_text(self) -> str: - """Default: root node's embedding text. Override per domain if needed.""" - return self.root().embedding_text() diff --git a/quantmind/knowledge/earnings.py b/quantmind/knowledge/earnings.py index 68a5e88..3c4aabe 100644 --- a/quantmind/knowledge/earnings.py +++ b/quantmind/knowledge/earnings.py @@ -19,7 +19,3 @@ class Earnings(FlattenKnowledge): guidance: str | None = None surprise_flags: list[str] = Field(default_factory=list) transcript_quote: str | None = None - - def embedding_text(self) -> str: - guidance = self.guidance or "" - return f"{self.ticker} {self.period} earnings\n{guidance}".strip() diff --git a/quantmind/knowledge/factor.py b/quantmind/knowledge/factor.py index ad81d54..5625bbd 100644 --- a/quantmind/knowledge/factor.py +++ b/quantmind/knowledge/factor.py @@ -17,7 +17,3 @@ class Factor(FlattenKnowledge): factor_name: str universe: str | None = None - - def embedding_text(self) -> str: - scope = self.universe or "unspecified" - return f"factor {self.factor_name} on {scope}" diff --git a/quantmind/knowledge/news.py b/quantmind/knowledge/news.py index 2bf2de9..a6d8083 100644 --- a/quantmind/knowledge/news.py +++ b/quantmind/knowledge/news.py @@ -19,7 +19,3 @@ class News(FlattenKnowledge): entities: list[str] = Field(default_factory=list) sentiment: Literal["positive", "neutral", "negative"] = "neutral" materiality: Literal["low", "medium", "high"] = "medium" - - def embedding_text(self) -> str: - entities = ", ".join(self.entities) - return f"{self.headline}\n{self.event_type}\n{entities}".strip() diff --git a/quantmind/knowledge/paper.py b/quantmind/knowledge/paper.py index 5345bdf..184f06c 100644 --- a/quantmind/knowledge/paper.py +++ b/quantmind/knowledge/paper.py @@ -1,52 +1,588 @@ -"""Paper knowledge schemas. - -A whole paper is `TreeKnowledge` (sections + subsections); the distilled -summary card is `PaperKnowledgeCard` (`FlattenKnowledge`). They are separate -items linked by ``PaperKnowledgeCard.paper_id == Paper.id``. - -`paper_flow` (PR5) typically produces a `Paper` first (sectioning + per-node -summarisation), then a `PaperKnowledgeCard` derived from the root summary. -""" +"""Source revisions and independently versioned paper artifacts.""" +import hashlib +import json +import re +from datetime import datetime from typing import Literal -from uuid import UUID +from uuid import NAMESPACE_URL, UUID, uuid5 -from pydantic import Field +from pydantic import ( + BaseModel, + ConfigDict, + Field, + field_validator, + model_validator, +) -from quantmind.knowledge._flatten import FlattenKnowledge +from quantmind.knowledge._base import SourceRef from quantmind.knowledge._tree import TreeKnowledge +PaperArtifactKind = Literal["paper_chunk_set", "paper_summary"] + + +def _stable_hash(value: object) -> str: + """Return a deterministic SHA-256 hash for one JSON-compatible value.""" + payload = json.dumps( + value, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _text_hash(value: str) -> str: + return hashlib.sha256(value.encode("utf-8")).hexdigest() + + +def _paper_source_id(content_hash: str) -> UUID: + return uuid5(NAMESPACE_URL, f"quantmind:paper-source:{content_hash}") + + +def _paper_artifact_id( + source_revision_id: UUID, + artifact_kind: PaperArtifactKind, + producer_config_hash: str, +) -> UUID: + return uuid5( + source_revision_id, + f"quantmind:{artifact_kind}:{producer_config_hash}", + ) + + +def _paper_asset_id( + source_revision_id: UUID, + *, + kind: str, + page_number: int | None, + content_hash: str, +) -> UUID: + page = page_number if page_number is not None else "document" + return uuid5( + source_revision_id, + f"quantmind:paper-asset:{kind}:{page}:{content_hash}", + ) + + +class PaperBoundingBox(BaseModel): + """One source rectangle in top-left-origin page coordinates.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + x0: float + y0: float + x1: float + y1: float + + @model_validator(mode="after") + def _coordinates_are_ordered(self) -> "PaperBoundingBox": + if self.x1 < self.x0 or self.y1 < self.y0: + raise ValueError("bounding-box coordinates must be ordered") + return self + + +class PaperAssetRef(BaseModel): + """Content-addressed source, screenshot, or extracted-image reference.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + asset_id: UUID + kind: Literal["raw", "screenshot", "image"] + media_type: str + content_hash: str = Field(pattern=r"^[0-9a-f]{64}$") + size_bytes: int = Field(ge=0) + page_number: int | None = Field(default=None, ge=1) + + +class PaperParsedBlock(BaseModel): + """One parser-owned text block in the durable source manifest.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + text: str + bbox: PaperBoundingBox + font_name: str | None = None + font_size: float | None = None + confidence: float | None = None + -class Paper(TreeKnowledge): - """A research paper as a tree of sections. +class PaperParsedPage(BaseModel): + """One physical page and its content-addressed visual references.""" - The tree's nodes carry per-section ``summary`` (for navigation) and, on - leaves, the section ``content`` (Markdown). Top-level metadata - (``arxiv_id``, ``authors``) lives here on the tree itself. + model_config = ConfigDict(extra="forbid", frozen=True) + + page_number: int = Field(ge=1) + width: float = Field(gt=0) + height: float = Field(gt=0) + text: str + blocks: tuple[PaperParsedBlock, ...] = () + screenshot_asset_id: UUID | None = None + image_asset_ids: tuple[UUID, ...] = () + + +class PaperParsedManifest(BaseModel): + """Complete page-aware parser output for one exact source revision.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + source_hash: str = Field(pattern=r"^[0-9a-f]{64}$") + parser_name: str + parser_version: str + cleanup_version: str + pages: tuple[PaperParsedPage, ...] = Field(min_length=1) + + @model_validator(mode="after") + def _pages_are_contiguous(self) -> "PaperParsedManifest": + expected = tuple(range(1, len(self.pages) + 1)) + actual = tuple(page.page_number for page in self.pages) + if actual != expected: + raise ValueError("parsed pages must be contiguous and 1-based") + return self + + +class PaperSourceRevision(BaseModel): + """Immutable source revision anchored to the exact fetched bytes. + + ``blobs`` carries exact bytes across the flow-to-library boundary. It is + excluded from canonical JSON; the local library persists every referenced + blob in a transactionally linked content-addressed table. """ - item_type: Literal["paper"] = "paper" + model_config = ConfigDict(extra="forbid", frozen=True) + + id: UUID + schema_version: Literal["1.0"] = "1.0" + source: SourceRef + as_of: datetime + available_at: datetime + published_at: datetime | None = None arxiv_id: str | None = None - authors: list[str] = Field(default_factory=list) - asset_classes: list[str] = Field(default_factory=list) + title: str | None = None + authors: tuple[str, ...] = () + parsed: PaperParsedManifest + raw_asset_id: UUID + assets: tuple[PaperAssetRef, ...] = Field(min_length=1) + blobs: dict[str, bytes] = Field( + default_factory=dict, + exclude=True, + repr=False, + ) + @field_validator("as_of", "available_at", "published_at") + @classmethod + def _timestamps_are_aware(cls, value: datetime | None) -> datetime | None: + if value is not None and ( + value.tzinfo is None or value.utcoffset() is None + ): + raise ValueError("paper source timestamps must be timezone-aware") + return value -class PaperKnowledgeCard(FlattenKnowledge): - """Distilled summary card of a `Paper`. + @model_validator(mode="after") + def _validate_revision(self) -> "PaperSourceRevision": + content_hash = self.source.content_hash + if content_hash is None: + raise ValueError("paper source content_hash is required") + if self.id != _paper_source_id(content_hash): + raise ValueError("paper source revision ID does not match content") + if self.parsed.source_hash != content_hash: + raise ValueError("parsed manifest does not match source content") + if self.source.kind == "arxiv" and ( + self.arxiv_id is None or re.search(r"v\d+$", self.arxiv_id) is None + ): + raise ValueError("arXiv paper sources require an exact revision") - The store layer keys this off ``paper_id`` so the card and its tree can - be retrieved together. The card is the right shape for tagging, - filtering, and dashboard surfaces; deep questions go to the tree. - """ + assets = {asset.asset_id: asset for asset in self.assets} + if len(assets) != len(self.assets): + raise ValueError("paper source asset IDs must be unique") + raw = assets.get(self.raw_asset_id) + if raw is None or raw.kind != "raw" or raw.content_hash != content_hash: + raise ValueError("raw paper asset does not match source content") + for asset in self.assets: + expected_id = _paper_asset_id( + self.id, + kind=asset.kind, + page_number=asset.page_number, + content_hash=asset.content_hash, + ) + if asset.asset_id != expected_id: + raise ValueError("paper asset ID does not match its content") + + for page in self.parsed.pages: + referenced = ( + (page.screenshot_asset_id,) + if page.screenshot_asset_id is not None + else () + ) + page.image_asset_ids + for asset_id in referenced: + asset = assets.get(asset_id) + if asset is None or asset.page_number != page.page_number: + raise ValueError("parsed page references an unknown asset") + + if self.blobs: + for content_hash_value, blob in self.blobs.items(): + if hashlib.sha256(blob).hexdigest() != content_hash_value: + raise ValueError("paper source blob hash mismatch") + for asset in self.assets: + blob = self.blobs.get(asset.content_hash) + if blob is None or len(blob) != asset.size_bytes: + raise ValueError("paper source is missing an asset blob") + return self + + def blob_for(self, asset_id: UUID) -> bytes: + """Return exact bytes for one referenced source asset.""" + asset = next( + (item for item in self.assets if item.asset_id == asset_id), + None, + ) + if asset is None: + raise KeyError(f"Paper asset '{asset_id}' not found") + try: + return self.blobs[asset.content_hash] + except KeyError as exc: + raise RuntimeError( + f"Paper asset '{asset_id}' bytes are not loaded" + ) from exc + + +class PaperSourceSpan(BaseModel): + """Exact character span and page evidence retained by one chunk.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + page_number: int = Field(ge=1) + start_char: int = Field(ge=0) + end_char: int = Field(gt=0) + block_boxes: tuple[PaperBoundingBox, ...] = () + asset_ids: tuple[UUID, ...] = () + + @model_validator(mode="after") + def _span_is_ordered(self) -> "PaperSourceSpan": + if self.end_char <= self.start_char: + raise ValueError("paper source span must be non-empty") + return self + + +def _paper_chunk_id( + chunk_set_id: UUID, + *, + position: int, + content_hash: str, + spans: tuple[PaperSourceSpan, ...], +) -> UUID: + span_hash = _stable_hash([span.model_dump(mode="json") for span in spans]) + return uuid5( + chunk_set_id, + f"quantmind:paper-chunk:{position}:{content_hash}:{span_hash}", + ) + + +class PaperChunk(BaseModel): + """Directly addressable page-aware member of a chunk-set artifact.""" - item_type: Literal["paper_card"] = "paper_card" + model_config = ConfigDict(extra="forbid", frozen=True) - paper_id: UUID - summary: str - methodology: str | None = None - key_findings: list[str] = Field(default_factory=list) - limitations: list[str] = Field(default_factory=list) + chunk_id: UUID + chunk_set_id: UUID + source_revision_id: UUID + position: int = Field(ge=0) + text: str = Field(min_length=1) + content_hash: str = Field(pattern=r"^[0-9a-f]{64}$") + source_spans: tuple[PaperSourceSpan, ...] = Field(min_length=1) + + @field_validator("text") + @classmethod + def _text_is_not_blank(cls, value: str) -> str: + if not value.strip(): + raise ValueError("paper chunk text must not be blank") + return value + + @model_validator(mode="after") + def _identity_matches_content(self) -> "PaperChunk": + if self.content_hash != _text_hash(self.text): + raise ValueError("paper chunk content hash mismatch") + expected = _paper_chunk_id( + self.chunk_set_id, + position=self.position, + content_hash=self.content_hash, + spans=self.source_spans, + ) + if self.chunk_id != expected: + raise ValueError("paper chunk ID does not match its content") + return self + + +class PaperChunkingConfig(BaseModel): + """Exact LlamaIndex splitter identity and configuration.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + splitter: Literal["llama-index-sentence-splitter"] = ( + "llama-index-sentence-splitter" + ) + splitter_version: str + chunk_size: int = Field(gt=0) + chunk_overlap: int = Field(ge=0) + + @model_validator(mode="after") + def _overlap_is_smaller_than_chunk(self) -> "PaperChunkingConfig": + if self.chunk_overlap >= self.chunk_size: + raise ValueError("chunk_overlap must be smaller than chunk_size") + return self + + +def _paper_chunk_set_content_hash(chunks: tuple[PaperChunk, ...]) -> str: + return _stable_hash( + [ + { + "position": chunk.position, + "content_hash": chunk.content_hash, + "source_spans": [ + span.model_dump(mode="json") for span in chunk.source_spans + ], + } + for chunk in chunks + ] + ) + + +class PaperChunkSet(BaseModel): + """Versioned ordered chunks produced from one source revision.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + id: UUID + artifact_kind: Literal["paper_chunk_set"] = "paper_chunk_set" + schema_version: Literal["1.0"] = "1.0" + source_revision_id: UUID + producer: PaperChunkingConfig + producer_config_hash: str = Field(pattern=r"^[0-9a-f]{64}$") + content_hash: str = Field(pattern=r"^[0-9a-f]{64}$") + chunks: tuple[PaperChunk, ...] = Field(min_length=1) + + @model_validator(mode="after") + def _validate_artifact(self) -> "PaperChunkSet": + expected_config_hash = _stable_hash( + self.producer.model_dump(mode="json") + ) + if self.producer_config_hash != expected_config_hash: + raise ValueError("paper chunk-set producer hash mismatch") + expected_id = _paper_artifact_id( + self.source_revision_id, + self.artifact_kind, + self.producer_config_hash, + ) + if self.id != expected_id: + raise ValueError("paper chunk-set ID does not match its producer") + if tuple(chunk.position for chunk in self.chunks) != tuple( + range(len(self.chunks)) + ): + raise ValueError("paper chunks must have contiguous positions") + if len({chunk.chunk_id for chunk in self.chunks}) != len(self.chunks): + raise ValueError("paper chunk IDs must be unique") + if any( + chunk.chunk_set_id != self.id + or chunk.source_revision_id != self.source_revision_id + for chunk in self.chunks + ): + raise ValueError("paper chunk membership does not match chunk set") + if self.content_hash != _paper_chunk_set_content_hash(self.chunks): + raise ValueError("paper chunk-set content hash mismatch") + return self + + +def _validate_chunk_set_source( + source: PaperSourceRevision, + chunk_set: PaperChunkSet, +) -> None: + """Validate every chunk span against its exact source manifest.""" + if chunk_set.source_revision_id != source.id: + raise ValueError("paper chunk set belongs to another source") + pages = {page.page_number: page for page in source.parsed.pages} + assets = {asset.asset_id: asset for asset in source.assets} + for chunk in chunk_set.chunks: + for span in chunk.source_spans: + page = pages.get(span.page_number) + if page is None: + raise ValueError("paper chunk span references an unknown page") + if span.end_char > len(page.text): + raise ValueError("paper chunk span exceeds its source page") + for asset_id in span.asset_ids: + asset = assets.get(asset_id) + if ( + asset is None + or asset.page_number != span.page_number + or asset.kind == "raw" + ): + raise ValueError( + "paper chunk span references an unknown page asset" + ) + + +class ArtifactLocator(BaseModel): + """Stable address for a paper artifact or one of its members.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + source_revision_id: UUID | None + artifact_id: UUID + artifact_kind: str = Field(min_length=1) + member_id: UUID | None = None + + +class PaperCitation(BaseModel): + """Code-resolved citation from a summary to one chunk and source page.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + chunk_set_id: UUID + chunk_id: UUID + page_number: int = Field(ge=1) + quote: str | None = Field(default=None, max_length=500) + + +class PaperSummaryProducer(BaseModel): + """Exact model, prompt, and output-bound identity for a summary.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + model: str + prompt_version: str + input_chunk_set_id: UUID + instructions_hash: str = Field(pattern=r"^[0-9a-f]{64}$") + max_output_tokens: int = Field(gt=0) + + +def _paper_summary_content_hash( + summary: str, + citations: tuple[PaperCitation, ...], +) -> str: + return _stable_hash( + { + "summary": summary, + "citations": [ + citation.model_dump(mode="json") for citation in citations + ], + } + ) + + +class PaperGlobalSummary(BaseModel): + """One independently versioned cited global-summary artifact.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + id: UUID + artifact_kind: Literal["paper_summary"] = "paper_summary" + schema_version: Literal["1.0"] = "1.0" + source_revision_id: UUID + producer: PaperSummaryProducer + producer_config_hash: str = Field(pattern=r"^[0-9a-f]{64}$") + content_hash: str = Field(pattern=r"^[0-9a-f]{64}$") + summary: str = Field(min_length=1) + citations: tuple[PaperCitation, ...] = Field(min_length=1) + derived_from: tuple[ArtifactLocator, ...] = Field(min_length=1) + + @field_validator("summary") + @classmethod + def _summary_is_not_blank(cls, value: str) -> str: + if not value.strip(): + raise ValueError("paper global summary must not be blank") + return value + + @model_validator(mode="after") + def _validate_artifact(self) -> "PaperGlobalSummary": + expected_config_hash = _stable_hash( + self.producer.model_dump(mode="json") + ) + if self.producer_config_hash != expected_config_hash: + raise ValueError("paper summary producer hash mismatch") + expected_id = _paper_artifact_id( + self.source_revision_id, + self.artifact_kind, + self.producer_config_hash, + ) + if self.id != expected_id: + raise ValueError("paper summary ID does not match its producer") + if self.content_hash != _paper_summary_content_hash( + self.summary, + self.citations, + ): + raise ValueError("paper summary content hash mismatch") + if any( + locator.source_revision_id != self.source_revision_id + or locator.member_id is not None + for locator in self.derived_from + ): + raise ValueError("paper summary lineage has an invalid locator") + if not any( + locator.artifact_kind == "paper_chunk_set" + and locator.artifact_id == self.producer.input_chunk_set_id + for locator in self.derived_from + ): + raise ValueError( + "paper summary producer input is missing from lineage" + ) + return self + + +class PaperFlowResult(BaseModel): + """Validated V1 result containing source, chunks, and cited summary.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + source_revision: PaperSourceRevision + chunk_set: PaperChunkSet + global_summary: PaperGlobalSummary + + @model_validator(mode="after") + def _validate_cross_artifact_links(self) -> "PaperFlowResult": + source_id = self.source_revision.id + if ( + self.chunk_set.source_revision_id != source_id + or self.global_summary.source_revision_id != source_id + ): + raise ValueError("paper artifacts do not share their source") + _validate_chunk_set_source(self.source_revision, self.chunk_set) + chunk_set_locator = ArtifactLocator( + source_revision_id=source_id, + artifact_id=self.chunk_set.id, + artifact_kind="paper_chunk_set", + ) + if chunk_set_locator not in self.global_summary.derived_from: + raise ValueError("paper summary is missing chunk-set lineage") + + chunks = {chunk.chunk_id: chunk for chunk in self.chunk_set.chunks} + for citation in self.global_summary.citations: + chunk = chunks.get(citation.chunk_id) + if chunk is None or citation.chunk_set_id != self.chunk_set.id: + raise ValueError("paper summary cites an unknown chunk") + pages = {span.page_number for span in chunk.source_spans} + if citation.page_number not in pages: + raise ValueError("paper summary citation page is not in chunk") + if citation.quote and citation.quote not in chunk.text: + raise ValueError("paper summary citation quote is not in chunk") + return self + + @property + def source(self) -> PaperSourceRevision: + """Return the exact source revision using a concise compatibility name.""" + return self.source_revision + + @property + def summary(self) -> PaperGlobalSummary: + """Return the global-summary artifact using a concise name.""" + return self.global_summary + + +class LegacyPaper(TreeKnowledge): + """Pre-V1 paper tree retained only for explicit database compatibility.""" + + item_type: Literal["paper"] = "paper" + arxiv_id: str | None = None + authors: list[str] = Field(default_factory=list) asset_classes: list[str] = Field(default_factory=list) - def embedding_text(self) -> str: - return f"{self.summary}\n{' '.join(self.key_findings)}" + +PaperArtifact = PaperChunkSet | PaperGlobalSummary +ResolvedPaperArtifact = PaperArtifact | PaperChunk diff --git a/quantmind/knowledge/thesis.py b/quantmind/knowledge/thesis.py index 940e5b3..5862ef1 100644 --- a/quantmind/knowledge/thesis.py +++ b/quantmind/knowledge/thesis.py @@ -16,6 +16,3 @@ class Thesis(FlattenKnowledge): item_type: Literal["thesis"] = "thesis" claim: str - - def embedding_text(self) -> str: - return self.claim diff --git a/quantmind/library/__init__.py b/quantmind/library/__init__.py index edfde4b..b5d6453 100644 --- a/quantmind/library/__init__.py +++ b/quantmind/library/__init__.py @@ -1,6 +1,15 @@ """Local semantic retrieval for canonical QuantMind knowledge.""" -from quantmind.library._types import SemanticHit, SemanticQuery +from quantmind.library._types import ( + SearchProjection, + SemanticHit, + SemanticQuery, +) from quantmind.library.local import LocalKnowledgeLibrary -__all__ = ["LocalKnowledgeLibrary", "SemanticHit", "SemanticQuery"] +__all__ = [ + "LocalKnowledgeLibrary", + "SearchProjection", + "SemanticHit", + "SemanticQuery", +] diff --git a/quantmind/library/_internal/llamaindex_retriever.py b/quantmind/library/_internal/llamaindex_retriever.py index a6e29f2..b79ca7d 100644 --- a/quantmind/library/_internal/llamaindex_retriever.py +++ b/quantmind/library/_internal/llamaindex_retriever.py @@ -20,9 +20,16 @@ class _IndexRecord: """Durable metadata and vector bytes for one searchable target.""" target_id: str + owner_kind: str item_id: UUID node_id: UUID | None item_type: str + source_revision_id: UUID | None + artifact_kind: str + projection_kind: str + projection_version: str + embedding_model: str + projection_hash: str matched_text: str as_of: float available_at: float | None @@ -154,10 +161,18 @@ def filter(self, query: SemanticQuery) -> tuple[_IndexRecord, ...]: else None ) item_types = set(query.item_types) if query.item_types else None + artifact_kinds = ( + set(query.artifact_kinds) if query.artifact_kinds else None + ) source_kinds = set(query.source_kinds) if query.source_kinds else None required_tags = set(query.tags) if query.tags else None selected: list[_IndexRecord] = [] for record in self._records: + if ( + artifact_kinds is not None + and record.artifact_kind not in artifact_kinds + ): + continue if item_types is not None and record.item_type not in item_types: continue if ( diff --git a/quantmind/library/_internal/retrieval_targets.py b/quantmind/library/_internal/retrieval_targets.py index 5cc1726..d36d753 100644 --- a/quantmind/library/_internal/retrieval_targets.py +++ b/quantmind/library/_internal/retrieval_targets.py @@ -4,7 +4,17 @@ from dataclasses import dataclass from uuid import UUID -from quantmind.knowledge import BaseKnowledge, FlattenKnowledge, TreeKnowledge +from quantmind.knowledge import ( + BaseKnowledge, + Citation, + Earnings, + Factor, + FlattenKnowledge, + News, + PaperFlowResult, + Thesis, + TreeKnowledge, +) _PROJECTION_SCHEMA_VERSION = "1" @@ -14,10 +24,14 @@ class _RetrievalTarget: """One stable item or non-root node target to embed.""" target_id: str + artifact_id: UUID + artifact_kind: str + source_revision_id: UUID | None node_id: UUID | None text: str projection_hash: str tree_id: UUID | None + citations: tuple[Citation, ...] = () def _project_knowledge(item: BaseKnowledge) -> list[_RetrievalTarget]: @@ -27,13 +41,16 @@ def _project_knowledge(item: BaseKnowledge) -> list[_RetrievalTarget]: "LocalKnowledgeLibrary supports FlattenKnowledge and TreeKnowledge" ) - root_text = item.embedding_text() + root_text = _knowledge_text(item) if not root_text: - raise ValueError("knowledge embedding_text() must not be empty") + raise ValueError("knowledge text projection must not be empty") tree_id = item.id if isinstance(item, TreeKnowledge) else None targets = [ _RetrievalTarget( target_id=f"item:{item.id}", + artifact_id=item.id, + artifact_kind=item.item_type, + source_revision_id=None, node_id=None, text=root_text, projection_hash=hashlib.sha256( @@ -53,12 +70,15 @@ def _project_knowledge(item: BaseKnowledge) -> list[_RetrievalTarget]: raise ValueError("tree node map key does not match node_id") if node_id == item.root_node_id: continue - text = node.embedding_text() + text = f"{node.title}\n{node.summary}" if not text: - raise ValueError("tree node embedding_text() must not be empty") + raise ValueError("tree node text projection must not be empty") targets.append( _RetrievalTarget( target_id=f"node:{item.id}:{node_id}", + artifact_id=item.id, + artifact_kind=item.item_type, + source_revision_id=None, node_id=node_id, text=text, projection_hash=hashlib.sha256( @@ -68,3 +88,83 @@ def _project_knowledge(item: BaseKnowledge) -> list[_RetrievalTarget]: ) ) return targets + + +def _project_paper(result: PaperFlowResult) -> list[_RetrievalTarget]: + """Project a cited summary and every non-empty chunk for text search.""" + source_id = result.source_revision.id + summary = result.global_summary + targets = [ + _RetrievalTarget( + target_id=f"artifact:{summary.id}", + artifact_id=summary.id, + artifact_kind=summary.artifact_kind, + source_revision_id=source_id, + node_id=None, + text=summary.summary, + projection_hash=hashlib.sha256( + summary.summary.encode("utf-8") + ).hexdigest(), + tree_id=None, + citations=tuple( + Citation( + source_id=str(source_id), + page=citation.page_number, + quote=citation.quote, + ) + for citation in summary.citations + ), + ) + ] + for chunk in result.chunk_set.chunks: + if not chunk.text.strip(): + continue + targets.append( + _RetrievalTarget( + target_id=f"artifact-member:{result.chunk_set.id}:{chunk.chunk_id}", + artifact_id=result.chunk_set.id, + artifact_kind=result.chunk_set.artifact_kind, + source_revision_id=source_id, + node_id=chunk.chunk_id, + text=chunk.text, + projection_hash=hashlib.sha256( + chunk.text.encode("utf-8") + ).hexdigest(), + tree_id=None, + citations=tuple( + Citation( + source_id=str(source_id), + page=span.page_number, + char_offset=span.start_char, + quote=chunk.text[:500], + ) + for span in chunk.source_spans + ), + ) + ) + if len(targets) != len(result.chunk_set.chunks) + 1: + raise ValueError( + "every paper chunk must have a required text projection" + ) + return targets + + +def _knowledge_text(item: BaseKnowledge) -> str: + """Select searchable text at the indexing boundary, not in the model.""" + if isinstance(item, News): + entities = ", ".join(item.entities) + return f"{item.headline}\n{item.event_type}\n{entities}".strip() + if isinstance(item, Earnings): + guidance = item.guidance or "" + return f"{item.ticker} {item.period} earnings\n{guidance}".strip() + if isinstance(item, Factor): + scope = item.universe or "unspecified" + return f"factor {item.factor_name} on {scope}" + if isinstance(item, Thesis): + return item.claim + if isinstance(item, TreeKnowledge): + root = item.root() + return f"{root.title}\n{root.summary}" + raise TypeError( + f"Unsupported knowledge projection type '{type(item).__name__}'" + ) diff --git a/quantmind/library/_internal/sqlite_store.py b/quantmind/library/_internal/sqlite_store.py index cbe4ebd..dc0512d 100644 --- a/quantmind/library/_internal/sqlite_store.py +++ b/quantmind/library/_internal/sqlite_store.py @@ -12,22 +12,30 @@ from pydantic import ValidationError from quantmind.knowledge import ( + ArtifactLocator, BaseKnowledge, + Citation, Earnings, Factor, + LegacyPaper, News, - Paper, - PaperKnowledgeCard, + PaperArtifact, + PaperChunkSet, + PaperFlowResult, + PaperGlobalSummary, + PaperSourceRevision, + ResolvedPaperArtifact, Thesis, TreeKnowledge, ) +from quantmind.knowledge.paper import _validate_chunk_set_source from quantmind.library._internal.llamaindex_retriever import _IndexRecord from quantmind.library._internal.retrieval_targets import ( _PROJECTION_SCHEMA_VERSION, _RetrievalTarget, ) -_DATABASE_SCHEMA_VERSION = 2 +_DATABASE_SCHEMA_VERSION = 3 _KNOWLEDGE_CLASSES: dict[str, type[BaseKnowledge]] = { f"{knowledge_type.__module__}:{knowledge_type.__qualname__}": knowledge_type @@ -35,11 +43,11 @@ Earnings, Factor, News, - Paper, - PaperKnowledgeCard, + LegacyPaper, Thesis, ) } +_KNOWLEDGE_CLASSES["quantmind.knowledge.paper:Paper"] = LegacyPaper @dataclass(frozen=True) @@ -90,6 +98,37 @@ class _PreparedPut: existing_embeddings: dict[str, _StoredEmbedding] +@dataclass(frozen=True) +class _CanonicalPaperMember: + """One separately addressable paper-artifact member.""" + + member_id: UUID + position: int + payload: str + content_hash: str + + +@dataclass(frozen=True) +class _CanonicalPaperArtifact: + """One aggregate artifact plus normalized member rows.""" + + artifact: PaperArtifact + payload: str + canonical_hash: str + members: tuple[_CanonicalPaperMember, ...] + + +@dataclass(frozen=True) +class _PreparedPaperPut: + """Validated source/artifact write plus reusable search projections.""" + + result: PaperFlowResult + source_payload: str + source_canonical_hash: str + artifacts: tuple[_CanonicalPaperArtifact, ...] + existing_embeddings: dict[str, _StoredEmbedding] + + def _json_payload(value: object) -> str: """Encode canonical JSON with stable ordering and separators.""" return json.dumps( @@ -144,6 +183,38 @@ def _canonical_payload(item: BaseKnowledge) -> _CanonicalDocument: ) +def _canonical_paper_artifact( + artifact: PaperArtifact, +) -> _CanonicalPaperArtifact: + """Normalize a paper artifact without losing its aggregate hash.""" + full_payload = _json_payload(artifact.model_dump(mode="json")) + canonical_hash = hashlib.sha256(full_payload.encode("utf-8")).hexdigest() + if isinstance(artifact, PaperGlobalSummary): + return _CanonicalPaperArtifact( + artifact=artifact, + payload=full_payload, + canonical_hash=canonical_hash, + members=(), + ) + members = tuple( + _CanonicalPaperMember( + member_id=chunk.chunk_id, + position=chunk.position, + payload=(payload := _json_payload(chunk.model_dump(mode="json"))), + content_hash=hashlib.sha256(payload.encode("utf-8")).hexdigest(), + ) + for chunk in artifact.chunks + ) + return _CanonicalPaperArtifact( + artifact=artifact, + payload=_json_payload( + artifact.model_dump(mode="json", exclude={"chunks"}) + ), + canonical_hash=canonical_hash, + members=members, + ) + + def _assemble_canonical_payload( *, item_id: str, @@ -257,17 +328,123 @@ def _timestamp(value: datetime, field_name: str) -> float: return value.astimezone(timezone.utc).timestamp() +_PAPER_TABLES_SQL = """ +CREATE TABLE paper_sources ( + source_revision_id TEXT PRIMARY KEY, + schema_version TEXT NOT NULL, + source_content_hash TEXT NOT NULL UNIQUE, + payload_json TEXT NOT NULL, + canonical_hash TEXT NOT NULL, + asset_count INTEGER NOT NULL CHECK (asset_count > 0) +); + +CREATE TABLE paper_source_assets ( + asset_id TEXT PRIMARY KEY, + source_revision_id TEXT NOT NULL, + kind TEXT NOT NULL, + page_number INTEGER, + media_type TEXT NOT NULL, + content_hash TEXT NOT NULL, + size_bytes INTEGER NOT NULL CHECK (size_bytes >= 0), + blob BLOB NOT NULL, + FOREIGN KEY (source_revision_id) REFERENCES paper_sources(source_revision_id) + ON DELETE CASCADE +); + +CREATE TABLE paper_artifacts ( + artifact_id TEXT PRIMARY KEY, + source_revision_id TEXT NOT NULL, + artifact_kind TEXT NOT NULL, + schema_version TEXT NOT NULL, + producer_config_hash TEXT NOT NULL, + payload_json TEXT NOT NULL, + canonical_hash TEXT NOT NULL, + member_count INTEGER NOT NULL CHECK (member_count >= 0), + target_count INTEGER NOT NULL CHECK (target_count > 0), + UNIQUE (source_revision_id, artifact_kind, producer_config_hash), + FOREIGN KEY (source_revision_id) REFERENCES paper_sources(source_revision_id) + ON DELETE CASCADE +); + +CREATE TABLE paper_artifact_members ( + artifact_id TEXT NOT NULL, + member_id TEXT NOT NULL, + position INTEGER NOT NULL CHECK (position >= 0), + payload_json TEXT NOT NULL, + content_hash TEXT NOT NULL, + PRIMARY KEY (artifact_id, member_id), + UNIQUE (artifact_id, position), + FOREIGN KEY (artifact_id) REFERENCES paper_artifacts(artifact_id) + ON DELETE CASCADE +); + +CREATE TABLE paper_artifact_lineage ( + artifact_id TEXT NOT NULL, + input_artifact_id TEXT NOT NULL, + relation TEXT NOT NULL, + PRIMARY KEY (artifact_id, input_artifact_id, relation), + FOREIGN KEY (artifact_id) REFERENCES paper_artifacts(artifact_id) + ON DELETE CASCADE, + FOREIGN KEY (input_artifact_id) REFERENCES paper_artifacts(artifact_id) + ON DELETE RESTRICT +); + +CREATE TABLE paper_projections ( + target_id TEXT PRIMARY KEY, + source_revision_id TEXT NOT NULL, + artifact_id TEXT NOT NULL, + member_id TEXT, + artifact_kind TEXT NOT NULL, + matched_text TEXT NOT NULL, + as_of REAL NOT NULL, + available_at REAL NOT NULL, + source_kind TEXT NOT NULL, + citations_json TEXT NOT NULL, + projection_kind TEXT NOT NULL, + modality TEXT NOT NULL, + embedding_model TEXT NOT NULL, + dimension INTEGER NOT NULL CHECK (dimension > 0), + projection_hash TEXT NOT NULL, + source_content_hash TEXT NOT NULL, + artifact_schema_version TEXT NOT NULL, + projection_schema_version TEXT NOT NULL, + artifact_canonical_hash TEXT NOT NULL, + embedding BLOB NOT NULL, + FOREIGN KEY (source_revision_id) REFERENCES paper_sources(source_revision_id) + ON DELETE CASCADE, + FOREIGN KEY (artifact_id) REFERENCES paper_artifacts(artifact_id) + ON DELETE CASCADE, + FOREIGN KEY (artifact_id, member_id) + REFERENCES paper_artifact_members(artifact_id, member_id) + ON DELETE CASCADE +); + +CREATE INDEX paper_artifacts_source_kind + ON paper_artifacts(source_revision_id, artifact_kind); +CREATE INDEX paper_projections_filters + ON paper_projections(artifact_kind, source_kind, source_revision_id); +""" + + def _initialize_schema(db: sqlite3.Connection) -> None: """Create the current schema or reject an incompatible local database.""" version_row = db.execute("PRAGMA user_version").fetchone() version = int(version_row[0]) - if version not in (0, _DATABASE_SCHEMA_VERSION): + if version not in (0, 2, _DATABASE_SCHEMA_VERSION): raise RuntimeError( "Stale knowledge library schema: database version " f"{version}, expected {_DATABASE_SCHEMA_VERSION}" ) if version == _DATABASE_SCHEMA_VERSION: return + if version == 2: + migration_sql = _PAPER_TABLES_SQL.replace( + "CREATE TABLE ", "CREATE TABLE IF NOT EXISTS " + ).replace("CREATE INDEX ", "CREATE INDEX IF NOT EXISTS ") + db.executescript( + f"{migration_sql}\nPRAGMA user_version = {_DATABASE_SCHEMA_VERSION};" + ) + return db.executescript( f""" CREATE TABLE knowledge_items ( @@ -325,6 +502,8 @@ def _initialize_schema(db: sqlite3.Connection) -> None: CREATE INDEX knowledge_nodes_parent ON knowledge_nodes(item_id, parent_id, position); + {_PAPER_TABLES_SQL} + PRAGMA user_version = {_DATABASE_SCHEMA_VERSION}; """ ) @@ -406,6 +585,241 @@ def prepare_put(self, item: BaseKnowledge) -> _PreparedPut: existing_embeddings=existing, ) + def prepare_put_paper(self, result: PaperFlowResult) -> _PreparedPaperPut: + """Validate paper blobs and load projections eligible for reuse.""" + source = result.source_revision + for asset in source.assets: + blob = source.blobs.get(asset.content_hash) + if blob is None: + raise ValueError( + f"Paper source is missing blob for asset '{asset.asset_id}'" + ) + if ( + len(blob) != asset.size_bytes + or hashlib.sha256(blob).hexdigest() != asset.content_hash + ): + raise ValueError( + f"Paper source blob for asset '{asset.asset_id}' is invalid" + ) + artifact_ids = (str(result.chunk_set.id), str(result.global_summary.id)) + rows = self._db.execute( + """ + SELECT * FROM paper_projections + WHERE artifact_id IN (?, ?) + """, + artifact_ids, + ).fetchall() + existing = { + str(row["target_id"]): _StoredEmbedding( + target_id=str(row["target_id"]), + embedding_model=str(row["embedding_model"]), + dimension=int(row["dimension"]), + projection_hash=str(row["projection_hash"]), + source_content_hash=str(row["source_content_hash"]), + knowledge_schema_version=str(row["artifact_schema_version"]), + projection_schema_version=str(row["projection_schema_version"]), + embedding=bytes(row["embedding"]), + ) + for row in rows + } + source_payload = _json_payload(source.model_dump(mode="json")) + return _PreparedPaperPut( + result=result, + source_payload=source_payload, + source_canonical_hash=hashlib.sha256( + source_payload.encode("utf-8") + ).hexdigest(), + artifacts=( + _canonical_paper_artifact(result.chunk_set), + _canonical_paper_artifact(result.global_summary), + ), + existing_embeddings=existing, + ) + + def put_paper( + self, + prepared: _PreparedPaperPut, + targets: Sequence[_RetrievalTarget], + vectors: dict[str, tuple[bytes, int]], + *, + embedding_model: str, + ) -> None: + """Atomically persist one source, two artifacts, lineage, and vectors.""" + result = prepared.result + source = result.source_revision + canonical_by_id = { + artifact.artifact.id: artifact for artifact in prepared.artifacts + } + targets_by_artifact: dict[UUID, list[_RetrievalTarget]] = {} + for target in targets: + targets_by_artifact.setdefault(target.artifact_id, []).append( + target + ) + if set(targets_by_artifact) != set(canonical_by_id): + raise ValueError("paper artifacts do not have complete projections") + try: + self._db.execute("BEGIN IMMEDIATE") + self._db.execute( + """ + INSERT INTO paper_sources ( + source_revision_id, schema_version, source_content_hash, + payload_json, canonical_hash, asset_count + ) VALUES (?, ?, ?, ?, ?, ?) + ON CONFLICT(source_revision_id) DO NOTHING + """, + ( + str(source.id), + source.schema_version, + source.source.content_hash, + prepared.source_payload, + prepared.source_canonical_hash, + len(source.assets), + ), + ) + for asset in source.assets: + self._db.execute( + """ + INSERT INTO paper_source_assets ( + asset_id, source_revision_id, kind, page_number, + media_type, content_hash, size_bytes, blob + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(asset_id) DO NOTHING + """, + ( + str(asset.asset_id), + str(source.id), + asset.kind, + asset.page_number, + asset.media_type, + asset.content_hash, + asset.size_bytes, + source.blobs[asset.content_hash], + ), + ) + for canonical in prepared.artifacts: + artifact = canonical.artifact + artifact_targets = targets_by_artifact[artifact.id] + self._db.execute( + """ + INSERT INTO paper_artifacts ( + artifact_id, source_revision_id, artifact_kind, + schema_version, producer_config_hash, payload_json, + canonical_hash, member_count, target_count + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(artifact_id) DO UPDATE SET + payload_json = excluded.payload_json, + canonical_hash = excluded.canonical_hash, + member_count = excluded.member_count, + target_count = excluded.target_count + """, + ( + str(artifact.id), + str(artifact.source_revision_id), + artifact.artifact_kind, + artifact.schema_version, + artifact.producer_config_hash, + canonical.payload, + canonical.canonical_hash, + len(canonical.members), + len(artifact_targets), + ), + ) + self._db.execute( + "DELETE FROM paper_projections WHERE artifact_id = ?", + (str(artifact.id),), + ) + self._db.execute( + "DELETE FROM paper_artifact_members WHERE artifact_id = ?", + (str(artifact.id),), + ) + self._db.execute( + "DELETE FROM paper_artifact_lineage WHERE artifact_id = ?", + (str(artifact.id),), + ) + for member in canonical.members: + self._db.execute( + """ + INSERT INTO paper_artifact_members ( + artifact_id, member_id, position, + payload_json, content_hash + ) VALUES (?, ?, ?, ?, ?) + """, + ( + str(artifact.id), + str(member.member_id), + member.position, + member.payload, + member.content_hash, + ), + ) + for locator in result.global_summary.derived_from: + self._db.execute( + """ + INSERT INTO paper_artifact_lineage ( + artifact_id, input_artifact_id, relation + ) VALUES (?, ?, ?) + """, + ( + str(result.global_summary.id), + str(locator.artifact_id), + "generated_from", + ), + ) + for target in targets: + blob, dimension = vectors[target.target_id] + canonical = canonical_by_id[target.artifact_id] + self._db.execute( + """ + INSERT INTO paper_projections ( + target_id, source_revision_id, artifact_id, member_id, + artifact_kind, matched_text, as_of, available_at, + source_kind, citations_json, projection_kind, modality, + embedding_model, dimension, projection_hash, + source_content_hash, artifact_schema_version, + projection_schema_version, artifact_canonical_hash, + embedding + ) VALUES ( + ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, + ?, ? + ) + """, + ( + target.target_id, + str(source.id), + str(target.artifact_id), + str(target.node_id) if target.node_id else None, + target.artifact_kind, + target.text, + _timestamp(source.as_of, "PaperSourceRevision.as_of"), + _timestamp( + source.available_at, + "PaperSourceRevision.available_at", + ), + source.source.kind, + _json_payload( + [ + citation.model_dump(mode="json") + for citation in target.citations + ] + ), + "text_embedding", + "text", + embedding_model, + dimension, + target.projection_hash, + source.source.content_hash, + canonical.artifact.schema_version, + _PROJECTION_SCHEMA_VERSION, + canonical.canonical_hash, + blob, + ), + ) + self._db.execute("COMMIT") + except Exception: + if self._db.in_transaction: + self._db.execute("ROLLBACK") + raise + def put( self, prepared: _PreparedPut, @@ -587,6 +1001,316 @@ def get(self, item_id: UUID) -> BaseKnowledge: canonical_hash=str(row["canonical_hash"]), ) + def get_paper_source(self, source_revision_id: UUID) -> PaperSourceRevision: + """Rehydrate one exact source revision and all referenced blobs.""" + row = self._db.execute( + "SELECT * FROM paper_sources WHERE source_revision_id = ?", + (str(source_revision_id),), + ).fetchone() + if row is None: + raise KeyError( + f"Paper source revision '{source_revision_id}' not found" + ) + payload = str(row["payload_json"]) + if hashlib.sha256(payload.encode("utf-8")).hexdigest() != str( + row["canonical_hash"] + ): + raise RuntimeError( + f"Stale paper source '{source_revision_id}': content hash mismatch" + ) + asset_rows = self._db.execute( + """ + SELECT * FROM paper_source_assets + WHERE source_revision_id = ? ORDER BY asset_id + """, + (str(source_revision_id),), + ).fetchall() + if len(asset_rows) != int(row["asset_count"]): + raise RuntimeError( + f"Stale paper source '{source_revision_id}': expected " + f"{row['asset_count']} assets, found {len(asset_rows)}" + ) + try: + parsed_payload = json.loads(payload) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"Stale paper source '{source_revision_id}': invalid JSON" + ) from exc + if not isinstance(parsed_payload, dict): + raise RuntimeError( + f"Stale paper source '{source_revision_id}': invalid payload" + ) + blobs: dict[str, bytes] = {} + for asset_row in asset_rows: + blob = bytes(asset_row["blob"]) + content_hash = str(asset_row["content_hash"]) + if ( + len(blob) != int(asset_row["size_bytes"]) + or hashlib.sha256(blob).hexdigest() != content_hash + ): + raise RuntimeError( + f"Corrupt paper asset '{asset_row['asset_id']}'" + ) + blobs[content_hash] = blob + parsed_payload["blobs"] = blobs + try: + source = PaperSourceRevision.model_validate(parsed_payload) + except ValidationError as exc: + raise RuntimeError( + f"Stale paper source '{source_revision_id}': payload no longer " + "validates" + ) from exc + if ( + source.id != source_revision_id + or source.schema_version != str(row["schema_version"]) + or source.source.content_hash != str(row["source_content_hash"]) + ): + raise RuntimeError( + f"Stale paper source '{source_revision_id}': identity mismatch" + ) + try: + stored_assets = { + UUID(str(asset_row["asset_id"])): asset_row + for asset_row in asset_rows + } + except ValueError as exc: + raise RuntimeError( + f"Stale paper source '{source_revision_id}': invalid asset ID" + ) from exc + canonical_assets = {asset.asset_id: asset for asset in source.assets} + if set(stored_assets) != set(canonical_assets): + raise RuntimeError( + f"Stale paper source '{source_revision_id}': asset identity " + "mismatch" + ) + for asset_id, asset in canonical_assets.items(): + stored = stored_assets[asset_id] + if any( + ( + str(stored["source_revision_id"]) != str(source.id), + str(stored["kind"]) != asset.kind, + stored["page_number"] != asset.page_number, + str(stored["media_type"]) != asset.media_type, + str(stored["content_hash"]) != asset.content_hash, + int(stored["size_bytes"]) != asset.size_bytes, + ) + ): + raise RuntimeError( + f"Stale paper asset '{asset_id}': metadata mismatch" + ) + return source + + def get_paper_artifact(self, artifact_id: UUID) -> PaperArtifact: + """Rehydrate one canonical summary or normalized chunk-set artifact.""" + row = self._db.execute( + "SELECT * FROM paper_artifacts WHERE artifact_id = ?", + (str(artifact_id),), + ).fetchone() + if row is None: + raise KeyError(f"Paper artifact '{artifact_id}' not found") + member_rows = self._db.execute( + """ + SELECT * FROM paper_artifact_members + WHERE artifact_id = ? ORDER BY position + """, + (str(artifact_id),), + ).fetchall() + if len(member_rows) != int(row["member_count"]): + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': expected " + f"{row['member_count']} members, found {len(member_rows)}" + ) + try: + payload_value = json.loads(str(row["payload_json"])) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': invalid JSON" + ) from exc + if not isinstance(payload_value, dict): + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': invalid payload" + ) + members: list[object] = [] + for member_row in member_rows: + member_payload = str(member_row["payload_json"]) + if hashlib.sha256( + member_payload.encode("utf-8") + ).hexdigest() != str(member_row["content_hash"]): + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': member content " + "hash mismatch" + ) + try: + parsed_member = json.loads(member_payload) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': invalid member JSON" + ) from exc + if ( + not isinstance(parsed_member, dict) + or parsed_member.get("chunk_id") != member_row["member_id"] + or parsed_member.get("position") != member_row["position"] + ): + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': member metadata " + "mismatch" + ) + members.append(parsed_member) + artifact_kind = str(row["artifact_kind"]) + if artifact_kind == "paper_chunk_set": + payload_value["chunks"] = members + model: type[PaperChunkSet] | type[PaperGlobalSummary] = ( + PaperChunkSet + ) + elif artifact_kind == "paper_summary": + if members: + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': summary has members" + ) + model = PaperGlobalSummary + else: + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': unsupported kind " + f"'{artifact_kind}'" + ) + full_payload = _json_payload(payload_value) + if hashlib.sha256(full_payload.encode("utf-8")).hexdigest() != str( + row["canonical_hash"] + ): + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': canonical hash mismatch" + ) + try: + artifact = model.model_validate(payload_value) + except ValidationError as exc: + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': payload no longer " + "validates" + ) from exc + try: + stored_source_revision_id = UUID(str(row["source_revision_id"])) + except ValueError as exc: + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': invalid source ID" + ) from exc + if ( + artifact.id != artifact_id + or artifact.source_revision_id != stored_source_revision_id + or artifact.artifact_kind != artifact_kind + or artifact.schema_version != str(row["schema_version"]) + or artifact.producer_config_hash != str(row["producer_config_hash"]) + ): + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': identity mismatch" + ) + if isinstance(artifact, PaperChunkSet): + try: + _validate_chunk_set_source( + self.get_paper_source(artifact.source_revision_id), + artifact, + ) + except ValueError as exc: + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': source span " + "mismatch" + ) from exc + lineage_rows = self._db.execute( + """ + SELECT input_artifact_id, relation + FROM paper_artifact_lineage + WHERE artifact_id = ? + ORDER BY input_artifact_id, relation + """, + (str(artifact_id),), + ).fetchall() + expected_lineage = ( + { + (str(locator.artifact_id), "generated_from") + for locator in artifact.derived_from + } + if isinstance(artifact, PaperGlobalSummary) + else set() + ) + stored_lineage = { + (str(lineage["input_artifact_id"]), str(lineage["relation"])) + for lineage in lineage_rows + } + if stored_lineage != expected_lineage: + raise RuntimeError( + f"Stale paper artifact '{artifact_id}': lineage mismatch" + ) + return artifact + + def get_paper_result( + self, + source_revision_id: UUID, + *, + chunk_set_id: UUID | None, + summary_id: UUID | None, + ) -> PaperFlowResult: + """Resolve one unambiguous V1 source/chunk-set/summary combination.""" + source = self.get_paper_source(source_revision_id) + + def select(kind: str, selected_id: UUID | None) -> PaperArtifact: + if selected_id is not None: + artifact = self.get_paper_artifact(selected_id) + if ( + artifact.artifact_kind != kind + or artifact.source_revision_id != source_revision_id + ): + raise KeyError( + f"Paper artifact '{selected_id}' does not belong to " + f"source '{source_revision_id}'" + ) + return artifact + rows = self._db.execute( + """ + SELECT artifact_id FROM paper_artifacts + WHERE source_revision_id = ? AND artifact_kind = ? + ORDER BY artifact_id + """, + (str(source_revision_id), kind), + ).fetchall() + if len(rows) != 1: + raise ValueError( + f"Paper source '{source_revision_id}' has {len(rows)} " + f"'{kind}' artifacts; specify an artifact ID" + ) + return self.get_paper_artifact(UUID(str(rows[0]["artifact_id"]))) + + chunk_set = select("paper_chunk_set", chunk_set_id) + summary = select("paper_summary", summary_id) + if not isinstance(chunk_set, PaperChunkSet) or not isinstance( + summary, PaperGlobalSummary + ): + raise RuntimeError("Stored paper artifact types are inconsistent") + return PaperFlowResult( + source_revision=source, + chunk_set=chunk_set, + global_summary=summary, + ) + + def resolve_paper_locator( + self, locator: ArtifactLocator + ) -> ResolvedPaperArtifact: + """Resolve an artifact locator to its canonical aggregate or member.""" + artifact = self.get_paper_artifact(locator.artifact_id) + if ( + artifact.source_revision_id != locator.source_revision_id + or artifact.artifact_kind != locator.artifact_kind + ): + raise KeyError( + "Paper artifact locator metadata does not match storage" + ) + if locator.member_id is None: + return artifact + if not isinstance(artifact, PaperChunkSet): + raise KeyError("Paper summary artifacts do not have members") + for chunk in artifact.chunks: + if chunk.chunk_id == locator.member_id: + return chunk + raise KeyError(f"Paper chunk '{locator.member_id}' not found") + def delete(self, item_id: UUID) -> None: """Transactionally remove canonical knowledge and every child record.""" exists = self._db.execute( @@ -760,9 +1484,16 @@ def load_index_records( try: record = _IndexRecord( target_id=target_id, + owner_kind="knowledge", item_id=UUID(str(row["item_id"])), node_id=UUID(str(node_value)) if node_value else None, item_type=str(row["item_type"]), + source_revision_id=None, + artifact_kind=str(row["item_type"]), + projection_kind="text_embedding", + projection_version=str(row["projection_schema_version"]), + embedding_model=str(row["embedding_model"]), + projection_hash=str(row["projection_hash"]), matched_text=str(row["matched_text"]), as_of=float(row["as_of"]), available_at=( @@ -782,6 +1513,238 @@ def load_index_records( f"Corrupt index data for target '{target_id}': invalid metadata" ) from exc records.append(record) + + orphan_paper = self._db.execute( + """ + SELECT p.target_id FROM paper_projections AS p + LEFT JOIN paper_sources AS s + ON s.source_revision_id = p.source_revision_id + LEFT JOIN paper_artifacts AS a ON a.artifact_id = p.artifact_id + WHERE s.source_revision_id IS NULL OR a.artifact_id IS NULL + LIMIT 1 + """ + ).fetchone() + if orphan_paper is not None: + raise RuntimeError( + f"Stale paper projection '{orphan_paper['target_id']}': " + "source or artifact is missing" + ) + orphan_member = self._db.execute( + """ + SELECT p.target_id FROM paper_projections AS p + LEFT JOIN paper_artifact_members AS m + ON m.artifact_id = p.artifact_id + AND m.member_id = p.member_id + WHERE p.member_id IS NOT NULL AND m.member_id IS NULL + LIMIT 1 + """ + ).fetchone() + if orphan_member is not None: + raise RuntimeError( + f"Stale paper projection '{orphan_member['target_id']}': " + "artifact member is missing" + ) + incomplete_artifact = self._db.execute( + """ + SELECT a.artifact_id, a.member_count, + COUNT(DISTINCT m.member_id) AS actual_members, + a.target_count, + COUNT(DISTINCT p.target_id) AS actual_targets + FROM paper_artifacts AS a + LEFT JOIN paper_artifact_members AS m + ON m.artifact_id = a.artifact_id + LEFT JOIN paper_projections AS p + ON p.artifact_id = a.artifact_id + GROUP BY a.artifact_id, a.member_count, a.target_count + HAVING actual_members != a.member_count + OR actual_targets != a.target_count + LIMIT 1 + """ + ).fetchone() + if incomplete_artifact is not None: + raise RuntimeError( + f"Stale paper artifact '{incomplete_artifact['artifact_id']}': " + f"expected {incomplete_artifact['member_count']} members and " + f"{incomplete_artifact['target_count']} projections" + ) + paper_rows = self._db.execute( + """ + SELECT p.*, a.canonical_hash AS current_canonical_hash, + a.schema_version AS current_schema_version, + s.source_content_hash AS current_source_content_hash + FROM paper_projections AS p + JOIN paper_artifacts AS a ON a.artifact_id = p.artifact_id + JOIN paper_sources AS s + ON s.source_revision_id = p.source_revision_id + ORDER BY p.target_id + """ + ).fetchall() + paper_artifacts: dict[UUID, PaperArtifact] = {} + paper_sources: dict[UUID, PaperSourceRevision] = {} + for row in paper_rows: + target_id = str(row["target_id"]) + dimension = int(row["dimension"]) + if str(row["embedding_model"]) != embedding_model: + raise RuntimeError( + f"Stale paper projection '{target_id}': embedding model " + "changed; re-put the paper result" + ) + if ( + embedding_dimensions is not None + and dimension != embedding_dimensions + ): + raise RuntimeError( + f"Stale paper projection '{target_id}': embedding dimension " + "changed; re-put the paper result" + ) + if ( + str(row["projection_kind"]) != "text_embedding" + or str(row["modality"]) != "text" + or str(row["projection_schema_version"]) + != _PROJECTION_SCHEMA_VERSION + or str(row["artifact_schema_version"]) + != str(row["current_schema_version"]) + or str(row["artifact_canonical_hash"]) + != str(row["current_canonical_hash"]) + or str(row["source_content_hash"]) + != str(row["current_source_content_hash"]) + ): + raise RuntimeError( + f"Stale paper projection '{target_id}': projection, source, " + "or artifact metadata changed; re-put the paper result" + ) + try: + artifact_id = UUID(str(row["artifact_id"])) + source_id = UUID(str(row["source_revision_id"])) + member_id = ( + UUID(str(row["member_id"])) + if row["member_id"] is not None + else None + ) + citations_value = json.loads(str(row["citations_json"])) + if not isinstance(citations_value, list): + raise ValueError + stored_citations = [ + Citation.model_validate(citation_value) + for citation_value in citations_value + ] + except (TypeError, ValueError, ValidationError) as exc: + raise RuntimeError( + f"Corrupt paper projection '{target_id}': invalid metadata" + ) from exc + + artifact = paper_artifacts.get(artifact_id) + if artifact is None: + artifact = self.get_paper_artifact(artifact_id) + paper_artifacts[artifact_id] = artifact + source = paper_sources.get(source_id) + if source is None: + source = self.get_paper_source(source_id) + paper_sources[source_id] = source + artifact_kind = str(row["artifact_kind"]) + if ( + artifact.source_revision_id != source_id + or artifact.artifact_kind != artifact_kind + or source.source.kind != str(row["source_kind"]) + or _timestamp(source.as_of, "PaperSourceRevision.as_of") + != float(row["as_of"]) + or _timestamp( + source.available_at, + "PaperSourceRevision.available_at", + ) + != float(row["available_at"]) + ): + raise RuntimeError( + f"Stale paper projection '{target_id}': canonical locator " + "or source evidence mismatch" + ) + + if member_id is None: + if not isinstance(artifact, PaperGlobalSummary): + raise RuntimeError( + f"Stale paper projection '{target_id}': chunk-set " + "aggregate is not searchable" + ) + expected_target_id = f"artifact:{artifact.id}" + expected_text = artifact.summary + expected_citations = [ + Citation( + source_id=str(source_id), + page=citation.page_number, + quote=citation.quote, + ) + for citation in artifact.citations + ] + else: + if not isinstance(artifact, PaperChunkSet): + raise RuntimeError( + f"Stale paper projection '{target_id}': summary has a " + "search member" + ) + chunk = next( + ( + value + for value in artifact.chunks + if value.chunk_id == member_id + ), + None, + ) + if chunk is None: + raise RuntimeError( + f"Stale paper projection '{target_id}': canonical " + "chunk is missing" + ) + expected_target_id = ( + f"artifact-member:{artifact.id}:{chunk.chunk_id}" + ) + expected_text = chunk.text + expected_citations = [ + Citation( + source_id=str(source_id), + page=span.page_number, + char_offset=span.start_char, + quote=chunk.text[:500], + ) + for span in chunk.source_spans + ] + expected_projection_hash = hashlib.sha256( + expected_text.encode("utf-8") + ).hexdigest() + if any( + ( + target_id != expected_target_id, + str(row["matched_text"]) != expected_text, + str(row["projection_hash"]) != expected_projection_hash, + stored_citations != expected_citations, + ) + ): + raise RuntimeError( + f"Stale paper projection '{target_id}': canonical text, " + "hash, or citations mismatch" + ) + record = _IndexRecord( + target_id=target_id, + owner_kind="paper", + item_id=artifact_id, + node_id=member_id, + item_type=artifact_kind, + source_revision_id=source_id, + artifact_kind=artifact_kind, + projection_kind=str(row["projection_kind"]), + projection_version=str(row["projection_schema_version"]), + embedding_model=str(row["embedding_model"]), + projection_hash=expected_projection_hash, + matched_text=expected_text, + as_of=float(row["as_of"]), + available_at=float(row["available_at"]), + source_kind=str(row["source_kind"]), + confidence="high", + tags=frozenset(), + tree_id=None, + dimension=dimension, + embedding=bytes(row["embedding"]), + ) + records.append(record) return records def close(self) -> None: diff --git a/quantmind/library/_types.py b/quantmind/library/_types.py index 9fd970e..9c3db23 100644 --- a/quantmind/library/_types.py +++ b/quantmind/library/_types.py @@ -6,7 +6,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator -from quantmind.knowledge import Citation, SourceRef +from quantmind.knowledge import ArtifactLocator, Citation, SourceRef class SemanticQuery(BaseModel): @@ -15,6 +15,7 @@ class SemanticQuery(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) text: str = Field(min_length=1) + artifact_kinds: list[str] | None = None item_types: list[str] | None = None source_kinds: ( list[ @@ -55,11 +56,26 @@ def _cutoffs_must_be_timezone_aware( return value +class SearchProjection(BaseModel): + """Rebuildable projection details used to rank one semantic hit.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + kind: Literal["text_embedding"] = "text_embedding" + version: str + modality: Literal["text"] = "text" + model: str + dimensions: int = Field(ge=1) + content_hash: str = Field(pattern=r"^[0-9a-f]{64}$") + + class SemanticHit(BaseModel): """Auditable evidence returned by semantic ranking.""" model_config = ConfigDict(extra="forbid", frozen=True) + locator: ArtifactLocator + projection: SearchProjection item_id: UUID node_id: UUID | None item_type: str diff --git a/quantmind/library/local.py b/quantmind/library/local.py index a2a8c58..9034ec9 100644 --- a/quantmind/library/local.py +++ b/quantmind/library/local.py @@ -6,7 +6,18 @@ from typing_extensions import Self -from quantmind.knowledge import BaseKnowledge, TreeKnowledge +from quantmind.knowledge import ( + ArtifactLocator, + BaseKnowledge, + Citation, + PaperChunkSet, + PaperFlowResult, + PaperGlobalSummary, + PaperSourceRevision, + ResolvedPaperArtifact, + TreeKnowledge, + TreeNode, +) from quantmind.library._internal.index_embeddings import ( _EmbeddingProvider, _OpenAIEmbeddingProvider, @@ -20,10 +31,15 @@ from quantmind.library._internal.retrieval_targets import ( _PROJECTION_SCHEMA_VERSION, _project_knowledge, + _project_paper, _RetrievalTarget, ) from quantmind.library._internal.sqlite_store import _SQLiteStore -from quantmind.library._types import SemanticHit, SemanticQuery +from quantmind.library._types import ( + SearchProjection, + SemanticHit, + SemanticQuery, +) class LocalKnowledgeLibrary: @@ -141,6 +157,78 @@ async def put(self, item: BaseKnowledge) -> None: ) self._index = None + async def put_paper(self, result: PaperFlowResult) -> None: + """Persist one source-first paper result and all required projections. + + Embeddings are prepared before the SQLite transaction. A provider + failure therefore leaves no partial source, artifact, lineage, or + required projection records. + """ + async with self._lock: + store = self._store + if store is None: + raise RuntimeError("LocalKnowledgeLibrary is closed") + prepared = store.prepare_put_paper(result) + targets = _project_paper(result) + artifacts = { + result.chunk_set.id: result.chunk_set, + result.global_summary.id: result.global_summary, + } + affected: list[_RetrievalTarget] = [] + vectors: dict[str, tuple[bytes, int]] = {} + for target in targets: + stored = prepared.existing_embeddings.get(target.target_id) + artifact = artifacts[target.artifact_id] + needs_embedding = stored is None + if stored is not None: + needs_embedding = any( + ( + stored.embedding_model != self._embedding_model, + stored.projection_hash != target.projection_hash, + stored.source_content_hash + != result.source_revision.source.content_hash, + stored.knowledge_schema_version + != artifact.schema_version, + stored.projection_schema_version + != _PROJECTION_SCHEMA_VERSION, + self._embedding_dimensions is not None + and stored.dimension != self._embedding_dimensions, + ) + ) + if needs_embedding: + affected.append(target) + continue + assert stored is not None + _decode_stored_vector( + stored.embedding, + stored.dimension, + target.target_id, + ) + vectors[target.target_id] = ( + stored.embedding, + stored.dimension, + ) + if affected: + provider_values = await self._embedding_provider.embed( + [target.text for target in affected], + model=self._embedding_model, + dimensions=self._embedding_dimensions, + ) + generated = _coerce_provider_vectors( + provider_values, + expected_count=len(affected), + expected_dimensions=self._embedding_dimensions, + ) + for target, vector in zip(affected, generated, strict=True): + vectors[target.target_id] = _encode_vector(vector) + store.put_paper( + prepared, + targets, + vectors, + embedding_model=self._embedding_model, + ) + self._index = None + async def get(self, item_id: UUID) -> BaseKnowledge: """Return validated canonical knowledge or report not-found/stale data.""" async with self._lock: @@ -149,6 +237,58 @@ async def get(self, item_id: UUID) -> BaseKnowledge: raise RuntimeError("LocalKnowledgeLibrary is closed") return store.get(item_id) + async def get_paper( + self, + source_revision_id: UUID, + *, + chunk_set_id: UUID | None = None, + summary_id: UUID | None = None, + ) -> PaperFlowResult: + """Return one unambiguous source/chunk-set/summary combination.""" + async with self._lock: + store = self._store + if store is None: + raise RuntimeError("LocalKnowledgeLibrary is closed") + return store.get_paper_result( + source_revision_id, + chunk_set_id=chunk_set_id, + summary_id=summary_id, + ) + + async def get_artifact( + self, artifact_id: UUID + ) -> PaperChunkSet | PaperGlobalSummary: + """Return one validated canonical paper artifact aggregate.""" + async with self._lock: + store = self._store + if store is None: + raise RuntimeError("LocalKnowledgeLibrary is closed") + return store.get_paper_artifact(artifact_id) + + async def resolve( + self, locator: ArtifactLocator + ) -> BaseKnowledge | TreeNode | ResolvedPaperArtifact: + """Resolve a semantic locator to its canonical aggregate or member.""" + async with self._lock: + store = self._store + if store is None: + raise RuntimeError("LocalKnowledgeLibrary is closed") + if locator.source_revision_id is not None: + return store.resolve_paper_locator(locator) + item = store.get(locator.artifact_id) + if locator.artifact_kind != item.item_type: + raise KeyError("Knowledge artifact locator metadata mismatch") + if locator.member_id is None: + return item + if not isinstance(item, TreeKnowledge): + raise KeyError("Flat knowledge artifacts do not have members") + try: + return item.nodes[locator.member_id] + except KeyError as exc: + raise KeyError( + f"Knowledge node '{locator.member_id}' not found" + ) from exc + async def search(self, query: SemanticQuery) -> list[SemanticHit]: """Rank filtered targets through the private LlamaIndex retriever.""" async with self._lock: @@ -181,47 +321,130 @@ async def search(self, query: SemanticQuery) -> list[SemanticHit]: ) canonical: dict[UUID, BaseKnowledge] = {} + paper_artifacts: dict[UUID, PaperChunkSet | PaperGlobalSummary] = {} + paper_sources: dict[UUID, PaperSourceRevision] = {} hits: list[SemanticHit] = [] for result in ranked: record = result.record - item = canonical.get(record.item_id) - if item is None: - try: - item = store.get(record.item_id) - except KeyError as exc: + locator = ArtifactLocator( + source_revision_id=record.source_revision_id, + artifact_id=record.item_id, + artifact_kind=record.artifact_kind, + member_id=record.node_id, + ) + if record.owner_kind == "paper": + source_id = record.source_revision_id + if source_id is None: raise RuntimeError( - f"Stale index data for item '{record.item_id}': " - "canonical knowledge is missing" - ) from exc - canonical[record.item_id] = item - if isinstance(item, TreeKnowledge): + f"Stale paper projection '{record.target_id}': " + "source locator is missing" + ) + source = paper_sources.get(source_id) + if source is None: + source = store.get_paper_source(source_id) + paper_sources[source_id] = source + artifact = paper_artifacts.get(record.item_id) + if artifact is None: + artifact = store.get_paper_artifact(record.item_id) + paper_artifacts[record.item_id] = artifact if record.node_id is None: - citations = item.root().citations - else: - node = item.nodes.get(record.node_id) - if node is None: + if not isinstance(artifact, PaperGlobalSummary): raise RuntimeError( - f"Stale index data for target " - f"'{record.target_id}': canonical node is missing" + f"Stale paper projection '{record.target_id}': " + "chunk-set aggregate is not searchable" ) - citations = node.citations - else: - if record.node_id is not None: + citations = [ + Citation( + source_id=str(source_id), + page=citation.page_number, + quote=citation.quote, + ) + for citation in artifact.citations + ] + elif isinstance(artifact, PaperChunkSet): + chunk = next( + ( + value + for value in artifact.chunks + if value.chunk_id == record.node_id + ), + None, + ) + if chunk is None: + raise RuntimeError( + f"Stale paper projection '{record.target_id}': " + "canonical chunk is missing" + ) + citations = [ + Citation( + source_id=str(source_id), + page=span.page_number, + char_offset=span.start_char, + quote=chunk.text[:500], + ) + for span in chunk.source_spans + ] + else: raise RuntimeError( - f"Stale index data for target '{record.target_id}': " - "node target belongs to non-tree knowledge" + f"Stale paper projection '{record.target_id}': " + "summary unexpectedly has a member" ) - citations = item.citations + as_of = source.as_of + available_at = source.available_at + source_ref = source.source + item_type = artifact.artifact_kind + else: + item = canonical.get(record.item_id) + if item is None: + try: + item = store.get(record.item_id) + except KeyError as exc: + raise RuntimeError( + f"Stale index data for item '{record.item_id}': " + "canonical knowledge is missing" + ) from exc + canonical[record.item_id] = item + if isinstance(item, TreeKnowledge): + if record.node_id is None: + citations = item.root().citations + else: + node = item.nodes.get(record.node_id) + if node is None: + raise RuntimeError( + f"Stale index data for target " + f"'{record.target_id}': canonical node is " + "missing" + ) + citations = node.citations + else: + if record.node_id is not None: + raise RuntimeError( + f"Stale index data for target " + f"'{record.target_id}': node target belongs to " + "non-tree knowledge" + ) + citations = item.citations + as_of = item.as_of + available_at = item.available_at + source_ref = item.source + item_type = item.item_type hits.append( SemanticHit( + locator=locator, + projection=SearchProjection( + version=record.projection_version, + model=record.embedding_model, + dimensions=record.dimension, + content_hash=record.projection_hash, + ), item_id=record.item_id, node_id=record.node_id, - item_type=item.item_type, + item_type=item_type, score=result.score, matched_text=record.matched_text, - as_of=item.as_of, - available_at=item.available_at, - source=item.source, + as_of=as_of, + available_at=available_at, + source=source_ref, citations=citations, ) ) diff --git a/quantmind/preprocess/fetch/_types.py b/quantmind/preprocess/fetch/_types.py index bee1e92..122773b 100644 --- a/quantmind/preprocess/fetch/_types.py +++ b/quantmind/preprocess/fetch/_types.py @@ -52,5 +52,6 @@ class RawPaper(Fetched): authors: tuple[str, ...] = () abstract: str | None = None published_at: datetime | None = None + updated_at: datetime | None = None primary_category: str | None = None categories: tuple[str, ...] = () diff --git a/quantmind/preprocess/fetch/arxiv.py b/quantmind/preprocess/fetch/arxiv.py index 25b2baf..959067b 100644 --- a/quantmind/preprocess/fetch/arxiv.py +++ b/quantmind/preprocess/fetch/arxiv.py @@ -93,6 +93,10 @@ async def fetch_arxiv(id_or_url: str) -> RawPaper: arxiv_id = _extract_arxiv_id(id_or_url) result = await asyncio.to_thread(_fetch_metadata_sync, arxiv_id) + get_short_id = getattr(result, "get_short_id", None) + resolved_arxiv_id = ( + str(get_short_id()) if callable(get_short_id) else arxiv_id + ) pdf_url = result.pdf_url if not pdf_url: raise LookupError(f"arxiv result has no pdf_url for {arxiv_id!r}") @@ -101,17 +105,21 @@ async def fetch_arxiv(id_or_url: str) -> RawPaper: response = await client.get(pdf_url, headers=headers) response.raise_for_status() pdf_bytes = response.content + fetched_at = datetime.now(timezone.utc) return RawPaper( bytes=pdf_bytes, content_type="application/pdf", source_url=pdf_url, headers={}, - arxiv_id=arxiv_id, + resolved_url=str(response.url), + fetched_at=fetched_at, + arxiv_id=resolved_arxiv_id, title=result.title, authors=tuple(str(a) for a in result.authors), abstract=result.summary, published_at=_to_utc(result.published), + updated_at=_to_utc(getattr(result, "updated", None)), primary_category=result.primary_category, categories=tuple(str(c) for c in result.categories), ) diff --git a/quantmind/rag/document.py b/quantmind/rag/document.py index 5810dbb..4f5c3b7 100644 --- a/quantmind/rag/document.py +++ b/quantmind/rag/document.py @@ -1,5 +1,6 @@ """Page-aware document chunking and retrieval through LlamaIndex.""" +import hashlib import json from dataclasses import dataclass from typing import Any @@ -32,6 +33,8 @@ class ParsedChunk: text: str source_hash: str page_number: int + start_char: int + end_char: int block_boxes: tuple[BoundingBox, ...] screenshot_path: str | None image_paths: tuple[str, ...] @@ -83,14 +86,29 @@ def _to_llama_documents(document: ParsedDocument) -> list[Document]: def _node_to_chunk(node: BaseNode) -> ParsedChunk: metadata = node.metadata + text = node.get_content(metadata_mode=MetadataMode.NONE) + start_value = getattr(node, "start_char_idx", None) + end_value = getattr(node, "end_char_idx", None) + start_char = int(start_value) if start_value is not None else 0 + end_char = ( + int(end_value) if end_value is not None else start_char + len(text) + ) + identity = hashlib.sha256( + ( + f"{metadata['source_hash']}:{metadata['page_number']}:" + f"{start_char}:{end_char}:{text}" + ).encode("utf-8") + ).hexdigest() boxes = tuple( BoundingBox(*values) for values in json.loads(metadata["block_boxes"]) ) return ParsedChunk( - chunk_id=node.node_id, - text=node.get_content(metadata_mode=MetadataMode.NONE), + chunk_id=identity, + text=text, source_hash=str(metadata["source_hash"]), page_number=int(metadata["page_number"]), + start_char=start_char, + end_char=end_char, block_boxes=boxes, screenshot_path=str(metadata["screenshot_path"]) or None, image_paths=tuple(json.loads(metadata["image_paths"])), diff --git a/scripts/examples/build_ai_infrastructure_bundle.py b/scripts/examples/build_ai_infrastructure_bundle.py index b68e0f3..5cfe32f 100644 --- a/scripts/examples/build_ai_infrastructure_bundle.py +++ b/scripts/examples/build_ai_infrastructure_bundle.py @@ -7,7 +7,7 @@ from dotenv import load_dotenv -from quantmind.knowledge import BaseKnowledge, Earnings, News, Paper +from quantmind.knowledge import BaseKnowledge, Earnings, LegacyPaper, News from quantmind.library import LocalKnowledgeLibrary _ROOT = Path(__file__).parents[2] @@ -20,7 +20,7 @@ _KNOWLEDGE_TYPES: dict[str, type[BaseKnowledge]] = { "earnings": Earnings, "news": News, - "paper": Paper, + "paper": LegacyPaper, } diff --git a/scripts/verify_pdf_rag_e2e.py b/scripts/verify_pdf_rag_e2e.py index 673c22c..1abe3dc 100644 --- a/scripts/verify_pdf_rag_e2e.py +++ b/scripts/verify_pdf_rag_e2e.py @@ -1,54 +1,245 @@ #!/usr/bin/env python3 -"""Run the bounded live Transformer PDF parsing and retrieval smoke test.""" +"""Run the bounded live *Attention Is All You Need* Paper Flow V1 slice.""" import asyncio +import os +import tempfile +from pathlib import Path +from typing import Any -from quantmind.preprocess.fetch import fetch_arxiv -from quantmind.preprocess.format import parse_pdf -from quantmind.rag import ( - SentenceSplitterConfig, - chunk_parsed_document, - retrieve_parsed_document, -) +from dotenv import load_dotenv + +from quantmind.configs import PaperFlowCfg +from quantmind.configs.paper import ArxivIdentifier +from quantmind.flows import paper_flow +from quantmind.knowledge import PaperChunk, PaperGlobalSummary +from quantmind.library import LocalKnowledgeLibrary, SemanticQuery _ARXIV_ID = "1706.03762v7" +_EMBEDDING_MODEL = "text-embedding-3-small" +_SUMMARY_QUERY = "What is the paper's central contribution?" +_CHUNK_QUERY = "How does multi-head attention work?" -async def main() -> int: - """Fetch the pinned paper and verify page-aware retrieval.""" - try: - paper = await asyncio.wait_for(fetch_arxiv(_ARXIV_ID), timeout=120) - document = await asyncio.wait_for(parse_pdf(paper.bytes), timeout=120) - chunks = chunk_parsed_document( - document, - config=SentenceSplitterConfig(chunk_size=512, chunk_overlap=64), +async def _search_and_resolve( + library: LocalKnowledgeLibrary, +) -> tuple[list[Any], list[Any], list[Any]]: + summary_hits = await library.search( + SemanticQuery( + text=_SUMMARY_QUERY, + artifact_kinds=["paper_summary"], + top_k=3, ) - hits = retrieve_parsed_document( - chunks, - "How does multi-head attention work?", + ) + chunk_hits = await library.search( + SemanticQuery( + text=_CHUNK_QUERY, + artifact_kinds=["paper_chunk_set"], top_k=5, ) + ) + resolved = [ + await library.resolve(hit.locator) + for hit in (*summary_hits, *chunk_hits) + ] + return summary_hits, chunk_hits, resolved + + +async def _run_vertical_slice() -> dict[str, Any]: + with tempfile.TemporaryDirectory(prefix="quantmind-paper-v1-") as directory: + root = Path(directory) + result = await paper_flow( + ArxivIdentifier(id=_ARXIV_ID), + cfg=PaperFlowCfg( + model="gpt-4o-mini", + output_dir=str(root / "assets"), + timeout_seconds=240, + max_summary_tool_calls=12, + max_summary_concurrency=2, + max_summary_input_tokens=120_000, + max_summary_output_tokens=4_096, + min_summary_citations=3, + min_summary_pages=2, + ), + ) + database = root / "library.db" + library = await LocalKnowledgeLibrary.open( + database, + embedding_model=_EMBEDDING_MODEL, + ) + try: + await library.put_paper(result) + ( + first_summary_hits, + first_chunk_hits, + first_resolved, + ) = await _search_and_resolve(library) + finally: + await library.close() + + reopened = await LocalKnowledgeLibrary.open( + database, + embedding_model=_EMBEDDING_MODEL, + ) + try: + restored = await reopened.get_paper( + result.source_revision.id, + chunk_set_id=result.chunk_set.id, + summary_id=result.global_summary.id, + ) + ( + second_summary_hits, + second_chunk_hits, + second_resolved, + ) = await _search_and_resolve(reopened) + finally: + await reopened.close() + + return { + "arxiv_id": result.source_revision.arxiv_id, + "page_count": len(result.source_revision.parsed.pages), + "text_page_count": sum( + bool(page.text.strip()) + for page in result.source_revision.parsed.pages + ), + "asset_count": len(result.source_revision.assets), + "screenshot_pages": [ + page.page_number + for page in result.source_revision.parsed.pages + if page.screenshot_asset_id is not None + ], + "chunk_asset_reference_count": sum( + len(span.asset_ids) + for chunk in result.chunk_set.chunks + for span in chunk.source_spans + ), + "chunk_count": len(result.chunk_set.chunks), + "summary": result.global_summary.summary, + "citation_count": len(result.global_summary.citations), + "citation_pages": sorted( + { + citation.page_number + for citation in result.global_summary.citations + } + ), + "first_summary_scores": [hit.score for hit in first_summary_hits], + "first_chunk_scores": [hit.score for hit in first_chunk_hits], + "second_summary_scores": [hit.score for hit in second_summary_hits], + "second_chunk_scores": [hit.score for hit in second_chunk_hits], + "first_summary_resolved": any( + isinstance(value, PaperGlobalSummary) for value in first_resolved + ), + "second_summary_resolved": any( + isinstance(value, PaperGlobalSummary) for value in second_resolved + ), + "first_multi_head_pages": [ + value.source_spans[0].page_number + for value in first_resolved + if isinstance(value, PaperChunk) + and "multi-head attention" in value.text.lower() + ], + "second_multi_head_pages": [ + value.source_spans[0].page_number + for value in second_resolved + if isinstance(value, PaperChunk) + and "multi-head attention" in value.text.lower() + ], + "restored": restored.chunk_set.id == result.chunk_set.id + and restored.global_summary.id == result.global_summary.id, + "embedding_model": _EMBEDDING_MODEL, + } + + +def _summary_has_required_coverage(summary: str) -> bool: + text = summary.lower() + attention_only = "attention" in text and any( + marker in text + for marker in ( + "attention-only", + "attention based", + "attention-based", + "solely", + "entirely", + "exclusively", + "purely", + "eliminat", + "without recurrence", + "remov", + "replac", + "eschew", + ) + ) + return ( + all( + any(term in text for term in group) + for group in ( + ("recurrence", "recurrent"), + ("convolution", "convolutional"), + ( + "encoder-decoder", + "encoder", + "decoder", + "multi-head attention", + "multihead attention", + ), + ("translation", "training efficiency", "training time"), + ) + ) + and attention_only + ) + + +def _passed(snapshot: dict[str, Any]) -> bool: + return bool( + snapshot["arxiv_id"] == _ARXIV_ID + and snapshot["page_count"] == 15 + and snapshot["text_page_count"] == 15 + and snapshot["asset_count"] > 1 + and len(snapshot["screenshot_pages"]) == 15 + and snapshot["chunk_asset_reference_count"] > 0 + and snapshot["chunk_count"] > 0 + and snapshot["summary"] + and snapshot["citation_count"] >= 3 + and len(snapshot["citation_pages"]) >= 2 + and snapshot["first_summary_scores"] + and snapshot["second_summary_scores"] + and snapshot["first_multi_head_pages"] + and snapshot["second_multi_head_pages"] + and snapshot["first_summary_resolved"] + and snapshot["second_summary_resolved"] + and snapshot["restored"] + and snapshot["embedding_model"] == _EMBEDDING_MODEL + and _summary_has_required_coverage(snapshot["summary"]) + ) + + +async def main() -> int: + """Run and report the timeout-bounded live vertical slice.""" + load_dotenv() + if not os.getenv("OPENAI_API_KEY"): + print("[FAIL] paper-flow-v1: OPENAI_API_KEY is not set") + return 1 + try: + snapshot = await asyncio.wait_for(_run_vertical_slice(), timeout=480) except Exception as exc: - print(f"[FAIL] pdf-rag: {type(exc).__name__}: {exc}") + print(f"[FAIL] paper-flow-v1: {type(exc).__name__}: {exc}") return 1 - relevant = [ - hit for hit in hits if "multi-head attention" in hit.chunk.text.lower() - ] - passed = ( - paper.arxiv_id == _ARXIV_ID - and len(document.pages) == 15 - and bool(chunks) - and bool(relevant) - and all(hit.chunk.page_number >= 1 for hit in hits) - ) - state = "PASS" if passed else "FAIL" + state = "PASS" if _passed(snapshot) else "FAIL" print( - f"[{state}] pdf-rag: arxiv={paper.arxiv_id} " - f"pages={len(document.pages)} chunks={len(chunks)} " - f"top_pages={[hit.chunk.page_number for hit in hits]}" + f"[{state}] paper-flow-v1: arxiv={snapshot['arxiv_id']} " + f"pages={snapshot['page_count']} " + f"text_pages={snapshot['text_page_count']} " + f"chunks={snapshot['chunk_count']} assets={snapshot['asset_count']} " + f"chunk_asset_refs={snapshot['chunk_asset_reference_count']} " + f"asset_pages={snapshot['screenshot_pages']} " + f"citation_pages={snapshot['citation_pages']} " + f"summary_scores={snapshot['second_summary_scores']} " + f"chunk_scores={snapshot['second_chunk_scores']} " + f"multi_head_pages={snapshot['second_multi_head_pages']}" ) - return 0 if passed else 1 + print(f"Summary:\n{snapshot['summary']}") + return 0 if state == "PASS" else 1 if __name__ == "__main__": diff --git a/tests/configs/test_paper.py b/tests/configs/test_paper.py index 3a9baf9..ea350b8 100644 --- a/tests/configs/test_paper.py +++ b/tests/configs/test_paper.py @@ -19,10 +19,18 @@ class PaperFlowCfgTests(unittest.TestCase): def test_defaults(self): cfg = PaperFlowCfg() - self.assertEqual(cfg.model, "gpt-4o") - self.assertTrue(cfg.extract_methodology) - self.assertTrue(cfg.extract_limitations) - self.assertIsNone(cfg.asset_class_hint) + self.assertEqual(cfg.model, "gpt-4o-mini") + self.assertEqual(cfg.max_turns, 16) + self.assertEqual(cfg.chunk_size, 512) + self.assertEqual(cfg.chunk_overlap, 64) + self.assertEqual(cfg.max_summary_tool_calls, 12) + self.assertEqual(cfg.min_summary_citations, 3) + + def test_invalid_overlap_or_coverage_bounds_are_rejected(self): + with self.assertRaises(ValidationError): + PaperFlowCfg(chunk_size=64, chunk_overlap=64) + with self.assertRaises(ValidationError): + PaperFlowCfg(min_summary_citations=1, min_summary_pages=2) class PaperInputDiscriminatedTests(unittest.TestCase): diff --git a/tests/flows/test_paper.py b/tests/flows/test_paper.py index 2a143c7..5219a87 100644 --- a/tests/flows/test_paper.py +++ b/tests/flows/test_paper.py @@ -1,13 +1,12 @@ -"""Tests for ``quantmind.flows.paper``.""" +"""Offline tests for source-first ``paper_flow`` behavior.""" +import asyncio import unittest from datetime import datetime, timezone from pathlib import Path -from typing import Any -from unittest.mock import AsyncMock, MagicMock, patch -from uuid import uuid4 +from unittest.mock import AsyncMock, patch -from agents import RunHooks +from agents import ModelSettings from quantmind.configs import PaperFlowCfg from quantmind.configs.paper import ( @@ -17,343 +16,321 @@ LocalFilePath, RawText, ) +from quantmind.flows._paper_summary import ( + PaperSummaryBudgetExceeded, + PaperSummaryCitationDraft, + PaperSummaryDraft, + _bounded_model_settings, + _SummaryBudget, +) from quantmind.flows.paper import ( + PaperCitationValidationError, UnsupportedContentTypeError, - _compose_instructions, - _fetch_and_format, - _format_by_content_type, - _format_input, + _build_global_summary, + _fetch_paper_source, paper_flow, ) -from quantmind.knowledge import Paper, SourceRef, TreeNode from quantmind.preprocess.fetch import Fetched, RawPaper - - -def _stub_paper() -> Paper: - root_id = uuid4() - root = TreeNode(node_id=root_id, title="root", summary="stub") - return Paper( - as_of=datetime(2026, 5, 7, tzinfo=timezone.utc), - source=SourceRef( - kind="arxiv", - uri="arxiv:2604.12345", - fetched_at=datetime(2026, 5, 7, tzinfo=timezone.utc), - ), - root_node_id=root_id, - nodes={root_id: root}, - ) - - -def _patch_runner(return_value: Any) -> Any: - return patch( - "quantmind.flows.paper.run_with_observability", - new=AsyncMock(return_value=return_value), - ) - - -class FormatByContentTypeTests(unittest.IsolatedAsyncioTestCase): - async def test_pdf_dispatches_to_pdf_to_markdown(self) -> None: - raw = Fetched(bytes=b"%PDF-x", content_type="application/pdf") - with patch( - "quantmind.flows.paper.pdf_to_markdown", - new=AsyncMock(return_value="MD"), - ) as pdf_mock: - md = await _format_by_content_type(raw) - pdf_mock.assert_awaited_once_with(b"%PDF-x") - self.assertEqual(md, "MD") - - async def test_html_dispatches_to_html_to_markdown(self) -> None: - raw = Fetched( - bytes="hi".encode("utf-8"), - content_type="text/html; charset=utf-8", - ) - with patch( - "quantmind.flows.paper.html_to_markdown", - new=AsyncMock(return_value="HTML-MD"), - ) as html_mock: - md = await _format_by_content_type(raw) - html_mock.assert_awaited_once_with("hi") - self.assertEqual(md, "HTML-MD") - - async def test_markdown_passes_through(self) -> None: - raw = Fetched( - bytes=b"# heading\n\nbody", - content_type="text/markdown", +from tests.paper_helpers import build_paper_result + +_FIXTURE = ( + Path(__file__).resolve().parents[1] + / "fixtures" + / "paper" + / "golden" + / "paper.pdf" +) +_WHEN = datetime(2017, 12, 6, tzinfo=timezone.utc) + + +class _FakeSummaryProvider: + def __init__(self, *, fail: Exception | None = None) -> None: + self.fail = fail + self.calls = [] + + async def summarize(self, source, chunk_set, *, cfg): + self.calls.append((source, chunk_set, cfg)) + if self.fail is not None: + raise self.fail + selected = [] + pages: set[int] = set() + for chunk in chunk_set.chunks: + page = chunk.source_spans[0].page_number + if page not in pages or len(selected) < cfg.min_summary_citations: + selected.append((chunk, page)) + pages.add(page) + if ( + len(selected) >= cfg.min_summary_citations + and len(pages) >= cfg.min_summary_pages + ): + break + return PaperSummaryDraft( + summary=( + "The paper studies a source-first test methodology and reports " + "deterministic evidence across physical pages." + ), + citations=tuple( + PaperSummaryCitationDraft( + chunk_index=chunk.position, + page_number=page, + ) + for chunk, page in selected + ), ) - md = await _format_by_content_type(raw) - self.assertEqual(md, "# heading\n\nbody") - async def test_plain_text_passes_through(self) -> None: - raw = Fetched(bytes=b"plain", content_type="text/plain") - md = await _format_by_content_type(raw) - self.assertEqual(md, "plain") - async def test_unsupported_content_type_raises(self) -> None: - raw = Fetched(bytes=b"\x00\x00", content_type="application/zip") - with self.assertRaises(UnsupportedContentTypeError): - await _format_by_content_type(raw) - - -class FetchAndFormatTests(unittest.IsolatedAsyncioTestCase): - async def test_arxiv_branch(self) -> None: - raw_paper = RawPaper( - bytes=b"%PDF", - content_type="application/pdf", - source_url="http://arxiv.org/pdf/2604.12345.pdf", - arxiv_id="2604.12345", - title="Momentum", - authors=("Alice", "Bob"), - ) - with ( - patch( - "quantmind.flows.paper.fetch_arxiv", - new=AsyncMock(return_value=raw_paper), - ) as fetch_mock, - patch( - "quantmind.flows.paper.pdf_to_markdown", - new=AsyncMock(return_value="MARKDOWN"), - ) as fmt_mock, - ): - md, meta = await _fetch_and_format(ArxivIdentifier(id="2604.12345")) - fetch_mock.assert_awaited_once_with("2604.12345") - fmt_mock.assert_awaited_once_with(b"%PDF") - self.assertEqual(md, "MARKDOWN") - self.assertEqual(meta["source"], "arxiv") - self.assertEqual(meta["arxiv_id"], "2604.12345") - self.assertEqual(meta["title"], "Momentum") - self.assertEqual(meta["authors"], ["Alice", "Bob"]) - - async def test_http_pdf_branch(self) -> None: - raw = Fetched( - bytes=b"%PDF", - content_type="application/pdf", - source_url="http://example/x.pdf", +class PaperFlowTests(unittest.IsolatedAsyncioTestCase): + async def test_local_pdf_builds_chunks_before_summary(self) -> None: + provider = _FakeSummaryProvider() + cfg = PaperFlowCfg(chunk_size=256, chunk_overlap=32) + + result = await paper_flow( + LocalFilePath(path=_FIXTURE), + cfg=cfg, + _summary_provider=provider, ) - with ( - patch( - "quantmind.flows.paper.fetch_url", - new=AsyncMock(return_value=raw), - ) as fetch_mock, - patch( - "quantmind.flows.paper.pdf_to_markdown", - new=AsyncMock(return_value="PDFMD"), + + self.assertEqual(len(provider.calls), 1) + source_seen, chunk_set_seen, _ = provider.calls[0] + self.assertIs(source_seen, result.source_revision) + self.assertIs(chunk_set_seen, result.chunk_set) + self.assertEqual(len(result.source_revision.parsed.pages), 4) + self.assertTrue(result.chunk_set.chunks) + self.assertTrue(result.global_summary.summary) + self.assertGreaterEqual(len(result.global_summary.citations), 3) + self.assertGreaterEqual( + len( + { + citation.page_number + for citation in result.global_summary.citations + } ), - ): - md, meta = await _fetch_and_format( - HttpUrl(url="http://example/x.pdf") + 2, + ) + self.assertTrue( + all( + chunk.source_revision_id == result.source_revision.id + and chunk.source_spans + for chunk in result.chunk_set.chunks ) - fetch_mock.assert_awaited_once_with("http://example/x.pdf") - self.assertEqual(md, "PDFMD") - self.assertEqual(meta["source"], "web") - self.assertEqual(meta["content_type"], "application/pdf") + ) - async def test_local_file_branch(self) -> None: - raw = Fetched( - bytes=b"## body", - content_type="text/markdown", - source_url="file:///tmp/p.md", + async def test_exact_arxiv_source_facts_are_code_owned(self) -> None: + raw = RawPaper( + bytes=_FIXTURE.read_bytes(), + content_type="application/pdf", + source_url="https://arxiv.org/pdf/1706.03762v7.pdf", + resolved_url="https://arxiv.org/pdf/1706.03762v7.pdf", + fetched_at=_WHEN, + arxiv_id="1706.03762v7", + title="Attention Is All You Need", + authors=("Ashish Vaswani",), + published_at=_WHEN, + updated_at=_WHEN, ) with patch( - "quantmind.flows.paper.read_local_file", + "quantmind.flows.paper.fetch_arxiv", new=AsyncMock(return_value=raw), - ) as read_mock: - md, meta = await _fetch_and_format( - LocalFilePath(path=Path("/tmp/p.md")) + ): + result = await paper_flow( + ArxivIdentifier(id="1706.03762v7"), + cfg=PaperFlowCfg(chunk_size=256, chunk_overlap=32), + _summary_provider=_FakeSummaryProvider(), ) - read_mock.assert_awaited_once_with(Path("/tmp/p.md")) - self.assertEqual(md, "## body") - self.assertEqual(meta["source"], "local") - self.assertEqual(meta["path"], "/tmp/p.md") - self.assertEqual(meta["content_type"], "text/markdown") - - async def test_raw_text_branch(self) -> None: - md, meta = await _fetch_and_format(RawText(text="hello")) - self.assertEqual(md, "hello") - self.assertEqual(meta, {"source": "inline"}) - - async def test_doi_branch_raises_not_implemented(self) -> None: - with self.assertRaises(NotImplementedError) as ctx: - await _fetch_and_format(DoiIdentifier(doi="10.1234/abcd")) - self.assertIn("DOI", str(ctx.exception)) - -class ComposeInstructionsTests(unittest.TestCase): - def test_default_renders_cfg_flags(self) -> None: - cfg = PaperFlowCfg( - extract_methodology=False, - extract_limitations=True, - asset_class_hint="equities", + self.assertEqual(result.source_revision.arxiv_id, "1706.03762v7") + self.assertEqual(result.source_revision.source.kind, "arxiv") + self.assertEqual( + result.source_revision.title, "Attention Is All You Need" ) - out = _compose_instructions( - "go {extract_methodology} {extract_limitations} " - "{asset_class_hint!r}", - None, - cfg, + self.assertEqual( + result.source_revision.source.content_hash, + result.source_revision.parsed.source_hash, ) - self.assertEqual(out, "go False True 'equities'") + self.assertFalse(hasattr(result.global_summary, "root_node_id")) - def test_extra_appended(self) -> None: - cfg = PaperFlowCfg() - out = _compose_instructions("BASE", "USER", cfg) - self.assertIn("BASE", out) - self.assertIn("Additional instructions:", out) - self.assertIn("USER", out) + async def test_summary_failure_prevents_flow_success(self) -> None: + provider = _FakeSummaryProvider(fail=RuntimeError("summary failed")) + with self.assertRaisesRegex(RuntimeError, "summary failed"): + await paper_flow( + LocalFilePath(path=_FIXTURE), + cfg=PaperFlowCfg(chunk_size=256, chunk_overlap=32), + _summary_provider=provider, + ) -class FormatInputTests(unittest.TestCase): - def test_tuple_authors_join_as_csv(self) -> None: - out = _format_input( - "BODY", - {"authors": ("Alice", "Bob"), "title": "x"}, + async def test_same_pdf_and_configs_have_idempotent_ids(self) -> None: + cfg = PaperFlowCfg(chunk_size=256, chunk_overlap=32) + first = await paper_flow( + LocalFilePath(path=_FIXTURE), + cfg=cfg, + _summary_provider=_FakeSummaryProvider(), + ) + second = await paper_flow( + LocalFilePath(path=_FIXTURE), + cfg=cfg, + _summary_provider=_FakeSummaryProvider(), ) - self.assertIn("authors: Alice, Bob", out) - self.assertIn("title: x", out) - self.assertIn("BODY", out) - def test_none_values_skipped(self) -> None: - out = _format_input("BODY", {"a": "1", "b": None}) - self.assertIn("a: 1", out) - self.assertNotIn("b:", out) + self.assertEqual(first.source_revision.id, second.source_revision.id) + self.assertEqual(first.chunk_set.id, second.chunk_set.id) + self.assertEqual(first.global_summary.id, second.global_summary.id) + self.assertEqual( + [chunk.chunk_id for chunk in first.chunk_set.chunks], + [chunk.chunk_id for chunk in second.chunk_set.chunks], + ) -class PaperFlowTests(unittest.IsolatedAsyncioTestCase): - async def test_happy_path_arxiv(self) -> None: - raw_paper = RawPaper( +class SourceDispatchTests(unittest.IsolatedAsyncioTestCase): + async def test_http_pdf_preserves_resolved_url_and_fetch_time(self) -> None: + raw = Fetched( bytes=b"%PDF", content_type="application/pdf", - arxiv_id="2604.12345", + source_url="https://example.test/input.pdf", + resolved_url="https://cdn.example.test/exact.pdf", + fetched_at=_WHEN, ) - stub = _stub_paper() - with ( - patch( - "quantmind.flows.paper.fetch_arxiv", - new=AsyncMock(return_value=raw_paper), - ), - patch( - "quantmind.flows.paper.pdf_to_markdown", - new=AsyncMock(return_value="MD"), - ), - _patch_runner(stub) as runner, + with patch( + "quantmind.flows.paper.fetch_url", + new=AsyncMock(return_value=raw), ): - out = await paper_flow(ArxivIdentifier(id="2604.12345")) - self.assertIs(out, stub) - runner.assert_awaited_once() - - async def test_extra_instructions_passed_to_agent(self) -> None: - seen: dict[str, Any] = {} + fetched = await _fetch_paper_source( + HttpUrl(url="https://example.test/input.pdf") + ) - def _capture_agent(*_a: Any, **kwargs: Any) -> Any: - seen.update(kwargs) - return MagicMock(name="agent", _name="paper_extractor") + self.assertEqual(fetched.uri, "https://cdn.example.test/exact.pdf") + self.assertEqual(fetched.fetched_at, _WHEN) - stub = _stub_paper() - with ( - patch("quantmind.flows.paper.Agent", side_effect=_capture_agent), - _patch_runner(stub), + async def test_non_pdf_and_raw_text_are_rejected_before_summary( + self, + ) -> None: + html = Fetched(bytes=b"", content_type="text/html") + with patch( + "quantmind.flows.paper.fetch_url", + new=AsyncMock(return_value=html), ): - await paper_flow( - RawText(text="hello"), - extra_instructions="EXTRA-USER-DIRECTIVE", - ) - self.assertIn("EXTRA-USER-DIRECTIVE", seen["instructions"]) - self.assertIn("structured QuantMind", seen["instructions"]) + with self.assertRaises(UnsupportedContentTypeError): + await _fetch_paper_source(HttpUrl(url="https://example.test")) + with self.assertRaises(UnsupportedContentTypeError): + await _fetch_paper_source(RawText(text="no physical pages")) - async def test_output_type_override_propagated(self) -> None: - seen: dict[str, Any] = {} + async def test_doi_requires_an_exact_pdf_resolver(self) -> None: + with self.assertRaises(NotImplementedError): + await _fetch_paper_source(DoiIdentifier(doi="10.1000/test")) - class MyPaper(Paper): - pass - def _capture_agent(*_a: Any, **kwargs: Any) -> Any: - seen.update(kwargs) - return MagicMock() +class CitationValidationTests(unittest.TestCase): + def test_unknown_chunk_page_and_quote_are_rejected(self) -> None: + result = build_paper_result() + cfg = PaperFlowCfg( + min_summary_citations=1, + min_summary_pages=1, + ) + invalid_drafts = ( + PaperSummaryDraft( + summary="x", + citations=( + PaperSummaryCitationDraft( + chunk_index=99, + page_number=1, + ), + ), + ), + PaperSummaryDraft( + summary="x", + citations=( + PaperSummaryCitationDraft( + chunk_index=0, + page_number=2, + ), + ), + ), + PaperSummaryDraft( + summary="x", + citations=( + PaperSummaryCitationDraft( + chunk_index=0, + page_number=1, + quote="not in source chunk", + ), + ), + ), + ) - with ( - patch("quantmind.flows.paper.Agent", side_effect=_capture_agent), - _patch_runner(_stub_paper()), - ): - await paper_flow(RawText(text="x"), output_type=MyPaper) - self.assertIs(seen["output_type"], MyPaper) - - async def test_extra_tools_and_guardrails_forwarded(self) -> None: - seen: dict[str, Any] = {} - - def _capture_agent(*_a: Any, **kwargs: Any) -> Any: - seen.update(kwargs) - return MagicMock() - - sentinel_tool = MagicMock(name="tool") - in_g = MagicMock() - out_g = MagicMock() - with ( - patch("quantmind.flows.paper.Agent", side_effect=_capture_agent), - _patch_runner(_stub_paper()), - ): - await paper_flow( - RawText(text="x"), - extra_tools=[sentinel_tool], - extra_input_guardrails=[in_g], - extra_output_guardrails=[out_g], - ) - self.assertEqual(seen["tools"], [sentinel_tool]) - self.assertEqual(seen["input_guardrails"], [in_g]) - self.assertEqual(seen["output_guardrails"], [out_g]) - - async def test_memory_accepted_as_no_op(self) -> None: - with ( - patch( - "quantmind.flows.paper.Agent", - return_value=MagicMock(), + for draft in invalid_drafts: + with self.subTest(draft=draft): + with self.assertRaises(PaperCitationValidationError): + _build_global_summary( + result.source_revision, + result.chunk_set, + draft, + cfg, + ) + + def test_configured_citation_and_page_coverage_is_enforced(self) -> None: + result = build_paper_result() + draft = PaperSummaryDraft( + summary="x", + citations=( + PaperSummaryCitationDraft(chunk_index=0, page_number=1), ), - _patch_runner(_stub_paper()) as runner, + ) + + with self.assertRaisesRegex( + PaperCitationValidationError, + "fewer citations", ): - await paper_flow(RawText(text="x"), memory=object()) - # The runner sees the memory placeholder forwarded. - self.assertIsNotNone(runner.await_args.kwargs["memory"]) - - async def test_extra_run_hooks_forwarded(self) -> None: - class _H(RunHooks[Any]): - pass - - hook = _H() - with ( - patch( - "quantmind.flows.paper.Agent", - return_value=MagicMock(), - ), - _patch_runner(_stub_paper()) as runner, + _build_global_summary( + result.source_revision, + result.chunk_set, + draft, + PaperFlowCfg(), + ) + + +class SummaryBudgetTests(unittest.IsolatedAsyncioTestCase): + async def test_tool_call_and_input_budgets_are_enforced(self) -> None: + result = build_paper_result() + cfg = PaperFlowCfg( + max_summary_tool_calls=1, + max_summary_input_tokens=10_000, + ) + budget = _SummaryBudget(cfg, "manifest") + + await budget.read(result.chunk_set, start=0, count=1) + with self.assertRaisesRegex( + PaperSummaryBudgetExceeded, + "tool_calls", ): - await paper_flow(RawText(text="x"), extra_run_hooks=[hook]) - self.assertEqual(runner.await_args.kwargs["extra_run_hooks"], [hook]) + await budget.read(result.chunk_set, start=1, count=1) - async def test_model_settings_forwarded_when_set(self) -> None: - seen: dict[str, Any] = {} + async def test_concurrent_reads_remain_bounded_and_valid(self) -> None: + result = build_paper_result() + cfg = PaperFlowCfg( + max_summary_tool_calls=3, + max_summary_concurrency=1, + max_summary_input_tokens=10_000, + ) + budget = _SummaryBudget(cfg, "manifest") - def _capture_agent(*_a: Any, **kwargs: Any) -> Any: - seen.update(kwargs) - return MagicMock() + values = await asyncio.gather( + budget.read(result.chunk_set, start=0, count=1), + budget.read(result.chunk_set, start=1, count=1), + budget.read(result.chunk_set, start=2, count=1), + ) + self.assertEqual(len(values), 3) - from agents import ModelSettings + def test_model_output_tokens_are_capped(self) -> None: + cfg = PaperFlowCfg( + max_summary_output_tokens=256, + model_settings=ModelSettings(max_tokens=1024), + ) - ms = ModelSettings(temperature=0.42) - cfg = PaperFlowCfg(model_settings=ms) - with ( - patch("quantmind.flows.paper.Agent", side_effect=_capture_agent), - _patch_runner(_stub_paper()), - ): - await paper_flow(RawText(text="x"), cfg=cfg) - self.assertIs(seen["model_settings"], ms) + settings = _bounded_model_settings(cfg) - async def test_model_settings_omitted_when_none(self) -> None: - seen: dict[str, Any] = {} + self.assertEqual(settings.max_tokens, 256) + self.assertTrue(settings.parallel_tool_calls) - def _capture_agent(*_a: Any, **kwargs: Any) -> Any: - seen.update(kwargs) - return MagicMock() - with ( - patch("quantmind.flows.paper.Agent", side_effect=_capture_agent), - _patch_runner(_stub_paper()), - ): - await paper_flow(RawText(text="x")) - self.assertNotIn("model_settings", seen) +if __name__ == "__main__": + unittest.main() diff --git a/tests/knowledge/test_base.py b/tests/knowledge/test_base.py index f1fc840..d442922 100644 --- a/tests/knowledge/test_base.py +++ b/tests/knowledge/test_base.py @@ -75,14 +75,11 @@ def test_minimal(self): class _ConcreteKnowledge(BaseKnowledge): - """Test fixture: concrete subclass that overrides embedding_text.""" + """Test fixture for shared canonical metadata behavior.""" item_type: str = "test" # pyright: ignore[reportIncompatibleVariableOverride] payload: str = "" - def embedding_text(self) -> str: - return self.payload - class BaseKnowledgeTests(unittest.TestCase): def test_as_of_required(self): @@ -139,19 +136,9 @@ def test_extra_forbidden(self): unexpected_field=1, # type: ignore[call-arg] ) - def test_embedding_text_default_raises(self): - # BaseKnowledge.embedding_text raises NotImplementedError; subclasses - # must override. We test via a class that doesn't override. - class _NoEmbed(BaseKnowledge): - item_type: str = "no_embed" # pyright: ignore[reportIncompatibleVariableOverride] - - item = _NoEmbed(as_of=_now(), source=_src()) - with self.assertRaises(NotImplementedError): - item.embedding_text() - - def test_embedding_text_override(self): + def test_embedding_selection_is_not_canonical_domain_behavior(self): item = _ConcreteKnowledge(as_of=_now(), source=_src(), payload="hello") - self.assertEqual(item.embedding_text(), "hello") + self.assertFalse(hasattr(item, "embedding_text")) def test_extraction_optional(self): item = _ConcreteKnowledge(as_of=_now(), source=_src()) @@ -211,9 +198,12 @@ def test_top_level_imports(self): Factor, FlattenKnowledge, GraphKnowledge, + LegacyPaper, News, - Paper, - PaperKnowledgeCard, + PaperChunkSet, + PaperFlowResult, + PaperGlobalSummary, + PaperSourceRevision, SourceRef, Thesis, TreeKnowledge, @@ -225,10 +215,13 @@ def test_top_level_imports(self): self.assertTrue(issubclass(GraphKnowledge, BaseKnowledge)) self.assertTrue(issubclass(News, FlattenKnowledge)) self.assertTrue(issubclass(Earnings, FlattenKnowledge)) - self.assertTrue(issubclass(PaperKnowledgeCard, FlattenKnowledge)) self.assertTrue(issubclass(Factor, FlattenKnowledge)) self.assertTrue(issubclass(Thesis, FlattenKnowledge)) - self.assertTrue(issubclass(Paper, TreeKnowledge)) + self.assertTrue(issubclass(LegacyPaper, TreeKnowledge)) + self.assertEqual(PaperSourceRevision.__name__, "PaperSourceRevision") + self.assertEqual(PaperChunkSet.__name__, "PaperChunkSet") + self.assertEqual(PaperGlobalSummary.__name__, "PaperGlobalSummary") + self.assertEqual(PaperFlowResult.__name__, "PaperFlowResult") # Ensure side-imports are real classes self.assertEqual(Citation.__name__, "Citation") self.assertEqual(SourceRef.__name__, "SourceRef") diff --git a/tests/knowledge/test_earnings.py b/tests/knowledge/test_earnings.py index 3025365..271346f 100644 --- a/tests/knowledge/test_earnings.py +++ b/tests/knowledge/test_earnings.py @@ -42,7 +42,7 @@ def test_full(self): self.assertEqual(e.revenue, 120.0) self.assertEqual(e.surprise_flags, ["eps_beat", "revenue_beat"]) - def test_embedding_text(self): + def test_retrieval_projection_is_not_a_domain_method(self): e = Earnings( as_of=_now(), source=_src(), @@ -50,10 +50,7 @@ def test_embedding_text(self): period="2026Q1", guidance="Raised FY guide", ) - self.assertEqual( - e.embedding_text(), - "AAPL 2026Q1 earnings\nRaised FY guide", - ) + self.assertFalse(hasattr(e, "embedding_text")) if __name__ == "__main__": diff --git a/tests/knowledge/test_factor.py b/tests/knowledge/test_factor.py index 11d620e..d123e75 100644 --- a/tests/knowledge/test_factor.py +++ b/tests/knowledge/test_factor.py @@ -22,18 +22,14 @@ def test_minimal(self): self.assertEqual(f.factor_name, "momentum_12_1") self.assertIsNone(f.universe) - def test_embedding_text(self): + def test_retrieval_projection_is_not_a_domain_method(self): f = Factor( as_of=_now(), source=_src(), factor_name="value", universe="us_equities", ) - self.assertEqual(f.embedding_text(), "factor value on us_equities") - - def test_embedding_text_default_universe(self): - f = Factor(as_of=_now(), source=_src(), factor_name="size") - self.assertEqual(f.embedding_text(), "factor size on unspecified") + self.assertFalse(hasattr(f, "embedding_text")) if __name__ == "__main__": diff --git a/tests/knowledge/test_news.py b/tests/knowledge/test_news.py index c52283c..84a65a6 100644 --- a/tests/knowledge/test_news.py +++ b/tests/knowledge/test_news.py @@ -42,7 +42,7 @@ def test_sentiment_enum(self): sentiment="ecstatic", # type: ignore[arg-type] ) - def test_embedding_text(self): + def test_retrieval_projection_is_not_a_domain_method(self): n = News( as_of=_now(), source=_src(), @@ -51,10 +51,7 @@ def test_embedding_text(self): timestamp=_now(), entities=["FOMC", "USD"], ) - self.assertEqual( - n.embedding_text(), - "Fed holds rates\nmonetary_policy\nFOMC, USD", - ) + self.assertFalse(hasattr(n, "embedding_text")) if __name__ == "__main__": diff --git a/tests/knowledge/test_paper.py b/tests/knowledge/test_paper.py index 30b2a86..5ab4312 100644 --- a/tests/knowledge/test_paper.py +++ b/tests/knowledge/test_paper.py @@ -1,109 +1,172 @@ -"""Tests for knowledge.paper — Paper (Tree) + PaperKnowledgeCard (Flatten).""" +"""Tests for source-first paper revisions and independent artifacts.""" import unittest -from datetime import datetime, timezone from uuid import uuid4 -from quantmind.knowledge._base import SourceRef -from quantmind.knowledge._tree import TreeNode -from quantmind.knowledge.paper import Paper, PaperKnowledgeCard - - -def _now() -> datetime: - return datetime(2026, 4, 1, tzinfo=timezone.utc) - - -def _src() -> SourceRef: - return SourceRef(kind="arxiv", uri="arxiv:2604.12345") - - -def _single_node_paper(**overrides) -> Paper: - root_id = uuid4() - root = TreeNode( - node_id=root_id, - parent_id=None, - position=0, - title="On Cross-Sectional Momentum", - summary="Methodology for momentum factor on US equities.", - ) - return Paper( - as_of=_now(), - source=_src(), - root_node_id=root_id, - nodes={root_id: root}, - **overrides, - ) - - -class PaperTreeTests(unittest.TestCase): - def test_minimal(self): - p = _single_node_paper() - self.assertEqual(p.item_type, "paper") - self.assertIsNone(p.arxiv_id) - self.assertEqual(p.authors, []) - self.assertEqual(p.asset_classes, []) - self.assertEqual(p.root().title, "On Cross-Sectional Momentum") - - def test_metadata(self): - p = _single_node_paper( - arxiv_id="2604.12345", - authors=["A. Smith", "B. Jones"], - asset_classes=["equities"], +from pydantic import ValidationError + +from quantmind.knowledge import ( + PaperCitation, + PaperFlowResult, + PaperSourceRevision, + PaperSourceSpan, +) +from quantmind.knowledge.paper import ( + _paper_chunk_id, + _paper_chunk_set_content_hash, + _paper_summary_content_hash, +) +from tests.paper_helpers import build_paper_result + + +class PaperArtifactTests(unittest.TestCase): + def test_result_has_one_shared_source_and_independent_artifacts( + self, + ) -> None: + result = build_paper_result() + + self.assertEqual( + result.chunk_set.source_revision_id, + result.source_revision.id, ) - self.assertEqual(p.arxiv_id, "2604.12345") - self.assertEqual(p.authors, ["A. Smith", "B. Jones"]) - self.assertEqual(p.asset_classes, ["equities"]) + self.assertEqual( + result.global_summary.source_revision_id, + result.source_revision.id, + ) + self.assertNotEqual(result.chunk_set.id, result.global_summary.id) + self.assertEqual(len(result.chunk_set.chunks), 3) + self.assertEqual(len(result.global_summary.citations), 3) + + def test_unchanged_producer_configs_create_stable_ids(self) -> None: + first = build_paper_result() + second = build_paper_result() - def test_embedding_text_uses_root(self): - p = _single_node_paper() + self.assertEqual(first.source_revision.id, second.source_revision.id) + self.assertEqual(first.chunk_set.id, second.chunk_set.id) + self.assertEqual(first.global_summary.id, second.global_summary.id) self.assertEqual( - p.embedding_text(), - "On Cross-Sectional Momentum\n" - "Methodology for momentum factor on US equities.", + [chunk.chunk_id for chunk in first.chunk_set.chunks], + [chunk.chunk_id for chunk in second.chunk_set.chunks], ) + def test_splitter_and_summary_configs_version_independently(self) -> None: + original = build_paper_result() + new_chunks = build_paper_result(chunk_size=256) + new_summary = build_paper_result(summary_model="another-model") -class PaperKnowledgeCardTests(unittest.TestCase): - def test_minimal(self): - paper_id = uuid4() - card = PaperKnowledgeCard( - as_of=_now(), - source=_src(), - paper_id=paper_id, - summary="A momentum study on US equities.", + self.assertEqual( + original.source_revision.id, new_chunks.source_revision.id ) - self.assertEqual(card.item_type, "paper_card") - self.assertEqual(card.paper_id, paper_id) - self.assertEqual(card.key_findings, []) - self.assertEqual(card.asset_classes, []) - - def test_full(self): - card = PaperKnowledgeCard( - as_of=_now(), - source=_src(), - paper_id=uuid4(), - summary="s", - methodology="m", - key_findings=["f1", "f2"], - limitations=["l1"], - asset_classes=["equities"], + self.assertNotEqual(original.chunk_set.id, new_chunks.chunk_set.id) + self.assertNotEqual( + original.global_summary.id, new_chunks.global_summary.id ) - self.assertEqual(card.methodology, "m") - self.assertEqual(card.key_findings, ["f1", "f2"]) - - def test_embedding_text(self): - card = PaperKnowledgeCard( - as_of=_now(), - source=_src(), - paper_id=uuid4(), - summary="momentum study", - key_findings=["beats SPX", "robust to costs"], + self.assertEqual(original.chunk_set.id, new_summary.chunk_set.id) + self.assertNotEqual( + original.global_summary.id, new_summary.global_summary.id ) - self.assertEqual( - card.embedding_text(), - "momentum study\nbeats SPX robust to costs", + + def test_canonical_models_have_no_embedding_vectors_or_text_method( + self, + ) -> None: + result = build_paper_result() + + for value in ( + result.source_revision, + result.chunk_set, + result.global_summary, + *result.chunk_set.chunks, + ): + self.assertNotIn("embedding", value.model_dump()) + self.assertFalse(hasattr(value, "embedding_text")) + + def test_source_json_excludes_blobs_but_preserves_manifest(self) -> None: + source = build_paper_result().source_revision + revived = PaperSourceRevision.model_validate_json( + source.model_dump_json() ) + self.assertEqual(revived.id, source.id) + self.assertEqual(revived.parsed, source.parsed) + self.assertEqual(revived.blobs, {}) + with self.assertRaisesRegex(RuntimeError, "bytes are not loaded"): + revived.blob_for(revived.raw_asset_id) + + def test_result_rejects_unknown_model_owned_citation_identity(self) -> None: + result = build_paper_result() + citation = result.global_summary.citations[0] + invalid = PaperCitation( + chunk_set_id=result.chunk_set.id, + chunk_id=uuid4(), + page_number=citation.page_number, + ) + invalid_citations = ( + invalid, + *result.global_summary.citations[1:], + ) + summary = result.global_summary.model_copy( + update={ + "citations": invalid_citations, + "content_hash": _paper_summary_content_hash( + result.global_summary.summary, + invalid_citations, + ), + } + ) + + with self.assertRaisesRegex(ValidationError, "unknown chunk"): + PaperFlowResult( + source_revision=result.source_revision, + chunk_set=result.chunk_set, + global_summary=summary, + ) + + def test_source_revision_rejects_content_owned_id_override(self) -> None: + source = build_paper_result().source_revision + payload = source.model_dump() + payload["id"] = uuid4() + payload["blobs"] = source.blobs + + with self.assertRaisesRegex(ValidationError, "ID does not match"): + PaperSourceRevision.model_validate(payload) + + def test_result_rejects_chunk_spans_outside_source_manifest(self) -> None: + result = build_paper_result() + chunk = result.chunk_set.chunks[0] + invalid_span = PaperSourceSpan( + page_number=1, + start_char=0, + end_char=len(result.source_revision.parsed.pages[0].text) + 1, + ) + invalid_chunk = chunk.model_copy( + update={ + "chunk_id": _paper_chunk_id( + result.chunk_set.id, + position=chunk.position, + content_hash=chunk.content_hash, + spans=(invalid_span,), + ), + "source_spans": (invalid_span,), + } + ) + invalid_chunks = ( + invalid_chunk, + *result.chunk_set.chunks[1:], + ) + invalid_chunk_set = result.chunk_set.model_copy( + update={ + "chunks": invalid_chunks, + "content_hash": _paper_chunk_set_content_hash(invalid_chunks), + } + ) + + with self.assertRaisesRegex(ValidationError, "exceeds its source page"): + PaperFlowResult( + source_revision=result.source_revision, + chunk_set=invalid_chunk_set, + global_summary=result.global_summary, + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/knowledge/test_roundtrip.py b/tests/knowledge/test_roundtrip.py index 3d0433d..6548ec8 100644 --- a/tests/knowledge/test_roundtrip.py +++ b/tests/knowledge/test_roundtrip.py @@ -14,13 +14,15 @@ Earnings, ExtractionRef, Factor, + LegacyPaper, News, - Paper, - PaperKnowledgeCard, + PaperChunkSet, + PaperGlobalSummary, SourceRef, Thesis, TreeNode, ) +from tests.paper_helpers import build_paper_result def _now() -> datetime: @@ -62,19 +64,6 @@ def test_earnings(self): revived = Earnings.model_validate_json(e.model_dump_json()) self.assertEqual(e, revived) - def test_paper_knowledge_card(self): - card = PaperKnowledgeCard( - as_of=_now(), - source=_src("arxiv"), - paper_id=uuid4(), - summary="momentum study", - methodology="cross-sectional", - key_findings=["beats SPX"], - asset_classes=["equities"], - ) - revived = PaperKnowledgeCard.model_validate_json(card.model_dump_json()) - self.assertEqual(card, revived) - def test_factor(self): f = Factor( as_of=_now(), @@ -92,7 +81,7 @@ def test_thesis(self): class TreeRoundTripTests(unittest.TestCase): - def _build_paper(self) -> Paper: + def _build_paper(self) -> LegacyPaper: leaf_id = uuid4() root_id = uuid4() leaf = TreeNode( @@ -111,7 +100,7 @@ def _build_paper(self) -> Paper: summary="paper-level summary", children_ids=[leaf_id], ) - return Paper( + return LegacyPaper( as_of=_now(), source=_src("arxiv"), extraction=_ext(), @@ -124,7 +113,7 @@ def _build_paper(self) -> Paper: def test_paper_dict_uuid_keys_round_trip(self): p = self._build_paper() - revived = Paper.model_validate_json(p.model_dump_json()) + revived = LegacyPaper.model_validate_json(p.model_dump_json()) self.assertEqual(p, revived) # UUID keys are preserved through the JSON detour. self.assertEqual(set(revived.nodes.keys()), set(p.nodes.keys())) @@ -133,10 +122,25 @@ def test_paper_dict_uuid_keys_round_trip(self): def test_paper_walk_dfs_after_round_trip(self): p = self._build_paper() - revived = Paper.model_validate_json(p.model_dump_json()) + revived = LegacyPaper.model_validate_json(p.model_dump_json()) titles = [n.title for n in revived.walk_dfs()] self.assertEqual(titles, ["Cross-Sectional Momentum", "Methodology"]) +class PaperArtifactRoundTripTests(unittest.TestCase): + def test_chunk_set_and_summary_round_trip_independently(self) -> None: + result = build_paper_result() + + chunk_set = PaperChunkSet.model_validate_json( + result.chunk_set.model_dump_json() + ) + summary = PaperGlobalSummary.model_validate_json( + result.global_summary.model_dump_json() + ) + + self.assertEqual(chunk_set, result.chunk_set) + self.assertEqual(summary, result.global_summary) + + if __name__ == "__main__": unittest.main() diff --git a/tests/knowledge/test_thesis.py b/tests/knowledge/test_thesis.py index c231bfc..d539011 100644 --- a/tests/knowledge/test_thesis.py +++ b/tests/knowledge/test_thesis.py @@ -21,9 +21,9 @@ def test_minimal(self): self.assertEqual(t.item_type, "thesis") self.assertEqual(t.claim, "USD weakens in H2 2026") - def test_embedding_text_returns_claim(self): + def test_retrieval_projection_is_not_a_domain_method(self): t = Thesis(as_of=_now(), source=_src(), claim="rates higher for longer") - self.assertEqual(t.embedding_text(), "rates higher for longer") + self.assertFalse(hasattr(t, "embedding_text")) if __name__ == "__main__": diff --git a/tests/knowledge/test_tree.py b/tests/knowledge/test_tree.py index 825b04b..d78fa56 100644 --- a/tests/knowledge/test_tree.py +++ b/tests/knowledge/test_tree.py @@ -104,9 +104,9 @@ def test_frozen(self): with self.assertRaises(ValidationError): n.title = "z" # type: ignore[misc] - def test_default_embedding_text(self): + def test_retrieval_projection_is_not_a_node_method(self): n = TreeNode(title="Methodology", summary="We use X to do Y.") - self.assertEqual(n.embedding_text(), "Methodology\nWe use X to do Y.") + self.assertFalse(hasattr(n, "embedding_text")) class TreeKnowledgeTests(unittest.TestCase): @@ -143,9 +143,9 @@ def test_find_path_unknown(self): tree = _make_tree() self.assertEqual(tree.find_path(uuid4()), []) - def test_embedding_text_uses_root(self): + def test_retrieval_projection_is_not_a_tree_method(self): tree = _make_tree() - self.assertEqual(tree.embedding_text(), "Root\nroot summary") + self.assertFalse(hasattr(tree, "embedding_text")) if __name__ == "__main__": diff --git a/tests/library/test_example_bundle.py b/tests/library/test_example_bundle.py index a6d0b8c..2d92efd 100644 --- a/tests/library/test_example_bundle.py +++ b/tests/library/test_example_bundle.py @@ -9,8 +9,8 @@ from quantmind.knowledge import ( BaseKnowledge, Earnings, + LegacyPaper, News, - Paper, TreeKnowledge, ) from quantmind.library import LocalKnowledgeLibrary, SemanticQuery @@ -26,7 +26,7 @@ _KNOWLEDGE_TYPES: dict[str, type[BaseKnowledge]] = { "earnings": Earnings, "news": News, - "paper": Paper, + "paper": LegacyPaper, } @@ -65,7 +65,7 @@ def test_bundle_is_a_concrete_cross_shape_scenario(self): self.assertIn("AI", self.scenario) self.assertEqual( {type(item) for item in self.items}, - {News, Earnings, Paper}, + {News, Earnings, LegacyPaper}, ) self.assertTrue( all("ai-infrastructure" in item.tags for item in self.items) @@ -83,7 +83,9 @@ def test_sources_and_financial_times_are_auditable(self): self.assertTrue(item.citations) def test_paper_tree_has_stable_valid_navigation(self): - paper = next(item for item in self.items if isinstance(item, Paper)) + paper = next( + item for item in self.items if isinstance(item, LegacyPaper) + ) self.assertEqual(len(paper.nodes), 4) self.assertEqual( [node.position for node in paper.children_of(paper.root_node_id)], diff --git a/tests/library/test_local.py b/tests/library/test_local.py index 6e1785b..a0a506c 100644 --- a/tests/library/test_local.py +++ b/tests/library/test_local.py @@ -10,8 +10,8 @@ from quantmind.knowledge import ( Citation, FlattenKnowledge, + LegacyPaper, News, - Paper, SourceRef, TreeNode, ) @@ -63,9 +63,6 @@ async def close(self) -> None: class _UnsupportedKnowledge(FlattenKnowledge): item_type: str = "unsupported" - def embedding_text(self) -> str: - return "unsupported canonical knowledge" - def _time(day: int, hour: int = 0) -> datetime: return datetime(2026, 7, day, hour, tzinfo=timezone.utc) @@ -102,7 +99,15 @@ def _news( ) -def _paper() -> tuple[Paper, UUID, UUID]: +def _news_projection(item: News) -> str: + return f"{item.headline}\n{item.event_type}\n{', '.join(item.entities)}".strip() + + +def _node_projection(node: TreeNode) -> str: + return f"{node.title}\n{node.summary}" + + +def _paper() -> tuple[LegacyPaper, UUID, UUID]: root_id = uuid4() methods_id = uuid4() results_id = uuid4() @@ -127,7 +132,7 @@ def _paper() -> tuple[Paper, UUID, UUID]: title="Results", summary="Capital expenditure is expected to rise", ) - paper = Paper( + paper = LegacyPaper( as_of=_time(1), available_at=_time(2), source=SourceRef(kind="arxiv", content_hash="paper-v1"), @@ -160,8 +165,8 @@ async def test_flat_put_get_and_deterministic_best_first_search(self): query_text = "management increases capital expenditure" provider = _FakeEmbeddingProvider( vectors={ - alpha.embedding_text(): [1.0, 0.0], - beta.embedding_text(): [0.0, 1.0], + _news_projection(alpha): [1.0, 0.0], + _news_projection(beta): [0.0, 1.0], query_text: [1.0, 0.0], } ) @@ -178,7 +183,10 @@ async def test_flat_put_get_and_deterministic_best_first_search(self): self.assertEqual([hit.item_id for hit in hits], [alpha.id, beta.id]) self.assertAlmostEqual(hits[0].score, 1.0) self.assertIsNone(hits[0].node_id) - self.assertEqual(hits[0].matched_text, alpha.embedding_text()) + self.assertEqual(hits[0].matched_text, _news_projection(alpha)) + self.assertEqual(hits[0].locator.artifact_id, alpha.id) + self.assertIsNone(hits[0].locator.source_revision_id) + self.assertEqual(hits[0].projection.model, "fake-2d") self.assertEqual(hits[0].source, alpha.source) self.assertEqual(hits[0].citations, alpha.citations) self.assertEqual(hits[0].available_at, alpha.available_at) @@ -215,10 +223,10 @@ async def test_tree_root_and_non_root_nodes_use_exact_grain_and_filters( query_text = "capital expenditure outlook" provider = _FakeEmbeddingProvider( vectors={ - paper.embedding_text(): [1.0, 0.0], - paper.nodes[methods_id].embedding_text(): [0.8, 0.2], - paper.nodes[results_id].embedding_text(): [0.9, 0.1], - unrelated.embedding_text(): [0.0, 1.0], + _node_projection(paper.root()): [1.0, 0.0], + _node_projection(paper.nodes[methods_id]): [0.8, 0.2], + _node_projection(paper.nodes[results_id]): [0.9, 0.1], + _news_projection(unrelated): [0.0, 1.0], query_text: [1.0, 0.0], } ) @@ -270,20 +278,23 @@ async def test_tree_root_and_non_root_nodes_use_exact_grain_and_filters( {methods_id, results_id}, ) root_hit = next(hit for hit in hits if hit.node_id is None) - self.assertEqual(root_hit.matched_text, paper.embedding_text()) + self.assertEqual( + root_hit.matched_text, + _node_projection(paper.root()), + ) self.assertEqual(root_hit.citations, paper.root().citations) methods_hit = next(hit for hit in hits if hit.node_id == methods_id) self.assertEqual( methods_hit.matched_text, - paper.nodes[methods_id].embedding_text(), + _node_projection(paper.nodes[methods_id]), ) self.assertEqual( methods_hit.citations, paper.nodes[methods_id].citations, ) stored = await library.get(methods_hit.item_id) - self.assertIsInstance(stored, Paper) - assert isinstance(stored, Paper) + self.assertIsInstance(stored, LegacyPaper) + assert isinstance(stored, LegacyPaper) self.assertEqual( stored.find_path(methods_id)[-1].node_id, methods_id ) @@ -372,7 +383,7 @@ async def test_reput_is_idempotent_and_invalidates_only_changed_metadata( ) await library.put(changed_text) self.assertEqual( - provider.calls[-1][2], (changed_text.embedding_text(),) + provider.calls[-1][2], (_news_projection(changed_text),) ) changed_source = changed_text.model_copy( @@ -413,7 +424,7 @@ async def test_tree_projection_change_invalidates_only_one_node(self): changed_paper = paper.model_copy(update={"nodes": changed_nodes}) await library.put(changed_paper) self.assertEqual( - provider.calls[-1][2], (changed_node.embedding_text(),) + provider.calls[-1][2], (_node_projection(changed_node),) ) schema_change = changed_paper.model_copy( diff --git a/tests/library/test_paper.py b/tests/library/test_paper.py new file mode 100644 index 0000000..4c0441a --- /dev/null +++ b/tests/library/test_paper.py @@ -0,0 +1,408 @@ +"""Offline source/artifact/projection tests for paper library persistence.""" + +import hashlib +import sqlite3 +import tempfile +import unittest +from collections.abc import Sequence +from pathlib import Path + +from quantmind.knowledge import PaperChunk, PaperGlobalSummary +from quantmind.library import LocalKnowledgeLibrary, SemanticQuery +from tests.paper_helpers import build_paper_result + + +class _FakeEmbeddingProvider: + def __init__(self, vectors: dict[str, list[float]] | None = None) -> None: + self.vectors = vectors or {} + self.calls: list[tuple[str, ...]] = [] + + async def embed( + self, + texts: Sequence[str], + *, + model: str, + dimensions: int | None, + ) -> list[list[float]]: + del model + self.calls.append(tuple(texts)) + size = dimensions or 2 + return [ + self.vectors.get( + text, + [ + float((sum(map(ord, text)) + offset) % 17 + 1) + for offset in range(size) + ], + ) + for text in texts + ] + + async def close(self) -> None: + """Release no resources.""" + + +class _FailingEmbeddingProvider(_FakeEmbeddingProvider): + async def embed( + self, + texts: Sequence[str], + *, + model: str, + dimensions: int | None, + ) -> list[list[float]]: + del texts, model, dimensions + raise RuntimeError("embedding unavailable") + + +class PaperLibraryTests(unittest.IsolatedAsyncioTestCase): + def setUp(self) -> None: + self._temporary_directory = tempfile.TemporaryDirectory() + self.db_path = Path(self._temporary_directory.name) / "paper.db" + + def tearDown(self) -> None: + self._temporary_directory.cleanup() + + async def test_put_persists_explicit_source_artifact_and_projection_layers( + self, + ) -> None: + result = build_paper_result() + provider = _FakeEmbeddingProvider() + library = await LocalKnowledgeLibrary.open( + self.db_path, + embedding_model="fake-2d", + embedding_dimensions=2, + _embedding_provider=provider, + ) + try: + await library.put_paper(result) + finally: + await library.close() + + self.assertEqual(len(provider.calls), 1) + self.assertEqual(len(provider.calls[0]), 4) + with sqlite3.connect(self.db_path) as db: + self.assertEqual(db.execute("PRAGMA user_version").fetchone()[0], 3) + self.assertEqual( + db.execute("SELECT COUNT(*) FROM paper_sources").fetchone()[0], + 1, + ) + self.assertEqual( + db.execute( + "SELECT COUNT(*) FROM paper_source_assets" + ).fetchone()[0], + 1, + ) + self.assertEqual( + db.execute("SELECT COUNT(*) FROM paper_artifacts").fetchone()[ + 0 + ], + 2, + ) + self.assertEqual( + db.execute( + "SELECT COUNT(*) FROM paper_artifact_members" + ).fetchone()[0], + 3, + ) + self.assertEqual( + db.execute( + "SELECT COUNT(*) FROM paper_artifact_lineage" + ).fetchone()[0], + 1, + ) + self.assertEqual( + db.execute("SELECT COUNT(*) FROM paper_projections").fetchone()[ + 0 + ], + 4, + ) + payloads = [ + row[0] + for row in db.execute( + "SELECT payload_json FROM paper_artifacts" + ).fetchall() + ] + self.assertTrue( + all("embedding" not in payload for payload in payloads) + ) + + async def test_reopen_round_trip_reuses_vectors_and_resolves_hits( + self, + ) -> None: + result = build_paper_result() + multi_head = result.chunk_set.chunks[1] + summary_query = "What is the paper's central contribution?" + chunk_query = "How does multi-head attention work?" + first = _FakeEmbeddingProvider( + vectors={ + result.global_summary.summary: [1.0, 0.0], + multi_head.text: [1.0, 0.0], + result.chunk_set.chunks[0].text: [0.0, 1.0], + result.chunk_set.chunks[2].text: [0.0, 1.0], + } + ) + library = await LocalKnowledgeLibrary.open( + self.db_path, + embedding_model="fake-2d", + embedding_dimensions=2, + _embedding_provider=first, + ) + await library.put_paper(result) + await library.close() + + second = _FakeEmbeddingProvider( + vectors={ + result.global_summary.summary: [1.0, 0.0], + multi_head.text: [1.0, 0.0], + summary_query: [1.0, 0.0], + chunk_query: [1.0, 0.0], + } + ) + library = await LocalKnowledgeLibrary.open( + self.db_path, + embedding_model="fake-2d", + embedding_dimensions=2, + _embedding_provider=second, + ) + try: + restored = await library.get_paper(result.source_revision.id) + self.assertEqual(restored.chunk_set, result.chunk_set) + self.assertEqual(restored.global_summary, result.global_summary) + self.assertEqual( + restored.source_revision.blob_for( + restored.source_revision.raw_asset_id + ), + result.source_revision.blob_for( + result.source_revision.raw_asset_id + ), + ) + self.assertEqual(second.calls, []) + + summary_hits = await library.search( + SemanticQuery( + text=summary_query, + artifact_kinds=["paper_summary"], + top_k=3, + ) + ) + self.assertEqual(len(summary_hits), 1) + self.assertEqual( + summary_hits[0].locator.artifact_id, + result.global_summary.id, + ) + self.assertEqual(summary_hits[0].projection.model, "fake-2d") + summary = await library.resolve(summary_hits[0].locator) + self.assertIsInstance(summary, PaperGlobalSummary) + + chunk_hits = await library.search( + SemanticQuery( + text=chunk_query, + artifact_kinds=["paper_chunk_set"], + top_k=5, + ) + ) + self.assertIn( + multi_head.chunk_id, [hit.node_id for hit in chunk_hits] + ) + matching_hit = next( + hit for hit in chunk_hits if hit.node_id == multi_head.chunk_id + ) + chunk = await library.resolve(matching_hit.locator) + self.assertIsInstance(chunk, PaperChunk) + assert isinstance(chunk, PaperChunk) + self.assertEqual(chunk.source_spans[0].page_number, 2) + self.assertEqual(matching_hit.citations[0].page, 2) + finally: + await library.close() + + async def test_reput_and_changed_summary_selectively_rebuild_projections( + self, + ) -> None: + original = build_paper_result() + changed = build_paper_result(summary_text="A refreshed cited summary.") + provider = _FakeEmbeddingProvider() + library = await LocalKnowledgeLibrary.open( + self.db_path, + embedding_model="fake-2d", + embedding_dimensions=2, + _embedding_provider=provider, + ) + try: + await library.put_paper(original) + self.assertEqual(len(provider.calls), 1) + await library.put_paper(original) + self.assertEqual(len(provider.calls), 1) + + await library.put_paper(changed) + self.assertEqual( + provider.calls[-1], (changed.global_summary.summary,) + ) + self.assertEqual(len(provider.calls), 2) + finally: + await library.close() + + async def test_multiple_chunk_and_summary_versions_coexist(self) -> None: + first = build_paper_result(chunk_size=128) + second = build_paper_result(chunk_size=256) + provider = _FakeEmbeddingProvider() + library = await LocalKnowledgeLibrary.open( + self.db_path, + embedding_model="fake-2d", + embedding_dimensions=2, + _embedding_provider=provider, + ) + try: + await library.put_paper(first) + await library.put_paper(second) + with self.assertRaisesRegex(ValueError, "specify an artifact ID"): + await library.get_paper(first.source_revision.id) + restored = await library.get_paper( + first.source_revision.id, + chunk_set_id=second.chunk_set.id, + summary_id=second.global_summary.id, + ) + self.assertEqual(restored.chunk_set.id, second.chunk_set.id) + self.assertEqual( + restored.global_summary.id, second.global_summary.id + ) + finally: + await library.close() + with sqlite3.connect(self.db_path) as db: + self.assertEqual( + db.execute("SELECT COUNT(*) FROM paper_artifacts").fetchone()[ + 0 + ], + 4, + ) + self.assertEqual( + db.execute("SELECT COUNT(*) FROM paper_projections").fetchone()[ + 0 + ], + 8, + ) + + async def test_required_projection_failure_is_atomic(self) -> None: + library = await LocalKnowledgeLibrary.open( + self.db_path, + embedding_model="failed", + embedding_dimensions=2, + _embedding_provider=_FailingEmbeddingProvider(), + ) + try: + with self.assertRaisesRegex(RuntimeError, "embedding unavailable"): + await library.put_paper(build_paper_result()) + finally: + await library.close() + with sqlite3.connect(self.db_path) as db: + self.assertEqual( + db.execute("SELECT COUNT(*) FROM paper_sources").fetchone()[0], + 0, + ) + self.assertEqual( + db.execute("SELECT COUNT(*) FROM paper_artifacts").fetchone()[ + 0 + ], + 0, + ) + self.assertEqual( + db.execute("SELECT COUNT(*) FROM paper_projections").fetchone()[ + 0 + ], + 0, + ) + + async def test_rehydrate_rejects_asset_metadata_drift(self) -> None: + result = build_paper_result() + library = await LocalKnowledgeLibrary.open( + self.db_path, + embedding_model="fake-2d", + embedding_dimensions=2, + _embedding_provider=_FakeEmbeddingProvider(), + ) + await library.put_paper(result) + await library.close() + + with sqlite3.connect(self.db_path) as db: + db.execute( + "UPDATE paper_source_assets SET media_type = ?", + ("application/tampered",), + ) + + library = await LocalKnowledgeLibrary.open( + self.db_path, + embedding_model="fake-2d", + embedding_dimensions=2, + _embedding_provider=_FakeEmbeddingProvider(), + ) + try: + with self.assertRaisesRegex(RuntimeError, "metadata mismatch"): + await library.get_paper(result.source_revision.id) + finally: + await library.close() + + async def test_rehydrate_rejects_missing_summary_lineage(self) -> None: + result = build_paper_result() + library = await LocalKnowledgeLibrary.open( + self.db_path, + embedding_model="fake-2d", + embedding_dimensions=2, + _embedding_provider=_FakeEmbeddingProvider(), + ) + await library.put_paper(result) + await library.close() + + with sqlite3.connect(self.db_path) as db: + db.execute( + "DELETE FROM paper_artifact_lineage WHERE artifact_id = ?", + (str(result.global_summary.id),), + ) + + library = await LocalKnowledgeLibrary.open( + self.db_path, + embedding_model="fake-2d", + embedding_dimensions=2, + _embedding_provider=_FakeEmbeddingProvider(), + ) + try: + with self.assertRaisesRegex(RuntimeError, "lineage mismatch"): + await library.get_artifact(result.global_summary.id) + finally: + await library.close() + + async def test_search_rejects_projection_text_drift(self) -> None: + result = build_paper_result() + library = await LocalKnowledgeLibrary.open( + self.db_path, + embedding_model="fake-2d", + embedding_dimensions=2, + _embedding_provider=_FakeEmbeddingProvider(), + ) + await library.put_paper(result) + await library.close() + + tampered = "tampered summary projection" + with sqlite3.connect(self.db_path) as db: + db.execute( + """ + UPDATE paper_projections + SET matched_text = ?, projection_hash = ? + WHERE artifact_kind = 'paper_summary' + """, + (tampered, hashlib.sha256(tampered.encode()).hexdigest()), + ) + + library = await LocalKnowledgeLibrary.open( + self.db_path, + embedding_model="fake-2d", + embedding_dimensions=2, + _embedding_provider=_FakeEmbeddingProvider(), + ) + try: + with self.assertRaisesRegex(RuntimeError, "canonical text"): + await library.search(SemanticQuery(text="summary")) + finally: + await library.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/library/test_types.py b/tests/library/test_types.py index 2476fb0..cb6ddc8 100644 --- a/tests/library/test_types.py +++ b/tests/library/test_types.py @@ -11,7 +11,12 @@ class PublicTypeTests(unittest.TestCase): def test_package_exports_only_domain_retrieval_types(self): self.assertEqual( set(library.__all__), - {"LocalKnowledgeLibrary", "SemanticQuery", "SemanticHit"}, + { + "LocalKnowledgeLibrary", + "SearchProjection", + "SemanticQuery", + "SemanticHit", + }, ) def test_query_rejects_blank_text(self): diff --git a/tests/paper_helpers.py b/tests/paper_helpers.py new file mode 100644 index 0000000..d62d5bd --- /dev/null +++ b/tests/paper_helpers.py @@ -0,0 +1,200 @@ +"""Deterministic paper artifacts shared by offline tests.""" + +import hashlib +from datetime import datetime, timezone + +from quantmind.knowledge import ( + ArtifactLocator, + PaperAssetRef, + PaperChunk, + PaperChunkingConfig, + PaperChunkSet, + PaperCitation, + PaperFlowResult, + PaperGlobalSummary, + PaperParsedManifest, + PaperParsedPage, + PaperSourceRevision, + PaperSourceSpan, + PaperSummaryProducer, + SourceRef, +) +from quantmind.knowledge.paper import ( + _paper_artifact_id, + _paper_asset_id, + _paper_chunk_id, + _paper_chunk_set_content_hash, + _paper_source_id, + _paper_summary_content_hash, + _stable_hash, + _text_hash, +) + +_RAW_BYTES = b"deterministic paper source revision" +_WHEN = datetime(2017, 12, 6, tzinfo=timezone.utc) + + +def build_paper_result( + *, + chunk_size: int = 128, + summary_model: str = "fake-summary", + summary_text: str = ( + "The Transformer replaces recurrence and convolution with an " + "encoder-decoder attention architecture, uses multi-head attention, " + "and improves translation quality with efficient training." + ), +) -> PaperFlowResult: + """Build one valid two-page result without parsing, network, or models.""" + source_hash = hashlib.sha256(_RAW_BYTES).hexdigest() + source_id = _paper_source_id(source_hash) + raw_asset = PaperAssetRef( + asset_id=_paper_asset_id( + source_id, + kind="raw", + page_number=None, + content_hash=source_hash, + ), + kind="raw", + media_type="application/pdf", + content_hash=source_hash, + size_bytes=len(_RAW_BYTES), + ) + source = PaperSourceRevision( + id=source_id, + source=SourceRef( + kind="arxiv", + uri="https://arxiv.org/pdf/1706.03762v7.pdf", + fetched_at=_WHEN, + content_hash=source_hash, + ), + as_of=_WHEN, + available_at=_WHEN, + published_at=_WHEN, + arxiv_id="1706.03762v7", + title="Attention Is All You Need", + authors=("Ashish Vaswani",), + parsed=PaperParsedManifest( + source_hash=source_hash, + parser_name="fake-parser", + parser_version="1", + cleanup_version="1", + pages=( + PaperParsedPage( + page_number=1, + width=612, + height=792, + text=( + "The Transformer removes recurrence and convolution " + "from sequence transduction." + ), + ), + PaperParsedPage( + page_number=2, + width=612, + height=792, + text=( + "Multi-head attention uses parallel projections. " + "The model improves translation and training efficiency." + ), + ), + ), + ), + raw_asset_id=raw_asset.asset_id, + assets=(raw_asset,), + blobs={source_hash: _RAW_BYTES}, + ) + + producer = PaperChunkingConfig( + splitter_version="fake-llama-index", + chunk_size=chunk_size, + chunk_overlap=min(16, chunk_size - 1), + ) + producer_hash = _stable_hash(producer.model_dump(mode="json")) + chunk_set_id = _paper_artifact_id( + source_id, + "paper_chunk_set", + producer_hash, + ) + chunk_values = ( + ( + 1, + "The Transformer removes recurrence and convolution.", + ), + (2, "Multi-head attention uses parallel learned projections."), + (2, "The model improves translation and training efficiency."), + ) + chunks: list[PaperChunk] = [] + for position, (page, text) in enumerate(chunk_values): + span = PaperSourceSpan( + page_number=page, + start_char=position * 10, + end_char=position * 10 + len(text), + ) + text_hash = _text_hash(text) + chunks.append( + PaperChunk( + chunk_id=_paper_chunk_id( + chunk_set_id, + position=position, + content_hash=text_hash, + spans=(span,), + ), + chunk_set_id=chunk_set_id, + source_revision_id=source_id, + position=position, + text=text, + content_hash=text_hash, + source_spans=(span,), + ) + ) + chunk_tuple = tuple(chunks) + chunk_set = PaperChunkSet( + id=chunk_set_id, + source_revision_id=source_id, + producer=producer, + producer_config_hash=producer_hash, + content_hash=_paper_chunk_set_content_hash(chunk_tuple), + chunks=chunk_tuple, + ) + + summary_producer = PaperSummaryProducer( + model=summary_model, + prompt_version="test-v1", + input_chunk_set_id=chunk_set_id, + instructions_hash=hashlib.sha256(b"test instructions").hexdigest(), + max_output_tokens=512, + ) + summary_config_hash = _stable_hash(summary_producer.model_dump(mode="json")) + citations = tuple( + PaperCitation( + chunk_set_id=chunk_set_id, + chunk_id=chunk.chunk_id, + page_number=chunk.source_spans[0].page_number, + ) + for chunk in chunk_tuple + ) + summary = PaperGlobalSummary( + id=_paper_artifact_id( + source_id, + "paper_summary", + summary_config_hash, + ), + source_revision_id=source_id, + producer=summary_producer, + producer_config_hash=summary_config_hash, + content_hash=_paper_summary_content_hash(summary_text, citations), + summary=summary_text, + citations=citations, + derived_from=( + ArtifactLocator( + source_revision_id=source_id, + artifact_id=chunk_set_id, + artifact_kind="paper_chunk_set", + ), + ), + ) + return PaperFlowResult( + source_revision=source, + chunk_set=chunk_set, + global_summary=summary, + ) diff --git a/tests/preprocess/fetch/test_arxiv.py b/tests/preprocess/fetch/test_arxiv.py index 4197753..48d559f 100644 --- a/tests/preprocess/fetch/test_arxiv.py +++ b/tests/preprocess/fetch/test_arxiv.py @@ -27,6 +27,7 @@ def _stub_arxiv_result( authors=["Alice Smith", "Bob Jones"], summary="Abstract goes here.", published=datetime(2024, 4, 15, tzinfo=timezone.utc), + updated=datetime(2024, 4, 16, tzinfo=timezone.utc), primary_category="q-fin.ST", categories=["q-fin.ST", "stat.ML"], ) @@ -89,9 +90,13 @@ async def test_returns_raw_paper(self): self.assertEqual(result.primary_category, "q-fin.ST") self.assertEqual(result.categories, ("q-fin.ST", "stat.ML")) self.assertEqual(result.source_url, pdf_url) + self.assertEqual(result.resolved_url, pdf_url) + self.assertIsNotNone(result.fetched_at) assert result.published_at is not None self.assertEqual(result.published_at.tzinfo, timezone.utc) self.assertEqual(result.published_at.year, 2024) + assert result.updated_at is not None + self.assertEqual(result.updated_at.day, 16) async def test_naive_published_promoted_to_utc(self): naive_published = datetime(2024, 4, 15, 12, 0) diff --git a/tests/rag/test_document.py b/tests/rag/test_document.py index 3ca0167..f80dd51 100644 --- a/tests/rag/test_document.py +++ b/tests/rag/test_document.py @@ -31,6 +31,17 @@ async def test_chunks_and_bm25_hits_keep_page_evidence(self): all(chunk.source_hash == document.source_hash for chunk in chunks) ) self.assertTrue(all(chunk.block_boxes for chunk in chunks)) + self.assertTrue( + all(0 <= chunk.start_char < chunk.end_char for chunk in chunks) + ) + repeated = chunk_parsed_document( + document, + config=SentenceSplitterConfig(chunk_size=256, chunk_overlap=32), + ) + self.assertEqual( + [chunk.chunk_id for chunk in chunks], + [chunk.chunk_id for chunk in repeated], + ) hits = retrieve_parsed_document( chunks, diff --git a/tests/test_verify_pdf_rag_e2e.py b/tests/test_verify_pdf_rag_e2e.py index fa41431..5ac4411 100644 --- a/tests/test_verify_pdf_rag_e2e.py +++ b/tests/test_verify_pdf_rag_e2e.py @@ -1,91 +1,87 @@ import io import unittest from contextlib import redirect_stdout -from types import SimpleNamespace from unittest.mock import AsyncMock, patch -from quantmind.preprocess.format import ParsedDocument, ParsedPage -from quantmind.rag import ( - ParsedChunk, - ParsedDocumentHit, -) from scripts import verify_pdf_rag_e2e -def _document(page_count: int = 15) -> ParsedDocument: - return ParsedDocument( - source_hash="hash", - parser_name="liteparse", - parser_version="2.6.0", - cleanup_version="1", - pages=tuple( - ParsedPage( - page_number=page, - width=612, - height=792, - text=f"page {page}", - blocks=(), - ) - for page in range(1, page_count + 1) +def _snapshot(**overrides): + value = { + "arxiv_id": "1706.03762v7", + "page_count": 15, + "text_page_count": 15, + "asset_count": 16, + "screenshot_pages": list(range(1, 16)), + "chunk_asset_reference_count": 42, + "chunk_count": 42, + "summary": ( + "An attention-only encoder-decoder removes recurrence and " + "convolution, uses multi-head attention, and improves translation " + "with greater training efficiency." ), - ) - - -def _chunk(text: str, page: int = 5) -> ParsedChunk: - return ParsedChunk( - chunk_id=f"chunk-{page}", - text=text, - source_hash="hash", - page_number=page, - block_boxes=(), - screenshot_path=None, - image_paths=(), - ) - - -class VerifyPdfRagE2ETests(unittest.IsolatedAsyncioTestCase): - async def test_main_passes_pinned_page_and_relevance_checks(self): - paper = SimpleNamespace( - arxiv_id="1706.03762v7", - bytes=b"pdf", + "citation_count": 3, + "citation_pages": [2, 4], + "first_summary_scores": [0.9], + "first_chunk_scores": [0.8], + "second_summary_scores": [0.9], + "second_chunk_scores": [0.8], + "first_summary_resolved": True, + "second_summary_resolved": True, + "first_multi_head_pages": [5], + "second_multi_head_pages": [5], + "restored": True, + "embedding_model": "text-embedding-3-small", + } + value.update(overrides) + return value + + +class VerifyPaperFlowE2ETests(unittest.IsolatedAsyncioTestCase): + def test_summary_coverage_accepts_exclusive_self_attention_wording(self): + summary = ( + "The encoder-decoder relies exclusively on self-attention instead " + "of recurrent and convolutional networks. Multi-head attention " + "improves translation with less training time." + ) + + self.assertTrue( + verify_pdf_rag_e2e._summary_has_required_coverage(summary) + ) + + def test_summary_coverage_accepts_pure_attention_wording(self): + summary = ( + "The encoder-decoder relies purely on attention, eschewing " + "recurrence and convolutions. It uses multi-head attention and " + "improves translation with less training time." + ) + + self.assertTrue( + verify_pdf_rag_e2e._summary_has_required_coverage(summary) ) - document = _document() - chunks = (_chunk("Multi-head attention projects queries."),) - hits = (ParsedDocumentHit(chunk=chunks[0], score=1.0),) + + async def test_main_passes_complete_vertical_slice(self): with ( + patch.dict("os.environ", {"OPENAI_API_KEY": "test"}), patch.object( verify_pdf_rag_e2e, - "fetch_arxiv", - new=AsyncMock(return_value=paper), - ), - patch.object( - verify_pdf_rag_e2e, - "parse_pdf", - new=AsyncMock(return_value=document), - ), - patch.object( - verify_pdf_rag_e2e, - "chunk_parsed_document", - return_value=chunks, - ), - patch.object( - verify_pdf_rag_e2e, - "retrieve_parsed_document", - return_value=hits, + "_run_vertical_slice", + new=AsyncMock(return_value=_snapshot()), ), redirect_stdout(io.StringIO()) as output, ): exit_code = await verify_pdf_rag_e2e.main() self.assertEqual(exit_code, 0) - self.assertIn("[PASS] pdf-rag", output.getvalue()) - self.assertIn("top_pages=[5]", output.getvalue()) + self.assertIn("[PASS] paper-flow-v1", output.getvalue()) + self.assertIn("multi_head_pages=[5]", output.getvalue()) async def test_main_reports_upstream_failure(self): with ( + patch.dict("os.environ", {"OPENAI_API_KEY": "test"}), patch.object( verify_pdf_rag_e2e, - "fetch_arxiv", + "_run_vertical_slice", new=AsyncMock(side_effect=TimeoutError("bounded timeout")), ), redirect_stdout(io.StringIO()) as output, @@ -93,43 +89,27 @@ async def test_main_reports_upstream_failure(self): exit_code = await verify_pdf_rag_e2e.main() self.assertEqual(exit_code, 1) - self.assertIn("[FAIL] pdf-rag: TimeoutError", output.getvalue()) + self.assertIn("[FAIL] paper-flow-v1: TimeoutError", output.getvalue()) - async def test_main_rejects_wrong_page_count_or_irrelevant_hits(self): - paper = SimpleNamespace( - arxiv_id="1706.03762v7", - bytes=b"pdf", - ) - document = _document(page_count=14) - chunks = (_chunk("Unrelated passage"),) - hits = (ParsedDocumentHit(chunk=chunks[0], score=1.0),) + async def test_main_rejects_missing_reopen_or_summary_coverage(self): with ( + patch.dict("os.environ", {"OPENAI_API_KEY": "test"}), patch.object( verify_pdf_rag_e2e, - "fetch_arxiv", - new=AsyncMock(return_value=paper), - ), - patch.object( - verify_pdf_rag_e2e, - "parse_pdf", - new=AsyncMock(return_value=document), - ), - patch.object( - verify_pdf_rag_e2e, - "chunk_parsed_document", - return_value=chunks, - ), - patch.object( - verify_pdf_rag_e2e, - "retrieve_parsed_document", - return_value=hits, + "_run_vertical_slice", + new=AsyncMock( + return_value=_snapshot( + restored=False, + summary="An unrelated summary.", + ) + ), ), redirect_stdout(io.StringIO()) as output, ): exit_code = await verify_pdf_rag_e2e.main() self.assertEqual(exit_code, 1) - self.assertIn("[FAIL] pdf-rag", output.getvalue()) + self.assertIn("[FAIL] paper-flow-v1", output.getvalue()) if __name__ == "__main__": From fd38db3da5850d242e12c432f7f8eb5af1011c82 Mon Sep 17 00:00:00 2001 From: pkuwkl Date: Sat, 18 Jul 2026 15:29:03 +0800 Subject: [PATCH 2/4] fix(ci): skip paper smoke without API key --- .github/workflows/e2e.yml | 12 ++++++++++++ docs/README.md | 3 +++ 2 files changed, 15 insertions(+) diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index 7376606..74cf989 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -137,5 +137,17 @@ jobs: - name: Install project runtime dependencies run: uv pip install --python .venv/bin/python -e . + - name: Check OpenAI credential + id: openai + shell: bash + run: | + if [[ -n "${OPENAI_API_KEY}" ]]; then + echo "available=true" >> "$GITHUB_OUTPUT" + else + echo "available=false" >> "$GITHUB_OUTPUT" + echo "::notice::Skipping Paper Flow V1 E2E because OPENAI_API_KEY is not configured." + fi + - name: Run live Paper Flow V1 E2E + if: steps.openai.outputs.available == 'true' run: .venv/bin/python scripts/verify_pdf_rag_e2e.py diff --git a/docs/README.md b/docs/README.md index 5758058..dbf8aa3 100644 --- a/docs/README.md +++ b/docs/README.md @@ -44,6 +44,9 @@ chunk projections with `text-embedding-3-small`, reopens the database, searches both artifact kinds, and resolves every hit. It runs daily, manually, and on pull requests that change its dependency paths, and it remains non-required because arXiv and model providers are public-network dependencies. +When the repository `OPENAI_API_KEY` secret is unavailable, the job emits an +explicit skip notice instead of reporting an implementation failure; the +catalog command remains the direct way to run the same bounded slice locally. ## Verification From ed998e55a51763620e1acde92de1d09d8b534489 Mon Sep 17 00:00:00 2001 From: pkuwkl Date: Sun, 19 Jul 2026 00:01:34 +0800 Subject: [PATCH 3/4] refactor(paper): add bounded research subagents --- contexts/design/flow/paper.md | 29 +- contexts/design/knowledge/paper.md | 6 +- contexts/design/library/local.md | 2 +- docs/library.md | 6 +- examples/flows/paper.py | 14 +- quantmind/configs/paper.py | 18 +- quantmind/flows/_paper_summary.py | 454 ++++++++++++++++++++++++----- quantmind/flows/paper.py | 4 + quantmind/knowledge/paper.py | 28 +- quantmind/library/_types.py | 9 +- scripts/verify_pdf_rag_e2e.py | 16 +- tests/configs/test_paper.py | 14 + tests/flows/test_paper.py | 159 +++++++++- tests/library/test_paper.py | 10 +- tests/library/test_types.py | 17 ++ tests/test_verify_pdf_rag_e2e.py | 1 + 16 files changed, 672 insertions(+), 115 deletions(-) diff --git a/contexts/design/flow/paper.md b/contexts/design/flow/paper.md index 4894810..4bca46c 100644 --- a/contexts/design/flow/paper.md +++ b/contexts/design/flow/paper.md @@ -53,9 +53,11 @@ The operation has a strict order: 3. Build and validate `PaperSourceRevision`, including content-addressed asset references and blobs. 4. Chunk each page with LlamaIndex while retaining page and character spans. 5. Build and validate `PaperChunkSet` with code-owned IDs and membership. -6. Invoke the bounded summarizer using only the validated chunk set. -7. Resolve model-returned chunk/page coordinates into canonical citations in code. -8. Build and validate `PaperGlobalSummary` and the cross-artifact `PaperFlowResult`. +6. Give a coordinator agent the bounded chunk manifest and a research-agent tool. +7. Require the coordinator to delegate complete chunk coverage to one or more bounded research subagents, with independent ranges eligible for parallel execution. +8. Let the coordinator synthesize typed worker reports and make bounded overlapping follow-up calls when needed. +9. Resolve model-returned chunk/page coordinates into canonical citations in code. +10. Build and validate `PaperGlobalSummary` and the cross-artifact `PaperFlowResult`. A summarization failure occurs after source and chunks exist in memory, but `paper_flow` returns no partial success value. Persistence is a separate explicit operation. @@ -97,17 +99,20 @@ Code then creates `PaperCitation` values and a `PaperGlobalSummary`. Its produce ## Bounded Model Calls -`PaperFlowCfg` makes runtime and usage bounds explicit: +`PaperFlowCfg` makes coordinator and nested-research bounds explicit: -- `max_summary_tool_calls` caps adaptive chunk reads; -- `max_summary_concurrency` bounds simultaneous chunk-read tools; -- `max_summary_input_tokens` caps estimated manifest and chunk-read input; -- `max_summary_output_tokens` caps the structured response and validated summary size; -- `max_turns` caps Agents SDK turns and defaults to 16 for the paper summarizer; +- `max_summary_tool_calls` caps research-subagent invocations, including focused follow-ups; +- `max_summary_concurrency` bounds simultaneous research-subagent runs; +- `max_summary_worker_turns` caps each nested agent run; +- `max_summary_worker_output_tokens` caps each structured worker report; +- `max_summary_input_tokens` caps the manifest, worker chunk inputs, and worker reports returned to the coordinator; +- `max_summary_output_tokens` caps the coordinator's final structured response; +- `max_summary_total_output_tokens` caps worker reports and final output together; +- `max_turns` caps coordinator turns and defaults to 16; - `timeout_seconds` bounds the complete summarization operation; - `summary_prompt_version` and `summary_instructions` version the semantic producer. -The summarizer receives a bounded manifest and may adaptively read up to eight consecutive chunks per tool call. A shared concurrency-safe budget is reserved before every read. The operation rejects work that would exceed a call or token limit rather than silently running without a bound. +The coordinator uses the Agents SDK manager-style pattern: `paper_chunk_researcher` is exposed through `Agent.as_tool()`, so the coordinator retains responsibility for the final summary while each specialist receives only its requested range. The manifest provides code-generated groups of at most eight consecutive chunks. Independent calls may run in parallel, and the coordinator may request bounded overlapping follow-up research. Code rejects missing full-chunk coverage, out-of-scope worker citations, invalid pages or quotes, and any call, concurrency, token, or runtime excess. ## Failure Semantics @@ -116,7 +121,7 @@ The summarizer receives a bounded manifest and may adaptively read up to eight c - Fetching, parsing, or missing parser assets raise their source error and produce no result. - Empty chunk output is invalid. - Invalid or insufficient summary citations raise `PaperCitationValidationError`. -- Call, token, output, concurrency, and timeout violations fail the summary operation. +- Missing worker coverage, invalid worker evidence, and call, token, output, concurrency, or timeout violations fail the summary operation. - Any canonical identity, content hash, membership, lineage, or cross-artifact mismatch fails Pydantic validation. No failure is converted into a partially valid `PaperFlowResult`. Callers may retry with the same source and producer settings; stable IDs make successful repeated runs idempotent. @@ -137,7 +142,7 @@ The bounded live slice is: python scripts/verify_pdf_rag_e2e.py ``` -It fetches exact arXiv revision `1706.03762v7`, parses at least 15 physical pages, creates multiple page-aware chunks, generates a cited global summary with bounded `gpt-4o-mini` calls, persists with `text-embedding-3-small`, closes and reopens the database, runs summary and chunk searches, resolves every returned locator, and prints useful summary, page, citation, and score diagnostics. The dedicated `paper-flow` job in `.github/workflows/e2e.yml` owns this non-required public-network check. +It fetches exact arXiv revision `1706.03762v7`, parses at least 15 physical pages, creates multiple page-aware chunks, delegates complete coverage to bounded `gpt-4o-mini` research subagents, synthesizes a cited global summary, persists with `text-embedding-3-small`, closes and reopens the database, runs summary and chunk searches, resolves every returned locator, and prints useful summary, page, citation, and score diagnostics. The dedicated `paper-flow` job in `.github/workflows/e2e.yml` owns this non-required public-network check. ## Out of Scope diff --git a/contexts/design/knowledge/paper.md b/contexts/design/knowledge/paper.md index 1139078..c90be9c 100644 --- a/contexts/design/knowledge/paper.md +++ b/contexts/design/knowledge/paper.md @@ -62,9 +62,11 @@ Every chunk span is also checked against that manifest: its page must exist, its - model identity; - prompt version; +- manager/research-agent orchestration version; - exact input chunk-set ID; -- summary-instruction hash; -- maximum output tokens. +- coordinator/research instructions hash; +- coordinator output limit; +- research call, concurrency, turn, and per-worker output limits. Changing any producer field creates a distinct artifact ID. Multiple chunk sets and summaries may coexist for one source revision. Loading a complete `PaperFlowResult` without explicit artifact IDs is allowed only when one unambiguous linked pair exists. diff --git a/contexts/design/library/local.md b/contexts/design/library/local.md index 9ee8ba5..f698ef4 100644 --- a/contexts/design/library/local.md +++ b/contexts/design/library/local.md @@ -94,7 +94,7 @@ Every stored projection records target identity, exact projection text and hash, ## Search and Resolution -`SemanticQuery.artifact_kinds` filters artifact granularity, including `paper_summary` and `paper_chunk_set`. Existing `item_types`, source, confidence, tag, tree, `as_of`, and `available_at` filters continue to apply where meaningful. Filtering happens before ranking. +`SemanticQuery.artifact_kinds` is a Pydantic-validated list of `PaperArtifactKind` values rather than an open string list. `PaperArtifactKind.GLOBAL_SUMMARY` selects summary aggregates and `PaperArtifactKind.CHUNK_SET` selects chunk members owned by chunk-set artifacts. JSON and YAML callers may still use the enum values `paper_summary` and `paper_chunk_set`; unknown values fail validation. Existing `item_types`, source, confidence, tag, tree, `as_of`, and `available_at` filters continue to apply where meaningful. Filtering happens before ranking. Each `SemanticHit` carries two complementary values: diff --git a/docs/library.md b/docs/library.md index c427cd0..2cb804f 100644 --- a/docs/library.md +++ b/docs/library.md @@ -31,6 +31,8 @@ Putting the same result again is safe and reuses valid vectors. A changed splitt Opening a library performs no embedding or network request. Search embeds only the query when stored projections are reusable: ```python +from quantmind.knowledge import PaperArtifactKind + library = await LocalKnowledgeLibrary.open( ".quantmind/library.db", embedding_model="text-embedding-3-small", @@ -39,14 +41,14 @@ try: summary_hits = await library.search( SemanticQuery( text="What is the paper's central contribution?", - artifact_kinds=["paper_summary"], + artifact_kinds=[PaperArtifactKind.GLOBAL_SUMMARY], top_k=3, ) ) chunk_hits = await library.search( SemanticQuery( text="How does multi-head attention work?", - artifact_kinds=["paper_chunk_set"], + artifact_kinds=[PaperArtifactKind.CHUNK_SET], top_k=5, ) ) diff --git a/examples/flows/paper.py b/examples/flows/paper.py index e4dad7f..2c1e4c4 100644 --- a/examples/flows/paper.py +++ b/examples/flows/paper.py @@ -10,7 +10,11 @@ from quantmind.configs import PaperFlowCfg from quantmind.configs.paper import ArxivIdentifier from quantmind.flows import paper_flow -from quantmind.knowledge import PaperChunk, PaperGlobalSummary +from quantmind.knowledge import ( + PaperArtifactKind, + PaperChunk, + PaperGlobalSummary, +) from quantmind.library import ( LocalKnowledgeLibrary, SemanticHit, @@ -32,14 +36,14 @@ async def _search_and_resolve( summary_hits = await library.search( SemanticQuery( text="What is the paper's central contribution?", - artifact_kinds=["paper_summary"], + artifact_kinds=[PaperArtifactKind.GLOBAL_SUMMARY], top_k=3, ) ) chunk_hits = await library.search( SemanticQuery( text="How does multi-head attention work?", - artifact_kinds=["paper_chunk_set"], + artifact_kinds=[PaperArtifactKind.CHUNK_SET], top_k=5, ) ) @@ -98,6 +102,10 @@ async def main() -> None: ) = await _search_and_resolve(library) print(restored.global_summary.summary) + print( + "summary_orchestration=" + f"{restored.global_summary.producer.orchestration}" + ) print( f"chunks={len(restored.chunk_set.chunks)} " f"source_pages={len(restored.source_revision.parsed.pages)}" diff --git a/quantmind/configs/paper.py b/quantmind/configs/paper.py index 18c3a22..2df83be 100644 --- a/quantmind/configs/paper.py +++ b/quantmind/configs/paper.py @@ -61,12 +61,15 @@ class PaperFlowCfg(BaseFlowCfg): max_turns: int = Field(default=16, ge=1) chunk_size: int = Field(default=512, gt=0) chunk_overlap: int = Field(default=64, ge=0) - summary_prompt_version: str = "paper-summary-v1" + summary_prompt_version: str = "paper-summary-v2" summary_instructions: str | None = None max_summary_tool_calls: int = Field(default=12, ge=1) max_summary_concurrency: int = Field(default=2, ge=1) + max_summary_worker_turns: int = Field(default=4, ge=1) + max_summary_worker_output_tokens: int = Field(default=1_536, ge=1) max_summary_input_tokens: int = Field(default=120_000, ge=1) max_summary_output_tokens: int = Field(default=4_096, ge=1) + max_summary_total_output_tokens: int = Field(default=20_000, ge=1) min_summary_citations: int = Field(default=3, ge=1) min_summary_pages: int = Field(default=2, ge=1) @@ -78,4 +81,17 @@ def _validate_paper_bounds(self) -> "PaperFlowCfg": raise ValueError( "min_summary_pages cannot exceed min_summary_citations" ) + if self.max_summary_concurrency > self.max_summary_tool_calls: + raise ValueError( + "max_summary_concurrency cannot exceed max_summary_tool_calls" + ) + minimum_output_budget = ( + self.max_summary_output_tokens + + self.max_summary_worker_output_tokens + ) + if self.max_summary_total_output_tokens < minimum_output_budget: + raise ValueError( + "max_summary_total_output_tokens must reserve one worker " + "report and the final summary" + ) return self diff --git a/quantmind/flows/_paper_summary.py b/quantmind/flows/_paper_summary.py index 179b8b7..f5c7f2b 100644 --- a/quantmind/flows/_paper_summary.py +++ b/quantmind/flows/_paper_summary.py @@ -1,26 +1,42 @@ -"""Bounded Agents SDK synthesis for one paper chunk-set artifact.""" +"""Bounded multi-agent synthesis for one paper chunk-set artifact.""" import asyncio import hashlib import json +from collections.abc import Awaitable, Callable from dataclasses import replace -from typing import Any, Protocol +from typing import Any, Literal, Protocol -from agents import Agent, ModelSettings, function_tool +from agents import Agent, FunctionTool, ModelSettings from pydantic import BaseModel, ConfigDict, Field, field_validator from quantmind.configs import PaperFlowCfg from quantmind.flows._runner import run_with_observability from quantmind.knowledge import PaperChunkSet, PaperSourceRevision +_MAX_RESEARCH_GROUP_SIZE = 8 +_ORCHESTRATION_VERSION = "manager-research-agents-v1" + _SUMMARY_INSTRUCTIONS = """\ -Write one accurate global summary of the supplied paper. Use -`read_chunk_group` to inspect source chunks before writing. Cover the central -contribution, architecture or methodology, principal results, and important -limitations. Return only summary prose plus citations. Each citation uses the -zero-based chunk index and a physical page owned by that chunk. Never invent -IDs, source metadata, storage links, pages, or quotations. Read useful -consecutive chunks in groups of up to eight; do not spend one call per chunk. +Act as the paper-summary coordinator. Delegate every suggested research group +to `research_chunk_group` before writing the final summary. You may call +independent groups in parallel and may make bounded overlapping follow-up calls +when a worker report exposes an ambiguity. Synthesize the worker reports into +one accurate global summary covering the central contribution, architecture or +methodology, principal results, and important limitations. Return only summary +prose plus citations. Each citation uses the zero-based chunk index and a +physical page owned by that chunk. Never invent IDs, source metadata, storage +links, pages, or quotations. Set a final citation quote to null unless it is an +exact contiguous substring copied character-for-character from the chunk. +""" + +_RESEARCH_INSTRUCTIONS = """\ +Act as a bounded paper research specialist. Analyze only the supplied chunk +range and requested focus. Return a concise scope summary plus structured +findings. Classify every finding as context, contribution, method, result, or +limitation and support it with a chunk index and physical page. Do not infer +canonical IDs, source metadata, or claims that are not supported by the +supplied chunks. """ @@ -38,6 +54,74 @@ class PaperSummaryCitationDraft(BaseModel): quote: str | None = Field(default=None, max_length=500) +class PaperResearchRequest(BaseModel): + """Structured coordinator request for one bounded research subagent.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + start: int = Field(ge=0) + count: int = Field(ge=1, le=_MAX_RESEARCH_GROUP_SIZE) + focus: str = Field(min_length=1, max_length=500) + + @field_validator("focus") + @classmethod + def _focus_is_not_blank(cls, value: str) -> str: + stripped = value.strip() + if not stripped: + raise ValueError("paper research focus must not be blank") + return stripped + + +class PaperResearchCitationDraft(BaseModel): + """Chunk and page coordinates returned by a research subagent.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + chunk_index: int = Field(ge=0) + page_number: int = Field(ge=1) + + +class PaperResearchFindingDraft(BaseModel): + """One subagent finding with non-canonical source coordinates.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + kind: Literal[ + "context", + "contribution", + "method", + "result", + "limitation", + ] + claim: str = Field(min_length=1) + citation: PaperResearchCitationDraft + + @field_validator("claim") + @classmethod + def _claim_is_not_blank(cls, value: str) -> str: + stripped = value.strip() + if not stripped: + raise ValueError("paper research claim must not be blank") + return stripped + + +class PaperResearchDraft(BaseModel): + """Typed evidence report returned by one research subagent.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + scope_summary: str = Field(min_length=1) + findings: tuple[PaperResearchFindingDraft, ...] = Field(min_length=1) + + @field_validator("scope_summary") + @classmethod + def _scope_summary_is_not_blank(cls, value: str) -> str: + stripped = value.strip() + if not stripped: + raise ValueError("paper research scope summary must not be blank") + return stripped + + class PaperSummaryDraft(BaseModel): """Limited model output containing prose and non-canonical coordinates.""" @@ -49,9 +133,10 @@ class PaperSummaryDraft(BaseModel): @field_validator("summary") @classmethod def _summary_is_not_blank(cls, value: str) -> str: - if not value.strip(): + stripped = value.strip() + if not stripped: raise ValueError("paper summary draft must not be blank") - return value + return stripped class _PaperSummaryProvider(Protocol): @@ -72,64 +157,216 @@ def _estimated_tokens(text: str) -> int: return max(1, (len(text) + 3) // 4) +def _research_payload( + source: PaperSourceRevision, + chunk_set: PaperChunkSet, + request: PaperResearchRequest, +) -> str: + if request.start >= len(chunk_set.chunks): + raise ValueError("research start is outside the chunk-set manifest") + selected = chunk_set.chunks[request.start : request.start + request.count] + if len(selected) != request.count: + raise ValueError("research range exceeds the chunk-set manifest") + return json.dumps( + { + "title": source.title, + "authors": source.authors, + "focus": request.focus, + "start": request.start, + "count": request.count, + "chunks": [ + { + "chunk_index": chunk.position, + "pages": sorted( + {span.page_number for span in chunk.source_spans} + ), + "text": chunk.text, + } + for chunk in selected + ], + }, + ensure_ascii=False, + ) + + +def _coerce_research_draft(value: Any) -> PaperResearchDraft: + if isinstance(value, str): + return PaperResearchDraft.model_validate_json(value) + return PaperResearchDraft.model_validate(value) + + +def _validate_research_draft( + chunk_set: PaperChunkSet, + request: PaperResearchRequest, + draft: PaperResearchDraft, +) -> None: + allowed = set(range(request.start, request.start + request.count)) + for finding in draft.findings: + citation = finding.citation + if citation.chunk_index not in allowed: + raise ValueError("research finding cites a chunk outside its scope") + chunk = chunk_set.chunks[citation.chunk_index] + pages = {span.page_number for span in chunk.source_spans} + if citation.page_number not in pages: + raise ValueError("research finding cites a page outside its chunk") + + class _SummaryBudget: - """Concurrency-safe accounting for adaptive chunk reads.""" + """Concurrency-safe accounting for nested research-agent calls.""" - def __init__(self, cfg: PaperFlowCfg, initial_text: str) -> None: + def __init__( + self, + cfg: PaperFlowCfg, + initial_text: str, + *, + chunk_count: int, + ) -> None: initial_tokens = _estimated_tokens(initial_text) if initial_tokens > cfg.max_summary_input_tokens: raise PaperSummaryBudgetExceeded( "paper summary manifest exceeds max_summary_input_tokens" ) + required_calls = ( + chunk_count + _MAX_RESEARCH_GROUP_SIZE - 1 + ) // _MAX_RESEARCH_GROUP_SIZE + if required_calls > cfg.max_summary_tool_calls: + raise PaperSummaryBudgetExceeded( + "max_summary_tool_calls cannot cover every paper chunk" + ) self._max_calls = cfg.max_summary_tool_calls self._max_input_tokens = cfg.max_summary_input_tokens + self._max_worker_output_tokens = cfg.max_summary_worker_output_tokens + self._max_final_output_tokens = cfg.max_summary_output_tokens + self._max_total_output_tokens = cfg.max_summary_total_output_tokens self._calls = 0 self._input_tokens = initial_tokens + self._output_tokens = 0 + self._reserved_input_tokens = 0 + self._reserved_output_tokens = 0 + self._covered_chunks: set[int] = set() self._lock = asyncio.Lock() self._semaphore = asyncio.Semaphore(cfg.max_summary_concurrency) - async def read( + @property + def research_calls(self) -> int: + """Return the number of research calls reserved so far.""" + return self._calls + + @property + def covered_chunks(self) -> frozenset[int]: + """Return chunk positions covered by successful worker reports.""" + return frozenset(self._covered_chunks) + + async def invoke_research( self, + source: PaperSourceRevision, chunk_set: PaperChunkSet, - *, - start: int, - count: int, + request: PaperResearchRequest, + operation: Callable[[], Awaitable[Any]], ) -> str: - """Return one bounded group while reserving calls and input tokens.""" - if start < 0 or count < 1 or count > 8: - raise ValueError( - "start must be non-negative and count must be 1..8" - ) - if start >= len(chunk_set.chunks): - raise ValueError("start is outside the chunk-set manifest") - async with self._semaphore: - selected = chunk_set.chunks[start : start + count] - payload = json.dumps( - [ - { - "chunk_index": chunk.position, - "pages": sorted( - {span.page_number for span in chunk.source_spans} - ), - "text": chunk.text, - } - for chunk in selected - ], - ensure_ascii=False, - ) - tokens = _estimated_tokens(payload) + """Run one agent tool call while reserving all shared budgets.""" + payload = _research_payload(source, chunk_set, request) + payload_tokens = _estimated_tokens(payload) + input_reservation = payload_tokens + self._max_worker_output_tokens + await self._semaphore.acquire() + reserved = False + try: async with self._lock: if self._calls >= self._max_calls: raise PaperSummaryBudgetExceeded( "paper summary exceeded max_summary_tool_calls" ) - if self._input_tokens + tokens > self._max_input_tokens: + projected_input = ( + self._input_tokens + + self._reserved_input_tokens + + input_reservation + ) + if projected_input > self._max_input_tokens: raise PaperSummaryBudgetExceeded( "paper summary exceeded max_summary_input_tokens" ) + projected_output = ( + self._output_tokens + + self._reserved_output_tokens + + self._max_worker_output_tokens + + self._max_final_output_tokens + ) + if projected_output > self._max_total_output_tokens: + raise PaperSummaryBudgetExceeded( + "paper summary exceeded max_summary_total_output_tokens" + ) self._calls += 1 - self._input_tokens += tokens - return payload + self._reserved_input_tokens += input_reservation + self._reserved_output_tokens += self._max_worker_output_tokens + reserved = True + + raw_output = await operation() + draft = _coerce_research_draft(raw_output) + _validate_research_draft(chunk_set, request, draft) + serialized = draft.model_dump_json() + output_tokens = _estimated_tokens(serialized) + if output_tokens > self._max_worker_output_tokens: + raise PaperSummaryBudgetExceeded( + "paper research worker exceeded " + "max_summary_worker_output_tokens" + ) + + async with self._lock: + self._reserved_input_tokens -= input_reservation + self._reserved_output_tokens -= self._max_worker_output_tokens + self._input_tokens += payload_tokens + output_tokens + self._output_tokens += output_tokens + self._covered_chunks.update( + range(request.start, request.start + request.count) + ) + reserved = False + return serialized + finally: + if reserved: + async with self._lock: + self._reserved_input_tokens -= input_reservation + self._reserved_output_tokens -= ( + self._max_worker_output_tokens + ) + self._semaphore.release() + + async def accept_final_output(self, draft: PaperSummaryDraft) -> None: + """Validate final output and require complete worker coverage.""" + serialized = draft.model_dump_json() + output_tokens = _estimated_tokens(serialized) + if output_tokens > self._max_final_output_tokens: + raise PaperSummaryBudgetExceeded( + "paper summary exceeded max_summary_output_tokens" + ) + async with self._lock: + if ( + self._output_tokens + output_tokens + > self._max_total_output_tokens + ): + raise PaperSummaryBudgetExceeded( + "paper summary exceeded max_summary_total_output_tokens" + ) + self._output_tokens += output_tokens + + def require_complete_coverage(self, chunk_count: int) -> None: + """Reject a coordinator that did not delegate every source chunk.""" + missing = set(range(chunk_count)) - self._covered_chunks + if missing: + preview = ", ".join(str(index) for index in sorted(missing)[:8]) + raise PaperSummaryBudgetExceeded( + "paper summary research did not cover every chunk; missing " + f"{preview}" + ) + + +def _suggested_research_groups(chunk_count: int) -> list[dict[str, int]]: + return [ + { + "start": start, + "count": min(_MAX_RESEARCH_GROUP_SIZE, chunk_count - start), + } + for start in range(0, chunk_count, _MAX_RESEARCH_GROUP_SIZE) + ] def _summary_manifest( @@ -153,6 +390,9 @@ def _summary_manifest( "page_count": len(source.parsed.pages), "required_citations": cfg.min_summary_citations, "required_source_pages": cfg.min_summary_pages, + "required_research_groups": _suggested_research_groups( + len(chunk_set.chunks) + ), "chunks": manifest, }, ensure_ascii=False, @@ -170,9 +410,25 @@ def _summary_instructions(cfg: PaperFlowCfg) -> str: def _summary_instructions_hash(cfg: PaperFlowCfg) -> str: - return hashlib.sha256( - _summary_instructions(cfg).encode("utf-8") - ).hexdigest() + payload = json.dumps( + { + "orchestration": _ORCHESTRATION_VERSION, + "coordinator_instructions": _summary_instructions(cfg), + "research_instructions": _RESEARCH_INSTRUCTIONS, + "max_research_calls": cfg.max_summary_tool_calls, + "max_research_concurrency": cfg.max_summary_concurrency, + "research_max_turns": cfg.max_summary_worker_turns, + "research_max_output_tokens": ( + cfg.max_summary_worker_output_tokens + ), + "max_total_input_tokens": cfg.max_summary_input_tokens, + "max_total_output_tokens": cfg.max_summary_total_output_tokens, + }, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() def _bounded_model_settings(cfg: PaperFlowCfg) -> ModelSettings: @@ -185,8 +441,69 @@ def _bounded_model_settings(cfg: PaperFlowCfg) -> ModelSettings: ) +def _bounded_research_model_settings(cfg: PaperFlowCfg) -> ModelSettings: + settings = cfg.model_settings or ModelSettings() + configured = settings.max_tokens or cfg.max_summary_worker_output_tokens + return replace( + settings, + max_tokens=min(configured, cfg.max_summary_worker_output_tokens), + parallel_tool_calls=False, + ) + + +def _build_research_agent_tool( + source: PaperSourceRevision, + chunk_set: PaperChunkSet, + cfg: PaperFlowCfg, + budget: _SummaryBudget, +) -> FunctionTool: + """Expose one bounded research agent to the summary coordinator.""" + researcher: Agent[Any] = Agent( + name="paper_chunk_researcher", + instructions=_RESEARCH_INSTRUCTIONS, + model=cfg.model, + model_settings=_bounded_research_model_settings(cfg), + output_type=PaperResearchDraft, + ) + + def build_input(options: Any) -> str: + request = PaperResearchRequest.model_validate(options["params"]) + return _research_payload(source, chunk_set, request) + + async def extract_output(result: Any) -> str: + draft = PaperResearchDraft.model_validate(result.final_output) + return draft.model_dump_json() + + tool = researcher.as_tool( + tool_name="research_chunk_group", + tool_description=( + "Delegate one bounded chunk range to a paper research subagent. " + "Call every required research group before final synthesis and " + "use extra overlapping calls only for focused follow-up." + ), + custom_output_extractor=extract_output, + max_turns=cfg.max_summary_worker_turns, + parameters=PaperResearchRequest, + input_builder=build_input, + failure_error_function=None, + ) + invoke_agent = tool.on_invoke_tool + + async def invoke_bounded_agent(context: Any, input_json: str) -> str: + request = PaperResearchRequest.model_validate_json(input_json) + return await budget.invoke_research( + source, + chunk_set, + request, + lambda: invoke_agent(context, input_json), + ) + + tool.on_invoke_tool = invoke_bounded_agent + return tool + + class _AgentsPaperSummaryProvider: - """Use one SDK agent with bounded adaptive access to chunk groups.""" + """Coordinate bounded research subagents and synthesize one summary.""" async def summarize( self, @@ -196,34 +513,29 @@ async def summarize( cfg: PaperFlowCfg, ) -> PaperSummaryDraft: manifest = _summary_manifest(source, chunk_set, cfg) - budget = _SummaryBudget(cfg, manifest) - - @function_tool - async def read_chunk_group(start: int, count: int) -> str: - """Read up to eight consecutive paper chunks. - - Args: - start: Zero-based first chunk index. - count: Number of chunks to read, from one through eight. - """ - return await budget.read( - chunk_set, - start=start, - count=count, - ) - - agent: Agent[Any] = Agent( - name="paper_global_summarizer", + budget = _SummaryBudget( + cfg, + manifest, + chunk_count=len(chunk_set.chunks), + ) + research_tool = _build_research_agent_tool( + source, + chunk_set, + cfg, + budget, + ) + coordinator: Agent[Any] = Agent( + name="paper_summary_coordinator", instructions=_summary_instructions(cfg), model=cfg.model, model_settings=_bounded_model_settings(cfg), - tools=[read_chunk_group], + tools=[research_tool], output_type=PaperSummaryDraft, ) try: output = await asyncio.wait_for( run_with_observability( - agent, + coordinator, manifest, cfg=cfg, memory=None, @@ -236,8 +548,6 @@ async def read_chunk_group(start: int, count: int) -> str: "paper summary exceeded timeout_seconds" ) from exc draft = PaperSummaryDraft.model_validate(output) - if _estimated_tokens(draft.summary) > cfg.max_summary_output_tokens: - raise PaperSummaryBudgetExceeded( - "paper summary exceeded max_summary_output_tokens" - ) + budget.require_complete_coverage(len(chunk_set.chunks)) + await budget.accept_final_output(draft) return draft diff --git a/quantmind/flows/paper.py b/quantmind/flows/paper.py index 6638dee..2e14b8e 100644 --- a/quantmind/flows/paper.py +++ b/quantmind/flows/paper.py @@ -482,6 +482,10 @@ def _build_global_summary( input_chunk_set_id=chunk_set.id, instructions_hash=_summary_instructions_hash(cfg), max_output_tokens=cfg.max_summary_output_tokens, + max_research_calls=cfg.max_summary_tool_calls, + max_research_concurrency=cfg.max_summary_concurrency, + research_max_turns=cfg.max_summary_worker_turns, + research_max_output_tokens=cfg.max_summary_worker_output_tokens, ) producer_hash = _stable_hash(producer.model_dump(mode="json")) artifact_id = _paper_artifact_id( diff --git a/quantmind/knowledge/paper.py b/quantmind/knowledge/paper.py index 184f06c..fe24944 100644 --- a/quantmind/knowledge/paper.py +++ b/quantmind/knowledge/paper.py @@ -4,6 +4,7 @@ import json import re from datetime import datetime +from enum import Enum from typing import Literal from uuid import NAMESPACE_URL, UUID, uuid5 @@ -18,7 +19,12 @@ from quantmind.knowledge._base import SourceRef from quantmind.knowledge._tree import TreeKnowledge -PaperArtifactKind = Literal["paper_chunk_set", "paper_summary"] + +class PaperArtifactKind(str, Enum): + """Closed paper-artifact discriminator accepted by Pydantic and JSON.""" + + CHUNK_SET = "paper_chunk_set" + GLOBAL_SUMMARY = "paper_summary" def _stable_hash(value: object) -> str: @@ -42,12 +48,13 @@ def _paper_source_id(content_hash: str) -> UUID: def _paper_artifact_id( source_revision_id: UUID, - artifact_kind: PaperArtifactKind, + artifact_kind: PaperArtifactKind | str, producer_config_hash: str, ) -> UUID: + kind = PaperArtifactKind(artifact_kind) return uuid5( source_revision_id, - f"quantmind:{artifact_kind}:{producer_config_hash}", + f"quantmind:{kind.value}:{producer_config_hash}", ) @@ -352,7 +359,9 @@ class PaperChunkSet(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) id: UUID - artifact_kind: Literal["paper_chunk_set"] = "paper_chunk_set" + artifact_kind: Literal[PaperArtifactKind.CHUNK_SET] = ( + PaperArtifactKind.CHUNK_SET + ) schema_version: Literal["1.0"] = "1.0" source_revision_id: UUID producer: PaperChunkingConfig @@ -448,9 +457,16 @@ class PaperSummaryProducer(BaseModel): model: str prompt_version: str + orchestration: Literal["manager-research-agents-v1"] = ( + "manager-research-agents-v1" + ) input_chunk_set_id: UUID instructions_hash: str = Field(pattern=r"^[0-9a-f]{64}$") max_output_tokens: int = Field(gt=0) + max_research_calls: int = Field(default=12, ge=1) + max_research_concurrency: int = Field(default=2, ge=1) + research_max_turns: int = Field(default=4, ge=1) + research_max_output_tokens: int = Field(default=1_536, ge=1) def _paper_summary_content_hash( @@ -473,7 +489,9 @@ class PaperGlobalSummary(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) id: UUID - artifact_kind: Literal["paper_summary"] = "paper_summary" + artifact_kind: Literal[PaperArtifactKind.GLOBAL_SUMMARY] = ( + PaperArtifactKind.GLOBAL_SUMMARY + ) schema_version: Literal["1.0"] = "1.0" source_revision_id: UUID producer: PaperSummaryProducer diff --git a/quantmind/library/_types.py b/quantmind/library/_types.py index 9c3db23..5f902c1 100644 --- a/quantmind/library/_types.py +++ b/quantmind/library/_types.py @@ -6,7 +6,12 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator -from quantmind.knowledge import ArtifactLocator, Citation, SourceRef +from quantmind.knowledge import ( + ArtifactLocator, + Citation, + PaperArtifactKind, + SourceRef, +) class SemanticQuery(BaseModel): @@ -15,7 +20,7 @@ class SemanticQuery(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) text: str = Field(min_length=1) - artifact_kinds: list[str] | None = None + artifact_kinds: list[PaperArtifactKind] | None = None item_types: list[str] | None = None source_kinds: ( list[ diff --git a/scripts/verify_pdf_rag_e2e.py b/scripts/verify_pdf_rag_e2e.py index 1abe3dc..7ebe2b9 100644 --- a/scripts/verify_pdf_rag_e2e.py +++ b/scripts/verify_pdf_rag_e2e.py @@ -12,7 +12,11 @@ from quantmind.configs import PaperFlowCfg from quantmind.configs.paper import ArxivIdentifier from quantmind.flows import paper_flow -from quantmind.knowledge import PaperChunk, PaperGlobalSummary +from quantmind.knowledge import ( + PaperArtifactKind, + PaperChunk, + PaperGlobalSummary, +) from quantmind.library import LocalKnowledgeLibrary, SemanticQuery _ARXIV_ID = "1706.03762v7" @@ -27,14 +31,14 @@ async def _search_and_resolve( summary_hits = await library.search( SemanticQuery( text=_SUMMARY_QUERY, - artifact_kinds=["paper_summary"], + artifact_kinds=[PaperArtifactKind.GLOBAL_SUMMARY], top_k=3, ) ) chunk_hits = await library.search( SemanticQuery( text=_CHUNK_QUERY, - artifact_kinds=["paper_chunk_set"], + artifact_kinds=[PaperArtifactKind.CHUNK_SET], top_k=5, ) ) @@ -56,8 +60,11 @@ async def _run_vertical_slice() -> dict[str, Any]: timeout_seconds=240, max_summary_tool_calls=12, max_summary_concurrency=2, + max_summary_worker_turns=4, + max_summary_worker_output_tokens=1_536, max_summary_input_tokens=120_000, max_summary_output_tokens=4_096, + max_summary_total_output_tokens=20_000, min_summary_citations=3, min_summary_pages=2, ), @@ -115,6 +122,7 @@ async def _run_vertical_slice() -> dict[str, Any]: ), "chunk_count": len(result.chunk_set.chunks), "summary": result.global_summary.summary, + "summary_orchestration": (result.global_summary.producer.orchestration), "citation_count": len(result.global_summary.citations), "citation_pages": sorted( { @@ -199,6 +207,7 @@ def _passed(snapshot: dict[str, Any]) -> bool: and snapshot["chunk_asset_reference_count"] > 0 and snapshot["chunk_count"] > 0 and snapshot["summary"] + and snapshot["summary_orchestration"] == "manager-research-agents-v1" and snapshot["citation_count"] >= 3 and len(snapshot["citation_pages"]) >= 2 and snapshot["first_summary_scores"] @@ -232,6 +241,7 @@ async def main() -> int: f"text_pages={snapshot['text_page_count']} " f"chunks={snapshot['chunk_count']} assets={snapshot['asset_count']} " f"chunk_asset_refs={snapshot['chunk_asset_reference_count']} " + f"orchestration={snapshot['summary_orchestration']} " f"asset_pages={snapshot['screenshot_pages']} " f"citation_pages={snapshot['citation_pages']} " f"summary_scores={snapshot['second_summary_scores']} " diff --git a/tests/configs/test_paper.py b/tests/configs/test_paper.py index ea350b8..3247356 100644 --- a/tests/configs/test_paper.py +++ b/tests/configs/test_paper.py @@ -24,6 +24,9 @@ def test_defaults(self): self.assertEqual(cfg.chunk_size, 512) self.assertEqual(cfg.chunk_overlap, 64) self.assertEqual(cfg.max_summary_tool_calls, 12) + self.assertEqual(cfg.max_summary_worker_turns, 4) + self.assertEqual(cfg.max_summary_worker_output_tokens, 1_536) + self.assertEqual(cfg.max_summary_total_output_tokens, 20_000) self.assertEqual(cfg.min_summary_citations, 3) def test_invalid_overlap_or_coverage_bounds_are_rejected(self): @@ -31,6 +34,17 @@ def test_invalid_overlap_or_coverage_bounds_are_rejected(self): PaperFlowCfg(chunk_size=64, chunk_overlap=64) with self.assertRaises(ValidationError): PaperFlowCfg(min_summary_citations=1, min_summary_pages=2) + with self.assertRaises(ValidationError): + PaperFlowCfg( + max_summary_tool_calls=1, + max_summary_concurrency=2, + ) + with self.assertRaises(ValidationError): + PaperFlowCfg( + max_summary_output_tokens=4_096, + max_summary_worker_output_tokens=1_536, + max_summary_total_output_tokens=5_000, + ) class PaperInputDiscriminatedTests(unittest.TestCase): diff --git a/tests/flows/test_paper.py b/tests/flows/test_paper.py index 5219a87..6bf86b0 100644 --- a/tests/flows/test_paper.py +++ b/tests/flows/test_paper.py @@ -17,10 +17,17 @@ RawText, ) from quantmind.flows._paper_summary import ( + PaperResearchCitationDraft, + PaperResearchDraft, + PaperResearchFindingDraft, + PaperResearchRequest, PaperSummaryBudgetExceeded, PaperSummaryCitationDraft, PaperSummaryDraft, + _AgentsPaperSummaryProvider, _bounded_model_settings, + _bounded_research_model_settings, + _build_research_agent_tool, _SummaryBudget, ) from quantmind.flows.paper import ( @@ -289,47 +296,181 @@ def test_configured_citation_and_page_coverage_is_enforced(self) -> None: class SummaryBudgetTests(unittest.IsolatedAsyncioTestCase): + @staticmethod + def _research_draft( + result, + request: PaperResearchRequest, + ) -> PaperResearchDraft: + chunk = result.chunk_set.chunks[request.start] + page = chunk.source_spans[0].page_number + return PaperResearchDraft( + scope_summary=f"Reviewed chunk {request.start}.", + findings=( + PaperResearchFindingDraft( + kind="method", + claim=f"Finding from chunk {request.start}.", + citation=PaperResearchCitationDraft( + chunk_index=request.start, + page_number=page, + ), + ), + ), + ) + async def test_tool_call_and_input_budgets_are_enforced(self) -> None: result = build_paper_result() cfg = PaperFlowCfg( max_summary_tool_calls=1, + max_summary_concurrency=1, max_summary_input_tokens=10_000, ) - budget = _SummaryBudget(cfg, "manifest") + budget = _SummaryBudget( + cfg, + "manifest", + chunk_count=len(result.chunk_set.chunks), + ) + first = PaperResearchRequest(start=0, count=1, focus="method") - await budget.read(result.chunk_set, start=0, count=1) + async def first_operation(): + return self._research_draft(result, first) + + await budget.invoke_research( + result.source_revision, + result.chunk_set, + first, + first_operation, + ) + second = PaperResearchRequest(start=1, count=1, focus="result") with self.assertRaisesRegex( PaperSummaryBudgetExceeded, "tool_calls", ): - await budget.read(result.chunk_set, start=1, count=1) + await budget.invoke_research( + result.source_revision, + result.chunk_set, + second, + first_operation, + ) - async def test_concurrent_reads_remain_bounded_and_valid(self) -> None: + async def test_concurrent_subagents_cover_every_chunk(self) -> None: result = build_paper_result() cfg = PaperFlowCfg( max_summary_tool_calls=3, - max_summary_concurrency=1, + max_summary_concurrency=2, max_summary_input_tokens=10_000, ) - budget = _SummaryBudget(cfg, "manifest") + budget = _SummaryBudget( + cfg, + "manifest", + chunk_count=len(result.chunk_set.chunks), + ) + + async def run(index: int) -> str: + request = PaperResearchRequest( + start=index, + count=1, + focus=f"chunk {index}", + ) + + async def operation(): + await asyncio.sleep(0) + return self._research_draft(result, request) + + return await budget.invoke_research( + result.source_revision, + result.chunk_set, + request, + operation, + ) values = await asyncio.gather( - budget.read(result.chunk_set, start=0, count=1), - budget.read(result.chunk_set, start=1, count=1), - budget.read(result.chunk_set, start=2, count=1), + *(run(index) for index in range(len(result.chunk_set.chunks))) ) self.assertEqual(len(values), 3) + self.assertEqual(budget.research_calls, 3) + self.assertEqual(budget.covered_chunks, frozenset({0, 1, 2})) + budget.require_complete_coverage(len(result.chunk_set.chunks)) + + async def test_missing_subagent_coverage_is_rejected(self) -> None: + result = build_paper_result() + cfg = PaperFlowCfg(max_summary_input_tokens=10_000) + budget = _SummaryBudget( + cfg, + "manifest", + chunk_count=len(result.chunk_set.chunks), + ) + + with self.assertRaisesRegex( + PaperSummaryBudgetExceeded, + "did not cover every chunk", + ): + budget.require_complete_coverage(len(result.chunk_set.chunks)) + + async def test_coordinator_without_delegation_is_rejected(self) -> None: + result = build_paper_result() + provider = _AgentsPaperSummaryProvider() + draft = PaperSummaryDraft( + summary="A summary that skipped its research workers.", + citations=( + PaperSummaryCitationDraft(chunk_index=0, page_number=1), + ), + ) + cfg = PaperFlowCfg( + min_summary_citations=1, + min_summary_pages=1, + ) + + with patch( + "quantmind.flows._paper_summary.run_with_observability", + new=AsyncMock(return_value=draft), + ): + with self.assertRaisesRegex( + PaperSummaryBudgetExceeded, + "did not cover every chunk", + ): + await provider.summarize( + result.source_revision, + result.chunk_set, + cfg=cfg, + ) + + def test_researcher_is_an_agents_sdk_agent_tool(self) -> None: + result = build_paper_result() + cfg = PaperFlowCfg() + budget = _SummaryBudget( + cfg, + "manifest", + chunk_count=len(result.chunk_set.chunks), + ) + + tool = _build_research_agent_tool( + result.source_revision, + result.chunk_set, + cfg, + budget, + ) + + self.assertEqual(tool.name, "research_chunk_group") + self.assertTrue(tool._is_agent_tool) + self.assertEqual( + set(tool.params_json_schema["properties"]), + {"start", "count", "focus"}, + ) def test_model_output_tokens_are_capped(self) -> None: cfg = PaperFlowCfg( max_summary_output_tokens=256, + max_summary_worker_output_tokens=128, model_settings=ModelSettings(max_tokens=1024), ) settings = _bounded_model_settings(cfg) + research_settings = _bounded_research_model_settings(cfg) self.assertEqual(settings.max_tokens, 256) self.assertTrue(settings.parallel_tool_calls) + self.assertEqual(research_settings.max_tokens, 128) + self.assertFalse(research_settings.parallel_tool_calls) if __name__ == "__main__": diff --git a/tests/library/test_paper.py b/tests/library/test_paper.py index 4c0441a..183b8df 100644 --- a/tests/library/test_paper.py +++ b/tests/library/test_paper.py @@ -7,7 +7,11 @@ from collections.abc import Sequence from pathlib import Path -from quantmind.knowledge import PaperChunk, PaperGlobalSummary +from quantmind.knowledge import ( + PaperArtifactKind, + PaperChunk, + PaperGlobalSummary, +) from quantmind.library import LocalKnowledgeLibrary, SemanticQuery from tests.paper_helpers import build_paper_result @@ -181,7 +185,7 @@ async def test_reopen_round_trip_reuses_vectors_and_resolves_hits( summary_hits = await library.search( SemanticQuery( text=summary_query, - artifact_kinds=["paper_summary"], + artifact_kinds=[PaperArtifactKind.GLOBAL_SUMMARY], top_k=3, ) ) @@ -197,7 +201,7 @@ async def test_reopen_round_trip_reuses_vectors_and_resolves_hits( chunk_hits = await library.search( SemanticQuery( text=chunk_query, - artifact_kinds=["paper_chunk_set"], + artifact_kinds=[PaperArtifactKind.CHUNK_SET], top_k=5, ) ) diff --git a/tests/library/test_types.py b/tests/library/test_types.py index cb6ddc8..efba2be 100644 --- a/tests/library/test_types.py +++ b/tests/library/test_types.py @@ -4,6 +4,7 @@ from pydantic import ValidationError import quantmind.library as library +from quantmind.knowledge import PaperArtifactKind from quantmind.library import SemanticQuery @@ -34,6 +35,22 @@ def test_query_rejects_non_positive_top_k(self): with self.assertRaises(ValidationError): SemanticQuery(text="rates", top_k=0) + def test_query_parses_typed_paper_artifact_kinds(self): + query = SemanticQuery( + text="attention", + artifact_kinds=["paper_summary"], + ) + + self.assertEqual( + query.artifact_kinds, + [PaperArtifactKind.GLOBAL_SUMMARY], + ) + with self.assertRaises(ValidationError): + SemanticQuery( + text="attention", + artifact_kinds=["unknown_artifact"], + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_verify_pdf_rag_e2e.py b/tests/test_verify_pdf_rag_e2e.py index 5ac4411..92f0720 100644 --- a/tests/test_verify_pdf_rag_e2e.py +++ b/tests/test_verify_pdf_rag_e2e.py @@ -20,6 +20,7 @@ def _snapshot(**overrides): "convolution, uses multi-head attention, and improves translation " "with greater training efficiency." ), + "summary_orchestration": "manager-research-agents-v1", "citation_count": 3, "citation_pages": [2, 4], "first_summary_scores": [0.9], From f56c9decf2595dcd6de2354f9fda6aba49dd34a3 Mon Sep 17 00:00:00 2001 From: pkuwkl Date: Sun, 19 Jul 2026 01:56:51 +0800 Subject: [PATCH 4/4] refactor(paper): identity in knowledge constructors; deterministic summary map-reduce Rework Paper Flow V1 so identity and orchestration sit at the right altitude, without changing the source-first data model, stored schema, or the public paper_flow(input, *, cfg) -> PaperFlowResult signature. Identity vs translation: - Add PaperSourceRevision.from_parsed / PaperChunkSet.from_parsed_chunks / PaperGlobalSummary.from_draft smart constructors that mint every ID, content and producer hash, and resolve/validate citations internally. - flows/paper.py now imports zero private _paper_* helpers and computes no ID. It only fetches, parses, reads asset bytes, and maps preprocess/rag values to knowledge-native inputs. `knowledge` stays an import-linter leaf. Summarization: - Replace the manager/worker coordinator + Agent.as_tool() + hand-rolled concurrency-safe _SummaryBudget with a deterministic map-reduce: code tiles the chunk set (coverage guaranteed by construction), fans out one research agent per group via asyncio.gather + Semaphore, and runs one reducer. - Bounds are delegated to the SDK (ModelSettings.max_tokens, output_type) and asyncio (wait_for). PaperSummaryProducer identity becomes map-reduce-v1 with research_group_size; drop the six budget knobs and two validator clauses from PaperFlowCfg, add summary_research_group_size / summary_concurrency. Tests and docs: - Rebuild tests/paper_helpers.build_paper_result on the new constructors and rewrite the summary tests as map-reduce tests; update the e2e verifier and the paper example. - Document the identity-vs-translation and agentic-vs-deterministic-map-reduce principles in contexts/design/operations/orchestration.md and refresh the paper flow/knowledge design pages. verify.sh: ruff, basedpyright (0 errors), 7 import contracts, 334 tests, 84% cov. Co-Authored-By: Claude Opus 4.8 (1M context) --- contexts/design/README.md | 1 + contexts/design/flow/paper.md | 38 +- contexts/design/knowledge/paper.md | 8 +- contexts/design/operations/orchestration.md | 75 ++++ quantmind/configs/paper.py | 35 +- quantmind/flows/__init__.py | 7 +- quantmind/flows/_paper_summary.py | 466 ++++++-------------- quantmind/flows/paper.py | 416 +++++------------ quantmind/knowledge/__init__.py | 12 + quantmind/knowledge/paper.py | 338 +++++++++++++- scripts/verify_pdf_rag_e2e.py | 10 +- tests/configs/test_paper.py | 19 +- tests/flows/test_paper.py | 282 ++++++------ tests/paper_helpers.py | 213 +++------ tests/test_verify_pdf_rag_e2e.py | 2 +- 15 files changed, 905 insertions(+), 1017 deletions(-) create mode 100644 contexts/design/operations/orchestration.md diff --git a/contexts/design/README.md b/contexts/design/README.md index 188079a..d055829 100644 --- a/contexts/design/README.md +++ b/contexts/design/README.md @@ -31,6 +31,7 @@ implementation must preserve. | RAG | [Page-aware document chunking and retrieval](rag/document.md) | | Library | [Local knowledge storage and meaning-based search](library/local.md) | | Operations | [Public operation naming](operations/naming.md) | +| Operations | [Orchestration and construction altitude](operations/orchestration.md) | ## Organization Rules diff --git a/contexts/design/flow/paper.md b/contexts/design/flow/paper.md index 4bca46c..63cbda8 100644 --- a/contexts/design/flow/paper.md +++ b/contexts/design/flow/paper.md @@ -53,12 +53,14 @@ The operation has a strict order: 3. Build and validate `PaperSourceRevision`, including content-addressed asset references and blobs. 4. Chunk each page with LlamaIndex while retaining page and character spans. 5. Build and validate `PaperChunkSet` with code-owned IDs and membership. -6. Give a coordinator agent the bounded chunk manifest and a research-agent tool. -7. Require the coordinator to delegate complete chunk coverage to one or more bounded research subagents, with independent ranges eligible for parallel execution. -8. Let the coordinator synthesize typed worker reports and make bounded overlapping follow-up calls when needed. +6. Tile the chunk set into fixed-size research groups in code, so every chunk is covered exactly once. +7. Fan out one bounded research agent per group (bounded concurrency), each returning typed findings for its own range only. +8. Run one reducer agent over the collected findings to synthesize the summary draft. 9. Resolve model-returned chunk/page coordinates into canonical citations in code. 10. Build and validate `PaperGlobalSummary` and the cross-artifact `PaperFlowResult`. +Steps 3, 5, 9, and 10 mint no IDs in the flow: the flow calls the knowledge-layer smart constructors `PaperSourceRevision.from_parsed`, `PaperChunkSet.from_parsed_chunks`, and `PaperGlobalSummary.from_draft`, which own every ID, content/producer hash, and citation resolution. The flow only fetches, parses, reads asset bytes, and maps those path-based artifacts into knowledge-native inputs. See [orchestration principles](../operations/orchestration.md). + A summarization failure occurs after source and chunks exist in memory, but `paper_flow` returns no partial success value. Persistence is a separate explicit operation. ## Source Revision @@ -92,27 +94,25 @@ Code accepts the draft only when: - every chunk index exists; - every cited page is present in that chunk's source spans; - every supplied quote occurs verbatim in the cited chunk; -- citation count meets `summary_min_citations`; -- distinct cited-page count meets `summary_min_pages`. +- citation count meets `min_summary_citations`; +- distinct cited-page count meets `min_summary_pages`. -Code then creates `PaperCitation` values and a `PaperGlobalSummary`. Its producer identity includes model, prompt version, input chunk-set ID, instructions hash, and maximum output tokens. Its lineage contains the exact input chunk-set locator. Summary citations must resolve through that chunk set to source pages. +`PaperGlobalSummary.from_draft` performs this resolution and coverage check and then mints the artifact; the flow supplies only the draft and the coverage policy. Its producer identity includes model, prompt version, input chunk-set ID, instructions hash, maximum output tokens, and research group size. Its lineage contains the exact input chunk-set locator. Summary citations must resolve through that chunk set to source pages. ## Bounded Model Calls -`PaperFlowCfg` makes coordinator and nested-research bounds explicit: - -- `max_summary_tool_calls` caps research-subagent invocations, including focused follow-ups; -- `max_summary_concurrency` bounds simultaneous research-subagent runs; -- `max_summary_worker_turns` caps each nested agent run; -- `max_summary_worker_output_tokens` caps each structured worker report; -- `max_summary_input_tokens` caps the manifest, worker chunk inputs, and worker reports returned to the coordinator; -- `max_summary_output_tokens` caps the coordinator's final structured response; -- `max_summary_total_output_tokens` caps worker reports and final output together; -- `max_turns` caps coordinator turns and defaults to 16; -- `timeout_seconds` bounds the complete summarization operation; +Summarization is a deterministic map-reduce, not an autonomous coordinator. Code — not a model — decides the decomposition: it tiles the chunk set into `summary_research_group_size` groups and runs one research agent per group, so complete chunk coverage is guaranteed by construction and never needs to be reconciled afterward. + +`PaperFlowCfg` exposes only structural bounds: + +- `summary_research_group_size` sets how many consecutive chunks each research agent receives; +- `summary_concurrency` bounds simultaneous research-agent runs; +- `max_summary_output_tokens` caps each agent's output through `ModelSettings.max_tokens`; +- `timeout_seconds` bounds the reducer call via `asyncio.wait_for`; +- `min_summary_citations` and `min_summary_pages` set the accepted-summary coverage policy; - `summary_prompt_version` and `summary_instructions` version the semantic producer. -The coordinator uses the Agents SDK manager-style pattern: `paper_chunk_researcher` is exposed through `Agent.as_tool()`, so the coordinator retains responsibility for the final summary while each specialist receives only its requested range. The manifest provides code-generated groups of at most eight consecutive chunks. Independent calls may run in parallel, and the coordinator may request bounded overlapping follow-up research. Code rejects missing full-chunk coverage, out-of-scope worker citations, invalid pages or quotes, and any call, concurrency, token, or runtime excess. +Bounding is delegated to the Agents SDK (per-agent `max_tokens`, structured `output_type`) and to `asyncio` (a `Semaphore` for concurrency, `wait_for` for the timeout). There is no hand-rolled token accountant, no coordinator agent, and no `Agent.as_tool()` research tool. The removed manager/worker design paid for an autonomous coordinator and then suppressed its nondeterminism with a concurrency-safe budget — exactly the runtime the SDK already provides. [Orchestration principles](../operations/orchestration.md) records when the agentic pattern is warranted instead. ## Failure Semantics @@ -121,7 +121,7 @@ The coordinator uses the Agents SDK manager-style pattern: `paper_chunk_research - Fetching, parsing, or missing parser assets raise their source error and produce no result. - Empty chunk output is invalid. - Invalid or insufficient summary citations raise `PaperCitationValidationError`. -- Missing worker coverage, invalid worker evidence, and call, token, output, concurrency, or timeout violations fail the summary operation. +- A research finding that cites outside its assigned group is rejected in code; a reducer timeout raises `PaperSummaryError`. - Any canonical identity, content hash, membership, lineage, or cross-artifact mismatch fails Pydantic validation. No failure is converted into a partially valid `PaperFlowResult`. Callers may retry with the same source and producer settings; stable IDs make successful repeated runs idempotent. diff --git a/contexts/design/knowledge/paper.md b/contexts/design/knowledge/paper.md index c90be9c..e1182a9 100644 --- a/contexts/design/knowledge/paper.md +++ b/contexts/design/knowledge/paper.md @@ -62,11 +62,11 @@ Every chunk span is also checked against that manifest: its page must exist, its - model identity; - prompt version; -- manager/research-agent orchestration version; +- map-reduce orchestration version; - exact input chunk-set ID; -- coordinator/research instructions hash; -- coordinator output limit; -- research call, concurrency, turn, and per-worker output limits. +- reducer/research instructions hash; +- per-agent output limit; +- research group size. Changing any producer field creates a distinct artifact ID. Multiple chunk sets and summaries may coexist for one source revision. Loading a complete `PaperFlowResult` without explicit artifact IDs is allowed only when one unambiguous linked pair exists. diff --git a/contexts/design/operations/orchestration.md b/contexts/design/operations/orchestration.md new file mode 100644 index 0000000..e696314 --- /dev/null +++ b/contexts/design/operations/orchestration.md @@ -0,0 +1,75 @@ +# Orchestration and Construction Altitude + +## Quick Summary + +- **Purpose**: Decide where construction/identity logic lives, and when a step should be an autonomous agent versus deterministic code. +- **Read when**: Adding or refactoring any flow that builds knowledge models from preprocess/rag values, or that calls a model more than once. +- **Status**: Derived from the Paper Flow V1 refactor. Applies to every flow. +- **Core rule**: Identity belongs in the knowledge layer; translation belongs in the flow. Use an autonomous agent only when the decomposition must be decided at runtime. + +## Contents + +- [Two Principles](#two-principles) +- [Principle 1: Identity vs Translation Altitude](#principle-1-identity-vs-translation-altitude) +- [Principle 2: Agentic vs Deterministic Orchestration](#principle-2-agentic-vs-deterministic-orchestration) +- [Worked Example: Paper Flow V1](#worked-example-paper-flow-v1) +- [Checklist](#checklist) + +## Two Principles + +Both principles came from one observation: a flow file had become large and was reaching into another module's private helpers. The cause was not bad code locally — it was two responsibilities placed at the wrong altitude. + +1. **Identity vs translation altitude** — who mints IDs and hashes, and who maps one type system onto another. +2. **Agentic vs deterministic orchestration** — whether a model decides *how* the work is decomposed, or code does. + +## Principle 1: Identity vs Translation Altitude + +"Construction" is two separable concerns that belong in different layers: + +| Concern | Depends on | Home | +|---|---|---| +| **Identity** — mint IDs, content/producer hashes, resolve and validate citations | knowledge types only | **knowledge layer** (smart constructors) | +| **Translation** — map `preprocess`/`rag` value objects to knowledge inputs, read asset bytes (IO) | preprocess/rag **and** knowledge | **flow layer** (the only layer that imports both) | + +The `knowledge` package is a leaf: an import-linter contract forbids it from importing `preprocess`, `rag`, `flows`, or `library`. So a knowledge model **cannot** name a `ParsedDocument` or `ParsedChunk`, and translation from those types genuinely has to live in the flow. That constraint is real and correct. + +But it does **not** justify the flow computing identity. The rule: + +- **Knowledge models expose smart constructors** (`from_parsed`, `from_parsed_chunks`, `from_draft`) that take knowledge-native inputs and mint every ID and hash internally. A model that validates its own identity in a `model_validator` must also be able to *construct* it; do not leave construction as free private functions for callers to import. +- **The flow adapts and calls.** It maps the ephemeral, path-based preprocess/rag types into knowledge-native inputs (plain value objects carrying geometry, text, and pre-read bytes) and passes them to the constructors. It imports **zero** private `_paper_*_id` / `_*_hash` helpers and computes no ID itself. + +**Smell test.** A module importing another module's underscore-prefixed helpers across a package boundary is almost always identity logic living at the wrong altitude. Move the logic behind a public constructor on the model, not the import into the caller. + +## Principle 2: Agentic vs Deterministic Orchestration + +The Agents SDK's manager/worker pattern (a coordinator agent that calls other agents via `Agent.as_tool()` and decides how many calls to make) earns its cost only when **the decomposition is decided at runtime** — when the agent must look at intermediate results and choose what to do next. + +Ask one question: + +> Is the decomposition something code can compute now, or must a model decide it after seeing intermediate results? + +- **Code can compute it now** → deterministic map-reduce. Tile the work in code, fan out stateless agents with `asyncio.gather` + a `Semaphore`, run one reducer. Coverage is guaranteed by construction; bounds come from the SDK (`ModelSettings.max_tokens`, `output_type`) and `asyncio` (`wait_for`). No coordinator, no coverage police, no token accountant. +- **A model must decide at runtime** → agentic coordinator with `as_tool`. And in that case do *not* also bolt a deterministic "must cover everything" check on top; that check contradicts the autonomy you just paid for. + +**The incoherent middle to avoid.** Running an autonomous coordinator *and* then enforcing a deterministic budget/coverage accountant pays for nondeterminism and immediately suppresses it — both costs, neither benefit. If you find yourself writing a concurrency-safe accountant to make an agent behave deterministically, the task wanted deterministic orchestration from the start. + +Using the SDK does **not** force the agentic pattern. Single-agent `Runner.run(agent, output_type=...)` gives structured output, tracing, retries, and guardrails inside a fully deterministic map-reduce — you do not drop to a raw completion API to get determinism. + +**Do not build a framework for it.** The fan-out primitive is a function, not a base class: `asyncio.gather` over code-computed items with a `Semaphore`. Keep it inline until a second flow needs it; only then extract a small `map_reduce(items, map_fn, reduce_fn, *, concurrency)` helper. A `BaseFlow`/`ParallelWorkflow` base class is the framework this project explicitly avoids. + +## Worked Example: Paper Flow V1 + +- **Identity** lives on `PaperSourceRevision.from_parsed`, `PaperChunkSet.from_parsed_chunks`, and `PaperGlobalSummary.from_draft`. `quantmind/flows/paper.py` imports no `_paper_*` helper and mints no ID. +- **Translation** stays in the flow: `_fetch_paper_source` (IO), `_adapt_pages` (reads screenshot/image bytes, maps to `PaperPageInput`), `_adapt_chunks` (maps `ParsedChunk` to `PaperChunkInput`). +- **Summarization** is a deterministic map-reduce (`_chunk_groups` tiles the chunk set; one research agent per group; one reducer). The earlier manager/worker coordinator plus its `_SummaryBudget` concurrency accountant were removed — see [Bounded Model Calls](../flow/paper.md#bounded-model-calls). + +## Checklist + +When adding a flow that turns preprocess/rag values into knowledge models: + +- [ ] Does the flow import any `_`-prefixed helper from `knowledge`? If so, move that logic into a smart constructor. +- [ ] Does each knowledge artifact have a `from_*` constructor that mints its own IDs/hashes? +- [ ] Is the flow's remaining work limited to fetch, parse, read bytes, and structural mapping? +- [ ] For multi-call model steps: can code compute the decomposition? If yes, use map-reduce, not a coordinator. +- [ ] Are you enforcing coverage/budget on an autonomous agent? If yes, prefer deterministic decomposition instead. +- [ ] Is the fan-out a plain function (`gather` + `Semaphore`), not a base class? diff --git a/quantmind/configs/paper.py b/quantmind/configs/paper.py index 2df83be..ed6e842 100644 --- a/quantmind/configs/paper.py +++ b/quantmind/configs/paper.py @@ -55,21 +55,25 @@ class DoiIdentifier(BaseInput): class PaperFlowCfg(BaseFlowCfg): - """Chunking, summarization, and hard-budget controls for ``paper_flow``.""" + """Chunking and summarization controls for ``paper_flow``. + + Summarization is a deterministic map-reduce: code tiles the chunk set into + ``summary_research_group_size`` groups (so coverage is guaranteed by + construction), fans out one research agent per group with + ``summary_concurrency`` parallelism, then runs one reducer. Per-agent output + is bounded by ``max_summary_output_tokens`` through ``ModelSettings``; there + is no hand-rolled token accountant. + """ model: str = "gpt-4o-mini" max_turns: int = Field(default=16, ge=1) chunk_size: int = Field(default=512, gt=0) chunk_overlap: int = Field(default=64, ge=0) - summary_prompt_version: str = "paper-summary-v2" + summary_prompt_version: str = "paper-summary-v3" summary_instructions: str | None = None - max_summary_tool_calls: int = Field(default=12, ge=1) - max_summary_concurrency: int = Field(default=2, ge=1) - max_summary_worker_turns: int = Field(default=4, ge=1) - max_summary_worker_output_tokens: int = Field(default=1_536, ge=1) - max_summary_input_tokens: int = Field(default=120_000, ge=1) - max_summary_output_tokens: int = Field(default=4_096, ge=1) - max_summary_total_output_tokens: int = Field(default=20_000, ge=1) + summary_research_group_size: int = Field(default=8, ge=1) + summary_concurrency: int = Field(default=4, ge=1) + max_summary_output_tokens: int = Field(default=4_096, gt=0) min_summary_citations: int = Field(default=3, ge=1) min_summary_pages: int = Field(default=2, ge=1) @@ -81,17 +85,4 @@ def _validate_paper_bounds(self) -> "PaperFlowCfg": raise ValueError( "min_summary_pages cannot exceed min_summary_citations" ) - if self.max_summary_concurrency > self.max_summary_tool_calls: - raise ValueError( - "max_summary_concurrency cannot exceed max_summary_tool_calls" - ) - minimum_output_budget = ( - self.max_summary_output_tokens - + self.max_summary_worker_output_tokens - ) - if self.max_summary_total_output_tokens < minimum_output_budget: - raise ValueError( - "max_summary_total_output_tokens must reserve one worker " - "report and the final summary" - ) return self diff --git a/quantmind/flows/__init__.py b/quantmind/flows/__init__.py index 8889955..dbd53be 100644 --- a/quantmind/flows/__init__.py +++ b/quantmind/flows/__init__.py @@ -13,11 +13,8 @@ from quantmind.flows.batch import BatchResult, batch_run from quantmind.flows.news import collect_news -from quantmind.flows.paper import ( - PaperCitationValidationError, - UnsupportedContentTypeError, - paper_flow, -) +from quantmind.flows.paper import UnsupportedContentTypeError, paper_flow +from quantmind.knowledge import PaperCitationValidationError __all__ = [ "BatchResult", diff --git a/quantmind/flows/_paper_summary.py b/quantmind/flows/_paper_summary.py index f5c7f2b..c7738bd 100644 --- a/quantmind/flows/_paper_summary.py +++ b/quantmind/flows/_paper_summary.py @@ -1,75 +1,58 @@ -"""Bounded multi-agent synthesis for one paper chunk-set artifact.""" +"""Deterministic map-reduce summarization for one paper chunk-set artifact. + +The chunk set is tiled into fixed-size groups by code, so every chunk is +covered exactly once. There is no coordinator agent deciding *how* to +decompose the work, and therefore no need to police coverage or reconcile a +shared token budget after the fact. One research agent runs per group (bounded +fan-out), and one reducer agent synthesizes the group reports into a single +cited global summary. + +Bounding is delegated to the Agents SDK: per-agent output is capped through +``ModelSettings.max_tokens`` and the reducer is wrapped in ``asyncio.wait_for``. +The removed manager/worker design instead ran an autonomous coordinator and +then fought its nondeterminism with a hand-rolled concurrency-safe accountant; +that is precisely the runtime the SDK already provides. +""" import asyncio import hashlib import json -from collections.abc import Awaitable, Callable -from dataclasses import replace +from collections.abc import Sequence +from dataclasses import dataclass, replace from typing import Any, Literal, Protocol -from agents import Agent, FunctionTool, ModelSettings +from agents import Agent, ModelSettings from pydantic import BaseModel, ConfigDict, Field, field_validator from quantmind.configs import PaperFlowCfg from quantmind.flows._runner import run_with_observability from quantmind.knowledge import PaperChunkSet, PaperSourceRevision -_MAX_RESEARCH_GROUP_SIZE = 8 -_ORCHESTRATION_VERSION = "manager-research-agents-v1" +_ORCHESTRATION_VERSION = "map-reduce-v1" _SUMMARY_INSTRUCTIONS = """\ -Act as the paper-summary coordinator. Delegate every suggested research group -to `research_chunk_group` before writing the final summary. You may call -independent groups in parallel and may make bounded overlapping follow-up calls -when a worker report exposes an ambiguity. Synthesize the worker reports into -one accurate global summary covering the central contribution, architecture or -methodology, principal results, and important limitations. Return only summary -prose plus citations. Each citation uses the zero-based chunk index and a -physical page owned by that chunk. Never invent IDs, source metadata, storage -links, pages, or quotations. Set a final citation quote to null unless it is an -exact contiguous substring copied character-for-character from the chunk. +Act as the paper-summary reducer. You are given research reports that together +cover every chunk of the paper. Synthesize them into one accurate global +summary covering the central contribution, architecture or methodology, +principal results, and important limitations. Return only summary prose plus +citations. Each citation uses the zero-based chunk index and a physical page +reported by a research finding. Set a citation quote to null unless it is +copied verbatim from a research finding's quote for that same chunk. Never +invent chunk indices, pages, or quotations. """ _RESEARCH_INSTRUCTIONS = """\ Act as a bounded paper research specialist. Analyze only the supplied chunk -range and requested focus. Return a concise scope summary plus structured -findings. Classify every finding as context, contribution, method, result, or -limitation and support it with a chunk index and physical page. Do not infer -canonical IDs, source metadata, or claims that are not supported by the -supplied chunks. +range. Return a concise scope summary plus structured findings. Classify every +finding as context, contribution, method, result, or limitation and support it +with a chunk index and physical page drawn from the supplied chunks. When you +quote, copy an exact contiguous substring from the chunk text; otherwise leave +the quote null. Do not infer canonical IDs or claims unsupported by the chunks. """ -class PaperSummaryBudgetExceeded(RuntimeError): - """The summary exceeded a configured runtime or usage boundary.""" - - -class PaperSummaryCitationDraft(BaseModel): - """Model-owned citation coordinates before code resolves canonical IDs.""" - - model_config = ConfigDict(extra="forbid", frozen=True) - - chunk_index: int = Field(ge=0) - page_number: int = Field(ge=1) - quote: str | None = Field(default=None, max_length=500) - - -class PaperResearchRequest(BaseModel): - """Structured coordinator request for one bounded research subagent.""" - - model_config = ConfigDict(extra="forbid", frozen=True) - - start: int = Field(ge=0) - count: int = Field(ge=1, le=_MAX_RESEARCH_GROUP_SIZE) - focus: str = Field(min_length=1, max_length=500) - - @field_validator("focus") - @classmethod - def _focus_is_not_blank(cls, value: str) -> str: - stripped = value.strip() - if not stripped: - raise ValueError("paper research focus must not be blank") - return stripped +class PaperSummaryError(RuntimeError): + """The summary exceeded a configured runtime boundary.""" class PaperResearchCitationDraft(BaseModel): @@ -95,6 +78,7 @@ class PaperResearchFindingDraft(BaseModel): ] claim: str = Field(min_length=1) citation: PaperResearchCitationDraft + quote: str | None = Field(default=None, max_length=500) @field_validator("claim") @classmethod @@ -122,6 +106,16 @@ def _scope_summary_is_not_blank(cls, value: str) -> str: return stripped +class PaperSummaryCitationDraft(BaseModel): + """Model-owned citation coordinates before code resolves canonical IDs.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + chunk_index: int = Field(ge=0) + page_number: int = Field(ge=1) + quote: str | None = Field(default=None, max_length=500) + + class PaperSummaryDraft(BaseModel): """Limited model output containing prose and non-canonical coordinates.""" @@ -153,27 +147,34 @@ async def summarize( ... -def _estimated_tokens(text: str) -> int: - return max(1, (len(text) + 3) // 4) +@dataclass(frozen=True) +class _ChunkGroup: + """A contiguous, code-chosen range of chunk positions.""" + + start: int + count: int + + +def _chunk_groups(chunk_count: int, size: int) -> list[_ChunkGroup]: + """Tile ``[0, chunk_count)`` so every chunk lands in exactly one group.""" + return [ + _ChunkGroup(start=start, count=min(size, chunk_count - start)) + for start in range(0, chunk_count, size) + ] def _research_payload( source: PaperSourceRevision, chunk_set: PaperChunkSet, - request: PaperResearchRequest, + group: _ChunkGroup, ) -> str: - if request.start >= len(chunk_set.chunks): - raise ValueError("research start is outside the chunk-set manifest") - selected = chunk_set.chunks[request.start : request.start + request.count] - if len(selected) != request.count: - raise ValueError("research range exceeds the chunk-set manifest") + selected = chunk_set.chunks[group.start : group.start + group.count] return json.dumps( { "title": source.title, "authors": source.authors, - "focus": request.focus, - "start": request.start, - "count": request.count, + "start": group.start, + "count": group.count, "chunks": [ { "chunk_index": chunk.position, @@ -189,211 +190,52 @@ def _research_payload( ) -def _coerce_research_draft(value: Any) -> PaperResearchDraft: - if isinstance(value, str): - return PaperResearchDraft.model_validate_json(value) - return PaperResearchDraft.model_validate(value) - - def _validate_research_draft( chunk_set: PaperChunkSet, - request: PaperResearchRequest, + group: _ChunkGroup, draft: PaperResearchDraft, ) -> None: - allowed = set(range(request.start, request.start + request.count)) + allowed = set(range(group.start, group.start + group.count)) for finding in draft.findings: citation = finding.citation if citation.chunk_index not in allowed: - raise ValueError("research finding cites a chunk outside its scope") + raise ValueError("research finding cites a chunk outside its group") chunk = chunk_set.chunks[citation.chunk_index] pages = {span.page_number for span in chunk.source_spans} if citation.page_number not in pages: raise ValueError("research finding cites a page outside its chunk") - - -class _SummaryBudget: - """Concurrency-safe accounting for nested research-agent calls.""" - - def __init__( - self, - cfg: PaperFlowCfg, - initial_text: str, - *, - chunk_count: int, - ) -> None: - initial_tokens = _estimated_tokens(initial_text) - if initial_tokens > cfg.max_summary_input_tokens: - raise PaperSummaryBudgetExceeded( - "paper summary manifest exceeds max_summary_input_tokens" - ) - required_calls = ( - chunk_count + _MAX_RESEARCH_GROUP_SIZE - 1 - ) // _MAX_RESEARCH_GROUP_SIZE - if required_calls > cfg.max_summary_tool_calls: - raise PaperSummaryBudgetExceeded( - "max_summary_tool_calls cannot cover every paper chunk" - ) - self._max_calls = cfg.max_summary_tool_calls - self._max_input_tokens = cfg.max_summary_input_tokens - self._max_worker_output_tokens = cfg.max_summary_worker_output_tokens - self._max_final_output_tokens = cfg.max_summary_output_tokens - self._max_total_output_tokens = cfg.max_summary_total_output_tokens - self._calls = 0 - self._input_tokens = initial_tokens - self._output_tokens = 0 - self._reserved_input_tokens = 0 - self._reserved_output_tokens = 0 - self._covered_chunks: set[int] = set() - self._lock = asyncio.Lock() - self._semaphore = asyncio.Semaphore(cfg.max_summary_concurrency) - - @property - def research_calls(self) -> int: - """Return the number of research calls reserved so far.""" - return self._calls - - @property - def covered_chunks(self) -> frozenset[int]: - """Return chunk positions covered by successful worker reports.""" - return frozenset(self._covered_chunks) - - async def invoke_research( - self, - source: PaperSourceRevision, - chunk_set: PaperChunkSet, - request: PaperResearchRequest, - operation: Callable[[], Awaitable[Any]], - ) -> str: - """Run one agent tool call while reserving all shared budgets.""" - payload = _research_payload(source, chunk_set, request) - payload_tokens = _estimated_tokens(payload) - input_reservation = payload_tokens + self._max_worker_output_tokens - await self._semaphore.acquire() - reserved = False - try: - async with self._lock: - if self._calls >= self._max_calls: - raise PaperSummaryBudgetExceeded( - "paper summary exceeded max_summary_tool_calls" - ) - projected_input = ( - self._input_tokens - + self._reserved_input_tokens - + input_reservation - ) - if projected_input > self._max_input_tokens: - raise PaperSummaryBudgetExceeded( - "paper summary exceeded max_summary_input_tokens" - ) - projected_output = ( - self._output_tokens - + self._reserved_output_tokens - + self._max_worker_output_tokens - + self._max_final_output_tokens - ) - if projected_output > self._max_total_output_tokens: - raise PaperSummaryBudgetExceeded( - "paper summary exceeded max_summary_total_output_tokens" - ) - self._calls += 1 - self._reserved_input_tokens += input_reservation - self._reserved_output_tokens += self._max_worker_output_tokens - reserved = True - - raw_output = await operation() - draft = _coerce_research_draft(raw_output) - _validate_research_draft(chunk_set, request, draft) - serialized = draft.model_dump_json() - output_tokens = _estimated_tokens(serialized) - if output_tokens > self._max_worker_output_tokens: - raise PaperSummaryBudgetExceeded( - "paper research worker exceeded " - "max_summary_worker_output_tokens" - ) - - async with self._lock: - self._reserved_input_tokens -= input_reservation - self._reserved_output_tokens -= self._max_worker_output_tokens - self._input_tokens += payload_tokens + output_tokens - self._output_tokens += output_tokens - self._covered_chunks.update( - range(request.start, request.start + request.count) - ) - reserved = False - return serialized - finally: - if reserved: - async with self._lock: - self._reserved_input_tokens -= input_reservation - self._reserved_output_tokens -= ( - self._max_worker_output_tokens - ) - self._semaphore.release() - - async def accept_final_output(self, draft: PaperSummaryDraft) -> None: - """Validate final output and require complete worker coverage.""" - serialized = draft.model_dump_json() - output_tokens = _estimated_tokens(serialized) - if output_tokens > self._max_final_output_tokens: - raise PaperSummaryBudgetExceeded( - "paper summary exceeded max_summary_output_tokens" - ) - async with self._lock: - if ( - self._output_tokens + output_tokens - > self._max_total_output_tokens - ): - raise PaperSummaryBudgetExceeded( - "paper summary exceeded max_summary_total_output_tokens" - ) - self._output_tokens += output_tokens - - def require_complete_coverage(self, chunk_count: int) -> None: - """Reject a coordinator that did not delegate every source chunk.""" - missing = set(range(chunk_count)) - self._covered_chunks - if missing: - preview = ", ".join(str(index) for index in sorted(missing)[:8]) - raise PaperSummaryBudgetExceeded( - "paper summary research did not cover every chunk; missing " - f"{preview}" + if finding.quote is not None and finding.quote not in chunk.text: + raise ValueError( + "research finding quote is not present in its chunk" ) -def _suggested_research_groups(chunk_count: int) -> list[dict[str, int]]: - return [ - { - "start": start, - "count": min(_MAX_RESEARCH_GROUP_SIZE, chunk_count - start), - } - for start in range(0, chunk_count, _MAX_RESEARCH_GROUP_SIZE) - ] - - -def _summary_manifest( +def _reduce_payload( source: PaperSourceRevision, chunk_set: PaperChunkSet, - cfg: PaperFlowCfg, + reports: Sequence[PaperResearchDraft], ) -> str: - manifest = [ - { - "chunk_index": chunk.position, - "pages": sorted({span.page_number for span in chunk.source_spans}), - "characters": len(chunk.text), - "preview": chunk.text[:160], - } - for chunk in chunk_set.chunks - ] return json.dumps( { "title": source.title, "authors": source.authors, - "page_count": len(source.parsed.pages), - "required_citations": cfg.min_summary_citations, - "required_source_pages": cfg.min_summary_pages, - "required_research_groups": _suggested_research_groups( - len(chunk_set.chunks) - ), - "chunks": manifest, + "chunk_count": len(chunk_set.chunks), + "reports": [ + { + "scope_summary": report.scope_summary, + "findings": [ + { + "kind": finding.kind, + "claim": finding.claim, + "chunk_index": finding.citation.chunk_index, + "page_number": finding.citation.page_number, + "quote": finding.quote, + } + for finding in report.findings + ], + } + for report in reports + ], }, ensure_ascii=False, ) @@ -413,16 +255,10 @@ def _summary_instructions_hash(cfg: PaperFlowCfg) -> str: payload = json.dumps( { "orchestration": _ORCHESTRATION_VERSION, - "coordinator_instructions": _summary_instructions(cfg), + "reducer_instructions": _summary_instructions(cfg), "research_instructions": _RESEARCH_INSTRUCTIONS, - "max_research_calls": cfg.max_summary_tool_calls, - "max_research_concurrency": cfg.max_summary_concurrency, - "research_max_turns": cfg.max_summary_worker_turns, - "research_max_output_tokens": ( - cfg.max_summary_worker_output_tokens - ), - "max_total_input_tokens": cfg.max_summary_input_tokens, - "max_total_output_tokens": cfg.max_summary_total_output_tokens, + "research_group_size": cfg.summary_research_group_size, + "max_output_tokens": cfg.max_summary_output_tokens, }, ensure_ascii=False, separators=(",", ":"), @@ -431,79 +267,17 @@ def _summary_instructions_hash(cfg: PaperFlowCfg) -> str: return hashlib.sha256(payload.encode("utf-8")).hexdigest() -def _bounded_model_settings(cfg: PaperFlowCfg) -> ModelSettings: +def _summary_model_settings(cfg: PaperFlowCfg) -> ModelSettings: settings = cfg.model_settings or ModelSettings() configured = settings.max_tokens or cfg.max_summary_output_tokens return replace( settings, max_tokens=min(configured, cfg.max_summary_output_tokens), - parallel_tool_calls=True, ) -def _bounded_research_model_settings(cfg: PaperFlowCfg) -> ModelSettings: - settings = cfg.model_settings or ModelSettings() - configured = settings.max_tokens or cfg.max_summary_worker_output_tokens - return replace( - settings, - max_tokens=min(configured, cfg.max_summary_worker_output_tokens), - parallel_tool_calls=False, - ) - - -def _build_research_agent_tool( - source: PaperSourceRevision, - chunk_set: PaperChunkSet, - cfg: PaperFlowCfg, - budget: _SummaryBudget, -) -> FunctionTool: - """Expose one bounded research agent to the summary coordinator.""" - researcher: Agent[Any] = Agent( - name="paper_chunk_researcher", - instructions=_RESEARCH_INSTRUCTIONS, - model=cfg.model, - model_settings=_bounded_research_model_settings(cfg), - output_type=PaperResearchDraft, - ) - - def build_input(options: Any) -> str: - request = PaperResearchRequest.model_validate(options["params"]) - return _research_payload(source, chunk_set, request) - - async def extract_output(result: Any) -> str: - draft = PaperResearchDraft.model_validate(result.final_output) - return draft.model_dump_json() - - tool = researcher.as_tool( - tool_name="research_chunk_group", - tool_description=( - "Delegate one bounded chunk range to a paper research subagent. " - "Call every required research group before final synthesis and " - "use extra overlapping calls only for focused follow-up." - ), - custom_output_extractor=extract_output, - max_turns=cfg.max_summary_worker_turns, - parameters=PaperResearchRequest, - input_builder=build_input, - failure_error_function=None, - ) - invoke_agent = tool.on_invoke_tool - - async def invoke_bounded_agent(context: Any, input_json: str) -> str: - request = PaperResearchRequest.model_validate_json(input_json) - return await budget.invoke_research( - source, - chunk_set, - request, - lambda: invoke_agent(context, input_json), - ) - - tool.on_invoke_tool = invoke_bounded_agent - return tool - - class _AgentsPaperSummaryProvider: - """Coordinate bounded research subagents and synthesize one summary.""" + """Fan out one research agent per chunk group, then reduce to a summary.""" async def summarize( self, @@ -512,31 +286,46 @@ async def summarize( *, cfg: PaperFlowCfg, ) -> PaperSummaryDraft: - manifest = _summary_manifest(source, chunk_set, cfg) - budget = _SummaryBudget( - cfg, - manifest, - chunk_count=len(chunk_set.chunks), + groups = _chunk_groups( + len(chunk_set.chunks), cfg.summary_research_group_size ) - research_tool = _build_research_agent_tool( - source, - chunk_set, - cfg, - budget, + model_settings = _summary_model_settings(cfg) + researcher: Agent[Any] = Agent( + name="paper_chunk_researcher", + instructions=_RESEARCH_INSTRUCTIONS, + model=cfg.model, + model_settings=model_settings, + output_type=PaperResearchDraft, ) - coordinator: Agent[Any] = Agent( - name="paper_summary_coordinator", + semaphore = asyncio.Semaphore(cfg.summary_concurrency) + + async def study(group: _ChunkGroup) -> PaperResearchDraft: + async with semaphore: + output = await run_with_observability( + researcher, + _research_payload(source, chunk_set, group), + cfg=cfg, + memory=None, + extra_run_hooks=[], + ) + draft = PaperResearchDraft.model_validate(output) + _validate_research_draft(chunk_set, group, draft) + return draft + + reports = await asyncio.gather(*(study(group) for group in groups)) + + reducer: Agent[Any] = Agent( + name="paper_summary_reducer", instructions=_summary_instructions(cfg), model=cfg.model, - model_settings=_bounded_model_settings(cfg), - tools=[research_tool], + model_settings=model_settings, output_type=PaperSummaryDraft, ) try: output = await asyncio.wait_for( run_with_observability( - coordinator, - manifest, + reducer, + _reduce_payload(source, chunk_set, reports), cfg=cfg, memory=None, extra_run_hooks=[], @@ -544,10 +333,7 @@ async def summarize( timeout=cfg.timeout_seconds, ) except asyncio.TimeoutError as exc: - raise PaperSummaryBudgetExceeded( + raise PaperSummaryError( "paper summary exceeded timeout_seconds" ) from exc - draft = PaperSummaryDraft.model_validate(output) - budget.require_complete_coverage(len(chunk_set.chunks)) - await budget.accept_final_output(draft) - return draft + return PaperSummaryDraft.model_validate(output) diff --git a/quantmind/flows/paper.py b/quantmind/flows/paper.py index 2e14b8e..d17fcc1 100644 --- a/quantmind/flows/paper.py +++ b/quantmind/flows/paper.py @@ -1,8 +1,15 @@ -"""Source-first paper flow that returns chunks before a cited summary.""" +"""Source-first paper flow that returns chunks before a cited summary. + +The flow is deliberately thin. It owns only what genuinely needs both the +preprocess/rag value objects and IO: fetching bytes, parsing the PDF, reading +page-asset bytes off disk, and mapping those ephemeral, path-based artifacts +into knowledge-native inputs. All identity (IDs, content and producer hashes, +citation resolution) lives on the knowledge models' ``from_*`` constructors, +so this module imports no private ID helpers and computes no paper ID itself. +""" -import hashlib import mimetypes -from dataclasses import dataclass +from collections.abc import Sequence from datetime import datetime, timezone from importlib.metadata import version from pathlib import Path @@ -24,32 +31,19 @@ _summary_instructions_hash, ) from quantmind.knowledge import ( - ArtifactLocator, - PaperAssetRef, + PaperAssetInput, PaperBoundingBox, - PaperChunk, PaperChunkingConfig, + PaperChunkInput, PaperChunkSet, - PaperCitation, + PaperCitationDraft, PaperFlowResult, PaperGlobalSummary, + PaperPageInput, PaperParsedBlock, - PaperParsedManifest, - PaperParsedPage, + PaperSourceFacts, PaperSourceRevision, - PaperSourceSpan, PaperSummaryProducer, - SourceRef, -) -from quantmind.knowledge.paper import ( - _paper_artifact_id, - _paper_asset_id, - _paper_chunk_id, - _paper_chunk_set_content_hash, - _paper_source_id, - _paper_summary_content_hash, - _stable_hash, - _text_hash, ) from quantmind.preprocess.fetch import ( Fetched, @@ -59,33 +53,17 @@ read_local_file, ) from quantmind.preprocess.format import ParsedDocument, parse_pdf -from quantmind.rag import SentenceSplitterConfig, chunk_parsed_document +from quantmind.rag import ( + ParsedChunk, + SentenceSplitterConfig, + chunk_parsed_document, +) class UnsupportedContentTypeError(ValueError): """The source is not a page-aware PDF supported by Paper Flow V1.""" -class PaperCitationValidationError(ValueError): - """A generated summary did not provide valid source coverage.""" - - -@dataclass(frozen=True) -class _FetchedPaperSource: - """Exact fetched bytes and code-owned source facts before parsing.""" - - bytes: bytes - media_type: str - kind: Literal["arxiv", "http", "local"] - uri: str - fetched_at: datetime - available_at: datetime - published_at: datetime | None = None - arxiv_id: str | None = None - title: str | None = None - authors: tuple[str, ...] = () - - async def paper_flow( input: PaperInput, *, @@ -95,12 +73,12 @@ async def paper_flow( """Build a page-aware chunk set and one cited global summary. IDs, source metadata, artifact membership, lineage, and citation links are - created and validated by code. The model returns only summary prose and - chunk/page coordinates through a bounded summarization seam. + minted and validated by the knowledge-layer constructors. The model returns + only summary prose and chunk/page coordinates through a bounded seam. Args: input: Typed paper source. V1 requires a PDF-backed input. - cfg: Splitter, summary model, and explicit usage/runtime limits. + cfg: Splitter, summary model, and usage/runtime limits. Returns: The exact source revision, one chunk-set artifact, and one cited @@ -108,17 +86,21 @@ async def paper_flow( Raises: UnsupportedContentTypeError: If the resolved content is not a PDF. - PaperCitationValidationError: If generated citations are invalid or - do not meet configured source-coverage requirements. + PaperCitationValidationError: If generated citations are invalid or do + not meet the configured source-coverage policy. NotImplementedError: If a DOI input has no exact open PDF resolver. """ cfg = cfg or PaperFlowCfg() - fetched = await _fetch_paper_source(input) - parsed = await parse_pdf( - fetched.bytes, - artifact_dir=cfg.output_dir, + facts = await _fetch_paper_source(input) + parsed = await parse_pdf(facts.raw_bytes, artifact_dir=cfg.output_dir) + source = PaperSourceRevision.from_parsed( + facts=facts, + source_hash=parsed.source_hash, + parser_name=parsed.parser_name, + parser_version=parsed.parser_version, + cleanup_version=parsed.cleanup_version, + pages=_adapt_pages(parsed), ) - source = _build_source_revision(fetched, parsed) parsed_chunks = chunk_parsed_document( parsed, config=SentenceSplitterConfig( @@ -126,10 +108,19 @@ async def paper_flow( chunk_overlap=cfg.chunk_overlap, ), ) - chunk_set = _build_chunk_set(source, parsed_chunks, cfg) + producer = PaperChunkingConfig( + splitter_version=version("llama-index-core"), + chunk_size=cfg.chunk_size, + chunk_overlap=cfg.chunk_overlap, + ) + chunk_set = PaperChunkSet.from_parsed_chunks( + source, + _adapt_chunks(parsed_chunks, source), + producer=producer, + ) provider = _summary_provider or _AgentsPaperSummaryProvider() draft = await provider.summarize(source, chunk_set, cfg=cfg) - summary = _build_global_summary(source, chunk_set, draft, cfg) + summary = _build_summary(chunk_set, draft, cfg) return PaperFlowResult( source_revision=source, chunk_set=chunk_set, @@ -154,7 +145,8 @@ def _require_pdf(raw: Fetched) -> None: ) -async def _fetch_paper_source(input: PaperInput) -> _FetchedPaperSource: +async def _fetch_paper_source(input: PaperInput) -> PaperSourceFacts: + """Fetch and normalize one input into code-owned source facts (IO).""" if isinstance(input, ArxivIdentifier): raw_paper: RawPaper = await fetch_arxiv(input.id) _require_pdf(raw_paper) @@ -165,11 +157,11 @@ async def _fetch_paper_source(input: PaperInput) -> _FetchedPaperSource: uri = raw_paper.resolved_url or raw_paper.source_url if uri is None: raise ValueError("resolved arXiv paper is missing its source URL") - return _FetchedPaperSource( - bytes=raw_paper.bytes, - media_type="application/pdf", + return PaperSourceFacts( kind="arxiv", uri=uri, + media_type="application/pdf", + raw_bytes=raw_paper.bytes, fetched_at=fetched_at, available_at=available_at, published_at=raw_paper.published_at, @@ -181,11 +173,11 @@ async def _fetch_paper_source(input: PaperInput) -> _FetchedPaperSource: raw_http = await fetch_url(input.url) _require_pdf(raw_http) fetched_at = _aware_or_now(raw_http.fetched_at) - return _FetchedPaperSource( - bytes=raw_http.bytes, - media_type="application/pdf", + return PaperSourceFacts( kind="http", uri=raw_http.resolved_url or raw_http.source_url or input.url, + media_type="application/pdf", + raw_bytes=raw_http.bytes, fetched_at=fetched_at, available_at=fetched_at, ) @@ -194,11 +186,11 @@ async def _fetch_paper_source(input: PaperInput) -> _FetchedPaperSource: _require_pdf(raw_local) observed_at = _aware_or_now(raw_local.fetched_at) path = Path(input.path).expanduser().resolve() - return _FetchedPaperSource( - bytes=raw_local.bytes, - media_type="application/pdf", + return PaperSourceFacts( kind="local", uri=raw_local.source_url or path.as_uri(), + media_type="application/pdf", + raw_bytes=raw_local.bytes, fetched_at=observed_at, available_at=observed_at, ) @@ -215,15 +207,13 @@ async def _fetch_paper_source(input: PaperInput) -> _FetchedPaperSource: raise TypeError(f"Unsupported PaperInput variant: {type(input)!r}") -def _asset_from_path( +def _read_page_asset( path_value: str, *, - source_revision_id, kind: Literal["screenshot", "image"], page_number: int, - assets: dict, - blobs: dict[str, bytes], -) -> PaperAssetRef: +) -> PaperAssetInput: + """Read one parser-written page asset off disk into knowledge input (IO).""" path = Path(path_value) try: content = path.read_bytes() @@ -231,286 +221,114 @@ def _asset_from_path( raise RuntimeError( f"Parser asset for page {page_number} is missing: {path}" ) from exc - content_hash = hashlib.sha256(content).hexdigest() media_type = ( mimetypes.guess_type(path.name)[0] or "application/octet-stream" ) - asset = PaperAssetRef( - asset_id=_paper_asset_id( - source_revision_id, - kind=kind, - page_number=page_number, - content_hash=content_hash, - ), - kind=kind, - media_type=media_type, - content_hash=content_hash, - size_bytes=len(content), - page_number=page_number, - ) - assets[asset.asset_id] = asset - blobs[content_hash] = content - return asset + return PaperAssetInput(kind=kind, content=content, media_type=media_type) -def _build_source_revision( - fetched: _FetchedPaperSource, - parsed: ParsedDocument, -) -> PaperSourceRevision: - source_id = _paper_source_id(parsed.source_hash) - blobs = {parsed.source_hash: fetched.bytes} - raw = PaperAssetRef( - asset_id=_paper_asset_id( - source_id, - kind="raw", - page_number=None, - content_hash=parsed.source_hash, - ), - kind="raw", - media_type=fetched.media_type, - content_hash=parsed.source_hash, - size_bytes=len(fetched.bytes), - ) - assets = {raw.asset_id: raw} - pages: list[PaperParsedPage] = [] +def _adapt_pages(parsed: ParsedDocument) -> list[PaperPageInput]: + """Map path-based parsed pages into knowledge-native page inputs.""" + pages: list[PaperPageInput] = [] for page in parsed.pages: + blocks = tuple( + PaperParsedBlock( + text=block.text, + bbox=PaperBoundingBox( + x0=block.bbox.x0, + y0=block.bbox.y0, + x1=block.bbox.x1, + y1=block.bbox.y1, + ), + font_name=block.font_name, + font_size=block.font_size, + confidence=block.confidence, + ) + for block in page.blocks + ) screenshot = ( - _asset_from_path( + _read_page_asset( page.screenshot_path, - source_revision_id=source_id, kind="screenshot", page_number=page.page_number, - assets=assets, - blobs=blobs, ) if page.screenshot_path is not None else None ) images = tuple( - _asset_from_path( + _read_page_asset( image_path, - source_revision_id=source_id, kind="image", page_number=page.page_number, - assets=assets, - blobs=blobs, ) for image_path in page.image_paths ) pages.append( - PaperParsedPage( + PaperPageInput( page_number=page.page_number, width=page.width, height=page.height, text=page.text, - blocks=tuple( - PaperParsedBlock( - text=block.text, - bbox=PaperBoundingBox( - x0=block.bbox.x0, - y0=block.bbox.y0, - x1=block.bbox.x1, - y1=block.bbox.y1, - ), - font_name=block.font_name, - font_size=block.font_size, - confidence=block.confidence, - ) - for block in page.blocks - ), - screenshot_asset_id=( - screenshot.asset_id if screenshot is not None else None - ), - image_asset_ids=tuple(image.asset_id for image in images), + blocks=blocks, + screenshot=screenshot, + images=images, ) ) - source_ref = SourceRef( - kind=fetched.kind, - uri=fetched.uri, - fetched_at=fetched.fetched_at, - content_hash=parsed.source_hash, - ) - return PaperSourceRevision( - id=source_id, - source=source_ref, - as_of=fetched.published_at or fetched.available_at, - available_at=fetched.available_at, - published_at=fetched.published_at, - arxiv_id=fetched.arxiv_id, - title=fetched.title, - authors=fetched.authors, - parsed=PaperParsedManifest( - source_hash=parsed.source_hash, - parser_name=parsed.parser_name, - parser_version=parsed.parser_version, - cleanup_version=parsed.cleanup_version, - pages=tuple(pages), - ), - raw_asset_id=raw.asset_id, - assets=tuple(assets.values()), - blobs=blobs, - ) + return pages -def _build_chunk_set( +def _adapt_chunks( + parsed_chunks: Sequence[ParsedChunk], source: PaperSourceRevision, - parsed_chunks, - cfg: PaperFlowCfg, -) -> PaperChunkSet: - producer = PaperChunkingConfig( - splitter_version=version("llama-index-core"), - chunk_size=cfg.chunk_size, - chunk_overlap=cfg.chunk_overlap, - ) - producer_hash = _stable_hash(producer.model_dump(mode="json")) - artifact_id = _paper_artifact_id( - source.id, - "paper_chunk_set", - producer_hash, - ) - page_assets = { - page.page_number: tuple( - asset_id - for asset_id in ( - (page.screenshot_asset_id,) - if page.screenshot_asset_id is not None - else () - ) - + page.image_asset_ids - ) - for page in source.parsed.pages - } - chunks: list[PaperChunk] = [] - for position, parsed_chunk in enumerate(parsed_chunks): - if parsed_chunk.source_hash != source.source.content_hash: +) -> list[PaperChunkInput]: + """Map splitter chunks into knowledge-native chunk inputs.""" + content_hash = source.source.content_hash + chunks: list[PaperChunkInput] = [] + for parsed_chunk in parsed_chunks: + if parsed_chunk.source_hash != content_hash: raise ValueError("parsed chunk belongs to another source revision") - span = PaperSourceSpan( - page_number=parsed_chunk.page_number, - start_char=parsed_chunk.start_char, - end_char=parsed_chunk.end_char, - block_boxes=tuple( - PaperBoundingBox( - x0=box.x0, - y0=box.y0, - x1=box.x1, - y1=box.y1, - ) - for box in parsed_chunk.block_boxes - ), - asset_ids=page_assets.get(parsed_chunk.page_number, ()), - ) - content_hash = _text_hash(parsed_chunk.text) chunks.append( - PaperChunk( - chunk_id=_paper_chunk_id( - artifact_id, - position=position, - content_hash=content_hash, - spans=(span,), + PaperChunkInput( + page_number=parsed_chunk.page_number, + start_char=parsed_chunk.start_char, + end_char=parsed_chunk.end_char, + block_boxes=tuple( + PaperBoundingBox(x0=box.x0, y0=box.y0, x1=box.x1, y1=box.y1) + for box in parsed_chunk.block_boxes ), - chunk_set_id=artifact_id, - source_revision_id=source.id, - position=position, text=parsed_chunk.text, - content_hash=content_hash, - source_spans=(span,), ) ) - if not chunks: - raise ValueError("paper source produced no non-empty chunks") - chunk_tuple = tuple(chunks) - return PaperChunkSet( - id=artifact_id, - source_revision_id=source.id, - producer=producer, - producer_config_hash=producer_hash, - content_hash=_paper_chunk_set_content_hash(chunk_tuple), - chunks=chunk_tuple, - ) + return chunks -def _build_global_summary( - source: PaperSourceRevision, +def _build_summary( chunk_set: PaperChunkSet, draft: PaperSummaryDraft, cfg: PaperFlowCfg, ) -> PaperGlobalSummary: - citations: list[PaperCitation] = [] - seen: set[tuple[int, int, str | None]] = set() - for draft_citation in draft.citations: - try: - chunk = chunk_set.chunks[draft_citation.chunk_index] - except IndexError as exc: - raise PaperCitationValidationError( - "paper summary cites an unknown chunk index" - ) from exc - pages = {span.page_number for span in chunk.source_spans} - if draft_citation.page_number not in pages: - raise PaperCitationValidationError( - "paper summary citation page is not owned by its chunk" - ) - quote = draft_citation.quote - if quote is not None and quote not in chunk.text: - raise PaperCitationValidationError( - "paper summary citation quote is not present in its chunk" - ) - key = (draft_citation.chunk_index, draft_citation.page_number, quote) - if key in seen: - continue - seen.add(key) - citations.append( - PaperCitation( - chunk_set_id=chunk_set.id, - chunk_id=chunk.chunk_id, - page_number=draft_citation.page_number, - quote=quote, - ) - ) - cited_pages = {citation.page_number for citation in citations} - if len(citations) < cfg.min_summary_citations: - raise PaperCitationValidationError( - "paper summary has fewer citations than min_summary_citations" - ) - if len(cited_pages) < cfg.min_summary_pages: - raise PaperCitationValidationError( - "paper summary has fewer source pages than min_summary_pages" - ) - + """Assemble the summary producer and delegate identity to the model.""" producer = PaperSummaryProducer( model=cfg.model, prompt_version=cfg.summary_prompt_version, input_chunk_set_id=chunk_set.id, instructions_hash=_summary_instructions_hash(cfg), max_output_tokens=cfg.max_summary_output_tokens, - max_research_calls=cfg.max_summary_tool_calls, - max_research_concurrency=cfg.max_summary_concurrency, - research_max_turns=cfg.max_summary_worker_turns, - research_max_output_tokens=cfg.max_summary_worker_output_tokens, + research_group_size=cfg.summary_research_group_size, ) - producer_hash = _stable_hash(producer.model_dump(mode="json")) - artifact_id = _paper_artifact_id( - source.id, - "paper_summary", - producer_hash, + citations = tuple( + PaperCitationDraft( + chunk_index=citation.chunk_index, + page_number=citation.page_number, + quote=citation.quote, + ) + for citation in draft.citations ) - citation_tuple = tuple(citations) - summary_text = draft.summary.strip() - return PaperGlobalSummary( - id=artifact_id, - source_revision_id=source.id, + return PaperGlobalSummary.from_draft( + chunk_set, producer=producer, - producer_config_hash=producer_hash, - content_hash=_paper_summary_content_hash( - summary_text, - citation_tuple, - ), - summary=summary_text, - citations=citation_tuple, - derived_from=( - ArtifactLocator( - source_revision_id=source.id, - artifact_id=chunk_set.id, - artifact_kind="paper_chunk_set", - ), - ), + summary=draft.summary, + citations=citations, + min_citations=cfg.min_summary_citations, + min_pages=cfg.min_summary_pages, ) diff --git a/quantmind/knowledge/__init__.py b/quantmind/knowledge/__init__.py index 9e904c5..37bd5ac 100644 --- a/quantmind/knowledge/__init__.py +++ b/quantmind/knowledge/__init__.py @@ -29,17 +29,23 @@ LegacyPaper, PaperArtifact, PaperArtifactKind, + PaperAssetInput, PaperAssetRef, PaperBoundingBox, PaperChunk, PaperChunkingConfig, + PaperChunkInput, PaperChunkSet, PaperCitation, + PaperCitationDraft, + PaperCitationValidationError, PaperFlowResult, PaperGlobalSummary, + PaperPageInput, PaperParsedBlock, PaperParsedManifest, PaperParsedPage, + PaperSourceFacts, PaperSourceRevision, PaperSourceSpan, PaperSummaryProducer, @@ -66,17 +72,23 @@ "LegacyPaper", "PaperArtifact", "PaperArtifactKind", + "PaperAssetInput", "PaperAssetRef", "PaperBoundingBox", "PaperChunk", + "PaperChunkInput", "PaperChunkingConfig", "PaperChunkSet", "PaperCitation", + "PaperCitationDraft", + "PaperCitationValidationError", "PaperFlowResult", "PaperGlobalSummary", + "PaperPageInput", "PaperParsedBlock", "PaperParsedManifest", "PaperParsedPage", + "PaperSourceFacts", "PaperSourceRevision", "PaperSourceSpan", "PaperSummaryProducer", diff --git a/quantmind/knowledge/paper.py b/quantmind/knowledge/paper.py index fe24944..0064536 100644 --- a/quantmind/knowledge/paper.py +++ b/quantmind/knowledge/paper.py @@ -3,6 +3,8 @@ import hashlib import json import re +from collections.abc import Sequence +from dataclasses import dataclass from datetime import datetime from enum import Enum from typing import Literal @@ -72,6 +74,40 @@ def _paper_asset_id( ) +class PaperCitationValidationError(ValueError): + """A generated summary did not provide valid source coverage.""" + + +@dataclass(frozen=True) +class PaperSourceFacts: + """Code-owned source facts normalized by the flow before construction. + + Everything here comes from fetching and IO rather than the parser or the + model, so identity construction can stay inside the knowledge layer while + fetching stays in the flow. + """ + + kind: Literal["arxiv", "http", "local"] + uri: str + media_type: str + raw_bytes: bytes + fetched_at: datetime + available_at: datetime + published_at: datetime | None = None + arxiv_id: str | None = None + title: str | None = None + authors: tuple[str, ...] = () + + +@dataclass(frozen=True) +class PaperAssetInput: + """Raw bytes for one page-level visual asset, pre-read by the flow.""" + + kind: Literal["screenshot", "image"] + content: bytes + media_type: str + + class PaperBoundingBox(BaseModel): """One source rectangle in top-left-origin page coordinates.""" @@ -89,6 +125,17 @@ def _coordinates_are_ordered(self) -> "PaperBoundingBox": return self +@dataclass(frozen=True) +class PaperChunkInput: + """One splitter-produced chunk in knowledge-native form.""" + + page_number: int + start_char: int + end_char: int + block_boxes: tuple[PaperBoundingBox, ...] + text: str + + class PaperAssetRef(BaseModel): """Content-addressed source, screenshot, or extracted-image reference.""" @@ -114,6 +161,19 @@ class PaperParsedBlock(BaseModel): confidence: float | None = None +@dataclass(frozen=True) +class PaperPageInput: + """One parsed page plus pre-read visual-asset bytes for construction.""" + + page_number: int + width: float + height: float + text: str + blocks: tuple[PaperParsedBlock, ...] = () + screenshot: PaperAssetInput | None = None + images: tuple[PaperAssetInput, ...] = () + + class PaperParsedPage(BaseModel): """One physical page and its content-addressed visual references.""" @@ -251,6 +311,114 @@ def blob_for(self, asset_id: UUID) -> bytes: f"Paper asset '{asset_id}' bytes are not loaded" ) from exc + @classmethod + def from_parsed( + cls, + *, + facts: PaperSourceFacts, + source_hash: str, + parser_name: str, + parser_version: str, + cleanup_version: str, + pages: Sequence[PaperPageInput], + ) -> "PaperSourceRevision": + """Build a source revision, minting every content-addressed ID. + + The flow supplies normalized ``facts`` and pre-read page bytes; this + constructor owns all identity (source ID, asset IDs) so callers never + compute a paper ID themselves. + """ + source_id = _paper_source_id(source_hash) + blobs: dict[str, bytes] = {source_hash: facts.raw_bytes} + raw = PaperAssetRef( + asset_id=_paper_asset_id( + source_id, + kind="raw", + page_number=None, + content_hash=source_hash, + ), + kind="raw", + media_type=facts.media_type, + content_hash=source_hash, + size_bytes=len(facts.raw_bytes), + ) + assets: dict[UUID, PaperAssetRef] = {raw.asset_id: raw} + parsed_pages: list[PaperParsedPage] = [] + for page in pages: + screenshot_id: UUID | None = None + if page.screenshot is not None: + screenshot_id = cls._register_asset( + source_id, page.page_number, page.screenshot, assets, blobs + ).asset_id + image_ids = tuple( + cls._register_asset( + source_id, page.page_number, image, assets, blobs + ).asset_id + for image in page.images + ) + parsed_pages.append( + PaperParsedPage( + page_number=page.page_number, + width=page.width, + height=page.height, + text=page.text, + blocks=page.blocks, + screenshot_asset_id=screenshot_id, + image_asset_ids=image_ids, + ) + ) + return cls( + id=source_id, + source=SourceRef( + kind=facts.kind, + uri=facts.uri, + fetched_at=facts.fetched_at, + content_hash=source_hash, + ), + as_of=facts.published_at or facts.available_at, + available_at=facts.available_at, + published_at=facts.published_at, + arxiv_id=facts.arxiv_id, + title=facts.title, + authors=facts.authors, + parsed=PaperParsedManifest( + source_hash=source_hash, + parser_name=parser_name, + parser_version=parser_version, + cleanup_version=cleanup_version, + pages=tuple(parsed_pages), + ), + raw_asset_id=raw.asset_id, + assets=tuple(assets.values()), + blobs=blobs, + ) + + @staticmethod + def _register_asset( + source_id: UUID, + page_number: int, + asset: PaperAssetInput, + assets: dict[UUID, PaperAssetRef], + blobs: dict[str, bytes], + ) -> PaperAssetRef: + content_hash = hashlib.sha256(asset.content).hexdigest() + ref = PaperAssetRef( + asset_id=_paper_asset_id( + source_id, + kind=asset.kind, + page_number=page_number, + content_hash=content_hash, + ), + kind=asset.kind, + media_type=asset.media_type, + content_hash=content_hash, + size_bytes=len(asset.content), + page_number=page_number, + ) + assets[ref.asset_id] = ref + blobs[content_hash] = asset.content + return ref + class PaperSourceSpan(BaseModel): """Exact character span and page evidence retained by one chunk.""" @@ -399,6 +567,73 @@ def _validate_artifact(self) -> "PaperChunkSet": raise ValueError("paper chunk-set content hash mismatch") return self + @classmethod + def from_parsed_chunks( + cls, + source: PaperSourceRevision, + chunks: Sequence[PaperChunkInput], + *, + producer: PaperChunkingConfig, + ) -> "PaperChunkSet": + """Build a chunk-set artifact, minting the artifact and chunk IDs. + + The flow supplies knowledge-native chunk inputs; this constructor owns + the artifact ID, per-chunk content hashes, chunk IDs, and the set + content hash. + """ + producer_hash = _stable_hash(producer.model_dump(mode="json")) + artifact_id = _paper_artifact_id( + source.id, PaperArtifactKind.CHUNK_SET, producer_hash + ) + page_assets = { + page.page_number: ( + ( + (page.screenshot_asset_id,) + if page.screenshot_asset_id is not None + else () + ) + + page.image_asset_ids + ) + for page in source.parsed.pages + } + built: list[PaperChunk] = [] + for position, chunk in enumerate(chunks): + span = PaperSourceSpan( + page_number=chunk.page_number, + start_char=chunk.start_char, + end_char=chunk.end_char, + block_boxes=chunk.block_boxes, + asset_ids=page_assets.get(chunk.page_number, ()), + ) + content_hash = _text_hash(chunk.text) + built.append( + PaperChunk( + chunk_id=_paper_chunk_id( + artifact_id, + position=position, + content_hash=content_hash, + spans=(span,), + ), + chunk_set_id=artifact_id, + source_revision_id=source.id, + position=position, + text=chunk.text, + content_hash=content_hash, + source_spans=(span,), + ) + ) + if not built: + raise ValueError("paper source produced no non-empty chunks") + chunk_tuple = tuple(built) + return cls( + id=artifact_id, + source_revision_id=source.id, + producer=producer, + producer_config_hash=producer_hash, + content_hash=_paper_chunk_set_content_hash(chunk_tuple), + chunks=chunk_tuple, + ) + def _validate_chunk_set_source( source: PaperSourceRevision, @@ -450,6 +685,15 @@ class PaperCitation(BaseModel): quote: str | None = Field(default=None, max_length=500) +@dataclass(frozen=True) +class PaperCitationDraft: + """Model-proposed citation coordinates before code resolves canonical IDs.""" + + chunk_index: int + page_number: int + quote: str | None = None + + class PaperSummaryProducer(BaseModel): """Exact model, prompt, and output-bound identity for a summary.""" @@ -457,16 +701,11 @@ class PaperSummaryProducer(BaseModel): model: str prompt_version: str - orchestration: Literal["manager-research-agents-v1"] = ( - "manager-research-agents-v1" - ) + orchestration: Literal["map-reduce-v1"] = "map-reduce-v1" input_chunk_set_id: UUID instructions_hash: str = Field(pattern=r"^[0-9a-f]{64}$") max_output_tokens: int = Field(gt=0) - max_research_calls: int = Field(default=12, ge=1) - max_research_concurrency: int = Field(default=2, ge=1) - research_max_turns: int = Field(default=4, ge=1) - research_max_output_tokens: int = Field(default=1_536, ge=1) + research_group_size: int = Field(ge=1) def _paper_summary_content_hash( @@ -543,6 +782,91 @@ def _validate_artifact(self) -> "PaperGlobalSummary": ) return self + @classmethod + def from_draft( + cls, + chunk_set: PaperChunkSet, + *, + producer: PaperSummaryProducer, + summary: str, + citations: Sequence[PaperCitationDraft], + min_citations: int, + min_pages: int, + ) -> "PaperGlobalSummary": + """Resolve model-proposed citations and mint the summary artifact. + + The model returns only prose and ``(chunk_index, page, quote)`` + coordinates. This constructor resolves canonical chunk IDs, validates + every citation against the chunk set, enforces the configured coverage + policy, and mints the artifact identity and lineage. + """ + if producer.input_chunk_set_id != chunk_set.id: + raise ValueError("summary producer input does not match chunk set") + resolved: list[PaperCitation] = [] + seen: set[tuple[int, int, str | None]] = set() + for draft in citations: + try: + chunk = chunk_set.chunks[draft.chunk_index] + except IndexError as exc: + raise PaperCitationValidationError( + "paper summary cites an unknown chunk index" + ) from exc + pages = {span.page_number for span in chunk.source_spans} + if draft.page_number not in pages: + raise PaperCitationValidationError( + "paper summary citation page is not owned by its chunk" + ) + if draft.quote is not None and draft.quote not in chunk.text: + raise PaperCitationValidationError( + "paper summary citation quote is not present in its chunk" + ) + key = (draft.chunk_index, draft.page_number, draft.quote) + if key in seen: + continue + seen.add(key) + resolved.append( + PaperCitation( + chunk_set_id=chunk_set.id, + chunk_id=chunk.chunk_id, + page_number=draft.page_number, + quote=draft.quote, + ) + ) + if len(resolved) < min_citations: + raise PaperCitationValidationError( + "paper summary has fewer citations than min_summary_citations" + ) + if len({citation.page_number for citation in resolved}) < min_pages: + raise PaperCitationValidationError( + "paper summary has fewer source pages than min_summary_pages" + ) + producer_hash = _stable_hash(producer.model_dump(mode="json")) + artifact_id = _paper_artifact_id( + chunk_set.source_revision_id, + PaperArtifactKind.GLOBAL_SUMMARY, + producer_hash, + ) + summary_text = summary.strip() + citation_tuple = tuple(resolved) + return cls( + id=artifact_id, + source_revision_id=chunk_set.source_revision_id, + producer=producer, + producer_config_hash=producer_hash, + content_hash=_paper_summary_content_hash( + summary_text, citation_tuple + ), + summary=summary_text, + citations=citation_tuple, + derived_from=( + ArtifactLocator( + source_revision_id=chunk_set.source_revision_id, + artifact_id=chunk_set.id, + artifact_kind="paper_chunk_set", + ), + ), + ) + class PaperFlowResult(BaseModel): """Validated V1 result containing source, chunks, and cited summary.""" diff --git a/scripts/verify_pdf_rag_e2e.py b/scripts/verify_pdf_rag_e2e.py index 7ebe2b9..8eeb54c 100644 --- a/scripts/verify_pdf_rag_e2e.py +++ b/scripts/verify_pdf_rag_e2e.py @@ -58,13 +58,9 @@ async def _run_vertical_slice() -> dict[str, Any]: model="gpt-4o-mini", output_dir=str(root / "assets"), timeout_seconds=240, - max_summary_tool_calls=12, - max_summary_concurrency=2, - max_summary_worker_turns=4, - max_summary_worker_output_tokens=1_536, - max_summary_input_tokens=120_000, + summary_research_group_size=8, + summary_concurrency=2, max_summary_output_tokens=4_096, - max_summary_total_output_tokens=20_000, min_summary_citations=3, min_summary_pages=2, ), @@ -207,7 +203,7 @@ def _passed(snapshot: dict[str, Any]) -> bool: and snapshot["chunk_asset_reference_count"] > 0 and snapshot["chunk_count"] > 0 and snapshot["summary"] - and snapshot["summary_orchestration"] == "manager-research-agents-v1" + and snapshot["summary_orchestration"] == "map-reduce-v1" and snapshot["citation_count"] >= 3 and len(snapshot["citation_pages"]) >= 2 and snapshot["first_summary_scores"] diff --git a/tests/configs/test_paper.py b/tests/configs/test_paper.py index 3247356..6247521 100644 --- a/tests/configs/test_paper.py +++ b/tests/configs/test_paper.py @@ -23,28 +23,17 @@ def test_defaults(self): self.assertEqual(cfg.max_turns, 16) self.assertEqual(cfg.chunk_size, 512) self.assertEqual(cfg.chunk_overlap, 64) - self.assertEqual(cfg.max_summary_tool_calls, 12) - self.assertEqual(cfg.max_summary_worker_turns, 4) - self.assertEqual(cfg.max_summary_worker_output_tokens, 1_536) - self.assertEqual(cfg.max_summary_total_output_tokens, 20_000) + self.assertEqual(cfg.summary_research_group_size, 8) + self.assertEqual(cfg.summary_concurrency, 4) + self.assertEqual(cfg.max_summary_output_tokens, 4_096) self.assertEqual(cfg.min_summary_citations, 3) + self.assertEqual(cfg.min_summary_pages, 2) def test_invalid_overlap_or_coverage_bounds_are_rejected(self): with self.assertRaises(ValidationError): PaperFlowCfg(chunk_size=64, chunk_overlap=64) with self.assertRaises(ValidationError): PaperFlowCfg(min_summary_citations=1, min_summary_pages=2) - with self.assertRaises(ValidationError): - PaperFlowCfg( - max_summary_tool_calls=1, - max_summary_concurrency=2, - ) - with self.assertRaises(ValidationError): - PaperFlowCfg( - max_summary_output_tokens=4_096, - max_summary_worker_output_tokens=1_536, - max_summary_total_output_tokens=5_000, - ) class PaperInputDiscriminatedTests(unittest.TestCase): diff --git a/tests/flows/test_paper.py b/tests/flows/test_paper.py index 6bf86b0..8a48ecc 100644 --- a/tests/flows/test_paper.py +++ b/tests/flows/test_paper.py @@ -1,6 +1,7 @@ """Offline tests for source-first ``paper_flow`` behavior.""" import asyncio +import json import unittest from datetime import datetime, timezone from pathlib import Path @@ -20,23 +21,22 @@ PaperResearchCitationDraft, PaperResearchDraft, PaperResearchFindingDraft, - PaperResearchRequest, - PaperSummaryBudgetExceeded, PaperSummaryCitationDraft, PaperSummaryDraft, + PaperSummaryError, _AgentsPaperSummaryProvider, - _bounded_model_settings, - _bounded_research_model_settings, - _build_research_agent_tool, - _SummaryBudget, + _chunk_groups, + _ChunkGroup, + _summary_model_settings, + _validate_research_draft, ) from quantmind.flows.paper import ( - PaperCitationValidationError, UnsupportedContentTypeError, - _build_global_summary, + _build_summary, _fetch_paper_source, paper_flow, ) +from quantmind.knowledge import PaperCitationValidationError from quantmind.preprocess.fetch import Fetched, RawPaper from tests.paper_helpers import build_paper_result @@ -267,8 +267,7 @@ def test_unknown_chunk_page_and_quote_are_rejected(self) -> None: for draft in invalid_drafts: with self.subTest(draft=draft): with self.assertRaises(PaperCitationValidationError): - _build_global_summary( - result.source_revision, + _build_summary( result.chunk_set, draft, cfg, @@ -287,191 +286,166 @@ def test_configured_citation_and_page_coverage_is_enforced(self) -> None: PaperCitationValidationError, "fewer citations", ): - _build_global_summary( - result.source_revision, + _build_summary( result.chunk_set, draft, PaperFlowCfg(), ) -class SummaryBudgetTests(unittest.IsolatedAsyncioTestCase): - @staticmethod - def _research_draft( - result, - request: PaperResearchRequest, - ) -> PaperResearchDraft: - chunk = result.chunk_set.chunks[request.start] - page = chunk.source_spans[0].page_number - return PaperResearchDraft( - scope_summary=f"Reviewed chunk {request.start}.", +def _fake_research_run(cfg, source_pages): + """A run_with_observability stub: real worker drafts, one reducer draft.""" + + async def fake_run(agent, payload, *, cfg, memory, extra_run_hooks): + data = json.loads(payload) + if agent.name == "paper_chunk_researcher": + return PaperResearchDraft( + scope_summary="reviewed the assigned chunk group", + findings=tuple( + PaperResearchFindingDraft( + kind="method", + claim=f"finding for chunk {chunk['chunk_index']}", + citation=PaperResearchCitationDraft( + chunk_index=chunk["chunk_index"], + page_number=chunk["pages"][0], + ), + ) + for chunk in data["chunks"] + ), + ) + return PaperSummaryDraft( + summary="synthesized global summary across physical pages", + citations=( + PaperSummaryCitationDraft(chunk_index=0, page_number=1), + ), + ) + + return fake_run + + +class SummaryMapReduceTests(unittest.IsolatedAsyncioTestCase): + def test_chunk_groups_tile_every_chunk_exactly_once(self) -> None: + groups = _chunk_groups(7, 3) + self.assertEqual( + [(group.start, group.count) for group in groups], + [(0, 3), (3, 3), (6, 1)], + ) + covered = [ + index + for group in groups + for index in range(group.start, group.start + group.count) + ] + self.assertEqual(covered, list(range(7))) + + def test_research_finding_outside_its_group_is_rejected(self) -> None: + result = build_paper_result() + draft = PaperResearchDraft( + scope_summary="reviewed the first chunk only", findings=( PaperResearchFindingDraft( kind="method", - claim=f"Finding from chunk {request.start}.", + claim="a finding pointing outside the assigned group", citation=PaperResearchCitationDraft( - chunk_index=request.start, - page_number=page, + chunk_index=2, + page_number=2, ), ), ), ) + with self.assertRaisesRegex(ValueError, "outside its group"): + _validate_research_draft( + result.chunk_set, + _ChunkGroup(start=0, count=1), + draft, + ) - async def test_tool_call_and_input_budgets_are_enforced(self) -> None: - result = build_paper_result() - cfg = PaperFlowCfg( - max_summary_tool_calls=1, - max_summary_concurrency=1, - max_summary_input_tokens=10_000, - ) - budget = _SummaryBudget( - cfg, - "manifest", - chunk_count=len(result.chunk_set.chunks), - ) - first = PaperResearchRequest(start=0, count=1, focus="method") - - async def first_operation(): - return self._research_draft(result, first) - - await budget.invoke_research( - result.source_revision, - result.chunk_set, - first, - first_operation, + def test_worker_and_reducer_output_is_capped(self) -> None: + capped = _summary_model_settings( + PaperFlowCfg( + max_summary_output_tokens=256, + model_settings=ModelSettings(max_tokens=1024), + ) ) - second = PaperResearchRequest(start=1, count=1, focus="result") - with self.assertRaisesRegex( - PaperSummaryBudgetExceeded, - "tool_calls", - ): - await budget.invoke_research( - result.source_revision, - result.chunk_set, - second, - first_operation, + self.assertEqual(capped.max_tokens, 256) + lower = _summary_model_settings( + PaperFlowCfg( + max_summary_output_tokens=256, + model_settings=ModelSettings(max_tokens=128), ) + ) + self.assertEqual(lower.max_tokens, 128) - async def test_concurrent_subagents_cover_every_chunk(self) -> None: + async def test_map_reduce_fans_out_one_worker_per_group(self) -> None: result = build_paper_result() cfg = PaperFlowCfg( - max_summary_tool_calls=3, - max_summary_concurrency=2, - max_summary_input_tokens=10_000, - ) - budget = _SummaryBudget( - cfg, - "manifest", - chunk_count=len(result.chunk_set.chunks), + summary_research_group_size=1, + summary_concurrency=2, + min_summary_citations=1, + min_summary_pages=1, ) - - async def run(index: int) -> str: - request = PaperResearchRequest( - start=index, - count=1, - focus=f"chunk {index}", + run_mock = AsyncMock( + side_effect=_fake_research_run( + cfg, result.source_revision.parsed.pages ) - - async def operation(): - await asyncio.sleep(0) - return self._research_draft(result, request) - - return await budget.invoke_research( + ) + with patch( + "quantmind.flows._paper_summary.run_with_observability", + new=run_mock, + ): + draft = await _AgentsPaperSummaryProvider().summarize( result.source_revision, result.chunk_set, - request, - operation, + cfg=cfg, ) - values = await asyncio.gather( - *(run(index) for index in range(len(result.chunk_set.chunks))) - ) - self.assertEqual(len(values), 3) - self.assertEqual(budget.research_calls, 3) - self.assertEqual(budget.covered_chunks, frozenset({0, 1, 2})) - budget.require_complete_coverage(len(result.chunk_set.chunks)) - - async def test_missing_subagent_coverage_is_rejected(self) -> None: - result = build_paper_result() - cfg = PaperFlowCfg(max_summary_input_tokens=10_000) - budget = _SummaryBudget( - cfg, - "manifest", - chunk_count=len(result.chunk_set.chunks), - ) - - with self.assertRaisesRegex( - PaperSummaryBudgetExceeded, - "did not cover every chunk", - ): - budget.require_complete_coverage(len(result.chunk_set.chunks)) + self.assertIsInstance(draft, PaperSummaryDraft) + worker_calls = [ + call.args[0].name + for call in run_mock.await_args_list + if call.args[0].name == "paper_chunk_researcher" + ] + self.assertEqual(len(worker_calls), len(result.chunk_set.chunks)) + self.assertEqual(run_mock.await_count, len(result.chunk_set.chunks) + 1) - async def test_coordinator_without_delegation_is_rejected(self) -> None: + async def test_summary_timeout_raises(self) -> None: result = build_paper_result() - provider = _AgentsPaperSummaryProvider() - draft = PaperSummaryDraft( - summary="A summary that skipped its research workers.", - citations=( - PaperSummaryCitationDraft(chunk_index=0, page_number=1), - ), - ) cfg = PaperFlowCfg( + summary_research_group_size=8, + timeout_seconds=0.01, min_summary_citations=1, min_summary_pages=1, ) + async def fake_run(agent, payload, *, cfg, memory, extra_run_hooks): + data = json.loads(payload) + if agent.name == "paper_chunk_researcher": + return PaperResearchDraft( + scope_summary="reviewed the assigned chunk group", + findings=tuple( + PaperResearchFindingDraft( + kind="method", + claim=f"finding for chunk {chunk['chunk_index']}", + citation=PaperResearchCitationDraft( + chunk_index=chunk["chunk_index"], + page_number=chunk["pages"][0], + ), + ) + for chunk in data["chunks"] + ), + ) + await asyncio.sleep(10) + with patch( "quantmind.flows._paper_summary.run_with_observability", - new=AsyncMock(return_value=draft), + new=AsyncMock(side_effect=fake_run), ): - with self.assertRaisesRegex( - PaperSummaryBudgetExceeded, - "did not cover every chunk", - ): - await provider.summarize( + with self.assertRaises(PaperSummaryError): + await _AgentsPaperSummaryProvider().summarize( result.source_revision, result.chunk_set, cfg=cfg, ) - def test_researcher_is_an_agents_sdk_agent_tool(self) -> None: - result = build_paper_result() - cfg = PaperFlowCfg() - budget = _SummaryBudget( - cfg, - "manifest", - chunk_count=len(result.chunk_set.chunks), - ) - - tool = _build_research_agent_tool( - result.source_revision, - result.chunk_set, - cfg, - budget, - ) - - self.assertEqual(tool.name, "research_chunk_group") - self.assertTrue(tool._is_agent_tool) - self.assertEqual( - set(tool.params_json_schema["properties"]), - {"start", "count", "focus"}, - ) - - def test_model_output_tokens_are_capped(self) -> None: - cfg = PaperFlowCfg( - max_summary_output_tokens=256, - max_summary_worker_output_tokens=128, - model_settings=ModelSettings(max_tokens=1024), - ) - - settings = _bounded_model_settings(cfg) - research_settings = _bounded_research_model_settings(cfg) - - self.assertEqual(settings.max_tokens, 256) - self.assertTrue(settings.parallel_tool_calls) - self.assertEqual(research_settings.max_tokens, 128) - self.assertFalse(research_settings.parallel_tool_calls) - if __name__ == "__main__": unittest.main() diff --git a/tests/paper_helpers.py b/tests/paper_helpers.py index d62d5bd..7e021ca 100644 --- a/tests/paper_helpers.py +++ b/tests/paper_helpers.py @@ -1,37 +1,36 @@ -"""Deterministic paper artifacts shared by offline tests.""" +"""Deterministic paper artifacts shared by offline tests. + +The fixture builds through the same knowledge-layer smart constructors that +``paper_flow`` uses, so tests exercise the real identity path instead of +re-deriving IDs and hashes with the private helpers. +""" import hashlib from datetime import datetime, timezone from quantmind.knowledge import ( - ArtifactLocator, - PaperAssetRef, - PaperChunk, PaperChunkingConfig, + PaperChunkInput, PaperChunkSet, - PaperCitation, + PaperCitationDraft, PaperFlowResult, PaperGlobalSummary, - PaperParsedManifest, - PaperParsedPage, + PaperPageInput, + PaperSourceFacts, PaperSourceRevision, - PaperSourceSpan, PaperSummaryProducer, - SourceRef, -) -from quantmind.knowledge.paper import ( - _paper_artifact_id, - _paper_asset_id, - _paper_chunk_id, - _paper_chunk_set_content_hash, - _paper_source_id, - _paper_summary_content_hash, - _stable_hash, - _text_hash, ) _RAW_BYTES = b"deterministic paper source revision" _WHEN = datetime(2017, 12, 6, tzinfo=timezone.utc) +_PAGE_ONE = ( + "The Transformer removes recurrence and convolution " + "from sequence transduction." +) +_PAGE_TWO = ( + "Multi-head attention uses parallel projections. " + "The model improves translation and training efficiency." +) def build_paper_result( @@ -46,152 +45,78 @@ def build_paper_result( ) -> PaperFlowResult: """Build one valid two-page result without parsing, network, or models.""" source_hash = hashlib.sha256(_RAW_BYTES).hexdigest() - source_id = _paper_source_id(source_hash) - raw_asset = PaperAssetRef( - asset_id=_paper_asset_id( - source_id, - kind="raw", - page_number=None, - content_hash=source_hash, - ), - kind="raw", - media_type="application/pdf", - content_hash=source_hash, - size_bytes=len(_RAW_BYTES), - ) - source = PaperSourceRevision( - id=source_id, - source=SourceRef( + source = PaperSourceRevision.from_parsed( + facts=PaperSourceFacts( kind="arxiv", uri="https://arxiv.org/pdf/1706.03762v7.pdf", + media_type="application/pdf", + raw_bytes=_RAW_BYTES, fetched_at=_WHEN, - content_hash=source_hash, + available_at=_WHEN, + published_at=_WHEN, + arxiv_id="1706.03762v7", + title="Attention Is All You Need", + authors=("Ashish Vaswani",), ), - as_of=_WHEN, - available_at=_WHEN, - published_at=_WHEN, - arxiv_id="1706.03762v7", - title="Attention Is All You Need", - authors=("Ashish Vaswani",), - parsed=PaperParsedManifest( - source_hash=source_hash, - parser_name="fake-parser", - parser_version="1", - cleanup_version="1", - pages=( - PaperParsedPage( - page_number=1, - width=612, - height=792, - text=( - "The Transformer removes recurrence and convolution " - "from sequence transduction." - ), - ), - PaperParsedPage( - page_number=2, - width=612, - height=792, - text=( - "Multi-head attention uses parallel projections. " - "The model improves translation and training efficiency." - ), - ), + source_hash=source_hash, + parser_name="fake-parser", + parser_version="1", + cleanup_version="1", + pages=( + PaperPageInput( + page_number=1, width=612, height=792, text=_PAGE_ONE + ), + PaperPageInput( + page_number=2, width=612, height=792, text=_PAGE_TWO ), ), - raw_asset_id=raw_asset.asset_id, - assets=(raw_asset,), - blobs={source_hash: _RAW_BYTES}, ) - producer = PaperChunkingConfig( - splitter_version="fake-llama-index", - chunk_size=chunk_size, - chunk_overlap=min(16, chunk_size - 1), - ) - producer_hash = _stable_hash(producer.model_dump(mode="json")) - chunk_set_id = _paper_artifact_id( - source_id, - "paper_chunk_set", - producer_hash, - ) chunk_values = ( - ( - 1, - "The Transformer removes recurrence and convolution.", - ), + (1, "The Transformer removes recurrence and convolution."), (2, "Multi-head attention uses parallel learned projections."), (2, "The model improves translation and training efficiency."), ) - chunks: list[PaperChunk] = [] - for position, (page, text) in enumerate(chunk_values): - span = PaperSourceSpan( + chunk_inputs = tuple( + PaperChunkInput( page_number=page, start_char=position * 10, end_char=position * 10 + len(text), + block_boxes=(), + text=text, ) - text_hash = _text_hash(text) - chunks.append( - PaperChunk( - chunk_id=_paper_chunk_id( - chunk_set_id, - position=position, - content_hash=text_hash, - spans=(span,), - ), - chunk_set_id=chunk_set_id, - source_revision_id=source_id, - position=position, - text=text, - content_hash=text_hash, - source_spans=(span,), - ) - ) - chunk_tuple = tuple(chunks) - chunk_set = PaperChunkSet( - id=chunk_set_id, - source_revision_id=source_id, - producer=producer, - producer_config_hash=producer_hash, - content_hash=_paper_chunk_set_content_hash(chunk_tuple), - chunks=chunk_tuple, + for position, (page, text) in enumerate(chunk_values) ) - - summary_producer = PaperSummaryProducer( - model=summary_model, - prompt_version="test-v1", - input_chunk_set_id=chunk_set_id, - instructions_hash=hashlib.sha256(b"test instructions").hexdigest(), - max_output_tokens=512, - ) - summary_config_hash = _stable_hash(summary_producer.model_dump(mode="json")) - citations = tuple( - PaperCitation( - chunk_set_id=chunk_set_id, - chunk_id=chunk.chunk_id, - page_number=chunk.source_spans[0].page_number, - ) - for chunk in chunk_tuple + chunk_set = PaperChunkSet.from_parsed_chunks( + source, + chunk_inputs, + producer=PaperChunkingConfig( + splitter_version="fake-llama-index", + chunk_size=chunk_size, + chunk_overlap=min(16, chunk_size - 1), + ), ) - summary = PaperGlobalSummary( - id=_paper_artifact_id( - source_id, - "paper_summary", - summary_config_hash, + + summary = PaperGlobalSummary.from_draft( + chunk_set, + producer=PaperSummaryProducer( + model=summary_model, + prompt_version="test-v1", + input_chunk_set_id=chunk_set.id, + instructions_hash=hashlib.sha256(b"test instructions").hexdigest(), + max_output_tokens=512, + research_group_size=8, ), - source_revision_id=source_id, - producer=summary_producer, - producer_config_hash=summary_config_hash, - content_hash=_paper_summary_content_hash(summary_text, citations), summary=summary_text, - citations=citations, - derived_from=( - ArtifactLocator( - source_revision_id=source_id, - artifact_id=chunk_set_id, - artifact_kind="paper_chunk_set", - ), + citations=tuple( + PaperCitationDraft( + chunk_index=index, + page_number=chunk.source_spans[0].page_number, + ) + for index, chunk in enumerate(chunk_set.chunks) ), + min_citations=1, + min_pages=1, ) return PaperFlowResult( source_revision=source, diff --git a/tests/test_verify_pdf_rag_e2e.py b/tests/test_verify_pdf_rag_e2e.py index 92f0720..175e8bc 100644 --- a/tests/test_verify_pdf_rag_e2e.py +++ b/tests/test_verify_pdf_rag_e2e.py @@ -20,7 +20,7 @@ def _snapshot(**overrides): "convolution, uses multi-head attention, and improves translation " "with greater training efficiency." ), - "summary_orchestration": "manager-research-agents-v1", + "summary_orchestration": "map-reduce-v1", "citation_count": 3, "citation_pages": [2, 4], "first_summary_scores": [0.9],