fix(document): complete the schema for namedscore (#65)

docarray · Jan 18, 2022 · cbd5fd0 · cbd5fd0
1 parent 79c13a0
commit cbd5fd0
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/docarray/document/pydantic_model.py b/docarray/document/pydantic_model.py
@@ -20,6 +20,13 @@ def _convert_ndarray_to_list(v: 'ArrayType'):
         return to_list(v)
 
 
+class _NamedScore(BaseModel):
+    value: Optional[float] = None
+    op_name: Optional[str] = None
+    description: Optional[str] = None
+    ref_id: Optional[str] = None
+
+
 class PydanticDocument(BaseModel):
     id: str
     parent_id: Optional[str]
@@ -36,8 +43,8 @@ class PydanticDocument(BaseModel):
     location: Optional[List[float]]
     embedding: Optional[Any]
     modality: Optional[str]
-    evaluations: Optional[Dict[str, Dict[str, '_StructValueType']]]
-    scores: Optional[Dict[str, Dict[str, '_StructValueType']]]
+    evaluations: Optional[Dict[str, '_NamedScore']]
+    scores: Optional[Dict[str, '_NamedScore']]
     chunks: Optional[List['PydanticDocument']]
     matches: Optional[List['PydanticDocument']]
 

diff --git a/tests/unit/test_pydantic.py b/tests/unit/test_pydantic.py
@@ -111,6 +111,7 @@ def test_match_to_from_pydantic():
     dap = da.to_pydantic_model()
     da_r = DocumentArray.from_pydantic_model(dap)
     assert da_r[0].matches[0].scores['cosine']
+    assert isinstance(da_r[0].matches[0].scores['cosine'], NamedScore)
     assert isinstance(da_r[0].matches[0].scores, defaultdict)
     assert isinstance(da_r[0].matches[0].scores['random_score'], NamedScore)