fix knn setup

Author: sethbuff

aredis_om/__init__.py

@@ -16,25 +16,3 @@
     RedisModelError,
     VectorFieldOptions,
 )
-
-
-__all__ = [
-    "redis",
-    "get_redis_connection",
-    "Field",
-    "HashModel",
-    "JsonModel",
-    "EmbeddedJsonModel",
-    "RedisModel",
-    "FindQuery",
-    "KNNExpression",
-    "VectorFieldOptions",
-    "has_redis_json",
-    "has_redisearch",
-    "MigrationError",
-    "Migrator",
-    "RedisModelError",
-    "NotFoundError",
-    "QueryNotSupportedError",
-    "QuerySyntaxError",
-]

aredis_om/async_redis.py

@@ -1,4 +1 @@
 from redis import asyncio as redis
-
-
-__all__ = ["redis"]

aredis_om/model/encoders.py

@@ -35,6 +35,7 @@
 from pydantic.deprecated.json import ENCODERS_BY_TYPE
 from pydantic_core import PydanticUndefined
 
+
 SetIntStr = Set[Union[int, str]]
 DictIntStrAny = Dict[Union[int, str], Any]
 

aredis_om/model/model.py

@@ -279,20 +279,24 @@ def tree(self):
 @dataclasses.dataclass
 class KNNExpression:
     k: int
-    vector_field_name: str
-    score_field_name: str
+    vector_field: "ExpressionProxy"
+    score_field: "ExpressionProxy"
     reference_vector: bytes
 
     def __str__(self):
-        return f"KNN $K @{self.vector_field_name} $knn_ref_vector"
+        return f"KNN $K @{self.vector_field_name} $knn_ref_vector AS {self.score_field_name}"
 
     @property
     def query_params(self) -> Dict[str, Union[str, bytes]]:
         return {"K": str(self.k), "knn_ref_vector": self.reference_vector}
 
     @property
-    def score_field(self) -> str:
-        return  self.score_field_name or f"_{self.vector_field_name}_score"
+    def score_field_name(self) -> str:
+        return self.score_field.field.alias
+
+    @property
+    def vector_field_name(self) -> str:
+        return self.vector_field.field.alias
 
 
 ExpressionOrNegated = Union[Expression, NegatedExpression]
@@ -438,7 +442,7 @@ def __init__(
         if sort_fields:
             self.sort_fields = self.validate_sort_fields(sort_fields)
         elif self.knn:
-            self.sort_fields = [self.knn.score_field]
+            self.sort_fields = [self.knn.score_field_name]
         else:
             self.sort_fields = []
 
@@ -511,7 +515,7 @@ def query_params(self):
     def validate_sort_fields(self, sort_fields: List[str]):
         for sort_field in sort_fields:
             field_name = sort_field.lstrip("-")
-            if self.knn and field_name == self.knn.score_field:
+            if self.knn and field_name == self.knn.score_field_name:
                 continue
             if field_name not in self.model.model_fields:  # type: ignore
                 raise QueryNotSupportedError(
@@ -1307,12 +1311,6 @@ def __new__(cls, name, bases, attrs, **kwargs):  # noqa C901
                     new_class._meta.primary_key = PrimaryKey(
                         name=field_name, field=field
                     )
-                if getattr(field, "vector_options", None) is not None:
-                    score_attr = f"_{field_name}_score"
-                    setattr(new_class, score_attr, None)
-                    new_class.__annotations__[score_attr] = Union[float, None]
-
-            new_class.model_config["from_attributes"] = True
 
         if not getattr(new_class._meta, "global_key_prefix", None):
             new_class._meta.global_key_prefix = getattr(
@@ -1371,15 +1369,10 @@ def outer_type_or_annotation(field: FieldInfo):
 
 
 class RedisModel(BaseModel, abc.ABC, metaclass=ModelMeta):
-    pk: Optional[str] = Field(default=None, primary_key=True)
+    pk: Optional[str] = Field(default=None, primary_key=True, validate_default=True)
     Meta = DefaultMeta
 
-    model_config = ConfigDict(
-        from_attributes=True,
-        arbitrary_types_allowed=True,
-        extra="allow",
-        validate_default=True,
-    )
+    model_config = ConfigDict(from_attributes=True)
 
     def __init__(__pydantic_self__, **data: Any) -> None:
         __pydantic_self__.validate_primary_key()
@@ -1518,9 +1511,6 @@ def to_string(s):
             if fields.get("$"):
                 json_fields = json.loads(fields.pop("$"))
                 doc = cls(**json_fields)
-                for k, v in fields.items():
-                    if k.startswith("__") and k.endswith("_score"):
-                        setattr(doc, k[1:], float(v))
             else:
                 doc = cls(**fields)
 

aredis_om/sync_redis.py

@@ -1,4 +1 @@
 import redis
-
-
-__all__ = ["redis"]

tests/test_json_model.py

@@ -11,7 +11,6 @@
 
 import pytest
 import pytest_asyncio
-from pydantic import field_validator
 
 from aredis_om import (
     EmbeddedJsonModel,
@@ -1269,15 +1268,18 @@ class Child(Model):
     assert child_validate.bio is None
     assert child_validate.other_name == "Maria"
 
+
 @py_test_mark_asyncio
 async def test_model_raises_error_if_inherited_from_indexed_model():
     class Model(JsonModel, index=True):
-            pass
+        pass
 
     with pytest.raises(RedisModelError):
+
         class Child(Model):
             pass
 
+
 @py_test_mark_asyncio
 async def test_non_indexed_model_raises_error_on_save():
     class Model(JsonModel):

tests/test_knn_expression.py

@@ -0,0 +1,65 @@
+# type: ignore
+import abc
+import time
+
+import pytest_asyncio
+
+from aredis_om import Field, JsonModel, KNNExpression, Migrator, VectorFieldOptions
+
+from .conftest import py_test_mark_asyncio
+
+
+vector_field_options = VectorFieldOptions.flat(
+    type=VectorFieldOptions.TYPE.FLOAT32,
+    dimension=768,
+    distance_metric=VectorFieldOptions.DISTANCE_METRIC.COSINE,
+)
+
+
+@pytest_asyncio.fixture
+async def m(key_prefix, redis):
+    class BaseJsonModel(JsonModel, abc.ABC):
+        class Meta:
+            global_key_prefix = key_prefix
+            database = redis
+
+    class Member(BaseJsonModel, index=True):
+        name: str
+        embeddings: list[list[float]] = Field([], vector_options=vector_field_options)
+        embeddings_score: float | None = None
+
+    await Migrator().run()
+
+    return Member
+
+
+@pytest_asyncio.fixture
+async def embedding_bytes():
+    return b"\x00" * 3072
+
+
+@py_test_mark_asyncio
+async def test_vector_field(m: type[JsonModel], embedding_bytes):
+    # Create a new instance of the Member model
+    member = m(name="seth", embeddings=[[0.1, 0.2, 0.3]])
+
+    # Save the member to Redis
+    mt = await member.save()
+
+    assert m.get(mt.pk)
+
+    time.sleep(1)
+
+    knn = KNNExpression(
+        k=1,
+        vector_field=m.embeddings,
+        score_field=m.embeddings_score,
+        reference_vector=embedding_bytes,
+    )
+
+    query = m.find()
+
+    members = await query.all()
+
+    assert len(members) == 1
+    assert members[0].embeddings_score is not None

tests/test_pydantic_integrations.py

@@ -51,7 +51,6 @@ def test_email_str(m):
         )
 
 
-
 def test_validator_sets_value_on_init():
     value = "bar"
 
@@ -62,6 +61,6 @@ class ModelWithValidator(HashModel):
         def set_field(cls, v):
             return value
 
-    m = ModelWithValidator()
+    m = ModelWithValidator(field="foo")
 
     assert m.field == value