diff --git a/pinecone/data/dataclasses/__init__.py b/pinecone/data/dataclasses/__init__.py index fddb08b1..c481d1ab 100644 --- a/pinecone/data/dataclasses/__init__.py +++ b/pinecone/data/dataclasses/__init__.py @@ -1,2 +1,3 @@ from .sparse_values import SparseValues from .vector import Vector +from .fetch_response import FetchResponse diff --git a/pinecone/data/dataclasses/fetch_response.py b/pinecone/data/dataclasses/fetch_response.py new file mode 100644 index 00000000..e8d28069 --- /dev/null +++ b/pinecone/data/dataclasses/fetch_response.py @@ -0,0 +1,11 @@ +from dataclasses import dataclass +from typing import Dict + +from .vector import Vector + + +@dataclass +class FetchResponse: + namespace: str + vectors: Dict[str, Vector] + usage: Dict[str, int] diff --git a/pinecone/data/index.py b/pinecone/data/index.py index 999c85a0..23808e5f 100644 --- a/pinecone/data/index.py +++ b/pinecone/data/index.py @@ -10,7 +10,6 @@ from pinecone.core.openapi.db_data.api.vector_operations_api import VectorOperationsApi from pinecone.core.openapi.db_data import API_VERSION from pinecone.core.openapi.db_data.models import ( - FetchResponse, QueryRequest, QueryResponse, RpcStatus, @@ -24,7 +23,7 @@ DescribeIndexStatsRequest, ListResponse, ) -from .dataclasses import Vector, SparseValues +from .dataclasses import Vector, SparseValues, FetchResponse from .interfaces import IndexInterface from .request_factory import IndexRequestFactory from .features.bulk_import import ImportFeatureMixin @@ -234,7 +233,12 @@ def delete( @validate_and_convert_errors def fetch(self, ids: List[str], namespace: Optional[str] = None, **kwargs) -> FetchResponse: args_dict = parse_non_empty_args([("namespace", namespace)]) - return self._vector_api.fetch_vectors(ids=ids, **args_dict, **kwargs) + result = self._vector_api.fetch_vectors(ids=ids, **args_dict, **kwargs) + return FetchResponse( + namespace=namespace, + vectors={k: Vector.from_dict(v) for k, v in result.vectors}, + usage=result.usage, + ) @validate_and_convert_errors def query(