feat(sdk): Add support for record.status on filter

argilla-io · Jun 28, 2024 · d22dfd9 · d22dfd9
1 parent 5060360
commit d22dfd9
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/argilla/src/argilla/_models/_search.py b/argilla/src/argilla/_models/_search.py
@@ -17,6 +17,11 @@
 from pydantic import BaseModel, Field
 
 
+class RecordFilterScopeModel(BaseModel):
+    entity: Literal["record"] = "record"
+    property: Literal["status"] = "status"
+
+
 class ResponseFilterScopeModel(BaseModel):
     """Filter scope for filtering on a response entity."""
 
@@ -42,6 +47,7 @@ class MetadataFilterScopeModel(BaseModel):
 
 ScopeModel = Annotated[
     Union[
+        RecordFilterScopeModel,
         ResponseFilterScopeModel,
         SuggestionFilterScopeModel,
         MetadataFilterScopeModel,

diff --git a/argilla/src/argilla/records/_search.py b/argilla/src/argilla/records/_search.py
@@ -26,6 +26,7 @@
     FilterModel,
     AndFilterModel,
     QueryModel,
+    RecordFilterScopeModel,
 )
 
 
@@ -54,8 +55,9 @@ def model(self) -> FilterModel:
     @staticmethod
     def _extract_filter_scope(field: str) -> ScopeModel:
         field = field.strip()
-
         if field == "status":
+            return RecordFilterScopeModel(property="status")
+        elif field == "responses.status":
             return ResponseFilterScopeModel(property="status")
         elif "metadata" in field:
             _, md_property = field.split(".")