Repository: Quansight/ragna (python)

Issue:
metadata filter
### Feature description

Add the metadata filter abstraction proposed in #256. It does not have to be connected to the `ragna.core.Chat` as this can be done in a follow-up PR. I just want to have it well tested before moving forward with the rest of the design.

The `MetadataFilter` object should have the following properties:

- It needs to support the following operators:
  - logical and
  - logical or
  - equal
  - not equal
  - less than
  - less or equal than
  - greater than
  - greater or equal than
  - in
  - not in

  I have surveyed the ecosystem and this set of operators is supported by Chroma, LanceDB, Pinecone and Vectara. Meaning, this is a good base set.
- It needs to have an "escape hatch operator", which can be used to supply a raw filter to the source storage. This is useful in case the user only uses one source storage (production use case) and it supports more operations than just the base case. Meaning, we are not restricting users to the set of operators above.
- It needs to be fully JSON (de-)serializable

### Value and/or benefit

First building block for allowing users to have a corpus of documents instead of uploading new ones for every chat.

### Anything else?

I have a complete implementation already in https://gist.github.com/pmeier/38ee90be6c30ecdf9bbec086a0dabafe. Unless there are some concerns, it can just be ported.

Pull request:
add metadata filter
Closes #422.

Patch:
diff --git a/ragna/core/__init__.py b/ragna/core/__init__.py
index 0f4b4bd..0f46baa 100644
--- a/ragna/core/__init__.py
+++ b/ragna/core/__init__.py
@@ -11,6 +11,8 @@ __all__ = [
     "LocalDocument",
     "Message",
     "MessageRole",
+    "MetadataFilter",
+    "MetadataOperator",
     "PackageRequirement",
     "Page",
     "PdfDocumentHandler",
@@ -42,6 +44,7 @@ from ._document import (
     PlainTextDocumentHandler,
     PptxDocumentHandler,
 )
+from ._metadata_filter import MetadataFilter, MetadataOperator
 
 # isort: split
 
diff --git a/ragna/core/_metadata_filter.py b/ragna/core/_metadata_filter.py
new file mode 100644
index 0000000..206bef5
--- /dev/null
+++ b/ragna/core/_metadata_filter.py
@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+import enum
+import json
+import textwrap
+from typing import Any, Literal, Sequence, cast
+
+
+class MetadataOperator(enum.Enum):
+    RAW = enum.auto()
+    AND = enum.auto()
+    OR = enum.auto()
+    EQ = enum.auto()
+    NE = enum.auto()
+    LT = enum.auto()
+    LE = enum.auto()
+    GT = enum.auto()
+    GE = enum.auto()
+    IN = enum.auto()
+    NOT_IN = enum.auto()
+
+
+class MetadataFilter:
+    # These are just to be consistent. The actual values have no effect.
+    _RAW_KEY = "filter"
+    _CHILDREN_KEY = "children"
+
+    def __init__(self, operator: MetadataOperator, key: str, value: Any) -> None:
+        self.operator = operator
+        self.key = key
+        self.value = value
+
+    def __repr__(self) -> str:
+        if self.operator is MetadataOperator.RAW:
+            return f"{self.operator.name}({self.value!r})"
+        elif self.operator in {MetadataOperator.AND, MetadataOperator.OR}:
+            return "\n".join(
+                [
+                    f"{self.operator.name}(",
+                    *[
+                        f"{textwrap.indent(repr(child), prefix=' ' * 2)},"
+                        for child in self.value
+                    ],
+                    ")",
+                ]
+            )
+        else:
+            return f"{self.operator.name}({self.key!r}, {self.value!r})"
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, MetadataFilter):
+            return NotImplemented
+
+        if self.operator != other.operator:
+            return False
+
+        if self.operator in {MetadataOperator.AND, MetadataOperator.OR}:
+            if len(self.value) != len(other.value):
+                return False
+
+            for self_child, other_child in zip(self.value, other.value):
+                if self_child != other_child:
+                    return False
+
+            return True
+        else:
+            return (self.key == other.key) and (self.value == other.value)
+
+    def _to_json(self) -> dict[str, Any]:
+        if self.operator in {MetadataOperator.AND, MetadataOperator.OR}:
+            value = [child._to_json() for child in self.value]
+        else:
+            value = self.value
+
+        return {self.operator.name: {self.key: value}}
+
+    def to_json(self) -> str:
+        return json.dumps(self._to_json())
+
+    @classmethod
+    def _from_json(cls, json_obj: dict[str, Any]) -> MetadataFilter:
+        operator, key_value = next(iter(json_obj.items()))
+        operator = MetadataOperator.__members__[operator]
+        key_value = cast(dict[str, Any], key_value)
+        key, value = next(iter(key_value.items()))
+
+        if operator in {MetadataOperator.AND, MetadataOperator.OR}:
+            value = [cls._from_json(child) for child in value]
+
+        return cls(operator, key, value)
+
+    @classmethod
+    def from_json(cls, json_str: str) -> MetadataFilter:
+        return cls._from_json(json.loads(json_str))
+
+    @classmethod
+    def raw(cls, value: Any) -> MetadataFilter:
+        return cls(MetadataOperator.RAW, cls._RAW_KEY, value)
+
+    @staticmethod
+    def _flatten(
+        operator: Literal[MetadataOperator.OR, MetadataOperator.AND],
+        children: Sequence[MetadataFilter],
+    ) -> list[MetadataFilter]:
+        flat_children = []
+        for child in children:
+            if child.operator == operator:
+                flat_children.extend(child.value)
+            else:
+                flat_children.append(child)
+        return flat_children
+
+    @classmethod
+    def and_(cls, children: Sequence[MetadataFilter]) -> MetadataFilter:
+        return cls(
+            MetadataOperator.AND,
+            cls._CHILDREN_KEY,
+            cls._flatten(MetadataOperator.AND, children),
+        )
+
+    @classmethod
+    def or_(cls, children: list[MetadataFilter]) -> MetadataFilter:
+        return cls(
+            MetadataOperator.OR,
+            cls._CHILDREN_KEY,
+            cls._flatten(MetadataOperator.OR, children),
+        )
+
+    @classmethod
+    def eq(cls, key: str, value: Any) -> MetadataFilter:
+        return cls(MetadataOperator.EQ, key, value)
+
+    @classmethod
+    def ne(cls, key: str, value: Any) -> MetadataFilter:
+        return cls(MetadataOperator.NE, key, value)
+
+    @classmethod
+    def lt(cls, key: str, value: Any) -> MetadataFilter:
+        return cls(MetadataOperator.LT, key, value)
+
+    @classmethod
+    def le(cls, key: str, value: Any) -> MetadataFilter:
+        return cls(MetadataOperator.LE, key, value)
+
+    @classmethod
+    def gt(cls, key: str, value: Any) -> MetadataFilter:
+        return cls(MetadataOperator.GT, key, value)
+
+    @classmethod
+    def ge(cls, key: str, value: Any) -> MetadataFilter:
+        return cls(MetadataOperator.GE, key, value)
+
+    @classmethod
+    def in_(cls, key: str, value: Any) -> MetadataFilter:
+        return cls(MetadataOperator.IN, key, value)
+
+    @classmethod
+    def not_in(cls, key: str, value: Any) -> MetadataFilter:
+        return cls(MetadataOperator.NOT_IN, key, value)


Test patch:
diff --git a/tests/core/test_metadata_filter.py b/tests/core/test_metadata_filter.py
new file mode 100644
index 0000000..7b25070
--- /dev/null
+++ b/tests/core/test_metadata_filter.py
@@ -0,0 +1,184 @@
+import pytest
+
+from ragna.core import MetadataFilter, MetadataOperator
+
+metadata_filters = pytest.mark.parametrize(
+    "metadata_filter",
+    [
+        pytest.param(MetadataFilter.raw("raw"), id="raw"),
+        pytest.param(
+            MetadataFilter.and_(
+                [
+                    MetadataFilter.raw("raw"),
+                    MetadataFilter.eq("key", "value"),
+                ]
+            ),
+            id="and",
+        ),
+        pytest.param(
+            MetadataFilter.or_(
+                [
+                    MetadataFilter.raw("raw"),
+                    MetadataFilter.eq("key", "value"),
+                ]
+            ),
+            id="or",
+        ),
+        pytest.param(
+            MetadataFilter.and_(
+                [
+                    MetadataFilter.raw("raw"),
+                    MetadataFilter.or_(
+                        [
+                            MetadataFilter.eq("key", "value"),
+                            MetadataFilter.ne("other_key", "other_value"),
+                        ]
+                    ),
+                ]
+            ),
+            id="and-nested",
+        ),
+        pytest.param(
+            MetadataFilter.or_(
+                [
+                    MetadataFilter.raw("raw"),
+                    MetadataFilter.and_(
+                        [
+                            MetadataFilter.eq("key", "value"),
+                            MetadataFilter.ne("other_key", "other_value"),
+                        ]
+                    ),
+                ]
+            ),
+            id="or-nested",
+        ),
+        pytest.param(MetadataFilter.eq("key", "value"), id="eq"),
+        pytest.param(MetadataFilter.ne("key", "value"), id="ne"),
+        pytest.param(MetadataFilter.lt("key", 1), id="lt"),
+        pytest.param(MetadataFilter.le("key", 0), id="le"),
+        pytest.param(MetadataFilter.gt("key", 1), id="gt"),
+        pytest.param(MetadataFilter.ge("key", 0), id="ge"),
+        pytest.param(MetadataFilter.in_("key", ["foo", "bar"]), id="in"),
+        pytest.param(MetadataFilter.not_in("key", ["foo", "bar"]), id="not_in"),
+    ],
+)
+
+
+@metadata_filters
+def test_self_similarity(metadata_filter):
+    assert metadata_filter == metadata_filter
+
+
+@metadata_filters
+def test_json_roundtrip(metadata_filter):
+    assert MetadataFilter.from_json(metadata_filter.to_json()) == metadata_filter
+
+
+@metadata_filters
+def test_repr_smoke(metadata_filter):
+    repr(metadata_filter)
+
+
+def test_raw():
+    value = "SENTINEL"
+    metadata_filter = MetadataFilter.raw(value)
+    assert metadata_filter.operator is MetadataOperator.RAW
+    assert metadata_filter.value == value
+
+
+def test_and():
+    children = [MetadataFilter.raw("SENTINEL1"), MetadataFilter.raw("SENTINEL2")]
+    metadata_filter = MetadataFilter.and_(children)
+    assert metadata_filter.operator is MetadataOperator.AND
+    assert metadata_filter.value == children
+
+
+def test_or():
+    children = [MetadataFilter.raw("SENTINEL1"), MetadataFilter.raw("SENTINEL2")]
+    metadata_filter = MetadataFilter.or_(children)
+    assert metadata_filter.operator is MetadataOperator.OR
+    assert metadata_filter.value == children
+
+
+@pytest.mark.parametrize(
+    "fn",
+    [
+        pytest.param(MetadataFilter.and_, id="and"),
+        pytest.param(MetadataFilter.or_, id="or"),
+    ],
+)
+def test_flatten(fn):
+    children = [MetadataFilter.raw(f"SENTINEL{idx}") for idx in range(3)]
+    assert fn([children[0], fn(children[1:])]) == fn(children)
+
+
+def test_eq():
+    key = "KEY_SENTINEL"
+    value = "VALUE_SENTINEL"
+    metadata_filter = MetadataFilter.eq(key, value)
+    assert metadata_filter.operator is MetadataOperator.EQ
+    assert metadata_filter.key == key
+    assert metadata_filter.value == value
+
+
+def test_ne():
+    key = "KEY_SENTINEL"
+    value = "VALUE_SENTINEL"
+    metadata_filter = MetadataFilter.ne(key, value)
+    assert metadata_filter.operator is MetadataOperator.NE
+    assert metadata_filter.key == key
+    assert metadata_filter.value == value
+
+
+def test_lt():
+    key = "KEY_SENTINEL"
+    value = 31476
+    metadata_filter = MetadataFilter.lt(key, value)
+    assert metadata_filter.operator is MetadataOperator.LT
+    assert metadata_filter.key == key
+    assert metadata_filter.value == value
+
+
+def test_le():
+    key = "KEY_SENTINEL"
+    value = 31476
+    metadata_filter = MetadataFilter.le(key, value)
+    assert metadata_filter.operator is MetadataOperator.LE
+    assert metadata_filter.key == key
+    assert metadata_filter.value == value
+
+
+def test_gt():
+    key = "KEY_SENTINEL"
+    value = 31476
+    metadata_filter = MetadataFilter.gt(key, value)
+    assert metadata_filter.operator is MetadataOperator.GT
+    assert metadata_filter.key == key
+    assert metadata_filter.value == value
+
+
+def test_ge():
+    key = "KEY_SENTINEL"
+    value = 31476
+    metadata_filter = MetadataFilter.ge(key, value)
+    assert metadata_filter.operator is MetadataOperator.GE
+    assert metadata_filter.key == key
+    assert metadata_filter.value == value
+
+
+def test_in():
+    key = "KEY_SENTINEL"
+    value = ["foo", "bar"]
+    metadata_filter = MetadataFilter.in_(key, value)
+    assert metadata_filter.operator is MetadataOperator.IN
+    assert metadata_filter.key == key
+    assert metadata_filter.value == value
+
+
+def test_not_in():
+    key = "KEY_SENTINEL"
+    value = ["foo", "bar"]
+    metadata_filter = MetadataFilter.not_in(key, value)
+    assert metadata_filter.operator is MetadataOperator.NOT_IN
+    assert metadata_filter.key == key
+    assert metadata_filter.value == value