"""A-MEM adapter. RobustAgenticMemorySystem with the evaluation's LLM and embedding clients injected. Each visit is added as a note stamped with the visit date, with memory evolution. A failed LLM call leaves the note with heuristic keywords and no evolution, and the visit is reported as not successful. Retrieval: one keyword-generation call, then the top retrieve_num notes and the notes they link to. """ from __future__ import annotations import importlib import logging import sys import time from pathlib import Path from typing import Dict, List import numpy as np from .base import EMBED_ENC, BaseAgent, MemoryBuildResult, MemoryContext, as_list from ..llm_client import OpenAIEmbedder, create_llm_client logger = logging.getLogger(__name__) AMEM_DIR = (Path(__file__).resolve().parents[2] / "vendor" / "amem" / "A-mem") class _ControllerShim: """Drop-in for llm_controller.llm: callers use get_completion(prompt, temperature). Counts the calls that raised in `failures`, with the last error in `last_error`.""" SYSTEM_MSG = ("Follow the format specified in the prompt exactly. " "Do not add extra commentary.") def __init__(self, client, temperature: float): self._client = client self._temperature = temperature self.failures = 0 self.last_error = None def get_completion(self, prompt: str, temperature: float | None = None) -> str: try: return self._client.chat([{"role": "system", "content": self.SYSTEM_MSG}, {"role": "user", "content": prompt}], temperature=self._temperature if temperature is None else temperature).content except Exception as e: self.failures += 1 self.last_error = f"{type(e).__name__}: {e}"[:200] raise def _make_st_shim(embedder: OpenAIEmbedder, max_input_tokens: int): """Embedding class bound in place of memory_layer.SentenceTransformer.""" def _truncate(t: str) -> str: ids = EMBED_ENC.encode(t, disallowed_special=()) return EMBED_ENC.decode(ids[:max_input_tokens]) if len(ids) > max_input_tokens else t class _STShim: def __init__(self, _model_name: str): self._emb = embedder def encode(self, texts): single = isinstance(texts, str) batch = [_truncate(texts)] if single else [_truncate(t) for t in texts] vecs = np.asarray(self._emb.embed(batch), dtype="float32") return vecs[0] if single else vecs return _STShim class AMemAgent(BaseAgent): ANSWERS_INTERNALLY = False MEMORIZE_WITH_LLM = True def __init__(self, *, model: str | None = None, retrieve_num: int, amem_evo_threshold: int, max_tokens: int, temperature: float, memorize_chunk_tokens: int, embedding_max_input_tokens: int, embedding_model: str | None = None, embedding_batch_size: int | None = None, internal_base_url: str | None = None, working_dir: str | None = None, **params): super().__init__(**params) self.model = model self.retrieve_num = retrieve_num self._vec_seen: set[str] = set() self._kw_fallbacks = 0 self.memorize_chunk_tokens = memorize_chunk_tokens self.embedding_max_input_tokens = embedding_max_input_tokens if str(AMEM_DIR) not in sys.path: sys.path.append(str(AMEM_DIR)) ml = importlib.import_module("memory_layer") mlr = importlib.import_module("memory_layer_robust") self._embedder = OpenAIEmbedder(model=embedding_model, batch_size=embedding_batch_size) ml.SentenceTransformer = _make_st_shim(self._embedder, self.embedding_max_input_tokens) self._system = mlr.RobustAgenticMemorySystem( model_name="injected", llm_backend="openai", llm_model=model, evo_threshold=amem_evo_threshold, api_key="injected", check_connection=False) self._llm_shim = _ControllerShim(create_llm_client( model=model, max_tokens=max_tokens, base_url=internal_base_url), temperature=temperature) self._system.llm_controller.llm = self._llm_shim self._note_visit: Dict[str, set] = {} self._cur_visit: int | None = None self._trace_evolution() def _trace_evolution(self) -> None: """Record the neighbour lookup and the evolution decision.""" sysm = self._system if getattr(sysm, "__traced__", False): return orig_find = sysm.find_related_memories orig_proc = sysm.process_memory state = {"stage": "query"} def find(query, k=5): out = orig_find(query, k=k) try: mem_str, indices = out if isinstance(out, tuple) else (out, []) idx = [] if indices is None else list(indices) ids = list(self._system.memories.keys()) if state["stage"] == "write": mems = self._system.memories state["nb"] = {} for i in idx: if int(i) < len(ids): nb = mems.get(ids[int(i)]) state["nb"][ids[int(i)]] = ( str(getattr(nb, "tags", None)), str(getattr(nb, "context", None))) self._trace("neighbour_lookup", stage=state["stage"], query=query, k=k, n_found=len(idx), neighbours=[{"index": int(i), "visits": (sorted(self._note_visit.get(ids[int(i)], ())) if int(i) < len(ids) else None)} for i in idx], neighbour_text=mem_str) except Exception: pass return out def proc(note): state["stage"] = "write" state["nb"] = {} before = getattr(note, "content", "") kw_before = as_list(getattr(note, "keywords", None)) try: evolved, out_note = orig_proc(note) finally: state["stage"] = "query" cur = self._cur_visit merged: list = [] try: if cur is not None: mems = self._system.memories for nid, sig in (state.get("nb") or {}).items(): nb = mems.get(nid) if nb is None: continue if (str(getattr(nb, "tags", None)), str(getattr(nb, "context", None))) != sig: merged.append({"note_id": str(nid), "visits_before": sorted( self._note_visit.get(str(nid), ()))}) self._note_visit.setdefault(str(nid), set()).add(int(cur)) except Exception: pass try: self._trace("evolution", note_id=str(getattr(note, "id", "")), evolved=bool(evolved), visit=cur, neighbours_merged=merged, content_before=before, content_after=getattr(out_note, "content", ""), keywords_before=kw_before, keywords_after=as_list(getattr(out_note, "keywords", None)), links_after=as_list(getattr(out_note, "links", None))) except Exception: pass return evolved, out_note retr = getattr(sysm, "retriever", None) if retr is not None and not getattr(retr, "__traced__", False): orig_search = retr.search def search(query, k=5, *a, **kw): out = orig_search(query, k, *a, **kw) try: ids = list(self._system.memories.keys()) res = [] if out is None else list(out) self._trace("retrieval", stage=state["stage"], query=query, k=k, n_found=len(res), hits=[{"index": int(i), "visits": sorted(self._note_visit.get(ids[int(i)], ())) if int(i) < len(ids) else None, "content": (self._system.memories[ids[int(i)]].content if int(i) < len(ids) else "")[:300]} for i in res]) except Exception: pass return out retr.search = search retr.__traced__ = True sysm.find_related_memories = find sysm.process_memory = proc sysm.__traced__ = True def _split_chunks(self, text: str) -> List[str]: toks = self.count_tokens(text) if toks <= self.memorize_chunk_tokens: return [text] if text.strip() else [] from .base import ENC ids = ENC.encode(text) return [ENC.decode(ids[i:i + self.memorize_chunk_tokens]) for i in range(0, len(ids), self.memorize_chunk_tokens)] def memorize(self, text: str, visit_number: int, date: str) -> MemoryBuildResult: """Add the visit as notes. Not successful when add_note raises (the visit is rolled back) or when an LLM call of content analysis or evolution failed (the notes are kept as A-MEM stored them).""" t0 = time.time() failures_before = self._llm_shim.failures evo_before = self._system.evo_cnt pre_retr = self._system.retriever pre_corpus = len(pre_retr.corpus) pre_emb = (0 if getattr(pre_retr, "embeddings", None) is None else int(pre_retr.embeddings.shape[0])) pre_docids = dict(pre_retr.document_ids) pre_notes = len(self._system.memories) self._cur_visit = visit_number added = [] try: for chunk in self._split_chunks(text): note_id = self._system.add_note(content=chunk, time=date) self._note_visit.setdefault(str(note_id), set()).add(int(visit_number)) added.append(str(note_id)) except Exception as e: logger.error(f"[A-MEM] add_note failed on v{visit_number}: {e}") try: self._system.retriever = pre_retr del pre_retr.corpus[pre_corpus:] if pre_emb == 0: pre_retr.embeddings = None elif pre_retr.embeddings is not None: pre_retr.embeddings = pre_retr.embeddings[:pre_emb] pre_retr.document_ids = pre_docids mems = self._system.memories while len(mems) > pre_notes: orphan_id = next(reversed(mems)) mems.pop(orphan_id) self._note_visit.pop(str(orphan_id), None) self._system.evo_cnt = evo_before except Exception as rb_err: logger.error(f"[A-MEM] rollback failed, store may be inconsistent: " f"{rb_err}") return MemoryBuildResult(success=False, method="amem", action="add_note", visit_number=visit_number, chunk_count=len(self._system.memories), time_cost=time.time() - t0, extra={"error": str(e)[:200], "rolled_back_to_notes": pre_notes}) extra = {"evolutions": self._system.evo_cnt - evo_before, "evo_cnt_total": self._system.evo_cnt} llm_failures = self._llm_shim.failures - failures_before if llm_failures: logger.warning("[A-MEM] v%s: %d failed LLM call(s); A-MEM stored the visit without their " "result (%s)", visit_number, llm_failures, self._llm_shim.last_error) extra.update(llm_failures=llm_failures, error=self._llm_shim.last_error) return MemoryBuildResult(success=llm_failures == 0, method="amem", action="add_note", visit_number=visit_number, chunk_count=len(self._system.memories), added_chunks=len(added), time_cost=time.time() - t0, extra=extra) _KEYWORDS_PROMPT = ("Given the following question, generate several keywords " "separated by commas.\n\nQuestion: {question}\n\nKeywords:") def _keyword_retrieve(self, question: str): """Keyword query, then the top retrieve_num notes and the notes they link to.""" from llm_text_parsers import parse_keywords_response keywords = parse_keywords_response( self._system.llm_controller.llm.get_completion( self._KEYWORDS_PROMPT.format(question=question))) query = str(keywords).strip() if not query: self._kw_fallbacks += 1 query = question if not self._system.memories: return "", [], query indices = self._system.retriever.search(query, self.retrieve_num) all_memories = list(self._system.memories.values()) note_ids = list(self._system.memories.keys()) memory_str, hit_ids = "", [] for i in indices: j = 0 m = all_memories[i] memory_str += ("talk start time:" + m.timestamp + "memory content: " + m.content + "memory context: " + m.context + "memory keywords: " + str(m.keywords) + "memory tags: " + str(m.tags) + "\n") hit_ids.append(note_ids[i]) for neighbor in m.links: n = all_memories[neighbor] memory_str += ("talk start time:" + n.timestamp + "memory content: " + n.content + "memory context: " + n.context + "memory keywords: " + str(n.keywords) + "memory tags: " + str(n.tags) + "\n") hit_ids.append(note_ids[neighbor]) if j >= self.retrieve_num: break j += 1 return memory_str, hit_ids, query def build_memory_context(self, question: str, as_of: str = "", display_question: str | None = None) -> MemoryContext: try: memory_str, hit_ids, keywords = self._keyword_retrieve(question) except Exception as e: logger.error(f"[A-MEM] retrieval failed: {e}") return MemoryContext(mem="", n_tok=0, retrieved_count=0, error=True, extra={"method": "amem", "error": str(e)[:200]}) visits = sorted({v for h in hit_ids for v in self._note_visit.get(h, ())}) return MemoryContext( mem=memory_str, n_tok=self.count_tokens(memory_str), retrieved_count=len(hit_ids), retrieved_visits=visits, retrieved_memories=[{"visit": None, "memory": memory_str[:300]}], extra={"method": "amem", "query_keywords": str(keywords)[:200], "keyword_fallbacks": self._kw_fallbacks, "notes_total": len(self._system.memories), "evo_cnt": self._system.evo_cnt}) def dump_memory(self, out_dir): from .base import snapshot_jsonl_gz, snapshot_vectors_delta rows = [] for nid, note in self._system.memories.items(): d = {k: v for k, v in vars(note).items() if k not in ("embedding",) and not k.startswith("_")} d["note_id"] = str(nid) d["visits"] = sorted(self._note_visit.get(str(nid), ())) rows.append(d) n = snapshot_jsonl_gz(out_dir / "notes.jsonl.gz", rows) nv = 0 try: r = self._system.retriever corpus = list(getattr(r, "corpus", []) or []) embs = getattr(r, "embeddings", None) if embs is not None and len(corpus) == len(embs): nv = snapshot_vectors_delta( out_dir, ((f"pos{i}", embs[i]) for i in range(len(corpus))), self._vec_seen) snapshot_jsonl_gz(out_dir / "corpus_order.jsonl.gz", ({"pos": i, "text": t} for i, t in enumerate(corpus))) except Exception as e: logger.error("[A-MEM] vector snapshot failed (run continues): %s", e) return {"notes": n, "evo_cnt": self._system.evo_cnt, "corpus": len(getattr(self._system.retriever, "corpus", [])), "vectors_new": nv}