Update to pinecone-plugin-inference=2.0.0

pyproject.toml

@@ -52,7 +52,7 @@ lz4 = { version = ">=3.1.3", optional = true }
 protobuf = { version = "^5.28", optional = true }
 protoc-gen-openapiv2 = {version = "^0.0.1", optional = true }
 pinecone-plugin-interface = "^0.0.7"
-pinecone-plugin-inference = "^1.1.0"
+pinecone-plugin-inference = "^2.0.0"
 python-dateutil = ">=2.5.3"
 
 [tool.poetry.group.types]

tests/integration/inference/test_embed.py

@@ -1,5 +1,7 @@
+import pytest
 from pinecone import Pinecone
 from pinecone.grpc import PineconeGRPC
+from pinecone.exceptions import PineconeApiException
 
 
 class TestInferencePlugin:
@@ -32,3 +34,15 @@ def test_embed_grpc(self, api_key):
         assert len(embeddings.get("data")[0]["values"]) == 1024
         assert len(embeddings.get("data")[1]["values"]) == 1024
         assert embeddings.get("model") == embedding_model
+
+    def test_embed_exception(self, api_key):
+        pc = Pinecone(api_key=api_key)
+
+        with pytest.raises(PineconeApiException) as e_info:
+            embedding_model = "DOES NOT EXIST"
+            pc.inference.embed(
+                model=embedding_model,
+                inputs=["The quick brown fox jumps over the lazy dog.", "lorem ipsum"],
+                parameters={"input_type": "query", "truncate": "END"},
+            )
+        assert e_info.value.status == 404

tests/integration/inference/test_rerank.py

@@ -1,5 +1,7 @@
+import pytest
 from pinecone import Pinecone
 from pinecone.grpc import PineconeGRPC
+from pinecone.exceptions import PineconeApiException
 
 
 class TestInferencePluginRerank:
@@ -10,7 +12,11 @@ def test_rerank(self, api_key):
         result = pc.inference.rerank(
             model=model,
             query="i love dogs",
-            documents=["dogs are pretty cool", "everyone loves dogs", "I'm a cat person"],
+            documents=[
+                "dogs are pretty cool",
+                "everyone loves dogs",
+                "I'm a cat person",
+            ],
             top_n=1,
             return_documents=True,
         )
@@ -28,7 +34,11 @@ def test_rerank_grpc(self, api_key):
         result = pc.inference.rerank(
             model=model,
             query="i love dogs",
-            documents=["dogs are pretty cool", "everyone loves dogs", "I'm a cat person"],
+            documents=[
+                "dogs are pretty cool",
+                "everyone loves dogs",
+                "I'm a cat person",
+            ],
             top_n=1,
             return_documents=True,
         )
@@ -38,3 +48,20 @@ def test_rerank_grpc(self, api_key):
         assert result.model == model
         assert isinstance(result.usage.rerank_units, int)
         assert result.usage.rerank_units == 1
+
+    def test_rerank_exception(self, api_key):
+        pc = Pinecone(api_key=api_key)
+        with pytest.raises(PineconeApiException) as e_info:
+            pc.inference.rerank(
+                model="DOES NOT EXIST",
+                query="i love dogs",
+                documents=[
+                    "dogs are pretty cool",
+                    "everyone loves dogs",
+                    "I'm a cat person",
+                ],
+                rank_fields=["custom-field"],
+                top_n=1,
+                return_documents=True,
+            )
+        assert e_info.value.status == 404