diff --git a/docs/design/en/2026-07-07-paper-flow-draft-assembly.md b/docs/design/en/2026-07-07-paper-flow-draft-assembly.md new file mode 100644 index 0000000..c803063 --- /dev/null +++ b/docs/design/en/2026-07-07-paper-flow-draft-assembly.md @@ -0,0 +1,190 @@ +# Design: paper_flow draft-schema + code-assembled UUIDs + +Date: 2026-07-07 +Status: Approved (pending spec review) +Scope: `quantmind/flows/paper.py`, new `quantmind/flows/_paper_draft.py`, tests + +## Problem + +`paper_flow` could not produce a `Paper` end-to-end against a live model. +Three failure modes were found by actually running it (the test suite mocks +the SDK, so `scripts/verify.sh` stayed green and never caught them): + +1. **Strict JSON schema rejected `Paper`.** `Agent(output_type=Paper)` fails + at `Runner.run` with `UserError: Strict JSON schema is enabled, but the + output type is not valid` — key-independent. Root cause: `Paper` → + `TreeKnowledge.nodes: dict[UUID, TreeNode]` maps to an open-ended + `additionalProperties` object, which strict mode forbids. + +2. **The model cannot reliably emit UUIDs.** After wrapping the output type + non-strict, the model returned semantic ids (`"paper_001"`, + `"introduction"`) for the many `UUID`-typed fields (`id`, `root_node_id`, + `node_id`, the `nodes` dict keys, and `Citation.tree_id` / `node_id`), + which fail Pydantic validation. + +3. **Even when told to emit UUIDs, the model corrupts them.** Prompt-patching + reduced but never eliminated the errors: the model hallucinated an invalid + hex character (`...-bdle-...` instead of `...-bd1e-...`) and *reused* one + UUID for two different nodes (a duplicate JSON key that silently drops a + node — Pydantic does not even flag it). + +Conclusion: asking an LLM to hand-generate dozens of unique, valid, +cross-referenced RFC-4122 UUIDs is structurally unreliable. The fix is to +stop the LLM from ever touching an identifier. + +## Approach (chosen: A — nested, id-less draft) + +Introduce an LLM-facing **draft schema** whose nodes carry only content and +nest directly (no ids, no adjacency lists). Code walks the draft, mints all +UUIDs, wires the tree, and assembles the real `Paper`. Because the LLM never +writes an identifier, id typos, id reuse, and dangling references become +impossible by construction. The draft also uses a `list` (not a `dict`), so +the LLM-facing schema has no open-ended `additionalProperties`. + +Rejected alternative B (flat list with string ids): still asks the LLM to own +id consistency; only downgrades UUID errors to string errors instead of +removing the class entirely. + +## Components + +### New module `quantmind/flows/_paper_draft.py` + +LLM-facing schema (all `model_config = ConfigDict(extra="forbid")`): + +```python +class DraftCitation(BaseModel): + quote: str | None = Field(default=None, max_length=500) + page: int | None = None + char_offset: int | None = None + +class DraftNode(BaseModel): + title: str + summary: str + content: str | None = None + citations: list[DraftCitation] = Field(default_factory=list) + children: list["DraftNode"] = Field(default_factory=list) + +class DraftPaper(BaseModel): + # DraftPaper IS the root node: it carries the root's node-like fields + # directly (content/citations/children) plus paper metadata. A live run + # showed gpt-4o-mini emits top-level sections as `children`/`citations` + # rather than nesting them under a separate `root:` field, so the schema + # matches that natural shape instead of fighting it. + title: str + summary: str + published_date: date | None = None # model reads it from the PDF; feeds as_of + arxiv_id: str | None = None + authors: list[str] = Field(default_factory=list) + asset_classes: list[str] = Field(default_factory=list) + content: str | None = None + citations: list[DraftCitation] = Field(default_factory=list) + children: list["DraftNode"] = Field(default_factory=list) +``` + +Assembly (pure, deterministic, no I/O): + +```python +def assemble_paper( + draft: DraftPaper, + *, + source: SourceRef, + source_id: str, + as_of: datetime, + extraction: ExtractionRef, + out_type: type[Paper], +) -> Paper +``` + +Behaviour: +- Generate `paper_id = uuid4()` up front (needed for `Citation.tree_id`). +- Wrap the `DraftPaper`'s own node-like fields (`title`, `summary`, + `content`, `citations`, `children`) into a root `DraftNode`, then DFS from + it; for each `DraftNode` mint `uuid4()` and build a + `TreeNode(node_id, parent_id, position, title, summary, content, + children_ids, citations)`. + - `position` = index among its siblings. + - `children_ids` = the freshly minted uuids of that node's `children`. + - each `DraftCitation` → `Citation(source_id=source_id, tree_id=paper_id, + node_id=, quote, page, char_offset)`. +- `root_node_id` = root node's uuid; `nodes` = the flat `dict[UUID, TreeNode]`. +- Construct `out_type(id=paper_id, as_of=as_of, source=source, + extraction=extraction, root_node_id=..., nodes=..., arxiv_id=draft.arxiv_id, + authors=draft.authors, asset_classes=draft.asset_classes)`. `item_type` is + left to the model's class default (`"paper"`), not passed explicitly. +- Top-level `Paper.citations` stays empty; citations live on their nodes. + +### `quantmind/flows/paper.py` changes + +- Agent output type: `AgentOutputSchema(DraftPaper, strict_json_schema=False)` + (replaces the transitional `AgentOutputSchema(Paper, ...)` wrap). +- `_DEFAULT_INSTRUCTIONS` rewritten to describe the nested draft (title / + summary / content / citations / children). Remove all UUID guidance and the + `as_of` / `source` guidance (code owns provenance now). Keep the three cfg + flags (`extract_methodology`, `extract_limitations`, `asset_class_hint`). +- After `run_with_observability` returns a `DraftPaper`, paper_flow builds + provenance and calls `assemble_paper(...)`, returning the final `Paper`. +- Provenance derivation (in paper_flow, from `source_meta`): + - `source: SourceRef` — `kind` + `uri` from the fetch metadata (`arxiv` / + `web` / `local` / `inline`; `doi` remains NotImplemented). `fetched_at` + and `content_hash` are left `None` in this MVP (the current `source_meta` + does not carry them; a follow-up can populate the sha256 dedup hash). + - `source_id: str` — a stable string (arxiv id / url / path / `"inline"`). + - `as_of` — `draft.published_date` (as UTC datetime) if present, else + `datetime.now(timezone.utc)`. + - `extraction: ExtractionRef(flow="paper_flow", model=cfg.model, + run_id=None, extracted_at=now)`. +- `output_type` override keeps meaning "final Paper type": the Agent always + emits `DraftPaper`; `out_type` is used only in `assemble_paper`. + +## Data flow + +``` +input → _fetch_and_format → (markdown, source_meta) + → Agent(output_type=DraftPaper) via run_with_observability → DraftPaper + → assemble_paper(draft, source, source_id, as_of, extraction, out_type) + → Paper +``` + +## Error handling + +- Id typos / reuse / dangling refs are eliminated by construction (LLM never + emits ids). +- `assemble_paper` is pure and unit-testable; no network, no clock beyond + `as_of` fallback. +- `DraftPaper.title` / `summary` are required, so a malformed draft still + triggers the SDK's normal validation-retry path. + +## Testing (TDD) + +New `tests/flows/test_paper_draft.py`: +- single-root draft → `Paper` with one node; `root_node_id == node_id`; + `source` / `as_of` / `extraction` set from args. +- nested draft (root + 2 children + 1 grandchild) → correct `parent_id`, + `children_ids`, `position`; uuids unique; `walk_dfs` order preserved. +- citations on a node → `Citation.node_id` is that node's uuid, + `Citation.tree_id == paper_id`, `quote` / `page` preserved, + `source_id` set. +- `out_type=MyPaper` (a `Paper` subclass) → returns a `MyPaper` instance. + +Update `tests/flows/test_paper.py`: +- stub the runner to return a `DraftPaper` (replace `_stub_paper()` with a + `_stub_draft()` helper) and assert paper_flow returns the assembled `Paper`. +- rewrite `test_output_type_override_propagated` to assert + `isinstance(result, MyPaper)` after assembly (the Agent's `output_type` is + now always `DraftPaper`). + +## Architecture / boundaries + +- `_paper_draft.py` imports only `quantmind.knowledge` (Paper, TreeNode, + Citation, SourceRef, ExtractionRef) + stdlib. It sits inside the `flows` + apex; no `import-linter` contract changes. +- Follows repo conventions: Pydantic at the LLM boundary; frozen value types + (`TreeNode`, `Citation`) constructed internally; functions over classes. + +## Out of scope (YAGNI) + +- `news_flow` / `earnings_flow` (same pattern can be applied later). +- `PaperKnowledgeCard` generation. +- DOI → OA PDF (unpaywall) resolver — tracked separately. +- Seeded/deterministic uuids — tests assert structure, not specific ids. +``` diff --git a/quantmind/flows/_paper_draft.py b/quantmind/flows/_paper_draft.py new file mode 100644 index 0000000..2311aca --- /dev/null +++ b/quantmind/flows/_paper_draft.py @@ -0,0 +1,137 @@ +"""LLM-facing draft schema for paper_flow + assembly into ``Paper``. + +The Agent produces a ``DraftPaper``: a nested tree carrying only content +(no ids). ``assemble_paper`` then mints every UUID and builds the real +``Paper``. Keeping ids out of the LLM's hands makes id typos, id reuse, and +dangling references impossible by construction. +""" + +from datetime import date, datetime +from typing import TypeVar +from uuid import UUID, uuid4 + +from pydantic import BaseModel, ConfigDict, Field + +from quantmind.knowledge import ( + Citation, + ExtractionRef, + Paper, + SourceRef, + TreeNode, +) + + +class DraftCitation(BaseModel): + """A citation the LLM attaches to a node — no id references.""" + + model_config = ConfigDict(extra="forbid") + + quote: str | None = Field(default=None, max_length=500) + page: int | None = None + char_offset: int | None = None + + +class DraftNode(BaseModel): + """A draft section node. Nests children directly; carries no ids.""" + + model_config = ConfigDict(extra="forbid") + + title: str + summary: str + content: str | None = None + citations: list[DraftCitation] = Field(default_factory=list) + children: list["DraftNode"] = Field(default_factory=list) + + +class DraftPaper(BaseModel): + """Top-level LLM output — the paper as the root of a nested section tree. + + ``DraftPaper`` *is* the root node: it carries the paper's ``title`` / + ``summary`` (plus optional ``content`` / ``citations``) and its + ``children`` are the paper's top-level sections, each of which may nest + further. This matches how the model naturally emits a document, rather + than forcing a single ``root`` field. + """ + + model_config = ConfigDict(extra="forbid") + + title: str + summary: str + published_date: date | None = None + arxiv_id: str | None = None + authors: list[str] = Field(default_factory=list) + asset_classes: list[str] = Field(default_factory=list) + content: str | None = None + citations: list[DraftCitation] = Field(default_factory=list) + children: list["DraftNode"] = Field(default_factory=list) + + +P = TypeVar("P", bound=Paper) + + +def assemble_paper( + draft: DraftPaper, + *, + source: SourceRef, + source_id: str, + as_of: datetime, + extraction: ExtractionRef, + out_type: type[P], +) -> P: + """Walk ``draft``, mint UUIDs, and build a fully-wired ``Paper``. + + Every node gets a fresh ``uuid4()``; ``parent_id`` / ``children_ids`` / + ``root_node_id`` and each citation's ``node_id`` / ``tree_id`` are set by + this function so the LLM never has to emit an id. + """ + paper_id = uuid4() + nodes: dict[UUID, TreeNode] = {} + + def _build(node: DraftNode, parent_id: UUID | None, position: int) -> UUID: + node_id = uuid4() + children_ids = [ + _build(child, node_id, i) for i, child in enumerate(node.children) + ] + citations = [ + Citation( + source_id=source_id, + page=c.page, + char_offset=c.char_offset, + quote=c.quote, + tree_id=paper_id, + node_id=node_id, + ) + for c in node.citations + ] + nodes[node_id] = TreeNode( + node_id=node_id, + parent_id=parent_id, + position=position, + title=node.title, + summary=node.summary, + content=node.content, + citations=citations, + children_ids=children_ids, + ) + return node_id + + root_draft = DraftNode( + title=draft.title, + summary=draft.summary, + content=draft.content, + citations=draft.citations, + children=draft.children, + ) + root_id = _build(root_draft, None, 0) + + return out_type( + id=paper_id, + as_of=as_of, + source=source, + extraction=extraction, + root_node_id=root_id, + nodes=nodes, + arxiv_id=draft.arxiv_id, + authors=list(draft.authors), + asset_classes=list(draft.asset_classes), + ) diff --git a/quantmind/flows/paper.py b/quantmind/flows/paper.py index 5b8373b..1560dbb 100644 --- a/quantmind/flows/paper.py +++ b/quantmind/flows/paper.py @@ -3,17 +3,19 @@ `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. +``Agent(output_type=DraftPaper)`` to produce an id-less draft tree. +``assemble_paper`` then mints ids and provenance to build the typed +``Paper`` ``TreeKnowledge`` object. Customization happens through the configured ``PaperFlowCfg`` (Layer 1) or the keyword arguments on this function (Layer 2). To swap the whole flow, fork this file (Layer 3 — design doc §9). """ +from datetime import datetime, timezone from typing import Any, TypeVar -from agents import Agent, RunHooks, Tool +from agents import Agent, AgentOutputSchema, RunHooks, Tool from quantmind.configs import PaperFlowCfg from quantmind.configs.paper import ( @@ -24,8 +26,9 @@ PaperInput, RawText, ) +from quantmind.flows._paper_draft import DraftPaper, assemble_paper from quantmind.flows._runner import run_with_observability -from quantmind.knowledge import Paper +from quantmind.knowledge import ExtractionRef, Paper, SourceRef from quantmind.preprocess.fetch import ( Fetched, fetch_arxiv, @@ -37,23 +40,29 @@ 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. +You are extracting a research paper into a nested draft tree. The top-level +DraftPaper IS the root of the tree: give it the paper's ``title`` and a short +``summary``, and put the paper's top-level sections in its ``children``. Every +node (the paper itself and each section) has a ``title`` and a short +``summary``; put the section body Markdown in ``content`` on the nodes that +carry text. Nest subsections under their parent section via that section's +``children``. + +Do NOT invent identifiers of any kind — the system assigns all ids. Attach +supporting ``citations`` (quote + page) to the node they support; never +reference a node by id. + +Set ``published_date`` to the paper's publication date when the text states +it. Fill ``arxiv_id``, ``authors``, and ``asset_classes`` from the metadata +and content when available. 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. +- 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 ``asset_classes`` if the paper does not state one explicitly. """ @@ -98,20 +107,37 @@ async def paper_flow( ), "model": cfg.model, "tools": list(extra_tools or []), - "output_type": out_type, + # The Agent always emits the id-less `DraftPaper` (never the real + # `Paper`): the LLM must not mint UUIDs, so ids are assembled in + # code via `assemble_paper` below. Non-strict schema mode is kept + # so optional fields stay simple across SDK versions. + "output_type": AgentOutputSchema(DraftPaper, strict_json_schema=False), "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( + draft: DraftPaper = await run_with_observability( agent, _format_input(raw_md, source_meta), cfg=cfg, memory=memory, extra_run_hooks=list(extra_run_hooks or []), ) + return assemble_paper( + draft, + source=_source_ref(source_meta), + source_id=_source_id(source_meta), + as_of=_as_of(draft), + extraction=ExtractionRef( + flow="paper_flow", + model=cfg.model, + run_id=None, + extracted_at=datetime.now(timezone.utc), + ), + out_type=out_type, + ) async def _fetch_and_format( @@ -201,3 +227,37 @@ def _format_input(raw_md: str, source_meta: dict[str, Any]) -> str: return ( f"--- Source metadata ---\n{header}\n\n--- Paper content ---\n{raw_md}" ) + + +def _source_ref(source_meta: dict[str, Any]) -> SourceRef: + """Build a typed ``SourceRef`` from fetch metadata.""" + src = source_meta.get("source", "inline") + if src == "arxiv": + arxiv_id = source_meta.get("arxiv_id") + uri = f"https://arxiv.org/abs/{arxiv_id}" if arxiv_id else None + return SourceRef(kind="arxiv", uri=uri) + if src == "web": + return SourceRef(kind="http", uri=source_meta.get("url")) + if src == "local": + return SourceRef(kind="local", uri=source_meta.get("path")) + return SourceRef(kind="manual", uri=None) + + +def _source_id(source_meta: dict[str, Any]) -> str: + """Return a stable citation source id from fetch metadata.""" + src = source_meta.get("source", "inline") + if src == "arxiv": + return source_meta.get("arxiv_id") or "arxiv" + if src == "web": + return source_meta.get("url") or "web" + if src == "local": + return source_meta.get("path") or "local" + return "inline" + + +def _as_of(draft: DraftPaper) -> datetime: + """Publication date from the draft if present, else now (UTC).""" + if draft.published_date is not None: + d = draft.published_date + return datetime(d.year, d.month, d.day, tzinfo=timezone.utc) + return datetime.now(timezone.utc) diff --git a/tests/flows/test_paper.py b/tests/flows/test_paper.py index 2a143c7..2a29f3a 100644 --- a/tests/flows/test_paper.py +++ b/tests/flows/test_paper.py @@ -1,13 +1,12 @@ """Tests for ``quantmind.flows.paper``.""" import unittest -from datetime import datetime, timezone +from datetime import date, datetime, timezone from pathlib import Path from typing import Any from unittest.mock import AsyncMock, MagicMock, patch -from uuid import uuid4 -from agents import RunHooks +from agents import AgentOutputSchema, RunHooks from quantmind.configs import PaperFlowCfg from quantmind.configs.paper import ( @@ -17,31 +16,24 @@ LocalFilePath, RawText, ) +from quantmind.flows._paper_draft import DraftPaper from quantmind.flows.paper import ( UnsupportedContentTypeError, + _as_of, _compose_instructions, _fetch_and_format, _format_by_content_type, _format_input, + _source_id, + _source_ref, paper_flow, ) -from quantmind.knowledge import Paper, SourceRef, TreeNode +from quantmind.knowledge import Paper from quantmind.preprocess.fetch import Fetched, RawPaper -def _stub_paper() -> Paper: - root_id = uuid4() - root = TreeNode(node_id=root_id, title="root", summary="stub") - return Paper( - as_of=datetime(2026, 5, 7, tzinfo=timezone.utc), - source=SourceRef( - kind="arxiv", - uri="arxiv:2604.12345", - fetched_at=datetime(2026, 5, 7, tzinfo=timezone.utc), - ), - root_node_id=root_id, - nodes={root_id: root}, - ) +def _stub_draft() -> DraftPaper: + return DraftPaper(title="root", summary="stub") def _patch_runner(return_value: Any) -> Any: @@ -216,6 +208,39 @@ def test_none_values_skipped(self) -> None: self.assertNotIn("b:", out) +class ProvenanceHelperTests(unittest.TestCase): + def test_source_ref_web_maps_to_http(self) -> None: + ref = _source_ref({"source": "web", "url": "https://x/y"}) + self.assertEqual(ref.kind, "http") + self.assertEqual(ref.uri, "https://x/y") + + def test_source_ref_local(self) -> None: + ref = _source_ref({"source": "local", "path": "/tmp/p.pdf"}) + self.assertEqual(ref.kind, "local") + self.assertEqual(ref.uri, "/tmp/p.pdf") + + def test_source_ref_inline_fallback(self) -> None: + ref = _source_ref({"source": "inline"}) + self.assertEqual(ref.kind, "manual") + self.assertIsNone(ref.uri) + + def test_source_id_local_and_web(self) -> None: + self.assertEqual( + _source_id({"source": "local", "path": "/tmp/p.pdf"}), "/tmp/p.pdf" + ) + self.assertEqual( + _source_id({"source": "web", "url": "https://x"}), "https://x" + ) + + def test_as_of_uses_published_date(self) -> None: + draft = DraftPaper( + title="T", summary="s", published_date=date(2023, 12, 7) + ) + self.assertEqual( + _as_of(draft), datetime(2023, 12, 7, tzinfo=timezone.utc) + ) + + class PaperFlowTests(unittest.IsolatedAsyncioTestCase): async def test_happy_path_arxiv(self) -> None: raw_paper = RawPaper( @@ -223,7 +248,7 @@ async def test_happy_path_arxiv(self) -> None: content_type="application/pdf", arxiv_id="2604.12345", ) - stub = _stub_paper() + stub = _stub_draft() with ( patch( "quantmind.flows.paper.fetch_arxiv", @@ -236,7 +261,8 @@ async def test_happy_path_arxiv(self) -> None: _patch_runner(stub) as runner, ): out = await paper_flow(ArxivIdentifier(id="2604.12345")) - self.assertIs(out, stub) + self.assertIsInstance(out, Paper) + self.assertEqual(out.root().title, "root") runner.assert_awaited_once() async def test_extra_instructions_passed_to_agent(self) -> None: @@ -246,7 +272,7 @@ def _capture_agent(*_a: Any, **kwargs: Any) -> Any: seen.update(kwargs) return MagicMock(name="agent", _name="paper_extractor") - stub = _stub_paper() + stub = _stub_draft() with ( patch("quantmind.flows.paper.Agent", side_effect=_capture_agent), _patch_runner(stub), @@ -256,24 +282,18 @@ def _capture_agent(*_a: Any, **kwargs: Any) -> Any: extra_instructions="EXTRA-USER-DIRECTIVE", ) self.assertIn("EXTRA-USER-DIRECTIVE", seen["instructions"]) - self.assertIn("structured QuantMind", seen["instructions"]) + self.assertIn("nested draft tree", 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() - with ( - patch("quantmind.flows.paper.Agent", side_effect=_capture_agent), - _patch_runner(_stub_paper()), + patch("quantmind.flows.paper.Agent", return_value=MagicMock()), + _patch_runner(_stub_draft()), ): - await paper_flow(RawText(text="x"), output_type=MyPaper) - self.assertIs(seen["output_type"], MyPaper) + out = await paper_flow(RawText(text="x"), output_type=MyPaper) + self.assertIsInstance(out, MyPaper) async def test_extra_tools_and_guardrails_forwarded(self) -> None: seen: dict[str, Any] = {} @@ -287,7 +307,7 @@ def _capture_agent(*_a: Any, **kwargs: Any) -> Any: out_g = MagicMock() with ( patch("quantmind.flows.paper.Agent", side_effect=_capture_agent), - _patch_runner(_stub_paper()), + _patch_runner(_stub_draft()), ): await paper_flow( RawText(text="x"), @@ -305,7 +325,7 @@ async def test_memory_accepted_as_no_op(self) -> None: "quantmind.flows.paper.Agent", return_value=MagicMock(), ), - _patch_runner(_stub_paper()) as runner, + _patch_runner(_stub_draft()) as runner, ): await paper_flow(RawText(text="x"), memory=object()) # The runner sees the memory placeholder forwarded. @@ -321,7 +341,7 @@ class _H(RunHooks[Any]): "quantmind.flows.paper.Agent", return_value=MagicMock(), ), - _patch_runner(_stub_paper()) as runner, + _patch_runner(_stub_draft()) as runner, ): await paper_flow(RawText(text="x"), extra_run_hooks=[hook]) self.assertEqual(runner.await_args.kwargs["extra_run_hooks"], [hook]) @@ -339,7 +359,7 @@ def _capture_agent(*_a: Any, **kwargs: Any) -> Any: cfg = PaperFlowCfg(model_settings=ms) with ( patch("quantmind.flows.paper.Agent", side_effect=_capture_agent), - _patch_runner(_stub_paper()), + _patch_runner(_stub_draft()), ): await paper_flow(RawText(text="x"), cfg=cfg) self.assertIs(seen["model_settings"], ms) @@ -353,7 +373,24 @@ def _capture_agent(*_a: Any, **kwargs: Any) -> Any: with ( patch("quantmind.flows.paper.Agent", side_effect=_capture_agent), - _patch_runner(_stub_paper()), + _patch_runner(_stub_draft()), ): await paper_flow(RawText(text="x")) self.assertNotIn("model_settings", seen) + + async def test_agent_output_type_is_nonstrict_draftpaper(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_draft()), + ): + await paper_flow(RawText(text="x")) + schema = seen["output_type"] + self.assertIsInstance(schema, AgentOutputSchema) + self.assertIs(schema.output_type, DraftPaper) + self.assertFalse(schema.is_strict_json_schema()) diff --git a/tests/flows/test_paper_draft.py b/tests/flows/test_paper_draft.py new file mode 100644 index 0000000..f47cacf --- /dev/null +++ b/tests/flows/test_paper_draft.py @@ -0,0 +1,147 @@ +"""Tests for ``quantmind.flows._paper_draft``.""" + +import unittest +from datetime import date, datetime, timezone +from uuid import UUID + +from pydantic import ValidationError + +from quantmind.flows._paper_draft import ( + DraftCitation, + DraftNode, + DraftPaper, + assemble_paper, +) +from quantmind.knowledge import ExtractionRef, Paper, SourceRef + + +class DraftSchemaTests(unittest.TestCase): + def test_parses_nested_tree(self) -> None: + draft = DraftPaper.model_validate( + { + "title": "Momentum", + "summary": "top", + "published_date": "2023-12-07", + "authors": ["Alice"], + "children": [ + {"title": "Intro", "summary": "s", "content": "body"} + ], + } + ) + self.assertEqual(draft.published_date, date(2023, 12, 7)) + self.assertEqual(draft.children[0].title, "Intro") + self.assertEqual(draft.children[0].content, "body") + + def test_extra_fields_forbidden(self) -> None: + with self.assertRaises(ValidationError): + DraftNode.model_validate( + {"title": "x", "summary": "y", "node_id": "abc"} + ) + + def test_citation_has_no_id_fields(self) -> None: + cite = DraftCitation(quote="q", page=3) + self.assertEqual(cite.page, 3) + self.assertNotIn("node_id", DraftCitation.model_fields) + self.assertNotIn("tree_id", DraftCitation.model_fields) + + +def _args() -> dict: + return { + "source": SourceRef(kind="local", uri="/tmp/p.pdf"), + "source_id": "p.pdf", + "as_of": datetime(2023, 12, 7, tzinfo=timezone.utc), + "extraction": ExtractionRef( + flow="paper_flow", + model="gpt-4o-mini", + extracted_at=datetime(2023, 12, 7, tzinfo=timezone.utc), + ), + "out_type": Paper, + } + + +class AssemblePaperTests(unittest.TestCase): + def test_single_root(self) -> None: + draft = DraftPaper(title="Root", summary="rs", content="body") + paper = assemble_paper(draft, **_args()) + self.assertIsInstance(paper, Paper) + self.assertEqual(len(paper.nodes), 1) + self.assertEqual(paper.root_node_id, paper.root().node_id) + self.assertEqual(paper.root().title, "Root") + self.assertEqual(paper.root().content, "body") + self.assertEqual(paper.source.kind, "local") + self.assertEqual( + paper.as_of, datetime(2023, 12, 7, tzinfo=timezone.utc) + ) + self.assertIsInstance(paper.root().node_id, UUID) + + def test_nested_wiring_and_positions(self) -> None: + draft = DraftPaper( + title="Root", + summary="rs", + children=[ + DraftNode( + title="A", + summary="a", + children=[DraftNode(title="A1", summary="a1")], + ), + DraftNode(title="B", summary="b"), + ], + ) + paper = assemble_paper(draft, **_args()) + self.assertEqual(len(paper.nodes), 4) + titles = [n.title for n in paper.walk_dfs()] + self.assertEqual(titles, ["Root", "A", "A1", "B"]) + root = paper.root() + self.assertEqual(len(root.children_ids), 2) + a = paper.children_of(root.node_id)[0] + self.assertEqual(a.parent_id, root.node_id) + self.assertEqual(a.position, 0) + b = paper.children_of(root.node_id)[1] + self.assertEqual(b.position, 1) + # uuids unique across the tree + ids = [n.node_id for n in paper.walk_dfs()] + self.assertEqual(len(ids), len(set(ids))) + + def test_grandchild_parent_wiring(self) -> None: + draft = DraftPaper( + title="Root", + summary="rs", + children=[ + DraftNode( + title="A", + summary="a", + children=[DraftNode(title="A1", summary="a1")], + ), + ], + ) + paper = assemble_paper(draft, **_args()) + root = paper.root() + a = paper.children_of(root.node_id)[0] + a1 = paper.children_of(a.node_id)[0] + self.assertEqual(a1.title, "A1") + self.assertEqual(a1.parent_id, a.node_id) + self.assertEqual(a1.children_ids, []) + + def test_citations_wired_to_node(self) -> None: + draft = DraftPaper( + title="Root", + summary="rs", + citations=[DraftCitation(quote="q", page=5)], + ) + paper = assemble_paper(draft, **_args()) + cite = paper.root().citations[0] + self.assertEqual(cite.quote, "q") + self.assertEqual(cite.page, 5) + self.assertEqual(cite.source_id, "p.pdf") + self.assertEqual(cite.node_id, paper.root().node_id) + self.assertEqual(cite.tree_id, paper.id) + + def test_out_type_subclass(self) -> None: + class MyPaper(Paper): + pass + + draft = DraftPaper(title="Root", summary="rs") + args = _args() + args["out_type"] = MyPaper + paper = assemble_paper(draft, **args) + self.assertIsInstance(paper, MyPaper)