diff --git a/pinecone/data/__init__.py b/pinecone/data/__init__.py index 02a67767..a6a791fc 100644 --- a/pinecone/data/__init__.py +++ b/pinecone/data/__init__.py @@ -1,4 +1,12 @@ -from .index import * +from .index import ( + Index as _Index, + FetchResponse, + QueryResponse, + DescribeIndexStatsResponse, + UpsertResponse, + SparseValues, + Vector, +) from .import_error import Index, IndexClientInstantiationError from .errors import ( VectorDictionaryMissingKeysError, diff --git a/pinecone/data/index.py b/pinecone/data/index.py index 90f7f827..9f52589c 100644 --- a/pinecone/data/index.py +++ b/pinecone/data/index.py @@ -11,18 +11,10 @@ from pinecone.core.openapi.db_data import API_VERSION from pinecone.core.openapi.db_data.models import ( FetchResponse, - QueryRequest, QueryResponse, - RpcStatus, - ScoredVector, - SingleQueryResults, IndexDescription as DescribeIndexStatsResponse, - UpsertRequest, UpsertResponse, Vector, - DeleteRequest, - UpdateRequest, - DescribeIndexStatsRequest, ListResponse, SparseValues, ) @@ -45,32 +37,13 @@ logger = logging.getLogger(__name__) -__all__ = [ - "_Index", - "FetchResponse", - "QueryRequest", - "QueryResponse", - "RpcStatus", - "ScoredVector", - "SingleQueryResults", - "DescribeIndexStatsResponse", - "UpsertRequest", - "UpsertResponse", - "UpdateRequest", - "Vector", - "DeleteRequest", - "UpdateRequest", - "DescribeIndexStatsRequest", - "SparseValues", -] - def parse_query_response(response: QueryResponse): response._data_store.pop("results", None) return response -class _Index(IndexInterface, ImportFeatureMixin): +class Index(IndexInterface, ImportFeatureMixin): """ A client for interacting with a Pinecone index via REST API. For improved performance, use the Pinecone GRPC index client. diff --git a/tests/unit/test_index.py b/tests/unit/test_index.py index d3c8ff7c..fe45108c 100644 --- a/tests/unit/test_index.py +++ b/tests/unit/test_index.py @@ -1,16 +1,9 @@ import pandas as pd import pytest -import pinecone from pinecone.data import _Index -from pinecone import UpsertRequest, Vector -from pinecone import ( - DescribeIndexStatsRequest, - ScoredVector, - QueryResponse, - UpsertResponse, - SparseValues, -) +import pinecone.core.openapi.db_data.models as oai_models +from pinecone import QueryResponse, UpsertResponse, SparseValues, Vector class TestRestIndex: @@ -39,7 +32,7 @@ def test_upsert_tuplesOfIdVec_UpserWithoutMD(self, mocker): mocker.patch.object(self.index._vector_api, "upsert_vectors", autospec=True) self.index.upsert([("vec1", self.vals1), ("vec2", self.vals2)], namespace="ns") self.index._vector_api.upsert_vectors.assert_called_once_with( - UpsertRequest( + oai_models.UpsertRequest( vectors=[ Vector(id="vec1", values=self.vals1, metadata={}), Vector(id="vec2", values=self.vals2, metadata={}), @@ -52,7 +45,7 @@ def test_upsert_tuplesOfIdVecMD_UpsertVectorsWithMD(self, mocker): mocker.patch.object(self.index._vector_api, "upsert_vectors", autospec=True) self.index.upsert([("vec1", self.vals1, self.md1), ("vec2", self.vals2, self.md2)]) self.index._vector_api.upsert_vectors.assert_called_once_with( - UpsertRequest( + oai_models.UpsertRequest( vectors=[ Vector(id="vec1", values=self.vals1, metadata=self.md1), Vector(id="vec2", values=self.vals2, metadata=self.md2), @@ -69,7 +62,7 @@ def test_upsert_dictOfIdVecMD_UpsertVectorsWithMD(self, mocker): ] ) self.index._vector_api.upsert_vectors.assert_called_once_with( - UpsertRequest( + oai_models.UpsertRequest( vectors=[ Vector(id="vec1", values=self.vals1, metadata=self.md1), Vector(id="vec2", values=self.vals2, metadata=self.md2), @@ -83,7 +76,7 @@ def test_upsert_dictOfIdVecMD_UpsertVectorsWithoutMD(self, mocker): [{"id": self.id1, "values": self.vals1}, {"id": self.id2, "values": self.vals2}] ) self.index._vector_api.upsert_vectors.assert_called_once_with( - UpsertRequest( + oai_models.UpsertRequest( vectors=[Vector(id="vec1", values=self.vals1), Vector(id="vec2", values=self.vals2)] ) ) @@ -97,7 +90,7 @@ def test_upsert_dictOfIdVecMD_UpsertVectorsWithSparseValues(self, mocker): ] ) self.index._vector_api.upsert_vectors.assert_called_once_with( - UpsertRequest( + oai_models.UpsertRequest( vectors=[ Vector(id="vec1", values=self.vals1, sparse_values=SparseValues(**self.sv1)), Vector(id="vec2", values=self.vals2, sparse_values=SparseValues(**self.sv2)), @@ -115,7 +108,7 @@ def test_upsert_vectors_upsertInputVectors(self, mocker): namespace="ns", ) self.index._vector_api.upsert_vectors.assert_called_once_with( - UpsertRequest( + oai_models.UpsertRequest( vectors=[ Vector(id="vec1", values=self.vals1, metadata=self.md1), Vector(id="vec2", values=self.vals2, metadata=self.md2), @@ -141,7 +134,7 @@ def test_upsert_parallelUpsert_callUpsertParallel(self, mocker): [async_result.get() for async_result in async_results] index._vector_api.upsert_vectors.assert_any_call( - UpsertRequest( + oai_models.UpsertRequest( vectors=[Vector(id="vec1", values=self.vals1, metadata=self.md1)], namespace="ns", ), @@ -149,7 +142,7 @@ def test_upsert_parallelUpsert_callUpsertParallel(self, mocker): ) index._vector_api.upsert_vectors.assert_any_call( - UpsertRequest( + oai_models.UpsertRequest( vectors=[Vector(id="vec2", values=self.vals2, metadata=self.md2)], namespace="ns", ), @@ -177,13 +170,13 @@ def test_upsert_vectorListIsMultiplyOfBatchSize_vectorsUpsertedInBatches(self, m ) self.index._vector_api.upsert_vectors.assert_any_call( - UpsertRequest( + oai_models.UpsertRequest( vectors=[Vector(id="vec1", values=self.vals1, metadata=self.md1)], namespace="ns" ) ) self.index._vector_api.upsert_vectors.assert_any_call( - UpsertRequest( + oai_models.UpsertRequest( vectors=[Vector(id="vec2", values=self.vals2, metadata=self.md2)], namespace="ns" ) ) @@ -211,7 +204,7 @@ def test_upsert_vectorListNotMultiplyOfBatchSize_vectorsUpsertedInBatches(self, ) self.index._vector_api.upsert_vectors.assert_any_call( - UpsertRequest( + oai_models.UpsertRequest( vectors=[ Vector(id="vec1", values=self.vals1, metadata=self.md1), Vector(id="vec2", values=self.vals2, metadata=self.md2), @@ -221,7 +214,7 @@ def test_upsert_vectorListNotMultiplyOfBatchSize_vectorsUpsertedInBatches(self, ) self.index._vector_api.upsert_vectors.assert_any_call( - UpsertRequest( + oai_models.UpsertRequest( vectors=[Vector(id="vec3", values=self.vals1, metadata=self.md1)], namespace="ns" ) ) @@ -249,7 +242,7 @@ def test_upsert_vectorListSmallerThanBatchSize_vectorsUpsertedInBatches(self, mo ) self.index._vector_api.upsert_vectors.assert_called_once_with( - UpsertRequest( + oai_models.UpsertRequest( vectors=[ Vector(id="vec1", values=self.vals1, metadata=self.md1), Vector(id="vec2", values=self.vals2, metadata=self.md2), @@ -282,7 +275,7 @@ def test_upsert_tuplesList_vectorsUpsertedInBatches(self, mocker): ) self.index._vector_api.upsert_vectors.assert_any_call( - UpsertRequest( + oai_models.UpsertRequest( vectors=[ Vector(id="vec1", values=self.vals1, metadata=self.md1), Vector(id="vec2", values=self.vals2, metadata=self.md2), @@ -292,7 +285,7 @@ def test_upsert_tuplesList_vectorsUpsertedInBatches(self, mocker): ) self.index._vector_api.upsert_vectors.assert_any_call( - UpsertRequest( + oai_models.UpsertRequest( vectors=[Vector(id="vec3", values=self.vals1, metadata=self.md1)], namespace="ns" ) ) @@ -315,7 +308,7 @@ def test_upsert_dataframe(self, mocker): self.index.upsert_from_dataframe(df) self.index._vector_api.upsert_vectors.assert_called_once_with( - UpsertRequest( + oai_models.UpsertRequest( vectors=[ Vector(id="vec1", values=self.vals1, metadata=self.md1), Vector(id="vec2", values=self.vals2, metadata=self.md2), @@ -356,7 +349,7 @@ def test_upsert_useBatchSizeAndAsyncReq_valueErrorRaised(self): def test_query_byVectorNoFilter_queryVectorNoFilter(self, mocker): response = QueryResponse( results=[], - matches=[ScoredVector(id="1", score=0.9, values=[0.0], metadata={"a": 2})], + matches=[oai_models.ScoredVector(id="1", score=0.9, values=[0.0], metadata={"a": 2})], namespace="test", ) mocker.patch.object( @@ -366,10 +359,10 @@ def test_query_byVectorNoFilter_queryVectorNoFilter(self, mocker): actual = self.index.query(top_k=10, vector=self.vals1) self.index._vector_api.query_vectors.assert_called_once_with( - pinecone.QueryRequest(top_k=10, vector=self.vals1) + oai_models.QueryRequest(top_k=10, vector=self.vals1) ) expected = QueryResponse( - matches=[ScoredVector(id="1", score=0.9, values=[0.0], metadata={"a": 2})], + matches=[oai_models.ScoredVector(id="1", score=0.9, values=[0.0], metadata={"a": 2})], namespace="test", ) assert expected.to_dict() == actual.to_dict() @@ -378,14 +371,18 @@ def test_query_byVectorWithFilter_queryVectorWithFilter(self, mocker): mocker.patch.object(self.index._vector_api, "query_vectors", autospec=True) self.index.query(top_k=10, vector=self.vals1, filter=self.filter1, namespace="ns") self.index._vector_api.query_vectors.assert_called_once_with( - pinecone.QueryRequest(top_k=10, vector=self.vals1, filter=self.filter1, namespace="ns") + oai_models.QueryRequest( + top_k=10, vector=self.vals1, filter=self.filter1, namespace="ns" + ) ) def test_query_byVecId_queryByVecId(self, mocker): mocker.patch.object(self.index._vector_api, "query_vectors", autospec=True) self.index.query(top_k=10, id="vec1", include_metadata=True, include_values=False) self.index._vector_api.query_vectors.assert_called_once_with( - pinecone.QueryRequest(top_k=10, id="vec1", include_metadata=True, include_values=False) + oai_models.QueryRequest( + top_k=10, id="vec1", include_metadata=True, include_values=False + ) ) def test_query_rejects_both_id_and_vector(self): @@ -408,21 +405,21 @@ def test_delete_byIds_deleteByIds(self, mocker): mocker.patch.object(self.index._vector_api, "delete_vectors", autospec=True) self.index.delete(ids=["vec1", "vec2"]) self.index._vector_api.delete_vectors.assert_called_once_with( - pinecone.DeleteRequest(ids=["vec1", "vec2"]) + oai_models.DeleteRequest(ids=["vec1", "vec2"]) ) def test_delete_deleteAllByFilter_deleteAllByFilter(self, mocker): mocker.patch.object(self.index._vector_api, "delete_vectors", autospec=True) self.index.delete(delete_all=True, filter=self.filter1, namespace="ns") self.index._vector_api.delete_vectors.assert_called_once_with( - pinecone.DeleteRequest(delete_all=True, filter=self.filter1, namespace="ns") + oai_models.DeleteRequest(delete_all=True, filter=self.filter1, namespace="ns") ) def test_delete_deleteAllNoFilter_deleteNoFilter(self, mocker): mocker.patch.object(self.index._vector_api, "delete_vectors", autospec=True) self.index.delete(delete_all=True) self.index._vector_api.delete_vectors.assert_called_once_with( - pinecone.DeleteRequest(delete_all=True) + oai_models.DeleteRequest(delete_all=True) ) # endregion @@ -449,14 +446,14 @@ def test_update_byIdAnValues_updateByIdAndValues(self, mocker): mocker.patch.object(self.index._vector_api, "update_vector", autospec=True) self.index.update(id="vec1", values=self.vals1, namespace="ns") self.index._vector_api.update_vector.assert_called_once_with( - pinecone.UpdateRequest(id="vec1", values=self.vals1, namespace="ns") + oai_models.UpdateRequest(id="vec1", values=self.vals1, namespace="ns") ) def test_update_byIdAnValuesAndMetadata_updateByIdAndValuesAndMetadata(self, mocker): mocker.patch.object(self.index._vector_api, "update_vector", autospec=True) self.index.update("vec1", values=self.vals1, metadata=self.md1) self.index._vector_api.update_vector.assert_called_once_with( - pinecone.UpdateRequest(id="vec1", values=self.vals1, metadata=self.md1) + oai_models.UpdateRequest(id="vec1", values=self.vals1, metadata=self.md1) ) # endregion @@ -467,14 +464,14 @@ def test_describeIndexStats_callWithoutFilter_CalledWithoutFilter(self, mocker): mocker.patch.object(self.index._vector_api, "describe_index_stats", autospec=True) self.index.describe_index_stats() self.index._vector_api.describe_index_stats.assert_called_once_with( - DescribeIndexStatsRequest() + oai_models.DescribeIndexStatsRequest() ) def test_describeIndexStats_callWithFilter_CalledWithFilter(self, mocker): mocker.patch.object(self.index._vector_api, "describe_index_stats", autospec=True) self.index.describe_index_stats(filter=self.filter1) self.index._vector_api.describe_index_stats.assert_called_once_with( - DescribeIndexStatsRequest(filter=self.filter1) + oai_models.DescribeIndexStatsRequest(filter=self.filter1) ) # endregion