From 2d00db9a35f0e615c9ae8c359a8f2fad9ad34faf Mon Sep 17 00:00:00 2001
From: Jen Hamon <jhamon@pinecone.io>
Date: Fri, 13 Dec 2024 15:46:05 -0500
Subject: [PATCH] Adjust fetch return type

---
 pinecone/data/dataclasses/__init__.py       |  1 +
 pinecone/data/dataclasses/fetch_response.py | 11 +++++++++++
 pinecone/data/index.py                      | 10 +++++++---
 3 files changed, 19 insertions(+), 3 deletions(-)
 create mode 100644 pinecone/data/dataclasses/fetch_response.py

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 818cf101..d1fb1bce 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(