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