"""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)