| """ReMem adapter (ReMem-I, the agentic variant). |
| |
| Visits accumulate; at each checkpoint the engine is rebuilt over all visits so far from the |
| cached gist and fact extractions, so memory at a checkpoint equals ReMem run once over the |
| visits up to it. Questions go through rag_for_qa: a prompted agent chooses among |
| semantic, lexical, gist and entity retrieval tools for up to `agent_max_steps` turns and gives |
| the answer. Each patient has its own working directory and OpenIE cache. ReMem writes its log |
| files (they contain the question text) to `remem_logs` under the working directory, unless the |
| LOG_DIR environment variable already names a folder. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import logging |
| import os |
| import re |
| import shutil |
| import time |
| import traceback |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Any, Dict, Optional |
|
|
| from .base import BaseAgent, MemoryBuildResult, MemoryContext |
| from .engine_bridge import EngineChat, EngineEmbedder, add_engine_source, read_openie_docs |
| from ..llm_client import USAGE, create_llm_client |
| from ..prompts import ANSWER_NOCTX_NOJSON |
| from .hipporag_agent import _unwrap_answer |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def _install_ordered_openie_loader(remem) -> None: |
| """Return cached and new OpenIE entries in document order, matching the positional index loop.""" |
| orig_load = type(remem).load_existing_openie |
|
|
| def ordered_load(chunk_keys): |
| all_openie_info, _ = orig_load(remem, chunk_keys) |
| by_idx = {info["idx"]: info for info in all_openie_info |
| if info.get("gists") is not None and info.get("facts") is not None} |
| keys = list(chunk_keys) |
| n = 0 |
| while n < len(keys) and keys[n] in by_idx: |
| n += 1 |
| return [by_idx[k] for k in keys[:n]], keys[n:] |
|
|
| remem.load_existing_openie = ordered_load |
|
|
|
|
| class OpenIECacheError(RuntimeError): |
| """The OpenIE cache does not line up with the indexed documents.""" |
|
|
|
|
| _RETRIEVAL_QUERIES: Dict[str, str] = {} |
| """Agent prompt -> question text that the ReMem retrieval tools search with instead of the prompt.""" |
| _RETRIEVAL_TOOLS = ("semantic_retrieve_tool.SemanticRetrieveTool", |
| "lexical_retrieve_tool.LexicalRetrieveTool", |
| "find_gist_contexts_tool.FindGistContextsTool", |
| "find_entity_contexts_tool.FindEntityContextsTool") |
|
|
|
|
| def _install_retrieval_queries() -> None: |
| """Make the ReMem retrieval tools search with the entry of _RETRIEVAL_QUERIES for the agent's prompt.""" |
| import importlib |
| for spec in _RETRIEVAL_TOOLS: |
| mod_name, cls_name = spec.split(".") |
| cls = getattr(importlib.import_module(f"remem.agent.tools.{mod_name}"), cls_name) |
| if getattr(cls, "__memebench_retrieval_queries__", False): |
| continue |
| orig_execute = cls.execute |
|
|
| def execute(self, context, *a, _orig=orig_execute, **kw): |
| raw = _RETRIEVAL_QUERIES.get(getattr(context, "query", None)) |
| if raw is None: |
| return _orig(self, context, *a, **kw) |
| shown = context.query |
| context.query = raw |
| try: |
| return _orig(self, context, *a, **kw) |
| finally: |
| context.query = shown |
|
|
| cls.execute = execute |
| cls.__memebench_retrieval_queries__ = True |
|
|
|
|
| _VISIT_HEADER_RE = re.compile(r"=== Visit v(\d+) · ") |
|
|
|
|
| @dataclass |
| class _Patient: |
| """Per-patient state. `docs` / `doc_visits` hold every visit so far (each rebuild indexes all of them); |
| `texts` / `visits` the visits received since the last build.""" |
| docs: list = field(default_factory=list) |
| doc_visits: list = field(default_factory=list) |
| texts: list = field(default_factory=list) |
| visits: list = field(default_factory=list) |
| indexed: list = field(default_factory=list) |
| received: int = 0 |
| built: bool = False |
| error: str | None = None |
| engine: Any = None |
|
|
|
|
| class ReMemAgent(BaseAgent): |
| MEMORIZE_WITH_LLM = True |
| ANSWERS_INTERNALLY = True |
|
|
| def __init__( |
| self, |
| *, |
| model: str | None = None, |
| temperature: float, |
| max_tokens: int, |
| provider: str = "openai", |
| api_key: Optional[str] = None, |
| base_url: Optional[str] = None, |
| internal_base_url: Optional[str] = None, |
| extract_method: str, |
| dataset: str, |
| is_directed_graph: bool, |
| synonymy_edge_sim_threshold: float, |
| synonymy_edge_topk: int, |
| retrieval_top_k: int, |
| qa_top_k: int, |
| linking_top_k: int, |
| damping: float, |
| passage_node_weight: float, |
| agent_fixed_tools: bool, |
| agent_max_steps: int, |
| save_openie: bool, |
| embedding_provider: str, |
| embedding_model: str | None = None, |
| embedding_dim: Optional[int] = None, |
| embedding_api_key: Optional[str] = None, |
| embedding_base_url: Optional[str] = None, |
| embedding_batch_size: int | None = None, |
| embedding_max_seq_len: int, |
| text_preprocessor_class_name: str = "TextPreprocessor", |
| max_context_tokens: int, |
| working_dir: Optional[str] = None, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| if embedding_provider not in ("openai", "api"): |
| raise ValueError(f"embedding_provider must be 'openai' or 'api', not {embedding_provider!r}") |
| base_url = base_url or internal_base_url |
| self.extract_method = extract_method |
| self.qa_top_k = qa_top_k |
| self.retrieval_top_k = retrieval_top_k |
| self.max_context_tokens = max_context_tokens |
| self.working_dir = working_dir or str(Path(__file__).resolve().parents[2] / "outputs" / "_remem_store") |
| os.environ.setdefault("LOG_DIR", os.path.abspath(os.path.join(self.working_dir, "remem_logs"))) |
| self._store_cleared = False |
| self._engine_settings = dict( |
| dataset=dataset, |
| llm_name=model, |
| llm_base_url=base_url or os.environ.get("OPENAI_BASE_URL") or "https://api.openai.com/v1", |
| llm_infer_mode="online", |
| temperature=temperature, |
| extract_llm_label=model, |
| extract_method=extract_method, |
| is_directed_graph=is_directed_graph, |
| synonymy_edge_sim_threshold=synonymy_edge_sim_threshold, |
| synonymy_edge_topk=synonymy_edge_topk, |
| graph_type="facts_and_sim", |
| retrieval_top_k=retrieval_top_k, |
| qa_top_k=qa_top_k, |
| linking_top_k=linking_top_k, |
| damping=damping, |
| passage_node_weight=passage_node_weight, |
| agent_fixed_tools=agent_fixed_tools, |
| agent_max_steps=agent_max_steps, |
| force_index_from_scratch=False, |
| force_openie_from_scratch=False, |
| save_openie=save_openie, |
| embedding_model_name=embedding_model, |
| embedding_batch_size=embedding_batch_size, |
| embedding_max_seq_len=embedding_max_seq_len, |
| embedding_return_as_normalized=True, |
| text_preprocessor_class_name=text_preprocessor_class_name, |
| preprocess_chunk_func="none", |
| do_eval_qa=False, |
| do_eval_retrieval=False, |
| qa_passage_prefix="- ", |
| ) |
| client = create_llm_client(provider=provider, model=model, temperature=temperature, |
| max_tokens=max_tokens, api_key=api_key, base_url=base_url) |
| self._chat = EngineChat(client, temperature=temperature, echo=True) |
| self._embedder = EngineEmbedder(model=embedding_model, batch_size=embedding_batch_size, dim=embedding_dim, |
| api_key=embedding_api_key, base_url=embedding_base_url) |
| self._patients: dict[Any, _Patient] = {} |
| self._last_build: dict = {} |
| add_engine_source("REMem") |
| self._engine_imported = False |
|
|
| def _patient(self) -> tuple[Any, _Patient]: |
| key = self._context_id if self._context_id is not None else 0 |
| if key not in self._patients: |
| self._patients[key] = _Patient() |
| return key, self._patients[key] |
|
|
| def _store_dir(self, key) -> str: |
| """The patient's engine directory; at the first use in this run an earlier store is removed.""" |
| path = os.path.join(self.working_dir, f"context_{key}") |
| if not self._store_cleared: |
| self._store_cleared = True |
| if os.path.isdir(path): |
| shutil.rmtree(path) |
| logger.info("[ReMem] removed the store left at %s", path) |
| os.makedirs(path, exist_ok=True) |
| return path |
|
|
| def _engine(self, key, state: _Patient): |
| """The patient's ReMem instance, created on first use with the evaluation's chat and embedding clients.""" |
| if state.engine is not None: |
| return state.engine |
| from remem.remem import ReMem |
| from remem.utils.config_utils import BaseConfig |
| if not self._engine_imported: |
| _install_retrieval_queries() |
| self._engine_imported = True |
|
|
| store = self._store_dir(key) |
| engine = ReMem(global_config=BaseConfig(save_dir=store, **self._engine_settings), working_dir=store, |
| llm=self._chat, extract_llm=self._chat, qa_llm=self._chat) |
| label = str(engine.global_config.extract_llm_label).replace("/", "_") |
| engine.openie_results_path = os.path.join(store, f"openie_results_ner_{label}.json") |
| engine.set_embedding_model(self._embedder) |
| _install_ordered_openie_loader(engine) |
| logger.info("[ReMem] engine ready for %s (extract_method=%s, store=%s)", key, self.extract_method, store) |
| state.engine = engine |
| return engine |
|
|
| def dump_memory(self, out_dir): |
| """Per-checkpoint export: the store tree with the visit of each row and the build state.""" |
| from .base import snapshot_jsonl_gz, snapshot_tree |
| stats = {} |
| for cid, state in self._patients.items(): |
| if state.engine is None and not state.received: |
| continue |
| iv = state.indexed |
| stats[f"build_{cid}"] = { |
| "build_ok": self._last_build.get("build_ok"), |
| "build_error": state.error, |
| "n_indexed_visits": len(iv), |
| "indexed_visits_max": max(iv) if iv else None, |
| "pending_unbuilt_visits": list(state.visits), |
| "qa_top_k": self.qa_top_k, "retrieval_top_k": self.retrieval_top_k, |
| } |
| for cid, state in self._patients.items(): |
| rm = state.engine |
| if rm is None: |
| continue |
| try: |
| src = self._store_dir(cid) |
| if os.path.isdir(src): |
| stats[f"store_files_{cid}"] = snapshot_tree(src, out_dir / f"store_{cid}") |
| except Exception as e: |
| logger.error("[ReMem] store snapshot failed (run continues): %s", e) |
| rows = [] |
| stores = getattr(rm, "episodic_embedding_stores", None) or {} |
| for name, store in stores.items(): |
| try: |
| all_rows = store.get_hash_id_to_row_readonly() |
| except Exception: |
| continue |
| for hid, row in all_rows.items(): |
| md = row.get("metadata") if isinstance(row, dict) else None |
| md = md if isinstance(md, dict) else {} |
| rows.append({"store": name, "hash_id": str(hid), |
| "source_visit": md.get("source_visit"), |
| "content": str(row.get("content", ""))[:300]}) |
| stats[f"rows_{cid}"] = snapshot_jsonl_gz( |
| out_dir / f"memory_rows_{cid}.jsonl.gz", rows) |
| g = getattr(rm, "graph", None) |
| if g is not None: |
| try: |
| stats[f"graph_{cid}"] = {"vertices": g.vcount(), "edges": g.ecount()} |
| except Exception: |
| pass |
| return stats or None |
|
|
| def _reset_derived_artifacts(self, key, state: _Patient) -> None: |
| """Remove the embedding stores and graph so a fresh engine rebuilds them; keep the OpenIE cache.""" |
| d = self._store_dir(key) |
| removed = [] |
| for name in sorted(os.listdir(d)): |
| path = os.path.join(d, name) |
| if os.path.isdir(path) and name.endswith("_embeddings"): |
| shutil.rmtree(path) |
| removed.append(name) |
| elif name == "graph.pkl": |
| os.remove(path) |
| removed.append(name) |
| state.engine = None |
| logger.info(f"[ReMem] rebuild: cleared {removed} (OpenIE cache kept)") |
|
|
| def memorize(self, text: str, visit_number: int, date: str) -> MemoryBuildResult: |
| """Accumulate one visit; the build runs at the next checkpoint.""" |
| USAGE.set_phase("memorize") |
| t0 = time.time() |
| key, state = self._patient() |
| if date: |
| text = f"Date: {date}\n{text}" |
| state.texts.append(text) |
| state.visits.append(visit_number) |
| if text.strip(): |
| state.docs.append(text.strip()) |
| state.doc_visits.append(visit_number) |
| state.received += 1 |
| return MemoryBuildResult( |
| success=True, method="remem", action="accumulate", visit_number=visit_number, |
| chunk_count=state.received, added_chunks=1, time_cost=time.time() - t0, |
| extra={"context_id": key, "extract_method": self.extract_method, "session_number": state.received, |
| "pending_sessions": len(state.texts), "total_accumulated_chars": sum(map(len, state.texts))}) |
|
|
| def _build_if_pending(self, key, state: _Patient) -> bool: |
| """Build at a checkpoint if visits are pending; return True if a build ran.""" |
| if not state.texts: |
| return False |
| USAGE.set_phase("memorize") |
| build = self._rebuild(key, state) |
| self._last_build = { |
| "documents_indexed": build.get("num_documents"), |
| "sessions_since_last_build": build.get("session_count"), |
| "indexed_visits_total": len(state.indexed), |
| "chunks_in_store": build.get("chunks_in_store"), |
| "graph_nodes": (build.get("graph_info") or {}).get("num_total_nodes"), |
| "build_time_s": round(build.get("time_cost", 0.0), 1), |
| "build_ok": build.get("success", False), |
| **({"build_error": build["error"]} if build.get("error") else {}), |
| } |
| logger.info(f"[ReMem] indexed: {self._last_build}") |
| return True |
|
|
| def consolidate(self) -> dict | None: |
| """Build at the checkpoint, so the checkpoint snapshot holds the store its questions see.""" |
| if not self._build_if_pending(*self._patient()): |
| return None |
| out = {"built": bool(self._last_build.get("build_ok")), |
| **{k: v for k, v in self._last_build.items() |
| if k in ("documents_indexed", "sessions_since_last_build", |
| "graph_nodes", "build_time_s", "build_ok")}} |
| if not out["built"]: |
| out["error"] = f"build failed: {self._last_build.get('build_error', '?')}" |
| return out |
|
|
| def _rebuild(self, key, state: _Patient) -> dict: |
| """Rebuild the patient's engine over every visit so far, reusing the cached extractions. Returns the |
| build summary; on failure the patient is marked unbuilt and the summary carries the error.""" |
| t0 = time.time() |
| waiting = len(state.texts) |
| try: |
| self._reset_derived_artifacts(key, state) |
| engine = self._engine(key, state) |
| docs = list(state.docs) |
| config = engine.global_config |
| saved, config.force_index_from_scratch = config.force_index_from_scratch, True |
| try: |
| engine.index(docs, doc_visits=state.doc_visits) |
| finally: |
| config.force_index_from_scratch = saved |
| if not os.path.exists(os.path.join(self._store_dir(key), "graph.pkl")): |
| raise RuntimeError("[ReMem] rebuild finished without writing graph.pkl") |
| engine.ready_to_retrieve = False |
| graph = engine.get_graph_info() |
| cached = read_openie_docs(engine.openie_results_path) |
| if cached: |
| if len(cached) != len(docs): |
| raise OpenIECacheError(f"[ReMem OpenIE cache] misaligned: {len(cached)} entries vs " |
| f"{len(docs)} docs") |
| probes = (0, len(cached) // 2, len(cached) - 1) |
| off = [k for k in probes if cached[k].get("passage", cached[k].get("verbatim")) != docs[k]] |
| if off: |
| raise OpenIECacheError(f"[ReMem OpenIE cache] misaligned: mismatches at {off}") |
| except OpenIECacheError: |
| raise |
| except Exception as e: |
| logger.error("[ReMem] rebuild failed: %s", e) |
| traceback.print_exc() |
| state.built = False |
| state.error = str(e)[:300] |
| return {"success": False, "session_count": waiting, "error": str(e), "time_cost": time.time() - t0} |
|
|
| state.built = True |
| state.error = None |
| state.indexed.extend(state.visits) |
| try: |
| chunks = len(engine.chunk_embedding_store.hash_id_to_row) |
| except Exception: |
| chunks = None |
| state.texts, state.visits = [], [] |
| logger.info("[ReMem] rebuilt over %d visits in %.1fs; graph %s", len(docs), time.time() - t0, graph) |
| return {"success": True, "session_count": waiting, "num_documents": len(docs), "chunks_in_store": chunks, |
| "graph_info": graph, "openie_count": len(cached) if cached else 0, "time_cost": time.time() - t0} |
|
|
| def build_memory_context(self, question: str, |
| as_of: str = "", display_question: str | None = None) -> MemoryContext: |
| """Build if visits are pending, then retrieve and answer through rag_for_qa.""" |
| key, state = self._patient() |
| self._build_if_pending(key, state) |
| if state.error: |
| return MemoryContext(mem="", n_tok=0, retrieved_count=0, error=True, answer="[err build_failed]", |
| extra={"method": "remem", "error": "build_failed", "build_error": state.error, |
| **self._last_build}) |
|
|
| USAGE.set_phase("query") |
| t0 = time.time() |
| if not state.built: |
| logger.warning("[ReMem] nothing is indexed for %s", key) |
| return MemoryContext(mem="", n_tok=0, retrieved_count=0, answer="[err no_data_indexed]", error=True, |
| extra={"method": "remem", "error": "no_data_indexed"}) |
|
|
| found, notes, answer, failed = [], {}, "", False |
| try: |
| prompt = ANSWER_NOCTX_NOJSON.format(as_of=as_of, qtext=display_question or question) |
| answer, found, notes = self._ask(self._engine(key, state), prompt, question) |
| except Exception as e: |
| logger.error("[ReMem] query failed: %s", e) |
| traceback.print_exc() |
| notes, answer, failed = {"error": str(e)}, f"[err {str(e)[:60]}]", True |
|
|
| kept, used = [], 0 |
| for doc in found: |
| size = self.count_tokens(doc["content"]) |
| if used + size > self.max_context_tokens: |
| break |
| kept.append(doc) |
| used += size |
|
|
| mem = "\n\n".join(f"[Memory {i}]\n{d['content']}" for i, d in enumerate(kept, 1)) |
| visits = sorted({int(d["source_visit"]) for d in kept if d.get("source_visit") is not None}) or \ |
| sorted({int(v) for d in kept for v in _VISIT_HEADER_RE.findall(d["content"])}) |
| return MemoryContext( |
| mem=mem, |
| n_tok=self.count_tokens(mem), |
| retrieved_count=len(kept), |
| retrieved_visits=visits, |
| retrieved_memories=[{"visit": None, "memory": d["content"][:300], "score": d.get("score")} |
| for d in kept], |
| context_truncated=len(kept) < len(found), |
| answer=_unwrap_answer(answer), error=failed, |
| extra={"method": "remem", "extract_method": self.extract_method, |
| "query_time_s": round(time.time() - t0, 2), **notes, **self._last_build}, |
| ) |
|
|
| def _ask(self, engine, prompt: str, question: str) -> tuple[str, list, dict]: |
| """Run ReMem's rag_for_qa agent on `prompt`; its retrieval tools search with `question`. |
| Returns the answer, the passages of the solution (rank, visit, content, score) and the graph seeds.""" |
| if question and question != prompt: |
| _RETRIEVAL_QUERIES[prompt] = question |
| try: |
| solutions, _, metadata, _, retrieval = engine.rag_for_qa( |
| queries=[prompt], gold_docs=None, gold_answers=None, metrics=(), to_save=False) |
| finally: |
| _RETRIEVAL_QUERIES.pop(prompt, None) |
| try: |
| self._trace("query_retrieve", question=prompt, retrieval_question=question or prompt, |
| retrieval=retrieval, metadata=metadata) |
| except Exception: |
| pass |
| if not solutions: |
| return "No answer generated", [], {} |
|
|
| solution = solutions[0] |
| answer = solution.answer if hasattr(solution, "answer") else str(solution) |
| docs = list(getattr(solution, "docs", None) or [])[:self.qa_top_k] |
| scores = getattr(solution, "doc_scores", None) |
| scores = list(scores) if scores is not None and len(scores) > 0 else [] |
| metas = getattr(solution, "doc_metadata", None) or [] |
| found = [] |
| for i, doc in enumerate(docs): |
| score = scores[i] if i < len(scores) else None |
| found.append({"rank": i + 1, |
| "source_visit": ((metas[i] if i < len(metas) else None) or {}).get("source_visit"), |
| "content": doc if isinstance(doc, str) else str(doc), |
| "score": float(score) if score is not None else None}) |
| seeds = [] |
| if getattr(solution, "graph_seeds", None) is not None: |
| try: |
| seeds = [str(s) for s in solution.graph_seeds[:5]] |
| except (TypeError, ValueError): |
| seeds = [] |
| return answer, found, {"graph_seeds": seeds} |
|
|
