"""Adapters between the evaluation clients and the pinned MemOS v2.0.11 engine.""" from __future__ import annotations import copy import json import math import re import time import uuid from collections import deque from contextlib import nullcontext, contextmanager from contextvars import ContextVar from functools import wraps from types import SimpleNamespace from typing import Any from ..llm_client import MAX_ATTEMPTS def install_recording_context(): """Carry the recording scope into the tasks submitted to MemOS thread pools.""" from memos.context.context import ContextThreadPoolExecutor from .memos_recording import Recorder if getattr(ContextThreadPoolExecutor.submit, "_recording_context", False): return original = ContextThreadPoolExecutor.submit @wraps(original) def submit(executor, function, *args, **kwargs): recorder = Recorder.current() if recorder is not None: function = recorder.wrap_callable(function) return original(executor, function, *args, **kwargs) submit._recording_context = True ContextThreadPoolExecutor.submit = submit class ExtractionError(ValueError): def __init__(self, kind: str, message: str): super().__init__(message) self.kind = kind def validate_extraction_response(content: str, finish_reason: str | None = None) -> dict: """Parse and validate one MemOS extraction response.""" if finish_reason in {"length", "content_filter"}: raise ExtractionError("truncated" if finish_reason == "length" else "filtered", f"Extraction terminated with {finish_reason}") if not isinstance(content, str) or not content.strip(): raise ExtractionError("empty_response", "Extraction returned no text") text = content.strip() if text.startswith("```"): text = re.sub(r"^```(?:json)?\s*", "", text, flags=re.I) text = re.sub(r"\s*```$", "", text) start = text.find("{") if start < 0: raise ExtractionError("parse_error", "Extraction contains no JSON object") try: value = json.loads(text[start:]) except json.JSONDecodeError as exc: raise ExtractionError("parse_error", f"Invalid extraction JSON at offset {exc.pos}") from exc if not isinstance(value, dict) or not isinstance(value.get("memory list"), list): raise ExtractionError("schema_error", "Expected an object containing a memory list array") if not isinstance(value.get("summary", ""), str): raise ExtractionError("schema_error", "Extraction summary must be text") for index, item in enumerate(value["memory list"]): if not isinstance(item, dict): raise ExtractionError("schema_error", f"Memory {index} must be an object") for key in ("value", "key"): if not isinstance(item.get(key), str) or not item[key].strip(): raise ExtractionError("schema_error", f"Memory {index} has invalid {key}") if item.get("memory_type") not in {"UserMemory", "LongTermMemory"}: raise ExtractionError("schema_error", f"Memory {index} has invalid memory_type") if not isinstance(item.get("tags"), list) or any(not isinstance(t, str) for t in item["tags"]): raise ExtractionError("schema_error", f"Memory {index} has invalid tags") return value def structured_messages(session: dict) -> list[dict]: """Build MemOS chat messages from one visit, keeping each turn's speaker and visit fields.""" date, visit = str(session["date"]), int(session["visit_number"]) department = str(session.get("department") or "") messages = [] for index, turn in enumerate(session["turns"]): speaker, content = str(turn["speaker"]), str(turn["text"]) folded = speaker.casefold() role = "user" if any(x in folded for x in ("patient",)) else "assistant" messages.append({"role": role, "content": f"{speaker} ({department}): {content}", "chat_time": date, "message_id": f"visit-{visit}-turn-{index}", "visit_number": visit, "date": date, "department": department, "speaker": speaker}) return messages def as_dict(item: Any) -> dict: return item.model_dump() if hasattr(item, "model_dump") else copy.deepcopy(item) def normalized_sources(metadata: dict) -> list[dict]: out = [] for source in metadata.get("sources") or []: source = source.model_dump() if hasattr(source, "model_dump") else source for _ in range(3): if not isinstance(source, str): break try: source = json.loads(source) except (ValueError, TypeError): break if isinstance(source, dict): out.append(source) return out def render_memory_items(items: list) -> tuple[str, list[int], list[dict]]: blocks, visits, previews = [], set(), [] for item in items: data = as_dict(item) metadata = data.get("metadata") or {} sources = normalized_sources(metadata) refs = {(s.get("visit_number"), s.get("date") or s.get("chat_time")) for s in sources} if metadata.get("visit_number") is not None: refs.add((metadata["visit_number"], metadata.get("date"))) refs = sorted(refs, key=lambda x: (str(x[1]), str(x[0]))) for visit, _ in refs: if visit is not None: visits.add(int(visit)) citations = "; ".join(f"recorded on {date or 'unknown date'}" + (f", visit {visit}" if visit is not None else "") for visit, date in refs) header = f"[Memory {len(blocks) + 1} | {metadata.get('memory_type', 'TextMemory')}" header += (f" | {citations}" if citations else "") + "]" memory = str(data.get("memory") or "").strip() if not memory: raise ValueError("Retrieved memory has empty content") blocks.append(f"{header}\n{memory}") previews.append({"id": data.get("id"), "memory": memory, "sources": sources, "type": metadata.get("memory_type"), "visits": [int(v) for v, _ in refs if v is not None]}) return "\n\n".join(blocks), sorted(visits), previews class TrackedLLM: def __init__(self, client, error_sink=None): self.chat_client = client self.model = getattr(client, "model", "unknown") self.model_name_or_path = self.model self.config = SimpleNamespace(model_name_or_path=self.model) self.error_sink = error_sink self._recovering_output = ContextVar("recovering_output", default=False) @contextmanager def defer_recoverable_output_errors(self): """Within this context, truncated or empty output is not reported to the error sink.""" token = self._recovering_output.set(True) try: yield finally: self._recovering_output.reset(token) def generate(self, messages, **kwargs): response = None try: response = self.chat_client.chat(messages, **kwargs) finish = getattr(response, "finish_reason", None) if finish in {"length", "content_filter"}: raise ExtractionError("truncated" if finish == "length" else "filtered", f"MemOS LLM stage terminated with {finish}") if not isinstance(response.content, str) or not response.content.strip(): raise ExtractionError("empty_response", "MemOS LLM stage returned no text") return response.content except Exception as exc: if isinstance(exc, ExtractionError): exc.raw_response = getattr(response, "content", None) exc.finish_reason = getattr(response, "finish_reason", None) recoverable = (self._recovering_output.get() and isinstance(exc, ExtractionError) and exc.kind in {"truncated", "empty_response"}) if self.error_sink and not recoverable: self.error_sink("llm", exc) raise class TrackedEmbedder: """Embedder for MemOS that accepts a string or a list and returns a list of vectors.""" def __init__(self, embedder, dimension: int, error_sink=None): self.client = embedder self.model = getattr(embedder, "model", "unknown") self.config = SimpleNamespace(model_name_or_path=self.model) self.dimension, self.error_sink = dimension, error_sink def embed(self, texts): values = [texts] if isinstance(texts, str) else list(texts) if not values: return [] recorder = getattr(self, "recorder", None) call_id = "embedding_" + uuid.uuid4().hex started = time.monotonic() if recorder is not None: recorder.emit("embedding_request", embedding_call_id=call_id, model=self.model, dimension=self.dimension, texts=list(values)) try: vectors = self.client.embed(values) if len(vectors) != len(values) or any(len(v) != self.dimension for v in vectors): raise ValueError("Embedding shape does not match configured dimension/input count") if recorder is not None: recorder.emit("embedding_result", embedding_call_id=call_id, model=self.model, dimension=self.dimension, count=len(vectors), vectors=vectors, seconds=time.monotonic() - started) return vectors except Exception as exc: if recorder is not None: recorder.emit("embedding_failure", embedding_call_id=call_id, error=f"{type(exc).__name__}: {exc}"[:300], seconds=time.monotonic() - started) if self.error_sink: self.error_sink("embedding", exc) raise class ReaderWindowEmbeddings: """Embed a validated window in one batch and return its vectors to MemOS item construction in order.""" def __init__(self, backend): self.backend = backend self.pending = deque() def prepare(self, values): self.assert_consumed() if values: vectors = self.backend.embed(values) if len(vectors) != len(values): raise ExtractionError("embedding_alignment_error", "Embedding batch omitted a validated memory") self.pending.extend(zip(values, vectors, strict=True)) def embed(self, values): if not self.pending or list(values) != [self.pending[0][0]]: raise ExtractionError("embedding_alignment_error", "MemOS item embedding input differs from validated window") _, vector = self.pending.popleft() return [vector] def assert_consumed(self): if self.pending: raise ExtractionError("item_construction_error", "MemOS reader did not consume every validated embedding") def make_validated_reader(llm, embedder, max_attempts: int, window_tokens: int, trace=None, *, reader_workers: int, window_overlap_tokens: int, recorder=None): """Build the MemOS chat reader with the evaluation clients and response validation.""" from memos.mem_reader.simple_struct import SimpleStructMemReader, PROMPT_DICT from memos.mem_reader.read_multi_modal import detect_lang from memos.mem_reader.utils import count_tokens_text from memos.context.context import ContextThreadPoolExecutor if int(reader_workers) < 1: raise ValueError("reader_workers must be positive") recorder = recorder or getattr(llm, "recorder", None) or getattr(llm.chat_client, "recorder", None) if recorder is not None: install_recording_context() class ValidatedReader(SimpleStructMemReader): def __init__(self): self.config = SimpleNamespace(remove_prompt_example=False) self.llm = self.general_llm = llm self.embedder = ReaderWindowEmbeddings(embedder) self.chat_window_max_tokens = window_tokens self.chat_window_overlap_tokens = int(window_overlap_tokens) self.memory_max_length = 8000 self._count_tokens = count_tokens_text self.graph_db = self.searcher = None self.expected_items = 0 self.attempts = 0 self.outcomes = [] self.reader_workers = int(reader_workers) self.scene_index = None self.scene_count = 0 self.scene_diagnostics = [] self.window_index = -1 self.recording_context = {} def _emit(self, event, **fields): if trace: trace(event, scene_index=self.scene_index, window_index=self.window_index, **fields) elif recorder is not None: recorder.emit(event, scene_index=self.scene_index, window_index=self.window_index, **fields) def _scope(self, **fields): return recorder.scope(**fields) if recorder is not None else nullcontext() def _get_llm_response(self, mem_str, custom_tags): self.window_index += 1 with self._scope(stage="extraction", scene_index=self.scene_index, window_index=self.window_index): return self._generate_window_response(mem_str, custom_tags) def _generate_window_response(self, mem_str, custom_tags): lang = detect_lang(mem_str) prompt = PROMPT_DICT["chat"][lang].replace("${conversation}", mem_str) tag_prompt = (PROMPT_DICT["custom_tags"][lang].replace("{custom_tags}", str(custom_tags)) if custom_tags else "") prompt = prompt.replace("${custom_tags_prompt}", tag_prompt) last = None for attempt in range(1, max_attempts + 1): self.attempts += 1 response = None try: with self._scope(attempt=attempt): response = llm.chat_client.chat([{"role": "user", "content": prompt}]) result = validate_extraction_response(response.content, getattr(response, "finish_reason", None)) self.expected_items += len(result["memory list"]) self.outcomes.append("valid_nonempty" if result["memory list"] else "valid_empty") if trace: self._emit("extraction_response", attempt=attempt, outcome=self.outcomes[-1], finish_reason=getattr(response, "finish_reason", None), produced=len(result["memory list"])) self._emit("extraction_parsed", attempt=attempt, result=copy.deepcopy(result)) except Exception as exc: last = exc if trace: self._emit("extraction_response", attempt=attempt, outcome=getattr(exc, "kind", "api_error"), error=str(exc)[:300], finish_reason=getattr(response, "finish_reason", None), raw_response=getattr(response, "content", None)) else: values = [item["value"] for item in result["memory list"]] with self._scope(stage="extraction_embedding", attempt=attempt): self.embedder.prepare(values) if trace and values: self._emit("extraction_embedding_batch", values=len(values), requests=1) return result raise ExtractionError(getattr(last, "kind", "api_error"), f"Extraction failed after {max_attempts} attempts: {last}") from last def _iter_chat_windows(self, scene_data_info, max_tokens=None, overlap=None): overlap = self.chat_window_overlap_tokens if overlap is None else overlap for window in super()._iter_chat_windows(scene_data_info, max_tokens, overlap): for source in window["sources"]: original = scene_data_info[source["index"]] for key in ("message_id", "visit_number", "date", "department", "speaker"): if key in original: source[key] = original[key] yield window def _read_scene(self, group, info): with self._scope(**{**self.recording_context, "stage": "extraction", "scene_index": self.scene_index}): items = self._process_chat_data(group, dict(info), mode="fine") if len(items) != self.expected_items: raise ExtractionError("item_construction_error", "MemOS reader omitted one or more validated memories") self.embedder.assert_consumed() return items def read_messages(self, messages, info): self.expected_items, self.attempts, self.outcomes = 0, 0, [] self.window_index = -1 self.recording_context = recorder.context() if recorder is not None else {} self.embedder.pending.clear() groups = self.get_scene_data_info([messages], "chat") eligible = [message for message in messages if isinstance(message, dict) and str(message.get("role") or "").strip().lower() in {"user", "assistant", "system"} and isinstance(message.get("content"), str) and message["content"]] self.scene_count = len(groups) self.scene_diagnostics = [] offset = 0 for group in groups: for index, message in enumerate(group): original = eligible[offset + index] if message["content"] != original["content"]: raise ExtractionError("source_alignment_error", "MemOS message/source window alignment differs") for key in ("message_id", "visit_number", "date", "department", "speaker"): if key in original: message[key] = original[key] offset += len(group) - 2 if not groups: return [] if len(groups) == 1: self.scene_index = 0 items, error = [], None try: items = self._read_scene(groups[0], info) return items except Exception as exc: error = getattr(exc, "kind", "dependency_error") raise finally: self.scene_diagnostics = [{"scene_index": 0, "attempts": self.attempts, "outcomes": list(self.outcomes), "items": len(items), "error": error}] if trace: trace("extraction_scenes", scenes=1, workers=1, failed=int(error is not None), diagnostics=self.scene_diagnostics) readers = [] for index in range(len(groups)): child = ValidatedReader() child.graph_db, child.searcher = self.graph_db, self.searcher child.scene_index = index child.recording_context = dict(self.recording_context) readers.append(child) results, failures = [], [] with ContextThreadPoolExecutor(max_workers=min(self.reader_workers, len(groups))) as executor: futures = [executor.submit(child._read_scene, group, dict(info)) for child, group in zip(readers, groups, strict=True)] for index, (child, future) in enumerate(zip(readers, futures, strict=True)): try: items = future.result() results.extend(items) error = None except Exception as exc: failures.append(exc) items, error = [], getattr(exc, "kind", "dependency_error") self.scene_diagnostics.append({"scene_index": index, "attempts": child.attempts, "outcomes": list(child.outcomes), "items": len(items), "error": error}) self.expected_items = sum(child.expected_items for child in readers) self.attempts = sum(child.attempts for child in readers) self.outcomes = [outcome for child in readers for outcome in child.outcomes] if trace: trace("extraction_scenes", scenes=len(groups), workers=min(self.reader_workers, len(groups)), failed=len(failures), diagnostics=self.scene_diagnostics) if failures: first = failures[0] raise ExtractionError(getattr(first, "kind", "dependency_error"), f"{len(failures)} of {len(groups)} extraction scenes failed; " f"visit not committed: {first}") from first if len(results) != self.expected_items: raise ExtractionError("item_construction_error", "MemOS reader omitted one or more validated memories") return results return ValidatedReader() def make_strict_reranker(url: str, model: str, timeout: float, error_sink=None, trace=None, recorder=None): """Build the MemOS HTTP BGE reranker with response validation.""" import requests from memos.reranker.http_bge import HTTPBGEReranker, _TAG1 if recorder is not None: install_recording_context() def emit(event, **fields): if trace: trace(event, **fields) elif recorder is not None: recorder.emit(event, **fields) class StrictBGEReranker(HTTPBGEReranker): def rerank(self, query, graph_results, top_k, search_priority=None, **kwargs): context = recorder.context() if recorder is not None else {} rerank_id = "rerank_" + uuid.uuid4().hex fields = {"stage": "bge_rerank", "retrieval_stage": context.get("stage"), "rerank_id": rerank_id} with recorder.scope(**fields) if recorder is not None else nullcontext(): return self._rerank_observed(query, graph_results, top_k, search_priority, rerank_id, **kwargs) def _rerank_observed(self, query, graph_results, top_k, search_priority, rerank_id, **kwargs): if not graph_results: emit("rerank_result", rerank_id=rerank_id, query=query, candidates=[], selected=[], skipped_empty=True) return [] started = time.monotonic() try: documents = [_TAG1.sub("", str(as_dict(item)["memory"])) for item in graph_results] emit("rerank_request", rerank_id=rerank_id, query=query, model=self.model, top_k=top_k, search_priority=copy.deepcopy(search_priority), documents=[{"index": i, "id": as_dict(item).get("id"), "memory_type": as_dict(item).get("metadata", {}).get("memory_type"), "text": documents[i]} for i, item in enumerate(graph_results)]) for attempt in range(MAX_ATTEMPTS): try: response = requests.post(self.reranker_url, json={"model": self.model, "query": query, "documents": documents}, timeout=self.timeout) response.raise_for_status() break except requests.RequestException: if attempt == MAX_ATTEMPTS - 1: raise time.sleep(1.5 * (attempt + 1)) response_data = response.json() rows = response_data.get("results") if not isinstance(rows, list) or len(rows) != len(documents): raise ValueError("BGE must return one indexed score for every document") scores = {} for row in rows: index, score = row.get("index"), row.get("relevance_score") if type(index) is not int or not 0 <= index < len(documents) or index in scores: raise ValueError("BGE returned an invalid or duplicate document index") score = float(score) if not math.isfinite(score) or not 0 <= score <= 1: raise ValueError("BGE relevance score must be finite and in [0,1]") scores[index] = score scored = [(item, self._apply_boost_generic(item, scores[index], search_priority)) for index, item in enumerate(graph_results)] candidate_scores = [{"index": index, "id": as_dict(item).get("id"), "raw_score": scores[index], "score": score} for index, (item, score) in enumerate(scored)] scored.sort(key=lambda pair: pair[1], reverse=True) emit("rerank_result", rerank_id=rerank_id, query=query, raw_response=response_data, candidates=candidate_scores, selected=[{"id": as_dict(item).get("id"), "score": score} for item, score in scored[:top_k]], top_k=top_k, seconds=time.monotonic() - started) if trace: trace("rerank", candidates=len(documents), returned=min(top_k, len(scored)), model=self.model, seconds=time.monotonic() - started) return scored[:top_k] except Exception as exc: emit("rerank_failure", rerank_id=rerank_id, query=query, error=f"{type(exc).__name__}: {exc}"[:300]) if error_sink: error_sink("reranker", exc) raise return StrictBGEReranker(reranker_url=url, model=model, timeout=timeout) class ObservedProxy: """Report the exceptions of a dependency to the error sink before re-raising them.""" def __init__(self, target, error_sink, label): self._target, self._error_sink, self._label = target, error_sink, label def __getattr__(self, name): attr = getattr(self._target, name) if not callable(attr): return attr def checked(*args, **kwargs): try: return attr(*args, **kwargs) except Exception as exc: self._error_sink(f"{self._label}.{name}", exc) raise return checked