diff --git a/pinecone/db_data/dataclasses/text_query.py b/pinecone/db_data/dataclasses/text_query.py index d3b7997a..d61a7dc6 100644 --- a/pinecone/db_data/dataclasses/text_query.py +++ b/pinecone/db_data/dataclasses/text_query.py @@ -36,12 +36,12 @@ class TextQuery(DictLike): boost: float | None = None slop: int | None = None - def to_dict(self) -> dict: + def as_dict(self) -> dict: """Serialize to API format. :returns: Dictionary representation for the API. """ - result: dict = {"field": self.field, "query": self.query} + result: dict = {"type": "text", "field": self.field, "text_query": self.query} if self.boost is not None: result["boost"] = self.boost if self.slop is not None: diff --git a/pinecone/db_data/dataclasses/vector_query.py b/pinecone/db_data/dataclasses/vector_query.py index 6b76bc81..d61b51e0 100644 --- a/pinecone/db_data/dataclasses/vector_query.py +++ b/pinecone/db_data/dataclasses/vector_query.py @@ -45,12 +45,12 @@ class VectorQuery(DictLike): values: list[float] | None = None sparse_values: SparseValues | None = None - def to_dict(self) -> dict: + def as_dict(self) -> dict: """Serialize to API format. :returns: Dictionary representation for the API. """ - result: dict = {"field": self.field} + result: dict = {"type": "vector", "field": self.field} if self.values is not None: result["values"] = self.values if self.sparse_values is not None: diff --git a/tests/unit/data/test_query_classes.py b/tests/unit/data/test_query_classes.py index 3c456d75..85f5eb0c 100644 --- a/tests/unit/data/test_query_classes.py +++ b/tests/unit/data/test_query_classes.py @@ -11,25 +11,36 @@ def test_required_params(self): assert query.boost is None assert query.slop is None - def test_to_dict_minimal(self): + def test_as_dict_minimal(self): query = TextQuery(field="title", query="pink panther") - result = query.to_dict() - assert result == {"field": "title", "query": "pink panther"} + result = query.as_dict() + assert result == {"type": "text", "field": "title", "text_query": "pink panther"} - def test_to_dict_with_boost(self): + def test_as_dict_with_boost(self): query = TextQuery(field="title", query="pink panther", boost=2.0) - result = query.to_dict() - assert result == {"field": "title", "query": "pink panther", "boost": 2.0} + result = query.as_dict() + assert result == { + "type": "text", + "field": "title", + "text_query": "pink panther", + "boost": 2.0, + } - def test_to_dict_with_slop(self): + def test_as_dict_with_slop(self): query = TextQuery(field="title", query="pink panther", slop=2) - result = query.to_dict() - assert result == {"field": "title", "query": "pink panther", "slop": 2} + result = query.as_dict() + assert result == {"type": "text", "field": "title", "text_query": "pink panther", "slop": 2} - def test_to_dict_with_all_options(self): + def test_as_dict_with_all_options(self): query = TextQuery(field="title", query="pink panther", boost=1.5, slop=3) - result = query.to_dict() - assert result == {"field": "title", "query": "pink panther", "boost": 1.5, "slop": 3} + result = query.as_dict() + assert result == { + "type": "text", + "field": "title", + "text_query": "pink panther", + "boost": 1.5, + "slop": 3, + } def test_dict_like_access(self): query = TextQuery(field="title", query="pink panther", boost=2.0) @@ -51,30 +62,32 @@ def test_required_params(self): assert query.values is None assert query.sparse_values is None - def test_to_dict_minimal(self): + def test_as_dict_minimal(self): query = VectorQuery(field="embedding") - result = query.to_dict() - assert result == {"field": "embedding"} + result = query.as_dict() + assert result == {"type": "vector", "field": "embedding"} - def test_to_dict_with_values(self): + def test_as_dict_with_values(self): query = VectorQuery(field="embedding", values=[0.1, 0.2, 0.3]) - result = query.to_dict() - assert result == {"field": "embedding", "values": [0.1, 0.2, 0.3]} + result = query.as_dict() + assert result == {"type": "vector", "field": "embedding", "values": [0.1, 0.2, 0.3]} - def test_to_dict_with_sparse_values(self): + def test_as_dict_with_sparse_values(self): sparse = SparseValues(indices=[1, 5, 10], values=[0.5, 0.3, 0.2]) query = VectorQuery(field="sparse_embedding", sparse_values=sparse) - result = query.to_dict() + result = query.as_dict() assert result == { + "type": "vector", "field": "sparse_embedding", "sparse_values": {"indices": [1, 5, 10], "values": [0.5, 0.3, 0.2]}, } - def test_to_dict_with_both_values(self): + def test_as_dict_with_both_values(self): sparse = SparseValues(indices=[1, 2], values=[0.5, 0.5]) query = VectorQuery(field="hybrid", values=[0.1, 0.2, 0.3], sparse_values=sparse) - result = query.to_dict() + result = query.as_dict() assert result == { + "type": "vector", "field": "hybrid", "values": [0.1, 0.2, 0.3], "sparse_values": {"indices": [1, 2], "values": [0.5, 0.5]}, @@ -97,12 +110,14 @@ class TestQueryUsageExamples: def test_text_query_example(self): query = TextQuery(field="title", query='return "pink panther"') - result = query.to_dict() + result = query.as_dict() + assert result["type"] == "text" assert result["field"] == "title" - assert result["query"] == 'return "pink panther"' + assert result["text_query"] == 'return "pink panther"' def test_vector_query_example(self): query = VectorQuery(field="embedding", values=[0.1, 0.2, 0.3]) - result = query.to_dict() + result = query.as_dict() + assert result["type"] == "vector" assert result["field"] == "embedding" assert result["values"] == [0.1, 0.2, 0.3]