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