code_release/code/memebench/methods/lightmem_agent.py

"""LightMem adapter.

Dialogue turns are added one by one (topic segmentation, then LLM extraction of memory
entries); retrieval returns the top entries for a question and the runner answers. The
evaluation's LLM and embedding clients are injected, and entry timestamps come from
the visit date. With `offline_consolidate: true`, LightMem's sleep-time update runs at each
checkpoint. A visit with a failed extraction call is reported as not successful, and a
consolidation with failed update calls returns an `error` entry. The topic segmenter is loaded
from a pinned Hugging Face revision.
"""
from __future__ import annotations

import logging
import shutil
import sys
import threading
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List

from .base import BaseAgent, MemoryBuildResult, MemoryContext, as_list
from ..llm_client import ENDPOINTS, MAX_ATTEMPTS, api_key_for, create_llm_client, http_timeout

logger = logging.getLogger(__name__)

LIGHTMEM_SRC = (Path(__file__).resolve().parents[2]
                / "vendor" / "LightMem" / "src")

SEGMENTER_MODEL = "microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank"
SEGMENTER_REVISION = "5f0c82792b7ea14c6484e015b6a072009496b7f2"


def _ensure_engine_on_path() -> None:
    if str(LIGHTMEM_SRC) not in sys.path:
        sys.path.append(str(LIGHTMEM_SRC))
    for name in ("LightMemory", "lightmem"):
        logging.getLogger(name).setLevel(logging.WARNING)


def _pinned_model_path(name: str, revision: str | None) -> str:
    """Local folder of the Hugging Face model `name` at commit `revision`, downloaded on first use, so that
    the model and its tokenizer load from the same pinned files. A local folder, or no revision, is
    returned unchanged."""
    if not revision or Path(name).is_dir():
        return name
    from huggingface_hub import snapshot_download
    return snapshot_download(name, revision=revision)


class _TrackedManager:
    """Memory manager over the evaluation's LLM client (replaces instance.manager).
    Counts the extraction calls that failed (LightMem goes on with no entries from them) in
    `extract_failures`, with the last error in `last_error`."""

    def __init__(self, client, model: str, temperature: float, max_tokens: int):
        self._client = client
        self.extract_failures = 0
        self.last_error = None
        self._failure_lock = threading.Lock()
        _ensure_engine_on_path()
        from lightmem.memory.prompts import (EXTRACTION_PROMPTS,
                                             METADATA_GENERATE_PROMPT, UPDATE_PROMPT)
        from lightmem.memory.utils import clean_response
        self.EXTRACTION_PROMPTS = EXTRACTION_PROMPTS
        self.METADATA_GENERATE_PROMPT = METADATA_GENERATE_PROMPT
        self.UPDATE_PROMPT = UPDATE_PROMPT
        self.clean_response = clean_response
        self.temperature = temperature
        self.max_tokens = max_tokens

        class _Cfg:
            pass
        self.config = _Cfg()
        self.config.model = model
        self.config.temperature = temperature
        self.config.max_tokens = max_tokens

    def generate_response(self, messages, response_format=None, tools=None,
                          tool_choice: str = "auto"):
        kwargs = {"response_format": response_format} if response_format else {}
        resp = self._client.chat(messages, temperature=self.temperature,
                                 max_tokens=self.max_tokens, **kwargs)
        usage = {"prompt_tokens": resp.input_tokens,
                 "completion_tokens": resp.output_tokens,
                 "total_tokens": resp.input_tokens + resp.output_tokens}
        return resp.content, usage

    def meta_text_extract(self, extract_list, messages_use="user_only",
                          topic_id_mapping=None, extraction_mode="flat",
                          custom_prompts=None):
        if not extract_list:
            return []
        default_prompts = self.EXTRACTION_PROMPTS.get(extraction_mode, {})
        prompts = default_prompts if custom_prompts is None else {**default_prompts,
                                                                  **custom_prompts}
        if extraction_mode == "flat":
            return self._extract_with_prompt(
                prompts.get("factual", self.METADATA_GENERATE_PROMPT),
                extract_list, messages_use, topic_id_mapping, "factual")
        if extraction_mode == "event":
            fact = self._extract_with_prompt(prompts["factual"], extract_list,
                                             messages_use, topic_id_mapping, "factual")
            rel = self._extract_with_prompt(prompts["relational"], extract_list,
                                            messages_use, topic_id_mapping, "relational")
            return self._merge_dual(fact, rel)
        raise ValueError(f"Unknown extraction_mode: {extraction_mode}")

    def _merge_dual(self, factual_results, relational_results):
        merged_results = []
        for factual, relational in zip(factual_results, relational_results):
            if factual is None and relational is None:
                merged_results.append(None)
                continue
            merged = {"input_prompt": [], "output_prompt": "", "cleaned_result": [],
                      "usage": {"prompt_tokens": 0, "completion_tokens": 0,
                                "total_tokens": 0}}
            for part in (factual, relational):
                if part is not None:
                    merged["input_prompt"].extend(part.get("input_prompt", []))
                    merged["cleaned_result"].extend(part.get("cleaned_result", []))
                    if part.get("usage"):
                        for key in merged["usage"]:
                            merged["usage"][key] += part["usage"].get(key, 0)
            merged["output_prompt"] = (
                f"Factual: {factual.get('output_prompt', 'N/A') if factual else 'N/A'}\n"
                f"Relational: {relational.get('output_prompt', 'N/A') if relational else 'N/A'}")
            merged_results.append(merged)
        return merged_results

    def _extract_with_prompt(self, system_prompt, extract_list, messages_use,
                             topic_id_mapping, entry_type="factual"):
        import concurrent.futures

        def concatenate_messages(segment, messages_use):
            role_filter = {"user_only": {"user"}, "assistant_only": {"assistant"},
                           "hybrid": {"user", "assistant"}}
            if messages_use not in role_filter:
                raise ValueError(f"Invalid messages_use value: {messages_use}")
            allowed = role_filter[messages_use]
            lines = []
            for mes in segment:
                if mes.get("role") in allowed:
                    seq = mes.get("sequence_number", 0)
                    time_prefix = ""
                    if mes.get("time_stamp") and mes.get("weekday"):
                        time_prefix = f"[{mes['time_stamp']}, {mes['weekday']}] "
                    who = mes.get("speaker_name") or mes["role"]
                    lines.append(f"{time_prefix}{seq // 2}.{who}: {mes.get('content', '')}")
            return "\n".join(lines)

        def process(args):
            api_call_idx, api_call_segments = args
            try:
                parts, gids = [], []
                if topic_id_mapping and api_call_idx < len(topic_id_mapping):
                    gids = topic_id_mapping[api_call_idx]
                for topic_idx, topic_segment in enumerate(api_call_segments):
                    gid = gids[topic_idx] if topic_idx < len(gids) else topic_idx + 1
                    parts.append(f"--- Topic {gid} ---\n"
                                 + concatenate_messages(topic_segment, messages_use))
                msgs = [{"role": "system", "content": system_prompt},
                        {"role": "user", "content": "\n".join(parts)}]
                raw, usage = self.generate_response(
                    messages=msgs, response_format={"type": "json_object"})
                facts = self.clean_response(raw)
                for entry in facts:
                    entry["entry_type"] = entry_type
                return {"input_prompt": msgs, "output_prompt": raw,
                        "cleaned_result": facts, "usage": usage,
                        "entry_type": entry_type}
            except Exception as e:
                logger.error(f"[LightMem] extraction call {api_call_idx} failed: {e}")
                with self._failure_lock:
                    self.extract_failures += 1
                    self.last_error = f"{type(e).__name__}: {e}"[:200]
                return {"input_prompt": [], "output_prompt": "", "cleaned_result": [],
                        "usage": None, "entry_type": entry_type}

        with concurrent.futures.ThreadPoolExecutor(
                max_workers=min(len(extract_list), 5)) as ex:
            return list(ex.map(process, enumerate(extract_list)))

    def _call_update_llm(self, system_prompt, target_entry, candidate_sources):
        import json as _json
        user_prompt = (f"Target memory:{target_entry['payload']['memory']}\n"
                       "Candidate memories:\n"
                       + "\n".join(f"- {c['payload']['memory']}"
                                   for c in candidate_sources))
        text, usage = self.generate_response(
            [{"role": "system", "content": system_prompt},
             {"role": "user", "content": user_prompt}],
            response_format={"type": "json_object"})
        try:
            result = _json.loads(text)
            if "action" not in result:
                result = {"action": "ignore"}
            result["usage"] = usage
            return result
        except Exception:
            return {"action": "ignore", "usage": usage}


class LightMemAgent(BaseAgent):
    ANSWERS_INTERNALLY = False
    MEMORIZE_WITH_LLM = True

    def __init__(self, *, model: str | None = None, retrieve_num: int,
                 pre_compress: bool, topic_segment: bool,
                 index_strategy: str, retrieve_strategy: str,
                 update_mode: str, extraction_mode: str,
                 messages_use: str,
                 max_user_tokens: int, shortmem_max_tokens: int,
                 temperature: float, max_tokens: int,
                 max_context_tokens: int, answer_max_tokens: int,
                 context_reserve_tokens: int, truncation_ratio: float,
                 embedding_model: str | None = None,
                 embedding_dim: int | None = None,
                 segmenter_model: str = SEGMENTER_MODEL,
                 segmenter_revision: str | None = SEGMENTER_REVISION,
                 compress_rate: float,
                 compress_device: str,
                 offline_consolidate: bool,
                 consolidate_top_k: int, consolidate_keep_top_n: int,
                 consolidate_score_threshold: float,
                 consolidate_max_calls: int,
                 consolidate_max_workers: int,
                 internal_base_url: str | None = None,
                 working_dir: str | None = None, **params):
        super().__init__(**params)
        if not working_dir:
            raise ValueError("LightMemAgent requires working_dir (the runner sets it per run)")
        self.model = model
        self.retrieve_num = retrieve_num
        self.max_context_tokens = max_context_tokens
        self.answer_max_tokens = answer_max_tokens
        self.max_user_tokens = max_user_tokens
        self.shortmem_max_tokens = shortmem_max_tokens
        self.context_reserve_tokens = context_reserve_tokens
        self.truncation_ratio = truncation_ratio
        self._consolidate = bool(offline_consolidate)
        self._cons_topk = consolidate_top_k
        self._vec_seen: set[str] = set()
        self._cons_keep = consolidate_keep_top_n
        self._cons_thresh = consolidate_score_threshold
        self._cons_max_calls = consolidate_max_calls
        self._cons_workers = consolidate_max_workers
        self._cfg = dict(pre_compress=pre_compress, topic_segment=topic_segment,
                         index_strategy=index_strategy,
                         retrieve_strategy=retrieve_strategy, update_mode=update_mode,
                         extraction_mode=extraction_mode, messages_use=messages_use,
                         embedding_model=embedding_model, embedding_dim=embedding_dim,
                         segmenter_model=segmenter_model, segmenter_revision=segmenter_revision,
                         compress_rate=compress_rate, compress_device=compress_device)
        self._working_dir = Path(working_dir)
        self._internal_client = create_llm_client(
            model=model, temperature=temperature,
            max_tokens=max_tokens, base_url=internal_base_url)
        self._manager = _TrackedManager(self._internal_client, model,
                                        temperature, max_tokens)
        self._instance = None
        self._id2visit: Dict[str, set] = {}
        self._cur_visit: int | None = None
        self._last_hits: list = []
        self._last_rows: list = []

    def _get_instance(self):
        if self._instance is not None:
            return self._instance
        _ensure_engine_on_path()
        from lightmem.memory.lightmem import LightMemory
        if not hasattr(LightMemory, "compressor"):
            LightMemory.compressor = None
        segmenter = _pinned_model_path(self._cfg["segmenter_model"], self._cfg["segmenter_revision"])
        qdrant_path = self._working_dir / "qdrant"
        if qdrant_path.exists():
            shutil.rmtree(qdrant_path)
        qdrant_path.mkdir(parents=True, exist_ok=True)
        emb_url = ENDPOINTS["embedding"]
        config = {
            "pre_compress": self._cfg["pre_compress"],
            "topic_segment": self._cfg["topic_segment"],
            "precomp_topic_shared": bool(self._cfg["pre_compress"]),
            "messages_use": self._cfg["messages_use"],
            "index_strategy": self._cfg["index_strategy"],
            "retrieve_strategy": self._cfg["retrieve_strategy"],
            "update": self._cfg["update_mode"],
            "extraction_mode": self._cfg["extraction_mode"],
            "metadata_generate": True,
            "text_summary": True,
            "memory_manager": {"model_name": "openai", "configs": {
                "model": self.model, "api_key": "injected",
                "temperature": self._manager.config.temperature,
                "max_tokens": self._manager.config.max_tokens}},
            **({"pre_compressor": {"model_name": "llmlingua-2", "configs": {
                "llmlingua_config": {
                    "model_name": segmenter,
                    "device_map": self._cfg["compress_device"],
                    "use_llmlingua2": True},
                "compress_config": {
                    "instruction": "", "rate": self._cfg["compress_rate"],
                    "target_token": -1}}}} if self._cfg["pre_compress"] else {}),
            "topic_segmenter": {"model_name": "llmlingua-2", "configs": {
                "model_name": segmenter, "buffer_len": 512,
                "model_config": {"attn_implementation": "eager"}}},
            "text_embedder": {"model_name": "openai", "configs": {
                "model": self._cfg["embedding_model"], "api_key": api_key_for(emb_url),
                "openai_base_url": emb_url, "embedding_dims": self._cfg["embedding_dim"]}},
            "embedding_retriever": {"model_name": "qdrant", "configs": {
                "collection_name": "lightmem_p",
                "embedding_model_dims": self._cfg["embedding_dim"],
                "path": str(qdrant_path)}},
        }
        instance = LightMemory.from_config(config)
        if hasattr(instance, "shortmem_buffer_manager"):
            instance.shortmem_buffer_manager.max_tokens = self.shortmem_max_tokens
        instance.manager = self._manager
        if getattr(instance.text_embedder, "client", None) is not None:
            import openai
            instance.text_embedder.client = openai.OpenAI(api_key=api_key_for(emb_url), base_url=emb_url,
                                                          max_retries=MAX_ATTEMPTS - 1, timeout=http_timeout())
        self._trace_extract(instance)
        self._hook_visit_tracking(instance)
        self._track_embedder(instance)
        self._instance = instance
        return instance

    def _track_embedder(self, instance) -> None:
        from ..llm_client import USAGE
        emb, orig, model = instance.text_embedder, instance.text_embedder.embed, \
            self._cfg["embedding_model"]

        def tracked_embed(text):
            before = int(getattr(emb, "total_tokens", 0) or 0)
            vec = orig(text)
            USAGE.log(model, "embed",
                      max(0, int(getattr(emb, "total_tokens", 0) or 0) - before), 0)
            return vec

        instance.text_embedder.embed = tracked_embed

    @staticmethod
    def _point_count(lm) -> int:
        try:
            r = lm.embedding_retriever
            return int(r.client.count(collection_name=r.collection_name,
                                      exact=True).count)
        except Exception:
            return 0

    def _split_turn(self, content: str, max_tokens: int) -> List[str]:
        if not content:
            return [""]
        if self.count_tokens(content) <= max_tokens:
            return [content]
        chunks: List[str] = []
        marked = content
        for ch in (". ", "! ", "? "):
            marked = marked.replace(ch, ch.strip() + "\n")
        current, current_tok = "", 0
        for sentence in (s.strip() for s in marked.split("\n")):
            if not sentence:
                continue
            tok = self.count_tokens(sentence)
            if tok > max_tokens:
                if current:
                    chunks.append(current.strip())
                    current, current_tok = "", 0
                remaining = sentence
                while remaining:
                    part = self._find_max_chunk(remaining, max_tokens) or remaining[:100]
                    if len(part) < len(remaining):
                        chunks.append(part.strip())
                        remaining = remaining[len(part):]
                    else:
                        current, current_tok = part, self.count_tokens(part)
                        break
            elif current_tok + tok > max_tokens:
                if current:
                    chunks.append(current.strip())
                current, current_tok = sentence, tok
            else:
                current += " " + sentence if current else sentence
                current_tok += tok
        if current.strip():
            chunks.append(current.strip())
        return chunks or [content[:1000]]

    def _find_max_chunk(self, text: str, max_tokens: int) -> str:
        if not text or self.count_tokens(text) <= max_tokens:
            return text
        left, right, result = 1, len(text), ""
        while left <= right:
            mid = (left + right) // 2
            if self.count_tokens(text[:mid]) <= max_tokens:
                result, left = text[:mid], mid + 1
            else:
                right = mid - 1
        if result and len(result) < len(text):
            for i in range(len(result) - 1, max(0, len(result) - 20), -1):
                if result[i] in " \t\n.,!?;:":
                    return result[:i + 1]
        return result

    @staticmethod
    def _session_timestamp(date: str) -> str:
        d = datetime.strptime(date[:10], "%Y-%m-%d")
        return f"{d.strftime('%Y/%m/%d')} ({d.strftime('%a')}) 00:00"

    def _format_messages(self, text: str, timestamp: str) -> List[Dict[str, Any]]:
        turns: List[tuple] = []
        role, buf = None, []
        for line in (l.strip() for l in text.split("\n")):
            if not line:
                continue
            if line.startswith("Patient:"):
                if role and buf:
                    turns.append((role, "\n".join(buf).strip()))
                role, buf = "user", [line[len("Patient:"):].strip()]
            elif line.startswith("Doctor:"):
                if role and buf:
                    turns.append((role, "\n".join(buf).strip()))
                role, buf = "assistant", [line[len("Doctor:"):].strip()]
            elif role is not None:
                buf.append(line)
        if role and buf:
            turns.append((role, "\n".join(buf).strip()))
        turns = [(r, c) for r, c in turns if c]
        if not turns:
            turns = [("user", text[:5000])]

        def msg(r, content):
            who = "Patient" if r == "user" else "Doctor"
            return {"time_stamp": timestamp, "role": r, "content": content,
                    "speaker_id": who.lower(), "speaker_name": who}

        messages: List[Dict[str, Any]] = []
        i = 0
        while i < len(turns):
            r, content = turns[i]
            if r == "user":
                paired = i + 1 < len(turns) and turns[i + 1][0] == "assistant"
                assistant = turns[i + 1][1] if paired else ""
                if self.count_tokens(content) > self.max_user_tokens:
                    parts = self._split_turn(content, self.max_user_tokens)
                    for j, part in enumerate(parts):
                        messages.append(msg("user", part))
                        messages.append(msg("assistant",
                                            assistant if j == len(parts) - 1 else ""))
                else:
                    messages.append(msg("user", content))
                    messages.append(msg("assistant", assistant))
                i += 2 if paired else 1
            else:
                messages.append(msg("user", ""))
                messages.append(msg("assistant", content))
                i += 1
        return messages

    def memorize(self, text: str, visit_number: int, date: str) -> MemoryBuildResult:
        """Add the visit turn by turn. Not successful when a turn raised, an extraction call failed, or no
        entry was stored."""
        t0 = time.time()
        failures_before = self._manager.extract_failures
        try:
            lm = self._get_instance()
            timestamp = self._session_timestamp(date)
            messages = self._format_messages(text, timestamp)
        except Exception as e:
            logger.error(f"[LightMem] setup failed on v{visit_number}: {e}")
            return MemoryBuildResult(success=False, method="lightmem",
                                     action="add_memory", visit_number=visit_number,
                                     time_cost=time.time() - t0,
                                     extra={"error": str(e)[:200]})
        self._cur_visit = visit_number
        total_turns = len(messages) // 2
        api_calls, errors = 0, 0
        points_before = self._point_count(lm)
        for turn_idx in range(total_turns):
            turn = messages[turn_idx * 2: turn_idx * 2 + 2]
            is_last = turn_idx == total_turns - 1
            try:
                result = lm.add_memory(messages=turn, force_segment=is_last,
                                       force_extract=is_last)
                api_calls += result.get("api_call_nums", 0) if isinstance(result, dict) else 0
            except Exception as e:
                errors += 1
                logger.error(f"[LightMem] add_memory turn {turn_idx} v{visit_number}: {e}")
        try:
            buf = lm.senmem_buffer_manager
            if buf.buffer:
                buf.buffer.clear()
                buf.token_count = 0
        except Exception:
            pass
        entries = max(0, self._point_count(lm) - points_before)
        stats = {}
        try:
            stats = lm.get_token_statistics().get("summary", {})
        except Exception:
            pass
        extract_failures = self._manager.extract_failures - failures_before
        if errors or extract_failures or not entries:
            logger.warning("[LightMem] v%s stored %d entries from %d turns with %d failed "
                           "turn(s) and %d failed extraction call(s); this visit's content is %s",
                           visit_number, entries, total_turns, errors, extract_failures,
                           "missing from memory" if not entries else "incomplete")
        extra = {"total_turns": total_turns, "api_calls": api_calls,
                 "extracted_entries": entries, "turn_errors": errors,
                 "engine_llm_calls_total": stats.get("total_llm_calls")}
        if extract_failures:
            extra.update(extract_failures=extract_failures, error=self._manager.last_error)
        return MemoryBuildResult(
            success=(errors == 0 and extract_failures == 0 and entries > 0), method="lightmem",
            action="add_memory", visit_number=visit_number, added_chunks=entries,
            time_cost=time.time() - t0, extra=extra)

    def build_memory_context(self, question: str, as_of: str = "",
                             display_question: str | None = None) -> MemoryContext:
        try:
            lm = self._get_instance()
            related = lm.retrieve(query=question, limit=self.retrieve_num)
            if isinstance(related, list):
                related = "\n".join(str(x) for x in related)
        except Exception as e:
            logger.error(f"[LightMem] retrieval failed: {e}")
            return MemoryContext(mem="", n_tok=0, retrieved_count=0, error=True,
                                 extra={"method": "lightmem", "error": str(e)[:200]})
        mem = (related or "").strip()
        lines = [l for l in mem.split("\n") if l.strip()]
        truncated = False
        budget = self.max_context_tokens - self.answer_max_tokens - self.context_reserve_tokens \
            - self.count_tokens(question)
        if mem and self.count_tokens(mem) > budget > 0:
            ratio = budget / self.count_tokens(mem)
            mem = mem[:int(len(mem) * ratio * self.truncation_ratio)] + "\n... [truncated]"
            truncated = True
        visits: set = set()
        for mid in self._last_hits:
            visits.update(self._id2visit.get(str(mid), ()))
        try:
            self._trace("query_retrieve", question=question, limit=self.retrieve_num,
                        truncated=truncated,
                        hits=[{**r, "visits": sorted(self._id2visit.get(r["id"], ()))}
                              for r in (self._last_rows or [])])
        except Exception:
            pass
        return MemoryContext(
            mem=mem, n_tok=self.count_tokens(mem), retrieved_count=len(lines),
            retrieved_visits=sorted(visits), context_truncated=truncated,
            retrieved_memories=[{"visit": None, "memory": l[:300]} for l in lines[:5]],
            extra={"method": "lightmem"})

    def dump_memory(self, out_dir):
        from .base import snapshot_jsonl_gz, snapshot_vectors_delta
        if self._instance is None:
            return None
        r = self._instance.embedding_retriever
        rows, vecs, offset = [], [], None
        while True:
            points, offset = r.client.scroll(collection_name=r.collection_name,
                                             limit=1000, offset=offset,
                                             with_payload=True, with_vectors=True)
            rows.extend({"id": str(pt.id),
                         "visits": sorted(self._id2visit.get(str(pt.id), ())),
                         **(pt.payload or {})} for pt in points)
            vecs.extend((str(pt.id), pt.vector) for pt in points if pt.vector is not None)
            if offset is None:
                break
        n = snapshot_jsonl_gz(out_dir / "entries.jsonl.gz", rows)
        nv = 0
        try:
            nv = snapshot_vectors_delta(out_dir, vecs, self._vec_seen)
        except Exception as e:
            logger.error("[LightMem] vector snapshot failed (run continues): %s", e)
        return {"entries": n, "vectors_new": nv}

    def _hook_visit_tracking(self, instance) -> None:
        """Record entry-id -> visit at write, and the ids returned by each search."""
        er = getattr(instance, "embedding_retriever", None)
        if er is None or getattr(er, "__memebench_visit_tracking__", False):
            return
        agent = self

        orig_insert = er.insert
        def insert(vectors, payloads=None, ids=None, *a, **kw):
            out = orig_insert(vectors, payloads=payloads, ids=ids, *a, **kw)
            try:
                if agent._cur_visit is not None:
                    for i in (ids or []):
                        agent._id2visit.setdefault(str(i), set()).add(int(agent._cur_visit))
            except Exception:
                pass
            return out
        er.insert = insert

        orig_search = er.search
        def search(*a, **kw):
            out = orig_search(*a, **kw)
            try:
                rows = [r for r in (out or [])
                        if isinstance(r, dict) and r.get("id") is not None]
                agent._last_hits = [str(r["id"]) for r in rows]
                agent._last_rows = [
                    {"id": str(r["id"]),
                     "score": r.get("score"),
                     "memory": str((r.get("payload") or {}).get("memory", ""))[:300]}
                    for r in rows]
            except Exception:
                agent._last_hits, agent._last_rows = [], []
            return out
        er.search = search
        er.__memebench_visit_tracking__ = True

    def _trace_extract(self, instance) -> None:
        """Record the entries each extraction call produced."""
        mgr = getattr(instance, "manager", None)
        if mgr is None or getattr(mgr, "__traced__", False):
            return
        orig = mgr.meta_text_extract

        def traced(*a, **kw):
            out = orig(*a, **kw)
            try:
                _o = as_list(out)
                self._trace("extract", n=len(_o), entries=_o)
            except Exception:
                pass
            return out

        mgr.meta_text_extract = traced
        mgr.__traced__ = True

    def consolidate(self) -> dict | None:
        """Offline update at each checkpoint when `offline_consolidate` is set.

        LightMem builds the update queues. Every entry that is in another entry's queue with a score of at least
        `consolidate_score_threshold` gets one update call with LightMem's UPDATE_PROMPT (at most
        `consolidate_max_calls` calls), and is deleted or updated as the reply says. When update calls failed, the
        result carries an `error` entry.
        """
        if not self._consolidate:
            return None
        t0 = time.time()
        lm = self._get_instance()
        try:
            from lightmem.memory.prompts import UPDATE_PROMPT
        except Exception as e:
            return {"error": f"UPDATE_PROMPT import failed: {e}"}
        lm.construct_update_queue_all_entries(top_k=self._cons_topk,
                                              keep_top_n=self._cons_keep)
        entries = lm.embedding_retriever.get_all()
        if not entries:
            return {"action": "consolidate", "entries": 0}
        back: dict[str, list] = {}
        for other in entries:
            for cand in (other.get("payload", {}) or {}).get("update_queue", []) or []:
                if cand.get("score", 0) >= self._cons_thresh:
                    back.setdefault(cand.get("id"), []).append(other)
        from concurrent.futures import ThreadPoolExecutor
        lock = threading.Lock()
        write_lock = threading.Lock()
        stats = {"upd": 0, "dele": 0, "calls": 0, "err": 0, "capped": 0}

        def one(entry):
            eid = entry["id"]
            sources = back.get(eid)
            if not sources:
                return
            with lock:
                if stats["calls"] >= self._cons_max_calls:
                    stats["capped"] += 1
                    return
                stats["calls"] += 1
            try:
                out = lm.manager._call_update_llm(UPDATE_PROMPT, entry, sources)
            except Exception:
                with lock:
                    stats["err"] += 1
                return
            if not out:
                return
            act = out.get("action")
            try:
                self._trace("consolidate_op", op=act, memory_id=str(eid),
                            visits_before=sorted(self._id2visit.get(str(eid), ())),
                            target=str((entry.get("payload") or {}).get("memory", ""))[:300],
                            new_memory=str(out.get("new_memory") or "")[:300],
                            sources=[{"id": str(c.get("id")),
                                      "visits": sorted(self._id2visit.get(str(c.get("id")), ())),
                                      "text": str((c.get("payload") or {}).get("memory", ""))[:200]}
                                     for c in (sources or [])][:10])
            except Exception:
                pass
            if act == "delete":
                with write_lock:
                    lm.embedding_retriever.delete(eid)
                with lock:
                    self._id2visit.pop(str(eid), None)
                    stats["dele"] += 1
            elif act == "update":
                pay = dict(entry.get("payload") or {})
                pay["memory"] = out.get("new_memory")
                with write_lock:
                    lm.embedding_retriever.update(vector_id=eid, vector=entry.get("vector"),
                                                  payload=pay)
                with lock:
                    merged = self._id2visit.setdefault(str(eid), set())
                    for src in sources:
                        merged |= self._id2visit.get(str(src.get("id")), set())
                    stats["upd"] += 1

        with ThreadPoolExecutor(max_workers=self._cons_workers) as ex:
            list(ex.map(one, entries))
        upd, dele, calls, err, capped = (stats["upd"], stats["dele"], stats["calls"],
                                         stats["err"], stats["capped"])
        out = {"action": "consolidate", "entries": len(entries),
               "with_candidates": len(back), "llm_calls": calls,
               "updated": upd, "deleted": dele, "errors": err,
               "capped_out": capped, "cap": self._cons_max_calls,
               "time_cost": round(time.time() - t0, 1)}
        if err:
            out["error"] = f"{err} of {calls} update call(s) failed"
        return out