[Fix] Return type from embed() should be iterator (#454)

## Problem

When migrating the `embed` and `rerank` over from the plugin, I forgot
to include these custom return objects.
 
## Solution

Add custom return types, adjust tests to ensure result is iterable.

## Type of Change

- [x] Bug fix (non-breaking change which fixes an issue)

Author: jhamon

pinecone/data/features/inference/inference.py

@@ -3,7 +3,7 @@
 
 from pinecone.openapi_support import ApiClient
 from pinecone.core.openapi.inference.apis import InferenceApi
-from pinecone.core.openapi.inference.models import EmbeddingsList, RerankResult
+from .models import EmbeddingsList, RerankResult
 from pinecone.core.openapi.inference import API_VERSION
 from pinecone.utils import setup_openapi_client, PluginAware
 
@@ -84,7 +84,8 @@ def embed(
         request_body = InferenceRequestBuilder.embed_request(
             model=model, inputs=inputs, parameters=parameters
         )
-        return self.__inference_api.embed(embed_request=request_body)
+        resp = self.__inference_api.embed(embed_request=request_body)
+        return EmbeddingsList(resp)
 
     def rerank(
         self,
@@ -162,4 +163,5 @@ def rerank(
             top_n=top_n,
             parameters=parameters,
         )
-        return self.__inference_api.rerank(rerank_request=rerank_request)
+        resp = self.__inference_api.rerank(rerank_request=rerank_request)
+        return RerankResult(resp)

pinecone/data/features/inference/inference_asyncio.py

@@ -1,7 +1,7 @@
 from typing import Optional, Dict, List, Union, Any
 
 from pinecone.core.openapi.inference.api.inference_api import AsyncioInferenceApi
-from pinecone.core.openapi.inference.models import EmbeddingsList, RerankResult
+from .models import EmbeddingsList, RerankResult
 
 from .inference_request_builder import (
     InferenceRequestBuilder,
@@ -64,7 +64,8 @@ async def embed(
         request_body = InferenceRequestBuilder.embed_request(
             model=model, inputs=inputs, parameters=parameters
         )
-        return await self.__inference_api.embed(embed_request=request_body)
+        resp = await self.__inference_api.embed(embed_request=request_body)
+        return EmbeddingsList(resp)
 
     async def rerank(
         self,
@@ -142,4 +143,5 @@ async def rerank(
             top_n=top_n,
             parameters=parameters,
         )
-        return await self.__inference_api.rerank(rerank_request=rerank_request)
+        resp = await self.__inference_api.rerank(rerank_request=rerank_request)
+        return RerankResult(resp)

pinecone/data/features/inference/models/__init__.py

@@ -0,0 +1,2 @@
+from .embedding_list import EmbeddingsList
+from .rerank_result import RerankResult

pinecone/data/features/inference/models/embedding_list.py

@@ -0,0 +1,29 @@
+from pinecone.core.openapi.inference.models import EmbeddingsList as OpenAPIEmbeddingsList
+
+
+class EmbeddingsList:
+    """
+    A list of embeddings.
+    """
+
+    def __init__(self, embeddings_list: OpenAPIEmbeddingsList):
+        self.embeddings_list = embeddings_list
+        self.current = 0
+
+    def __getitem__(self, index):
+        return self.embeddings_list.get("data")[index]
+
+    def __len__(self):
+        return len(self.embeddings_list.get("data"))
+
+    def __iter__(self):
+        return iter(self.embeddings_list.get("data"))
+
+    def __str__(self):
+        return str(self.embeddings_list)
+
+    def __repr__(self):
+        return repr(self.embeddings_list)
+
+    def __getattr__(self, attr):
+        return getattr(self.embeddings_list, attr)

pinecone/data/features/inference/models/rerank_result.py

@@ -0,0 +1,19 @@
+from pinecone.core.openapi.inference.models import RerankResult as OpenAPIRerankResult
+
+
+class RerankResult:
+    """
+    A wrapper around OpenAPIRerankResult.
+    """
+
+    def __init__(self, rerank_result: OpenAPIRerankResult):
+        self.rerank_result = rerank_result
+
+    def __str__(self):
+        return str(self.rerank_result)
+
+    def __repr__(self):
+        return repr(self.rerank_result)
+
+    def __getattr__(self, attr):
+        return getattr(self.rerank_result, attr)

tests/integration/inference/test_asyncio_inference.py

@@ -19,23 +19,32 @@ async def test_create_embeddings(self, model_input, model_output):
             parameters={"input_type": "query", "truncate": "END"},
         )
         assert embeddings.vector_type == "dense"
+        assert embeddings.get("vector_type") == "dense"
         assert embeddings.model == model_output
+        assert embeddings.get("model") == model_output
         assert len(embeddings.data) == 2
-        assert len(embeddings.data[0].values) == 1024
-        assert len(embeddings.data[1].values) == 1024
+        assert len(embeddings.get("data")) == 2
+        assert embeddings.usage is not None
 
-        # Dict-style bracket accessors
-        assert embeddings["vector_type"] == "dense"
-        assert embeddings["model"] == model_output
-        assert len(embeddings["data"]) == 2
+        individual_embedding = embeddings[0]
+        assert len(individual_embedding.values) == 1024
+        assert individual_embedding.vector_type.value == "dense"
+        assert len(individual_embedding["values"]) == 1024
 
-        # Dict-style get accessors for embeddings
-        assert embeddings.get("vector_type") == "dense"
-        assert embeddings.get("model") == model_output
-        assert len(embeddings.get("data")) == 2
-        assert len(embeddings.get("data")[0]["values"]) == 1024
-        assert len(embeddings.get("data")[1]["values"]) == 1024
-        assert embeddings.get("model") == model_output
+        await pc.close()
+
+    async def test_embedding_result_is_iterable(self):
+        pc = PineconeAsyncio()
+        embeddings = await pc.inference.embed(
+            model=EmbedModel.Multilingual_E5_Large,
+            inputs=["The quick brown fox jumps over the lazy dog.", "lorem ipsum"],
+            parameters={"input_type": "query", "truncate": "END"},
+        )
+        iter_count = 0
+        for embedding in embeddings:
+            iter_count += 1
+            assert len(embedding.values) == 1024
+        assert iter_count == 2
         await pc.close()
 
     @pytest.mark.parametrize(

tests/integration/inference/test_inference.py

@@ -18,23 +18,30 @@ def test_create_embeddings(self, model_input, model_output):
             parameters={"input_type": "query", "truncate": "END"},
         )
         assert embeddings.vector_type == "dense"
+        assert embeddings.get("vector_type") == "dense"
         assert embeddings.model == model_output
+        assert embeddings.get("model") == model_output
         assert len(embeddings.data) == 2
-        assert len(embeddings.data[0].values) == 1024
-        assert len(embeddings.data[1].values) == 1024
+        assert len(embeddings.get("data")) == 2
+        assert embeddings.usage is not None
 
-        # Dict-style bracket accessors
-        assert embeddings["vector_type"] == "dense"
-        assert embeddings["model"] == model_output
-        assert len(embeddings["data"]) == 2
+        individual_embedding = embeddings[0]
+        assert len(individual_embedding.values) == 1024
+        assert individual_embedding.vector_type.value == "dense"
+        assert len(individual_embedding["values"]) == 1024
 
-        # Dict-style get accessors for embeddings
-        assert embeddings.get("vector_type") == "dense"
-        assert embeddings.get("model") == model_output
-        assert len(embeddings.get("data")) == 2
-        assert len(embeddings.get("data")[0]["values"]) == 1024
-        assert len(embeddings.get("data")[1]["values"]) == 1024
-        assert embeddings.get("model") == model_output
+    def test_embedding_result_is_iterable(self):
+        pc = Pinecone()
+        embeddings = pc.inference.embed(
+            model=EmbedModel.Multilingual_E5_Large,
+            inputs=["The quick brown fox jumps over the lazy dog.", "lorem ipsum"],
+            parameters={"input_type": "query", "truncate": "END"},
+        )
+        iter_count = 0
+        for embedding in embeddings:
+            iter_count += 1
+            assert len(embedding.values) == 1024
+        assert iter_count == 2
 
     @pytest.mark.parametrize(
         "model_input,model_output",