From cec82cdc70457b51123b7e2d53bcc0ddf90bf7e0 Mon Sep 17 00:00:00 2001 From: Han Xiao Date: Tue, 18 Jan 2022 21:14:41 +0100 Subject: [PATCH 1/2] fix(document): complete the schema for namedscore --- docarray/document/pydantic_model.py | 11 +++++++++-- tests/unit/test_pydantic.py | 3 ++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/docarray/document/pydantic_model.py b/docarray/document/pydantic_model.py index 596fa568b78..b332c578a7c 100644 --- 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 index cdd8a87da4d..4db8c2ede2e 100644 --- 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) @@ -158,4 +159,4 @@ def test_to_from_with_blob(protocol, to_fn, blob): assert d.blob == r_d.blob if d.blob: - assert isinstance(r_d.blob, bytes) + assert isinstance(r_d.blob, bytes) \ No newline at end of file From 0b68ad636bc46270cf714d6927174d91c1aba408 Mon Sep 17 00:00:00 2001 From: Han Xiao Date: Tue, 18 Jan 2022 21:14:56 +0100 Subject: [PATCH 2/2] fix(document): complete the schema for namedscore --- tests/unit/test_pydantic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_pydantic.py b/tests/unit/test_pydantic.py index 4db8c2ede2e..074f162852f 100644 --- a/tests/unit/test_pydantic.py +++ b/tests/unit/test_pydantic.py @@ -159,4 +159,4 @@ def test_to_from_with_blob(protocol, to_fn, blob): assert d.blob == r_d.blob if d.blob: - assert isinstance(r_d.blob, bytes) \ No newline at end of file + assert isinstance(r_d.blob, bytes)