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..74cf989 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,17 @@ jobs: - name: Install project runtime dependencies run: uv pip install --python .venv/bin/python -e . - - name: Run live PDF RAG E2E + - 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/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..d055829 100644 --- a/contexts/design/README.md +++ b/contexts/design/README.md @@ -24,12 +24,14 @@ 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) | | 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 de1b509..63cbda8 100644 --- a/contexts/design/flow/paper.md +++ b/contexts/design/flow/paper.md @@ -1,360 +1,155 @@ -# 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. 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`. -The planned pipeline accepts the existing `PaperInput` types with the following -behavior: +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). -| 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. | +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. -`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. +## Source Revision -### Explicitly unsupported inputs +`PaperSourceRevision` owns the exact fetched bytes and deterministic parser result. Its ID is derived from the source SHA-256 hash. It records: -V1 does not claim support for: +- 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. -- 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. +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. -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. +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. -### Fetching rules +## Chunk Set -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. +`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. -## Keep PDF Pages Separate +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. -PDF parsing must preserve ordered pages before any tree builder runs. The next -step receives data shaped like this: +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. -```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 -``` +## Cited Global Summary -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. +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. -## Retries and Failures by Step +Code accepts the draft only when: -| 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. | +- 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 `min_summary_citations`; +- distinct cited-page count meets `min_summary_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`. +`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. -## What Paper Extraction Does Not Do +## Bounded Model Calls -Paper extraction returns one validated `Paper`. Other components handle: +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. -- 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. +`PaperFlowCfg` exposes only structural bounds: -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. +- `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. -## How PageIndex Can Fit Later - -Future integration must preserve these decisions: +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. -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. +## Failure Semantics -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. +- 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`. +- 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. -## Fixed Paper Test Data +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. -The fixed test files live at: +## Persistence and Retrieval -```text -tests/fixtures/paper/golden/ -├── paper.pdf -└── expected.json +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. + +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. + +## 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, 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 + +- 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..e1182a9 --- /dev/null +++ b/contexts/design/knowledge/paper.md @@ -0,0 +1,95 @@ +# 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; +- map-reduce orchestration version; +- exact input chunk-set ID; +- 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. + +## 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..f698ef4 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` 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. -- `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/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/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..dbf8aa3 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,16 @@ 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. +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 Run the deterministic required verification for every change: diff --git a/docs/library.md b/docs/library.md index 3323d34..2cb804f 100644 --- a/docs/library.md +++ b/docs/library.md @@ -1,64 +1,98 @@ # 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 +from quantmind.knowledge import PaperArtifactKind + +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=[PaperArtifactKind.GLOBAL_SUMMARY], + top_k=3, + ) + ) + chunk_hits = await library.search( + SemanticQuery( + text="How does multi-head attention work?", + artifact_kinds=[PaperArtifactKind.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 -Representative output has this shape; scores depend on the selected embedding -model: +`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. -```text -Bundle: .../examples/library/data/ai_infrastructure.db -Query: What evidence shows demand for AI infrastructure is expanding? +`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. -1. score=0.812 type=paper - path: ... > ... - matched: ... - source: https://... - citation: ... +## Bundled Compatibility Example + +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 ``` -`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. +`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 +``` + +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..2c1e4c4 --- /dev/null +++ b/examples/flows/paper.py @@ -0,0 +1,154 @@ +"""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 ( + PaperArtifactKind, + 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=[PaperArtifactKind.GLOBAL_SUMMARY], + top_k=3, + ) + ) + chunk_hits = await library.search( + SemanticQuery( + text="How does multi-head attention work?", + artifact_kinds=[PaperArtifactKind.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( + "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)}" + ) + 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 10f4eed..6b8340b 100644 Binary files a/examples/library/data/ai_infrastructure.db and b/examples/library/data/ai_infrastructure.db differ diff --git a/quantmind/configs/paper.py b/quantmind/configs/paper.py index 8f4878b..ed6e842 100644 --- a/quantmind/configs/paper.py +++ b/quantmind/configs/paper.py @@ -1,17 +1,14 @@ """Paper-flow configuration + input discriminated union. -`PaperInput` is one of: -- `ArxivIdentifier`: arxiv id or full URL parsed by preprocess.fetch.arxiv -- `HttpUrl`: any web URL (PDF or HTML; routed by content-type) -- `LocalFilePath`: filesystem path to a PDF / HTML / Markdown file -- `RawText`: an inline string (for tests or LLM-pre-cleaned inputs) -- `DoiIdentifier`: a DOI to be resolved via preprocess.fetch.doi +`PaperInput` describes supported and reserved source variants. Paper Flow V1 +accepts PDF-backed arXiv, HTTP, and local inputs. It rejects non-PDF HTTP/local +content and raw text, and reserves DOI input until an exact PDF resolver exists. """ from pathlib import Path from typing import Annotated, Literal, Union -from pydantic import Field +from pydantic import Field, model_validator from quantmind.configs.base import BaseFlowCfg, BaseInput @@ -24,21 +21,21 @@ class ArxivIdentifier(BaseInput): class HttpUrl(BaseInput): - """Any web URL; PDF vs HTML is decided by content-type.""" + """A web URL that must resolve to a PDF in Paper Flow V1.""" type: Literal["http"] = "http" url: str class LocalFilePath(BaseInput): - """Filesystem path to a PDF / HTML / Markdown file.""" + """Filesystem path to a PDF for Paper Flow V1.""" type: Literal["local"] = "local" path: Path class RawText(BaseInput): - """Inline text input (tests / pre-cleaned content).""" + """Reserved inline text input rejected by page-aware Paper Flow V1.""" type: Literal["text"] = "text" text: str @@ -58,8 +55,34 @@ class DoiIdentifier(BaseInput): class PaperFlowCfg(BaseFlowCfg): - """Knobs specific to paper_flow.""" - - extract_methodology: bool = True - extract_limitations: bool = True - asset_class_hint: str | None = None + """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-v3" + summary_instructions: str | None = None + 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) + + @model_validator(mode="after") + def _validate_paper_bounds(self) -> "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..dbd53be 100644 --- a/quantmind/flows/__init__.py +++ b/quantmind/flows/__init__.py @@ -7,16 +7,18 @@ - ``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.knowledge import PaperCitationValidationError __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..c7738bd --- /dev/null +++ b/quantmind/flows/_paper_summary.py @@ -0,0 +1,339 @@ +"""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 Sequence +from dataclasses import dataclass, replace +from typing import Any, Literal, Protocol + +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 + +_ORCHESTRATION_VERSION = "map-reduce-v1" + +_SUMMARY_INSTRUCTIONS = """\ +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. 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 PaperSummaryError(RuntimeError): + """The summary exceeded a configured runtime boundary.""" + + +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 + quote: str | None = Field(default=None, max_length=500) + + @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 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: + stripped = value.strip() + if not stripped: + raise ValueError("paper summary draft must not be blank") + return stripped + + +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.""" + ... + + +@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, + group: _ChunkGroup, +) -> str: + selected = chunk_set.chunks[group.start : group.start + group.count] + return json.dumps( + { + "title": source.title, + "authors": source.authors, + "start": group.start, + "count": group.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 _validate_research_draft( + chunk_set: PaperChunkSet, + group: _ChunkGroup, + draft: PaperResearchDraft, +) -> None: + 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 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") + 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 _reduce_payload( + source: PaperSourceRevision, + chunk_set: PaperChunkSet, + reports: Sequence[PaperResearchDraft], +) -> str: + return json.dumps( + { + "title": source.title, + "authors": source.authors, + "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, + ) + + +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: + payload = json.dumps( + { + "orchestration": _ORCHESTRATION_VERSION, + "reducer_instructions": _summary_instructions(cfg), + "research_instructions": _RESEARCH_INSTRUCTIONS, + "research_group_size": cfg.summary_research_group_size, + "max_output_tokens": cfg.max_summary_output_tokens, + }, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +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), + ) + + +class _AgentsPaperSummaryProvider: + """Fan out one research agent per chunk group, then reduce to a summary.""" + + async def summarize( + self, + source: PaperSourceRevision, + chunk_set: PaperChunkSet, + *, + cfg: PaperFlowCfg, + ) -> PaperSummaryDraft: + groups = _chunk_groups( + len(chunk_set.chunks), cfg.summary_research_group_size + ) + 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, + ) + 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=model_settings, + output_type=PaperSummaryDraft, + ) + try: + output = await asyncio.wait_for( + run_with_observability( + reducer, + _reduce_payload(source, chunk_set, reports), + cfg=cfg, + memory=None, + extra_run_hooks=[], + ), + timeout=cfg.timeout_seconds, + ) + except asyncio.TimeoutError as exc: + raise PaperSummaryError( + "paper summary exceeded timeout_seconds" + ) from exc + return PaperSummaryDraft.model_validate(output) diff --git a/quantmind/flows/paper.py b/quantmind/flows/paper.py index 5b8373b..d17fcc1 100644 --- a/quantmind/flows/paper.py +++ b/quantmind/flows/paper.py @@ -1,19 +1,19 @@ -"""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). +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. """ -from typing import Any, TypeVar - -from agents import Agent, RunHooks, Tool +import mimetypes +from collections.abc import Sequence +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 +24,311 @@ 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 ( + PaperAssetInput, + PaperBoundingBox, + PaperChunkingConfig, + PaperChunkInput, + PaperChunkSet, + PaperCitationDraft, + PaperFlowResult, + PaperGlobalSummary, + PaperPageInput, + PaperParsedBlock, + PaperSourceFacts, + PaperSourceRevision, + PaperSummaryProducer, +) 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 - -P = TypeVar("P", bound=Paper) - -_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. - -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. -""" +from quantmind.preprocess.format import ParsedDocument, parse_pdf +from quantmind.rag import ( + ParsedChunk, + SentenceSplitterConfig, + chunk_parsed_document, +) class UnsupportedContentTypeError(ValueError): - """Fetched bytes have a content type paper_flow cannot route to a parser.""" + """The source is not a page-aware PDF supported by Paper Flow V1.""" 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 + 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 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 the configured source-coverage policy. + 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 + 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), + ) + 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 []), + ) + 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_summary(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) -> PaperSourceFacts: + """Fetch and normalize one input into code-owned source facts (IO).""" 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 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, + 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 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, + ) 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 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, + ) 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") - ) - 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}" +def _read_page_asset( + path_value: str, + *, + kind: Literal["screenshot", "image"], + page_number: int, +) -> PaperAssetInput: + """Read one parser-written page asset off disk into knowledge input (IO).""" + 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 + media_type = ( + mimetypes.guess_type(path.name)[0] or "application/octet-stream" ) + return PaperAssetInput(kind=kind, content=content, media_type=media_type) -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 _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 = ( + _read_page_asset( + page.screenshot_path, + kind="screenshot", + page_number=page.page_number, + ) + if page.screenshot_path is not None + else None + ) + images = tuple( + _read_page_asset( + image_path, + kind="image", + page_number=page.page_number, + ) + for image_path in page.image_paths + ) + pages.append( + PaperPageInput( + page_number=page.page_number, + width=page.width, + height=page.height, + text=page.text, + blocks=blocks, + screenshot=screenshot, + images=images, + ) + ) + return pages + + +def _adapt_chunks( + parsed_chunks: Sequence[ParsedChunk], + source: PaperSourceRevision, +) -> 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") + chunks.append( + 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 + ), + text=parsed_chunk.text, + ) + ) + return chunks + + +def _build_summary( + chunk_set: PaperChunkSet, + draft: PaperSummaryDraft, + cfg: PaperFlowCfg, +) -> PaperGlobalSummary: + """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, + research_group_size=cfg.summary_research_group_size, + ) + citations = tuple( + PaperCitationDraft( + chunk_index=citation.chunk_index, + page_number=citation.page_number, + quote=citation.quote, + ) + for citation in draft.citations ) - 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: - 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}" + return PaperGlobalSummary.from_draft( + chunk_set, + producer=producer, + 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 972a535..37bd5ac 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,33 @@ 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, + PaperAssetInput, + PaperAssetRef, + PaperBoundingBox, + PaperChunk, + PaperChunkingConfig, + PaperChunkInput, + PaperChunkSet, + PaperCitation, + PaperCitationDraft, + PaperCitationValidationError, + PaperFlowResult, + PaperGlobalSummary, + PaperPageInput, + PaperParsedBlock, + PaperParsedManifest, + PaperParsedPage, + PaperSourceFacts, + PaperSourceRevision, + PaperSourceSpan, + PaperSummaryProducer, + ResolvedPaperArtifact, +) from quantmind.knowledge.thesis import Thesis __all__ = [ @@ -42,7 +68,30 @@ "Earnings", "Factor", "News", - "Paper", - "PaperKnowledgeCard", + "ArtifactLocator", + "LegacyPaper", + "PaperArtifact", + "PaperArtifactKind", + "PaperAssetInput", + "PaperAssetRef", + "PaperBoundingBox", + "PaperChunk", + "PaperChunkInput", + "PaperChunkingConfig", + "PaperChunkSet", + "PaperCitation", + "PaperCitationDraft", + "PaperCitationValidationError", + "PaperFlowResult", + "PaperGlobalSummary", + "PaperPageInput", + "PaperParsedBlock", + "PaperParsedManifest", + "PaperParsedPage", + "PaperSourceFacts", + "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..0064536 100644 --- a/quantmind/knowledge/paper.py +++ b/quantmind/knowledge/paper.py @@ -1,52 +1,930 @@ -"""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 collections.abc import Sequence +from dataclasses import dataclass +from datetime import datetime +from enum import Enum 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 -class Paper(TreeKnowledge): - """A research paper as a tree of sections. +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: + """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 | str, + producer_config_hash: str, +) -> UUID: + kind = PaperArtifactKind(artifact_kind) + return uuid5( + source_revision_id, + f"quantmind:{kind.value}:{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 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. - 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. + 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. """ - item_type: Literal["paper"] = "paper" + 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 - authors: list[str] = Field(default_factory=list) - asset_classes: list[str] = Field(default_factory=list) + 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.""" + + 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 + + +@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.""" + model_config = ConfigDict(extra="forbid", frozen=True) -class PaperKnowledgeCard(FlattenKnowledge): - """Distilled summary card of a `Paper`. + 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) - 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. + +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 + + +@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.""" + + 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_card"] = "paper_card" + 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 + 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 + + @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") + + 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 + + @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.""" + + 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) + 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.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + 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[PaperArtifactKind.CHUNK_SET] = ( + PaperArtifactKind.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 + + @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, + 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) + + +@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.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + model: str + prompt_version: str + 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) + research_group_size: int = Field(ge=1) + + +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[PaperArtifactKind.GLOBAL_SUMMARY] = ( + PaperArtifactKind.GLOBAL_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 + + @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.""" + + 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..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 Citation, SourceRef +from quantmind.knowledge import ( + ArtifactLocator, + Citation, + PaperArtifactKind, + SourceRef, +) class SemanticQuery(BaseModel): @@ -15,6 +20,7 @@ class SemanticQuery(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) text: str = Field(min_length=1) + artifact_kinds: list[PaperArtifactKind] | None = None item_types: list[str] | None = None source_kinds: ( list[ @@ -55,11 +61,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..8eeb54c 100644 --- a/scripts/verify_pdf_rag_e2e.py +++ b/scripts/verify_pdf_rag_e2e.py @@ -1,54 +1,251 @@ #!/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 ( + PaperArtifactKind, + 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=[PaperArtifactKind.GLOBAL_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=[PaperArtifactKind.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, + summary_research_group_size=8, + summary_concurrency=2, + 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, + "summary_orchestration": (result.global_summary.producer.orchestration), + "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["summary_orchestration"] == "map-reduce-v1" + 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"orchestration={snapshot['summary_orchestration']} " + 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..6247521 100644 --- a/tests/configs/test_paper.py +++ b/tests/configs/test_paper.py @@ -19,10 +19,21 @@ 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.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) class PaperInputDiscriminatedTests(unittest.TestCase): diff --git a/tests/flows/test_paper.py b/tests/flows/test_paper.py index 2a143c7..8a48ecc 100644 --- a/tests/flows/test_paper.py +++ b/tests/flows/test_paper.py @@ -1,13 +1,13 @@ -"""Tests for ``quantmind.flows.paper``.""" +"""Offline tests for source-first ``paper_flow`` behavior.""" +import asyncio +import json 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 +17,435 @@ LocalFilePath, RawText, ) +from quantmind.flows._paper_summary import ( + PaperResearchCitationDraft, + PaperResearchDraft, + PaperResearchFindingDraft, + PaperSummaryCitationDraft, + PaperSummaryDraft, + PaperSummaryError, + _AgentsPaperSummaryProvider, + _chunk_groups, + _ChunkGroup, + _summary_model_settings, + _validate_research_draft, +) from quantmind.flows.paper import ( UnsupportedContentTypeError, - _compose_instructions, - _fetch_and_format, - _format_by_content_type, - _format_input, + _build_summary, + _fetch_paper_source, paper_flow, ) -from quantmind.knowledge import Paper, SourceRef, TreeNode +from quantmind.knowledge import PaperCitationValidationError from quantmind.preprocess.fetch import Fetched, RawPaper +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 + ), + ) -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 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, + ) + 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 + } + ), + 2, + ) + self.assertTrue( + all( + chunk.source_revision_id == result.source_revision.id + and chunk.source_spans + for chunk in result.chunk_set.chunks + ) + ) -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", + 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.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", + "quantmind.flows.paper.fetch_arxiv", + new=AsyncMock(return_value=raw), + ): + result = await paper_flow( + ArxivIdentifier(id="1706.03762v7"), + cfg=PaperFlowCfg(chunk_size=256, chunk_overlap=32), + _summary_provider=_FakeSummaryProvider(), + ) + + 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" + ) + self.assertEqual( + result.source_revision.source.content_hash, + result.source_revision.parsed.source_hash, ) - md = await _format_by_content_type(raw) - self.assertEqual(md, "# heading\n\nbody") + self.assertFalse(hasattr(result.global_summary, "root_node_id")) - 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_summary_failure_prevents_flow_success(self) -> None: + provider = _FakeSummaryProvider(fail=RuntimeError("summary failed")) - 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) + with self.assertRaisesRegex(RuntimeError, "summary failed"): + await paper_flow( + LocalFilePath(path=_FIXTURE), + cfg=PaperFlowCfg(chunk_size=256, chunk_overlap=32), + _summary_provider=provider, + ) + 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(), + ) -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"), + 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], ) - 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: + + +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", - source_url="http://example/x.pdf", - ) - 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"), - ), - ): - md, meta = await _fetch_and_format( - HttpUrl(url="http://example/x.pdf") - ) - 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", + source_url="https://example.test/input.pdf", + resolved_url="https://cdn.example.test/exact.pdf", + fetched_at=_WHEN, ) with patch( - "quantmind.flows.paper.read_local_file", + "quantmind.flows.paper.fetch_url", new=AsyncMock(return_value=raw), - ) as read_mock: - md, meta = await _fetch_and_format( - LocalFilePath(path=Path("/tmp/p.md")) + ): + fetched = await _fetch_paper_source( + HttpUrl(url="https://example.test/input.pdf") ) - 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"}) + self.assertEqual(fetched.uri, "https://cdn.example.test/exact.pdf") + self.assertEqual(fetched.fetched_at, _WHEN) - 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)) + 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), + ): + 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_doi_requires_an_exact_pdf_resolver(self) -> None: + with self.assertRaises(NotImplementedError): + await _fetch_paper_source(DoiIdentifier(doi="10.1000/test")) -class ComposeInstructionsTests(unittest.TestCase): - def test_default_renders_cfg_flags(self) -> None: + +class CitationValidationTests(unittest.TestCase): + def test_unknown_chunk_page_and_quote_are_rejected(self) -> None: + result = build_paper_result() cfg = PaperFlowCfg( - extract_methodology=False, - extract_limitations=True, - asset_class_hint="equities", + min_summary_citations=1, + min_summary_pages=1, ) - out = _compose_instructions( - "go {extract_methodology} {extract_limitations} " - "{asset_class_hint!r}", - None, - cfg, + 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", + ), + ), + ), ) - self.assertEqual(out, "go False True 'equities'") - - 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) - -class FormatInputTests(unittest.TestCase): - def test_tuple_authors_join_as_csv(self) -> None: - out = _format_input( - "BODY", - {"authors": ("Alice", "Bob"), "title": "x"}, + for draft in invalid_drafts: + with self.subTest(draft=draft): + with self.assertRaises(PaperCitationValidationError): + _build_summary( + 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), + ), ) - 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) + with self.assertRaisesRegex( + PaperCitationValidationError, + "fewer citations", + ): + _build_summary( + result.chunk_set, + draft, + PaperFlowCfg(), + ) -class PaperFlowTests(unittest.IsolatedAsyncioTestCase): - async def test_happy_path_arxiv(self) -> None: - raw_paper = RawPaper( - bytes=b"%PDF", - content_type="application/pdf", - arxiv_id="2604.12345", - ) - 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"), +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), ), - _patch_runner(stub) as runner, - ): - 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] = {} + return fake_run - def _capture_agent(*_a: Any, **kwargs: Any) -> Any: - seen.update(kwargs) - return MagicMock(name="agent", _name="paper_extractor") - stub = _stub_paper() - with ( - patch("quantmind.flows.paper.Agent", side_effect=_capture_agent), - _patch_runner(stub), - ): - await paper_flow( - RawText(text="hello"), - extra_instructions="EXTRA-USER-DIRECTIVE", +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="a finding pointing outside the assigned group", + citation=PaperResearchCitationDraft( + 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, ) - self.assertIn("EXTRA-USER-DIRECTIVE", seen["instructions"]) - self.assertIn("structured QuantMind", seen["instructions"]) - - async def test_output_type_override_propagated(self) -> None: - seen: dict[str, Any] = {} - - class MyPaper(Paper): - pass - def _capture_agent(*_a: Any, **kwargs: Any) -> Any: - seen.update(kwargs) - return MagicMock() + 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), + ) + ) + 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) - 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], + async def test_map_reduce_fans_out_one_worker_per_group(self) -> None: + result = build_paper_result() + cfg = PaperFlowCfg( + summary_research_group_size=1, + summary_concurrency=2, + min_summary_citations=1, + min_summary_pages=1, + ) + run_mock = AsyncMock( + side_effect=_fake_research_run( + cfg, result.source_revision.parsed.pages ) - 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(), - ), - _patch_runner(_stub_paper()) as runner, - ): - 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, + ) + with patch( + "quantmind.flows._paper_summary.run_with_observability", + new=run_mock, ): - await paper_flow(RawText(text="x"), extra_run_hooks=[hook]) - self.assertEqual(runner.await_args.kwargs["extra_run_hooks"], [hook]) - - async def test_model_settings_forwarded_when_set(self) -> None: - seen: dict[str, Any] = {} + draft = await _AgentsPaperSummaryProvider().summarize( + result.source_revision, + result.chunk_set, + cfg=cfg, + ) - def _capture_agent(*_a: Any, **kwargs: Any) -> Any: - seen.update(kwargs) - return MagicMock() + 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_summary_timeout_raises(self) -> None: + result = build_paper_result() + cfg = PaperFlowCfg( + summary_research_group_size=8, + timeout_seconds=0.01, + min_summary_citations=1, + min_summary_pages=1, + ) - from agents import ModelSettings + 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) - ms = ModelSettings(temperature=0.42) - cfg = PaperFlowCfg(model_settings=ms) - with ( - patch("quantmind.flows.paper.Agent", side_effect=_capture_agent), - _patch_runner(_stub_paper()), + with patch( + "quantmind.flows._paper_summary.run_with_observability", + new=AsyncMock(side_effect=fake_run), ): - await paper_flow(RawText(text="x"), cfg=cfg) - self.assertIs(seen["model_settings"], ms) + with self.assertRaises(PaperSummaryError): + await _AgentsPaperSummaryProvider().summarize( + result.source_revision, + result.chunk_set, + cfg=cfg, + ) - async def test_model_settings_omitted_when_none(self) -> None: - seen: dict[str, Any] = {} - 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..183b8df --- /dev/null +++ b/tests/library/test_paper.py @@ -0,0 +1,412 @@ +"""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 ( + PaperArtifactKind, + 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=[PaperArtifactKind.GLOBAL_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=[PaperArtifactKind.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..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 @@ -11,7 +12,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): @@ -29,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/paper_helpers.py b/tests/paper_helpers.py new file mode 100644 index 0000000..7e021ca --- /dev/null +++ b/tests/paper_helpers.py @@ -0,0 +1,125 @@ +"""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 ( + PaperChunkingConfig, + PaperChunkInput, + PaperChunkSet, + PaperCitationDraft, + PaperFlowResult, + PaperGlobalSummary, + PaperPageInput, + PaperSourceFacts, + PaperSourceRevision, + PaperSummaryProducer, +) + +_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( + *, + 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 = 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, + available_at=_WHEN, + published_at=_WHEN, + arxiv_id="1706.03762v7", + title="Attention Is All You Need", + authors=("Ashish Vaswani",), + ), + 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 + ), + ), + ) + + 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."), + ) + chunk_inputs = tuple( + PaperChunkInput( + page_number=page, + start_char=position * 10, + end_char=position * 10 + len(text), + block_boxes=(), + text=text, + ) + for position, (page, text) in enumerate(chunk_values) + ) + 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.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, + ), + summary=summary_text, + 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, + 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..175e8bc 100644 --- a/tests/test_verify_pdf_rag_e2e.py +++ b/tests/test_verify_pdf_rag_e2e.py @@ -1,91 +1,88 @@ 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", + "summary_orchestration": "map-reduce-v1", + "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 +90,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__":