mb-rag 1.1.43__py3-none-any.whl → 1.1.45__py3-none-any.whl

mb_rag/chatbot/basic.py

@@ -119,7 +119,7 @@ class ModelFactory:
         return ChatAnthropic(**kwargs)
 
     @classmethod
-    def create_google(cls, model_name: str = "gemini-1.5-flash", **kwargs) -> Any:
+    def create_google(cls, model_name: str = "gemini-2.0-flash", **kwargs) -> Any:
         """
         Create Google chatbot model
         Args:

mb_rag/rag/embeddings.py

@@ -180,6 +180,23 @@ class ModelProvider:
         kwargs["model"] = model_name
         return GoogleGenerativeAIEmbeddings(**kwargs)
 
+    @staticmethod
+    def get_rag_qwen(model_name: str = "qwen", **kwargs):
+        """
+        Load Qwen embedding model. 
+        Uses Transformers for embedding generation.
+
+        Args:
+            model_name (str): Model identifier (default: "qwen")
+            **kwargs: Additional arguments for model initialization
+
+        Returns:
+            QwenEmbeddings: Initialized Qwen embeddings model
+        """
+        from langchain.embeddings import HuggingFaceEmbeddings
+
+        return HuggingFaceEmbeddings(model_name="Qwen/Qwen3-Embedding-0.6B", **kwargs)
+
 def load_embedding_model(model_name: str = 'openai', model_type: str = "text-embedding-ada-002", **kwargs):
     """
     Load a RAG model based on provider and type.
@@ -206,6 +223,8 @@ def load_embedding_model(model_name: str = 'openai', model_type: str = "text-emb
             return ModelProvider.get_rag_google(model_type, **kwargs)
         elif model_name == 'anthropic':
             return ModelProvider.get_rag_anthropic(model_type, **kwargs)
+        elif model_name == 'qwen':
+            return ModelProvider.get_rag_qwen(model_type, **kwargs)
         else:
             raise ValueError(f"Invalid model name: {model_name}")
     except ImportError as e:

mb_rag/utils/all_data_extract.py

@@ -0,0 +1,65 @@
+## Docling data extract
+
+from typing import List
+from mb_rag.utils.extra import check_package
+
+__all__ = ['DocumentExtractor']
+
+class DocumentExtractor:
+    """
+    DocumentExtractor class for extracting data from documents using Docling.
+    """
+    
+    def __init__(self):
+        """
+        Initialize the DocumentExtractor class.
+        Checking for Docling package.
+        """
+        if not check_package("docling"):
+            raise ImportError("Docling package not found. Please install it using: pip install docling")
+        from docling import Docling
+        self.Docling = Docling
+    
+    def _extract_data(self, file_path: str, **kwargs) -> List[str]:
+        """
+        Extract data from a document using Docling.
+        """
+        try:
+            docling = self.Docling(file_path, **kwargs)
+            return docling.extract()
+        except Exception as e:
+            raise Exception(f"Error extracting data from {file_path}: {str(e)}")
+
+    def get_data(self,file_path: str, save_path: str = None, data_store_type: str = "markdown",**kwargs) -> List[str]:
+        """
+        Get data from a document using Docling.
+        Args:
+            file_path (str): Path to the document
+            save_path (str): Path to save the extracted data. Default is None. If None, data saved as Markdown file as docling_{file_name}.md
+            data_store_type (str): Saving document as markdown, txt or html. Default is markdown
+            **kwargs: Additional arguments for Docling
+        Returns:
+            List[str]: Extracted data
+        """
+        data = self._extract_data(file_path, **kwargs)   
+        if data_store_type == "markdown":
+            data_type = "md"
+        elif data_store_type == "txt":
+            data_type = "txt"
+        elif data_store_type == "html":
+            data_type = "html"
+        else:
+            print("Invalid data store type. Defaulting to text (txt)")
+            data_type = "txt"
+        if save_path is None:
+            save_path = f"docling_{file_path.split('/')[-1].split('.')[0]}.{data_type}"
+            print(f"Saving extracted data to {save_path}")
+        if data_store_type == "markdown":
+            data_with_type = data.document.export_to_markdown()
+        elif data_store_type == "txt":
+            data_with_type = data.document.export_to_text()
+        elif data_store_type == "html":
+            data_with_type = data.document.export_to_html()
+        with open(save_path, 'w') as f:
+            f.write(data_with_type)
+        return data

mb_rag/version.py

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

{mb_rag-1.1.43.dist-info → mb_rag-1.1.45.dist-info}/METADATA

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

{mb_rag-1.1.43.dist-info → mb_rag-1.1.45.dist-info}/RECORD

@@ -1,17 +1,18 @@
 mb_rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mb_rag/version.py,sha256=F0q66lXUV7IiL9pvn0fGJZImdHtQZkDF6FEwisuxspA,207
+mb_rag/version.py,sha256=9g4JnrnLgsbs9ZJE0iG3ErX8u7puBHMVjLiS08_wP_0,207
 mb_rag/chatbot/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mb_rag/chatbot/basic.py,sha256=jsoPp0b-JhYRYW15WRv73h4OoZzwT2LTtptEOHnSQJo,23825
+mb_rag/chatbot/basic.py,sha256=8tXU_3Yiqv0J-2Bnpw8p9sQaOlZHzX-Xenjs9GmWqes,23825
 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=H6dDkZi4Ez9NZQrzaYzMlQ66ILLZiOIM-k1LNFfcjSM,27603
+mb_rag/rag/embeddings.py,sha256=CI1tJnIUyGsZhFaqCCZ5xmKKJqdAT1ZAMRReUXLLt2k,28274
 mb_rag/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mb_rag/utils/all_data_extract.py,sha256=TL6O4vNc7mPW-OGK-LhXQQIkSr0o3_7BqNAD-YpTQMU,2532
 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.43.dist-info/METADATA,sha256=xzXHksEUFKJAEm48bYtz1U0uu7evE4lrnoYUX3zVje8,234
-mb_rag-1.1.43.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-mb_rag-1.1.43.dist-info/top_level.txt,sha256=FIK1eAa5uYnurgXZquBG-s3PIy-HDTC5yJBW4lTH_pM,7
-mb_rag-1.1.43.dist-info/RECORD,,
+mb_rag-1.1.45.dist-info/METADATA,sha256=o7mzyY2MJfPaopqUvup1i4ptZeHnMja0jphc4y7jylM,234
+mb_rag-1.1.45.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mb_rag-1.1.45.dist-info/top_level.txt,sha256=FIK1eAa5uYnurgXZquBG-s3PIy-HDTC5yJBW4lTH_pM,7
+mb_rag-1.1.45.dist-info/RECORD,,