mb-rag 1.1.37__py3-none-any.whl → 1.1.39__py3-none-any.whl

mb_rag/rag/embeddings.py

@@ -325,7 +325,8 @@ class embedding_generator:
         # Initialize generator
         gen = embedding_generator(
             model="openai",
-            model_type="text-embedding-3-small"
+            model_type="text-embedding-3-small",
+            collection_name='test'
         )
 
         # Generate embeddings
@@ -335,7 +336,7 @@ class embedding_generator:
         )
 
         # Load retriever
-        retriever = gen.load_retriever('./embeddings')
+        retriever = gen.load_retriever('./embeddings', collection_name='test')
 
         # Query embeddings
         results = gen.query_embeddings("What is this about?")
@@ -430,12 +431,13 @@ class embedding_generator:
         else:
             return "Vector store not found"
 
-    def load_embeddings(self, embeddings_folder_path: str):
+    def load_embeddings(self, embeddings_folder_path: str,collection_name: str = 'test'):
         """
         Load embeddings from folder.
 
         Args:
             embeddings_folder_path (str): Path to embeddings folder
+            collection_name (str): Name of the collection. Default: 'test'
 
         Returns:
             Optional[Chroma]: Loaded vector store or None if not found
@@ -443,7 +445,8 @@ class embedding_generator:
         if self.check_file(embeddings_folder_path):
             if self.vector_store_type == 'chroma':
                 return Chroma(persist_directory=embeddings_folder_path,
-                            embedding_function=self.model)
+                            embedding_function=self.model,
+                            collection_name=collection_name)
         else:
             if self.logger:
                 self.logger.info("Embeddings file not found")
@@ -451,7 +454,8 @@ class embedding_generator:
 
     def load_retriever(self, embeddings_folder_path: str,
                       search_type: List[str] = ["similarity_score_threshold"],
-                      search_params: List[Dict] = [{"k": 3, "score_threshold": 0.9}]):
+                      search_params: List[Dict] = [{"k": 3, "score_threshold": 0.9}],
+                      collection_name: str = 'test'):
         """
         Load retriever with search configuration.
 
@@ -459,6 +463,7 @@ class embedding_generator:
             embeddings_folder_path (str): Path to embeddings folder
             search_type (List[str]): List of search types
             search_params (List[Dict]): List of search parameters
+            collection_name (str): Name of the collection. Default: 'test'
 
         Returns:
             Union[Any, List[Any]]: Single retriever or list of retrievers
@@ -472,7 +477,7 @@ class embedding_generator:
             )
             ```
         """
-        db = self.load_embeddings(embeddings_folder_path)
+        db = self.load_embeddings(embeddings_folder_path, collection_name)
         if db is not None:
             if self.vector_store_type == 'chroma':
                 if len(search_type) != len(search_params):
@@ -508,7 +513,7 @@ class embedding_generator:
             chunk_overlap (int): Overlap between chunks
         """
         if self.vector_store_type == 'chroma':
-            db = self.load_embeddings(embeddings_folder_path)
+            db = self.load_embeddings(embeddings_folder_path, collection_name)
             if db is not None:
                 docs = self.tokenize(data, text_splitter_type, chunk_size, chunk_overlap)
                 db.add_documents(docs)

mb_rag/version.py

@@ -1,5 +1,5 @@
 MAJOR_VERSION = 1
 MINOR_VERSION = 1
-PATCH_VERSION = 37
+PATCH_VERSION = 39
 version = '{}.{}.{}'.format(MAJOR_VERSION, MINOR_VERSION, PATCH_VERSION)
 __all__  = ['MAJOR_VERSION', 'MINOR_VERSION', 'PATCH_VERSION', 'version']

{mb_rag-1.1.37.dist-info → mb_rag-1.1.39.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mb_rag
-Version: 1.1.37
+Version: 1.1.39
 Summary: RAG function file
 Author: ['Malav Bateriwala']
 Requires-Python: >=3.8

{mb_rag-1.1.37.dist-info → mb_rag-1.1.39.dist-info}/RECORD

@@ -1,17 +1,17 @@
 mb_rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mb_rag/version.py,sha256=KNBbKy-E6H597p6u0q5gyRGkgk6uZAjkIt_B5xgV65I,207
+mb_rag/version.py,sha256=HcQVlfDNzK6DIr0RRqDVdD0Tb_es_b1S3Ta7qd1VAWE,207
 mb_rag/chatbot/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 mb_rag/chatbot/basic.py,sha256=wZ9rZvLtNEZPhaAFkOv0yaPCYiki9goNr77Pyqwfn5g,23139
 mb_rag/chatbot/chains.py,sha256=vDbLX5R29sWN1pcFqJ5fyxJEgMCM81JAikunAEvMC9A,7223
 mb_rag/chatbot/prompts.py,sha256=n1PyiLbU-5fkslRv6aVOzt0dDlwya_cEdQ7kRnRhMuY,1749
 mb_rag/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mb_rag/rag/embeddings.py,sha256=eutlhc8Pm3zfXu3NTSGEm7A_NpvyxE99Ocn9Vw_IOrM,27109
+mb_rag/rag/embeddings.py,sha256=MhxG0n7Q0XwPGYIRjCpbSZTlZLyj6Shxo64n347BEWA,27497
 mb_rag/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 mb_rag/utils/bounding_box.py,sha256=G0hdDam8QmYtD9lfwMeDHGm-TTo6KZg-yK5ESFL9zaM,8366
 mb_rag/utils/document_extract.py,sha256=vZiFB1RYm1BIEaNA0MveJ5Zp-KEi0ngKjW8xEdtPqXA,12558
 mb_rag/utils/extra.py,sha256=spbFrGgdruNyYQ5PzgvpSIa6Nm0rn9bb4qc8W9g582o,2492
 mb_rag/utils/pdf_extract.py,sha256=cVeMyhnAU4XZxjIZHKMYhrktTjUNOjhx2r_LZKReOZE,15598
-mb_rag-1.1.37.dist-info/METADATA,sha256=mQoHe8LBmDwRVMcgGh-3B3SeTcR7dJm-T8uAq-5jlsw,234
-mb_rag-1.1.37.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-mb_rag-1.1.37.dist-info/top_level.txt,sha256=FIK1eAa5uYnurgXZquBG-s3PIy-HDTC5yJBW4lTH_pM,7
-mb_rag-1.1.37.dist-info/RECORD,,
+mb_rag-1.1.39.dist-info/METADATA,sha256=MlnqdDCIR0P88B8EiNDnjQQCrI4BMT-5pm1Dlu9B3Og,234
+mb_rag-1.1.39.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mb_rag-1.1.39.dist-info/top_level.txt,sha256=FIK1eAa5uYnurgXZquBG-s3PIy-HDTC5yJBW4lTH_pM,7
+mb_rag-1.1.39.dist-info/RECORD,,