230 lines
8.2 KiB
Python
230 lines
8.2 KiB
Python
"""Phase 0 user-provided material ingestion."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import os
|
||
import shutil
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import requests
|
||
|
||
|
||
DEFAULT_FIRERED_OCR_ENDPOINT = "http://192.168.50.100:8001"
|
||
DEFAULT_OCR_MAX_PAGES = 50
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class OcrResult:
|
||
text: str
|
||
pages_processed: int
|
||
output_path: Path
|
||
|
||
|
||
def safe_filename(path: Path) -> str:
|
||
name = path.name.strip()
|
||
return name or "material"
|
||
|
||
|
||
def extract_pdf_text(path: Path) -> tuple[str, int, bool]:
|
||
from pypdf import PdfReader
|
||
|
||
reader = PdfReader(str(path))
|
||
chunks: list[str] = []
|
||
for index, page in enumerate(reader.pages, start=1):
|
||
text = (page.extract_text() or "").strip()
|
||
if text:
|
||
chunks.append(f"\n\n## Page {index}\n\n{text}")
|
||
combined = "".join(chunks).strip()
|
||
ocr_required = len(combined) < max(20, len(reader.pages) * 20)
|
||
return combined, len(reader.pages), ocr_required
|
||
|
||
|
||
def ocr_endpoint_from_env() -> str:
|
||
return os.environ.get("DEEP_RESEARCH_OCR_ENDPOINT", DEFAULT_FIRERED_OCR_ENDPOINT).rstrip("/")
|
||
|
||
|
||
def ocr_max_pages_from_env() -> int:
|
||
raw = os.environ.get("DEEP_RESEARCH_OCR_MAX_PAGES")
|
||
if not raw:
|
||
return DEFAULT_OCR_MAX_PAGES
|
||
try:
|
||
return max(1, int(raw))
|
||
except ValueError:
|
||
return DEFAULT_OCR_MAX_PAGES
|
||
|
||
|
||
def render_pdf_pages(pdf_path: Path, output_dir: Path, *, max_pages: int) -> list[Path]:
|
||
import fitz
|
||
|
||
pages_dir = output_dir / f"{pdf_path.stem}.ocr-pages"
|
||
pages_dir.mkdir(parents=True, exist_ok=True)
|
||
image_paths: list[Path] = []
|
||
doc = fitz.open(pdf_path)
|
||
try:
|
||
for index, page in enumerate(doc[:max_pages], start=1):
|
||
pix = page.get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False)
|
||
image_path = pages_dir / f"page-{index:03d}.png"
|
||
pix.save(image_path)
|
||
image_paths.append(image_path)
|
||
finally:
|
||
doc.close()
|
||
return image_paths
|
||
|
||
|
||
def data_url_for_image(path: Path) -> str:
|
||
encoded = base64.b64encode(path.read_bytes()).decode("ascii")
|
||
return f"data:image/png;base64,{encoded}"
|
||
|
||
|
||
def call_firered_ocr(image_path: Path, *, endpoint: str) -> str:
|
||
payload = {
|
||
"model": "firered-ocr",
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": "请识别图片中的全部文字，保持原有顺序，只输出文字。"},
|
||
{"type": "image_url", "image_url": {"url": data_url_for_image(image_path)}},
|
||
],
|
||
}
|
||
],
|
||
"temperature": 0,
|
||
"max_tokens": 3000,
|
||
}
|
||
response = requests.post(f"{endpoint.rstrip('/')}/v1/chat/completions", json=payload, timeout=60)
|
||
if not response.ok:
|
||
raise RuntimeError(f"{response.status_code} {response.text[:500]}")
|
||
data = response.json()
|
||
return str(data["choices"][0]["message"].get("content") or "").strip()
|
||
|
||
|
||
def ocr_pdf_with_firered(*, pdf_path: Path, output_dir: Path, endpoint: str, max_pages: int) -> OcrResult:
|
||
image_paths = render_pdf_pages(pdf_path, output_dir, max_pages=max_pages)
|
||
chunks: list[str] = []
|
||
for index, image_path in enumerate(image_paths, start=1):
|
||
text = call_firered_ocr(image_path, endpoint=endpoint)
|
||
if text:
|
||
chunks.append(f"\n\n## OCR Page {index}\n\n{text}")
|
||
combined = "".join(chunks).strip()
|
||
output_path = output_dir / f"{pdf_path.stem}.ocr.md"
|
||
body = [
|
||
f"# OCR Material: {pdf_path.name}",
|
||
"",
|
||
f"- source_path: {pdf_path}",
|
||
f"- endpoint: {endpoint}",
|
||
f"- pages_processed: {len(image_paths)}",
|
||
"",
|
||
combined or "OCR 未返回可用文本。",
|
||
"",
|
||
]
|
||
output_path.write_text("\n".join(body), encoding="utf-8")
|
||
return OcrResult(text=combined, pages_processed=len(image_paths), output_path=output_path)
|
||
|
||
|
||
def ingest_input_materials(project_root: Path, materials: list[str] | None) -> list[dict[str, Any]]:
|
||
inventory: list[dict[str, Any]] = []
|
||
if not materials:
|
||
return inventory
|
||
|
||
inputs_dir = project_root / "phase0" / "inputs"
|
||
extracted_dir = project_root / "phase0" / "extracted"
|
||
ocr_endpoint = ocr_endpoint_from_env()
|
||
ocr_max_pages = ocr_max_pages_from_env()
|
||
inputs_dir.mkdir(parents=True, exist_ok=True)
|
||
extracted_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
for raw in materials:
|
||
source = Path(raw).expanduser()
|
||
if not source.exists():
|
||
inventory.append({"kind": "note", "note": raw})
|
||
continue
|
||
|
||
copied = inputs_dir / safe_filename(source)
|
||
shutil.copy2(source, copied)
|
||
item: dict[str, Any] = {
|
||
"kind": source.suffix.lower().lstrip(".") or "file",
|
||
"source_path": str(source),
|
||
"copied_to": str(copied.relative_to(project_root)),
|
||
"size_bytes": source.stat().st_size,
|
||
}
|
||
|
||
if source.suffix.lower() == ".pdf":
|
||
text, pages, ocr_required = extract_pdf_text(source)
|
||
extracted = extracted_dir / f"{source.stem}.md"
|
||
ocr_result: OcrResult | None = None
|
||
ocr_error: str | None = None
|
||
if ocr_required:
|
||
try:
|
||
ocr_result = ocr_pdf_with_firered(
|
||
pdf_path=source,
|
||
output_dir=extracted_dir,
|
||
endpoint=ocr_endpoint,
|
||
max_pages=min(pages, ocr_max_pages),
|
||
)
|
||
if ocr_result.text:
|
||
text = "\n\n".join(part for part in [text, ocr_result.text] if part)
|
||
except Exception as exc: # noqa: BLE001 - ingestion should not block project init.
|
||
ocr_error = str(exc)
|
||
|
||
body = [
|
||
f"# Extracted Material: {source.name}",
|
||
"",
|
||
f"- source_path: {source}",
|
||
f"- copied_to: {copied.relative_to(project_root)}",
|
||
f"- pages: {pages}",
|
||
f"- ocr_required: {str(ocr_required).lower()}",
|
||
f"- ocr_status: {'completed' if ocr_result else 'failed' if ocr_error else 'not_required'}",
|
||
"",
|
||
text or "未能从 PDF 直接抽取文本；该材料可能需要 OCR。",
|
||
"",
|
||
]
|
||
if ocr_error:
|
||
body.extend(["## OCR Error", "", ocr_error, ""])
|
||
extracted.write_text("\n".join(body), encoding="utf-8")
|
||
item.update(
|
||
{
|
||
"pages": pages,
|
||
"extracted_to": str(extracted.relative_to(project_root)),
|
||
"text_chars": len(text),
|
||
"ocr_required": ocr_required,
|
||
"ocr_status": "completed" if ocr_result else "failed" if ocr_error else "not_required",
|
||
}
|
||
)
|
||
if ocr_result:
|
||
item.update(
|
||
{
|
||
"ocr_endpoint": ocr_endpoint,
|
||
"ocr_pages_processed": ocr_result.pages_processed,
|
||
"ocr_extracted_to": str(ocr_result.output_path.relative_to(project_root)),
|
||
"ocr_text_chars": len(ocr_result.text),
|
||
}
|
||
)
|
||
if ocr_error:
|
||
item["ocr_error"] = ocr_error
|
||
else:
|
||
item["ocr_required"] = source.suffix.lower() in {".png", ".jpg", ".jpeg", ".tif", ".tiff"}
|
||
inventory.append(item)
|
||
|
||
return inventory
|
||
|
||
|
||
def render_material_inventory(inventory: list[dict[str, Any]]) -> str:
|
||
if not inventory:
|
||
return "- 暂无；可通过 `--input-material` 加入审计报告、问题清单或内部记录。"
|
||
lines: list[str] = []
|
||
for item in inventory:
|
||
if item.get("kind") == "note":
|
||
lines.append(f"- 备注：{item.get('note', '')}")
|
||
continue
|
||
marker = "；需要 OCR" if item.get("ocr_required") else ""
|
||
ocr = f"；OCR：{item.get('ocr_status')}" if item.get("ocr_status") else ""
|
||
extracted = item.get("extracted_to")
|
||
extra = f"；抽取文本：{extracted}" if extracted else ""
|
||
lines.append(
|
||
f"- {item.get('copied_to')}（{item.get('kind')}，{item.get('size_bytes', 0)} bytes{extra}{marker}{ocr}）"
|
||
)
|
||
return "\n".join(lines)
|