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

mb_rag/rag/embeddings.py

@@ -61,7 +61,9 @@ from langchain.text_splitter import (
     CharacterTextSplitter,
     RecursiveCharacterTextSplitter,
     SentenceTransformersTokenTextSplitter,
-    TokenTextSplitter)
+    TokenTextSplitter,
+    MarkdownHeaderTextSplitter,
+    SemanticChunker)
 from langchain_community.document_loaders import TextLoader, FireCrawlLoader
 from langchain_chroma import Chroma
 from ..utils.extra import load_env_file
@@ -69,6 +71,8 @@ from langchain.chains import create_history_aware_retriever, create_retrieval_ch
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain_community.document_compressors import FlashrankRerank
 
 load_env_file()
 
@@ -181,13 +185,13 @@ class ModelProvider:
         return GoogleGenerativeAIEmbeddings(**kwargs)
 
     @staticmethod
-    def get_rag_qwen(model_name: str = "qwen", **kwargs):
+    def get_rag_qwen(model_name: str = "Qwen/Qwen3-Embedding-0.6B", **kwargs):
         """
         Load Qwen embedding model. 
         Uses Transformers for embedding generation.
 
         Args:
-            model_name (str): Model identifier (default: "qwen")
+            model_name (str): Model identifier (default: "Qwen/Qwen3-Embedding-0.6B")
             **kwargs: Additional arguments for model initialization
 
         Returns:
@@ -195,7 +199,7 @@ class ModelProvider:
         """
         from langchain.embeddings import HuggingFaceEmbeddings
 
-        return HuggingFaceEmbeddings(model_name="Qwen/Qwen3-Embedding-0.6B", **kwargs)
+        return HuggingFaceEmbeddings(model_name=model_name, **kwargs)
 
 def load_embedding_model(model_name: str = 'openai', model_type: str = "text-embedding-ada-002", **kwargs):
     """
@@ -308,6 +312,14 @@ class TextProcessor:
             'token': TokenTextSplitter(
                 chunk_size=chunk_size,
                 chunk_overlap=chunk_overlap
+            ),
+            'markdown_header': MarkdownHeaderTextSplitter(
+                chunk_size=chunk_size,
+                chunk_overlap=chunk_overlap
+            ),
+            'semantic_chunker': SemanticChunker(
+                chunk_size=chunk_size,
+                chunk_overlap=chunk_overlap
             )
         }
 
@@ -323,6 +335,7 @@ class TextProcessor:
             print(f"Text data splitted into {len(docs)} chunks")
         return docs
 
+
 class embedding_generator:
     """
     Main class for generating embeddings and managing RAG operations.
@@ -374,6 +387,7 @@ class embedding_generator:
         self.vector_store = self.load_vectorstore(**(vector_store_kwargs or {}))
         self.collection_name = collection_name
         self.text_processor = TextProcessor(logger)
+        self.compression_retriever = None
 
     def check_file(self, file_path: str) -> bool:
         """Check if file exists."""
@@ -570,6 +584,49 @@ class embedding_generator:
             retriever = self.retriever
         return retriever.get_relevant_documents(query)
 
+    def load_flashrank_compression_retriever(self, base_retriever=None, model_name: str = "flashrank/flashrank-base", top_n: int = 5):
+        """
+        Load a ContextualCompressionRetriever using FlashrankRerank.
+
+        Args:
+            base_retriever: Existing retriever (if None, uses self.retriever)
+            model_name (str): Flashrank model identifier (default: "flashrank/flashrank-base")
+            top_n (int): Number of top documents to return after reranking
+
+        Returns:
+            ContextualCompressionRetriever: A compression-based retriever using Flashrank
+        """
+        if base_retriever is None:
+            base_retriever = self.retriever
+        if base_retriever is None:
+            raise ValueError("Base retriever is required.")
+
+        compressor = FlashrankRerank(model=model_name, top_n=top_n)
+        self.compression_retriever = ContextualCompressionRetriever(
+            base_compressor=compressor,
+            base_retriever=base_retriever
+        )
+
+        if self.logger:
+            self.logger.info("Loaded Flashrank compression retriever.")
+        return self.compression_retriever
+
+    def compression_invoke(self, query: str):
+        """
+        Invoke compression retriever. Only one compression retriever (Reranker) added right now. 
+
+        Args:
+            query (str): Query string
+
+        Returns:
+            Any: Query results
+        """
+
+        if self.compression_retriever is None:
+            self.compression_retriever = self.load_flashrank_compression_retriever(base_retriever=self.retriever)
+            print("Compression retriever loaded.")
+        return self.compression_retriever.invoke(query)
+
     def generate_rag_chain(self, context_prompt: str = None, retriever=None, llm=None):
         """
         Generate RAG chain for conversation.

mb_rag/version.py

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

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

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

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

@@ -1,18 +1,18 @@
 mb_rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mb_rag/version.py,sha256=9g4JnrnLgsbs9ZJE0iG3ErX8u7puBHMVjLiS08_wP_0,207
+mb_rag/version.py,sha256=5oZ2NODxfaDs2uAEuvKTpFLPnK9Lnk_JL9tUdtfhFAc,207
 mb_rag/chatbot/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 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=CI1tJnIUyGsZhFaqCCZ5xmKKJqdAT1ZAMRReUXLLt2k,28274
+mb_rag/rag/embeddings.py,sha256=uP7dlEtvI7UE7aUdFHdsRax6HaWKMMMdV5LZiG4CIZY,30515
 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.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,,
+mb_rag-1.1.47.dist-info/METADATA,sha256=S2KejHwSkHGmx_UzZG46mCnHjLvbkYq5K6INyqBWAtk,234
+mb_rag-1.1.47.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mb_rag-1.1.47.dist-info/top_level.txt,sha256=FIK1eAa5uYnurgXZquBG-s3PIy-HDTC5yJBW4lTH_pM,7
+mb_rag-1.1.47.dist-info/RECORD,,