Spaces:
Running
Running
| """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") | |
| 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 ------------------------------------------- | |
| 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 | |
| ) | |