v0.20 alpha skill-driven python core
This commit is contained in:
kai
@@ -0,0 +1,229 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user