mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
feat: add structure tree builder for paper navigation
This commit is contained in:
parent
b02e762ce4
commit
12b930c31d
8 changed files with 299 additions and 0 deletions
|
|
@ -17,6 +17,7 @@ _COMMAND_REGISTRY: dict[str, str] = {
|
|||
"paper-status": "paperforge.commands.paper_status",
|
||||
"agent-context": "paperforge.commands.agent_context",
|
||||
"prune": "paperforge.commands.prune",
|
||||
"paper-navigation": "paperforge.commands.paper_navigation",
|
||||
"reading-log": "paperforge.commands.reading_log",
|
||||
}
|
||||
|
||||
|
|
|
|||
60
paperforge/commands/paper_navigation.py
Normal file
60
paperforge/commands/paper_navigation.py
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
|
||||
from paperforge import __version__ as PF_VERSION
|
||||
from paperforge.core.result import PFResult
|
||||
from paperforge.core.io import read_json
|
||||
from paperforge.retrieval.structure_tree import summarize_role_index
|
||||
|
||||
|
||||
def _resolve_paper_root(vault_path, query: str):
|
||||
"""Resolve a paper query to its OCR root directory."""
|
||||
from paperforge.memory.db import get_connection, get_memory_db_path
|
||||
from paperforge.memory.query import lookup_paper
|
||||
from paperforge.worker._utils import pipeline_paths
|
||||
|
||||
db_path = get_memory_db_path(vault_path)
|
||||
if not db_path or not db_path.exists():
|
||||
return None
|
||||
conn = get_connection(db_path, read_only=True)
|
||||
try:
|
||||
entries = lookup_paper(conn, query)
|
||||
if entries:
|
||||
ocr_root = pipeline_paths(vault_path)["ocr"]
|
||||
return ocr_root / entries[0]["zotero_key"]
|
||||
finally:
|
||||
conn.close()
|
||||
return None
|
||||
|
||||
|
||||
def run(args: argparse.Namespace) -> int:
|
||||
paper_root = _resolve_paper_root(args.vault_path, args.query)
|
||||
if paper_root is None:
|
||||
data = {"mode": "not_found", "paper_id": args.query, "error": "Paper not found in vault"}
|
||||
if args.json:
|
||||
print(PFResult(ok=False, command="paper-navigation", version=PF_VERSION, data=data).to_json())
|
||||
else:
|
||||
print(json.dumps(data, ensure_ascii=False, indent=2))
|
||||
return 1
|
||||
|
||||
tree_path = paper_root / "index" / "structure-tree.json"
|
||||
if tree_path.exists():
|
||||
tree = read_json(tree_path)
|
||||
payload = {"mode": "structure_tree", "paper_id": tree.get("paper_id", ""), "nodes": tree.get("nodes", [])}
|
||||
else:
|
||||
role_index_path = paper_root / "index" / "role-index.json"
|
||||
if role_index_path.exists():
|
||||
role_index = read_json(role_index_path)
|
||||
summary = summarize_role_index(role_index)
|
||||
payload = {"mode": "role_index_summary", "paper_id": args.query, "summary": summary}
|
||||
else:
|
||||
payload = {"mode": "no_index", "paper_id": args.query}
|
||||
|
||||
result = PFResult(ok=True, command="paper-navigation", version=PF_VERSION, data=payload)
|
||||
if args.json:
|
||||
print(result.to_json())
|
||||
else:
|
||||
print(json.dumps(payload, ensure_ascii=False, indent=2))
|
||||
return 0
|
||||
3
paperforge/retrieval/__init__.py
Normal file
3
paperforge/retrieval/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
"""paperforge.retrieval — retrieval substrate for paper navigation and search."""
|
||||
|
||||
from __future__ import annotations
|
||||
67
paperforge/retrieval/gateway.py
Normal file
67
paperforge/retrieval/gateway.py
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
"""paperforge.retrieval.gateway — Layer 4 gateway command routing.
|
||||
|
||||
Provides the core `route_gateway()` function that all gateway commands
|
||||
route through. Currently delegates to existing capabilities; later tasks
|
||||
will upgrade routing to use real Layer 4 artifacts.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from paperforge import __version__ as PF_VERSION
|
||||
from paperforge.core.result import PFResult
|
||||
from paperforge.query_planning import build_query_plan, enrich_query_plan_with_runtime
|
||||
|
||||
# Map gateway intent names to build_query_plan intent strings
|
||||
INTENTS: dict[str, str] = {
|
||||
"paper-lookup": "known-paper",
|
||||
"content-discovery": "content",
|
||||
"paper-navigation": "known-paper",
|
||||
"scoped-fetch": "known-paper",
|
||||
}
|
||||
|
||||
|
||||
def route_gateway(
|
||||
vault: Path,
|
||||
intent: str,
|
||||
query: str,
|
||||
*,
|
||||
json_mode: bool,
|
||||
limit: int = 5,
|
||||
) -> PFResult:
|
||||
"""Route a gateway command through the query planning pipeline.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
vault : Path
|
||||
PaperForge vault root.
|
||||
intent : str
|
||||
Gateway intent name (e.g. ``"paper-lookup"``).
|
||||
query : str
|
||||
Free-text or structured query.
|
||||
json_mode : bool
|
||||
Whether the caller expects JSON output.
|
||||
limit : int, optional
|
||||
Maximum result count (default 5).
|
||||
|
||||
Returns
|
||||
-------
|
||||
PFResult
|
||||
Result packet with the route plan and intent metadata.
|
||||
"""
|
||||
plan = enrich_query_plan_with_runtime(
|
||||
build_query_plan(query, INTENTS[intent]),
|
||||
vault,
|
||||
)
|
||||
return PFResult(
|
||||
ok=True,
|
||||
command=intent,
|
||||
version=PF_VERSION,
|
||||
data={
|
||||
"intent": intent,
|
||||
"query": query,
|
||||
"route_plan": plan,
|
||||
"limit": limit,
|
||||
},
|
||||
)
|
||||
44
paperforge/retrieval/structure_tree.py
Normal file
44
paperforge/retrieval/structure_tree.py
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from paperforge.core.io import write_json
|
||||
|
||||
|
||||
def build_structure_tree(structured_blocks: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
nodes: list[dict[str, Any]] = []
|
||||
current_section: dict[str, Any] | None = None
|
||||
for block in structured_blocks:
|
||||
role = block.get("role")
|
||||
text = str(block.get("text", "")).strip()
|
||||
if role in {"section_heading", "subsection_heading", "introduction_heading", "abstract_heading"} and text:
|
||||
current_section = {
|
||||
"node_id": f"sec:{block.get('block_id')}",
|
||||
"kind": "section",
|
||||
"title": text,
|
||||
"level": 1 if role != "subsection_heading" else 2,
|
||||
"section_path": [text] if role != "subsection_heading" else [nodes[-1]["title"], text] if nodes else [text],
|
||||
"page_span": [block.get("page", 0), block.get("page", 0)],
|
||||
"block_span": [[block.get("page", 0), block.get("block_id", "")]],
|
||||
"children": [],
|
||||
"objects": [],
|
||||
}
|
||||
nodes.append(current_section)
|
||||
elif current_section is not None:
|
||||
current_section["page_span"][1] = block.get("page", current_section["page_span"][1])
|
||||
current_section["block_span"].append([block.get("page", 0), block.get("block_id", "")])
|
||||
return {"paper_id": structured_blocks[0].get("paper_id", "") if structured_blocks else "", "nodes": nodes}
|
||||
|
||||
|
||||
def write_structure_tree(index_root: Path, tree: dict[str, Any]) -> None:
|
||||
index_root.mkdir(parents=True, exist_ok=True)
|
||||
write_json(index_root / "structure-tree.json", tree)
|
||||
|
||||
|
||||
def summarize_role_index(role_index: dict[str, Any]) -> dict[str, Any]:
|
||||
role_counts: dict[str, int] = {}
|
||||
for key, entries in role_index.items():
|
||||
if isinstance(entries, list):
|
||||
role_counts[key] = len(entries)
|
||||
return {"role_counts": role_counts}
|
||||
|
|
@ -2029,6 +2029,11 @@ def postprocess_ocr_result(vault: Path, key: str, all_results: list[dict]) -> tu
|
|||
resolved_metadata=resolved,
|
||||
)
|
||||
write_role_index(ocr_root / "index", role_indexes)
|
||||
# --- Phase 6: structure tree ---
|
||||
from paperforge.retrieval.structure_tree import build_structure_tree, write_structure_tree
|
||||
|
||||
structure_tree = build_structure_tree(structured)
|
||||
write_structure_tree(ocr_root / "index", structure_tree)
|
||||
|
||||
# Update meta.json with version payloads
|
||||
ocr_model = meta.get("ocr_model", meta.get("ocr_provider", "PaddleOCR"))
|
||||
|
|
|
|||
|
|
@ -398,6 +398,11 @@ def run_derived_rebuild_for_keys(vault: Path, keys: list[str], progress_bar=None
|
|||
resolved_metadata=resolved,
|
||||
)
|
||||
write_role_index(paper_root / "index", role_indexes)
|
||||
# Rebuild structure tree
|
||||
from paperforge.retrieval.structure_tree import build_structure_tree, write_structure_tree
|
||||
|
||||
structure_tree = build_structure_tree(structured)
|
||||
write_structure_tree(paper_root / "index", structure_tree)
|
||||
|
||||
# Update version state in meta.json
|
||||
meta = read_json(artifacts.meta_json) if artifacts.meta_json.exists() else {}
|
||||
|
|
|
|||
114
tests/test_layer4_structure_tree.py
Normal file
114
tests/test_layer4_structure_tree.py
Normal file
|
|
@ -0,0 +1,114 @@
|
|||
"""Tests for Layer 4 structure tree builder and paper navigation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from paperforge.retrieval.structure_tree import (
|
||||
build_structure_tree,
|
||||
summarize_role_index,
|
||||
write_structure_tree,
|
||||
)
|
||||
|
||||
|
||||
def test_build_structure_tree_creates_section_nodes_from_headings():
|
||||
structured_blocks = [
|
||||
{"paper_id": "ABCD1234", "page": 1, "block_id": "b1", "role": "section_heading", "text": "Methods"},
|
||||
{"paper_id": "ABCD1234", "page": 1, "block_id": "b2", "role": "body_paragraph", "text": "We recruited 30 patients."},
|
||||
]
|
||||
tree = build_structure_tree(structured_blocks)
|
||||
assert tree["paper_id"] == "ABCD1234"
|
||||
assert len(tree["nodes"]) == 1
|
||||
assert tree["nodes"][0]["title"] == "Methods"
|
||||
assert tree["nodes"][0]["section_path"] == ["Methods"]
|
||||
assert tree["nodes"][0]["level"] == 1
|
||||
assert tree["nodes"][0]["page_span"] == [1, 1]
|
||||
|
||||
|
||||
def test_build_structure_tree_handles_subsection():
|
||||
structured_blocks = [
|
||||
{"paper_id": "ABCD1234", "page": 1, "block_id": "b1", "role": "section_heading", "text": "Methods"},
|
||||
{"paper_id": "ABCD1234", "page": 1, "block_id": "b2", "role": "body_paragraph", "text": "We recruited 30 patients."},
|
||||
{"paper_id": "ABCD1234", "page": 2, "block_id": "b3", "role": "subsection_heading", "text": "Statistical Analysis"},
|
||||
{"paper_id": "ABCD1234", "page": 2, "block_id": "b4", "role": "body_paragraph", "text": "We used t-tests."},
|
||||
]
|
||||
tree = build_structure_tree(structured_blocks)
|
||||
assert tree["paper_id"] == "ABCD1234"
|
||||
assert len(tree["nodes"]) == 2
|
||||
assert tree["nodes"][1]["title"] == "Statistical Analysis"
|
||||
assert tree["nodes"][1]["level"] == 2
|
||||
assert tree["nodes"][1]["section_path"] == ["Methods", "Statistical Analysis"]
|
||||
|
||||
|
||||
def test_build_structure_tree_extends_page_span():
|
||||
structured_blocks = [
|
||||
{"paper_id": "X", "page": 1, "block_id": "s1", "role": "section_heading", "text": "Introduction"},
|
||||
{"paper_id": "X", "page": 1, "block_id": "p1", "role": "body_paragraph", "text": "Para 1"},
|
||||
{"paper_id": "X", "page": 2, "block_id": "p2", "role": "body_paragraph", "text": "Para 2"},
|
||||
{"paper_id": "X", "page": 3, "block_id": "p3", "role": "body_paragraph", "text": "Para 3"},
|
||||
]
|
||||
tree = build_structure_tree(structured_blocks)
|
||||
assert tree["nodes"][0]["page_span"] == [1, 3]
|
||||
|
||||
|
||||
def test_build_structure_tree_handles_introduction_and_abstract_headings():
|
||||
blocks = [
|
||||
{"paper_id": "Y", "page": 1, "block_id": "a1", "role": "abstract_heading", "text": "Abstract"},
|
||||
{"paper_id": "Y", "page": 1, "block_id": "b1", "role": "abstract_body", "text": "Summary here."},
|
||||
{"paper_id": "Y", "page": 1, "block_id": "i1", "role": "introduction_heading", "text": "Introduction"},
|
||||
{"paper_id": "Y", "page": 2, "block_id": "b2", "role": "body_paragraph", "text": "Background."},
|
||||
]
|
||||
tree = build_structure_tree(blocks)
|
||||
assert len(tree["nodes"]) == 2
|
||||
assert tree["nodes"][0]["title"] == "Abstract"
|
||||
assert tree["nodes"][1]["title"] == "Introduction"
|
||||
assert tree["nodes"][0]["kind"] == "section"
|
||||
|
||||
|
||||
def test_build_structure_tree_ignores_heading_roles_with_empty_text():
|
||||
blocks = [
|
||||
{"paper_id": "Z", "page": 1, "block_id": "h1", "role": "section_heading", "text": ""},
|
||||
{"paper_id": "Z", "page": 1, "block_id": "b1", "role": "body_paragraph", "text": "Some text."},
|
||||
]
|
||||
tree = build_structure_tree(blocks)
|
||||
assert len(tree["nodes"]) == 0
|
||||
|
||||
|
||||
def test_build_structure_tree_empty_input():
|
||||
tree = build_structure_tree([])
|
||||
assert tree["paper_id"] == ""
|
||||
assert tree["nodes"] == []
|
||||
|
||||
|
||||
def test_write_structure_tree_creates_file(tmp_path):
|
||||
tree = {"paper_id": "X", "nodes": []}
|
||||
write_structure_tree(tmp_path, tree)
|
||||
target = tmp_path / "structure-tree.json"
|
||||
assert target.exists()
|
||||
import json
|
||||
assert json.loads(target.read_text(encoding="utf-8")) == tree
|
||||
|
||||
|
||||
def test_summarize_role_index_counts_roles():
|
||||
role_index = {
|
||||
"headings": [{"paper_id": "X"}, {"paper_id": "X"}],
|
||||
"body_paragraphs": [{"paper_id": "X"} for _ in range(5)],
|
||||
"figure_captions": [{"paper_id": "X"}],
|
||||
}
|
||||
summary = summarize_role_index(role_index)
|
||||
assert summary["role_counts"]["headings"] == 2
|
||||
assert summary["role_counts"]["body_paragraphs"] == 5
|
||||
assert summary["role_counts"]["figure_captions"] == 1
|
||||
|
||||
|
||||
def test_build_structure_tree_block_span_accumulates():
|
||||
blocks = [
|
||||
{"paper_id": "P", "page": 1, "block_id": "sec1", "role": "section_heading", "text": "Results"},
|
||||
{"paper_id": "P", "page": 1, "block_id": "r1", "role": "body_paragraph", "text": "Result A"},
|
||||
{"paper_id": "P", "page": 2, "block_id": "r2", "role": "body_paragraph", "text": "Result B"},
|
||||
]
|
||||
tree = build_structure_tree(blocks)
|
||||
assert len(tree["nodes"][0]["block_span"]) == 3
|
||||
assert [1, "sec1"] in tree["nodes"][0]["block_span"]
|
||||
assert [2, "r2"] in tree["nodes"][0]["block_span"]
|
||||
Loading…
Reference in a new issue