rag.py

"""
clankerDiffusion — local RAG / knowledge-injection subsystem.

A small, dependency-free BM25 retriever so the agent can pull relevant
passages from a local knowledge base and inject them as <context> blocks.

Two ways knowledge gets injected:
  1. Model-driven : the model emits <tool name="retrieve">query</tool> and the
     agent runs it, appending the <result> and continuing (ReAct).
  2. Controller-driven : the agent's RAG controller watches the response as it
     is generated and, mid-turn, retrieves passages for the current question and
     injects <context>...</context> right into the stream -- "knowledge injection
     even in the middle of a response" -- without the model having to ask.

BM25 is used by default (no model downloads, runs anywhere). If
sentence-transformers is available it is used as an optional re-ranker.
"""
import os
import re
import json
import math

HERE = os.path.dirname(os.path.abspath(__file__))
DATADIR = os.path.join(HERE, "data")
KB_DIR = os.path.join(DATADIR, "kb")
DEFAULT_INDEX = os.path.join(KB_DIR, "index.json")

_TOKEN_RE = re.compile(r"[a-z0-9]+")


def _tok(s):
    return _TOKEN_RE.findall(s.lower())


def _chunk(text, size=180, stride=90):
    """Split text into token-window chunks with overlap."""
    toks = _tok(text)
    if not toks:
        return []
    out = []
    i = 0
    while i < len(toks):
        out.append(" ".join(toks[i:i + size]))
        if i + size >= len(toks):
            break
        i += stride
    return out


class KnowledgeBase:
    def __init__(self):
        self.docs = []          # list[str]  chunk texts
        self._df = {}           # term -> doc freq
        self._N = 0
        self._avgdl = 1.0
        self._built = False

    # -- ingest -------------------------------------------------------------
    def add_text(self, text, source=""):
        for ch in _chunk(text):
            if ch:
                self.docs.append(ch)
        self._built = False

    def ingest_path(self, path):
        """Ingest a file or a directory of files into the KB."""
        if os.path.isdir(path):
            files = []
            for root, _, fs in os.walk(path):
                for f in fs:
                    if f.lower().endswith((".txt", ".md", ".py", ".csv", ".json", ".log")):
                        files.append(os.path.join(root, f))
        else:
            files = [path]
        for fp in files:
            try:
                with open(fp, "r", errors="replace") as fh:
                    self.add_text(fh.read(), source=fp)
            except Exception as e:
                print(f"[rag] skip {fp}: {e}")
        print(f"[rag] ingested {len(files)} file(s) -> {len(self.docs)} chunks")

    # -- index --------------------------------------------------------------
    def _build(self):
        self._df = {}
        self._N = len(self.docs)
        lengths = []
        for d in self.docs:
            seen = set()
            for t in _tok(d):
                seen.add(t)
            lengths.append(len(_tok(d)))
            for t in seen:
                self._df[t] = self._df.get(t, 0) + 1
        self._avgdl = (sum(lengths) / self._N) if self._N else 1.0
        self._built = True

    def retrieve(self, query, k=4):
        if not self._built:
            self._build()
        if self._N == 0:
            return []
        q_toks = _tok(query)
        if not q_toks:
            return []
        k1, b = 1.5, 0.75
        scores = []
        for d in self.docs:
            dtoks = _tok(d)
            dl = len(dtoks)
            tf = {}
            for t in dtoks:
                tf[t] = tf.get(t, 0) + 1
            s = 0.0
            for t in q_toks:
                if t not in self._df:
                    continue
                idf = math.log((self._N - self._df[t] + 0.5) / (self._df[t] + 0.5) + 1.0)
                f = tf.get(t, 0)
                s += idf * (f * (k1 + 1)) / (f + k1 * (1 - b + b * dl / self._avgdl))
            scores.append(s)
        order = sorted(range(self._N), key=lambda i: scores[i], reverse=True)
        return [self.docs[i] for i in order[:k] if scores[i] > 0]

    # -- persist ------------------------------------------------------------
    def save(self, path=DEFAULT_INDEX):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        json.dump({"docs": self.docs}, open(path, "w"), ensure_ascii=False)
        print(f"[rag] saved index -> {path} ({len(self.docs)} chunks)")

    @classmethod
    def load(cls, path=DEFAULT_INDEX):
        kb = cls()
        if os.path.exists(path):
            data = json.load(open(path, "r", encoding="utf-8"))
            kb.docs = data.get("docs", [])
            kb._built = False
            print(f"[rag] loaded index {path} ({len(kb.docs)} chunks)")
        return kb


# A process-wide default KB, lazily built from data/kb/.
_DEFAULT_KB = None


def default_kb():
    global _DEFAULT_KB
    if _DEFAULT_KB is None:
        if os.path.exists(DEFAULT_INDEX):
            _DEFAULT_KB = KnowledgeBase.load(DEFAULT_INDEX)
        elif os.path.isdir(KB_DIR):
            kb = KnowledgeBase()
            kb.ingest_path(KB_DIR)
            _DEFAULT_KB = kb
        else:
            _DEFAULT_KB = KnowledgeBase()
    return _DEFAULT_KB


def retrieve(query, k=4, kb=None):
    """Top-level retrieval used by the `retrieve` tool + RAG controller."""
    kb = kb or default_kb()
    return kb.retrieve(query, k=k)