fix(vector_stores/turbopuffer): apply all filter operators instead of dropping them (#6564)
This commit is contained in:
@@ -131,6 +131,18 @@ class TurbopufferDB(VectorStoreBase):
|
|||||||
))
|
))
|
||||||
return results
|
return results
|
||||||
|
|
||||||
|
# Maps mem0 filter operators to their Turbopuffer equivalents.
|
||||||
|
OPERATOR_MAP = {
|
||||||
|
"eq": "Eq",
|
||||||
|
"ne": "NotEq",
|
||||||
|
"gt": "Gt",
|
||||||
|
"gte": "Gte",
|
||||||
|
"lt": "Lt",
|
||||||
|
"lte": "Lte",
|
||||||
|
"in": "In",
|
||||||
|
"nin": "NotIn",
|
||||||
|
}
|
||||||
|
|
||||||
def _convert_filters(self, filters: Optional[Dict]):
|
def _convert_filters(self, filters: Optional[Dict]):
|
||||||
"""
|
"""
|
||||||
Convert mem0 filters to Turbopuffer filter format.
|
Convert mem0 filters to Turbopuffer filter format.
|
||||||
@@ -143,10 +155,14 @@ class TurbopufferDB(VectorStoreBase):
|
|||||||
conditions = []
|
conditions = []
|
||||||
for key, value in filters.items():
|
for key, value in filters.items():
|
||||||
if isinstance(value, dict):
|
if isinstance(value, dict):
|
||||||
if "gte" in value:
|
for op, operand in value.items():
|
||||||
conditions.append((key, "Gte", value["gte"]))
|
tpuf_op = self.OPERATOR_MAP.get(op)
|
||||||
if "lte" in value:
|
if tpuf_op is None:
|
||||||
conditions.append((key, "Lte", value["lte"]))
|
raise ValueError(
|
||||||
|
f"Unsupported filter operator '{op}' for field '{key}'. "
|
||||||
|
f"Supported operators: {sorted(self.OPERATOR_MAP)}"
|
||||||
|
)
|
||||||
|
conditions.append((key, tpuf_op, operand))
|
||||||
else:
|
else:
|
||||||
conditions.append((key, "Eq", value))
|
conditions.append((key, "Eq", value))
|
||||||
|
|
||||||
|
|||||||
@@ -276,6 +276,40 @@ class TestConvertFilters:
|
|||||||
assert ("user_id", "Eq", "u1") in conditions
|
assert ("user_id", "Eq", "u1") in conditions
|
||||||
assert ("score", "Gte", 0.5) in conditions
|
assert ("score", "Gte", 0.5) in conditions
|
||||||
|
|
||||||
|
def test_gt_operator_not_dropped(self, db):
|
||||||
|
"""Regression: {"gt": ...} was silently dropped, returning unfiltered results."""
|
||||||
|
result = db._convert_filters({"age": {"gt": 18}})
|
||||||
|
assert result == ("age", "Gt", 18)
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"op,expected_token",
|
||||||
|
[
|
||||||
|
("eq", "Eq"),
|
||||||
|
("ne", "NotEq"),
|
||||||
|
("gt", "Gt"),
|
||||||
|
("gte", "Gte"),
|
||||||
|
("lt", "Lt"),
|
||||||
|
("lte", "Lte"),
|
||||||
|
("in", "In"),
|
||||||
|
("nin", "NotIn"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_all_operators_mapped(self, db, op, expected_token):
|
||||||
|
operand = [1, 2] if op in ("in", "nin") else 5
|
||||||
|
result = db._convert_filters({"age": {op: operand}})
|
||||||
|
assert result == ("age", expected_token, operand)
|
||||||
|
|
||||||
|
def test_multiple_operators_on_one_field(self, db):
|
||||||
|
result = db._convert_filters({"age": {"gt": 18, "lt": 65}})
|
||||||
|
assert result[0] == "And"
|
||||||
|
conditions = result[1]
|
||||||
|
assert ("age", "Gt", 18) in conditions
|
||||||
|
assert ("age", "Lt", 65) in conditions
|
||||||
|
|
||||||
|
def test_unknown_operator_raises(self, db):
|
||||||
|
with pytest.raises(ValueError, match="Unsupported filter operator"):
|
||||||
|
db._convert_filters({"age": {"between": [1, 2]}})
|
||||||
|
|
||||||
|
|
||||||
# ── search ───────────────────────────────────────────────────────────
|
# ── search ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user