diff --git a/src/memos/multi_mem_cube/single_cube.py b/src/memos/multi_mem_cube/single_cube.py index f84fc60e1..1ea4be4e5 100644 --- a/src/memos/multi_mem_cube/single_cube.py +++ b/src/memos/multi_mem_cube/single_cube.py @@ -37,6 +37,93 @@ logger = get_logger(__name__) +# Cap on how many memories per bucket appear in the search log summary. The +# full count is still reported via ``summary["counts"]``; this just bounds the +# INFO-line size when top_k is large. +_LOG_SAMPLE_MAX_PER_BUCKET = 5 + +# Buckets we surface counts/samples for in the log summary. ``act_mem`` and +# ``para_mem`` are intentionally excluded — they are not populated by the text +# search path and would just add noise. +_LOG_SUMMARY_BUCKETS = ("text_mem", "pref_mem", "tool_mem", "skill_mem") + + +def _summarize_search_result_for_log( + result: dict[str, Any], +) -> dict[str, Any]: + """Build a compact, embedding-free summary of a ``MOSSearchResult`` for + logging. + + The previous ``search_memories`` INFO log dumped the full result dict, + which under the default ``dedup="mmr"`` (or ``dedup="sim"``) branch keeps + the raw ``metadata.embedding`` vector on every memory. That leaked + high-dimensional floats into the log stream, caused log explosion, and + raised privacy concerns (issue #2103). + + This helper produces a bounded dict containing: + * ``counts``: total memories per bucket (``text_mem`` / ``pref_mem`` / + ``tool_mem`` / ``skill_mem``). + * ``total``: sum of counts across all buckets. + * ``samples``: up to ``_LOG_SAMPLE_MAX_PER_BUCKET`` items per bucket, + exposing only ``id`` / ``memory_type`` / ``relativity`` — never the + embedding. + * ``has_embedding``: True iff at least one memory carries a non-empty + embedding list. The boolean is safe to log; the vector itself is not. + + The helper degrades gracefully on malformed input (missing buckets, + non-dict memory items). + """ + + counts: dict[str, int] = dict.fromkeys(_LOG_SUMMARY_BUCKETS, 0) + samples: dict[str, list[dict[str, Any]]] = {bucket: [] for bucket in _LOG_SUMMARY_BUCKETS} + has_embedding = False + + for bucket in _LOG_SUMMARY_BUCKETS: + groups = result.get(bucket) or [] + if not isinstance(groups, list): + continue + + for group in groups: + if not isinstance(group, dict): + continue + memories = group.get("memories") or [] + if not isinstance(memories, list): + continue + + for mem in memories: + if not isinstance(mem, dict): + continue + + # Count only validated dict entries so ``counts`` matches the + # docstring's "total memories per bucket" contract. Non-dict + # items (strings, ``None``, etc.) are silently ignored rather + # than inflating the reported total. + counts[bucket] += 1 + + metadata = mem.get("metadata") or {} + if isinstance(metadata, dict): + embedding = metadata.get("embedding") + if isinstance(embedding, list) and len(embedding) > 0: + has_embedding = True + + if len(samples[bucket]) < _LOG_SAMPLE_MAX_PER_BUCKET: + metadata_dict = metadata if isinstance(metadata, dict) else {} + samples[bucket].append( + { + "id": mem.get("id"), + "memory_type": metadata_dict.get("memory_type"), + "relativity": metadata_dict.get("relativity"), + } + ) + + return { + "counts": counts, + "total": sum(counts.values()), + "samples": samples, + "has_embedding": has_embedding, + } + + if TYPE_CHECKING: from memos.api.product_models import APIADDRequest, APIFeedbackRequest, APISearchRequest from memos.mem_cube.navie import NaiveMemCube @@ -129,8 +216,9 @@ def search_memories(self, search_req: APISearchRequest) -> dict[str, Any]: self.cube_id, ) - self.logger.info(f"Search memories result: {memories_result}") - self.logger.info(f"Search {len(memories_result)} memories.") + self.logger.info( + "Search memories summary: %s", _summarize_search_result_for_log(memories_result) + ) return memories_result @timed diff --git a/tests/multi_mem_cube/__init__.py b/tests/multi_mem_cube/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/multi_mem_cube/test_search_log_sanitization.py b/tests/multi_mem_cube/test_search_log_sanitization.py new file mode 100644 index 000000000..31053d035 --- /dev/null +++ b/tests/multi_mem_cube/test_search_log_sanitization.py @@ -0,0 +1,313 @@ +""" +Regression tests for the sanitized ``search_memories`` INFO log emitted by +:mod:`memos.multi_mem_cube.single_cube`. + +Issue #2103: the previous implementation printed the full ``MOSSearchResult`` +dict, which under the default ``dedup="mmr"`` (or ``dedup="sim"``) branch keeps +``metadata.embedding`` as a high-dimensional float list. That leaks the raw +vectors into the INFO log stream, causing log explosion and privacy risk. + +The fix introduces :func:`_summarize_search_result_for_log` which produces a +bounded, embedding-free summary. These tests pin its contract so any regression +that re-enables the leak fails fast. + +NOTE: ``_summarize_search_result_for_log`` is imported lazily inside each test +to work around a known circular import in ``memos.api.handlers.__init__`` (see +``tests/test_add_stage_logging.py`` for the same workaround). +""" + +from __future__ import annotations + +import json + +import pytest + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _import_helper(): + # Preload memos.api.handlers first — importing single_cube directly would + # trigger a partially-initialised module during its own re-entrant import + # chain (see the circular-import workaround in tests/test_add_stage_logging.py). + import memos.api.handlers # noqa: F401 + + from memos.multi_mem_cube.single_cube import _summarize_search_result_for_log + + return _summarize_search_result_for_log + + +def _make_memory( + mem_id: str, + memory_type: str, + *, + embedding: list[float] | None = None, + relativity: float = 0.5, + cube_id: str = "cube-A", +) -> dict: + """Construct a single formatted memory dict matching what + ``format_memory_item`` produces downstream.""" + return { + "id": mem_id, + "memory": f"memory text for {mem_id}", + "ref_id": f"[{mem_id[:8]}]", + "metadata": { + "id": mem_id, + "memory_type": memory_type, + "relativity": relativity, + "embedding": embedding if embedding is not None else [], + "sources": [], + "usage": [], + }, + "cube_id": cube_id, + } + + +def _make_result( + text_mems: list[dict] | None = None, + pref_mems: list[dict] | None = None, + tool_mems: list[dict] | None = None, + skill_mems: list[dict] | None = None, + cube_id: str = "cube-A", +) -> dict: + """Build a ``MOSSearchResult``-shaped dict for a single cube.""" + return { + "text_mem": [ + { + "cube_id": cube_id, + "memories": text_mems or [], + "total_nodes": len(text_mems or []), + } + ], + "act_mem": [], + "para_mem": [], + "pref_mem": [ + { + "cube_id": cube_id, + "memories": pref_mems or [], + "total_nodes": len(pref_mems or []), + } + ], + "pref_note": "", + "tool_mem": [ + { + "cube_id": cube_id, + "memories": tool_mems or [], + "total_nodes": len(tool_mems or []), + } + ], + "skill_mem": [ + { + "cube_id": cube_id, + "memories": skill_mems or [], + "total_nodes": len(skill_mems or []), + } + ], + } + + +# --------------------------------------------------------------------------- +# Core sanitization contract +# --------------------------------------------------------------------------- + + +class TestNoEmbeddingLeak: + def test_embedding_floats_are_stripped_from_summary(self): + """The raw embedding float list must never appear in the summary.""" + summarize = _import_helper() + big_vector = [0.0123, -0.0456, 0.7890, -0.1234, 0.5678] * 50 # 250 floats + result = _make_result( + text_mems=[ + _make_memory("id-1", "UserMemory", embedding=big_vector), + _make_memory("id-2", "LongTermMemory", embedding=big_vector), + ] + ) + + summary = summarize(result) + rendered = json.dumps(summary, default=str) + + # None of the vector components may appear in the rendered log payload. + for value in (0.0123, -0.0456, 0.7890, 0.5678): + assert str(value) not in rendered, ( + f"Embedding float {value} leaked into log summary: {rendered!r}" + ) + + def test_summary_flags_embedding_presence_as_boolean(self): + """We still want to know *whether* embeddings were populated, but only + as a bool — not the vector itself.""" + summarize = _import_helper() + big_vector = [0.11, 0.22, 0.33] + result = _make_result(text_mems=[_make_memory("a", "UserMemory", embedding=big_vector)]) + + summary = summarize(result) + + assert summary["has_embedding"] is True + + def test_summary_marks_absent_embedding(self): + summarize = _import_helper() + result = _make_result(text_mems=[_make_memory("a", "UserMemory", embedding=[])]) + + summary = summarize(result) + + assert summary["has_embedding"] is False + + +# --------------------------------------------------------------------------- +# Debug value preservation +# --------------------------------------------------------------------------- + + +class TestSummaryContent: + def test_per_bucket_counts_are_reported(self): + summarize = _import_helper() + result = _make_result( + text_mems=[ + _make_memory("t1", "UserMemory"), + _make_memory("t2", "LongTermMemory"), + _make_memory("t3", "WorkingMemory"), + ], + pref_mems=[_make_memory("p1", "PreferenceMemory")], + tool_mems=[ + _make_memory("tool1", "ToolSchemaMemory"), + _make_memory("tool2", "ToolTrajectoryMemory"), + ], + skill_mems=[_make_memory("s1", "SkillMemory")], + ) + + summary = summarize(result) + counts = summary["counts"] + + assert counts["text_mem"] == 3 + assert counts["pref_mem"] == 1 + assert counts["tool_mem"] == 2 + assert counts["skill_mem"] == 1 + + def test_total_count_is_reported(self): + summarize = _import_helper() + result = _make_result( + text_mems=[_make_memory(f"t{i}", "UserMemory") for i in range(4)], + pref_mems=[_make_memory("p1", "PreferenceMemory")], + ) + + summary = summarize(result) + + # total should be sum across all four buckets + assert summary["total"] == 5 + + def test_sample_memory_ids_are_included_for_debug(self): + summarize = _import_helper() + result = _make_result( + text_mems=[ + _make_memory("id-alpha", "UserMemory", relativity=0.9), + _make_memory("id-beta", "LongTermMemory", relativity=0.8), + ] + ) + + summary = summarize(result) + + # We expect the per-bucket sample to expose ids + memory_type + relativity + text_samples = summary["samples"]["text_mem"] + assert len(text_samples) == 2 + ids = {item["id"] for item in text_samples} + assert ids == {"id-alpha", "id-beta"} + types = {item["memory_type"] for item in text_samples} + assert types == {"UserMemory", "LongTermMemory"} + + def test_sample_size_is_bounded(self): + """Samples must be capped to keep the log line small even when top_k + is large. The cap value itself is an implementation detail; here we + just assert it is bounded to <= 5 per bucket.""" + summarize = _import_helper() + text_mems = [_make_memory(f"id-{i}", "UserMemory") for i in range(20)] + result = _make_result(text_mems=text_mems) + + summary = summarize(result) + text_samples = summary["samples"]["text_mem"] + + assert len(text_samples) <= 5 + # But bucket count still reports the true value. + assert summary["counts"]["text_mem"] == 20 + + +# --------------------------------------------------------------------------- +# Dedup branches — regression coverage for issue #2103 +# --------------------------------------------------------------------------- + + +class TestDedupBranches: + @pytest.mark.parametrize( + ("embedding_value", "expected_has_embedding"), + [ + ([0.001, 0.002, 0.003, 0.004, 0.005], True), # dedup="mmr" / "sim" + ([], False), # dedup="no" (or anything else that drops embeddings) + ], + ) + def test_both_branches_produce_safe_output(self, embedding_value, expected_has_embedding): + """Whether embeddings are populated (mmr/sim) or empty (no), the log + summary must remain safe and structurally identical.""" + summarize = _import_helper() + result = _make_result( + text_mems=[_make_memory("m1", "UserMemory", embedding=embedding_value)] + ) + + summary = summarize(result) + rendered = json.dumps(summary, default=str) + + # These specific floats must not leak. + for value in embedding_value: + assert str(value) not in rendered + + # ``has_embedding`` is the one boolean signal we deliberately expose; + # pin it per branch so the empty-embedding case (where the loop above + # is a no-op) still carries a real, non-trivial assertion. + assert summary["has_embedding"] is expected_has_embedding + + # Pin the top-level structure *exactly* — using ``==`` rather than + # ``>=``. A future accidental leak under a different key name + # (e.g. ``"embed"`` / ``"vector"`` / ``"raw_embedding"``) would + # add an extra key here and fail this assertion, whereas the earlier + # ``'"embedding"' not in rendered`` check was trivially true since the + # helper never emits that literal key at all. + assert set(summary.keys()) == {"counts", "total", "samples", "has_embedding"} + + +# --------------------------------------------------------------------------- +# Robustness +# --------------------------------------------------------------------------- + + +class TestRobustness: + def test_missing_bucket_keys_do_not_raise(self): + """The helper must tolerate partially-populated results — e.g. when a + bucket is missing entirely (older callers, edge cases).""" + summarize = _import_helper() + result = {"text_mem": [], "pref_mem": []} + summary = summarize(result) + + assert summary["counts"]["text_mem"] == 0 + assert summary["counts"]["pref_mem"] == 0 + assert summary["counts"]["tool_mem"] == 0 + assert summary["counts"]["skill_mem"] == 0 + assert summary["total"] == 0 + + def test_non_dict_memory_items_do_not_raise(self): + """The helper should degrade gracefully rather than crash the whole + search path if a memory item is not the expected shape.""" + summarize = _import_helper() + result = _make_result() + # inject a malformed memory + result["text_mem"][0]["memories"] = [ + {"id": "ok", "metadata": {"memory_type": "UserMemory", "relativity": 0.1}}, + "not-a-dict", + None, + ] + + # Must not raise + summary = summarize(result) + + # Only validated dict entries contribute to the reported bucket count; + # the ``"not-a-dict"`` string and ``None`` are silently skipped so the + # summary matches the docstring's "total memories per bucket" contract. + assert summary["counts"]["text_mem"] == 1