Add support for filter comparisons (MemMachine#676)

Author: o-love

src/memmachine/common/episode_store/episode_sqlalchemy_store.py

@@ -30,7 +30,6 @@
 from sqlalchemy.sql import Select
 from sqlalchemy.sql.elements import ColumnElement
 
-from memmachine.common.data_types import FilterablePropertyValue
 from memmachine.common.episode_store.episode_model import Episode as EpisodeE
 from memmachine.common.episode_store.episode_model import EpisodeEntry, EpisodeType
 from memmachine.common.episode_store.episode_storage import EpisodeIdT, EpisodeStorage
@@ -45,6 +44,7 @@
 from memmachine.common.filter.filter_parser import (
     Or as FilterOr,
 )
+from memmachine.common.filter.sql_filter_util import parse_sql_filter
 
 
 class BaseEpisodeStore(DeclarativeBase):
@@ -246,7 +246,9 @@ def _apply_episode_filter(
         filters: list[ColumnElement[bool]] = []
 
         if filter_expr is not None:
-            filters.append(self._compile_episode_filter_expr(filter_expr))
+            parsed_filter = self._compile_episode_filter_expr(filter_expr)
+            if parsed_filter is not None:
+                filters.append(parsed_filter)
 
         if start_time is not None:
             filters.append(Episode.created_at >= start_time)
@@ -266,40 +268,18 @@ def _apply_episode_filter(
     def _compile_episode_comparison_expr(
         self,
         expr: FilterComparison,
-    ) -> ColumnElement[bool]:
+    ) -> ColumnElement[bool] | None:
         column, is_metadata = self._resolve_episode_field(expr.field)
 
-        if column is None:
-            raise ValueError(f"Unsupported episode filter field: {expr.field}")
-
-        if expr.op == "=":
-            value = expr.value
-            if isinstance(value, list):
-                raise ValueError("'=' comparison cannot accept list values")
-            if is_metadata:
-                value = self._normalize_metadata_value(value)
-                return column == value
-            return column == value
-
-        if expr.op == "in":
-            if not isinstance(expr.value, list):
-                raise ValueError("IN comparison requires a list of values")
-
-            values = expr.value
-            if is_metadata:
-                values = [self._normalize_metadata_value(v) for v in values]
-
-            return column.in_(values)
-
-        if expr.op == "is_null":
-            return column.is_(None)
-
-        if expr.op == "is_not_null":
-            return column.is_not(None)
-
-        raise ValueError(f"Unsupported operator: {expr.op}")
+        return parse_sql_filter(
+            column=column,
+            is_metadata=is_metadata,
+            expr=expr,
+        )
 
-    def _compile_episode_filter_expr(self, expr: FilterExpr) -> ColumnElement[bool]:
+    def _compile_episode_filter_expr(
+        self, expr: FilterExpr
+    ) -> ColumnElement[bool] | None:
         if isinstance(expr, FilterComparison):
             return self._compile_episode_comparison_expr(expr)
 
@@ -315,14 +295,6 @@ def _compile_episode_filter_expr(self, expr: FilterExpr) -> ColumnElement[bool]:
 
         raise TypeError(f"Unsupported filter expression type: {type(expr)!r}")
 
-    @staticmethod
-    def _normalize_metadata_value(
-        value: FilterablePropertyValue | list[FilterablePropertyValue],
-    ) -> str:
-        if isinstance(value, bool):
-            return "true" if value else "false"
-        return "" if value is None else str(value)
-
     @staticmethod
     def _resolve_episode_field(
         field: str,

src/memmachine/common/filter/filter_parser.py

@@ -20,7 +20,7 @@ class Comparison(FilterExpr):
     """Filter comparison of a field against a value or list of values."""
 
     field: str
-    op: str  # "=", "in", "is_null", "is_not_null"
+    op: str  # "=", "in", ">", "<", ">=", "<=", "is_null", "is_not_null"
     value: FilterablePropertyValue | list[FilterablePropertyValue]
 
 
@@ -57,7 +57,11 @@ class Token(NamedTuple):
     ("LPAREN", r"\("),
     ("RPAREN", r"\)"),
     ("COMMA", r","),
+    ("GE", r">="),
+    ("LE", r"<="),
     ("EQ", r"="),
+    ("GT", r">"),
+    ("LT", r"<"),
     ("STRING", r"'[^']*'"),
     ("IDENT", r"[A-Za-z0-9_\.]+"),
     ("WS", r"\s+"),
@@ -155,10 +159,12 @@ def _parse_comparison(self) -> FilterExpr:
         field_tok = self._expect("IDENT")
         field = field_tok.value
 
-        if self._accept("EQ"):
-            # field = value
+        op_tok = self._accept("EQ", "GE", "LE", "GT", "LT")
+        if op_tok:
+            # field =/>=/>/</<= value
             value = self._parse_value()
-            return Comparison(field=field, op="=", value=value)
+            op = {"EQ": "=", "GE": ">=", "LE": "<=", "GT": ">", "LT": "<"}[op_tok.type]
+            return Comparison(field=field, op=op, value=value)
 
         if self._accept("IN"):
             self._expect("LPAREN")
@@ -179,7 +185,9 @@ def _parse_comparison(self) -> FilterExpr:
             op = "is_not_null" if negate else "is_null"
             return Comparison(field=field, op=op, value=None)
 
-        raise FilterParseError(f"Expected '=' or IN after field {field}")
+        raise FilterParseError(
+            f"Expected comparison operator (=, IN, >, <, >=, <=, IS) after field {field}"
+        )
 
     def _parse_value(self) -> FilterablePropertyValue:
         tok = self._expect("IDENT", "STRING")

src/memmachine/common/filter/sql_filter_util.py

@@ -0,0 +1,72 @@
+"""SQLAlchemy utilities for FilterExpr."""
+
+import logging
+from typing import Any
+
+from sqlalchemy import ColumnElement
+from sqlalchemy.orm import InstrumentedAttribute, MappedColumn
+
+from memmachine.common.data_types import FilterablePropertyValue
+from memmachine.common.filter.filter_parser import Comparison
+
+logger = logging.getLogger(__name__)
+
+
+def _normalize_metadata_value(
+    value: FilterablePropertyValue | list[FilterablePropertyValue],
+) -> str:
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    return "" if value is None else str(value)
+
+
+def _ensure_scalar_value(
+    value: FilterablePropertyValue | list[FilterablePropertyValue], op: str
+) -> FilterablePropertyValue:
+    if isinstance(value, list):
+        raise TypeError(f"'{op}' comparison cannot accept list values")
+    return value
+
+
+def _ensure_list_value(
+    value: FilterablePropertyValue | list[FilterablePropertyValue],
+) -> list[FilterablePropertyValue]:
+    if not isinstance(value, list):
+        raise TypeError("IN comparison requires a list of values")
+    return value
+
+
+def parse_sql_filter(
+    column: MappedColumn[Any] | InstrumentedAttribute[Any] | None,
+    is_metadata: bool,
+    expr: Comparison,
+) -> ColumnElement[bool] | None:
+    """Parse a FilterExpr comparison into an SQLAlchemy boolean expression."""
+    if column is None:
+        logger.warning("Unsupported feature filter field: %s", expr.field)
+        return None
+
+    op = expr.op
+    normalize = _normalize_metadata_value if is_metadata else lambda v: v
+
+    match op:
+        case "is_null":
+            return column.is_(None)
+        case "is_not_null":
+            return column.is_not(None)
+        case "in":
+            values = _ensure_list_value(expr.value)
+            if is_metadata:
+                values = [normalize(v) for v in values]
+            return column.in_(values)
+        case ">" | "<" | ">=" | "<=" | "=":
+            value = normalize(_ensure_scalar_value(expr.value, op))
+            return {
+                ">": column > value,
+                "<": column < value,
+                ">=": column >= value,
+                "<=": column <= value,
+                "=": column == value,
+            }[op]
+        case _:
+            raise ValueError(f"Unsupported operator: {expr.op}")

src/memmachine/common/vector_graph_store/neo4j_vector_graph_store.py

@@ -1616,10 +1616,14 @@ def _render_filter_expr(
             params: dict[
                 str, FilterablePropertyValue | list[FilterablePropertyValue]
             ] = {}
-            if expr.op == "=":
+            if expr.op in (">", "<", ">=", "<=", "="):
                 if isinstance(expr.value, list):
-                    raise ValueError("'=' comparison cannot accept list values")
-                condition = f"{field_ref} = ${query_value_parameter}.{param_name}"
+                    raise ValueError(
+                        f"'{expr.op}' comparison cannot accept list values"
+                    )
+                condition = (
+                    f"{field_ref} {expr.op} ${query_value_parameter}.{param_name}"
+                )
                 params[param_name] = expr.value
             elif expr.op == "in":
                 if not isinstance(expr.value, list):

src/memmachine/semantic_memory/storage/neo4j_semantic_storage.py

@@ -915,28 +915,7 @@ def _render_filter_expr(
     ) -> tuple[str, dict[str, Any]]:
         if isinstance(expr, FilterComparison):
             field_ref = self._resolve_field_reference(alias, expr.field)
-            params: dict[str, Any]
-            if expr.op == "=":
-                if isinstance(expr.value, list):
-                    raise ValueError("'=' comparison cannot accept list values")
-                param_name = self._next_filter_param()
-                condition = f"{field_ref} = ${param_name}"
-                params = {param_name: expr.value}
-            elif expr.op == "in":
-                if not isinstance(expr.value, list):
-                    raise ValueError("IN comparison requires a list of values")
-                param_name = self._next_filter_param()
-                condition = f"{field_ref} IN ${param_name}"
-                params = {param_name: expr.value}
-            elif expr.op == "is_null":
-                condition = f"{field_ref} IS NULL"
-                params = {}
-            elif expr.op == "is_not_null":
-                condition = f"{field_ref} IS NOT NULL"
-                params = {}
-            else:
-                raise ValueError(f"Unsupported operator: {expr.op}")
-            return condition, params
+            return self._render_comparison_condition(field_ref, expr)
         if isinstance(expr, FilterAnd):
             left_cond, left_params = self._render_filter_expr(alias, expr.left)
             right_cond, right_params = self._render_filter_expr(alias, expr.right)
@@ -951,6 +930,32 @@ def _render_filter_expr(
             return condition, left_params
         raise TypeError(f"Unsupported filter expression type: {type(expr)!r}")
 
+    def _render_comparison_condition(
+        self, field_ref: str, expr: FilterComparison
+    ) -> tuple[str, dict[str, Any]]:
+        op = expr.op
+        params: dict[str, Any] = {}
+
+        if op == "in":
+            if not isinstance(expr.value, list):
+                raise ValueError("IN comparison requires a list of values")
+            param = self._next_filter_param()
+            return f"{field_ref} IN ${param}", {param: expr.value}
+
+        if op in (">", "<", ">=", "<=", "="):
+            if isinstance(expr.value, list):
+                raise ValueError(f"'{op}' comparison cannot accept list values")
+            param = self._next_filter_param()
+            return f"{field_ref} {op} ${param}", {param: expr.value}
+
+        if op == "is_null":
+            return f"{field_ref} IS NULL", params
+
+        if op == "is_not_null":
+            return f"{field_ref} IS NOT NULL", params
+
+        raise ValueError(f"Unsupported operator: {op}")
+
     def _resolve_field_reference(self, alias: str, field: str) -> str:
         if field.startswith(("m.", "metadata.")):
             key = field.split(".", 1)[1]

src/memmachine/semantic_memory/storage/sqlalchemy_pgvector_semantic.py

@@ -30,10 +30,15 @@
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.engine import Connection
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
-from sqlalchemy.orm import DeclarativeBase, aliased, mapped_column
+from sqlalchemy.orm import (
+    DeclarativeBase,
+    InstrumentedAttribute,
+    MappedColumn,
+    aliased,
+    mapped_column,
+)
 from sqlalchemy.sql import Delete, Select, func
 
-from memmachine.common.data_types import FilterablePropertyValue
 from memmachine.common.episode_store.episode_model import EpisodeIdT
 from memmachine.common.errors import InvalidArgumentError, ResourceNotFoundError
 from memmachine.common.filter.filter_parser import (
@@ -48,6 +53,7 @@
 from memmachine.common.filter.filter_parser import (
     Or as FilterOr,
 )
+from memmachine.common.filter.sql_filter_util import parse_sql_filter
 from memmachine.semantic_memory.semantic_model import SemanticFeature, SetIdT
 from memmachine.semantic_memory.storage.storage_base import (
     FeatureIdT,
@@ -616,35 +622,11 @@ def _compile_feature_comparison_expr(
     ) -> ColumnElement[bool] | None:
         column, is_metadata = self._resolve_feature_field(table, expr.field)
 
-        if column is None:
-            logger.warning("Unsupported feature filter field: %s", expr.field)
-            return None
-
-        if expr.op == "=":
-            value = expr.value
-            if isinstance(value, list):
-                raise ValueError("'=' comparison cannot accept list values")
-            if is_metadata:
-                value = self._normalize_metadata_value(value)
-                return column == value
-            return column == value
-
-        if expr.op == "in":
-            if not isinstance(expr.value, list):
-                raise ValueError("IN comparison requires a list of values")
-
-            values = expr.value
-            if is_metadata:
-                values = [self._normalize_metadata_value(v) for v in values]
-            return column.in_(values)
-
-        if expr.op == "is_null":
-            return column.is_(None)
-
-        if expr.op == "is_not_null":
-            return column.is_not(None)
-
-        raise ValueError(f"Unsupported operator: {expr.op}")
+        return parse_sql_filter(
+            column=column,
+            is_metadata=is_metadata,
+            expr=expr,
+        )
 
     def _compile_feature_filter_expr(
         self,
@@ -668,19 +650,13 @@ def _compile_feature_filter_expr(
 
         raise TypeError(f"Unsupported filter expression type: {type(expr)!r}")
 
-    @staticmethod
-    def _normalize_metadata_value(
-        value: FilterablePropertyValue | list[FilterablePropertyValue],
-    ) -> str:
-        if isinstance(value, bool):
-            return "true" if value else "false"
-        return "" if value is None else str(value)
-
     @staticmethod
     def _resolve_feature_field(
         table: type[Feature],
         field: str,
-    ) -> tuple[Any, bool] | tuple[None, bool]:
+    ) -> (
+        tuple[MappedColumn[Any] | InstrumentedAttribute[Any], bool] | tuple[None, bool]
+    ):
         normalized = field
         field_mapping = {
             "set_id": table.set_id,