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