"""MemOS (v2.0.11) tree text memory adapter.""" from __future__ import annotations import gzip import json import logging import os import re import sys import threading import time from collections import Counter from contextlib import nullcontext from pathlib import Path from types import SimpleNamespace from .base import BaseAgent, MemoryBuildResult, MemoryContext from .memos_adapters import ( ExtractionError, ObservedProxy, TrackedEmbedder, TrackedLLM, as_dict, make_strict_reranker, make_validated_reader, normalized_sources, render_memory_items, structured_messages, ) from ..llm_client import HTTP_TIMEOUT_S, MAX_ATTEMPTS, OpenAIEmbedder, create_llm_client from .memos_recording import Recorder, RecordingError PIN = "27c9e71912e33a2f8c6bea6d79c08dbae282e4b8" DEFAULT_SOURCE = Path(__file__).resolve().parents[2] / "vendor/memos/src" class MemOSAgent(BaseAgent): METHOD_NAME = "memos" ANSWERS_INTERNALLY = False MEMORIZE_WITH_LLM = True FAIL_CLOSED_ON_MEMORY_ERROR = True def __init__(self, *, model=None, retrieve_num, working_memory_size, embedding_model=None, embedding_dim=None, embedding_batch_size=None, extractor_temperature, extractor_max_tokens, extraction_attempts=MAX_ATTEMPTS, chat_window_tokens, chat_window_overlap_tokens, reader_workers, internal_base_url=None, working_dir=None, source_path, neo4j_uri, neo4j_user, neo4j_password=None, neo4j_password_file=None, neo4j_db_name, qdrant_url, reranker_url, reranker_model="bge-reranker-v2-m3", reranker_timeout=HTTP_TIMEOUT_S, relativity, max_memory_tokens, llm_client=None, embedder=None, **params): super().__init__(**params) unsupported = sorted(params) if unsupported: raise ValueError("Unsupported memos parameters: " + ", ".join(unsupported)) if not working_dir: raise ValueError("memos requires a per-run working_dir") if int(retrieve_num) < 1: raise ValueError("memos retrieve_num must be positive") self.working_dir = Path(working_dir).resolve() self.working_dir.mkdir(parents=True, exist_ok=True) self.source_path = Path(source_path or DEFAULT_SOURCE).resolve() self.model, self.retrieve_num = model, int(retrieve_num) self.working_memory_size = int(working_memory_size) self.embedding_model, self.embedding_dim = embedding_model, int(embedding_dim) self.max_memory_tokens = int(max_memory_tokens) self.extraction_attempts, self.chat_window_tokens = int(extraction_attempts), int(chat_window_tokens) self.chat_window_overlap_tokens = int(chat_window_overlap_tokens) self.reader_workers = int(reader_workers) if min(self.extraction_attempts, self.chat_window_tokens, self.working_memory_size, self.reader_workers) < 1: raise ValueError("Extraction attempts, chat window, reader workers, and working capacity " "must be positive") self.relativity = float(relativity) self.neo4j_uri = neo4j_uri or os.environ.get("MEMEBENCH_MEMOS_NEO4J_URI", "bolt://127.0.0.1:17687") self.neo4j_user, self.neo4j_db_name = neo4j_user, neo4j_db_name self.neo4j_password = neo4j_password or os.environ.get("MEMEBENCH_MEMOS_NEO4J_PASSWORD", "") if not self.neo4j_password and neo4j_password_file: self.neo4j_password = Path(neo4j_password_file).expanduser().read_text().strip() self.qdrant_url = qdrant_url or os.environ.get("MEMEBENCH_MEMOS_QDRANT_URL", "http://127.0.0.1:16333") self.reranker_url = (reranker_url or os.environ.get("MEMEBENCH_MEMOS_RERANKER_URL", "http://127.0.0.1:18086/rerank")) self.reranker_model, self.reranker_timeout = reranker_model, float(reranker_timeout) self._errors = [] self.recorder = Recorder(self.working_dir.parent) self.recorder.add_secret(self.neo4j_password) client = llm_client or create_llm_client(model=model, temperature=extractor_temperature, max_tokens=extractor_max_tokens, base_url=internal_base_url) self.recorder.observe_client(client, stage="internal") self.llm = TrackedLLM(client, self._record_error) self.llm.recorder = self.recorder self.embedder = TrackedEmbedder( embedder or OpenAIEmbedder(model=embedding_model, batch_size=embedding_batch_size), self.embedding_dim, self._record_error) self.embedder.recorder = self.recorder self._initialized = False self._closed = False self._stored_ids = set() self._visits = {} self._last_visit = None self._last_date = None self._trace_log = (self.working_dir / "trace.jsonl").open("a", encoding="utf-8") self._event_lock = threading.Lock() self.resolved_config = { "source_commit": PIN, "profile": "memos_text_architecture", "reader": "simple_struct", "extraction_mode": "fine", "search_mode": "fast", "retrieve_num": self.retrieve_num, "working_memory_size": self.working_memory_size, "embedding_model": embedding_model, "embedding_dim": self.embedding_dim, "model": model, "extractor_max_tokens": extractor_max_tokens, "extractor_temperature": extractor_temperature, "extraction_attempts": self.extraction_attempts, "chat_window_tokens": self.chat_window_tokens, "reranker_model": reranker_model, "reader_workers": self.reader_workers, "relativity": self.relativity, "dedup": "mmr", "reorganize": False, "search_strategy": {"bm25": False, "cot": False, "fast_graph": False, "fulltext": False}, "graph_backend": "neo4j-community", "query_working_memory": "scheduler_isolated", "source_date_rendering": "recorded_on", "answer_owner": "benchmark_runner", } def _event(self, event, **fields): try: record = self.recorder.envelope(event, **{"visit_number": self._last_visit, **fields}) with self._event_lock: self._trace_log.write(json.dumps(record, ensure_ascii=False, default=str) + "\n") self._trace_log.flush() self._trace(event, **{key: value for key, value in record.items() if key != "event"}) return record except Exception as exc: raise RecordingError(f"MemOS algorithm trace failed: {type(exc).__name__}") from exc def _scope(self, **fields): recorder = getattr(self, "recorder", None) return recorder.scope(**fields) if recorder is not None else nullcontext() def set_context_id(self, context_id): super().set_context_id(context_id) self.recorder.bind_context(patient_id=self._context_id) def _record_error(self, stage, error): self._errors.append({"stage": stage, "error": str(error)[:300]}) def _assert_healthy(self): if self._errors: raise RuntimeError("MemOS dependency failure: " + json.dumps(self._errors[:3])) def _load_source(self): if not (self.source_path / "memos/__init__.py").is_file(): raise ImportError(f"Pinned MemOS source missing: {self.source_path}") loaded = sys.modules.get("memos") if loaded is not None and not Path(loaded.__file__).resolve().is_relative_to(self.source_path): raise RuntimeError("A different MemOS source is already imported; run memos in a fresh process") os.environ["MEMOS_BASE_PATH"] = str(self.working_dir) if str(self.source_path) not in sys.path: sys.path.insert(0, str(self.source_path)) self.source_identity = {"repo": "https://github.com/MemTensor/MemOS", "tag": "v2.0.11", "commit": PIN} self.recorder.write_manifest(config=self.resolved_config, source_identity=self.source_identity, extra={"patient_id": self._context_id}) @staticmethod def _quiet_engine_logging(root, handlers, level): for name in ("memos", "SearchHandler"): logger = logging.getLogger(name) logger.handlers[:] = [logging.NullHandler()] logger.setLevel(logging.WARNING) logger.propagate = False root.handlers[:] = handlers root.setLevel(level) def _initialize_engine(self): if self._initialized: return if self._closed: raise RuntimeError("Agent is closed") if not self._context_id: raise ValueError("Set the patient context_id before MemOS memory initialization") if not self.neo4j_password: raise ValueError("Set MEMEBENCH_MEMOS_NEO4J_PASSWORD or neo4j_password") self._load_source() root = logging.getLogger() handlers, level = list(root.handlers), root.level try: from memos.configs.graph_db import Neo4jCommunityGraphDBConfig from memos.graph_dbs.neo4j_community import Neo4jCommunityGraphDB from memos.memories.textual.simple_tree import SimpleTreeTextMemory from memos.memories.textual.tree_text_memory.organize.manager import MemoryManager from memos.api.handlers.search_handler import SearchHandler from .memos_working import WorkingSet run_name = f"{self._context_id}_{Path(self.working_dir).parent.name}" self.namespace = "memebench_" + re.sub(r"[^A-Za-z0-9_]", "_", run_name) graph_config = Neo4jCommunityGraphDBConfig( uri=self.neo4j_uri, user=self.neo4j_user, password=self.neo4j_password, db_name=self.neo4j_db_name, auto_create=False, use_multi_db=False, user_name=self.namespace, embedding_dimension=self.embedding_dim, vec_config={"backend": "qdrant", "config": {"url": self.qdrant_url, "collection_name": self.namespace, "vector_dimension": self.embedding_dim, "distance_metric": "cosine"}}) class NormalizedCommunityGraph(Neo4jCommunityGraphDB): def delete_node(inner, id, user_name=None): effective = user_name or inner.config.user_name existing = inner.get_node(id, user_name=effective) if existing is None: return super().delete_node(id, user_name=effective) inner.vec_db.delete([id]) self._event("delete", namespace=effective, removed_ids=[id], removed=[existing], before_ids=[id], after_ids=[]) def remove_oldest_memory(inner, memory_type, keep_latest, user_name=None): effective = user_name or inner.config.user_name before_items = inner.get_all_memory_items(memory_type, user_name=effective) before = {item["id"] for item in before_items} super().remove_oldest_memory(memory_type, keep_latest, user_name=effective) after = {item["id"] for item in inner.get_all_memory_items(memory_type, user_name=effective)} if before - after: inner.vec_db.delete(sorted(before - after)) self._event("eviction", namespace=effective, memory_type=memory_type, keep_latest=keep_latest, before_ids=sorted(before), after_ids=sorted(after), removed_ids=sorted(before - after), removed=[item for item in before_items if item["id"] not in after]) def _parse_node(inner, node_data): node = super()._parse_node(node_data) node["metadata"]["sources"] = normalized_sources(node["metadata"]) return node def _parse_nodes(inner, nodes_data): nodes = super()._parse_nodes(nodes_data) for node in nodes: node["metadata"]["sources"] = normalized_sources(node["metadata"]) return nodes self.graph = NormalizedCommunityGraph(graph_config) self.graph.vec_db = ObservedProxy(self.graph.vec_db, self._record_error, "qdrant") self.graph.clear(user_name=self.namespace) self._assert_healthy() self.reranker = make_strict_reranker(self.reranker_url, self.reranker_model, self.reranker_timeout, error_sink=self._record_error, trace=self._event, recorder=self.recorder) self.reader = make_validated_reader(self.llm, self.embedder, self.extraction_attempts, self.chat_window_tokens, self._event, reader_workers=self.reader_workers, window_overlap_tokens=self.chat_window_overlap_tokens, recorder=self.recorder) self.reader.set_graph_db(self.graph) self.manager = MemoryManager(self.graph, self.embedder, self.llm, memory_size={"WorkingMemory": self.working_memory_size, "LongTermMemory": 1500, "UserMemory": 480}, is_reorganize=False) config = SimpleNamespace(mode="sync", search_strategy=self.resolved_config["search_strategy"], internet_retriever=None, memory_filename="graph.json") self.tree = SimpleTreeTextMemory(self.llm, self.embedder, self.reader, self.graph, self.reranker, self.manager, config, include_embedding=True) self.searcher = self.tree.get_searcher(manual_close_internet=True, process_llm=self.llm) self.reader.set_searcher(self.searcher) self.product = SearchHandler.__new__(SearchHandler) self.product.deps = SimpleNamespace(searcher=self.searcher, reranker=self.reranker, embedder=self.embedder) self.product.logger = logging.getLogger("memos.product") self.working = WorkingSet(self.tree, self.llm, self.embedder, self.namespace, capacity=self.working_memory_size, trace=self._event, recorder=self.recorder) self._initialized = True (self.working_dir / "resolved_config.json").write_text(json.dumps({ **self.resolved_config, "namespace": self.namespace, "source_identity": self.source_identity}, indent=2), encoding="utf-8") finally: self._quiet_engine_logging(root, handlers, level) def memorize_session(self, session, rendered_text=None): recorder = getattr(self, "recorder", None) reference = recorder.record_visit(session) if recorder is not None else {} with self._scope(stage="memorize", visit_number=int(session["visit_number"]), date=str(session["date"]), **{key: value for key, value in reference.items() if key != "visit_number"}): return self._memorize_session(session, rendered_text) def _memorize_session(self, session, rendered_text=None): started = time.monotonic() visit, date = int(session["visit_number"]), str(session["date"]) self._last_visit, self._last_date = visit, date extracted = [] extraction_event = None self._errors = [] try: self._initialize_engine() if visit in self._visits: raise ValueError(f"Visit {visit} was already ingested") if self._visits and date < max(self._visits.values()): raise ValueError("Visits must be ingested chronologically") messages = structured_messages(session) extracted = self.reader.read_messages(messages, { "user_id": self.namespace, "session_id": f"{self.namespace}:visit:{visit}", "patient_id": self._context_id, "visit_number": visit, "date": date, "department": str(session.get("department") or "")}) for item in extracted: item.metadata.visit_number = visit item.metadata.date = date item.metadata.origin_memory_id = item.id extraction_event = self._event("extract", produced=len(extracted), entries=[as_dict(item) for item in extracted]) before_ids = sorted(self._stored_ids) if extracted: ids = self.tree.add(extracted, user_name=self.namespace) self._assert_healthy() if set(ids) != {item.id for item in extracted}: raise RuntimeError("MemOS writer did not acknowledge every extracted memory") for item in extracted: stored = self.graph.get_node(item.id, include_embedding=True, user_name=self.namespace) if (stored is None or stored["metadata"].get("memory_type") != item.metadata.memory_type or not stored["metadata"].get("embedding")): raise RuntimeError("Persistent graph/vector/type verification failed") self._stored_ids.update(ids) self._assert_healthy() self._visits[visit] = date self._event("store_commit", committed=len(extracted), stored_memories=len(self._stored_ids), types=dict(Counter(item.metadata.memory_type for item in extracted)), committed_ids=[item.id for item in extracted], before_ids=before_ids, after_ids=sorted(self._stored_ids), extraction_event_id=(extraction_event or {}).get("event_id")) return MemoryBuildResult(success=True, method="memos", action="extract_add", visit_number=visit, chunk_count=len(self._stored_ids), added_chunks=len(extracted), time_cost=time.monotonic() - started, extra={"extracted_count": len(extracted), "outcome": "valid_nonempty" if extracted else "valid_empty", "attempts": self.reader.attempts, "window_outcomes": self.reader.outcomes, "reader_workers": self.reader_workers, "scene_count": self.reader.scene_count, "scene_workers": min(self.reader_workers, self.reader.scene_count)}) except Exception as exc: rollback_error = None if extracted and self._initialized: try: self._delete_ids([item.id for item in extracted], extraction_event_id=(extraction_event or {}).get("event_id")) except Exception as rollback: rollback_error = str(rollback)[:300] self._event("memorize_failure", error=str(exc)[:300], kind=getattr(exc, "kind", "persistence_or_api_error"), rollback_error=rollback_error) return MemoryBuildResult(success=False, method="memos", action="extract_add", visit_number=visit, chunk_count=len(self._stored_ids), time_cost=time.monotonic() - started, extra={"error": str(exc)[:300], "outcome": getattr(exc, "kind", "persistence_or_api_error"), "rollback_error": rollback_error, "scene_diagnostics": getattr(getattr(self, "reader", None), "scene_diagnostics", [])}) def memorize(self, text, visit_number, date): return self.memorize_session({"visit_number": visit_number, "date": date, "department": "", "turns": [{"speaker": "Patient record", "text": text}]}) def _delete_ids(self, ids, extraction_event_id=None): if not ids: return with self.graph.driver.session(database=self.graph.db_name) as session: session.run("MATCH (n:Memory) WHERE n.user_name=$namespace AND n.id IN $ids DETACH DELETE n", namespace=self.namespace, ids=list(ids)).consume() self.graph.vec_db.delete(list(ids)) self._event("store_rollback_delete", removed_ids=list(ids), extraction_event_id=extraction_event_id) def _product_search(self, question): items = self.searcher.search(question, top_k=3 * self.retrieve_num, info={"user_id": self.namespace}, mode="fast", memory_type="All", user_name=self.namespace, dedup="mmr", search_tool_memory=False, include_skill_memory=False, include_preference_memory=False) self._assert_healthy() records = [as_dict(item) for item in items] self._event("candidates", requested=3 * self.retrieve_num, returned=len(records), hits=[{"id": item["id"], "type": item["metadata"].get("memory_type"), "score": item["metadata"].get("relativity")} for item in records]) results = {"text_mem": [{"cube_id": self.namespace, "memories": records, "total_nodes": len(records)}], "pref_mem": []} results = self.product._apply_relativity_threshold(results, self.relativity) results = self.product._mmr_dedup_text_memories(results, self.retrieve_num, 0) selected = [item for bucket in results["text_mem"] for item in bucket["memories"]] if len(selected) > self.retrieve_num: raise RuntimeError(f"MemOS product MMR exceeded final k={self.retrieve_num}") return selected def build_memory_context(self, question, as_of="", display_question=None): with self._scope(stage="query", visit_number=getattr(self, "_last_visit", None), as_of=as_of): return self._build_memory_context(question, as_of, display_question) def _build_memory_context(self, question, as_of="", display_question=None): started = time.monotonic() self._errors = [] working_cleaned = False try: self._initialize_engine() if as_of and self._visits and max(self._visits.values()) > as_of: raise ValueError("Store contains visits after this checkpoint") self.working.reset() working_diag = self.working.refresh(question) self._assert_healthy() working_items = self.tree.get_working_memory(user_name=self.namespace) self._event("working_memory", count=len(working_items), diagnostics=working_diag, entries=[as_dict(item) for item in working_items]) if len(working_items) > self.working_memory_size: raise RuntimeError("Working memory exceeds configured capacity") selected = self._product_search(question) self._assert_healthy() mem, visits, previews = render_memory_items(selected) n_tok = self.count_tokens(mem) if n_tok > self.max_memory_tokens: raise ValueError("Selected memory context exceeds configured token budget") self._event("query_retrieve", question=question, as_of=as_of, final_k=self.retrieve_num, returned=len(selected), n_tok=n_tok, hits=previews, seconds=time.monotonic() - started) return MemoryContext(mem=mem, n_tok=n_tok, retrieved_count=len(selected), retrieved_visits=visits, retrieved_memories=previews, answer=None, extra={"method": "memos", "working_count": len(working_items), "stored_memories": len(self._stored_ids)}) except Exception as exc: from .memos_working import WorkingMemoryError recoverable = ((isinstance(exc, WorkingMemoryError) and getattr(exc, "recoverable_output", False) is True) or (isinstance(exc, ExtractionError) and exc.kind in {"truncated", "empty_response"})) error_detail = {"type": type(exc).__name__, "recoverable_question": False} if recoverable and self._initialized: self.working.reset() self._assert_healthy() working_cleaned = True error_detail["recoverable_question"] = True self._event("query_failure", error=str(exc)[:300], error_detail=error_detail) return MemoryContext(mem="", n_tok=0, retrieved_count=0, error=True, extra={"method": "memos", "error": str(exc)[:300], "error_detail": error_detail}) finally: if self._initialized and not working_cleaned: self.working.reset() def consolidate(self): with self._scope(stage="consolidate", visit_number=getattr(self, "_last_visit", None)): if self._initialized: self.manager.wait_reorganizer() self._assert_healthy() return None def dump_memory(self, out_dir): with self._scope(stage="snapshot", visit_number=getattr(self, "_last_visit", None)): return self._dump_memory(out_dir) def _dump_memory(self, out_dir): self._initialize_engine() self.working.reset() self.consolidate() out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) graph = self.graph.export_graph(include_embedding=True, user_name=self.namespace, status=["activated", "resolving", "archived", "deleted"]) vectors = [item.to_dict() if hasattr(item, "to_dict") else as_dict(item) for item in self.graph.vec_db.get_all()] self._assert_healthy() graph_ids = {node["id"] for node in graph["nodes"]} vector_ids = {item["id"] for item in vectors} if graph_ids != vector_ids: raise RuntimeError("Snapshot graph/vector ID sets differ") with gzip.open(out_dir / "graph.json.gz", "wt", encoding="utf-8") as handle: json.dump(graph, handle, ensure_ascii=False) with gzip.open(out_dir / "vectors.json.gz", "wt", encoding="utf-8") as handle: json.dump(vectors, handle, ensure_ascii=False) manifest = {"format": "memos_checkpoint", "source_commit": PIN, "resolved_config": self.resolved_config, "namespace": self.namespace, "patient_id": self._context_id, "visits": self._visits, "last_visit": self._last_visit, "last_date": self._last_date, "nodes": len(graph_ids), "vectors": len(vectors), "edges": len(graph.get("edges", [])), "types": dict(Counter(node["metadata"]["memory_type"] for node in graph["nodes"]))} (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8") return {key: manifest[key] for key in ("nodes", "vectors", "edges", "types")} def close(self): with self._scope(stage="close", visit_number=getattr(self, "_last_visit", None)): return self._close() def _close(self): if self._closed: return errors = [] def release(label, function): if callable(function): try: function() except Exception as exc: errors.append((label, exc)) working = getattr(self, "working", None) release("working_memory", getattr(working, "close", None) or getattr(working, "reset", None)) reorganizer = getattr(getattr(self, "manager", None), "reorganizer", None) release("reorganizer", getattr(reorganizer, "stop", None)) executor = getattr(getattr(self, "searcher", None), "_usage_executor", None) if executor is not None: release("search_executor", lambda: executor.shutdown(wait=True)) graph = getattr(self, "graph", None) vectors = getattr(graph, "vec_db", None) release("qdrant", getattr(getattr(vectors, "client", None), "close", None)) release("neo4j", getattr(getattr(graph, "driver", None), "close", None)) release("trace", getattr(getattr(self, "_trace_log", None), "close", None)) self._closed = True if errors: raise RuntimeError("MemOS cleanup failed: " + "; ".join( f"{label}: {exc}" for label, exc in errors)) from errors[0][1]