feat: add Pinecone cli support (#717)

* feat: add Pinecone cli support

* fix: replace deprecated search method

* chore: format cli.py

Author: Pitrick3141

install/requirements_py3.11.txt

@@ -1,7 +1,7 @@
 grpcio==1.53.2
 grpcio-tools==1.53.0
 qdrant-client
-pinecone-client
+pinecone
 weaviate-client
 elasticsearch==8.16.0
 pgvector

pyproject.toml

@@ -57,7 +57,7 @@ all = [
     "grpcio==1.53.0", # for qdrant-client and pymilvus
     "grpcio-tools==1.53.0", # for qdrant-client and pymilvus
     "qdrant-client",
-    "pinecone-client",
+    "pinecone",
     "weaviate-client",
     "elasticsearch",
     "sqlalchemy",
@@ -83,7 +83,7 @@ all = [
 ]
 
 qdrant          = [ "qdrant-client" ]
-pinecone        = [ "pinecone-client" ]
+pinecone        = [ "pinecone" ]
 weaviate        = [ "weaviate-client" ]
 elastic         = [ "elasticsearch" ]
 # For elastic and aliyun_elasticsearch

vectordb_bench/backend/clients/pinecone/cli.py

@@ -0,0 +1,54 @@
+from typing import Annotated, TypedDict, Unpack
+
+import click
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from .. import DB
+from ..api import EmptyDBCaseConfig
+
+
+class PineconeTypedDict(TypedDict):
+    api_key: Annotated[
+        str,
+        click.option("--api-key", type=str, help="Pinecone API key", required=True),
+    ]
+    index_name: Annotated[
+        str,
+        click.option("--index-name", type=str, help="Pinecone index name", required=True),
+    ]
+    version: Annotated[
+        str,
+        click.option("--version", type=str, help="Database version", default="", show_default=True),
+    ]
+    note: Annotated[
+        str,
+        click.option("--note", type=str, help="Additional notes", default="", show_default=True),
+    ]
+
+
+class PineconeIndexTypedDict(CommonTypedDict, PineconeTypedDict): ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PineconeIndexTypedDict)
+def Pinecone(**parameters: Unpack[PineconeIndexTypedDict]):
+    from .config import PineconeConfig
+
+    run(
+        db=DB.Pinecone,
+        db_config=PineconeConfig(
+            db_label=parameters["db_label"],
+            version=parameters["version"],
+            note=parameters["note"],
+            api_key=SecretStr(parameters["api_key"]),
+            index_name=parameters["index_name"],
+        ),
+        db_case_config=EmptyDBCaseConfig(),
+        **parameters,
+    )

vectordb_bench/backend/clients/qdrant_cloud/qdrant_cloud.py

@@ -198,15 +198,17 @@ def search_embedding(
         """
         assert self.qdrant_client is not None
 
-        res = self.qdrant_client.search(
+        points_res = self.qdrant_client.query_points(
             collection_name=self.collection_name,
-            query_vector=query,
+            query=query,
             limit=k,
             query_filter=self.query_filter,
             search_params=self.db_case_config.search_param(),
             with_payload=self.db_case_config.with_payload,
         )
 
+        res = points_res.points
+
         return [r.id for r in res]
 
     def prepare_filter(self, filters: Filter):

vectordb_bench/cli/vectordbbench.py

@@ -23,6 +23,7 @@
 from ..backend.clients.pgvecto_rs.cli import PgVectoRSHNSW, PgVectoRSIVFFlat
 from ..backend.clients.pgvector.cli import PgVectorHNSW
 from ..backend.clients.pgvectorscale.cli import PgVectorScaleDiskAnn
+from ..backend.clients.pinecone.cli import Pinecone
 from ..backend.clients.qdrant_cloud.cli import QdrantCloud
 from ..backend.clients.qdrant_local.cli import QdrantLocal
 from ..backend.clients.redis.cli import Redis
@@ -76,6 +77,7 @@
 cli.add_command(Chroma)
 cli.add_command(Zvec)
 cli.add_command(Endee)
+cli.add_command(Pinecone)
 
 
 if __name__ == "__main__":