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