From 09b89197740eb2f3a24ee06e5f915ebd683548e3 Mon Sep 17 00:00:00 2001 From: Ziyang Guo <121015044+RerankerGuo@users.noreply.github.com> Date: Wed, 1 Jul 2026 22:55:30 +0800 Subject: [PATCH 1/2] feat: route composite cube writes by responsibility --- src/memos/multi_mem_cube/composite_cube.py | 80 +++++++++++++++--- tests/multi_mem_cube/test_composite_cube.py | 94 +++++++++++++++++++++ 2 files changed, 161 insertions(+), 13 deletions(-) create mode 100644 tests/multi_mem_cube/test_composite_cube.py diff --git a/src/memos/multi_mem_cube/composite_cube.py b/src/memos/multi_mem_cube/composite_cube.py index 11c5245ba..f7a4ff9f2 100644 --- a/src/memos/multi_mem_cube/composite_cube.py +++ b/src/memos/multi_mem_cube/composite_cube.py @@ -19,21 +19,65 @@ class CompositeCubeView(MemCubeView): """ A composite view over multiple logical cubes. - For now (fast mode), it simply fan-out writes to all cubes; - later we can add smarter routing / slow mode here. + By default, writes keep the legacy fan-out behavior. If request metadata + explicitly names a target cube or responsibility, writes are routed only to + matching cube views and fall back to fan-out when nothing matches. """ cube_views: list[SingleCubeView] logger: Any + def _request_info(self, request: Any) -> dict[str, Any]: + info = getattr(request, "info", None) + return info if isinstance(info, dict) else {} + + def _view_responsibilities(self, view: SingleCubeView) -> set[str]: + values: set[str] = set() + for attr in ("responsibility", "responsibilities"): + raw = getattr(view, attr, None) + if isinstance(raw, str): + values.add(raw) + elif isinstance(raw, (list, tuple, set)): + values.update(str(item) for item in raw if item) + return values + + def _route_views(self, request: Any) -> list[SingleCubeView]: + info = self._request_info(request) + target_cube_id = info.get("target_cube_id") or info.get("cube_id") + if target_cube_id: + routed = [view for view in self.cube_views if view.cube_id == target_cube_id] + if routed: + return routed + self.logger.warning( + "[CompositeCubeView] target cube %s not found; fallback to fan-out", + target_cube_id, + ) + + responsibility = info.get("responsibility") or info.get("cube_responsibility") + if responsibility: + routed = [ + view + for view in self.cube_views + if responsibility in self._view_responsibilities(view) + ] + if routed: + return routed + self.logger.warning( + "[CompositeCubeView] responsibility %s not matched; fallback to fan-out", + responsibility, + ) + + return self.cube_views + def add_memories(self, add_req: APIADDRequest) -> list[dict[str, Any]]: all_results: list[dict[str, Any]] = [] - cube_count = len(self.cube_views) + target_views = self._route_views(add_req) + cube_count = len(target_views) with timed_stage("add", "multi_cube", cube_count=cube_count): - for idx, view in enumerate(self.cube_views): + for idx, view in enumerate(target_views): self.logger.info( - "[CompositeCubeView] fan-out add to cube=%s (%d/%d)", + "[CompositeCubeView] route add to cube=%s (%d/%d)", view.cube_id, idx + 1, cube_count, @@ -66,13 +110,22 @@ def _search_single_cube(view: SingleCubeView) -> dict[str, Any]: } for future in as_completed(future_to_view): + view = future_to_view[future] cube_result = future.result() - merged_results["text_mem"].extend(cube_result.get("text_mem", [])) - merged_results["act_mem"].extend(cube_result.get("act_mem", [])) - merged_results["para_mem"].extend(cube_result.get("para_mem", [])) - merged_results["pref_mem"].extend(cube_result.get("pref_mem", [])) - merged_results["tool_mem"].extend(cube_result.get("tool_mem", [])) - merged_results["skill_mem"].extend(cube_result.get("skill_mem", [])) + memory_keys = ( + "text_mem", + "act_mem", + "para_mem", + "pref_mem", + "tool_mem", + "skill_mem", + ) + for key in memory_keys: + memories = cube_result.get(key, []) + for memory in memories: + if isinstance(memory, dict): + memory.setdefault("cube_id", view.cube_id) + merged_results[key].extend(memories) note = cube_result.get("pref_note") if note: if merged_results["pref_note"]: @@ -84,9 +137,10 @@ def _search_single_cube(view: SingleCubeView) -> dict[str, Any]: def feedback_memories(self, feedback_req: APIFeedbackRequest) -> list[dict[str, Any]]: all_results: list[dict[str, Any]] = [] + target_views = self._route_views(feedback_req) - for view in self.cube_views: - self.logger.info(f"[CompositeCubeView] fan-out add to cube={view.cube_id}") + for view in target_views: + self.logger.info(f"[CompositeCubeView] route feedback to cube={view.cube_id}") results = view.feedback_memories(feedback_req) all_results.extend(results) diff --git a/tests/multi_mem_cube/test_composite_cube.py b/tests/multi_mem_cube/test_composite_cube.py new file mode 100644 index 000000000..3057a79c3 --- /dev/null +++ b/tests/multi_mem_cube/test_composite_cube.py @@ -0,0 +1,94 @@ +import logging + +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any + +from memos.multi_mem_cube.composite_cube import CompositeCubeView + + +@dataclass +class FakeCubeView: + cube_id: str + responsibility: str | None = None + result: list[dict[str, Any]] = field(default_factory=list) + search_result: dict[str, Any] = field(default_factory=dict) + add_calls: int = 0 + feedback_calls: int = 0 + + def add_memories(self, _add_req): + self.add_calls += 1 + return list(self.result) + + def feedback_memories(self, _feedback_req): + self.feedback_calls += 1 + return list(self.result) + + def search_memories(self, _search_req): + return dict(self.search_result) + + +def test_composite_routes_add_to_explicit_target_cube(): + general = FakeCubeView(cube_id="general", result=[{"memory": "g"}]) + product = FakeCubeView(cube_id="product", result=[{"memory": "p"}]) + composite = CompositeCubeView( + cube_views=[general, product], + logger=logging.getLogger("test.composite"), + ) + + result = composite.add_memories(SimpleNamespace(info={"target_cube_id": "product"})) + + assert result == [{"memory": "p"}] + assert general.add_calls == 0 + assert product.add_calls == 1 + + +def test_composite_routes_feedback_by_responsibility(): + general = FakeCubeView(cube_id="general", responsibility="general") + product = FakeCubeView(cube_id="product", responsibility="shopping", result=[{"ok": True}]) + composite = CompositeCubeView( + cube_views=[general, product], + logger=logging.getLogger("test.composite"), + ) + + result = composite.feedback_memories(SimpleNamespace(info={"responsibility": "shopping"})) + + assert result == [{"ok": True}] + assert general.feedback_calls == 0 + assert product.feedback_calls == 1 + + +def test_composite_falls_back_to_fanout_without_route_match(): + general = FakeCubeView(cube_id="general", result=[{"memory": "g"}]) + product = FakeCubeView(cube_id="product", result=[{"memory": "p"}]) + composite = CompositeCubeView( + cube_views=[general, product], + logger=logging.getLogger("test.composite"), + ) + + result = composite.add_memories(SimpleNamespace(info={"responsibility": "unknown"})) + + assert result == [{"memory": "g"}, {"memory": "p"}] + assert general.add_calls == 1 + assert product.add_calls == 1 + + +def test_composite_search_adds_missing_cube_provenance(): + general = FakeCubeView( + cube_id="general", + search_result={"text_mem": [{"memory": "g"}], "pref_note": "general note"}, + ) + product = FakeCubeView( + cube_id="product", + search_result={"text_mem": [{"memory": "p", "cube_id": "existing"}]}, + ) + composite = CompositeCubeView( + cube_views=[general, product], + logger=logging.getLogger("test.composite"), + ) + + result = composite.search_memories(SimpleNamespace()) + + by_memory = {item["memory"]: item for item in result["text_mem"]} + assert by_memory["g"]["cube_id"] == "general" + assert by_memory["p"]["cube_id"] == "existing" From 68583fdf04abf74d0415c1d1cde5b5e87a206d44 Mon Sep 17 00:00:00 2001 From: Ziyang Guo <121015044+RerankerGuo@users.noreply.github.com> Date: Thu, 2 Jul 2026 11:04:56 +0800 Subject: [PATCH 2/2] fix: satisfy ruff isinstance union check --- src/memos/multi_mem_cube/composite_cube.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/memos/multi_mem_cube/composite_cube.py b/src/memos/multi_mem_cube/composite_cube.py index f7a4ff9f2..67beed671 100644 --- a/src/memos/multi_mem_cube/composite_cube.py +++ b/src/memos/multi_mem_cube/composite_cube.py @@ -37,7 +37,7 @@ def _view_responsibilities(self, view: SingleCubeView) -> set[str]: raw = getattr(view, attr, None) if isinstance(raw, str): values.add(raw) - elif isinstance(raw, (list, tuple, set)): + elif isinstance(raw, list | tuple | set): values.update(str(item) for item in raw if item) return values