|
@@ -0,0 +1,215 @@
|
|
|
|
|
+"""
|
|
|
|
|
+lambdagent.ingest — citation-aware document ingestion (M2 Phase E).
|
|
|
|
|
+
|
|
|
|
|
+CR-20260607-001 rev.2 §6.2 (FR-004/FR-005) + Q7 (Obsidian import).
|
|
|
|
|
+
|
|
|
|
|
+The research agent packs (Phase F) promise "回答和报告中尽量保留来源文件、
|
|
|
|
|
+页码或段落" (FR-005). To deliver that, ingestion must preserve provenance:
|
|
|
|
|
+every chunk carries which file + which page it came from, so retrieval can
|
|
|
|
|
+render `[Source N | paper.pdf p.12]`.
|
|
|
|
|
+
|
|
|
|
|
+Three building blocks, all pure + dependency-light:
|
|
|
|
|
+ - extract_pdf_pages(path) → [(page_number, text), ...] (1-indexed)
|
|
|
|
|
+ - chunk_pages(pages, source) → [CitedChunk, ...] (carries page)
|
|
|
|
|
+ - import_obsidian_vault(dir) → [CitedChunk, ...] (markdown vault)
|
|
|
|
|
+
|
|
|
|
|
+PDF extraction prefers pdfplumber (better layout + reliable page split),
|
|
|
|
|
+falls back to PyPDF2, and raises a clear error if neither is installed.
|
|
|
|
|
+Obsidian/markdown needs no dependency.
|
|
|
|
|
+"""
|
|
|
|
|
+from __future__ import annotations
|
|
|
|
|
+
|
|
|
|
|
+import os
|
|
|
|
|
+from dataclasses import dataclass, field
|
|
|
|
|
+from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@dataclass
|
|
|
|
|
+class CitedChunk:
|
|
|
|
|
+ """A retrievable chunk that remembers where it came from."""
|
|
|
|
|
+ text: str
|
|
|
|
|
+ source: str # filename or relative path
|
|
|
|
|
+ page: Optional[int] = None # 1-indexed PDF page, or None for non-paged
|
|
|
|
|
+ chunk_index: int = 0 # position within the source
|
|
|
|
|
+ extra: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
+
|
|
|
|
|
+ def to_metadata(self) -> Dict[str, Any]:
|
|
|
|
|
+ """Metadata dict suitable for rag.SimpleVectorStore.add(metadata=...).
|
|
|
|
|
+ The keys `source` + `page` are what rag's citation formatter looks
|
|
|
|
|
+ for to render `[Source N | source p.page]`."""
|
|
|
|
|
+ md: Dict[str, Any] = {"source": self.source, "chunk_index": self.chunk_index}
|
|
|
|
|
+ if self.page is not None:
|
|
|
|
|
+ md["page"] = self.page
|
|
|
|
|
+ md.update(self.extra)
|
|
|
|
|
+ return md
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+# ── PDF page extraction ───────────────────────────────────────────────────
|
|
|
|
|
+
|
|
|
|
|
+def extract_pdf_pages(path: str) -> List[Tuple[int, str]]:
|
|
|
|
|
+ """Extract text per page. Returns [(page_number_1indexed, text), ...].
|
|
|
|
|
+
|
|
|
|
|
+ Prefers pdfplumber (cleaner layout-aware extraction), falls back to
|
|
|
|
|
+ PyPDF2. Pages with no extractable text are still returned (empty string)
|
|
|
|
|
+ so page numbering stays aligned with the physical document.
|
|
|
|
|
+ """
|
|
|
|
|
+ if not os.path.isfile(path):
|
|
|
|
|
+ raise FileNotFoundError(f"PDF not found: {path}")
|
|
|
|
|
+
|
|
|
|
|
+ # pdfplumber first
|
|
|
|
|
+ try:
|
|
|
|
|
+ import pdfplumber # type: ignore
|
|
|
|
|
+ out: List[Tuple[int, str]] = []
|
|
|
|
|
+ with pdfplumber.open(path) as pdf:
|
|
|
|
|
+ for i, page in enumerate(pdf.pages):
|
|
|
|
|
+ out.append((i + 1, page.extract_text() or ""))
|
|
|
|
|
+ return out
|
|
|
|
|
+ except ImportError:
|
|
|
|
|
+ pass
|
|
|
|
|
+
|
|
|
|
|
+ # PyPDF2 fallback
|
|
|
|
|
+ try:
|
|
|
|
|
+ import PyPDF2 # type: ignore
|
|
|
|
|
+ out = []
|
|
|
|
|
+ with open(path, "rb") as f:
|
|
|
|
|
+ reader = PyPDF2.PdfReader(f)
|
|
|
|
|
+ for i, page in enumerate(reader.pages):
|
|
|
|
|
+ out.append((i + 1, page.extract_text() or ""))
|
|
|
|
|
+ return out
|
|
|
|
|
+ except ImportError:
|
|
|
|
|
+ raise RuntimeError(
|
|
|
|
|
+ "PDF extraction needs pdfplumber or PyPDF2. "
|
|
|
|
|
+ "Install with: pip install pdfplumber"
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+# ── Page-preserving chunker ────────────────────────────────────────────────
|
|
|
|
|
+
|
|
|
|
|
+def chunk_text(text: str, chunk_size: int = 800, overlap: int = 100) -> List[str]:
|
|
|
|
|
+ """Split text into overlapping windows on paragraph/whitespace
|
|
|
|
|
+ boundaries. Pure helper used by chunk_pages."""
|
|
|
|
|
+ text = text.strip()
|
|
|
|
|
+ if not text:
|
|
|
|
|
+ return []
|
|
|
|
|
+ if len(text) <= chunk_size:
|
|
|
|
|
+ return [text]
|
|
|
|
|
+
|
|
|
|
|
+ chunks: List[str] = []
|
|
|
|
|
+ start = 0
|
|
|
|
|
+ n = len(text)
|
|
|
|
|
+ while start < n:
|
|
|
|
|
+ end = min(start + chunk_size, n)
|
|
|
|
|
+ # Try to break on a paragraph or sentence boundary near `end`.
|
|
|
|
|
+ if end < n:
|
|
|
|
|
+ window = text[start:end]
|
|
|
|
|
+ for sep in ("\n\n", "\n", ". ", "。", " "):
|
|
|
|
|
+ idx = window.rfind(sep)
|
|
|
|
|
+ if idx > chunk_size * 0.5: # don't break too early
|
|
|
|
|
+ end = start + idx + len(sep)
|
|
|
|
|
+ break
|
|
|
|
|
+ chunks.append(text[start:end].strip())
|
|
|
|
|
+ if end >= n:
|
|
|
|
|
+ break
|
|
|
|
|
+ start = max(end - overlap, start + 1)
|
|
|
|
|
+ return [c for c in chunks if c]
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def chunk_pages(
|
|
|
|
|
+ pages: List[Tuple[int, str]],
|
|
|
|
|
+ source: str,
|
|
|
|
|
+ chunk_size: int = 800,
|
|
|
|
|
+ overlap: int = 100,
|
|
|
|
|
+) -> List[CitedChunk]:
|
|
|
|
|
+ """Chunk per-page text into CitedChunks that remember their page.
|
|
|
|
|
+
|
|
|
|
|
+ A chunk never spans two pages (so the page number is always exact).
|
|
|
|
|
+ `chunk_index` is global across the document.
|
|
|
|
|
+ """
|
|
|
|
|
+ out: List[CitedChunk] = []
|
|
|
|
|
+ idx = 0
|
|
|
|
|
+ for page_no, page_text in pages:
|
|
|
|
|
+ for piece in chunk_text(page_text, chunk_size, overlap):
|
|
|
|
|
+ out.append(CitedChunk(
|
|
|
|
|
+ text=piece, source=source, page=page_no, chunk_index=idx,
|
|
|
|
|
+ ))
|
|
|
|
|
+ idx += 1
|
|
|
|
|
+ return out
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def ingest_pdf(path: str, chunk_size: int = 800, overlap: int = 100) -> List[CitedChunk]:
|
|
|
|
|
+ """Convenience: extract + chunk a PDF in one call. Source = basename."""
|
|
|
|
|
+ pages = extract_pdf_pages(path)
|
|
|
|
|
+ return chunk_pages(pages, os.path.basename(path), chunk_size, overlap)
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+# ── Obsidian / markdown vault import (Q7) ──────────────────────────────────
|
|
|
|
|
+
|
|
|
|
|
+import re as _re
|
|
|
|
|
+
|
|
|
|
|
+_WIKILINK_RE = _re.compile(r"\[\[([^\]|]+)(?:\|([^\]]+))?\]\]")
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def _strip_wikilinks(text: str) -> Tuple[str, List[str]]:
|
|
|
|
|
+ """Replace Obsidian `[[Note]]` / `[[Note|alias]]` with readable text and
|
|
|
|
|
+ return the referenced note targets (for building a link graph later).
|
|
|
|
|
+
|
|
|
|
|
+ The body keeps the human-facing display text — the alias when present,
|
|
|
|
|
+ else the target — so the chunk reads naturally. The link graph keeps the
|
|
|
|
|
+ real target (`extra['links']`)."""
|
|
|
|
|
+ links: List[str] = []
|
|
|
|
|
+
|
|
|
|
|
+ def _repl(m):
|
|
|
|
|
+ target = m.group(1).strip()
|
|
|
|
|
+ alias = (m.group(2) or "").strip()
|
|
|
|
|
+ links.append(target)
|
|
|
|
|
+ return alias or target
|
|
|
|
|
+
|
|
|
|
|
+ return _WIKILINK_RE.sub(_repl, text), links
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def import_obsidian_vault(
|
|
|
|
|
+ vault_dir: str,
|
|
|
|
|
+ chunk_size: int = 800,
|
|
|
|
|
+ overlap: int = 100,
|
|
|
|
|
+ include_subdirs: bool = True,
|
|
|
|
|
+) -> List[CitedChunk]:
|
|
|
|
|
+ """Walk an Obsidian vault (a folder of .md files) and return CitedChunks.
|
|
|
|
|
+
|
|
|
|
|
+ Q7 decision: Obsidian is a first-class import because a vault is just a
|
|
|
|
|
+ markdown folder — zero integration cost. `[[wikilinks]]` are flattened to
|
|
|
|
|
+ their display text in the chunk body, and the referenced note names are
|
|
|
|
|
+ kept in `extra['links']` so a future pass can build the link graph.
|
|
|
|
|
+
|
|
|
|
|
+ `source` is the path relative to the vault root (e.g. "notes/qgt.md").
|
|
|
|
|
+ Markdown has no pages, so `page` is None.
|
|
|
|
|
+ """
|
|
|
|
|
+ if not os.path.isdir(vault_dir):
|
|
|
|
|
+ raise NotADirectoryError(f"Obsidian vault not found: {vault_dir}")
|
|
|
|
|
+
|
|
|
|
|
+ vault_abs = os.path.abspath(vault_dir)
|
|
|
|
|
+ out: List[CitedChunk] = []
|
|
|
|
|
+
|
|
|
|
|
+ for root, dirs, files in os.walk(vault_abs):
|
|
|
|
|
+ # Skip Obsidian's own config + hidden dirs.
|
|
|
|
|
+ dirs[:] = [d for d in dirs if not d.startswith(".")]
|
|
|
|
|
+ if not include_subdirs and root != vault_abs:
|
|
|
|
|
+ continue
|
|
|
|
|
+ for fname in sorted(files):
|
|
|
|
|
+ if not fname.lower().endswith((".md", ".markdown")):
|
|
|
|
|
+ continue
|
|
|
|
|
+ fpath = os.path.join(root, fname)
|
|
|
|
|
+ rel = os.path.relpath(fpath, vault_abs)
|
|
|
|
|
+ try:
|
|
|
|
|
+ with open(fpath, encoding="utf-8") as f:
|
|
|
|
|
+ raw = f.read()
|
|
|
|
|
+ except (OSError, UnicodeDecodeError):
|
|
|
|
|
+ continue
|
|
|
|
|
+ clean, links = _strip_wikilinks(raw)
|
|
|
|
|
+ idx = 0
|
|
|
|
|
+ for piece in chunk_text(clean, chunk_size, overlap):
|
|
|
|
|
+ out.append(CitedChunk(
|
|
|
|
|
+ text=piece, source=rel, page=None, chunk_index=idx,
|
|
|
|
|
+ extra={"links": links} if links else {},
|
|
|
|
|
+ ))
|
|
|
|
|
+ idx += 1
|
|
|
|
|
+ return out
|