code_release/code/memebench/methods/memos_adapters.py

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