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