ai-parrot 0.8.3__cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

parrot/stores/embeddings/bge.py

@@ -0,0 +1,20 @@
+from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+from .base import BaseEmbed
+
+
+class BgeEmbed(BaseEmbed):
+    """A wrapper class for VertexAI embeddings."""
+    model_name: str = "BAAI/bge-large-en-v1.5"
+
+    def _create_embedding(self, model_name: str = None, **kwargs):
+        # Embedding Model:
+        device = self._get_device()
+        model_args = {
+            **self.model_kwargs,
+            'device': device,
+        }
+        return HuggingFaceBgeEmbeddings(
+            model_name=model_name or self.model_name,
+            model_kwargs=model_args,
+            encode_kwargs=self.encode_kwargs
+        )

parrot/stores/embeddings/fastembed.py

@@ -0,0 +1,17 @@
+from langchain_community.embeddings.fastembed import (
+    FastEmbedEmbeddings
+)
+from .abstract import AbstractEmbed
+
+
+class FastembedEmbed(AbstractEmbed):
+    """A wrapper class for FastEmbed embeddings."""
+    model_name: str = "BAAI/bge-large-en-v1.5"
+
+    def _create_embedding(self, model_name: str = None, **kwargs):
+        # Embedding Model:
+        return FastEmbedEmbeddings(
+            model_name=model_name,
+            max_length=1024,
+            threads=4
+        )

parrot/stores/embeddings/google.py

@@ -0,0 +1,18 @@
+from langchain_google_genai import (
+    GoogleGenerativeAIEmbeddings
+)
+from .abstract import AbstractEmbed
+from ...conf import GOOGLE_API_KEY
+
+class GoogleEmbed(AbstractEmbed):
+    """A wrapper class for VertexAI embeddings."""
+    model_name: str = "models/embedding-001"
+
+    def _create_embedding(self, model_name: str = None, api_key: str = None, **kwargs):
+        # Embedding Model:
+        api_key = api_key or GOOGLE_API_KEY
+        return GoogleGenerativeAIEmbeddings(
+            model=model_name or self.model_name,
+            google_api_key=api_key,
+            **kwargs
+        )

parrot/stores/embeddings/huggingface.py

@@ -0,0 +1,20 @@
+from langchain_huggingface import HuggingFaceEmbeddings
+from .base import BaseEmbed
+
+
+class HugginfaceEmbed(BaseEmbed):
+    """A wrapper class for HuggingFace embeddings."""
+    model_name: str = "sentence-transformers/all-mpnet-base-v2"
+
+    def _create_embedding(self, model_name: str = None, **kwargs):
+        # Embedding Model:
+        device = self._get_device()
+        model_args = {
+            **self.model_kwargs,
+            'device': device,
+        }
+        return HuggingFaceEmbeddings(
+            model_name=model_name or self.model_name,
+            model_kwargs=model_args,
+            encode_kwargs=self.encode_kwargs
+        )

parrot/stores/embeddings/ollama.py

@@ -0,0 +1,14 @@
+from langchain_ollama import OllamaEmbeddings
+from .base import BaseEmbed
+
+
+class OllamaEmbed(BaseEmbed):
+    """A wrapper class for Ollama embeddings."""
+    model_name: str = "llama3"
+
+    def _create_embedding(self, model_name: str = None, **kwargs):
+        # Embedding Model:
+        return OllamaEmbeddings(
+            model_name=model_name or self.model_name,
+            **kwargs
+        )

parrot/stores/embeddings/openai.py

@@ -0,0 +1,26 @@
+from langchain_openai import (  # pylint: disable=E0401, E0611
+    OpenAIEmbeddings
+)
+from .abstract import AbstractEmbed
+from ...conf import OPENAI_API_KEY, OPENAI_ORGANIZATION
+
+
+class OpenAIEmbed(AbstractEmbed):
+    """A wrapper class for OpenAI embeddings."""
+    model_name: str = "text-embedding-3-large"
+
+    def _create_embedding(
+        self,
+        model_name: str = None,
+        api_key: str = None,
+        organization: str = None,
+        **kwargs
+    ):
+        # Embedding
+        return OpenAIEmbeddings(
+            model=model_name or self.model_name,
+            dimensions=kwargs.get('dimensions', 512),
+            api_key=api_key or OPENAI_API_KEY,
+            organization=organization or OPENAI_ORGANIZATION,
+            max_retries=kwargs.get('max_retries', 4),
+        )

parrot/stores/embeddings/transformers.py

@@ -0,0 +1,21 @@
+from langchain_huggingface import HuggingFaceEmbeddings
+from .base import BaseEmbed
+
+
+class TransformersEmbed(BaseEmbed):
+    """A wrapper class for Transformers embeddings."""
+    model_name: str = "all-MiniLM-L6-v2"
+    #  "sentence-transformers/all-mpnet-base-v2"
+
+    def _create_embedding(self, model_name: str = None, **kwargs):
+        # Embedding Model:
+        device = self._get_device()
+        model_args = {
+            **self.model_kwargs,
+            'device': device,
+        }
+        return HuggingFaceEmbeddings(
+            model_name=model_name or self.model_name,
+            model_kwargs=model_args,
+            encode_kwargs=self.encode_kwargs
+        )

parrot/stores/embeddings/vertexai.py

@@ -0,0 +1,17 @@
+from langchain_google_vertexai import VertexAIEmbeddings
+from ...conf import VERTEX_PROJECT_ID, VERTEX_REGION
+from .abstract import AbstractEmbed
+
+class VertexAIEmbed(AbstractEmbed):
+    """A wrapper class for VertexAI embeddings."""
+    model_name: str = "text-embedding-004"
+
+    def _create_embedding(self, model_name: str = None, project_id: str = None, region: str = None):
+        # Embedding Model:
+        return VertexAIEmbeddings(
+            model_name=model_name or self.model_name,
+            project=project_id or VERTEX_PROJECT_ID,
+            location=region or VERTEX_REGION,
+            request_parallelism=5,
+            max_retries=4,
+        )

parrot/stores/empty.py

@@ -0,0 +1,10 @@
+class EmptyStore:
+    """
+    Empty Store reference, used on bots without Vector Store Support.
+    """
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        pass

parrot/stores/faiss.py

@@ -0,0 +1,160 @@
+from collections.abc import Callable
+from pathlib import PurePath
+from typing import Optional, Union
+from uuid import uuid4
+import faiss
+from langchain.docstore.document import Document
+from langchain.memory import VectorStoreRetrieverMemory
+from langchain_community.vectorstores import FAISS
+from langchain_community.docstore.in_memory import InMemoryDocstore
+from langchain_community.vectorstores.utils import DistanceStrategy
+from .abstract import AbstractStore
+
+
+class FaissStore(AbstractStore):
+    """FAISS DB Store Class.
+
+    Using FAISS as an in-memory 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.index_path: PurePath = kwargs.pop('index_path', 'faiss_index')
+        if not self.index_path.exists():
+            self.index_path.mkdir(exist_ok=True)
+
+    async def connection(self):
+        """Initialize FAISS vector store.
+
+        If an index exists, load it; otherwise, create a new FAISS store.
+        """
+        try:
+            index_file = self.index_path.joinpath("index.faiss")
+            if not index_file.exists():
+                raise FileNotFoundError
+            self._connection = FAISS.load_local(
+                folder_path=self.index_path,
+                embeddings=self._embed_.embedding,
+                allow_dangerous_deserialization=True,
+                distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE
+            )
+        except FileNotFoundError:
+            # Create a new FAISS index if none exists.
+            print(len(self._embed_.embedding.embed_query("test")))
+            index = faiss.IndexFlatL2(len(self._embed_.embedding.embed_query("test")))
+            self._connection = FAISS(
+                embedding_function=self._embed_.embedding,
+                index=index,
+                docstore=InMemoryDocstore(),
+                index_to_docstore_id={},
+                normalize_L2=True,
+                distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE
+            )
+            self._connection.save_local(self.index_path)  # Save the new index
+        self._connected = True
+        return self._connection
+
+    async def disconnect(self) -> None:
+        """Clears FAISS in-memory index."""
+        self._connection = None
+        self._connected = False
+
+    def get_vector(
+        self,
+        embedding: Optional[Callable] = None,
+    ) -> FAISS:
+        """Returns FAISS VectorStore instance."""
+        if embedding is not None:
+            _embed_ = embedding
+        else:
+            _embed_ = self.create_embedding(
+                embedding_model=self.embedding_model
+            )
+        return FAISS.load_local(
+            folder_path=self.index_path,
+            embeddings=_embed_,
+            allow_dangerous_deserialization=True,
+            distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE
+        )
+
+    async def from_documents(self, documents: list[Document], **kwargs):
+        """Save Documents as Vectors in FAISS."""
+        vectordb = await FAISS.afrom_documents(
+            documents=documents,
+            embedding=self._embed_.embedding,
+            allow_dangerous_deserialization=True,
+            distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE
+        )
+        vectordb.save_local(self.index_path)  # Persist FAISS index
+        return vectordb
+
+    async def add_documents(
+        self,
+        documents: list,
+        embedding: Optional[Callable] = None,
+    ) -> bool:
+        """Add Documents to FAISS."""
+        async with self:
+            vector_db = self.get_vector(embedding=embedding)
+            await vector_db.aadd_documents(documents=documents)
+            vector_db.save_local(self.index_path)  # Save updated index
+        return True
+
+    async def update_documents(
+        self,
+        documents: list,
+        embedding: Optional[Callable] = None,
+    ) -> bool:
+        """
+        Update Documents in FAISS (FAISS does not natively support updates).
+        """
+        async with self:
+            vector_db = self.get_vector(embedding=embedding)
+            if all('id' in doc for doc in documents):
+                ids = [doc.pop('id') for doc in documents]
+                vector_db.delete(ids)  # Remove old entries
+                await vector_db.aadd_documents(documents=documents)  # Add new versions
+                vector_db.save_local(self.index_path)
+                return True
+            return False
+
+    async def similarity_search(
+        self,
+        query: str,
+        embedding: Optional[Callable] = None,
+        limit: int = 2,
+    ) -> list:
+        """Performs similarity search in FAISS."""
+        async with self:
+            vector_db = self.get_vector(embedding=embedding)
+            return vector_db.similarity_search(query, k=limit)
+
+    def memory_retriever(
+        self,
+        documents: Optional[list] = None,
+        num_results: int = 5
+    ) -> VectorStoreRetrieverMemory:
+        """Retrieves stored memory-based documents."""
+        if not documents:
+            documents = []
+        vectordb = FAISS.from_documents(
+            documents=documents,
+            embedding=self._embed_.embedding,
+            allow_dangerous_deserialization=True,
+            distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE
+        )
+        retriever = FAISS.as_retriever(
+            vectordb,
+            search_kwargs=dict(k=num_results)
+        )
+        return VectorStoreRetrieverMemory(retriever=retriever)