diff --git a/python/tests/test_doc.py b/python/tests/test_doc.py index ecfd31e4d..7d042a6e5 100644 --- a/python/tests/test_doc.py +++ b/python/tests/test_doc.py @@ -13,6 +13,7 @@ # limitations under the License. from __future__ import annotations +import json import math import pytest @@ -83,6 +84,14 @@ def test_with_numpy_array(self): assert doc.vector("image") == [1, 2, 3] assert doc.vector("keys") == {1: 1.0, 2: 2.0, 3: 3.0} + def test_init_normalizes_numpy_vectors(self): + import numpy as np + + doc = Doc(id="1", vectors={"dense": np.array([1, 2, 3])}) + + assert doc.vector("dense") == [1, 2, 3] + assert json.loads(repr(doc))["vectors"]["dense"] == [1, 2, 3] + # ---------------------------- # CppDoc Test Case diff --git a/python/zvec/model/doc.py b/python/zvec/model/doc.py index 83f75cbef..e441410a3 100644 --- a/python/zvec/model/doc.py +++ b/python/zvec/model/doc.py @@ -72,7 +72,7 @@ def __init__( ): self.id = id self.score = score - self.vectors = vectors or {} + self.vectors = _normalize_vectors(vectors) self.fields = fields or {} def has_field(self, name: str) -> bool: @@ -167,15 +167,23 @@ def _from_tuple( vectors = data_tuple[3] if vectors is not None: - obj.vectors = { - name: (vec.tolist() if hasattr(vec, "tolist") else vec) - for name, vec in vectors.items() - } + obj.vectors = _normalize_vectors(vectors) else: obj.vectors = {} return obj +def _normalize_vectors( + vectors: Optional[dict[str, VectorType]], +) -> dict[str, VectorType]: + if vectors is None: + return {} + return { + name: (vec.tolist() if hasattr(vec, "tolist") else vec) + for name, vec in vectors.items() + } + + #: Type alias for query results: a list of documents returned by a single query route. DocList = list[Doc]