Pinecone re-indexing and Postgres deletion of data in test-connection issue fixes (#49)

* Supporting run_id

* Refactoring changes

* Update version

Author: gaya3-zipstack

src/unstract/adapters/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.15.1"
+__version__ = "0.16.0"
 
 
 import logging

src/unstract/adapters/base.py

@@ -1,15 +1,5 @@
 import logging
 from abc import ABC, abstractmethod
-from typing import Any, Union
-
-from llama_index.core import MockEmbedding
-from llama_index.core.embeddings import BaseEmbedding
-from llama_index.core.llms import LLM, MockLLM
-from llama_index.core.vector_stores import SimpleVectorStore
-from llama_index.core.vector_stores.types import (
-    BasePydanticVectorStore,
-    VectorStore,
-)
 
 from unstract.adapters.enums import AdapterTypes
 
@@ -50,21 +40,6 @@ def get_json_schema() -> str:
     def get_adapter_type() -> AdapterTypes:
         return ""
 
-    def get_llm_instance(self, llm_config: dict[str, Any]) -> LLM:
-        # Overriding implementations use llm_config
-        return MockLLM()
-
-    def get_vector_db_instance(
-        self, vector_db_config: dict[str, Any]
-    ) -> Union[BasePydanticVectorStore, VectorStore]:
-        # Overriding implementations use vector_db_config
-        return SimpleVectorStore()
-
-    def get_embedding_instance(
-        self, embed_config: dict[str, Any]
-    ) -> BaseEmbedding:
-        return MockEmbedding(embed_dim=1)
-
     @abstractmethod
     def test_connection(self) -> bool:
         """Override to test connection for a adapter.

src/unstract/adapters/vectordb/constants.py

@@ -2,5 +2,5 @@ class VectorDbConstants:
     VECTOR_DB_NAME = "collection_name"
     EMBEDDING_DIMENSION = "embedding_dimension"
     DEFAULT_VECTOR_DB_NAME = "unstract"
-    DEFAULT_EMBEDDING_SIZE = 1536
+    DEFAULT_EMBEDDING_SIZE = 1
     TEST_CONNECTION_EMBEDDING_SIZE = 1

src/unstract/adapters/vectordb/milvus/src/milvus.py

@@ -19,10 +19,11 @@ class Constants:
 
 class Milvus(VectorDBAdapter):
     def __init__(self, settings: dict[str, Any]):
-        super().__init__("Milvus")
-        self.config = settings
-        self.client: Optional[MilvusClient] = None
-        self.collection_name: str = VectorDbConstants.DEFAULT_VECTOR_DB_NAME
+        self._config = settings
+        self._client: Optional[MilvusClient] = None
+        self._collection_name: str = VectorDbConstants.DEFAULT_VECTOR_DB_NAME
+        self._vector_db_instance = self._get_vector_db_instance()
+        super().__init__("Milvus", self._vector_db_instance)
 
     @staticmethod
     def get_id() -> str:
@@ -48,36 +49,43 @@ def get_json_schema() -> str:
         return schema
 
     def get_vector_db_instance(self) -> VectorStore:
+        return self._vector_db_instance
+
+    def _get_vector_db_instance(self) -> VectorStore:
         try:
-            self.collection_name = VectorDBHelper.get_collection_name(
-                self.config.get(VectorDbConstants.VECTOR_DB_NAME),
-                self.config.get(VectorDbConstants.EMBEDDING_DIMENSION),
+            self._collection_name = VectorDBHelper.get_collection_name(
+                self._config.get(VectorDbConstants.VECTOR_DB_NAME),
+                self._config.get(VectorDbConstants.EMBEDDING_DIMENSION),
             )
-            dimension = self.config.get(
+            dimension = self._config.get(
                 VectorDbConstants.EMBEDDING_DIMENSION,
                 VectorDbConstants.DEFAULT_EMBEDDING_SIZE,
             )
             vector_db: VectorStore = MilvusVectorStore(
-                uri=self.config.get(Constants.URI, ""),
-                collection_name=self.collection_name,
-                token=self.config.get(Constants.TOKEN, ""),
+                uri=self._config.get(Constants.URI, ""),
+                collection_name=self._collection_name,
+                token=self._config.get(Constants.TOKEN, ""),
                 dim=dimension,
             )
             if vector_db is not None:
-                self.client = vector_db.client
+                self._client = vector_db.client
             return vector_db
         except Exception as e:
             raise AdapterError(str(e))
 
     def test_connection(self) -> bool:
-        self.config[VectorDbConstants.EMBEDDING_DIMENSION] = (
+        self._config[VectorDbConstants.EMBEDDING_DIMENSION] = (
             VectorDbConstants.TEST_CONNECTION_EMBEDDING_SIZE
         )
         vector_db = self.get_vector_db_instance()
         test_result: bool = VectorDBHelper.test_vector_db_instance(
             vector_store=vector_db
         )
         # Delete the collection that was created for testing
-        if self.client is not None:
-            self.client.drop_collection(self.collection_name)
+        if self._client is not None:
+            self._client.drop_collection(self._collection_name)
         return test_result
+
+    def close(self, **kwargs: Any) -> None:
+        if self._client:
+            self._client.close()

src/unstract/adapters/vectordb/pinecone/src/pinecone.py

@@ -2,6 +2,7 @@
 import os
 from typing import Any, Optional
 
+from llama_index.core.schema import BaseNode
 from llama_index.core.vector_stores.types import BasePydanticVectorStore
 from llama_index.vector_stores.pinecone import PineconeVectorStore
 from pinecone import NotFoundException
@@ -33,10 +34,11 @@ class Constants:
 
 class Pinecone(VectorDBAdapter):
     def __init__(self, settings: dict[str, Any]):
-        super().__init__("Pinecone")
-        self.config = settings
-        self.client: Optional[LLamaIndexPinecone] = None
-        self.collection_name: str = VectorDbConstants.DEFAULT_VECTOR_DB_NAME
+        self._config = settings
+        self._client: Optional[LLamaIndexPinecone] = None
+        self._collection_name: str = VectorDbConstants.DEFAULT_VECTOR_DB_NAME
+        self._vector_db_instance = self._get_vector_db_instance()
+        super().__init__("Pinecone", self._vector_db_instance)
 
     @staticmethod
     def get_id() -> str:
@@ -62,65 +64,101 @@ def get_json_schema() -> str:
         return schema
 
     def get_vector_db_instance(self) -> BasePydanticVectorStore:
-        try:
-            self.client = LLamaIndexPinecone(
-                api_key=str(self.config.get(Constants.API_KEY))
-            )
-            collection_name = VectorDBHelper.get_collection_name(
-                self.config.get(VectorDbConstants.VECTOR_DB_NAME),
-                self.config.get(VectorDbConstants.EMBEDDING_DIMENSION),
-            )
-            self.collection_name = collection_name.replace("_", "-").lower()
-            dimension = self.config.get(
-                VectorDbConstants.EMBEDDING_DIMENSION,
-                VectorDbConstants.DEFAULT_EMBEDDING_SIZE,
-            )
+        return self._vector_db_instance
 
-            specification = self.config.get(Constants.SPECIFICATION)
-            if specification == Constants.SPEC_POD:
-                environment = self.config.get(Constants.ENVIRONMENT)
-                spec = PodSpec(
-                    environment=environment,
-                    replicas=Constants.DEFAULT_SPEC_COUNT_VALUE,
-                    shards=Constants.DEFAULT_SPEC_COUNT_VALUE,
-                    pods=Constants.DEFAULT_SPEC_COUNT_VALUE,
-                    pod_type=Constants.DEFAULT_POD_TYPE,
-                )
-            elif specification == Constants.SPEC_SERVERLESS:
-                cloud = self.config.get(Constants.CLOUD)
-                region = self.config.get(Constants.REGION)
-                spec = ServerlessSpec(cloud=cloud, region=region)
-            logger.info(f"Setting up Pinecone spec for {spec}")
-            try:
-                self.client.describe_index(name=self.collection_name)
-            except NotFoundException:
-                logger.info(
-                    f"Index:{self.collection_name} does not exist. Creating it."
-                )
-                self.client.create_index(
-                    name=self.collection_name,
-                    dimension=dimension,
-                    metric=Constants.METRIC,
-                    spec=spec,
-                )
-            vector_db: BasePydanticVectorStore = PineconeVectorStore(
-                index_name=self.collection_name,
-                api_key=str(self.config.get(Constants.API_KEY)),
-                environment=str(self.config.get(Constants.ENVIRONMENT)),
+    def _get_vector_db_instance(self) -> BasePydanticVectorStore:
+
+        self._client = LLamaIndexPinecone(
+            api_key=str(self._config.get(Constants.API_KEY))
+        )
+        collection_name = VectorDBHelper.get_collection_name(
+            self._config.get(VectorDbConstants.VECTOR_DB_NAME),
+            self._config.get(VectorDbConstants.EMBEDDING_DIMENSION),
+        )
+        self._collection_name = collection_name.replace("_", "-").lower()
+        dimension = self._config.get(
+            VectorDbConstants.EMBEDDING_DIMENSION,
+            VectorDbConstants.DEFAULT_EMBEDDING_SIZE,
+        )
+
+        specification = self._config.get(Constants.SPECIFICATION)
+        if specification == Constants.SPEC_POD:
+            environment = self._config.get(Constants.ENVIRONMENT)
+            spec = PodSpec(
+                environment=environment,
+                replicas=Constants.DEFAULT_SPEC_COUNT_VALUE,
+                shards=Constants.DEFAULT_SPEC_COUNT_VALUE,
+                pods=Constants.DEFAULT_SPEC_COUNT_VALUE,
+                pod_type=Constants.DEFAULT_POD_TYPE,
             )
-            return vector_db
-        except Exception as e:
-            raise AdapterError(str(e))
+        elif specification == Constants.SPEC_SERVERLESS:
+            cloud = self._config.get(Constants.CLOUD)
+            region = self._config.get(Constants.REGION)
+            spec = ServerlessSpec(cloud=cloud, region=region)
+        logger.info(f"Setting up Pinecone spec for {spec}")
+        try:
+            self._client.describe_index(name=self._collection_name)
+        except NotFoundException:
+            logger.info(f"Index:{self._collection_name} does not exist. Creating it.")
+            self._client.create_index(
+                name=self._collection_name,
+                dimension=dimension,
+                metric=Constants.METRIC,
+                spec=spec,
+            )
+        self.vector_db: BasePydanticVectorStore = PineconeVectorStore(
+            index_name=self._collection_name,
+            api_key=str(self._config.get(Constants.API_KEY)),
+            environment=str(self._config.get(Constants.ENVIRONMENT)),
+        )
+        return self.vector_db
 
     def test_connection(self) -> bool:
-        self.config[VectorDbConstants.EMBEDDING_DIMENSION] = (
+        self._config[VectorDbConstants.EMBEDDING_DIMENSION] = (
             VectorDbConstants.TEST_CONNECTION_EMBEDDING_SIZE
         )
         vector_db = self.get_vector_db_instance()
         test_result: bool = VectorDBHelper.test_vector_db_instance(
             vector_store=vector_db
         )
         # Delete the collection that was created for testing
-        if self.client:
-            self.client.delete_index(self.collection_name)
+        if self._client:
+            self._client.delete_index(self._collection_name)
         return test_result
+
+    def close(self, **kwargs: Any) -> None:
+        # Close connection is not defined for this client
+        pass
+
+    def delete(self, ref_doc_id: str, **delete_kwargs: dict[Any, Any]) -> None:
+        specification = self._config.get(Constants.SPECIFICATION)
+        if specification == Constants.SPEC_SERVERLESS:
+            # To delete all records representing chunks of a single document,
+            # first list the record IDs based on their common ID prefix,
+            # and then delete the records by ID:
+            try:
+                index = self._client.Index(self._collection_name)  # type: ignore
+                # Get all record having the ref_doc_id and delete them
+                for ids in index.list(prefix=ref_doc_id):
+                    logger.info(ids)
+                    index.delete(ids=ids)
+            except Exception as e:
+                raise AdapterError(str(e))
+        elif specification == Constants.SPEC_POD:
+            if self.vector_db.environment == "gcp-starter":  # type: ignore
+                raise AdapterError(
+                    "Re-indexing is not supported on Starter indexes. "
+                    "Use Serverless or paid plan for Pod spec"
+                )
+            else:
+                super().delete(ref_doc_id=ref_doc_id, **delete_kwargs)
+
+    def add(
+        self,
+        ref_doc_id: str,
+        nodes: list[BaseNode],
+    ) -> list[str]:
+        for i, node in enumerate(nodes):
+            node_id = ref_doc_id + "-" + node.node_id
+            nodes[i].id_ = node_id
+        return self.vector_db.add(nodes=nodes)