| """Dense retrieval over visit chunks with a FAISS inner-product index. |
| |
| New chunks are embedded at each memorize(); chunking and the MEMORY block follow bm25_rag. |
| """ |
| from __future__ import annotations |
|
|
| import time |
|
|
| import numpy as np |
|
|
| from .base import BaseAgent, MemoryBuildResult, MemoryContext |
|
|
|
|
| class DenseRAGAgent(BaseAgent): |
| def __init__(self, embedder=None, *, top_k: int, chunk_size: int, chunk_overlap: int, |
| max_context_tokens: int, **params): |
| super().__init__(**params) |
| if embedder is None: |
| raise ValueError("DenseRAGAgent requires an embedder (llm_client.make_embedder())") |
| self.embedder = embedder |
| self.top_k = top_k |
| 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._index = None |
|
|
| @staticmethod |
| def _normalize(vecs: list[list[float]]) -> np.ndarray: |
| a = np.asarray(vecs, dtype="float32") |
| norms = np.linalg.norm(a, axis=1, keepdims=True) |
| norms[norms == 0] = 1.0 |
| return a / norms |
|
|
| 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) |
| if new_chunks: |
| vecs = self._normalize(self.embedder.embed(new_chunks)) |
| if self._index is None: |
| import faiss |
| self._index = faiss.IndexFlatIP(vecs.shape[1]) |
| self._index.add(vecs) |
| self._chunks.extend(new_chunks) |
| self._chunk_visits.extend([visit_number] * len(new_chunks)) |
| return MemoryBuildResult(success=True, method="dense_rag", action="add_to_faiss", |
| visit_number=visit_number, chunk_count=len(self._chunks), |
| added_chunks=len(new_chunks), time_cost=time.time() - t0, |
| extra={"embedding_model": self.embedder.model}) |
|
|
| def _rank(self, query: str) -> list[int]: |
| if self._index is None or not self._chunks: |
| return [] |
| qv = self._normalize(self.embedder.embed([query])) |
| k = min(self.top_k, len(self._chunks)) |
| _, idxs = self._index.search(qv, k) |
| return [int(i) for i in idxs[0] if 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), |
| "embedding_model": self.embedder.model}, |
| ) |
|
|
