| """ |
| 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 = [] |
| self._df = {} |
| self._N = 0 |
| self._avgdl = 1.0 |
| self._built = False |
|
|
| |
| 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") |
|
|
| |
| 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] |
|
|
| |
| 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 |
|
|
|
|
| |
| _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) |
|
|
