Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Fix] Rename Index class to fix plugin installation #425

Merged
merged 2 commits into from
Dec 16, 2024
+44 −66
Merged

[Fix] Rename Index class to fix plugin installation #425

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8e0de62
Rename _Index class back to Index
jhamon Dec 16, 2024
e477769
Fix unit tests
jhamon Dec 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion pinecone/data/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
29 changes: 1 addition & 28 deletions pinecone/data/index.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand All @@ -45,32 +37,13 @@

logger = logging.getLogger(__name__)

__all__ = [
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice thinking getting these out of here. 👍

"_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.
Expand Down
71 changes: 34 additions & 37 deletions tests/unit/test_index.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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:
Expand Down Expand Up @@ -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={}),
Expand All @@ -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),
Expand All @@ -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),
Expand All @@ -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)]
)
)
Expand All @@ -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)),
Expand All @@ -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),
Expand All @@ -141,15 +134,15 @@ 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",
),
async_req=True,
)

index._vector_api.upsert_vectors.assert_any_call(
UpsertRequest(
oai_models.UpsertRequest(
vectors=[Vector(id="vec2", values=self.vals2, metadata=self.md2)],
namespace="ns",
),
Expand Down Expand Up @@ -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"
)
)
Expand Down Expand Up @@ -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),
Expand All @@ -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"
)
)
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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),
Expand All @@ -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"
)
)
Expand All @@ -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),
Expand Down Expand Up @@ -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(
Expand All @@ -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()
Expand All @@ -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):
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Loading