File size: 3,386 Bytes
750b88d | 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 | """BM25 sparse retrieval over visit chunks.
Text is tokenized with nltk word_tokenize (split on whitespace when nltk is not available). Each visit is
chunked on its own into `chunk_size`-token chunks, and continuation chunks repeat the visit header. The top
chunks form the MEMORY block.
"""
from __future__ import annotations
import time
from .base import BaseAgent, MemoryBuildResult, MemoryContext
class BM25RAGAgent(BaseAgent):
def __init__(self, *, top_k: int, k1: float, b: float, chunk_size: int, chunk_overlap: int,
max_context_tokens: int, **params):
super().__init__(**params)
self.top_k = top_k
self.k1 = k1
self.b = b
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
self.max_context_tokens = max_context_tokens
self._chunks: list[str] = []
self._chunk_visits: list[int] = []
self._tokenized: list[list[str]] = []
self._bm25 = None
@staticmethod
def _words(text: str) -> list[str]:
try:
from nltk.tokenize import word_tokenize
return word_tokenize(text.lower())
except (ImportError, LookupError):
return text.lower().split()
def memorize(self, text: str, visit_number: int, date: str) -> MemoryBuildResult:
t0 = time.time()
header = text.splitlines()[0] if text else ""
new_chunks = self.split_tokens(text, self.chunk_size, self.chunk_overlap, header=header)
self._chunks.extend(new_chunks)
self._chunk_visits.extend([visit_number] * len(new_chunks))
self._tokenized.extend(self._words(c) for c in new_chunks)
from rank_bm25 import BM25Okapi
self._bm25 = BM25Okapi(self._tokenized, k1=self.k1, b=self.b)
return MemoryBuildResult(success=True, method="bm25_rag", action="build_bm25_index",
visit_number=visit_number, chunk_count=len(self._chunks),
added_chunks=len(new_chunks), time_cost=time.time() - t0)
def _rank(self, query: str) -> list[int]:
if self._bm25 is None or not self._chunks:
return []
scores = self._bm25.get_scores(self._words(query))
top = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:self.top_k]
return [i for i in top if scores[i] > 0]
def build_memory_context(self, question: str,
as_of: str = "", display_question: str | None = None) -> MemoryContext:
idxs = self._rank(question)
kept, used = [], 0
for i in idxs:
n = self.count_tokens(self._chunks[i])
if used + n > self.max_context_tokens:
break
kept.append(i)
used += n
mem = "\n\n".join(f"[Memory {j + 1}]\n{self._chunks[i]}" for j, i in enumerate(kept))
visits = [self._chunk_visits[i] for i in kept]
return MemoryContext(
mem=mem, n_tok=self.count_tokens(mem), retrieved_count=len(kept),
retrieved_visits=sorted(set(visits)),
retrieved_memories=[{"visit": self._chunk_visits[i],
"memory": self._chunks[i][:300]} for i in kept],
context_truncated=len(kept) < len(idxs),
extra={"original_retrieved_count": len(idxs), "k1": self.k1, "b": self.b},
)
|