add pinecone namespace (run-llama#1169)

joyxof · Apr 13, 2023 · 8bd08b6 · 8bd08b6
1 parent f91fde7
commit 8bd08b6
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/gpt_index/indices/vector_store/vector_indices.py b/gpt_index/indices/vector_store/vector_indices.py
@@ -210,6 +210,7 @@ def __init__(
         pinecone_index: Optional[Any] = None,
         index_name: Optional[str] = None,
         environment: Optional[str] = None,
+        namespace: Optional[str] = None,
         metadata_filters: Optional[Dict[str, Any]] = None,
         pinecone_kwargs: Optional[Dict] = None,
         insert_kwargs: Optional[Dict] = None,
@@ -230,6 +231,7 @@ def __init__(
                 pinecone_index=pinecone_index,
                 index_name=index_name,
                 environment=environment,
+                namespace=namespace,
                 metadata_filters=metadata_filters,
                 pinecone_kwargs=pinecone_kwargs,
                 insert_kwargs=insert_kwargs,

diff --git a/gpt_index/vector_stores/pinecone.py b/gpt_index/vector_stores/pinecone.py
@@ -4,20 +4,20 @@
 
 """
 
+import logging
 import os
-from typing import Any, Dict, List, Optional, cast, Callable
+from collections import Counter
 from functools import partial
+from typing import Any, Callable, Dict, List, Optional, cast
 
-from gpt_index.data_structs.node_v2 import Node, DocumentRelationship
+from gpt_index.data_structs.node_v2 import DocumentRelationship, Node
 from gpt_index.vector_stores.types import (
     NodeEmbeddingResult,
     VectorStore,
-    VectorStoreQueryResult,
     VectorStoreQuery,
     VectorStoreQueryMode,
+    VectorStoreQueryResult,
 )
-from collections import Counter
-import logging
 
 _logger = logging.getLogger(__name__)
 
@@ -128,6 +128,7 @@ def __init__(
         pinecone_index: Optional[Any] = None,
         index_name: Optional[str] = None,
         environment: Optional[str] = None,
+        namespace: Optional[str] = None,
         metadata_filters: Optional[Dict[str, Any]] = None,
         pinecone_kwargs: Optional[Dict] = None,
         insert_kwargs: Optional[Dict] = None,
@@ -148,6 +149,7 @@ def __init__(
 
         self._index_name = index_name
         self._environment = environment
+        self._namespace = namespace
         if pinecone_index is not None:
             self._pinecone_index = cast(pinecone.Index, pinecone_index)
             _logger.warn(
@@ -202,6 +204,7 @@ def config_dict(self) -> dict:
         return {
             "index_name": self._index_name,
             "environment": self._environment,
+            "namespace": self._namespace,
             "metadata_filters": self._metadata_filters,
             "pinecone_kwargs": self._pinecone_kwargs,
             "insert_kwargs": self._insert_kwargs,
@@ -264,7 +267,9 @@ def add(
                     [node.get_text()], self._tokenizer
                 )[0]
                 entry.update({"sparse_values": sparse_vector})
-            self._pinecone_index.upsert([entry], **self._pinecone_kwargs)
+            self._pinecone_index.upsert(
+                [entry], namespace=self._namespace, **self._pinecone_kwargs
+            )
             ids.append(new_id)
         return ids
 
@@ -312,6 +317,7 @@ def query(self, query: VectorStoreQuery) -> VectorStoreQueryResult:
             top_k=query.similarity_top_k,
             include_values=True,
             include_metadata=True,
+            namespace=self._namespace,
             filter=self._metadata_filters,
             **self._pinecone_kwargs,
         )

diff --git a/tests/indices/vector_store/utils.py b/tests/indices/vector_store/utils.py
@@ -29,6 +29,7 @@ def query(
         include_values: bool = True,
         include_metadata: bool = True,
         filter: Optional[Dict[str, Any]] = None,
+        namespace: Optional[str] = None,
     ) -> Any:
         """Mock query."""
         # index_mat is n x k