cleanup/core/content_extract.py
2026-09-05 15:24:52 +02:00

87 lines
2.6 KiB
Python

"""
Extracts a text snippet from a file so an LLM suggestion can be based on
what's actually in the file, not just its name — helps most for the
genuinely opaque cases (scans, "Download (3).jpg"-style names) that a
filename-only guess can't do anything with.
.md/.txt need nothing extra. .pdf/.docx/.xlsx each need one small
pip-installed library:
pip install pypdf python-docx openpyxl --break-system-packages
If a library isn't installed, that format quietly returns None (treated
as "no content available", falls back to filename-only) rather than
crashing anything — one missing library never blocks the others.
"""
from pathlib import Path
from typing import Optional
_MAX_CHARS = 2000 # a snippet is enough for a topic guess; keeps the LLM prompt small and fast
def extract_text(path: Path) -> Optional[str]:
ext = path.suffix.lower()
try:
if ext in (".md", ".txt", ".csv"):
return _read_plain(path)
if ext == ".pdf":
return _read_pdf(path)
if ext == ".docx":
return _read_docx(path)
if ext == ".xlsx":
return _read_xlsx(path)
except Exception:
return None # a corrupt/unreadable file falls back to filename-only, not a crash
return None # unsupported format
def _read_plain(path: Path) -> str:
return path.read_text(encoding="utf-8", errors="ignore")[:_MAX_CHARS]
def _read_pdf(path: Path) -> Optional[str]:
try:
from pypdf import PdfReader
except ImportError:
return None
reader = PdfReader(str(path))
chunks = []
total = 0
for page in reader.pages[:3]: # first few pages is plenty for a topic guess
text = page.extract_text() or ""
chunks.append(text)
total += len(text)
if total > _MAX_CHARS:
break
return "".join(chunks)[:_MAX_CHARS]
def _read_docx(path: Path) -> Optional[str]:
try:
import docx
except ImportError:
return None
document = docx.Document(str(path))
text = "\n".join(p.text for p in document.paragraphs)
return text[:_MAX_CHARS]
def _read_xlsx(path: Path) -> Optional[str]:
try:
import openpyxl
except ImportError:
return None
wb = openpyxl.load_workbook(str(path), read_only=True, data_only=True)
parts = []
total = 0
for ws in wb.worksheets[:2]:
for row in ws.iter_rows(max_row=20, values_only=True):
line = " ".join(str(c) for c in row if c is not None)
parts.append(line)
total += len(line)
if total > _MAX_CHARS:
break
if total > _MAX_CHARS:
break
return "\n".join(parts)[:_MAX_CHARS]