ai-parrot 0.8.3__cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

parrot/stores/abstract.py

@@ -0,0 +1,248 @@
+from abc import ABC, abstractmethod
+from typing import List, Union
+import importlib
+from collections.abc import Callable
+from langchain.docstore.document import Document
+from langchain_core.vectorstores import VectorStoreRetriever
+from navconfig.logging import logging
+from ..conf import (
+    EMBEDDING_DEFAULT_MODEL
+)
+from ..exceptions import ConfigError  # pylint: disable=E0611
+from .embeddings import supported_embeddings
+
+
+class AbstractStore(ABC):
+    """AbstractStore class.
+
+        Base class for all Database Vector Stores.
+    Args:
+        embeddings (str): Embedding name.
+
+    Supported Vector Stores:
+        - Qdrant
+        - Milvus
+        - Faiss
+        - Chroma
+        - PgVector
+    """
+
+    def __init__(
+        self,
+        embedding_model: Union[dict, str] = None,
+        embedding: Union[dict, Callable] = None,
+        **kwargs
+    ):
+        self.client: Callable = None
+        self.vector: Callable = None
+        self._embed_: Callable = None
+        self._connected: bool = False
+        if embedding_model is not None:
+            if isinstance(embedding_model, str):
+                self.embedding_model = {
+                    'model_name': embedding_model,
+                    'model_type': 'huggingface'
+                }
+            elif isinstance(embedding_model, dict):
+                self.embedding_model = embedding_model
+            else:
+                raise ValueError(
+                    "Embedding Model must be a string or a dictionary."
+                )
+        # Use or not connection to a vector database:
+        self._use_database: bool = kwargs.get('use_database', True)
+        # Database Information:
+        self.collection_name: str = kwargs.get('collection_name', 'my_collection')
+        self.dimension: int = kwargs.get("dimension", 768)
+        self._metric_type: str = kwargs.get("metric_type", 'COSINE')
+        self._index_type: str = kwargs.get("index_type", 'IVF_FLAT')
+        self.database: str = kwargs.get('database', '')
+        self.index_name = kwargs.get("index_name", "my_index")
+        if embedding is not None:
+            if isinstance(embedding, str):
+                self.embedding_model = {
+                    'model_name': embedding,
+                    'model_type': 'huggingface'
+                }
+            elif isinstance(embedding, dict):
+                self.embedding_model = embedding
+            else:
+                # is a callable:
+                self.embedding_model = {
+                    'model_name': EMBEDDING_DEFAULT_MODEL,
+                    'model_type': 'huggingface'
+                }
+                self._embed_ = embedding
+        self.logger = logging.getLogger(
+            f"Store.{__name__}"
+        )
+        # Client Connection (if required):
+        self._connection = None
+        # Create the Embedding Model:
+        self._embed_ = self.create_embedding(
+            embedding_model=self.embedding_model
+        )
+
+    @property
+    def connected(self) -> bool:
+        return self._connected
+
+    def is_connected(self):
+        return self._connected
+
+    @abstractmethod
+    async def connection(self) -> tuple:
+        pass
+
+    @abstractmethod
+    async def disconnect(self) -> None:
+        pass
+
+    # Async Context Manager
+    async def __aenter__(self):
+        if self._use_database:
+            if not self._connection:
+                await self.connection()
+        return self
+
+    async def _free_resources(self):
+        self._embed_.free()
+        self._embed_ = None
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        # closing Embedding
+        try:
+            await self.disconnect()
+        except RuntimeError:
+            pass
+
+    @abstractmethod
+    def get_vector(self, metric_type: str = None, **kwargs):
+        pass
+
+    def get_vectorstore(self):
+        return self.get_vector()
+
+    @abstractmethod
+    async def similarity_search(
+        self,
+        query: str,
+        collection: Union[str, None] = None,
+        limit: int = 2
+    ) -> list:  # noqa
+        pass
+
+    @abstractmethod
+    async def from_documents(
+        self,
+        documents: List[Document],
+        collection: Union[str, None] = None,
+        **kwargs
+    ) -> Callable:
+        """
+        Create Vector Store from Documents.
+
+        Args:
+            documents (List[Document]): List of Documents.
+            collection (str): Collection Name.
+            kwargs: Additional Arguments.
+
+        Returns:
+            Callable VectorStore.
+        """
+
+    @abstractmethod
+    async def add_documents(
+        self,
+        documents: List[Document],
+        collection: Union[str, None] = None,
+        **kwargs
+    ) -> None:
+        """
+        Add Documents to Vector Store.
+
+        Args:
+            documents (List[Document]): List of Documents.
+            collection (str): Collection Name.
+            kwargs: Additional Arguments.
+
+        Returns:
+            None.
+        """
+
+    def create_embedding(
+        self,
+
+        embedding_model: dict,
+        **kwargs
+    ):
+        """
+        Create Embedding Model.
+
+        Args:
+            embedding_model (dict): Embedding Model Configuration.
+            kwargs: Additional Arguments.
+
+        Returns:
+            Callable: Embedding Model.
+
+        """
+        model_type = embedding_model.get('model_type', 'huggingface')
+        model_name = embedding_model.get('model_name', EMBEDDING_DEFAULT_MODEL)
+        if model_type not in supported_embeddings:
+            raise ConfigError(
+                f"Embedding Model Type: {model_type} not supported."
+            )
+        embed_cls = supported_embeddings[model_type]
+        cls_path = f".embeddings.{model_type}"  # Relative module path
+        try:
+            embed_module = importlib.import_module(
+                cls_path,
+                package=__package__
+            )
+            embed_obj = getattr(embed_module, embed_cls)
+            return embed_obj(
+                model_name=model_name,
+                **kwargs
+            )
+        except ImportError as e:
+            raise ConfigError(
+                f"Error Importing Embedding Model: {model_type}"
+            ) from e
+
+    def get_default_embedding(self):
+        embed_model = {
+            'model_name': EMBEDDING_DEFAULT_MODEL,
+            'model_type': 'huggingface'
+        }
+        return self.create_embedding(
+            embedding_model=embed_model
+        )
+
+    def generate_embedding(self, documents: List[Document]):
+        if not self._embed_:
+            self._embed_ = self.get_default_embedding()
+
+        # Using the Embed Model to Generate Embeddings:
+        embeddings = self._embed_.embed_documents(documents)
+        return embeddings
+
+    def as_retriever(
+        self,
+        metric_type: str = 'COSINE',
+        index_type: str = 'IVF_FLAT',
+        search_type: str = 'similarity',
+        chain_type: str = 'stuff',
+        search_kwargs: dict = None
+    ) -> Callable:
+        vector = self.get_vector(metric_type=metric_type, index_type=index_type)
+        if not vector:
+            raise ConfigError(
+                "Vector Store is not connected. Check your connection."
+            )
+        return VectorStoreRetriever(
+            vectorstore=vector,
+            search_type=search_type,
+            chain_type=chain_type,
+            search_kwargs=search_kwargs
+        )

parrot/stores/chroma.py

@@ -0,0 +1,188 @@
+from collections.abc import Callable
+from typing import Optional, Union
+from uuid import uuid4
+import logging
+from langchain.docstore.document import Document
+from langchain.memory import VectorStoreRetrieverMemory
+import chromadb
+from langchain_chroma import Chroma
+from .abstract import AbstractStore
+from ..conf import CHROMADB_HOST, CHROMADB_PORT
+
+
+logging.getLogger('chromadb').setLevel(logging.INFO)
+
+
+class ChromaStore(AbstractStore):
+    """Chroma DB Store Class.
+
+    Using Chroma as Document Vector Store.
+
+    """
+
+    def __init__(
+        self,
+        embedding_model: Union[dict, str] = None,
+        embedding: Union[dict, Callable] = None,
+        **kwargs
+    ):
+        super().__init__(
+            embedding_model=embedding_model,
+            embedding=embedding,
+            **kwargs
+        )
+        self.database_path: str = kwargs.pop('database_path', 'chroma.db')
+        self._ephemeral: bool = kwargs.pop('ephemeral', False)
+        self._local: bool = kwargs.pop('local', False)
+        self.host = kwargs.pop("host", CHROMADB_HOST)
+        self.port = kwargs.pop("port", CHROMADB_PORT)
+        self._collection = None
+
+    async def connection(self):
+        """Connection to ChromaDB.
+
+        Args:
+            alias (str): Database alias.
+
+        Returns:
+            Callable: ChromaDB connection.
+
+        """
+        if self._ephemeral:
+            self._connection = chromadb.Client()
+        elif self._local:
+            self._connection = chromadb.PersistentClient(
+                path=self.database_path,
+                database=self.database,
+            )
+        else:
+            # Client-Server Connection:
+            self._connection = chromadb.HttpClient(
+                host=self.host,
+                port=self.port,
+                database=self.database,
+            )
+        self._collection = self._connection.get_or_create_collection(self.collection_name)
+        self._connected = True
+        return self._connection
+
+    async def disconnect(self) -> None:
+        """
+        Closing the Connection on ChromaDB
+        """
+        self._connection = None
+        self._connected = False
+
+    def get_vector(
+        self,
+        collection: Union[str, None] = None,
+        embedding: Optional[Callable] = None,
+    ) -> Chroma:
+
+        if not collection:
+            collection = self.collection_name
+        if embedding is not None:
+            _embed_ = embedding
+        else:
+            _embed_ = self._embed_ or self.create_embedding(
+                embedding_model=self.embedding_model
+            )
+        return Chroma(
+            collection_name=self.collection_name,
+            embedding_function=_embed_.embedding,
+            client=self._connection,
+            create_collection_if_not_exists=True,
+        )
+
+    async def from_documents(self, documents: list[Document], collection: str = None, **kwargs):
+        """
+        Save Documents as Vectors in Chroma.
+        """
+        vectordb = await Chroma.afrom_documents(
+            documents=documents,
+            embedding=self._embed_.embedding,
+            connection=self._connection,
+        )
+        return vectordb
+
+
+    async def add_texts(self, objects: list, collection: str = None):
+        """
+        Add Texts to ChromaDB
+        """
+        async with self:
+            collection = self._connection.get_or_create_collection(collection)
+            for i, doc in enumerate(objects):
+                collection.add(ids=[str(i)], documents=[doc])
+        return True
+
+    async def add_documents(
+        self,
+        documents: list,
+        collection: str = None,
+        embedding: Optional[Callable] = None,
+    ) -> bool:
+        """Add Documents to ChromaDB"""
+
+        if collection is None:
+            collection = self.collection_name
+
+        async with self:
+            collection_obj = self._connection.get_or_create_collection(collection)
+            uuids = [str(uuid4()) for _ in range(len(documents))]
+            vector_db = self.get_vector(collection=collection, embedding=embedding)
+            await vector_db.aadd_documents(documents=documents, ids=uuids)
+
+        return True
+
+
+    async def update_documents(
+        self,
+        documents: list,
+        collection: str = None,
+        embedding: Optional[Callable] = None,
+    )   -> bool:
+        """
+        Update Documents to ChromaDB
+        """
+        async with self:
+            collection = self._connection.get_or_create_collection(collection)
+            vector_db = self.get_vector(collection=collection, embedding=embedding)
+            # Split the documents into ids and documents
+            if all('id' in doc for doc in documents):
+                ids = [doc.pop('id') for doc in documents]
+                vector_db.update_documents(documents=documents, ids=ids)
+                return True
+            return False
+
+    async def similarity_search(
+        self,
+        query: str,
+        collection: Union[str, None] = None,
+        embedding: Optional[Callable] = None,
+        limit: int = 2,
+        filter: Optional[dict] = None,
+    ) -> list:
+        if collection is None:
+            collection = self.collection_name
+        async with self:
+            vector_db = self.get_vector(collection=collection, embedding=embedding)
+            return vector_db.similarity_search(query, k=limit, filter=filter)
+
+    def memory_retriever(
+        self,
+        documents: Optional[list] = None,
+        num_results: int  = 5
+    ) -> VectorStoreRetrieverMemory:
+        if not documents:
+            documents = []
+        vectordb = Chroma.from_documents(
+            documents=documents,
+            embedding=self._embed_.embedding,
+            connection=self._connection,
+        )
+        retriever = Chroma.as_retriever(
+            vectordb,
+            search_kwargs=dict(k=num_results)
+        )
+        return VectorStoreRetrieverMemory(retriever=retriever)

parrot/stores/duck.py

@@ -0,0 +1,162 @@
+from collections.abc import Callable
+from typing import Optional, Union
+import duckdb
+from langchain.docstore.document import Document
+from langchain.memory import VectorStoreRetrieverMemory
+from langchain_community.vectorstores import DuckDB
+from .abstract import AbstractStore
+
+
+class DuckDBStore(AbstractStore):
+    """DuckDB Store Class.
+
+    Using DuckDB as Document Vector Store.
+
+    """
+    default_config: dict ={
+        "enable_external_access": "false",
+        "autoinstall_known_extensions": "false",
+        "autoload_known_extensions": "false"
+    }
+
+    def __init__(
+        self,
+        embedding_model: Union[dict, str] = None,
+        embedding: Union[dict, Callable] = None,
+        **kwargs
+    ):
+        super().__init__(
+            embedding_model=embedding_model,
+            embedding=embedding,
+            **kwargs
+        )
+        self.credentials = {
+            "database": self.database,
+        }
+
+        config: dict = kwargs.pop("config", {})
+        self.config = {
+            **self.default_config,
+            **config
+        }
+
+    async def connection(self, alias: str = None):
+        """Connection to DuckDB.
+
+        Args:
+            alias (str): Database alias.
+
+        Returns:
+            Callable: DuckDB connection.
+
+        """
+        self._connection = duckdb.connect(**self.credentials, config=self.config)
+        self._connected = True
+        return self._connection
+
+    async def disconnect(self) -> None:
+        """
+        Closing the Connection on DuckDB
+        """
+        try:
+            if self._connection:
+                self._connection.close()
+        except Exception as err:
+            raise RuntimeError(
+                message=f"{__name__!s}: Closing Error: {err!s}"
+            ) from err
+        finally:
+            self._connection = None
+            self._connected = False
+
+    def get_vector(
+        self,
+        collection: Union[str, None] = None,
+        embedding: Optional[Callable] = None,
+        metadata_field: str = 'id',
+        text_field: str = 'text',
+        vector_key: str = 'vector',
+    ) -> DuckDB:
+
+        if not collection:
+            collection = self.collection_name
+        if embedding is not None:
+            _embed_ = embedding
+        else:
+            _embed_ = self.create_embedding(
+                embedding_model=self.embedding_model
+            )
+        return DuckDB(
+            connection=self._connection,
+            table_name=collection,
+            embedding=_embed_,
+            vector_key=vector_key,
+            text_key=text_field,
+            id_key=metadata_field
+        )
+
+
+    async def add_texts(self, objects: list, collection: str = None):
+        """
+        Add Texts to DuckDB
+        """
+        async with self:
+            store = self.get_vector(collection=collection)
+            store.add_texts(objects)
+        return True
+
+    async def similarity_search(
+        self,
+        query: str,
+        collection: Union[str, None] = None,
+        embedding: Optional[Callable] = None,
+        limit: int = 2,
+    ) -> list:
+        if collection is None:
+            collection = self.collection_name
+        async with self:
+            vector_db = self.get_vector(collection=collection, embedding=embedding)
+            return await vector_db.asimilarity_search(query, k=limit)
+
+    def memory_retriever(
+        self,
+        documents: Optional[list] = None,
+        num_results: int  = 5
+    ) -> VectorStoreRetrieverMemory:
+        if not documents:
+            documents = []
+        vectordb = DuckDB.from_documents(
+            documents=documents,
+            embedding=self._embed_.embedding,
+            connection=self._connection,
+        )
+        retriever = DuckDB.as_retriever(
+            vectordb,
+            search_kwargs=dict(k=num_results)
+        )
+        return VectorStoreRetrieverMemory(retriever=retriever)
+
+    async def from_documents(self, documents: list[Document], collection: str = None, **kwargs):
+        """
+        Save Documents as Vectors in DuckDB.
+        """
+        if not collection:
+            collection = self.collection_name
+        vectordb = await DuckDB.afrom_documents(
+            documents,
+            embedding=self._embed_.embedding,
+            connection=self._connection,
+        )
+        return vectordb
+
+    async def add_documents(self, documents: list[Document], collection: str = None, **kwargs):
+        """
+        Add Documents as Vectors in DuckDB.
+        """
+        if not collection:
+            collection = self.collection_name
+        vectordb = self.get_vector(collection=collection)
+        result = await vectordb.aadd_documents(
+            documents=documents
+        )
+        return result

parrot/stores/embeddings/__init__.py

@@ -0,0 +1,10 @@
+supported_embeddings = {
+    'openai': 'OpenAIEmbed',
+    'google': 'GoogleEmbed',
+    'vertexai': 'VertexAIEmbed',
+    'huggingface': 'HugginfaceEmbed',
+    'fastembed': 'FastembedEmbed',
+    'bge': 'BgeEmbed',
+    'ollama': 'OllamaEmbed',
+    'transformers': 'TransformersEmbed'
+}

parrot/stores/embeddings/abstract.py

@@ -0,0 +1,46 @@
+from abc import ABC, abstractmethod
+from ...conf import (
+    MAX_BATCH_SIZE,
+    EMBEDDING_DEFAULT_MODEL,
+    EMBEDDING_DEVICE
+)
+
+class AbstractEmbed(ABC):
+    """A wrapper class for Create embeddings."""
+    model_name: str = EMBEDDING_DEFAULT_MODEL
+    encode_kwargs: str = {
+        'normalize_embeddings': True,
+        "batch_size": MAX_BATCH_SIZE
+    }
+    model_kwargs = {
+        'device': EMBEDDING_DEVICE,
+        'trust_remote_code':True
+    }
+
+    def __init__(self, model_name: str = None, **kwargs):
+        self._embedding = self._create_embedding(model_name, **kwargs)
+
+    @property
+    def embedding(self):
+        return self._embedding
+
+    def free(self):
+        """
+        Free the resources.
+        """
+        pass
+
+    def _get_device(self):
+        return EMBEDDING_DEVICE
+
+    @abstractmethod
+    def _create_embedding(self, model_name: str = None, **kwargs):
+        """
+        Create Embedding Model.
+        Args:
+            model_name (str): The name of the model to use.
+
+        Returns:
+            Callable: Embedding Model.
+        """
+        pass

parrot/stores/embeddings/base.py

@@ -0,0 +1,52 @@
+from typing import Optional
+import gc
+import torch
+from .abstract import AbstractEmbed
+from ...conf import CUDA_DEFAULT_DEVICE, EMBEDDING_DEVICE
+
+
+class BaseEmbed(AbstractEmbed):
+    """A wrapper class for Base embeddings.
+
+    Use this class to Embedding Models that requires Torch/Transformers.
+    """
+    model_kwargs = {
+        'device': EMBEDDING_DEVICE,
+        'trust_remote_code':True
+    }
+    def _get_device(self, device_type: str = None, cuda_number: Optional[int] = None):
+        """Get Default device for Torch and transformers.
+
+        """
+        # torch.backends.cudnn.deterministic = True
+        if device_type is not None:
+            return torch.device(device_type)
+        if torch.cuda.is_available():
+            if CUDA_DEFAULT_DEVICE == 'cpu':
+                # Use CPU even if CUDA is available
+                return torch.device('cpu')
+            if cuda_number is not None:
+                # Use specified CUDA GPU
+                return torch.device(f'cuda:{cuda_number}')
+            # Use CUDA GPU if available
+            if cuda_number is None:
+                return torch.device(f'cuda:0')
+        if torch.backends.mps.is_available():
+            # Use CUDA Multi-Processing Service if available
+            return torch.device("mps")
+        if EMBEDDING_DEVICE == 'cuda':
+            if cuda_number is None and CUDA_DEFAULT_DEVICE != 'cpu':
+                # Use CUDA GPU if available
+                cuda_number = CUDA_DEFAULT_DEVICE
+                return torch.device(f'cuda:{cuda_number}')
+        return torch.device(EMBEDDING_DEVICE)
+
+    def free(self):
+        """
+        Free the resources.
+        """
+        try:
+            gc.collect()               # Run Python garbage collector to free unreferenced objects
+            torch.cuda.empty_cache()   # Release cached memory blocks back to the GPU
+        except Exception as e:
+            print(f"Error: {e}")