Update response model with sort_property (#75)

Co-authored-by: Dan Jones <[email protected]>

Author: danmichaeljones

test/query_agent/test_query_model.py

@@ -561,7 +561,18 @@ def json(self) -> dict:
                 "value": 0.0,
             },
             "collection": "test_collection",
-        }
+            "sort_property": None,
+        },
+        {
+            "query": None,
+            "filters": None,
+            "collection": "test_collection",
+            "sort_property": {
+                "property_name": "test_property",
+                "order": "ascending",
+                "tie_break": None,
+            },
+        },
     ],
     "usage": {
         "model_units": 1,

test/test_imports.py

@@ -80,6 +80,7 @@ def test_class_exports():
         QueryResult,
         QueryResultWithCollection,
         QueryResultWithCollectionNormalized,
+        QuerySort,
         QueryWithCollection,
         SearchModeResponseBase,
         Source,
@@ -142,6 +143,7 @@ def test_class_exports():
         QueryAgentResponse,
         QueryResult,
         QueryResultWithCollection,
+        QuerySort,
         QueryWithCollection,
         SearchModeResponseBase,
         Source,
@@ -222,6 +224,7 @@ def test_class_exports():
         "FilterAndOr",
         "QueryResultWithCollectionNormalized",
         "AggregationResultWithCollectionNormalized",
+        "QuerySort",
     ]
 
     assert hasattr(

weaviate_agents/classes/__init__.py

@@ -37,6 +37,7 @@
     QueryResult,
     QueryResultWithCollection,
     QueryResultWithCollectionNormalized,
+    QuerySort,
     QueryWithCollection,
     SearchModeResponseBase,
     Source,
@@ -110,4 +111,5 @@
     "FilterAndOr",
     "QueryResultWithCollectionNormalized",
     "AggregationResultWithCollectionNormalized",
+    "QuerySort",
 ]

weaviate_agents/classes/query.py

@@ -28,6 +28,7 @@
     QueryResult,
     QueryResultWithCollection,
     QueryResultWithCollectionNormalized,
+    QuerySort,
     QueryWithCollection,
     SearchModeResponseBase,
     Source,
@@ -82,4 +83,5 @@
     "FilterAndOr",
     "QueryResultWithCollectionNormalized",
     "AggregationResultWithCollectionNormalized",
+    "QuerySort",
 ]

weaviate_agents/query/classes/__init__.py

@@ -29,6 +29,7 @@
     QueryResult,
     QueryResultWithCollection,
     QueryResultWithCollectionNormalized,
+    QuerySort,
     QueryWithCollection,
     SearchModeResponseBase,
     Source,
@@ -82,4 +83,5 @@
     "FilterAndOr",
     "QueryResultWithCollectionNormalized",
     "AggregationResultWithCollectionNormalized",
+    "QuerySort",
 ]

weaviate_agents/query/classes/response.py

@@ -391,10 +391,17 @@ class FilterAndOr(BaseModel):
     filters: list[Union[PropertyFilter, FilterAndOr]]
 
 
+class QuerySort(BaseModel):
+    property_name: str
+    order: Literal["ascending", "descending"]
+    tie_break: Union[QuerySort, None]
+
+
 class QueryResultWithCollectionNormalized(BaseModel):
     query: Union[str, None]
     filters: Union[PropertyFilter, FilterAndOr, None]
     collection: str
+    sort_property: Union[QuerySort, None] = None
 
 
 class AggregationResultWithCollectionNormalized(BaseModel):