ai-parrot 0.1.0__cp311-cp311-manylinux_2_28_x86_64.whl → 0.2.1__cp311-cp311-manylinux_2_28_x86_64.whl

{ai_parrot-0.1.0.dist-info → ai_parrot-0.2.1.dist-info}/METADATA

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

{ai_parrot-0.1.0.dist-info → ai_parrot-0.2.1.dist-info}/RECORD

@@ -1,12 +1,12 @@
 parrot/__init__.py,sha256=oBpxNjpQcJpAHWZU6Hi0tBMGTpTaJEKDXHbfAhqZZAQ,333
-parrot/conf.py,sha256=6kiE38EgnPxtjJc1dde6QCacXZOByne7mWHI8_7_M-E,4150
+parrot/conf.py,sha256=JGCMuB6A_C_z78juMc_-UXCxlHb4rIB_wskSmCubFB0,4128
 parrot/exceptions.cpython-311-x86_64-linux-gnu.so,sha256=p2epCb17B2DCUg3-CZizDQIsIXBhRph07DeJ1Egp8Z0,297896
 parrot/manager.py,sha256=UShLbqfNyPvnGYHzprzQYwu_oj1L-Mhhk2qMRRnsfwo,5187
 parrot/models.py,sha256=GlyWMe2OCgM9dySm2BzK60yKOMHqFth7TYHidTWnioY,12513
 parrot/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-parrot/version.py,sha256=JhX3p58JOZ_NOGA2KeSy5_F1V-BPCf-j5PB_1iDw7yc,373
-parrot/chatbots/__init__.py,sha256=nsu6UPfb6ATv7U9IEHSe2HoGOE9LKLhiFm9gXWjjbMI,201
-parrot/chatbots/abstract.py,sha256=sm4zzapNG_CF_uZrYp91GHMo2SJrOgChX1L-7DHceOo,32968
+parrot/version.py,sha256=ySJiD_gpM6lKn3IqSBYrc1kqg43oubJCIQHgj6rU0Z8,373
+parrot/chatbots/__init__.py,sha256=ypskCnME0xUv6psBEGCEyXCrD0J0ULHSllpVmSxqb4A,200
+parrot/chatbots/abstract.py,sha256=sQsvcBmAnN9QRa2Ls-P72GOJixEoe5Pf6d4K8orM9d8,33278
 parrot/chatbots/asktroc.py,sha256=gyWzyvpAnmXwXd-3DEKoIJtAxt6NnP5mUZdZbkFky8s,604
 parrot/chatbots/base.py,sha256=fXqCzxeA2dU2wYlOV7jCMF4S9_8ohqGMdUfm2VQTfQQ,8054
 parrot/chatbots/basic.py,sha256=DIMTPoGc90BRSlokeOdnjlEXAAfZlIFqxXWaMyAX9uk,232
@@ -16,7 +16,8 @@ parrot/chatbots/copilot.py,sha256=lnlqVXzk4orm-f3qU5iRswwW7WRMV_BISUqw2o0Krzo,47
 parrot/chatbots/dataframe.py,sha256=CfZiLKIwnaku52nl2PNjciqRlH8m2lM4buO6xI7P408,3914
 parrot/chatbots/hragents.py,sha256=PyNIBJ2OH5CtfVydccgpY50V6GI3cLKuVdOMaa7sQz0,574
 parrot/chatbots/oddie.py,sha256=RMbANmJZP1_vLVGKRNBKmA8otyAiWPkvpA0rJ0U3tZk,796
-parrot/chatbots/retrievals/__init__.py,sha256=GOidRGK4O6d0msMC0czIVoscyRvC6YkR8qJGvqqCJec,17404
+parrot/chatbots/odoo.py,sha256=RMbANmJZP1_vLVGKRNBKmA8otyAiWPkvpA0rJ0U3tZk,796
+parrot/chatbots/retrievals/__init__.py,sha256=MfcCFvD4Z5mon7R9ZxGtbAONm36rsRRmYv-C_iwoGe8,19802
 parrot/chatbots/retrievals/constitutional.py,sha256=x7kFR0KpmkR0Wz6bXNOeU3Hesn_EXc6eOSd4TdXxb8o,600
 parrot/crew/__init__.py,sha256=oWLEWNNdYHB2AZUDUZVlSsdMOJiWamvM-jVZ8SzoWQE,77
 parrot/crew/tools/__init__.py,sha256=n2nXEVl8VeWmdI_NyGbiVTAq_S51ebNExuyRM7u9SeI,569
@@ -32,7 +33,7 @@ parrot/crew/tools/search.py,sha256=itWpJD_78IXVtBSKrQIAamCPLbejwmUj1oCBKqBxvO0,1
 parrot/crew/tools/url.py,sha256=m5v9XFQPotHRID7NmY-NQVmXtXEJvfly6tFMslL8U_0,682
 parrot/handlers/__init__.py,sha256=WKn2yg4pSi5-IivkQ34uHSfTG7f1nYj4jB8-4wNG2_A,64
 parrot/handlers/bots.py,sha256=GG2W8y_GC71f8_mMHxwAYP-ope7ecNP_GShL9c66_uY,5782
-parrot/handlers/chat.py,sha256=zgb-vW1SeUCHSwi_ov2m6pepvMvLhy-i8lmSTRtOzY0,5475
+parrot/handlers/chat.py,sha256=iqTChdNqcbhCR2meOqLhPx2XQiSiDzzFP8nRAcPH_7s,5437
 parrot/interfaces/__init__.py,sha256=m3FpxsFeEKq_hvQR9YRGAhZS6u8LZZBuzf6-njugiHE,69
 parrot/interfaces/database.py,sha256=PHy133KTqa9K0tALOQt1q0ufwUZ4n7rndZrhob6Wbk4,609
 parrot/llms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -77,8 +78,8 @@ parrot/loaders/handlers/data.py,sha256=olZ2p-wyUMGoazah7tgHY7V9buGX1FOeJ-cv2vGEo
 parrot/loaders/utils/__init__.py,sha256=SkDyK3MuPGhp0NM6kHvaxQDe97Gcl3n9t5A741OVh1c,28
 parrot/loaders/utils/models.py,sha256=BHcnFzS7oOKUGG-vPy_TlqeFNAC8puRgsIMX2mSgKxY,1059
 parrot/stores/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-parrot/stores/abstract.py,sha256=LdI7taO_gC7YIzFaKlHXY6Kfbd-dpksYCJBYhyd3kLs,5076
-parrot/stores/milvus.py,sha256=V9VkejSNQxhGS4ByiHIJKtKx_QSHdb1j_rBKuLm8oY0,17993
+parrot/stores/abstract.py,sha256=hbRI6wJGCShMlfDc0UkQx7gD9r3LJrBpif6tAKXeDAE,5106
+parrot/stores/milvus.py,sha256=s6BeFDxot9njxUkdB_pPZRjmrgZYa_xuNm--TlG_PjI,18930
 parrot/stores/qdrant.py,sha256=y-jG5sB6ICESyxdpyGBhQ-hgTDLxxPtM5xYZTWIRiAY,4720
 parrot/tools/__init__.py,sha256=lIPK9acoJE2wmxmfbFIQl6dEpUcsBAkapLWvnP_RWG0,669
 parrot/tools/abstract.py,sha256=pVSZw8MDpbVcQ-CHaGwP6CpqXHIs8hH8Oy1AqUuMmrw,1706
@@ -101,8 +102,8 @@ resources/users/handlers.py,sha256=BGzqBvPY_OaIF_nONWX4b_B5OyyBrdGuSihIsdlFwjk,2
 resources/users/models.py,sha256=glk7Emv7QCi6i32xRFDrGc8UwK23_LPg0XUOJoHnwRU,6799
 settings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 settings/settings.py,sha256=9ueEvyLNurUX-AaIeRPV8GKX1c4YjDLbksUAeqEq6Ck,1854
-ai_parrot-0.1.0.dist-info/LICENSE,sha256=vRKOoa7onTsLNvSzJtGtMaNhWWh8B3YAT733Tlu6M4o,1070
-ai_parrot-0.1.0.dist-info/METADATA,sha256=fwcnN1ALMLoPQXvSMDnjAPZ7xIFa2bMeJMgQeEYMZmk,9336
-ai_parrot-0.1.0.dist-info/WHEEL,sha256=czZm5XKxt3lM6KcAi6fwGhndrA0uZKYAkkvqX1vZAww,113
-ai_parrot-0.1.0.dist-info/top_level.txt,sha256=qHoO4BhYDfeTkyKnciZSQtn5FSLN3Q-P5xCTkyvbuxg,26
-ai_parrot-0.1.0.dist-info/RECORD,,
+ai_parrot-0.2.1.dist-info/LICENSE,sha256=vRKOoa7onTsLNvSzJtGtMaNhWWh8B3YAT733Tlu6M4o,1070
+ai_parrot-0.2.1.dist-info/METADATA,sha256=d1pyf6EdFGj4pmmk1wZFIFYKI6roPFZpgpSpf5f2_xw,9336
+ai_parrot-0.2.1.dist-info/WHEEL,sha256=czZm5XKxt3lM6KcAi6fwGhndrA0uZKYAkkvqX1vZAww,113
+ai_parrot-0.2.1.dist-info/top_level.txt,sha256=qHoO4BhYDfeTkyKnciZSQtn5FSLN3Q-P5xCTkyvbuxg,26
+ai_parrot-0.2.1.dist-info/RECORD,,

parrot/chatbots/__init__.py

@@ -3,5 +3,5 @@ from .basic import Chatbot
 from .asktroc import AskTROC
 from .hragents import HRAgent
 from .bose import BoseBot
-from .oddie import OddieBot
+from .odoo import OddieBot
 from .cody import Cody

parrot/chatbots/abstract.py

@@ -125,7 +125,8 @@ from ..conf import (
     MAX_VRAM_AVAILABLE,
     RAM_AVAILABLE,
     default_dsn,
-    REDIS_HISTORY_URL
+    REDIS_HISTORY_URL,
+    EMBEDDING_DEFAULT_MODEL
 )
 from ..interfaces import DBInterface
 from ..models import ChatbotModel
@@ -146,16 +147,17 @@ class AbstractChatbot(ABC, DBInterface):
         "You are {name}, an expert AI assistant and {role} Working at {company}.\n\n"
         "Your primary function is to {goal}\n"
         "Use the provided context of the documents you have processed or extracted from other provided tools or sources to provide informative, detailed and accurate responses.\n"
-        "I am here to help with {role}, {backstory}.\n\n"
+        "I am here to help with {role}.\n"
+        "**Backstory:**\n"
+        "{backstory}.\n\n"
         "Focus on answering the question directly but detailed. Do not include an introduction or greeting in your response.\n\n"
         "{company_information}\n\n"
+        "Here is a brief summary of relevant information:\n"
         "Context: {context}\n\n"
         "Given this information, please provide answers to the following question adding detailed and useful insights:\n\n"
-        "Chat History: {chat_history}\n"
-        "Human: {question}\n"
-        "Here is a brief summary of relevant information:\n"
-        "{summaries}\n\n"
-        "Assistant Answer:\n"
+        "**Chat History:** {chat_history}\n\n"
+        "**Human Question:** {question}\n"
+        "Assistant Answer:\n\n"
         "{rationale}\n"
         "You are a fluent speaker, you can talk and respond fluently in English and Spanish, and you must answer in the same language as the user's question. If the user's language is not English, you should translate your response into their language.\n"
     )
@@ -536,7 +538,7 @@ class AbstractChatbot(ABC, DBInterface):
         models = file_config.get('models', {})
         if not self.embedding_model_name:
             self.embedding_model_name = models.get(
-                'embedding_name', None,
+                'embedding', EMBEDDING_DEFAULT_MODEL
             )
         if not self.tokenizer_model_name:
             self.tokenizer_model_name = models.get('tokenizer')
@@ -624,8 +626,12 @@ class AbstractChatbot(ABC, DBInterface):
                 )
         elif vector_db == 'MilvusStore':
             if MILVUS_ENABLED is True:
+                print('EMBEDDINGS > ', self.embeddings)
+                print('AQUI MILVUS >> ', embed)
+                print('AQUI MILVUS >> ', self.embedding_model_name)
                 self._store = MilvusStore(
                     embeddings=embed,
+                    embedding_name=self.embedding_model_name,
                     use_bge=self.use_bge,
                     use_fastembed=self.use_fastembed,
                     **config

parrot/chatbots/odoo.py

@@ -0,0 +1,17 @@
+from .basic import Chatbot
+
+
+class OddieBot(Chatbot):
+    """Represents an agent in Navigator.
+
+        Each agent has a name, a role, a goal, a backstory,
+        and an optional language model (llm).
+    """
+    name: str = 'Oddie'
+    company: str = 'T-ROC Global'
+    company_website: str = 'https://www.trocglobal.com'
+    contact_information = 'communications@trocglobal.com'
+    contact_form = 'https://www.surveymonkey.com/r/TROC_Suggestion_Box'
+    role = "Odoo and ERP Specialist and Odoo Programmer"
+    goal = "To provide information and support on Odoo and ERP systems, help with troubleshooting, and answer any questions you may have about any Odoo and ERP systems implementation."
+    specialty_area = 'Bring useful information about Odoo ERP, documentation, usage, samples, etc.'

parrot/chatbots/retrievals/__init__.py

@@ -108,12 +108,14 @@ class RetrievalManager:
     def create_memory(
         self,
         key: str = 'chat_history',
-        input_key: str = 'question'
+        input_key: str = 'question',
+        output_key: str = 'answer'
     ):
         return ConversationBufferMemory(
             memory_key=key,
             return_messages=True,
-            input_key=input_key
+            input_key=input_key,
+            output_key=output_key
         )
 
     ### Different types of Retrieval
@@ -122,7 +124,7 @@ class RetrievalManager:
         question: str = None,
         chain_type: str = 'stuff',
         search_type: str = 'similarity',
-        search_kwargs: dict = {"k": 10, "fetch_k": 30, "lambda_mult": 0.86},
+        search_kwargs: dict = {"k": 5, "fetch_k": 20, "lambda_mult": 0.89},
         return_docs: bool = True,
         metric_type: str = None,
         memory: Any = None,
@@ -137,18 +139,27 @@ class RetrievalManager:
         vector = self.store.get_vector(
             metric_type=metric_type
         )
-        retriever = VectorStoreRetriever(
+        simil_retriever = VectorStoreRetriever(
             vectorstore=vector,
-            search_type=search_type,
+            search_type='similarity',
             chain_type=chain_type,
             search_kwargs=search_kwargs
         )
+        retriever = vector.as_retriever(
+            search_type=search_type,
+            search_kwargs=search_kwargs
+        )
         if self.kb:
             # Get a BM25 Retriever:
             b25_retriever = BM25Retriever.from_documents(self.kb)
             retriever = EnsembleRetriever(
-                retrievers=[retriever, b25_retriever],
-                weights=[0.9, 0.1]
+                retrievers=[simil_retriever, retriever, b25_retriever],
+                weights=[0.6, 0.3, 0.1]
+            )
+        else:
+            retriever = EnsembleRetriever(
+                retrievers=[simil_retriever, retriever],
+                weights=[0.7, 0.3]
             )
         custom_template = self.template.format_map(
             SafeDict(
@@ -209,7 +220,7 @@ class RetrievalManager:
         question: str = None,
         chain_type: str = 'stuff',
         search_type: str = 'mmr',
-        search_kwargs: dict = {"k": 10, "fetch_k": 20, "lambda_mult": 0.86},
+        search_kwargs: dict = {"k": 5, "fetch_k": 40, "lambda_mult": 0.85},
         return_docs: bool = True,
         metric_type: str = None,
         use_llm: str = None
@@ -235,14 +246,14 @@ class RetrievalManager:
             b25_retriever = BM25Retriever.from_documents(self.kb)
             retriever = EnsembleRetriever(
                 retrievers=[simil_retriever, retriever, b25_retriever],
-                weights=[0.3, 0.6, 0.1]
+                weights=[0.6, 0.3, 0.1]
             )
         else:
-            # retriever = EnsembleRetriever(
-            #     retrievers=[simil_retriever, retriever],
-            #     weights=[0.7, 0.3]
-            # )
-            retriever = simil_retriever
+            retriever = EnsembleRetriever(
+                retrievers=[simil_retriever, retriever],
+                weights=[0.7, 0.3]
+            )
+            # retriever = simil_retriever
         custom_template = self.template.format_map(
             SafeDict(
                 chat_history=''
@@ -284,13 +295,15 @@ class RetrievalManager:
             chain_type=chain_type,
             retriever=retriever,
             return_source_documents=return_docs,
+            verbose=True,
             chain_type_kwargs={
                 "prompt": PromptTemplate(
                     template=custom_template,
-                    input_variables=['summaries', 'question']
+                    input_variables=['context', 'question']
                 )
             },
         )
+        # Debug Code ::
         # print('=====================')
         # print(custom_template)
         # response = self.chain.invoke(question)
@@ -311,6 +324,7 @@ class RetrievalManager:
         #     distance = 'EMPTY'
         # print('DISTANCE > ', distance)
         # print('CHAIN > ', self.chain)
+
         return self
 
     def get_current_context(self):
@@ -328,7 +342,7 @@ class RetrievalManager:
             count = 0
             d = {}
             for source in source_documents:
-                if count >= 10:
+                if count >= 20:
                     break  # Exit loop after processing 10 documents
                 metadata = source.metadata
                 if 'url' in metadata:
@@ -433,42 +447,95 @@ class RetrievalManager:
                 )
 
 
-    def question(self, question):
-        # vector = self.store.get_vector()
-        # docs = vector.similarity_search(self._question, k=1)
-        # max_confidence_score = 0
-        # for doc_id, similarity_score in docs:
-        #     max_confidence_score = max(max_confidence_score, similarity_score)
-        try:
-            # Invoke the chain with the given question
-            response = self.chain.invoke(question)
-        except Exception as exc:
-            self.logger.error(
-                f"Error invoking chain: {exc}"
+    async def question(
+            self,
+            question: str = None,
+            chain_type: str = 'stuff',
+            search_type: str = 'similarity',
+            search_kwargs: dict = {"k": 5, "fetch_k": 20, "lambda_mult": 0.89},
+            return_docs: bool = True,
+            metric_type: str = None,
+            memory: Any = None,
+            **kwargs
+    ):
+        # Generating Vector:
+        # print('COLLECTION > ', self.store.collection, self.store.embedding_name)
+        # print('DB > ', self.store.database)
+        async with self.store as store:  #pylint: disable=E1101
+            vector = store.get_vector(metric_type=metric_type)
+            # docs = vector.similarity_search(question, k=1)
+            # print('DOCS > ', docs)
+            retriever = VectorStoreRetriever(
+                vectorstore=vector,
+                search_type=search_type,
+                chain_type=chain_type,
+                search_kwargs=search_kwargs
             )
-            return None, None
-        try:
-            qa_response = ChatResponse(**response)
-        except (ValueError, TypeError) as exc:
-            self.logger.error(
-                f"Error validating response: {exc}"
+            custom_template = self.template.format_map(
+                SafeDict(
+                    summaries=''
+                )
             )
-            return None, None
-        except ValidationError as exc:
-            self.logger.error(
-                f"Error on response: {exc.payload}"
+            _prompt = PromptTemplate(
+                template=custom_template,
+                input_variables=["context", "chat_history", "question"],
             )
-            return None, None
+            try:
+                chain = ConversationalRetrievalChain.from_llm(
+                    llm=self.model,
+                    retriever=retriever,
+                    chain_type=chain_type,
+                    verbose=False,
+                    memory=memory,
+                    return_source_documents=return_docs,
+                    return_generated_question=True,
+                    combine_docs_chain_kwargs={"prompt": _prompt},
+                    **kwargs
+                )
+                response = chain.invoke(
+                    question
+                )
+            except Exception as exc:
+                self.logger.error(
+                    f"Error invoking chain: {exc}"
+                )
+                return None, None
+            try:
+                qa_response = ChatResponse(**response)
+            except (ValueError, TypeError) as exc:
+                self.logger.error(
+                    f"Error validating response: {exc}"
+                )
+                return None, None
+            except ValidationError as exc:
+                self.logger.error(
+                    f"Error on response: {exc.payload}"
+                )
+                return None, None
         try:
-            return self.as_markdown(
+            qa_response.response = self.as_markdown(
                 qa_response
-            ), qa_response
+            )
+            # # saving question to Usage Log
+            # if self.request:
+            #     tasker = self.request.app['service_queue']
+            #     await tasker.put(
+            #         self.log_usage,
+            #         response=qa_response,
+            #         request=self.request
+            #     )
+            # else:
+            #     asyncio.create_task(
+            #         self.log_usage(response=qa_response)
+            #     )
+            return qa_response
         except Exception as exc:
             self.logger.exception(
                 f"Error on response: {exc}"
             )
             return None, None
 
+
     async def invoke(self, question):
         # Invoke the chain with the given question
         try:
@@ -481,7 +548,6 @@ class RetrievalManager:
             )
             return None, None
         try:
-            print('RESPONSE > ', response)
             qa_response = ChatResponse(**response)
         except (ValueError, TypeError) as exc:
             self.logger.error(

parrot/conf.py

@@ -85,7 +85,7 @@ QDRANT_CONN_TYPE = config.get('QDRANT_CONN_TYPE', fallback='server')
 EMBEDDING_DEVICE = config.get('EMBEDDING_DEVICE', fallback='cpu')
 EMBEDDING_DEFAULT_MODEL = config.get(
     'EMBEDDING_DEFAULT_MODEL',
-    fallback='sentence-transformers/sentence-t5-large'
+    fallback='thenlper/gte-base'
 )
 MAX_VRAM_AVAILABLE = config.get('MAX_VRAM_AVAILABLE', fallback=20000)
 RAM_AVAILABLE = config.get('RAM_AVAILABLE', fallback=819200)

parrot/handlers/chat.py

@@ -48,7 +48,6 @@ class ChatHandler(BaseView):
         post.
         description: Post method for ChatHandler.
         """
-        print('ARGS > ', args, kwargs)
         app = self.request.app
         name = self.request.match_info.get('chatbot_name', None)
         qs = self.query_parameters(self.request)
@@ -92,19 +91,20 @@ class ChatHandler(BaseView):
         session = self.request.session
         try:
             session_id = session.get('session_id', None)
-            # print('SESSION ID > ', session_id)
             memory_key = f'{session.session_id}_{name}_message_store'
-            # print('MEM STORAGE > ', memory_key)
             memory = chatbot.get_memory(session_id=memory_key)
-            # print('MEMORY >> ', memory)
-            async with chatbot.get_retrieval(request=self.request) as retrieval:
-                qa = retrieval.conversation(
+            with chatbot.get_retrieval(request=self.request) as retrieval:
+                result = await retrieval.question(
                     question=question,
-                    search_kwargs={"k": 10},
-                    use_llm=llm,
                     memory=memory
                 )
-                result = await qa.invoke(question)
+                # qa = retrieval.qa(
+                #     question=question,
+                #     search_kwargs={"k": 10},
+                #     use_llm=llm,
+                #     # memory=memory
+                # )
+                # result = await qa.invoke(question)
                 # Drop "memory" information:
                 result.chat_history = None
                 result.source_documents = None

parrot/stores/abstract.py

@@ -32,7 +32,7 @@ class AbstractStore(ABC):
         self._connected: bool = False
         self.use_bge: bool = kwargs.pop("use_bge", False)
         self.fastembed: bool = kwargs.pop("use_fastembed", False)
-        self.embedding_name: str = EMBEDDING_DEFAULT_MODEL
+        self.embedding_name: str = kwargs.pop('embedding_name', EMBEDDING_DEFAULT_MODEL)
         self.dimension: int = kwargs.pop("dimension", 768)
         self._metric_type: str = kwargs.pop("metric_type", 'COSINE')
         self._index_type: str = kwargs.pop("index_type", 'IVF_FLAT')

parrot/stores/milvus.py

@@ -106,7 +106,30 @@ class MilvusStore(AbstractStore):
             self.kwargs['db_name'] = self.database
             self.use_database(self.database, create=True)
 
-    def connect(self, client_id: str = None):
+    async def __aenter__(self):
+        try:
+            self.tensor = torch.randn(1000, 1000).cuda()
+        except RuntimeError:
+            self.tensor = None
+        if self._embed_ is None:
+            self._embed_ = self.create_embedding(
+                model_name=self.embedding_name
+            )
+        self._client_id = str(uuid.uuid4())
+        self._client, self.client_id = self.connect(client_id=self._client_id)
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        # closing Embedding
+        self._embed_ = None
+        del self.tensor
+        try:
+            self.close(client_id=self.client_id)
+            torch.cuda.empty_cache()
+        except RuntimeError:
+            pass
+
+    def connect(self, client_id: str = None) -> tuple:
         # 1. Set up a pyMilvus default connection
         # Unique connection:
         if not client_id:
@@ -118,6 +141,7 @@ class MilvusStore(AbstractStore):
         self._client = MilvusClient(
             **self.kwargs
         )
+        print('CONN > ', self._client, client_id, connections)
         self._connected = True
         return self._client, client_id
 
@@ -133,7 +157,6 @@ class MilvusStore(AbstractStore):
             "port": self.port,
             **kwargs
         }
-        print('ARGS >', args)
         try:
             conn = connections.connect(**args)
             db.create_database(db_name)
@@ -452,8 +475,9 @@ class MilvusStore(AbstractStore):
         self,
         collection: str = None,
         metric_type: str = None,
-        nprobe: int = 200,
-        metadata_field: str = None
+        nprobe: int = 32,
+        metadata_field: str = None,
+        consistency_level: str = 'session'
     ) -> Milvus:
         if not metric_type:
             metric_type = self._metric_type
@@ -462,7 +486,7 @@ class MilvusStore(AbstractStore):
         _search = {
             "search_params": {
                 "metric_type": metric_type,
-                "params": {"nprobe": nprobe},
+                "params": {"nprobe": nprobe, "nlist": 1024},
             }
         }
         if metadata_field:
@@ -474,7 +498,7 @@ class MilvusStore(AbstractStore):
         return Milvus(
             embedding_function=_embed_,
             collection_name=collection,
-            consistency_level='Bounded',
+            consistency_level=consistency_level,
             connection_args={
                 **self.kwargs
             },
@@ -484,7 +508,13 @@ class MilvusStore(AbstractStore):
             **_search
         )
 
-    def similarity_search(self, query: str, collection: str = None, limit: int = 2) -> list:
+    def similarity_search(
+        self,
+        query: str,
+        collection: str = None,
+        limit: int = 2,
+        consistency_level: str = 'Bounded'
+    ) -> list:
         if collection is None:
             collection = self.collection
         if self._embed_ is None:
@@ -496,7 +526,7 @@ class MilvusStore(AbstractStore):
         vector_db = Milvus(
             embedding_function=_embed_,
             collection_name=collection,
-            consistency_level='Bounded',
+            consistency_level=consistency_level,
             connection_args={
                 **self.kwargs
             },

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 aiohttp applications."
-__version__ = "0.1.0"
+__version__ = "0.2.1"
 __author__ = "Jesus Lara"
 __author_email__ = "jesuslarag@gmail.com"
 __license__ = "MIT"