"""Synchronous, question-local use of MemOS's working-memory algorithm.""" from __future__ import annotations import copy from contextlib import nullcontext from datetime import timedelta from types import MethodType, SimpleNamespace from typing import Any, Callable from ..llm_client import MAX_ATTEMPTS class WorkingMemoryError(RuntimeError): """A MemOS working-set stage did not complete successfully.""" class WorkingMemoryOutputError(WorkingMemoryError): """The model output stayed invalid after every attempt; no selection was committed.""" recoverable_output = True WORKING_OUTPUT_MAX_ATTEMPTS = MAX_ATTEMPTS def _origin_id(item: Any) -> str | None: """Return the id of the persistent memory that a working copy was made from.""" return ((item.metadata.info or {}).get("copied_from_id") or getattr(item.metadata, "copied_from_id", None)) class WorkingSet: """Question-local MemOS working memory; an invalid model output is retried with the validation error appended.""" def __init__( self, tree: Any, process_llm: Any, embedder: Any, namespace: str, *, capacity: int, trace: Callable[..., Any] | None = None, recorder: Any = None, ) -> None: if not namespace or capacity < 1: raise ValueError("Working memory requires a namespace and positive capacity") self.tree = tree self.process_llm = process_llm self.embedder = embedder self.namespace = namespace self.capacity = int(capacity) self.trace = trace self.recorder = recorder or getattr(process_llm, "recorder", None) self.last_diagnostics: dict[str, Any] = {} self._monitor = None self._engine = None self.tree.memory_manager.memory_size["WorkingMemory"] = self.capacity if self.recorder is not None: from .memos_adapters import install_recording_context install_recording_context() def _emit(self, event: str, **payload: Any) -> None: if self.trace is not None: self.trace(event, **payload) elif self.recorder is not None: self.recorder.emit(event, **payload) def _scope(self, **fields: Any) -> Any: return self.recorder.scope(**fields) if self.recorder is not None else nullcontext() def _validated_generate( self, llm: Any, args: tuple, kwargs: dict, *, stage: str, validate: Callable[[str], Any], schema: str, diagnostics: dict, ) -> tuple[str, Any]: """Call the model with the MemOS request and retry invalid output with the validation error appended.""" from .memos_adapters import ExtractionError call_args, call_kwargs = args, kwargs for attempt in range(1, WORKING_OUTPUT_MAX_ATTEMPTS + 1): response = None finish_reason = None failure = None defer = getattr(llm, "defer_recoverable_output_errors", nullcontext) with self._scope(working_stage=stage, working_stage_attempt=attempt, working_stage_max_attempts=WORKING_OUTPUT_MAX_ATTEMPTS, working_output_recovery=attempt > 1): try: with defer(): response = llm.generate(*call_args, **call_kwargs) except ExtractionError as exc: if exc.kind not in {"truncated", "empty_response"}: raise response = getattr(exc, "raw_response", None) finish_reason = getattr(exc, "finish_reason", None) failure = exc if failure is None: try: value = validate(response) except (ValueError, TypeError, KeyError) as exc: failure = exc diagnostics.setdefault("stage_attempts", {})[stage] = attempt if failure is None: if attempt > 1: recovery = {"stage": stage, "attempts": attempt, "max_attempts": WORKING_OUTPUT_MAX_ATTEMPTS} diagnostics.setdefault("output_recoveries", []).append(recovery) self._emit("working_stage_recovered", **recovery) return response, value diagnostics[stage if stage.endswith("_success") else stage + "_success"] = False self._emit("working_stage_failure", stage=stage, attempt=attempt, max_attempts=WORKING_OUTPUT_MAX_ATTEMPTS, recoverable=True, exhausted=attempt == WORKING_OUTPUT_MAX_ATTEMPTS, error_kind=getattr(failure, "kind", "invalid_schema"), error=str(failure)[:300], raw_response=response, finish_reason=finish_reason) if attempt == WORKING_OUTPUT_MAX_ATTEMPTS: raise WorkingMemoryOutputError( f"MemOS working-memory {stage} failed after {attempt} output attempts" ) from failure feedback = ( f"The previous response failed validation for this same {stage} task: " f"{str(failure)[:300]}. Regenerate the answer using the original " "query and candidates above. Preserve their original indices. " f"Required schema: {schema}. Return one complete valid JSON value " "only, with no Markdown or prose outside JSON. Escape all string " "contents correctly. Keep any reasoning field concise so the " "complete JSON fits the output limit. Do not repeat candidate " "indices or replace the task with a fallback." ) copied_args, copied_kwargs = list(copy.deepcopy(args)), copy.deepcopy(kwargs) messages = copied_args[0] if copied_args else copied_kwargs.get("messages") if not isinstance(messages, list): raise WorkingMemoryError( f"MemOS working-memory {stage} recovery requires message-list input" ) from failure if isinstance(response, str) and response: messages.append({"role": "assistant", "content": response}) messages.append({"role": "user", "content": feedback}) if copied_args: copied_args[0] = messages else: copied_kwargs["messages"] = messages call_args, call_kwargs = tuple(copied_args), copied_kwargs self._emit("working_stage_retry", stage=stage, next_attempt=attempt + 1, max_attempts=WORKING_OUTPUT_MAX_ATTEMPTS, feedback=feedback) raise AssertionError("Unreachable working-memory recovery state") def reset(self) -> None: """Delete this namespace's working copies and the per-query scheduler state.""" try: items = self.tree.get_working_memory(user_name=self.namespace) ids = [item.id for item in items] if ids: self.tree.delete(ids, user_name=self.namespace) if self.tree.get_working_memory(user_name=self.namespace): raise WorkingMemoryError("Working-memory cleanup left stored copies behind") self._emit("working_reset", namespace=self.namespace, deleted_ids=ids, remaining_count=0) finally: if self._monitor is not None: for attribute in ("query_monitors", "working_memory_monitors", "activation_memory_monitors"): managers = getattr(self._monitor, attribute) for cubes in managers.values(): for manager in cubes.values(): manager.release_locks(manager.user_id, manager.mem_cube_id) managers.clear() self._monitor = None if self._engine is not None: self._engine.dispose() self._engine = None close = reset def refresh(self, query: str) -> dict[str, Any]: """Rebuild the working set for one query from the persistent memories.""" if not isinstance(query, str) or not query.strip(): raise ValueError("Working-memory refresh requires a nonempty query") self.reset() from sqlalchemy import create_engine from sqlalchemy.pool import StaticPool from memos.configs.mem_scheduler import GeneralSchedulerConfig from memos.mem_scheduler.base_mixins.memory_ops import BaseSchedulerMemoryMixin from memos.mem_scheduler.memory_manage_modules.retriever import SchedulerRetriever from memos.mem_scheduler.monitors.general_monitor import SchedulerGeneralMonitor from memos.mem_scheduler.optimized_scheduler import OptimizedScheduler from memos.mem_scheduler.schemas.general_schemas import TreeTextMemory_SEARCH_METHOD from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem from memos.mem_scheduler.schemas.task_schemas import MEM_UPDATE_TASK_LABEL, QUERY_TASK_LABEL from memos.mem_scheduler.task_schedule_modules.context import ( SchedulerHandlerContext, SchedulerHandlerServices, ) from memos.mem_scheduler.task_schedule_modules.handlers.memory_update_handler import ( MemoryUpdateHandler, ) from memos.mem_scheduler.task_schedule_modules.handlers.query_handler import QueryMessageHandler from memos.mem_scheduler.utils.misc_utils import extract_json_obj config = GeneralSchedulerConfig( top_k=self.capacity, context_window_size=1, working_mem_monitor_capacity=self.capacity, enable_activation_memory=False, enable_parallel_dispatch=False, use_redis_queue=False, thread_pool_max_workers=2, ) self._engine = create_engine( "sqlite://", poolclass=StaticPool, connect_args={"check_same_thread": False} ) monitor = SchedulerGeneralMonitor(self.process_llm, config, db_engine=self._engine) self._monitor = monitor monitor.last_query_consume_time -= timedelta(seconds=monitor.query_trigger_interval + 1) retriever = SchedulerRetriever(self.process_llm, config) diagnostics: dict[str, Any] = { "algorithm": "query_memory_update_optimized_scheduler", "namespace": self.namespace, "capacity": self.capacity, "candidate_count": 0, "candidate_ids": [], "ranking_success": None, "filtering_success": None, "working_count": 0, "working_ids": [], "origin_ids": [], "query_count": 1, "transport": "synchronous_in_process", } self.last_diagnostics = diagnostics monitor_stage = None def checked_monitor_generate(*args: Any, **kwargs: Any) -> str: if monitor_stage is None: return self.process_llm.generate(*args, **kwargs) def validate(response: str) -> Any: value = extract_json_obj(response) if monitor_stage == "keywords": valid = isinstance(value, list) and all(isinstance(v, str) for v in value) else: valid = (isinstance(value, dict) and type(value.get("trigger_retrieval")) is bool and isinstance(value.get("missing_evidences"), list) and all(isinstance(v, str) for v in value["missing_evidences"])) if not valid: raise ValueError("Unexpected monitor response schema") return value schema = ("an array of keyword strings" if monitor_stage == "keywords" else "an object with boolean trigger_retrieval and array-of-strings missing_evidences") response, value = self._validated_generate( self.process_llm, args, kwargs, stage=monitor_stage, validate=validate, schema=schema, diagnostics=diagnostics) diagnostics[monitor_stage + "_success"] = True self._emit("working_monitor_decision", stage=monitor_stage, decision=copy.deepcopy(value)) return response def checked_monitor_method(stage: str, function: Callable[..., Any]) -> Callable[..., Any]: def call(*args: Any, **kwargs: Any) -> Any: nonlocal monitor_stage monitor_stage = stage try: with self._scope(stage="working_" + stage): self._emit("working_monitor_input", stage=stage, args=copy.deepcopy(args), kwargs=copy.deepcopy(kwargs)) return function(*args, **kwargs) finally: monitor_stage = None return call monitor._process_llm = SimpleNamespace(generate=checked_monitor_generate) monitor.extract_query_keywords = checked_monitor_method("keywords", monitor.extract_query_keywords) monitor.detect_intent = checked_monitor_method("intent", monitor.detect_intent) def checked_indices(component: Any, function: Callable[..., Any], key: str, count: int, stage: str, **kwargs: Any) -> Any: original_llm = component.process_llm def generate(*args: Any, **call_kwargs: Any) -> str: def validate(response: str) -> Any: parsed = extract_json_obj(response) indices = parsed.get(key) if isinstance(parsed, dict) else None if (not isinstance(indices, list) or any(type(i) is not int or not 0 <= i < count for i in indices) or len(indices) != len(set(indices))): raise ValueError(f"{key} must contain unique valid candidate indices") if "reasoning" not in parsed: raise ValueError("Required reasoning field is missing") return parsed response, parsed = self._validated_generate( original_llm, args, call_kwargs, stage=stage, validate=validate, schema=(f'an object containing "{key}": an array of unique integer ' f'indices from 0 through {count - 1}' if count else f'an object containing "{key}": an empty array') + ', and a concise "reasoning" field', diagnostics=diagnostics) self._emit("working_selection_decision", stage=stage, decision=copy.deepcopy(parsed)) return response component.process_llm = SimpleNamespace(generate=generate) try: with self._scope(stage="working_" + stage.removesuffix("_success")): return function(**kwargs) finally: component.process_llm = original_llm orig_rerank = retriever.rerank_pipeline.rerank_memories orig_filter = retriever.memory_filter.filter_unrelated_and_redundant_memories ranking_items: list[Any] = [] def rerank_with_valid_indices(queries: list, original_memories: list, top_k: int) -> Any: from memos.mem_scheduler.utils.filter_utils import transform_name_to_key mapping = {transform_name_to_key(item.memory): item.id for item in ranking_items} self._emit("working_ranking_candidates", queries=list(queries), top_k=top_k, candidates=[{"index": i, "id": mapping.get(transform_name_to_key(text)), "memory": text} for i, text in enumerate(original_memories)]) return checked_indices(retriever.rerank_pipeline, orig_rerank, "new_order", len(original_memories), "ranking_success", queries=queries, original_memories=original_memories, top_k=top_k) def filter_with_valid_indices(query_history: list, memories: list) -> Any: self._emit("working_filter_candidates", queries=list(query_history), candidates=[{"index": i, "id": item.id, "memory": item.memory} for i, item in enumerate(memories)]) return checked_indices(retriever.memory_filter, orig_filter, "kept_memories", len(memories), "filtering_success", query_history=query_history, memories=memories) retriever.rerank_pipeline.rerank_memories = rerank_with_valid_indices retriever.memory_filter.filter_unrelated_and_redundant_memories = filter_with_valid_indices def checked_stage(name: str, function: Callable[..., Any]) -> Callable[..., Any]: def call(*args: Any, **kwargs: Any) -> Any: nonlocal ranking_items if name == "ranking_success": original = kwargs.get("original_memory", args[1] if len(args) > 1 else []) new = kwargs.get("new_memory", args[2] if len(args) > 2 else []) ranking_items = list(original) + list(new) self._emit("working_ranking_input", original_ids=[item.id for item in original], new_ids=[item.id for item in new], entries=[{"id": item.id, "memory": item.memory, "memory_type": item.metadata.memory_type} for item in ranking_items]) result, success = function(*args, **kwargs) diagnostics[name] = bool(success) self._emit("working_selection_result", stage=name, success=bool(success), selected_ids=[item.id for item in result]) if not success: raise WorkingMemoryError(f"MemOS working-memory {name} failed") return result, success return call retriever.process_and_rerank_memories = checked_stage( "ranking_success", retriever.process_and_rerank_memories ) retriever.filter_unrelated_and_redundant_memories = checked_stage( "filtering_success", retriever.filter_unrelated_and_redundant_memories ) orig_search = retriever.search def observed_search(*args: Any, **kwargs: Any) -> Any: with self._scope(stage="working_recall"): self._emit("working_recall_input", query=kwargs.get("query", args[0] if args else None), top_k=kwargs.get("top_k"), method=kwargs.get("method"), search_args=copy.deepcopy(kwargs.get("search_args"))) result = orig_search(*args, **kwargs) self._emit("working_recall_result", query=kwargs.get("query", args[0] if args else None), selected_ids=[item.id for item in result]) return result retriever.search = observed_search cube = SimpleNamespace(text_mem=self.tree) scheduler = SimpleNamespace( monitor=monitor, retriever=retriever, top_k=self.capacity, log_working_memory_replacement=lambda **kwargs: None, _submit_web_logs=lambda *args, **kwargs: None, ) scheduler.transform_working_memories_to_monitors = MethodType( BaseSchedulerMemoryMixin.transform_working_memories_to_monitors, scheduler ) originals: dict[str, Any] = {} def replace_working_memory(**kwargs: Any) -> Any: candidates = kwargs["new_memory"] diagnostics["candidate_count"] = len(candidates) diagnostics["candidate_ids"] = [item.id for item in candidates] copied = [] for item in candidates: if item.metadata.memory_type not in ("UserMemory", "LongTermMemory"): raise WorkingMemoryError("Scheduler candidate is not a persistent text memory") originals[item.id] = self.tree.get(item.id, user_name=self.namespace).model_dump() candidate = item.model_copy(deep=True) info = dict(candidate.metadata.info or {}) info["copied_from_id"] = item.id info["copied_from_type"] = item.metadata.memory_type candidate.metadata.info = info copied.append(candidate) kwargs["new_memory"] = copied selected = OptimizedScheduler.replace_working_memory(scheduler, **kwargs) diagnostics["selected_origin_ids"] = [item.id for item in selected] return selected update_handler = None def submit_messages(messages: list[Any]) -> None: for message in messages: if message.label != MEM_UPDATE_TASK_LABEL: raise WorkingMemoryError(f"Unexpected scheduler task: {message.label}") update_handler.batch_handler(self.namespace, self.namespace, [message]) services = SchedulerHandlerServices( validate_messages=lambda *args, **kwargs: None, submit_messages=submit_messages, create_event_log=lambda **kwargs: SimpleNamespace(), submit_web_logs=lambda *args, **kwargs: None, map_memcube_name=lambda cube_id: cube_id, update_activation_memory_periodically=lambda **kwargs: None, replace_working_memory=replace_working_memory, transform_working_memories_to_monitors=scheduler.transform_working_memories_to_monitors, log_working_memory_replacement=lambda **kwargs: None, ) context = SchedulerHandlerContext( get_mem_cube=lambda: cube, get_monitor=lambda: monitor, get_retriever=lambda: retriever, get_mem_reader=lambda: None, get_feedback_server=lambda: None, get_search_method=lambda: TreeTextMemory_SEARCH_METHOD, get_top_k=lambda: self.capacity, get_enable_activation_memory=lambda: False, get_query_key_words_limit=lambda: 20, services=services, ) update_handler = MemoryUpdateHandler(context) query_handler = QueryMessageHandler(context) message = ScheduleMessageItem( user_id=self.namespace, mem_cube_id=self.namespace, user_name=self.namespace, session_id="independent_benchmark_question", label=QUERY_TASK_LABEL, content=query, ) try: with self._scope(stage="working_refresh", namespace=self.namespace): query_handler.batch_handler(self.namespace, self.namespace, [message]) items = self.tree.get_working_memory(user_name=self.namespace) if len(items) > self.capacity: raise WorkingMemoryError("MemOS working set exceeded configured capacity") expected = diagnostics.get("selected_origin_ids", [])[:self.capacity] actual = [_origin_id(item) for item in items] if sorted(expected) != sorted(actual): raise WorkingMemoryError("MemOS working-memory write was incomplete") if len(actual) != len(set(actual)): raise WorkingMemoryError("MemOS working-memory selection duplicated an origin") for item in items: origin = _origin_id(item) if not origin or origin == item.id or origin not in originals: raise WorkingMemoryError("Working copy lacks a distinct, traceable origin") for memory_id, before in originals.items(): after = self.tree.get(memory_id, user_name=self.namespace).model_dump() if before != after: raise WorkingMemoryError("Working refresh modified a persistent memory") diagnostics.update( working_count=len(items), working_ids=[item.id for item in items], origin_ids=[_origin_id(item) for item in items], ) self._emit("working_refresh", **copy.deepcopy(diagnostics)) return copy.deepcopy(diagnostics) except Exception: self.reset() raise