code_release/code/memebench/methods/remem_agent.py

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