klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

enable using custom Pinecone index name (#1172)

Browse Source 
Co-authored-by: Deven Patel <deven298@yahoo.com>
This commit is contained in:
Peter Jausovec
2024-01-25 00:00:10 -08:00
committed by GitHub
parent b7d365119c
commit 446d0975aa
4 changed files with 49 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
embedchain/config/vectordb/pinecone.py
View File

@@ -9,6 +9,7 @@ class PineconeDBConfig(BaseVectorDbConfig):
def __init__(
self,
collection_name: Optional[str] = None,
index_name: Optional[str] = None,
dir: Optional[str] = None,
vector_dimension: int = 1536,
metric: Optional[str] = "cosine",
@@ -17,4 +18,5 @@ class PineconeDBConfig(BaseVectorDbConfig):
self.metric = metric
self.vector_dimension = vector_dimension
self.extra_params = extra_params
self.index_name = index_name or f"{collection_name}-{vector_dimension}".lower().replace("_", "-")
super().__init__(collection_name=collection_name, dir=dir)

22
embedchain/vectordb/pinecone.py
View File

@@ -53,20 +53,21 @@ class PineconeDB(BaseVectorDB):
if not self.embedder:
raise ValueError("Embedder not set. Please set an embedder with `set_embedder` before initialization.")
# Loads the Pinecone index or creates it if not present.
def _setup_pinecone_index(self):
"""
Loads the Pinecone index or creates it if not present.
"""
pinecone.init(
api_key=os.environ.get("PINECONE_API_KEY"),
environment=os.environ.get("PINECONE_ENV"),
**self.config.extra_params,
)
self.index_name = self._get_index_name()
indexes = pinecone.list_indexes()
if indexes is None or self.index_name not in indexes:
if indexes is None or self.config.index_name not in indexes:
pinecone.create_index(
name=self.index_name, metric=self.config.metric, dimension=self.config.vector_dimension
name=self.config.index_name, metric=self.config.metric, dimension=self.config.vector_dimension
)
return pinecone.Index(self.index_name)
return pinecone.Index(self.config.index_name)
def get(self, ids: Optional[list[str]] = None, where: Optional[dict[str, any]] = None, limit: Optional[int] = None):
"""
@@ -193,18 +194,9 @@ class PineconeDB(BaseVectorDB):
Resets the database. Deletes all embeddings irreversibly.
"""
# Delete all data from the database
pinecone.delete_index(self.index_name)
pinecone.delete_index(self.config.index_name)
self._setup_pinecone_index()
# Pinecone only allows alphanumeric characters and "-" in the index name
def _get_index_name(self) -> str:
"""Get the Pinecone index for a collection
:return: Pinecone index
:rtype: str
"""
return f"{self.config.collection_name}-{self.config.vector_dimension}".lower().replace("_", "-")
@staticmethod
def _generate_filter(where: dict):
query = {}