scripts/pdf_ingest.py
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env python3
"""
PDF Ingest — XKB Active Recall Layer
批次將 PDF 論文資料夾轉為知識卡片並建立向量索引。
支援多欄位學術論文、自動抽取標題/作者、清除頁首頁尾雜訊。
Usage:
python3 pdf_ingest.py /path/to/papers/ # 整個資料夾
python3 pdf_ingest.py paper.pdf # 單一檔案
python3 pdf_ingest.py /path/to/papers/ --dry-run # 預覽不寫入
python3 pdf_ingest.py /path/to/papers/ --category research --tag ai
python3 pdf_ingest.py /path/to/papers/ --rebuild-index # 重建向量索引
"""
from __future__ import annotations
import argparse
import json
import os
import re
import sys
import hashlib
import urllib.request
import urllib.parse
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional
try:
import fitz # pymupdf
except ImportError:
print("Error: pymupdf not installed. Run: pip install pymupdf --break-system-packages")
sys.exit(1)
sys.path.insert(0, str(Path(__file__).resolve().parent))
import xkb_paths
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "tools"))
from runtime_config import runtime_env
WORKSPACE = xkb_paths.WORKSPACE
SCRIPTS_DIR = xkb_paths.SCRIPTS_DIR
def _get_gemini_key() -> str:
return runtime_env().get("GEMINI_API_KEY", "")
def _llm_bilingual_summary(title: str, abstract: str) -> tuple[str, list[str]]:
"""Call Gemini to generate bilingual summary + Chinese domain tags.
Returns (summary_str, tags_list). summary_str = 'zh | en'."""
key = _get_gemini_key()
if not key:
return "", []
prompt = (
f"Paper title: {title}\n\nAbstract (first 800 chars):\n{abstract[:800]}\n\n"
"Output exactly three lines:\n"
"ZH: <20-40字繁體中文摘要，說明這篇論文的核心發現>\n"
"EN: <15-30 word English summary of the core finding>\n"
"TAGS: <5-8個繁體中文醫療/技術關鍵詞，逗號分隔，例如：牙科,影像診斷,深度學習,心臟病>"
)
try:
url = f"https://generativelanguage.googleapis.com/v1/models/gemini-2.5-flash:generateContent?key={key}"
payload = json.dumps({"contents": [{"parts": [{"text": prompt}]}]}).encode()
req = urllib.request.Request(url, data=payload, headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=30) as r:
data = json.loads(r.read())
text = data["candidates"][0]["content"]["parts"][0]["text"].strip()
zh = re.search(r"ZH:\s*(.+)", text)
en = re.search(r"EN:\s*(.+)", text)
tags_m = re.search(r"TAGS:\s*(.+)", text)
summary_parts = [x.group(1).strip() for x in [zh, en] if x]
summary = " | ".join(summary_parts) if summary_parts else ""
zh_tags = [t.strip() for t in tags_m.group(1).split(",") if t.strip()] if tags_m else []
return summary, zh_tags[:8]
except Exception:
return "", []
INDEX_PATH = WORKSPACE / "memory" / "bookmarks" / "search_index.json"
INGESTED_LOG = WORKSPACE / "memory" / "x-knowledge-base" / "pdf-ingested.json"
# ── Text cleaning ──────────────────────────────────────────────────────────────
def _clean_text(text: str) -> str:
"""Remove common academic PDF noise: page numbers, headers, footers, URLs."""
lines = text.splitlines()
cleaned = []
for line in lines:
line = line.strip()
# Skip very short lines (likely page numbers or artifacts)
if len(line) <= 3:
continue
# Skip lines that are just numbers (page numbers)
if re.fullmatch(r"\d+", line):
continue
# Skip common header/footer patterns
if re.match(r"^(arXiv|doi:|https?://|©|Copyright|All rights reserved)", line, re.IGNORECASE):
continue
cleaned.append(line)
# Collapse multiple blank lines
text = "\n".join(cleaned)
text = re.sub(r"\n{3,}", "\n\n", text)
return text.strip()
def _extract_title_author(doc: fitz.Document, filename: str) -> tuple[str, str]:
"""Try to extract title and author from PDF metadata, fallback to filename."""
meta = doc.metadata or {}
title = (meta.get("title") or "").strip()
author = (meta.get("author") or "").strip()
# If no metadata title, try first page heuristic
if not title and doc.page_count > 0:
first_page_text = doc[0].get_text("text")
lines = [l.strip() for l in first_page_text.splitlines() if len(l.strip()) > 10]
if lines:
# First substantial line is often the title
title = lines[0][:150]
# Fallback to filename
if not title:
title = Path(filename).stem.replace("_", " ").replace("-", " ")
return title, author
# ── PDF extraction ─────────────────────────────────────────────────────────────
def extract_pdf(path: Path) -> Optional[dict]:
"""Extract text and metadata from a PDF. Returns None if extraction fails."""
try:
doc = fitz.open(str(path))
except Exception as e:
print(f" ✗ Failed to open {path.name}: {e}")
return None
if doc.page_count == 0:
print(f" ✗ Empty PDF: {path.name}")
return None
# Extract full text
pages_text = []
for page in doc:
pages_text.append(page.get_text("text"))
full_text = "\n\n".join(pages_text)
full_text = _clean_text(full_text)
if len(full_text) < 200:
print(f" ✗ Too little text extracted from {path.name} ({len(full_text)} chars)")
return None
title, author = _extract_title_author(doc, path.name)
doc.close()
# Generate a stable ID from file path
file_id = hashlib.md5(str(path).encode()).hexdigest()[:12]
# Summary: try bilingual LLM summary, fallback to first 500 chars
en_excerpt = full_text[:500].replace("\n", " ").strip()
if len(en_excerpt) > 480:
en_excerpt = en_excerpt[:480] + "…"
bilingual, llm_tags = _llm_bilingual_summary(title, en_excerpt)
summary = bilingual if bilingual else en_excerpt
return {
"id": file_id,
"title": title,
"author": author,
"source_file": str(path),
"filename": path.name,
"page_count": len(pages_text),
"char_count": len(full_text),
"full_text": full_text,
"summary": summary,
"llm_tags": llm_tags,
"extracted_at": datetime.now(timezone.utc).isoformat(),
}
# ── Text file extraction (for fetch_pubmed.py output) ─────────────────────────
def extract_text_file(path: Path) -> Optional[dict]:
"""Extract from markdown/txt files (e.g. from fetch_pubmed.py)."""
try:
content = path.read_text(encoding="utf-8")
except Exception as e:
print(f" ✗ Failed to read {path.name}: {e}")
return None
if len(content) < 200:
print(f" ✗ Too little content in {path.name}")
return None
# Try to extract title from first heading
lines = content.splitlines()
title = ""
for line in lines[:5]:
if line.startswith("# "):
title = line.lstrip("# ").strip()
break
if not title:
title = path.stem.replace("_", " ").replace("-", " ")
# Author from **Authors:** line
author = ""
for line in lines[:10]:
m = re.search(r'\*\*Authors?:\*\*\s*(.+)', line)
if m:
author = m.group(1).strip()
break
en_excerpt = content[:500].replace("\n", " ").strip()
if len(en_excerpt) > 480:
en_excerpt = en_excerpt[:480] + "…"
bilingual, llm_tags = _llm_bilingual_summary(title, en_excerpt)
summary = bilingual if bilingual else en_excerpt
file_id = hashlib.md5(str(path).encode()).hexdigest()[:12]
return {
"id": file_id,
"title": title,
"author": author,
"source_file": str(path),
"filename": path.name,
"page_count": 0,
"char_count": len(content),
"full_text": content,
"summary": summary,
"llm_tags": llm_tags,
"extracted_at": datetime.now(timezone.utc).isoformat(),
}
# ── Knowledge card builder ─────────────────────────────────────────────────────
def _build_card(extracted: dict, category: str, tags: list[str]) -> dict:
"""Convert extracted PDF data into XKB bookmarks index format (compatible with xkb_ask.py)."""
llm_tags = extracted.get("llm_tags") or []
all_tags = tags + llm_tags + (["author:" + extracted["author"]] if extracted["author"] else [])
source = extracted["source_file"]
# Prepend Chinese summary to searchable for better Chinese query matching
zh_summary = extracted["summary"].split(" | ")[0] if " | " in extracted["summary"] else ""
searchable = (zh_summary + " " if zh_summary else "") + extracted["full_text"][:8000]
return {
"path": source,
"relative_path": extracted["filename"],
"title": extracted["title"],
"category": category,
"tags": all_tags,
"summary": extracted["summary"],
"source_url": source,
"source_type": "pdf",
"searchable": searchable[:8200],
"mtime": extracted["extracted_at"],
"size": extracted["char_count"],
"enriched": True,
"_pdf_id": extracted["id"],
"author": extracted["author"],
"page_count": extracted["page_count"],
"quality_score": min(1.0, extracted["char_count"] / 10000),
}
# ── Index management ───────────────────────────────────────────────────────────
def _load_index() -> tuple[list, dict]:
"""Returns (items_list, metadata_dict) from bookmarks index."""
if INDEX_PATH.exists():
try:
data = json.loads(INDEX_PATH.read_text(encoding="utf-8"))
if isinstance(data, dict):
return data.get("items", []), {k: v for k, v in data.items() if k != "items"}
elif isinstance(data, list):
return data, {}
except Exception:
pass
return [], {}
def _save_index(cards: list, meta: dict) -> None:
INDEX_PATH.parent.mkdir(parents=True, exist_ok=True)
output = {**meta, "count": len(cards), "items": cards}
INDEX_PATH.write_text(json.dumps(output, ensure_ascii=False, indent=2), encoding="utf-8")
def _load_ingested_log() -> set:
if INGESTED_LOG.exists():
try:
return set(json.loads(INGESTED_LOG.read_text(encoding="utf-8")))
except Exception:
return set()
return set()
def _save_ingested_log(ingested: set) -> None:
INGESTED_LOG.parent.mkdir(parents=True, exist_ok=True)
INGESTED_LOG.write_text(json.dumps(sorted(ingested), ensure_ascii=False, indent=2), encoding="utf-8")
# ── Vector index rebuild ───────────────────────────────────────────────────────
def _rebuild_vector_index() -> None:
script = SCRIPTS_DIR / "build_vector_index.py"
if not script.exists():
print(" ⚠ build_vector_index.py not found, skipping vector index")
return
print("\n🔢 Rebuilding vector index (Gemini embeddings)...")
import subprocess
result = subprocess.run(
[sys.executable, str(script)],
capture_output=False,
env={**os.environ, "OPENCLAW_WORKSPACE": str(WORKSPACE)},
)
if result.returncode == 0:
print(" ✓ Vector index rebuilt")
else:
print(" ✗ Vector index rebuild failed (check GEMINI_API_KEY)")
# ── Main ───────────────────────────────────────────────────────────────────────
def ingest(
path: Path,
category: str = "research",
tags: list[str] | None = None,
dry_run: bool = False,
rebuild_index: bool = False,
force: bool = False,
) -> int:
tags = tags or []
all_files: list[Path] = []
if path.is_dir():
pdfs = sorted(path.glob("**/*.pdf"))
texts = sorted(path.glob("**/*.md")) + sorted(path.glob("**/*.txt"))
all_files = pdfs + texts
print(f"📂 Found {len(pdfs)} PDFs + {len(texts)} text/markdown files in {path}")
elif path.suffix.lower() in (".pdf", ".md", ".txt"):
all_files = [path]
else:
print(f"Error: {path} is not a PDF/MD/TXT or directory")
return 1
if not all_files:
print("No files found.")
return 0
pdf_files = all_files # keep variable name for compatibility
ingested_ids = _load_ingested_log()
existing_index, index_meta = _load_index()
existing_ids = {c.get("_pdf_id") for c in existing_index if c.get("_pdf_id")}
new_cards = []
skipped = 0
failed = 0
for i, pdf_path in enumerate(pdf_files, 1):
file_id = hashlib.md5(str(pdf_path).encode()).hexdigest()[:12]
print(f"\n[{i}/{len(pdf_files)}] {pdf_path.name}")
if not force and file_id in ingested_ids:
print(f" ↩ Already ingested, skipping (use --force to re-ingest)")
skipped += 1
continue
if pdf_path.suffix.lower() in (".md", ".txt"):
extracted = extract_text_file(pdf_path)
else:
extracted = extract_pdf(pdf_path)
if not extracted:
failed += 1
continue
card = _build_card(extracted, category, tags)
print(f" ✓ {extracted['title'][:70]}")
print(f" {extracted['page_count']} pages | {extracted['char_count']:,} chars | score={card['quality_score']:.2f}")
if not dry_run:
new_cards.append(card)
ingested_ids.add(file_id)
if dry_run:
print(f"\n[dry-run] Would add {len(pdf_files) - skipped - failed} cards")
print(f" Skipped (already ingested): {skipped}")
print(f" Failed: {failed}")
return 0
if new_cards:
# Merge into existing index (dedup by _pdf_id)
new_ids = {c["_pdf_id"] for c in new_cards}
updated_index = [c for c in existing_index if c.get("_pdf_id") not in new_ids]
updated_index.extend(new_cards)
_save_index(updated_index, index_meta)
_save_ingested_log(ingested_ids)
print(f"\n✅ Added {len(new_cards)} cards to index (total: {len(updated_index)})")
else:
print(f"\n⚠ No new cards added (skipped={skipped}, failed={failed})")
if rebuild_index or new_cards:
_rebuild_vector_index()
return 0
def main() -> int:
parser = argparse.ArgumentParser(description="PDF Ingest — XKB knowledge pipeline")
parser.add_argument("path", help="PDF file or folder of PDFs")
parser.add_argument("--category", default="research", help="Category tag (default: research)")
parser.add_argument("--tag", action="append", dest="tags", default=[], help="Additional tags (repeatable)")
parser.add_argument("--dry-run", action="store_true", help="Preview without writing")
parser.add_argument("--force", action="store_true", help="Re-ingest already processed files")
parser.add_argument("--rebuild-index", action="store_true", help="Force rebuild vector index after ingest")
args = parser.parse_args()
return ingest(
path=Path(args.path),
category=args.category,
tags=args.tags,
dry_run=args.dry_run,
rebuild_index=args.rebuild_index,
force=args.force,
)
import xkb_usage # noqa: E402 — 量測誰在跑，見 scripts/xkb_usage.py
if __name__ == "__main__":
xkb_usage.record(__file__)
raise SystemExit(main())