|
12 | 12 |
|
13 | 13 | import pytest
|
14 | 14 | import pytest_asyncio
|
| 15 | +from redis.exceptions import ResponseError |
15 | 16 |
|
16 | 17 | from aredis_om import (
|
17 | 18 | EmbeddedJsonModel,
|
|
21 | 22 | NotFoundError,
|
22 | 23 | QueryNotSupportedError,
|
23 | 24 | RedisModelError,
|
| 25 | + VectorFieldOptions, |
24 | 26 | )
|
25 | 27 |
|
26 | 28 | # We need to run this check as sync code (during tests) even in async mode
|
@@ -1173,3 +1175,82 @@ class Game(JsonModel):
|
1173 | 1175 | )
|
1174 | 1176 | print(q.query)
|
1175 | 1177 | assert q.query == "(@player1_username:{username})| (@player2_username:{username})"
|
| 1178 | + |
| 1179 | + |
| 1180 | +@py_test_mark_asyncio |
| 1181 | +def test_vector_field_definition(redis): |
| 1182 | + """ |
| 1183 | + Test the definition and behavior of a vector field in a JsonModel. |
| 1184 | + This test verifies: |
| 1185 | + 1. The model schema includes "VECTOR" for a vector field with specified options. |
| 1186 | + 2. Instances with vector fields can be saved and retrieved accurately. |
| 1187 | + 3. Vector field values remain consistent after persistence. |
| 1188 | +
|
| 1189 | + Args: |
| 1190 | + redis: Redis connection fixture for testing. |
| 1191 | + """ |
| 1192 | + |
| 1193 | + class Group(JsonModel): |
| 1194 | + articles: List[str] |
| 1195 | + tender_text: str = Field(index=False) |
| 1196 | + tender_embedding: List[float] = Field( |
| 1197 | + index=True, |
| 1198 | + vector_options=VectorFieldOptions( |
| 1199 | + algorithm=VectorFieldOptions.ALGORITHM.FLAT, |
| 1200 | + type=VectorFieldOptions.TYPE.FLOAT32, |
| 1201 | + dimension=3, |
| 1202 | + distance_metric=VectorFieldOptions.DISTANCE_METRIC.COSINE, |
| 1203 | + ), |
| 1204 | + ) |
| 1205 | + |
| 1206 | + schema = Group.redisearch_schema() |
| 1207 | + assert "VECTOR" in schema |
| 1208 | + |
| 1209 | + group = Group( |
| 1210 | + articles=["article_1", "article_2"], |
| 1211 | + tender_text="Sample text", |
| 1212 | + tender_embedding=[0.1, 0.2, 0.3], |
| 1213 | + ) |
| 1214 | + group.save() |
| 1215 | + |
| 1216 | + retrieved_group = await Group.get(group.pk) |
| 1217 | + assert retrieved_group.tender_embedding == [0.1, 0.2, 0.3] |
| 1218 | + |
| 1219 | + retrieved_group = Group.get(group.pk) |
| 1220 | + assert retrieved_group.tender_embedding == [0.1, 0.2, 0.3] |
| 1221 | + |
| 1222 | + |
| 1223 | +def test_vector_field_schema_debug(redis): |
| 1224 | + """ |
| 1225 | + Test and debug the schema definition for a vector field in a JsonModel. |
| 1226 | +
|
| 1227 | + This test ensures: |
| 1228 | + 1. The schema for a vector field is generated correctly. |
| 1229 | + 2. No syntax errors occur when saving a model instance. |
| 1230 | + 3. The Redis schema syntax is valid and debugged if issues arise. |
| 1231 | +
|
| 1232 | + Steps: |
| 1233 | + - Define a `TestModel` with a vector field `embedding`. |
| 1234 | + - Attempt to save an instance and print the schema. |
| 1235 | + - Handle and fail gracefully if Redis raises a ResponseError. |
| 1236 | +
|
| 1237 | + Args: |
| 1238 | + redis: Redis connection fixture for testing. |
| 1239 | + """ |
| 1240 | + |
| 1241 | + class TestModel(JsonModel): |
| 1242 | + embedding: List[float] = Field( |
| 1243 | + index=True, |
| 1244 | + vector_options=VectorFieldOptions( |
| 1245 | + algorithm=VectorFieldOptions.ALGORITHM.FLAT, |
| 1246 | + type=VectorFieldOptions.TYPE.FLOAT32, |
| 1247 | + dimension=3, |
| 1248 | + distance_metric=VectorFieldOptions.DISTANCE_METRIC.COSINE, |
| 1249 | + ), |
| 1250 | + ) |
| 1251 | + |
| 1252 | + try: |
| 1253 | + TestModel(embedding=[0.1, 0.2, 0.3]).save() |
| 1254 | + print(TestModel.redisearch_schema()) |
| 1255 | + except ResponseError as e: |
| 1256 | + pytest.fail(f"Redis rejected the schema with error: {e}") |
0 commit comments