Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion src/memos/multi_mem_cube/single_cube.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
)
from memos.memories.textual.item import TextualMemoryItem
from memos.multi_mem_cube.views import MemCubeView
from memos.search import search_text_memories
from memos.search import resolve_filter_for_cube, search_text_memories
from memos.templates.mem_reader_prompts import PROMPT_MAPPING
from memos.types.general_types import (
FINE_STRATEGY,
Expand Down Expand Up @@ -91,6 +91,13 @@ def search_memories(self, search_req: APISearchRequest) -> dict[str, Any]:
Unified memory search handling (text + preference memories).
Preference memories are now searched through the same _search_text flow.
"""
cube_filter = resolve_filter_for_cube(search_req.filter, self.cube_id)
if cube_filter is not search_req.filter:
import copy

search_req = copy.copy(search_req)
search_req.filter = cube_filter

# Create UserContext object
user_context = UserContext(
user_id=search_req.user_id,
Expand Down
14 changes: 12 additions & 2 deletions src/memos/search/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,14 @@
from .search_service import SearchContext, build_search_context, search_text_memories
from .search_service import (
SearchContext,
build_search_context,
resolve_filter_for_cube,
search_text_memories,
)


__all__ = ["SearchContext", "build_search_context", "search_text_memories"]
__all__ = [
"SearchContext",
"build_search_context",
"resolve_filter_for_cube",
"search_text_memories",
]
29 changes: 29 additions & 0 deletions src/memos/search/search_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,35 @@ def build_search_context(
)


def resolve_filter_for_cube(
raw_filter: dict[str, Any] | None, cube_id: str
) -> dict[str, Any] | None:
"""Resolve a multi-cube filter dict into the sub-filter for a single cube.

Supported forms:
- None → None (no filter)
- {"and": [...]} / {"or": [...]} → returned as-is (unified, all cubes share)
- {"cube_A": {...}, "cube_B": {...}} → return raw_filter[cube_id] or None
Mixed top-level (and/or + cube keys) is rejected.
"""
if raw_filter is None:
return None

has_logic_key = "and" in raw_filter or "or" in raw_filter
other_keys = {k for k in raw_filter if k not in ("and", "or")}

if has_logic_key and other_keys:
raise ValueError(
"Invalid filter: top-level 'and'/'or' cannot coexist with per-cube keys "
f"{other_keys}. Use either a unified filter or per-cube filter, not both."
)

if has_logic_key:
return raw_filter

return raw_filter.get(cube_id)


def search_text_memories(
text_mem: Any,
search_req: APISearchRequest,
Expand Down
78 changes: 78 additions & 0 deletions tests/search/test_resolve_filter_for_cube.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
import pytest

from memos.search.search_service import resolve_filter_for_cube


class TestResolveFilterForCube:
"""Tests for resolve_filter_for_cube — multi-cube filter routing."""

# ── None passthrough ──

def test_none_returns_none(self):
assert resolve_filter_for_cube(None, "cube_001") is None

# ── Unified filter (filter2): top-level and/or ──

def test_unified_and_returns_same_for_any_cube(self):
f = {"and": [{"tags": {"contains": "阅读"}}, {"created_at": {"gte": "2025-01-01"}}]}
assert resolve_filter_for_cube(f, "cube_001") is f
assert resolve_filter_for_cube(f, "cube_999") is f

def test_unified_or_returns_same_for_any_cube(self):
f = {"or": [{"tags": {"contains": "A"}}, {"tags": {"contains": "B"}}]}
assert resolve_filter_for_cube(f, "cube_001") is f

# ── Per-cube filter (filter1 / filter4) ──

def test_per_cube_returns_matching_sub_filter(self):
sub_a = {"and": [{"tags": {"contains": "阅读"}}]}
sub_b = {"and": [{"tags": {"contains": "工作"}}]}
f = {"cube_A": sub_a, "cube_B": sub_b}

assert resolve_filter_for_cube(f, "cube_A") is sub_a
assert resolve_filter_for_cube(f, "cube_B") is sub_b

def test_per_cube_missing_key_returns_none(self):
f = {
"cube_A": {"and": [{"tags": {"contains": "阅读"}}]},
"cube_B": {"and": [{"tags": {"contains": "工作"}}]},
}
assert resolve_filter_for_cube(f, "cube_C") is None

def test_per_cube_single_key(self):
sub = {"and": [{"created_at": {"gte": "2025-01-01"}}]}
f = {"cube_only": sub}
assert resolve_filter_for_cube(f, "cube_only") is sub
assert resolve_filter_for_cube(f, "other") is None

# ── Mixed (filter3): illegal ──

def test_mixed_and_with_cube_key_raises(self):
f = {
"and": [{"tags": {"contains": "阅读"}}],
"cube_A": {"and": [{"tags": {"contains": "工作"}}]},
}
with pytest.raises(ValueError, match="cannot coexist"):
resolve_filter_for_cube(f, "cube_A")

def test_mixed_or_with_cube_key_raises(self):
f = {
"or": [{"tags": {"contains": "阅读"}}],
"cube_B": {"and": [{"tags": {"contains": "工作"}}]},
}
with pytest.raises(ValueError, match="cannot coexist"):
resolve_filter_for_cube(f, "cube_B")

# ── Edge cases ──

def test_empty_dict_returns_none(self):
assert resolve_filter_for_cube({}, "cube_001") is None

def test_per_cube_with_empty_sub_filter(self):
f = {"cube_A": {}}
result = resolve_filter_for_cube(f, "cube_A")
assert result == {}

def test_unified_and_empty_list(self):
f = {"and": []}
assert resolve_filter_for_cube(f, "any") is f
Loading