From cfaccb3ec7664f70ac40d6911379fe5567c07d0e Mon Sep 17 00:00:00 2001 From: Maneet Gupta Date: Sat, 13 Jun 2026 12:32:42 +0530 Subject: [PATCH] Feat: Adding scripts for processing various files. Adding scripts to process various files, chunk them on rolling basis and pass through sentence transformers to make vector embeddings of dim 384 --- .gitignore | 3 +- app/example_usage.py | 126 ++++++++ app/processing/__init__.py | 27 ++ app/processing/embeddings.py | 109 +++++++ app/processing/extractors.py | 594 +++++++++++++++++++++++++++++++++++ app/processing/ocr.py | 99 ++++++ app/processing/pipeline.py | 161 ++++++++++ app/processing/traversal.py | 150 +++++++++ requirements.txt | Bin 920 -> 1428 bytes 9 files changed, 1268 insertions(+), 1 deletion(-) create mode 100644 app/example_usage.py create mode 100644 app/processing/__init__.py create mode 100644 app/processing/embeddings.py create mode 100644 app/processing/extractors.py create mode 100644 app/processing/ocr.py create mode 100644 app/processing/pipeline.py create mode 100644 app/processing/traversal.py diff --git a/.gitignore b/.gitignore index ced36aa..4898d55 100644 --- a/.gitignore +++ b/.gitignore @@ -51,4 +51,5 @@ docker-compose.override.yml # OS files # ====================== .DS_Store -Thumbs.db \ No newline at end of file +Thumbs.db +test_data/ \ No newline at end of file diff --git a/app/example_usage.py b/app/example_usage.py new file mode 100644 index 0000000..1971849 --- /dev/null +++ b/app/example_usage.py @@ -0,0 +1,126 @@ +""" +example_usage.py +---------------- +Shows how a FastAPI background task / CLI script would call the +processing package and persist results using SQLAlchemy. + +This file is NOT part of the processing package itself — it lives +alongside it as reference for the backend team. +""" + +# ── SQLAlchemy models (abbreviated) ───────────────────────────────────────── +# from app.models import File, FileContent +# from app.database import SessionLocal + +# ── Processing package ─────────────────────────────────────────────────────── +from processing import process + + +# ───────────────────────────────────────────────────────────────────────────── +# 1. Basic full-directory scan +# ───────────────────────────────────────────────────────────────────────────── + +def index_directory(root: str, db) -> None: + """ + Index every supported file under *root* and persist to the database. + + Parameters + ---------- + root : str + Folder or single file path to scan. + db : SQLAlchemy Session + Caller is responsible for lifecycle (commit / rollback / close). + """ + file_rows, content_rows = process(root) + + # ── Insert File rows ────────────────────────────────────────────────────── + for fr in file_rows: + # file_rows keys: file_path, file_hash, mime_type, last_modified, tags + existing = db.query(File).filter_by(file_path=fr["file_path"]).first() + if existing: + # Update hash / timestamp in case the file changed + for k, v in fr.items(): + setattr(existing, k, v) + else: + db.add(File(**fr)) + + db.flush() # ensure File PKs exist before FK inserts + + # ── Insert FileContent rows ─────────────────────────────────────────────── + for cr in content_rows: + file_path = cr.pop("file_path") # resolve FK + file_obj = db.query(File).filter_by(file_path=file_path).one() + + db.add(FileContent( + file_id = file_obj.id, + chunk_index = cr["chunk_index"], + content_text = cr["content_text"], + embedding = cr["embedding"], + # keyword_tokens is a server-generated TSVECTOR column; + # PostgreSQL fills it automatically — do not pass it here. + )) + + db.commit() + print(f"Indexed {len(file_rows)} file(s), {len(content_rows)} chunk(s).") + + +# ───────────────────────────────────────────────────────────────────────────── +# 2. Incremental re-index (skip unchanged files) +# ───────────────────────────────────────────────────────────────────────────── + +def reindex_directory(root: str, db) -> None: + """Like index_directory but skips files whose SHA-256 hash is unchanged.""" + existing_hashes = {row.file_hash for row in db.query(File.file_hash).all()} + file_rows, content_rows = process(root, skip_hashes=existing_hashes) + + for fr in file_rows: + db.merge(File(**fr)) # upsert on unique file_path + + db.flush() + + for cr in content_rows: + file_path = cr.pop("file_path") + file_obj = db.query(File).filter_by(file_path=file_path).one() + db.add(FileContent( + file_id = file_obj.id, + chunk_index = cr["chunk_index"], + content_text = cr["content_text"], + embedding = cr["embedding"], + )) + + db.commit() + + +# ───────────────────────────────────────────────────────────────────────────── +# 3. FastAPI background task example +# ───────────────────────────────────────────────────────────────────────────── + +# from fastapi import BackgroundTasks +# +# @app.post("/index") +# async def trigger_index(path: str, background_tasks: BackgroundTasks): +# background_tasks.add_task(index_directory, path, SessionLocal()) +# return {"status": "indexing started", "path": path} + + +# ───────────────────────────────────────────────────────────────────────────── +# 4. CLI quick-test +# ───────────────────────────────────────────────────────────────────────────── + +if __name__ == "__main__": + import sys, pprint + root = sys.argv[1] if len(sys.argv) > 1 else "." + file_rows, content_rows = process(root) + + print(f"\n{'='*60}") + print(f"Files found : {len(file_rows)}") + print(f"Chunks created: {len(content_rows)}") + print(f"{'='*60}\n") + + for cr in content_rows[:3]: + pprint.pprint({ + "file_path": cr["file_path"], + "chunk_index": cr["chunk_index"], + "content_text": cr["content_text"][:120] + "…", + "embedding_dim": len(cr["embedding"]), + }) diff --git a/app/processing/__init__.py b/app/processing/__init__.py new file mode 100644 index 0000000..2a5646d --- /dev/null +++ b/app/processing/__init__.py @@ -0,0 +1,27 @@ +""" +processing +---------- +WhereTF local-search processing system. + +Public surface +~~~~~~~~~~~~~~ + + from processing import process + + file_rows, content_rows = process("/path/to/scan") + + # file_rows → list of dicts matching the `File` SQLAlchemy model + # content_rows → list of dicts matching the `FileContent` model + +Lower-level modules +~~~~~~~~~~~~~~~~~~~ +* ``traversal`` – recursive file walker +* ``extractors`` – per-format text + embedded-OCR extractors +* ``ocr`` – EasyOCR singleton wrapper +* ``embeddings`` – batched sentence-transformer inference +* ``pipeline`` – top-level orchestrator (re-exported as ``process``) +""" + +from .pipeline import process + +__all__ = ["process"] diff --git a/app/processing/embeddings.py b/app/processing/embeddings.py new file mode 100644 index 0000000..cd9d027 --- /dev/null +++ b/app/processing/embeddings.py @@ -0,0 +1,109 @@ +""" +embeddings.py +------------- +Batched embedding pipeline for WhereTF. + +Takes a flat list of ``ChunkData`` dicts (produced by the extractors) and +mutates each one in-place by adding an ``"embedding"`` key whose value is a +plain Python ``list[float]`` of 384 dimensions — ready to be passed straight +into the ``FileContent`` SQLAlchemy model. + +The model is loaded once and cached for the lifetime of the process. +""" + +from __future__ import annotations + +import logging +from functools import lru_cache +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .extractors import ChunkData + +logger = logging.getLogger(__name__) + +MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2" +DEFAULT_BATCH_SIZE = 64 + + +# --------------------------------------------------------------------------- +# Lazy model singleton +# --------------------------------------------------------------------------- + +@lru_cache(maxsize=1) +def _get_model(): + """Load and cache the SentenceTransformer model (once per process).""" + try: + from sentence_transformers import SentenceTransformer # type: ignore + except ImportError as exc: + raise RuntimeError( + "sentence-transformers is required. " + "Install it with: pip install sentence-transformers" + ) from exc + + logger.info("[embeddings] Loading model '%s' …", MODEL_NAME) + model = SentenceTransformer(MODEL_NAME) + logger.info("[embeddings] Model loaded.") + return model + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def embed_chunks( + chunks: list["ChunkData"], + batch_size: int = DEFAULT_BATCH_SIZE, +) -> list["ChunkData"]: + """ + Add an ``"embedding"`` field to every chunk dict **in-place** and + return the same list for convenient chaining. + + Parameters + ---------- + chunks: + Output of :func:`~processing.extractors.extract` — one dict per + text chunk. Dicts that already have an ``"embedding"`` key are + skipped so this function is safe to call multiple times. + batch_size: + Number of texts to encode per forward pass. 64 is a good default + for CPU; increase to 256+ when a GPU is available. + + Returns + ------- + list[ChunkData] + The same list, each dict now containing: + ``{"chunk_index": int, "content_text": str, "embedding": list[float]}`` + """ + if not chunks: + return chunks + + # Separate chunks that still need embeddings + pending_indices = [i for i, c in enumerate(chunks) if "embedding" not in c] + if not pending_indices: + return chunks + + model = _get_model() + texts = [chunks[i]["content_text"] for i in pending_indices] + + logger.info( + "[embeddings] Encoding %d chunk(s) in batches of %d …", + len(texts), + batch_size, + ) + + # encode() returns a numpy array of shape (N, 384) + vectors = model.encode( + texts, + batch_size=batch_size, + show_progress_bar=len(texts) > batch_size, + convert_to_numpy=True, + normalize_embeddings=True, # unit-normalised → cosine ≡ dot product + ) + + for list_pos, chunk_idx in enumerate(pending_indices): + # Store as plain Python list[float] — JSON-serialisable, pgvector-ready + chunks[chunk_idx]["embedding"] = vectors[list_pos].tolist() + + logger.info("[embeddings] Done.") + return chunks diff --git a/app/processing/extractors.py b/app/processing/extractors.py new file mode 100644 index 0000000..67cda57 --- /dev/null +++ b/app/processing/extractors.py @@ -0,0 +1,594 @@ +""" +extractors.py +------------- +One extractor class / function per supported file type. + +Every extractor returns a list of ``ChunkData`` dicts: + + { + "chunk_index" : int, # 0-based position within the file + "content_text": str, # full text for this chunk (page / slide / block) + # "embedding" and "keyword_tokens" are added downstream + } + +Embedded-image OCR (PDF, DOCX, PPTX) appends OCR text to the same chunk +so that the text and its visual content share the same vector space. + +Rolling/overlapping chunking is used for all text-heavy formats so that +context is never silently cut at a hard boundary. + +No database sessions, no embedding calls, no side effects. +""" + +from __future__ import annotations + +import io +import json +import logging +import zipfile +from pathlib import Path +from typing import TypedDict + +from .ocr import ocr_image_bytes + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Shared type +# --------------------------------------------------------------------------- + +class ChunkData(TypedDict): + chunk_index: int + content_text: str + # embedding and keyword_tokens intentionally absent here; + # they are injected by the embedding pipeline. + + +# --------------------------------------------------------------------------- +# Chunking configuration +# --------------------------------------------------------------------------- + +# Rolling-window parameters (characters). +# CHUNK_SIZE : maximum characters per chunk. +# CHUNK_OVERLAP: how many characters from the end of chunk N are repeated +# at the start of chunk N+1, preserving cross-boundary context. +CHUNK_SIZE = 1_500 +CHUNK_OVERLAP = 200 + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + +def _clean(text: str) -> str: + """Collapse excessive whitespace without destroying newlines.""" + import re + text = re.sub(r"[ \t]+", " ", text) + text = re.sub(r"\n{3,}", "\n\n", text) + return text.strip() + + +def _rolling_chunks(text: str, size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> list[str]: + """ + Split *text* into overlapping sliding windows. + + Each window is at most *size* characters. The hop between windows is + exactly ``size - overlap`` characters, so every boundary region appears + in two successive chunks and no context is silently lost. + + The soft-break on the last newline inside each window only fires when + it falls within the hop zone ``[start + step, end]``, so it can never + shrink the step and cause degenerate micro-chunks. + + Returns a list of non-empty stripped strings. + """ + if not text: + return [] + + step = size - overlap # guaranteed minimum forward progress + chunks: list[str] = [] + start = 0 + length = len(text) + + while start < length: + end = min(start + size, length) + + # Soft-break: snap to last newline in the tail of the window, + # but only if that keeps the step at least `step` chars long. + if end < length: + nl = text.rfind("\n", start + step, end) + if nl != -1: + end = nl + 1 # include the newline in this chunk + + chunk = text[start:end].strip() + if chunk: + chunks.append(chunk) + + start += step # fixed-size hop, always progresses + + return chunks + + +def _zip_image_bytes(zf: zipfile.ZipFile, name: str) -> bytes: + """Read a member from an open ZipFile; return empty bytes on error.""" + try: + return zf.read(name) + except Exception: # noqa: BLE001 + return b"" + + +def _pdf_image_to_png_bytes(img_meta: dict) -> bytes: + """ + Convert a pdfplumber image-metadata dict to a valid PNG byte-string + that PIL can open. + + pdfplumber stores raw decoded pixel data (not a container format) in + img_meta["stream"].get_data(). We reconstruct a proper PNG from the + dimensions and colour-space information that pdfplumber also exposes. + + Supported colour spaces: DeviceRGB (3-channel), DeviceGray (1-channel), + and DeviceCMYK (4-channel, converted to RGB). + Falls back to trying the raw bytes directly (sometimes the stream IS + a JPEG or PNG already). + """ + try: + from PIL import Image # type: ignore + + stream = img_meta.get("stream") + if stream is None: + return b"" + + raw = stream.get_data() + + # --- fast path: bytes are already a valid image container ---------- + try: + Image.open(io.BytesIO(raw)).verify() + return raw # it worked — JPEG / PNG inside PDF + except Exception: + pass # fall through to pixel reconstruction + + # --- reconstruct from raw pixel buffer ------------------------------ + width = int(img_meta.get("width", img_meta.get("Width", 0))) + height = int(img_meta.get("height", img_meta.get("Height", 0))) + cs = str(img_meta.get("colorspace", img_meta.get("ColorSpace", "RGB"))) + + if width == 0 or height == 0: + return b"" + + if "gray" in cs.lower() or "grey" in cs.lower(): + mode, channels = "L", 1 + elif "cmyk" in cs.lower(): + mode, channels = "CMYK", 4 + else: + mode, channels = "RGB", 3 + + expected = width * height * channels + if len(raw) < expected: + return b"" + + pil_img = Image.frombytes(mode, (width, height), raw[:expected]) + + # CMYK → RGB so EasyOCR never sees CMYK + if mode == "CMYK": + pil_img = pil_img.convert("RGB") + + buf = io.BytesIO() + pil_img.save(buf, format="PNG") + return buf.getvalue() + + except Exception as exc: # noqa: BLE001 + logger.debug("[pdf] Could not convert image to PNG: %s", exc) + return b"" + + +# --------------------------------------------------------------------------- +# PDF extractor +# --------------------------------------------------------------------------- + +def extract_pdf(path: Path) -> list[ChunkData]: + """ + One page = one text block. The text blocks are then fed through the + rolling-window chunker so long pages produce overlapping chunks. + + Each page block also includes OCR output from any embedded raster + images on that page. + """ + try: + import pdfplumber # type: ignore + except ImportError as exc: + raise RuntimeError("pip install pdfplumber") from exc + + page_blocks: list[str] = [] + + with pdfplumber.open(str(path)) as pdf: + for page_idx, page in enumerate(pdf.pages): + parts: list[str] = [] + + # 1. Native text layer + native = page.extract_text() or "" + if native.strip(): + parts.append(native) + + # 2. Embedded images → PNG reconstruction → OCR + for img_meta in page.images: + try: + png_bytes = _pdf_image_to_png_bytes(img_meta) + if not png_bytes: + continue + ocr_text = ocr_image_bytes(png_bytes) + if ocr_text: + parts.append(f"[OCR] {ocr_text}") + except Exception as exc: # noqa: BLE001 + logger.debug("[pdf] OCR failed on page %d image: %s", page_idx, exc) + + block = _clean("\n".join(parts)) + if block: + page_blocks.append(block) + + # Join all page blocks then apply rolling chunker so cross-page context + # is preserved at page boundaries too. + full_text = "\n\n".join(page_blocks) + raw_chunks = _rolling_chunks(full_text) + + return [ + ChunkData(chunk_index=i, content_text=c) + for i, c in enumerate(raw_chunks) + ] + + +# --------------------------------------------------------------------------- +# DOCX extractor +# --------------------------------------------------------------------------- + +def extract_docx(path: Path) -> list[ChunkData]: + """ + All paragraphs are joined into a single document string, then split + via the rolling-window chunker. Embedded images are OCR-ed and + inserted at the end of the text before chunking. + """ + try: + import docx as _docx # type: ignore (python-docx) + except ImportError as exc: + raise RuntimeError("pip install python-docx") from exc + + doc = _docx.Document(str(path)) + + # ---- collect image bytes from word/media/ ---- + with zipfile.ZipFile(str(path), "r") as zf: + media_bytes: list[bytes] = [ + _zip_image_bytes(zf, name) + for name in zf.namelist() + if name.startswith("word/media/") + ] + + # ---- build full document text ---- + lines: list[str] = [] + for para in doc.paragraphs: + text = para.text.strip() + if text: + lines.append(text) + + # Append OCR from all embedded images + for img_bytes in media_bytes: + if not img_bytes: + continue + ocr_text = ocr_image_bytes(img_bytes) + if ocr_text: + lines.append(f"[OCR] {ocr_text}") + + full_text = _clean("\n".join(lines)) + raw_chunks = _rolling_chunks(full_text) + + return [ + ChunkData(chunk_index=i, content_text=c) + for i, c in enumerate(raw_chunks) + ] + + +# --------------------------------------------------------------------------- +# PPTX extractor +# --------------------------------------------------------------------------- + +def extract_pptx(path: Path) -> list[ChunkData]: + """ + One chunk per slide (slides are natural semantic units in presentations). + Slide text + speaker notes + OCR of every embedded image on that slide. + Long slides are further split by the rolling chunker. + """ + try: + from pptx import Presentation # type: ignore + except ImportError as exc: + raise RuntimeError("pip install python-pptx") from exc + + prs = Presentation(str(path)) + chunks: list[ChunkData] = [] + global_chunk_idx = 0 + + for slide in prs.slides: + parts: list[str] = [] + + # 1. Text from all shapes + for shape in slide.shapes: + if shape.has_text_frame: + for para in shape.text_frame.paragraphs: + line = " ".join(run.text for run in para.runs).strip() + if line: + parts.append(line) + + # 2. Speaker notes + if slide.has_notes_slide: + notes_tf = slide.notes_slide.notes_text_frame + notes_text = notes_tf.text.strip() if notes_tf else "" + if notes_text: + parts.append(f"[Notes] {notes_text}") + + # 3. Embedded images via shape relationships (most reliable path) + seen_blobs: set[int] = set() + for rel in slide.part.rels.values(): + try: + if "image" not in rel.reltype: + continue + blob = rel.target_part.blob + blob_id = id(blob) + if blob_id in seen_blobs: + continue + seen_blobs.add(blob_id) + ocr_text = ocr_image_bytes(blob) + if ocr_text: + parts.append(f"[OCR] {ocr_text}") + except Exception: # noqa: BLE001 + pass + + slide_text = _clean("\n".join(parts)) + if not slide_text: + continue + + for sub_chunk in _rolling_chunks(slide_text): + chunks.append(ChunkData(chunk_index=global_chunk_idx, content_text=sub_chunk)) + global_chunk_idx += 1 + + return chunks + + +# --------------------------------------------------------------------------- +# XLSX extractor +# --------------------------------------------------------------------------- + +def extract_xlsx(path: Path) -> list[ChunkData]: + """ + Convert each sheet to a readable text table (tab-separated). + Each sheet becomes one rolling-chunked text block so large sheets + produce overlapping chunks. + """ + try: + import openpyxl # type: ignore + except ImportError as exc: + raise RuntimeError("pip install openpyxl") from exc + + wb = openpyxl.load_workbook(str(path), read_only=True, data_only=True) + chunks: list[ChunkData] = [] + global_idx = 0 + + for sheet_name in wb.sheetnames: + ws = wb[sheet_name] + rows: list[str] = [] + + for row in ws.iter_rows(values_only=True): + # Skip completely empty rows + if all(cell is None for cell in row): + continue + cells = [str(cell) if cell is not None else "" for cell in row] + rows.append("\t".join(cells)) + + if not rows: + continue + + # Prepend sheet name as context header + sheet_text = _clean(f"[Sheet: {sheet_name}]\n" + "\n".join(rows)) + + for sub_chunk in _rolling_chunks(sheet_text): + chunks.append(ChunkData(chunk_index=global_idx, content_text=sub_chunk)) + global_idx += 1 + + wb.close() + return chunks + + +# --------------------------------------------------------------------------- +# Standalone image extractor (OCR only) +# --------------------------------------------------------------------------- + +def extract_image(path: Path) -> list[ChunkData]: + """Single chunk: full OCR output of the image file.""" + img_bytes = path.read_bytes() + ocr_text = ocr_image_bytes(img_bytes) + if not ocr_text: + return [] + # Images rarely exceed one chunk but run through rolling chunker for safety + return [ + ChunkData(chunk_index=i, content_text=c) + for i, c in enumerate(_rolling_chunks(_clean(ocr_text))) + ] + + +# --------------------------------------------------------------------------- +# Plain-text / code extractor +# --------------------------------------------------------------------------- + +def extract_text(path: Path) -> list[ChunkData]: + """ + Read a UTF-8 text file and split it into rolling overlapping chunks. + JSON files are pretty-printed before chunking. + CSV files get a column-header prefix on every chunk for LLM context. + """ + try: + raw = path.read_text(encoding="utf-8", errors="replace") + except OSError as exc: + logger.warning("[text] Cannot read %s: %s", path, exc) + return [] + + ext = path.suffix.lower() + + # --- JSON: pretty-print for readability --------------------------------- + if ext == ".json": + try: + raw = json.dumps(json.loads(raw), indent=2, ensure_ascii=False) + except json.JSONDecodeError: + pass + + # --- CSV: prepend header to every chunk for grounding ------------------- + if ext == ".csv": + lines = raw.splitlines() + header = lines[0] if lines else "" + body = "\n".join(lines[1:]) if len(lines) > 1 else "" + sub_chunks = _rolling_chunks(_clean(body)) + return [ + ChunkData(chunk_index=i, content_text=_clean(f"[Columns: {header}]\n{c}")) + for i, c in enumerate(sub_chunks) + ] if sub_chunks else [] + + raw_chunks = _rolling_chunks(_clean(raw)) + return [ + ChunkData(chunk_index=i, content_text=c) + for i, c in enumerate(raw_chunks) + ] + + +# --------------------------------------------------------------------------- +# Dispatch table +# --------------------------------------------------------------------------- + +#: Map every supported extension to its extractor function. +#: Add new extensions here; no other file needs changing. +_EXTENSION_MAP: dict[str, object] = { + # ── Documents ───────────────────────────────────────────────────────── + ".pdf": extract_pdf, + ".docx": extract_docx, + ".pptx": extract_pptx, + ".xlsx": extract_xlsx, + ".xls": extract_xlsx, # openpyxl handles legacy xls via compatibility + # ── Images (standalone OCR) ─────────────────────────────────────────── + ".png": extract_image, + ".jpg": extract_image, + ".jpeg": extract_image, + ".webp": extract_image, + ".bmp": extract_image, + ".tiff": extract_image, + ".tif": extract_image, + # ── Plain text & markup ─────────────────────────────────────────────── + ".txt": extract_text, + ".md": extract_text, + ".rst": extract_text, + ".tex": extract_text, + # ── Data formats ───────────────────────────────────────────────────── + ".json": extract_text, + ".csv": extract_text, + ".tsv": extract_text, + ".xml": extract_text, + ".yaml": extract_text, + ".yml": extract_text, + ".toml": extract_text, + # ── Python ─────────────────────────────────────────────────────────── + ".py": extract_text, + ".pyi": extract_text, + ".ipynb": extract_text, + # ── JavaScript / TypeScript ─────────────────────────────────────────── + ".js": extract_text, + ".jsx": extract_text, + ".ts": extract_text, + ".tsx": extract_text, + ".mjs": extract_text, + ".cjs": extract_text, + # ── Web ─────────────────────────────────────────────────────────────── + ".html": extract_text, + ".htm": extract_text, + ".css": extract_text, + ".scss": extract_text, + ".sass": extract_text, + # ── Systems / compiled languages ───────────────────────────────────── + ".c": extract_text, + ".h": extract_text, + ".cpp": extract_text, + ".cc": extract_text, + ".cxx": extract_text, + ".hpp": extract_text, + ".hxx": extract_text, + ".cs": extract_text, # C# + ".java": extract_text, + ".kt": extract_text, # Kotlin + ".swift": extract_text, + ".go": extract_text, + ".rs": extract_text, # Rust + ".zig": extract_text, + # ── Scripting / shell ───────────────────────────────────────────────── + ".sh": extract_text, + ".bash": extract_text, + ".zsh": extract_text, + ".fish": extract_text, + ".ps1": extract_text, # PowerShell + ".bat": extract_text, + ".cmd": extract_text, + # ── Ruby / PHP / others ─────────────────────────────────────────────── + ".rb": extract_text, + ".php": extract_text, + ".lua": extract_text, + ".pl": extract_text, # Perl + ".r": extract_text, # R + ".scala": extract_text, + ".ex": extract_text, # Elixir + ".exs": extract_text, + ".erl": extract_text, # Erlang + ".hs": extract_text, # Haskell + ".ml": extract_text, # OCaml + ".clj": extract_text, # Clojure + # ── Config / infra ──────────────────────────────────────────────────── + ".ini": extract_text, + ".cfg": extract_text, + ".conf": extract_text, + ".env": extract_text, + ".dockerfile": extract_text, + ".tf": extract_text, # Terraform + ".hcl": extract_text, + ".sql": extract_text, + ".graphql": extract_text, + ".proto": extract_text, # Protobuf +} + + +def extract(path: Path) -> list[ChunkData]: + """ + Dispatch to the correct extractor based on file extension. + + Parameters + ---------- + path: + Resolved :class:`~pathlib.Path` to the file. + + Returns + ------- + list[ChunkData] + Zero or more chunk dicts ready for the embedding pipeline. + """ + # Handle extensionless files often found in codebases (Makefile, Dockerfile…) + ext = path.suffix.lower() or path.name.lower() + extractor = _EXTENSION_MAP.get(ext) + + # Special-case common extensionless names + if extractor is None and path.name.lower() in { + "makefile", "dockerfile", "jenkinsfile", "vagrantfile", + "gemfile", "rakefile", "procfile", "brewfile", + ".gitignore", ".gitattributes", ".editorconfig", + "requirements", "pipfile", "cargo.lock", "go.sum", + }: + extractor = extract_text + + if extractor is None: + logger.debug("[extractor] No extractor for '%s' (%s)", ext, path.name) + return [] + + try: + return extractor(path) + except Exception as exc: # noqa: BLE001 + logger.error("[extractor] Failed to extract %s: %s", path, exc) + return [] \ No newline at end of file diff --git a/app/processing/ocr.py b/app/processing/ocr.py new file mode 100644 index 0000000..58c0d9f --- /dev/null +++ b/app/processing/ocr.py @@ -0,0 +1,99 @@ +""" +ocr.py +------ +Thin, lazy-initialised wrapper around EasyOCR. + +Design goals +~~~~~~~~~~~~ +* One global reader instance (EasyOCR is expensive to initialise). +* Accepts raw ``bytes`` so callers never need to touch the filesystem. +* Returns a single cleaned string; empty string on failure. +* GPU is used when available; falls back to CPU silently. +""" + +from __future__ import annotations + +import io +import logging +from functools import lru_cache +from typing import Sequence + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Lazy singleton +# --------------------------------------------------------------------------- + +@lru_cache(maxsize=1) +def _get_reader(languages: tuple[str, ...] = ("en",)): + """ + Initialise and cache a single EasyOCR Reader. + + The tuple argument is required so that lru_cache can hash it. + Import is deferred so the module loads fast even when EasyOCR is + installed but not yet needed. + """ + try: + import easyocr # type: ignore + logger.info("[ocr] Initialising EasyOCR (languages=%s) …", languages) + return easyocr.Reader(list(languages), gpu=_gpu_available()) + except ImportError as exc: + raise RuntimeError( + "EasyOCR is required for embedded image OCR. " + "Install it with: pip install easyocr" + ) from exc + + +def _gpu_available() -> bool: + try: + import torch # type: ignore + return torch.cuda.is_available() + except ImportError: + return False + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def ocr_image_bytes( + image_bytes: bytes, + languages: Sequence[str] = ("en",), +) -> str: + """ + Run EasyOCR on *image_bytes* and return the extracted text. + + Parameters + ---------- + image_bytes: + Raw bytes of any image format supported by EasyOCR + (PNG, JPEG, WEBP, BMP, TIFF, …). + languages: + ISO-639-1 language codes to pass to EasyOCR. + + Returns + ------- + str + Whitespace-joined OCR result, or ``""`` if nothing was detected + or an error occurred. + """ + if not image_bytes: + return "" + + try: + import numpy as np # type: ignore + from PIL import Image # type: ignore + + reader = _get_reader(tuple(languages)) + + # Convert bytes → PIL → numpy (EasyOCR accepts numpy arrays natively) + pil_img = Image.open(io.BytesIO(image_bytes)).convert("RGB") + np_img = np.array(pil_img) + + results = reader.readtext(np_img, detail=0, paragraph=True) + return " ".join(results).strip() + + except Exception as exc: # noqa: BLE001 + logger.warning("[ocr] Failed to OCR image chunk (%d bytes): %s", len(image_bytes), exc) + return "" diff --git a/app/processing/pipeline.py b/app/processing/pipeline.py new file mode 100644 index 0000000..5f9e383 --- /dev/null +++ b/app/processing/pipeline.py @@ -0,0 +1,161 @@ +""" +pipeline.py +----------- +Top-level orchestrator for the WhereTF processing system. + +Call :func:`process` with a folder or file path. It returns two lists +of plain dicts — one per SQLAlchemy model — that your backend can insert +directly without any parsing or transformation: + + file_rows, content_rows = process("/home/user/Documents") + + # Save them: + for fr in file_rows: + db.add(File(**fr)) + + for cr in content_rows: + db.add(FileContent(**cr)) + + db.commit() + +``FileContent`` rows reference their parent ``File`` by ``file_path`` +(the unique natural key) so you can look up the FK after inserting +``File`` rows. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any + +from .embeddings import embed_chunks +from .extractors import extract +from .traversal import FileEntry, iter_files + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Output types (plain dicts — no SQLAlchemy imports here) +# --------------------------------------------------------------------------- + +# Matches File model fields exactly. +FileRow = dict[str, Any] +""" +{ + "file_path" : str, + "file_hash" : str, + "mime_type" : str, + "last_modified" : datetime, + "tags" : list[str], +} +""" + +# Matches FileContent model fields exactly. +ContentRow = dict[str, Any] +""" +{ + "file_path" : str, # FK reference — not a SQLAlchemy column, + # but lets the backend do File lookup + "chunk_index" : int, + "content_text": str, + "embedding" : list[float], # 384 dims, unit-normalised + # "keyword_tokens" is a server-side generated column (to_tsvector), + # so we deliberately omit it; PostgreSQL fills it automatically. +} +""" + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def process( + root: str | Path, + embedding_batch_size: int = 64, + skip_hashes: set[str] | None = None, +) -> tuple[list[FileRow], list[ContentRow]]: + """ + Walk *root*, extract text (+ embedded-image OCR), embed everything, + and return two lists of insertion-ready dicts. + + Parameters + ---------- + root: + Path to a directory or a single file. + embedding_batch_size: + Passed through to :func:`~processing.embeddings.embed_chunks`. + Increase for GPU machines. + skip_hashes: + Optional set of SHA-256 hex digests for files already in the DB. + Matching files are skipped entirely (incremental indexing). + + Returns + ------- + file_rows : list[FileRow] + One dict per discovered file. + content_rows : list[ContentRow] + One dict per text chunk across all files. + Each dict carries ``"file_path"`` so the backend can resolve the FK. + """ + skip_hashes = skip_hashes or set() + + file_rows: list[FileRow] = [] + all_chunks: list[dict[str, Any]] = [] # ChunkData + "file_path" tag + + # ------------------------------------------------------------------------- + # Phase 1 – Traverse & Extract + # ------------------------------------------------------------------------- + for entry in iter_files(root): + file_path: str = entry["file_path"] + file_hash: str = entry["file_hash"] + + if file_hash in skip_hashes: + logger.debug("[pipeline] Skipping already-indexed file: %s", file_path) + continue + + logger.info("[pipeline] Processing: %s", file_path) + file_rows.append(dict(entry)) # shallow copy; tags list is mutable + + chunks = extract(Path(file_path)) + if not chunks: + logger.debug("[pipeline] No content extracted from: %s", file_path) + continue + + # Tag each chunk with the owning file path so we can split later + for chunk in chunks: + chunk["_file_path"] = file_path # temporary internal field + + all_chunks.extend(chunks) + + if not all_chunks: + logger.info("[pipeline] No chunks to embed.") + return file_rows, [] + + # ------------------------------------------------------------------------- + # Phase 2 – Batch Embed (single pass over all chunks for efficiency) + # ------------------------------------------------------------------------- + embed_chunks(all_chunks, batch_size=embedding_batch_size) + + # ------------------------------------------------------------------------- + # Phase 3 – Build ContentRow dicts + # ------------------------------------------------------------------------- + content_rows: list[ContentRow] = [] + for chunk in all_chunks: + content_rows.append( + { + "file_path": chunk.pop("_file_path"), # remove temp field + "chunk_index": chunk["chunk_index"], + "content_text": chunk["content_text"], + "embedding": chunk["embedding"], + # keyword_tokens: omitted — generated by PostgreSQL + } + ) + + logger.info( + "[pipeline] Finished. files=%d chunks=%d", + len(file_rows), + len(content_rows), + ) + return file_rows, content_rows diff --git a/app/processing/traversal.py b/app/processing/traversal.py new file mode 100644 index 0000000..f15e0d6 --- /dev/null +++ b/app/processing/traversal.py @@ -0,0 +1,150 @@ +""" +traversal.py +------------ +Recursive directory / single-file walker for WhereTF. + +Yields FileEntry dicts that carry the raw metadata needed to populate +the `File` SQLAlchemy model. No DB sessions are opened here. +""" + +from __future__ import annotations + +import hashlib +import mimetypes +from datetime import datetime +from pathlib import Path +from typing import Generator + +# --------------------------------------------------------------------------- +# Types +# --------------------------------------------------------------------------- + +#: Every extension the processing pipeline can handle. +#: Keep in sync with the dispatch table in extractors.py. +SUPPORTED_EXTENSIONS: frozenset[str] = frozenset( + { + # ── Documents ───────────────────────────────────────────────────── + ".pdf", ".docx", ".pptx", ".xlsx", ".xls", + # ── Images (standalone OCR) ─────────────────────────────────────── + ".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff", ".tif", + # ── Plain text & markup ─────────────────────────────────────────── + ".txt", ".md", ".rst", ".tex", + # ── Data formats ───────────────────────────────────────────────── + ".json", ".csv", ".tsv", ".xml", ".yaml", ".yml", ".toml", + # ── Python ─────────────────────────────────────────────────────── + ".py", ".pyi", ".ipynb", + # ── JavaScript / TypeScript ─────────────────────────────────────── + ".js", ".jsx", ".ts", ".tsx", ".mjs", ".cjs", + # ── Web ─────────────────────────────────────────────────────────── + ".html", ".htm", ".css", ".scss", ".sass", + # ── Systems / compiled languages ───────────────────────────────── + ".c", ".h", ".cpp", ".cc", ".cxx", ".hpp", ".hxx", + ".cs", ".java", ".kt", ".swift", ".go", ".rs", ".zig", + # ── Scripting / shell ───────────────────────────────────────────── + ".sh", ".bash", ".zsh", ".fish", ".ps1", ".bat", ".cmd", + # ── Ruby / PHP / others ─────────────────────────────────────────── + ".rb", ".php", ".lua", ".pl", ".r", ".scala", + ".ex", ".exs", ".erl", ".hs", ".ml", ".clj", + # ── Config / infra ──────────────────────────────────────────────── + ".ini", ".cfg", ".conf", ".env", + ".dockerfile", ".tf", ".hcl", ".sql", ".graphql", ".proto", + } +) + +#: Extensionless filenames that should still be processed as plain text. +SUPPORTED_NAMES: frozenset[str] = frozenset( + { + "makefile", "dockerfile", "jenkinsfile", "vagrantfile", + "gemfile", "rakefile", "procfile", "brewfile", + ".gitignore", ".gitattributes", ".editorconfig", + "requirements", "pipfile", "cargo.lock", "go.sum", + } +) + + +class FileEntry(dict): + """ + A plain dict subclass that remains fully JSON-serialisable. + + Keys match the ``File`` SQLAlchemy model fields: + file_path : str – absolute, POSIX-style path + file_hash : str – SHA-256 hex digest + mime_type : str – best-guess MIME type + last_modified : datetime + tags : list[str] – starts empty; callers may populate later + """ + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _sha256(path: Path, chunk_size: int = 1 << 20) -> str: + """Stream-hash a file in 1 MiB chunks — never loads the whole file.""" + h = hashlib.sha256() + with path.open("rb") as fh: + while True: + buf = fh.read(chunk_size) + if not buf: + break + h.update(buf) + return h.hexdigest() + + +def _mime(path: Path) -> str: + mime, _ = mimetypes.guess_type(str(path)) + return mime or "application/octet-stream" + + +def _is_supported(path: Path) -> bool: + """Return True if this path should be processed.""" + return ( + path.suffix.lower() in SUPPORTED_EXTENSIONS + or path.name.lower() in SUPPORTED_NAMES + ) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def iter_files(root: str | Path) -> Generator[FileEntry, None, None]: + """ + Yield one :class:`FileEntry` for every supported file found under + *root*. If *root* is a single file, yield just that file (provided + its extension or name is supported). + + Parameters + ---------- + root: + A directory path or a single file path. + + Yields + ------ + FileEntry + Populated with ``file_path``, ``file_hash``, ``mime_type``, + ``last_modified``, and an empty ``tags`` list. + """ + root = Path(root).resolve() + + if root.is_file(): + paths: list[Path] | Generator = [root] + elif root.is_dir(): + paths = (p for p in root.rglob("*") if p.is_file()) + else: + raise FileNotFoundError(f"Path does not exist or is not accessible: {root}") + + for path in paths: + if not _is_supported(path): + continue + try: + stat = path.stat() + yield FileEntry( + file_path=path.as_posix(), + file_hash=_sha256(path), + mime_type=_mime(path), + last_modified=datetime.fromtimestamp(stat.st_mtime), + tags=[], + ) + except (OSError, PermissionError) as exc: + print(f"[traversal] Skipping {path}: {exc}") \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index c0d5e2eb153765efd1147810345a59255a75fa77..1e2d2c9c3590850b66505fd6f146fb671a834eba 100644 GIT binary patch delta 501 zcmaJ-OAdli5NmW{qH*ok3t$BNtn>_CK@p;f0s)P>;t;yv7#`8~i4cXDmp7UBW;$&r z?>u_XyzR!=B0!7-&NyNN9|@LNV2lxL#$o;3qt+nva5a6l>1*ghA6?G9;e;I`gs}LE z`Fi6tf5z*;{b3N!#~LZI5xhc=c;;1SR;6Vo;0k+MIS_a$