File size: 5,446 Bytes
df43f42 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """
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)
|