149 lines
4.9 KiB
Python
149 lines
4.9 KiB
Python
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import logging
|
||
import subprocess
|
||
import tempfile
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
|
||
from django.conf import settings
|
||
from docx import Document
|
||
from openpyxl import load_workbook
|
||
from pypdf import PdfReader
|
||
from pptx import Presentation
|
||
|
||
from .rag_embedding import EmbeddingFunction
|
||
|
||
|
||
logger = logging.getLogger("review_agent.regulatory_review.rag_index")
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class TextChunk:
|
||
text: str
|
||
metadata: dict[str, object]
|
||
|
||
|
||
def chunk_text(text: str, *, source: str, chunk_size: int = 900, overlap: int = 120) -> list[TextChunk]:
|
||
normalized = "\n".join(line.strip() for line in text.splitlines() if line.strip())
|
||
if not normalized:
|
||
return []
|
||
chunks = []
|
||
start = 0
|
||
index = 0
|
||
step = max(1, chunk_size - overlap)
|
||
while start < len(normalized):
|
||
part = normalized[start : start + chunk_size].strip()
|
||
if part:
|
||
chunks.append(TextChunk(text=part, metadata={"source": source, "chunk_index": index}))
|
||
index += 1
|
||
start += step
|
||
return chunks
|
||
|
||
|
||
def extract_text_from_path(path: Path) -> str:
|
||
suffix = path.suffix.lower()
|
||
if suffix in {".txt", ".md"}:
|
||
return path.read_text(encoding="utf-8", errors="ignore")
|
||
if suffix == ".pdf":
|
||
return "\n".join(page.extract_text() or "" for page in PdfReader(str(path)).pages)
|
||
if suffix == ".docx":
|
||
return "\n".join(paragraph.text for paragraph in Document(str(path)).paragraphs)
|
||
if suffix == ".pptx":
|
||
presentation = Presentation(str(path))
|
||
lines = []
|
||
for slide in presentation.slides:
|
||
for shape in slide.shapes:
|
||
if hasattr(shape, "text"):
|
||
lines.append(shape.text)
|
||
return "\n".join(lines)
|
||
if suffix == ".xlsx":
|
||
workbook = load_workbook(path, data_only=True, read_only=True)
|
||
lines = []
|
||
for sheet in workbook.worksheets:
|
||
for row in sheet.iter_rows(values_only=True):
|
||
values = [str(cell) for cell in row if cell not in {None, ""}]
|
||
if values:
|
||
lines.append("\t".join(values))
|
||
return "\n".join(lines)
|
||
if suffix == ".doc":
|
||
return _extract_legacy_doc_with_libreoffice(path)
|
||
return ""
|
||
|
||
|
||
def _extract_legacy_doc_with_libreoffice(path: Path) -> str:
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
target_dir = Path(tmp_dir)
|
||
try:
|
||
subprocess.run(
|
||
[
|
||
"soffice",
|
||
"--headless",
|
||
"--convert-to",
|
||
"docx",
|
||
"--outdir",
|
||
str(target_dir),
|
||
str(path),
|
||
],
|
||
check=True,
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=60,
|
||
)
|
||
except (FileNotFoundError, subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
|
||
raise RuntimeError(f"无法通过 LibreOffice 转换法规 .doc 材料：{path.name}") from exc
|
||
converted = target_dir / f"{path.stem}.docx"
|
||
if not converted.exists():
|
||
raise RuntimeError(f"LibreOffice 未生成 docx：{path.name}")
|
||
return extract_text_from_path(converted)
|
||
|
||
|
||
def collect_source_chunks(source_dir: Path) -> list[TextChunk]:
|
||
chunks: list[TextChunk] = []
|
||
for path in sorted(source_dir.rglob("*")):
|
||
if not path.is_file():
|
||
continue
|
||
try:
|
||
text = extract_text_from_path(path)
|
||
except RuntimeError as exc:
|
||
logger.warning("Regulatory source extraction skipped", extra={"path": str(path), "error": str(exc)})
|
||
continue
|
||
chunks.extend(chunk_text(text, source=str(path.relative_to(source_dir))))
|
||
return chunks
|
||
|
||
|
||
def build_chroma_index(
|
||
*,
|
||
source_dir: Path,
|
||
embedding_provider: EmbeddingFunction,
|
||
persist_path: Path | None = None,
|
||
collection_name: str | None = None,
|
||
) -> int:
|
||
try:
|
||
import chromadb
|
||
except ImportError as exc:
|
||
raise RuntimeError("chromadb 未安装，请先安装 requirements.txt。") from exc
|
||
|
||
persist_path = persist_path or Path(settings.REGULATORY_RAG_CHROMA_PATH)
|
||
collection_name = collection_name or settings.REGULATORY_RAG_COLLECTION
|
||
persist_path.mkdir(parents=True, exist_ok=True)
|
||
chunks = collect_source_chunks(source_dir)
|
||
client = chromadb.PersistentClient(path=str(persist_path))
|
||
collection = client.get_or_create_collection(collection_name)
|
||
if not chunks:
|
||
return 0
|
||
texts = [chunk.text for chunk in chunks]
|
||
embeddings = embedding_provider(texts)
|
||
ids = [
|
||
hashlib.sha256(f"{chunk.metadata['source']}:{chunk.metadata['chunk_index']}".encode("utf-8")).hexdigest()
|
||
for chunk in chunks
|
||
]
|
||
collection.upsert(
|
||
ids=ids,
|
||
documents=texts,
|
||
metadatas=[chunk.metadata for chunk in chunks],
|
||
embeddings=embeddings,
|
||
)
|
||
return len(chunks)
|