"""Lightweight RAG engine. Scans a directory for ``.md`` and ``.txt`` files, performs basic semantic chunking (paragraph-based with a soft character cap), and provides a simple TF-IDF search over the resulting chunks. Dependency-free so it can run in constrained environments. Can be swapped for a vector store later without changing the public API (``ingest``, ``search``, ``stats``). """ from __future__ import annotations import math import re import threading from collections import Counter from dataclasses import dataclass, field from pathlib import Path from typing import Iterable from core.vector_store import DEFAULT_INDEX_DIR, FaissVectorStore DEFAULT_HARVEST_DIR = Path("data") / "master_harvest" _SUPPORTED_SUFFIXES = {".md", ".txt"} _TOKEN_RE = re.compile(r"[A-Za-z0-9_]+") _PARAGRAPH_SPLIT_RE = re.compile(r"\n\s*\n") @dataclass class Chunk: """A single searchable chunk of text.""" source: str index: int text: str tokens: Counter = field(default_factory=Counter) def to_dict(self) -> dict: return {"source": self.source, "index": self.index, "text": self.text} def _tokenize(text: str) -> list[str]: return [t.lower() for t in _TOKEN_RE.findall(text)] def _chunk_text(text: str, max_chars: int = 1200) -> list[str]: """Split text into semantic-ish chunks. Paragraph boundaries are preferred. Paragraphs longer than ``max_chars`` are split on sentence boundaries, then on hard character boundaries. """ chunks: list[str] = [] for paragraph in _PARAGRAPH_SPLIT_RE.split(text): paragraph = paragraph.strip() if not paragraph: continue if len(paragraph) <= max_chars: chunks.append(paragraph) continue sentences = re.split(r"(?<=[.!?])\s+", paragraph) buf = "" for sentence in sentences: if len(buf) + len(sentence) + 1 > max_chars and buf: chunks.append(buf.strip()) buf = "" if len(sentence) > max_chars: for i in range(0, len(sentence), max_chars): chunks.append(sentence[i : i + max_chars].strip()) continue buf = f"{buf} {sentence}".strip() if buf: chunks.append(buf.strip()) return [c for c in chunks if c] class RAGEngine: """Simple, thread-safe RAG engine over a local directory of docs.""" def __init__( self, harvest_dir: str | Path = DEFAULT_HARVEST_DIR, index_dir: str | Path | None = None, prefer_faiss: bool = True, ) -> None: self.harvest_dir = Path(harvest_dir) self._chunks: list[Chunk] = [] self._doc_freq: Counter = Counter() self._lock = threading.Lock() self._index_dir = Path(index_dir) if index_dir else DEFAULT_INDEX_DIR self._prefer_faiss = prefer_faiss and FaissVectorStore.is_available() self._vector_store: FaissVectorStore | None = None if self._prefer_faiss: self._vector_store = FaissVectorStore(index_dir=self._index_dir) # Best-effort load of any persisted index from a prior `build`. self._vector_store.load() # -- backend introspection ------------------------------------------- @property def backend(self) -> str: if self._vector_store is not None: return "faiss" return "tfidf" def vector_store(self) -> FaissVectorStore | None: return self._vector_store # -- ingestion -------------------------------------------------------- def ingest(self) -> dict: """Re-scan the harvest directory and rebuild the in-memory index.""" with self._lock: chunks: list[Chunk] = [] doc_freq: Counter = Counter() files = self._discover_files() for path in files: try: text = path.read_text(encoding="utf-8", errors="replace") except OSError: continue try: rel = str(path.relative_to(self.harvest_dir)) except ValueError: rel = path.name for idx, chunk_text in enumerate(_chunk_text(text)): tokens = Counter(_tokenize(chunk_text)) if not tokens: continue chunks.append( Chunk(source=rel, index=idx, text=chunk_text, tokens=tokens) ) for token in tokens: doc_freq[token] += 1 self._chunks = chunks self._doc_freq = doc_freq tfidf_stats = { "files_ingested": len(files), "chunks_indexed": len(chunks), "harvest_dir": str(self.harvest_dir), "backend": self.backend, } # Build the FAISS store from the same chunks so both backends stay # consistent. Done outside the lock since FaissVectorStore has its # own internal locking and embedding is slow. if self._vector_store is not None: try: self._vector_store.build( [ {"source": c.source, "index": c.index, "text": c.text} for c in chunks ] ) except Exception as exc: # pragma: no cover - defensive tfidf_stats["faiss_error"] = str(exc) return tfidf_stats def _discover_files(self) -> list[Path]: if not self.harvest_dir.exists(): return [] return sorted( p for p in self.harvest_dir.rglob("*") if p.is_file() and p.suffix.lower() in _SUPPORTED_SUFFIXES ) # -- search ----------------------------------------------------------- def search(self, query: str, top_k: int = 5) -> list[dict]: """Return the top ``k`` chunks matching ``query``. Uses FAISS if available, otherwise the in-memory TF-IDF path. """ if top_k <= 0: return [] if self._vector_store is not None: hits = self._vector_store.search(query, top_k=top_k) if hits: return [hit.to_dict() for hit in hits] # Fall through to TF-IDF if FAISS index is empty (e.g. no `ingest` # has been run yet but TF-IDF was warmed via tests). query_tokens = _tokenize(query) if not query_tokens: return [] with self._lock: total_chunks = len(self._chunks) if total_chunks == 0: return [] scored: list[tuple[float, Chunk]] = [] for chunk in self._chunks: score = 0.0 for token in query_tokens: tf = chunk.tokens.get(token, 0) if tf == 0: continue df = self._doc_freq.get(token, 0) idf = math.log((total_chunks + 1) / (df + 1)) + 1.0 score += tf * idf if score > 0: scored.append((score, chunk)) scored.sort(key=lambda item: item[0], reverse=True) results = [] for score, chunk in scored[:top_k]: result = chunk.to_dict() result["score"] = round(score, 4) results.append(result) return results # -- introspection ---------------------------------------------------- def stats(self) -> dict: with self._lock: base = { "chunks_indexed": len(self._chunks), "unique_tokens": len(self._doc_freq), "harvest_dir": str(self.harvest_dir), "backend": self.backend, } if self._vector_store is not None: base["vector_store"] = self._vector_store.stats() return base def iter_supported_files(root: Path) -> Iterable[Path]: """Public helper for callers that want to enumerate harvest files.""" if not root.exists(): return iter(()) return ( p for p in root.rglob("*") if p.is_file() and p.suffix.lower() in _SUPPORTED_SUFFIXES )