File size: 24,992 Bytes
750b88d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 | """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
|