feat: group private apps into andreknie-privat
This commit is contained in:
@@ -0,0 +1,51 @@
|
||||
"""Text extraction layer for various document formats."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from .base import ExtractionResult, TextExtractor
|
||||
from .docx import DocxExtractor
|
||||
from .ocr import OCRExtractor
|
||||
from .pdf import PDFExtractor
|
||||
from .text import PlainTextExtractor
|
||||
|
||||
|
||||
def get_extractor(file_path: Path, config=None):
|
||||
"""Get the appropriate extractor for a file.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file to extract.
|
||||
config: Optional IngestionConfig for LLM-based extraction.
|
||||
|
||||
Returns:
|
||||
An extractor instance that can handle the file, or None.
|
||||
"""
|
||||
suffix = file_path.suffix.lower()
|
||||
|
||||
if suffix in (".md", ".txt"):
|
||||
return PlainTextExtractor()
|
||||
elif suffix == ".pdf":
|
||||
if config:
|
||||
return PDFExtractor(
|
||||
ocr_enabled=config.ocr_enabled,
|
||||
ocr_language=config.ocr_language,
|
||||
llm_provider=config.llm_provider,
|
||||
llm_model=config.llm_model,
|
||||
mistral_api_key=config.mistral_api_key,
|
||||
)
|
||||
return PDFExtractor()
|
||||
elif suffix in (".jpg", ".jpeg", ".png"):
|
||||
return OCRExtractor()
|
||||
elif suffix == ".docx":
|
||||
return DocxExtractor()
|
||||
return None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ExtractionResult",
|
||||
"TextExtractor",
|
||||
"PlainTextExtractor",
|
||||
"PDFExtractor",
|
||||
"OCRExtractor",
|
||||
"DocxExtractor",
|
||||
"get_extractor",
|
||||
]
|
||||
@@ -0,0 +1,27 @@
|
||||
"""Base protocol and data structures for text extraction."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Protocol
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExtractionResult:
|
||||
"""Result of extracting text from a document."""
|
||||
|
||||
text: str
|
||||
source_file: Path
|
||||
extraction_method: str # "direct", "ocr", "mixed"
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class TextExtractor(Protocol):
|
||||
"""Protocol that all text extractors must implement."""
|
||||
|
||||
def can_handle(self, file_path: Path) -> bool:
|
||||
"""Return True if this extractor can process the given file."""
|
||||
...
|
||||
|
||||
def extract(self, file_path: Path) -> ExtractionResult:
|
||||
"""Extract text content from the file."""
|
||||
...
|
||||
@@ -0,0 +1,57 @@
|
||||
"""DOCX text extraction using python-docx.
|
||||
|
||||
Preserves heading structure and paragraph breaks.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from .base import ExtractionResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DocxExtractor:
|
||||
"""Extracts text from .docx files preserving structure."""
|
||||
|
||||
def can_handle(self, file_path: Path) -> bool:
|
||||
"""Return True for .docx files."""
|
||||
return file_path.suffix.lower() == ".docx"
|
||||
|
||||
def extract(self, file_path: Path) -> ExtractionResult:
|
||||
"""Extract text from DOCX preserving headings and paragraphs."""
|
||||
try:
|
||||
from docx import Document
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"python-docx is required for DOCX extraction. "
|
||||
"Install with: pip install python-docx"
|
||||
)
|
||||
|
||||
doc = Document(str(file_path))
|
||||
parts: list[str] = []
|
||||
|
||||
for paragraph in doc.paragraphs:
|
||||
text = paragraph.text.strip()
|
||||
if not text:
|
||||
continue
|
||||
|
||||
# Preserve heading structure with markdown-style headers
|
||||
style_name = paragraph.style.name.lower() if paragraph.style else ""
|
||||
if "heading 1" in style_name:
|
||||
parts.append(f"# {text}")
|
||||
elif "heading 2" in style_name:
|
||||
parts.append(f"## {text}")
|
||||
elif "heading 3" in style_name:
|
||||
parts.append(f"### {text}")
|
||||
elif "heading 4" in style_name:
|
||||
parts.append(f"#### {text}")
|
||||
else:
|
||||
parts.append(text)
|
||||
|
||||
return ExtractionResult(
|
||||
text="\n\n".join(parts),
|
||||
source_file=file_path,
|
||||
extraction_method="direct",
|
||||
metadata={"paragraph_count": len(parts)},
|
||||
)
|
||||
@@ -0,0 +1,48 @@
|
||||
"""OCR extraction for images using pytesseract.
|
||||
|
||||
Handles .jpg, .jpeg, .png files. Configured for German + English (deu+eng).
|
||||
Returns empty text (not error) when OCR cannot extract readable content.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from .base import ExtractionResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SUPPORTED_EXTENSIONS = {".jpg", ".jpeg", ".png"}
|
||||
|
||||
|
||||
class OCRExtractor:
|
||||
"""Extracts text from images using Tesseract OCR."""
|
||||
|
||||
def __init__(self, language: str = "deu+eng"):
|
||||
self.language = language
|
||||
|
||||
def can_handle(self, file_path: Path) -> bool:
|
||||
"""Return True for .jpg, .jpeg, .png files."""
|
||||
return file_path.suffix.lower() in SUPPORTED_EXTENSIONS
|
||||
|
||||
def extract(self, file_path: Path) -> ExtractionResult:
|
||||
"""Run OCR on an image file. Returns empty text on failure."""
|
||||
text = self._run_ocr(file_path)
|
||||
return ExtractionResult(
|
||||
text=text,
|
||||
source_file=file_path,
|
||||
extraction_method="ocr",
|
||||
metadata={"ocr_language": self.language},
|
||||
)
|
||||
|
||||
def _run_ocr(self, file_path: Path) -> str:
|
||||
"""Execute pytesseract on the image. Returns empty string on failure."""
|
||||
try:
|
||||
import pytesseract
|
||||
from PIL import Image
|
||||
|
||||
image = Image.open(file_path)
|
||||
text = pytesseract.image_to_string(image, lang=self.language)
|
||||
return text.strip()
|
||||
except Exception as e:
|
||||
logger.warning("OCR extraction failed for %s: %s", file_path, e)
|
||||
return ""
|
||||
@@ -0,0 +1,156 @@
|
||||
"""PDF text extraction using pdfplumber.
|
||||
|
||||
Detects scanned pages (no extractable text) and delegates to:
|
||||
1. Mistral Vision API (preferred — handles handwriting well)
|
||||
2. Tesseract OCR (fallback if no LLM configured)
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import pdfplumber
|
||||
|
||||
from .base import ExtractionResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PDFExtractor:
|
||||
"""Extracts text from PDF files, with Vision LLM or OCR fallback for scanned pages."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ocr_enabled: bool = True,
|
||||
ocr_language: str = "deu+eng",
|
||||
llm_provider: str = "mistral",
|
||||
llm_model: str = "mistral-small-latest",
|
||||
mistral_api_key: str = "",
|
||||
):
|
||||
self.ocr_enabled = ocr_enabled
|
||||
self.ocr_language = ocr_language
|
||||
self.llm_provider = llm_provider
|
||||
self.llm_model = llm_model
|
||||
self.mistral_api_key = mistral_api_key
|
||||
self._use_vision = bool(mistral_api_key)
|
||||
|
||||
def can_handle(self, file_path: Path) -> bool:
|
||||
"""Return True for .pdf files."""
|
||||
return file_path.suffix.lower() == ".pdf"
|
||||
|
||||
def extract(self, file_path: Path) -> ExtractionResult:
|
||||
"""Extract text from PDF, using Vision LLM or OCR for scanned pages."""
|
||||
pages_text: list[str] = []
|
||||
has_direct_text = False
|
||||
has_vision_text = False
|
||||
|
||||
with pdfplumber.open(file_path) as pdf:
|
||||
for page_num, page in enumerate(pdf.pages, start=1):
|
||||
text = page.extract_text() or ""
|
||||
text = text.strip()
|
||||
|
||||
if text:
|
||||
pages_text.append(text)
|
||||
has_direct_text = True
|
||||
elif self.ocr_enabled:
|
||||
# Scanned page — try Vision LLM first, then Tesseract
|
||||
logger.debug(
|
||||
"Page %d of %s has no text, attempting vision OCR",
|
||||
page_num, file_path,
|
||||
)
|
||||
ocr_text = self._vision_ocr_page(page, page_num, file_path)
|
||||
if ocr_text:
|
||||
pages_text.append(ocr_text)
|
||||
has_vision_text = True
|
||||
|
||||
# Determine extraction method
|
||||
if has_direct_text and has_vision_text:
|
||||
method = "mixed"
|
||||
elif has_vision_text:
|
||||
method = "vision-ocr"
|
||||
else:
|
||||
method = "direct"
|
||||
|
||||
return ExtractionResult(
|
||||
text="\n\n".join(pages_text),
|
||||
source_file=file_path,
|
||||
extraction_method=method,
|
||||
metadata={"page_count": len(pages_text)},
|
||||
)
|
||||
|
||||
def _vision_ocr_page(self, page, page_num: int, file_path: Path) -> str:
|
||||
"""OCR a PDF page using Mistral Vision API, with Tesseract fallback."""
|
||||
if self._use_vision:
|
||||
try:
|
||||
return self._mistral_vision_ocr(page, page_num)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Vision OCR failed for page %d of %s: %s, trying Tesseract",
|
||||
page_num, file_path, e,
|
||||
)
|
||||
|
||||
# Fallback to Tesseract
|
||||
return self._tesseract_ocr(page, page_num, file_path)
|
||||
|
||||
def _mistral_vision_ocr(self, page, page_num: int) -> str:
|
||||
"""Send page image to Mistral Vision for text extraction."""
|
||||
import httpx
|
||||
|
||||
# Convert page to PNG image
|
||||
image = page.to_image(resolution=200)
|
||||
img_buffer = io.BytesIO()
|
||||
image.original.save(img_buffer, format="PNG")
|
||||
img_bytes = img_buffer.getvalue()
|
||||
img_b64 = base64.b64encode(img_bytes).decode("utf-8")
|
||||
|
||||
# Call Mistral Vision API via LiteLLM-compatible format
|
||||
response = httpx.post(
|
||||
"https://api.mistral.ai/v1/chat/completions",
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.mistral_api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json={
|
||||
"model": self.llm_model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Extract ALL text from this image. This is a handwritten note or document. "
|
||||
"Transcribe everything you can read, preserving structure (paragraphs, lists, headings). "
|
||||
"If text is in German, keep it in German. Return ONLY the extracted text, no commentary.",
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{img_b64}",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
"max_tokens": 4000,
|
||||
},
|
||||
timeout=60.0,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
text = data["choices"][0]["message"]["content"]
|
||||
logger.info("Vision OCR page %d: extracted %d chars", page_num, len(text))
|
||||
return text.strip()
|
||||
|
||||
def _tesseract_ocr(self, page, page_num: int, file_path: Path) -> str:
|
||||
"""Fallback: OCR a PDF page using Tesseract."""
|
||||
try:
|
||||
import pytesseract
|
||||
image = page.to_image(resolution=300)
|
||||
text = pytesseract.image_to_string(image.original, lang=self.ocr_language)
|
||||
return text.strip()
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Tesseract OCR failed for page %d of %s: %s", page_num, file_path, e
|
||||
)
|
||||
return ""
|
||||
@@ -0,0 +1,57 @@
|
||||
"""Plain text and markdown file extractor.
|
||||
|
||||
Handles .md and .txt files. Reads as UTF-8 by default, with fallback
|
||||
charset detection for non-UTF-8 encoded files.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from .base import ExtractionResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SUPPORTED_EXTENSIONS = {".md", ".txt"}
|
||||
|
||||
|
||||
class PlainTextExtractor:
|
||||
"""Extracts text from plain text and markdown files."""
|
||||
|
||||
def can_handle(self, file_path: Path) -> bool:
|
||||
"""Return True for .md and .txt files."""
|
||||
return file_path.suffix.lower() in SUPPORTED_EXTENSIONS
|
||||
|
||||
def extract(self, file_path: Path) -> ExtractionResult:
|
||||
"""Read file as UTF-8, falling back to charset detection."""
|
||||
text = self._read_with_fallback(file_path)
|
||||
return ExtractionResult(
|
||||
text=text,
|
||||
source_file=file_path,
|
||||
extraction_method="direct",
|
||||
metadata={"encoding": "utf-8"},
|
||||
)
|
||||
|
||||
def _read_with_fallback(self, file_path: Path) -> str:
|
||||
"""Try UTF-8 first, then attempt charset detection."""
|
||||
try:
|
||||
return file_path.read_text(encoding="utf-8")
|
||||
except UnicodeDecodeError:
|
||||
logger.warning(
|
||||
"UTF-8 decode failed for %s, attempting charset detection", file_path
|
||||
)
|
||||
return self._read_with_detection(file_path)
|
||||
|
||||
def _read_with_detection(self, file_path: Path) -> str:
|
||||
"""Detect encoding and read file."""
|
||||
raw = file_path.read_bytes()
|
||||
|
||||
# Try common encodings before giving up
|
||||
for encoding in ("latin-1", "cp1252", "iso-8859-1"):
|
||||
try:
|
||||
return raw.decode(encoding)
|
||||
except (UnicodeDecodeError, LookupError):
|
||||
continue
|
||||
|
||||
# Last resort: decode with replacement characters
|
||||
logger.warning("Could not detect encoding for %s, using replacement", file_path)
|
||||
return raw.decode("utf-8", errors="replace")
|
||||
Reference in New Issue
Block a user