doclex-cli 0.1.1__py3-none-any.whl

doclex/extraction/extractor.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from kreuzberg import ExtractionConfig, LayoutDetectionConfig, OutputFormat, ResultFormat, extract_file_sync
+from kreuzberg.exceptions import KreuzbergError
+
+from doclex.metadata.models import ExtractedDoc
+
+
+def extract(path: Path) -> ExtractedDoc:
+    """Extract text from a file using Kreuzberg with layout detection (RT-DETR v2 / TATR).
+
+    Raises:
+        ExtractionError: if Kreuzberg fails to extract the file.
+    """
+    config = ExtractionConfig(
+        layout=LayoutDetectionConfig(),
+        output_format=OutputFormat.MARKDOWN,
+        result_format=ResultFormat.UNIFIED,
+    )
+
+    try:
+        result = extract_file_sync(str(path), None, config)
+    except (KreuzbergError, OSError) as exc:
+        raise ExtractionError(str(path), str(exc)) from exc
+    except (KeyboardInterrupt, SystemExit):
+        raise
+    except BaseException as exc:
+        # pyo3_runtime.PanicException inherits from BaseException, not Exception.
+        # Catch Kreuzberg Rust panics (e.g. UTF-8 char boundary bug in markdown renderer)
+        # and fall back to standard extraction so the file is not lost.
+        if "char boundary" in str(exc):
+            try:
+                result = extract_file_sync(str(path), None, ExtractionConfig())
+            except (KreuzbergError, OSError) as exc2:
+                raise ExtractionError(str(path), str(exc2)) from exc2
+        else:
+            raise ExtractionError(str(path), str(exc)) from exc
+
+    tables = list(result.tables) if result.tables else []
+
+    metadata = dict(result.metadata) if result.metadata else {}
+    if "pages" in metadata and isinstance(metadata["pages"], dict):
+        metadata["pages"].pop("boundaries", None)
+
+    return ExtractedDoc(
+        path=str(path),
+        text=result.content or "",
+        mime_type=result.mime_type or _mime_from_suffix(path),
+        metadata=metadata,
+        tables=tables,
+    )
+
+
+def _mime_from_suffix(path: Path) -> str:
+    mapping = {
+        ".pdf": "application/pdf",
+        ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+        ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+        ".csv": "text/csv",
+        ".txt": "text/plain",
+        ".md": "text/markdown",
+    }
+    return mapping.get(path.suffix.lower(), "application/octet-stream")
+
+
+class ExtractionError(Exception):
+    def __init__(self, path: str, reason: str) -> None:
+        super().__init__(f"Failed to extract {path!r}: {reason}")
+        self.path = path
+        self.reason = reason

doclex/index/bm25_index.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+from bm25x import BM25
+
+from doclex.metadata.models import Segment
+
+
+_SNIPPET_WINDOW = 200  # chars around the query match for snippet generation
+
+
+class BM25Index:
+    """Wrapper around BM25X with persistence and segment-level CRUD."""
+
+    def __init__(self, index_dir: Path) -> None:
+        self._dir = index_dir
+        self._index: BM25 | None = None
+
+    def load_or_create(self) -> None:
+        """Load existing index from disk or create an empty one."""
+        if self._dir.exists() and any(self._dir.iterdir()):
+            self._index = BM25.load(str(self._dir))
+        else:
+            self._dir.mkdir(parents=True, exist_ok=True)
+            self._index = BM25()
+
+    def rebuild_from_segments(self, segments: list[Segment]) -> list[int]:
+        """Rebuild the full index from a list of segments. Returns new bm25_doc_ids."""
+        self._dir.mkdir(parents=True, exist_ok=True)
+        self._index = BM25()
+        doc_ids: list[int] = []
+        if segments:
+            doc_ids = self._index.add([s.content for s in segments])
+        self.save()
+        return doc_ids
+
+    @property
+    def _bm25(self) -> BM25:
+        if self._index is None:
+            raise RuntimeError("Index not loaded. Call load_or_create() first.")
+        return self._index
+
+    def add(self, segments: list[Segment]) -> list[int]:
+        """Add segments to the index. Returns assigned bm25_doc_ids."""
+        if not segments:
+            return []
+        doc_ids: list[int] = self._bm25.add([s.content for s in segments])
+        return doc_ids
+
+    def delete(self, bm25_doc_ids: list[int]) -> None:
+        """Remove segments by their BM25X doc ids."""
+        if bm25_doc_ids:
+            self._bm25.delete(bm25_doc_ids)
+
+    def search(self, query: str, k: int) -> list[tuple[int, float]]:
+        """Search the index. Returns list of (bm25_doc_id, score)."""
+        if len(self._bm25) == 0:
+            return []
+        results = self._bm25.search(query, k)
+        return [(doc_id, float(score)) for doc_id, score in results]
+
+    def save(self) -> None:
+        self._bm25.save(str(self._dir))
+
+    def __len__(self) -> int:
+        return len(self._bm25) if self._index is not None else 0
+
+
+def make_snippet(content: str, query: str, window: int = _SNIPPET_WINDOW) -> str:
+    """Extract a snippet centered on the window with the highest density of query-term matches.
+
+    Query terms are split on non-alphanumeric characters (matching how BM25X
+    tokenizes documents, e.g. "GT/DPM" -> "gt", "dpm") instead of whitespace only,
+    so acronyms and punctuation-joined terms can still anchor the snippet. Among
+    all match positions, the one whose surrounding window covers the most distinct
+    query terms wins (ties broken by total match count, then earliest position) —
+    not just the first occurrence of the first term — so the snippet favors the
+    part of the segment where the query actually clusters. The window is then
+    snapped to word boundaries so it never cuts a word in half.
+    """
+    terms = [t for t in re.split(r"\W+", query.lower()) if t]
+
+    matches: list[tuple[int, int, str]] = [
+        (m.start(), m.end(), term)
+        for term in terms
+        for m in re.finditer(re.escape(term), content, re.IGNORECASE)
+    ]
+
+    if not matches:
+        return content[:window].rstrip() + ("…" if len(content) > window else "")
+
+    def density(match: tuple[int, int, str]) -> tuple[int, int, int]:
+        pos = match[0]
+        lo, hi = pos - window // 2, pos + window // 2
+        nearby = [t for s, _, t in matches if lo <= s <= hi]
+        return (len(set(nearby)), len(nearby), -pos)
+
+    match_start, match_end, _ = max(matches, key=density)
+
+    start = max(0, match_start - window // 2)
+    end = min(len(content), match_end + window // 2)
+    snippet = content[start:end]
+
+    if start > 0:
+        ws_positions = [p for p in (snippet.find(" "), snippet.find("\n"), snippet.find("\t")) if p != -1]
+        if ws_positions:
+            snippet = snippet[min(ws_positions) + 1 :]
+    if end < len(content):
+        last_ws = max(snippet.rfind(" "), snippet.rfind("\n"), snippet.rfind("\t"))
+        if last_ws > 0:
+            snippet = snippet[:last_ws]
+
+    snippet = snippet.strip() or content[start:end].strip()
+    if start > 0:
+        snippet = "…" + snippet
+    if end < len(content):
+        snippet = snippet + "…"
+    return snippet

doclex/ingestion/scanner.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass
+from pathlib import Path
+
+from doclex.config import SUPPORTED_EXTENSIONS
+from doclex.metadata.db import Database
+from doclex.metadata.models import FileRecord
+
+
+@dataclass
+class ChangeSet:
+    new: list[Path]
+    modified: list[Path]
+    deleted: list[str]  # paths of files no longer on disk
+
+
+def file_hash(path: Path) -> str:
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(65536), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def file_record_from_disk(path: Path) -> FileRecord:
+    stat = path.stat()
+    return FileRecord(
+        path=str(path),
+        hash=file_hash(path),
+        mtime=stat.st_mtime,
+        size=stat.st_size,
+        mime_type=_mime_from_suffix(path),
+        status="pending",
+    )
+
+
+async def reconcile(folder: Path, db: Database) -> ChangeSet:
+    """Compare files on disk against SQLite state and return what changed."""
+    # Files on disk
+    on_disk: dict[str, Path] = {}
+    for path in folder.rglob("*"):
+        if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS:
+            # Skip the .doclex directory itself
+            if ".doclex" in path.parts:
+                continue
+            on_disk[str(path)] = path
+
+    # Files in SQLite
+    db_records: dict[str, FileRecord] = {
+        r.path: r for r in await db.get_all_files()
+    }
+
+    new: list[Path] = []
+    modified: list[Path] = []
+    deleted: list[str] = []
+
+    for path_str, path in on_disk.items():
+        if path_str not in db_records:
+            new.append(path)
+        else:
+            db_rec = db_records[path_str]
+            stat = path.stat()
+            # Use mtime as cheap pre-check before computing hash
+            if stat.st_mtime != db_rec.mtime or stat.st_size != db_rec.size:
+                current_hash = file_hash(path)
+                if current_hash != db_rec.hash:
+                    modified.append(path)
+
+    for path_str in db_records:
+        if path_str not in on_disk:
+            deleted.append(path_str)
+
+    return ChangeSet(new=new, modified=modified, deleted=deleted)
+
+
+def _mime_from_suffix(path: Path) -> str:
+    mapping = {
+        ".pdf": "application/pdf",
+        ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+        ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+        ".csv": "text/csv",
+        ".txt": "text/plain",
+        ".md": "text/markdown",
+    }
+    return mapping.get(path.suffix.lower(), "application/octet-stream")

doclex/ingestion/watcher.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import asyncio
+import threading
+from collections.abc import Callable
+from pathlib import Path
+
+from watchdog.events import FileSystemEvent, FileSystemEventHandler
+from watchdog.observers import Observer
+
+from doclex.config import SUPPORTED_EXTENSIONS
+
+
+class _Handler(FileSystemEventHandler):
+    def __init__(
+        self,
+        debounce_secs: float,
+        on_change: Callable[[Path, str], None],
+        loop: asyncio.AbstractEventLoop,
+    ) -> None:
+        super().__init__()
+        self._debounce = debounce_secs
+        self._on_change = on_change
+        self._loop = loop
+        self._timers: dict[str, threading.Timer] = {}
+        self._lock = threading.Lock()
+
+    def _is_supported(self, path: str) -> bool:
+        return Path(path).suffix.lower() in SUPPORTED_EXTENSIONS
+
+    def _schedule(self, path: str, event_type: str) -> None:
+        if not self._is_supported(path):
+            return
+        if ".doclex" in Path(path).parts:
+            return
+        with self._lock:
+            if path in self._timers:
+                self._timers[path].cancel()
+            timer = threading.Timer(
+                self._debounce,
+                self._fire,
+                args=(path, event_type),
+            )
+            self._timers[path] = timer
+            timer.start()
+
+    def _fire(self, path: str, event_type: str) -> None:
+        with self._lock:
+            self._timers.pop(path, None)
+        self._loop.call_soon_threadsafe(self._on_change, Path(path), event_type)
+
+    def on_created(self, event: FileSystemEvent) -> None:
+        if not event.is_directory:
+            self._schedule(event.src_path, "created")
+
+    def on_modified(self, event: FileSystemEvent) -> None:
+        if not event.is_directory:
+            self._schedule(event.src_path, "modified")
+
+    def on_deleted(self, event: FileSystemEvent) -> None:
+        if not event.is_directory:
+            self._schedule(event.src_path, "deleted")
+
+    def on_moved(self, event: FileSystemEvent) -> None:
+        if not event.is_directory:
+            self._schedule(event.src_path, "deleted")
+            self._schedule(event.dest_path, "created")
+
+
+class FileWatcher:
+    """Watchdog-based folder monitor with debounce."""
+
+    def __init__(
+        self,
+        folder: Path,
+        debounce_secs: float,
+        on_change: Callable[[Path, str], None],
+    ) -> None:
+        self._folder = folder
+        self._debounce = debounce_secs
+        self._on_change = on_change
+        self._observer: Observer | None = None
+
+    def start(self, loop: asyncio.AbstractEventLoop) -> None:
+        handler = _Handler(self._debounce, self._on_change, loop)
+        self._observer = Observer()
+        self._observer.schedule(handler, str(self._folder), recursive=True)
+        self._observer.start()
+
+    def stop(self) -> None:
+        if self._observer:
+            self._observer.stop()
+            self._observer.join()
+            self._observer = None

doclex/metadata/db.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import aiosqlite
+
+from .models import FileRecord, Segment
+
+if TYPE_CHECKING:
+    pass
+
+SCHEMA = """
+CREATE TABLE IF NOT EXISTS files (
+    id          INTEGER PRIMARY KEY AUTOINCREMENT,
+    path        TEXT    NOT NULL UNIQUE,
+    hash        TEXT    NOT NULL,
+    mtime       REAL    NOT NULL,
+    size        INTEGER NOT NULL,
+    mime_type   TEXT    NOT NULL,
+    status      TEXT    NOT NULL DEFAULT 'pending',
+    indexed_at  REAL,
+    error_msg   TEXT
+);
+
+CREATE TABLE IF NOT EXISTS segments (
+    id              TEXT    PRIMARY KEY,
+    file_id         INTEGER NOT NULL REFERENCES files(id) ON DELETE CASCADE,
+    content         TEXT    NOT NULL,
+    location_type   TEXT    NOT NULL,
+    location_label  TEXT    NOT NULL,
+    location_index  INTEGER NOT NULL,
+    char_count      INTEGER NOT NULL,
+    bm25_doc_id     INTEGER
+);
+
+CREATE INDEX IF NOT EXISTS idx_files_path   ON files(path);
+CREATE INDEX IF NOT EXISTS idx_files_hash   ON files(hash);
+CREATE INDEX IF NOT EXISTS idx_segs_file_id ON segments(file_id);
+"""
+
+
+class Database:
+    def __init__(self, db_path: Path) -> None:
+        self._path = db_path
+        self._conn: aiosqlite.Connection | None = None
+
+    async def connect(self) -> None:
+        self._conn = await aiosqlite.connect(self._path)
+        self._conn.row_factory = aiosqlite.Row
+        await self._conn.execute("PRAGMA journal_mode=WAL")
+        await self._conn.execute("PRAGMA foreign_keys=ON")
+        await self._conn.executescript(SCHEMA)
+        await self._conn.commit()
+
+    async def close(self) -> None:
+        if self._conn:
+            await self._conn.close()
+            self._conn = None
+
+    @property
+    def conn(self) -> aiosqlite.Connection:
+        if self._conn is None:
+            raise RuntimeError("Database not connected. Call connect() first.")
+        return self._conn
+
+    # ------------------------------------------------------------------
+    # Files
+    # ------------------------------------------------------------------
+
+    async def upsert_file(self, record: FileRecord) -> int:
+        """Insert or update a file record. Returns the file id."""
+        async with self.conn.execute(
+            """
+            INSERT INTO files (path, hash, mtime, size, mime_type, status, indexed_at, error_msg)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+            ON CONFLICT(path) DO UPDATE SET
+                hash       = excluded.hash,
+                mtime      = excluded.mtime,
+                size       = excluded.size,
+                mime_type  = excluded.mime_type,
+                status     = excluded.status,
+                indexed_at = excluded.indexed_at,
+                error_msg  = excluded.error_msg
+            """,
+            (
+                record.path,
+                record.hash,
+                record.mtime,
+                record.size,
+                record.mime_type,
+                record.status,
+                record.indexed_at,
+                record.error_msg,
+            ),
+        ) as cursor:
+            await self.conn.commit()
+            # Fetch id after upsert
+            async with self.conn.execute(
+                "SELECT id FROM files WHERE path = ?", (record.path,)
+            ) as cur2:
+                row = await cur2.fetchone()
+                return row["id"]
+
+    async def get_file_by_path(self, path: str) -> FileRecord | None:
+        async with self.conn.execute(
+            "SELECT * FROM files WHERE path = ?", (path,)
+        ) as cursor:
+            row = await cursor.fetchone()
+            return _row_to_file(row) if row else None
+
+    async def get_all_files(self) -> list[FileRecord]:
+        async with self.conn.execute("SELECT * FROM files ORDER BY path") as cursor:
+            rows = await cursor.fetchall()
+            return [_row_to_file(r) for r in rows]
+
+    async def mark_file_indexed(self, file_id: int) -> None:
+        await self.conn.execute(
+            "UPDATE files SET status = 'indexed', indexed_at = ?, error_msg = NULL WHERE id = ?",
+            (time.time(), file_id),
+        )
+        await self.conn.commit()
+
+    async def mark_file_error(self, file_id: int, error_msg: str) -> None:
+        await self.conn.execute(
+            "UPDATE files SET status = 'error', error_msg = ? WHERE id = ?",
+            (error_msg, file_id),
+        )
+        await self.conn.commit()
+
+    async def delete_file(self, path: str) -> None:
+        """Deletes file and cascades to its segments."""
+        await self.conn.execute("DELETE FROM files WHERE path = ?", (path,))
+        await self.conn.commit()
+
+    # ------------------------------------------------------------------
+    # Segments
+    # ------------------------------------------------------------------
+
+    async def insert_segments(self, file_id: int, segments: list[Segment]) -> None:
+        await self.conn.executemany(
+            """
+            INSERT OR REPLACE INTO segments
+                (id, file_id, content, location_type, location_label, location_index, char_count, bm25_doc_id)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            [
+                (
+                    s.id,
+                    file_id,
+                    s.content,
+                    s.location_type,
+                    s.location_label,
+                    s.location_index,
+                    s.char_count,
+                    s.bm25_doc_id,
+                )
+                for s in segments
+            ],
+        )
+        await self.conn.commit()
+
+    async def get_segments_for_file(self, file_id: int) -> list[Segment]:
+        async with self.conn.execute(
+            """
+            SELECT s.*, f.path as file_path FROM segments s
+            JOIN files f ON f.id = s.file_id
+            WHERE s.file_id = ?
+            ORDER BY s.location_index
+            """,
+            (file_id,),
+        ) as cursor:
+            rows = await cursor.fetchall()
+            return [_row_to_segment(r) for r in rows]
+
+    async def get_segment_by_id(self, segment_id: str) -> Segment | None:
+        async with self.conn.execute(
+            """
+            SELECT s.*, f.path as file_path FROM segments s
+            JOIN files f ON f.id = s.file_id
+            WHERE s.id = ?
+            """,
+            (segment_id,),
+        ) as cursor:
+            row = await cursor.fetchone()
+            return _row_to_segment(row) if row else None
+
+    async def get_bm25_doc_ids_for_file(self, file_id: int) -> list[int]:
+        async with self.conn.execute(
+            "SELECT bm25_doc_id FROM segments WHERE file_id = ? AND bm25_doc_id IS NOT NULL",
+            (file_id,),
+        ) as cursor:
+            rows = await cursor.fetchall()
+            return [r["bm25_doc_id"] for r in rows]
+
+    async def delete_segments_for_file(self, file_id: int) -> None:
+        await self.conn.execute("DELETE FROM segments WHERE file_id = ?", (file_id,))
+        await self.conn.commit()
+
+    async def update_bm25_doc_ids(self, segments: list, doc_ids: list[int]) -> None:
+        """Update bm25_doc_id for each segment after an index rebuild."""
+        await self.conn.executemany(
+            "UPDATE segments SET bm25_doc_id = ? WHERE id = ?",
+            [(doc_id, seg.id) for seg, doc_id in zip(segments, doc_ids)],
+        )
+        await self.conn.commit()
+
+    async def get_all_segments(self) -> list[Segment]:
+        """Used to rebuild the BM25 index on startup."""
+        async with self.conn.execute(
+            """
+            SELECT s.*, f.path as file_path FROM segments s
+            JOIN files f ON f.id = s.file_id
+            ORDER BY s.bm25_doc_id
+            """
+        ) as cursor:
+            rows = await cursor.fetchall()
+            return [_row_to_segment(r) for r in rows]
+
+
+# ------------------------------------------------------------------
+# Row helpers
+# ------------------------------------------------------------------
+
+def _row_to_file(row: aiosqlite.Row) -> FileRecord:
+    return FileRecord(
+        id=row["id"],
+        path=row["path"],
+        hash=row["hash"],
+        mtime=row["mtime"],
+        size=row["size"],
+        mime_type=row["mime_type"],
+        status=row["status"],
+        indexed_at=row["indexed_at"],
+        error_msg=row["error_msg"],
+    )
+
+
+def _row_to_segment(row: aiosqlite.Row) -> Segment:
+    return Segment(
+        id=row["id"],
+        file_path=row["file_path"],
+        content=row["content"],
+        location_type=row["location_type"],
+        location_label=row["location_label"],
+        location_index=row["location_index"],
+        char_count=row["char_count"],
+        bm25_doc_id=row["bm25_doc_id"],
+    )

doclex/metadata/models.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class FileRecord:
+    path: str
+    hash: str
+    mtime: float
+    size: int
+    mime_type: str
+    status: str  # pending | indexed | error
+    id: int | None = None
+    indexed_at: float | None = None
+    error_msg: str | None = None
+
+
+@dataclass
+class ExtractedDoc:
+    path: str
+    text: str
+    mime_type: str
+    metadata: dict = field(default_factory=dict)
+    tables: list = field(default_factory=list)      # kreuzberg extracted tables
+
+
+@dataclass
+class Segment:
+    id: str           # sha256(path + location_label + str(location_index))
+    file_path: str
+    content: str
+    location_type: str   # page | sheet | section | block | table
+    location_label: str  # "Page 3", "Sheet: Ventas", "# Introducción"
+    location_index: int
+    char_count: int = 0
+    bm25_doc_id: int | None = None  # internal BM25X index assigned on add()
+    section_id: str | None = None   # reference to tree node id
+    level: int = 0                  # heading level (1-3); 0 for non-heading nodes
+
+    def __post_init__(self) -> None:
+        if not self.char_count:
+            self.char_count = len(self.content)
+
+
+@dataclass
+class SearchResult:
+    segment_id: str
+    score: float
+    file_path: str
+    mime_type: str
+    location_type: str
+    location_label: str
+    content: str
+    snippet: str  # ~200 chars alrededor del match