From f01d80a8650b0b3052c3f56fd8598feeebb30cd6 Mon Sep 17 00:00:00 2001 From: Jennifer Hamon Date: Mon, 16 Dec 2024 16:44:02 -0500 Subject: [PATCH] [Fix] Rename Index class to fix plugin installation (#425) ## Problem By renaming the `Index` class to `_Index` in https://github.com/pinecone-io/pinecone-python-client/pull/418, I accidentally broke plugin installation. ## Solution Rename the class back to `Index` and instead alias it to `_Index` through import statements. We still need to alias it within the `pinecone/data/__init__.py` module definition to avoid collision with the Index class designed to give users a more informative error when they incorrectly import the `Index` class from the `pinecone` module. I went ahead and also cleaned up some stuff where we were unnecessarily exposing some internal openapi details in the public interface. Users should never realistically need to construct openapi request objects. ## Type of Change - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] This change requires a documentation update - [ ] Infrastructure change (CI configs, etc) - [ ] Non-code change (docs, etc) - [ ] None of the above: (explain here) --- pinecone/data/__init__.py | 10 +++++- pinecone/data/index.py | 29 +--------------- tests/unit/test_index.py | 71 +++++++++++++++++++-------------------- 3 files changed, 44 insertions(+), 66 deletions(-) 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