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