| """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}, |
| ) |
|
|
