ai-parrot 0.5.8__tar.gz → 0.5.10__tar.gz

{ai_parrot-0.5.8 → ai_parrot-0.5.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ai-parrot
-Version: 0.5.8
+Version: 0.5.10
 Summary: Live Chatbots based on Langchain chatbots and Agents     Integrated into Navigator Framework or used into any aiohttp applications.
 Home-page: https://github.com/phenobarbital/ai-parrot
 Author: Jesus Lara

{ai_parrot-0.5.8 → ai_parrot-0.5.10}/parrot/bots/abstract.py

@@ -959,6 +959,7 @@ class AbstractBot(DBInterface, ABC):
                 return_messages=True
             )
         async with self.store as store:  #pylint: disable=E1101
+            print('STORES > ', self.stores, store)
             if self._use_vector:
                 if len(self.stores) > 1:
                     retriever = MultiVectorStoreRetriever(

{ai_parrot-0.5.8 → ai_parrot-0.5.10}/parrot/bots/chatbot.py

@@ -185,8 +185,8 @@ class Chatbot(AbstractBot):
         self._vector_store = bot.database
         self.configure_store()
         # after configuration, setup the chatbot
-        if bot.template_prompt:
-            self.template_prompt = bot.template_prompt
+        if bot.system_prompt_template:
+            self.system_prompt_template = bot.system_prompt_template
         self._define_prompt(
             config={}
         )

ai_parrot-0.5.10/parrot/bots/retrievals/multi.py

@@ -0,0 +1,102 @@
+from typing import List, Dict, Any, Optional
+from langchain_core.retrievers import BaseRetriever
+from langchain_core.callbacks import CallbackManagerForRetrieverRun
+from langchain_core.documents import Document
+
+
+class MultiVectorStoreRetriever(BaseRetriever):
+    """
+    This aggregator retriever queries multiple vector stores
+    and merges the results into a single list.
+    """
+
+    stores: List[Any]
+    search_kwargs: Dict[str, Any]
+    chain_type: str
+    search_type: str
+    metric_type: str
+
+    def __init__(
+        self,
+        stores: List[Any],
+        metric_type: str = 'COSINE',
+        chain_type: str = 'stuff',
+        search_type: str = 'similarity',
+        search_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize the MultiVectorStoreRetriever.
+
+        Args:
+            stores: List of vector stores that provide an as_retriever() method
+            metric_type: Similarity metric (COSINE, EUCLIDEAN, DOT_PRODUCT)
+            chain_type: Chain type for the retriever
+            search_type: Search type (similarity, mmr)
+            search_kwargs: Dict to pass to each store's retrieve method (like 'k' for top_k)
+        """
+        super().__init__()
+        self.stores = stores
+        self.search_kwargs = search_kwargs or {}
+        self.chain_type = chain_type
+        self.search_type = search_type
+        self.metric_type = metric_type
+
+    def _get_relevant_documents(
+        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+    ) -> List[Document]:
+        """
+        Get documents relevant to the query.
+        Implements the required method from BaseRetriever.
+
+        Args:
+            query: Query string
+            run_manager: CallbackManager for the run
+
+        Returns:
+            List of relevant documents
+        """
+        all_results = []
+        for store in self.stores:
+            retriever = store.as_retriever(
+                search_type=self.search_type,
+                search_kwargs=self.search_kwargs,
+                chain_type=self.chain_type
+            )
+            # Pass the run_manager to the sub-retrievers if they support it
+            if hasattr(retriever, "_get_relevant_documents"):
+                callback_manager = run_manager.get_child()
+                docs = retriever._get_relevant_documents(query, run_manager=callback_manager)
+            else:
+                docs = retriever.get_relevant_documents(query)
+            all_results.extend(docs)
+        return all_results
+
+    async def _aget_relevant_documents(
+        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+    ) -> List[Document]:
+        """
+        Asynchronously get documents relevant to the query.
+        Implements the required async method from BaseRetriever.
+
+        Args:
+            query: Query string
+            run_manager: CallbackManager for the run
+
+        Returns:
+            List of relevant documents
+        """
+        all_results = []
+        for store in self.stores:
+            retriever = store.as_retriever(
+                search_type=self.search_type,
+                search_kwargs=self.search_kwargs,
+                chain_type=self.chain_type
+            )
+            # Pass the run_manager to the sub-retrievers if they support it
+            if hasattr(retriever, "_aget_relevant_documents"):
+                callback_manager = run_manager.get_child()
+                docs = await retriever._aget_relevant_documents(query, run_manager=callback_manager)
+            else:
+                docs = await retriever.aget_relevant_documents(query)
+            all_results.extend(docs)
+        return all_results

{ai_parrot-0.5.8 → ai_parrot-0.5.10}/parrot/stores/abstract.py

@@ -230,11 +230,16 @@ class AbstractStore(ABC):
     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)
+        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,

{ai_parrot-0.5.8 → ai_parrot-0.5.10}/parrot/stores/milvus.py

@@ -251,6 +251,7 @@ class MilvusStore(AbstractStore):
         self,
         collection: Union[str, None] = None,
         metric_type: str = None,
+        index_type: str = None,
         nprobe: int = 32,
         metadata_field: str = None,
         consistency_level: str = 'session'
@@ -261,10 +262,11 @@ class MilvusStore(AbstractStore):
             collection = self.collection_name
         if not metric_type:
             metric_type = self._metric_type or 'COSINE'
+        _idx = index_type or self._index_type
         _search = {
             "search_params": {
                 "metric_type": metric_type,
-                "index_type": self._index_type,
+                "index_type": _idx,
                 "params": {"nprobe": nprobe, "nlist": 1024},
             }
         }

{ai_parrot-0.5.8 → ai_parrot-0.5.10}/parrot/version.py

@@ -3,7 +3,7 @@
 __title__ = "ai-parrot"
 __description__ = "Live Chatbots based on Langchain chatbots and Agents \
     Integrated into Navigator Framework or used into any aiohttp applications."
-__version__ = "0.5.8"
+__version__ = "0.5.10"
 __author__ = "Jesus Lara"
 __author_email__ = "jesuslarag@gmail.com"
 __license__ = "MIT"

ai_parrot-0.5.8/parrot/bots/retrievals/multi.py

@@ -1,50 +0,0 @@
-from typing import List
-from collections.abc import Callable
-from langchain.schema import Document
-
-class MultiVectorStoreRetriever:
-    """
-    This aggregator retriever queries multiple vector stores
-    and merges the results into a single list.
-    """
-
-    def __init__(
-        self,
-        stores: List[Callable],
-        metric_type: str = 'COSINE',
-        chain_type: str = 'stuff',
-        search_type: str = 'similarity',
-        search_kwargs: dict = None,
-    ):
-        """
-        stores: List of AbstractStore or anything that provides an as_retriever() method
-        search_kwargs: dict to pass on to each store's retrieve method (like 'k' for top_k)
-        """
-        self.stores = stores
-        self.search_kwargs = search_kwargs or {}
-        self.chain_type = chain_type
-        self.search_type = search_type
-        self.metric_type = metric_type
-
-    def get_relevant_documents(self, query: str) -> List[Document]:
-        all_results = []
-        for store in self.stores:
-            retriever = store.as_retriever(**self.search_kwargs)
-            docs = retriever.get_relevant_documents(query)
-            all_results.extend(docs)
-        return all_results
-
-    async def aget_relevant_documents(self, query: str) -> List[Document]:
-        # Async version
-        all_results = []
-        for store in self.stores:
-            async with store as st:
-                retriever = st.as_retriever(
-                    search_type=self.search_type,
-                    search_kwargs=self.search_kwargs,
-                    metric_type=self.metric_type,
-                    chain_type=self.chain_type,
-                )
-                docs = await retriever.aget_relevant_documents(query)
-                all_results.extend(docs)
-        return all_results