File size: 3,526 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 78 79 80 81 82 | """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},
)
|