rag.py

"""ARIA RAG engine.

Loads markdown knowledge files, chunks them by `##` heading (then by
~300-token windows with overlap), embeds with all-MiniLM-L6-v2, and serves
cosine-similarity search via a FAISS inner-product index over normalized
vectors. The index is built once at startup — the knowledge base is small,
so this takes seconds.
"""
from __future__ import annotations

import os
import re
from dataclasses import dataclass

KNOWLEDGE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "knowledge")
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
CHUNK_WORDS = 220   # ~300 tokens
OVERLAP_WORDS = 40  # ~50 tokens
TOP_K = 4


@dataclass
class Chunk:
    source: str    # filename, e.g. "genie.md"
    heading: str   # nearest ## heading ("" for file preamble)
    text: str      # chunk text with heading prepended


def _split_sections(markdown: str) -> list[tuple[str, str]]:
    """Split a markdown document into (heading, body) pairs on `##` headings."""
    parts = re.split(r"(?m)^##\s+", markdown)
    sections: list[tuple[str, str]] = []

    preamble = parts[0].strip()
    if preamble:
        # Strip the top-level `# Title` line but keep any intro text.
        preamble_body = re.sub(r"(?m)^#\s+.*$", "", preamble).strip()
        title = ""
        m = re.search(r"(?m)^#\s+(.*)$", preamble)
        if m:
            title = m.group(1).strip()
        if preamble_body:
            sections.append((title, preamble_body))

    for part in parts[1:]:
        lines = part.splitlines()
        heading = lines[0].strip() if lines else ""
        body = "\n".join(lines[1:]).strip()
        if body:
            sections.append((heading, body))
    return sections


def _windows(words: list[str], size: int, overlap: int):
    """Yield overlapping word windows of `size` with `overlap` words shared."""
    step = max(size - overlap, 1)
    start = 0
    while True:
        yield words[start:start + size]
        if start + size >= len(words):
            break
        start += step


def load_chunks() -> list[Chunk]:
    """Read every .md file in knowledge/ and return its chunks."""
    chunks: list[Chunk] = []
    for name in sorted(os.listdir(KNOWLEDGE_DIR)):
        if not name.endswith(".md"):
            continue
        with open(os.path.join(KNOWLEDGE_DIR, name), encoding="utf-8") as f:
            text = f.read()
        for heading, body in _split_sections(text):
            words = body.split()
            for piece in _windows(words, CHUNK_WORDS, OVERLAP_WORDS):
                if not piece:
                    continue
                chunk_text = " ".join(piece)
                if heading:
                    chunk_text = f"{heading}\n{chunk_text}"
                chunks.append(Chunk(source=name, heading=heading, text=chunk_text))
    return chunks


class RagIndex:
    """In-memory FAISS index over the knowledge chunks (cosine similarity)."""

    def __init__(self) -> None:
        # Heavy imports kept local so chunking can be unit-tested without them.
        import faiss
        import numpy as np
        from sentence_transformers import SentenceTransformer

        self._np = np
        self.chunks = load_chunks()
        if not self.chunks:
            raise RuntimeError(f"No knowledge chunks found in {KNOWLEDGE_DIR}")

        self.model = SentenceTransformer(EMBED_MODEL, device="cpu")
        vectors = self.model.encode(
            [c.text for c in self.chunks],
            normalize_embeddings=True,
            show_progress_bar=False,
        )
        vectors = np.asarray(vectors, dtype="float32")
        self.index = faiss.IndexFlatIP(vectors.shape[1])  # cosine via normalized IP
        self.index.add(vectors)

    def search(self, query: str, k: int = TOP_K) -> list[Chunk]:
        q = self.model.encode([query], normalize_embeddings=True)
        q = self._np.asarray(q, dtype="float32")
        _, idx = self.index.search(q, min(k, len(self.chunks)))
        return [self.chunks[i] for i in idx[0] if i != -1]