Repository: Quansight/ragna
Pull request #476 (2024-08-05 11:35:04)

support pydantic (de)serialization on MetadataFilter
The is a prerequisite for using `MetadataFilter` as part of FastAPI. I've scrapped the actual JSON serialization, because that would create JSON inside JSON with pydantic and pydantic knows how to serialize dicts anyway.

nenb: Also, I haven't looked into it, but a lot of tests seem to fail. Not sure what is going on...
pmeier: > It then occurred to me, do we really need the `children` key?

Nope

https://github.com/Quansight/ragna/blob/ad9c1307b31a8e67f09411acecc84e0b4a58ba57/ragna/core/_metadata_filter.py#L23-L26

Per comment, the same for the raw filter. So you can break down your example even more to

```
{'OR': [{'RAW': 'raw'}, {'GE': {'key': 1}}]}
```

That is as concise as it gets. However we also need to consider the use case of translating it. As discussed offline, I dislike the bits you highlighted as well. I invite everyone to come up with a better scheme and I'm happy to adopt it.

> Also, I haven't looked into it, but a lot of tests seem to fail. Not sure what is going on...

Just minor stuff that should be fixed by #461.
pmeier: Thinking about this more, I realized that there is no requirement for the primitive version to have a 1-to-1 relationship to the Python object. Meaning, we can go for the simplifications proposed in https://github.com/Quansight/ragna/pull/476#pullrequestreview-2219217627 and https://github.com/Quansight/ragna/pull/476#issuecomment-2269759736 and still keep one Python object that always has the `operator`, `key`, and `value` fields.

I implemented that in 70ac0b51fa2865868e79fb4e33667fcb67e32b8a. Here is the example from above:

```python
import json

from ragna.core import MetadataFilter

metadata_filter = MetadataFilter.or_(
    [
        MetadataFilter.raw("raw"),
        MetadataFilter.ge("key", 1),
    ]
)

print(metadata_filter)
print(json.dumps(metadata_filter.to_primitive()))
```

```
OR(
  RAW('raw'),
  GE('key', 1),
)
{"OR": [{"RAW": "raw"}, {"GE": {"key": 1}}]}
```

Patch:
diff --git a/ragna/core/_metadata_filter.py b/ragna/core/_metadata_filter.py
index 206bef5..89b3154 100644
--- a/ragna/core/_metadata_filter.py
+++ b/ragna/core/_metadata_filter.py
@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import enum
-import json
 import textwrap
-from typing import Any, Literal, Sequence, cast
+from typing import Any, Literal, Sequence, Union, cast
+
+import pydantic
+import pydantic_core
 
 
 class MetadataOperator(enum.Enum):
@@ -21,10 +23,6 @@ class MetadataOperator(enum.Enum):
 
 
 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
@@ -66,36 +64,71 @@ class MetadataFilter:
         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:
+    def to_primitive(self) -> dict[str, Any]:
+        if self.operator is MetadataOperator.RAW:
             value = self.value
+        elif self.operator in {MetadataOperator.AND, MetadataOperator.OR}:
+            value = [child.to_primitive() for child in self.value]
+        else:
+            value = {self.key: self.value}
 
-        return {self.operator.name: {self.key: value}}
-
-    def to_json(self) -> str:
-        return json.dumps(self._to_json())
+        return {self.operator.name: value}
 
     @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]
+    def from_primitive(cls, obj: dict[str, Any]) -> MetadataFilter:
+        operator, value = next(iter(obj.items()))
+        operator = MetadataOperator.__members__[operator.upper()]
+        if operator is MetadataOperator.RAW:
+            key = ""
+        elif operator in {MetadataOperator.AND, MetadataOperator.OR}:
+            key = ""
+            value = [cls.from_primitive(child) for child in value]
+        else:
+            key, value = next(iter(cast(dict[str, Any], value).items()))
 
         return cls(operator, key, value)
 
     @classmethod
-    def from_json(cls, json_str: str) -> MetadataFilter:
-        return cls._from_json(json.loads(json_str))
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
+    ) -> pydantic_core.CoreSchema:
+        def validate(value: Union[MetadataFilter, dict[str, Any]]) -> MetadataFilter:
+            if isinstance(value, MetadataFilter):
+                return value
+            else:
+                return cls.from_primitive(value)
+
+        def serialize(value: Union[MetadataFilter, dict[str, Any]]) -> dict[str, Any]:
+            if isinstance(value, MetadataFilter):
+                return value.to_primitive()
+            else:
+                return value
+
+        dict_schema = pydantic_core.core_schema.dict_schema(
+            keys_schema=pydantic_core.core_schema.literal_schema(
+                list(MetadataOperator.__members__)
+            ),
+        )
+        return pydantic_core.core_schema.no_info_after_validator_function(
+            # allowed input schemas
+            schema=pydantic_core.core_schema.union_schema(
+                [
+                    pydantic_core.core_schema.is_instance_schema(cls),
+                    dict_schema,
+                ]
+            ),
+            # function to be applied after the input schema check
+            function=validate,
+            serialization=pydantic_core.core_schema.plain_serializer_function_ser_schema(
+                function=serialize,
+                return_schema=dict_schema,
+                when_used="json",
+            ),
+        )
 
     @classmethod
     def raw(cls, value: Any) -> MetadataFilter:
-        return cls(MetadataOperator.RAW, cls._RAW_KEY, value)
+        return cls(MetadataOperator.RAW, "", value)
 
     @staticmethod
     def _flatten(
@@ -114,7 +147,7 @@ class MetadataFilter:
     def and_(cls, children: Sequence[MetadataFilter]) -> MetadataFilter:
         return cls(
             MetadataOperator.AND,
-            cls._CHILDREN_KEY,
+            "",
             cls._flatten(MetadataOperator.AND, children),
         )
 
@@ -122,7 +155,7 @@ class MetadataFilter:
     def or_(cls, children: list[MetadataFilter]) -> MetadataFilter:
         return cls(
             MetadataOperator.OR,
-            cls._CHILDREN_KEY,
+            "",
             cls._flatten(MetadataOperator.OR, children),
         )
 


Test patch:
diff --git a/tests/core/test_metadata_filter.py b/tests/core/test_metadata_filter.py
index 7b25070..ae142b5 100644
--- a/tests/core/test_metadata_filter.py
+++ b/tests/core/test_metadata_filter.py
@@ -1,3 +1,4 @@
+import pydantic
 import pytest
 
 from ragna.core import MetadataFilter, MetadataOperator
@@ -70,8 +71,22 @@ def test_self_similarity(metadata_filter):
 
 
 @metadata_filters
-def test_json_roundtrip(metadata_filter):
-    assert MetadataFilter.from_json(metadata_filter.to_json()) == metadata_filter
+def test_to_from_primitive_roundtrip(metadata_filter):
+    assert (
+        MetadataFilter.from_primitive(metadata_filter.to_primitive()) == metadata_filter
+    )
+
+
+@metadata_filters
+def test_pydantic(metadata_filter):
+    class Model(pydantic.BaseModel):
+        mf: MetadataFilter
+
+    assert Model(mf=metadata_filter).mf == metadata_filter
+    assert Model(mf=metadata_filter.to_primitive()).mf == metadata_filter
+    assert Model(mf=metadata_filter).model_dump(mode="json") == {
+        "mf": metadata_filter.to_primitive()
+    }
 
 
 @metadata_filters