Refactor document processing and retrieval workflow; add utility functions for markdown conversion and logging

app.py

@@ -4,8 +4,15 @@ from src.indexing.document_processing import DocumentProcessor
 from src.indexing.vectore_store import VectorStoreManager
 from src.tools_retrieval.retriever import RetrieverManager
 from src.workflow import RAGWorkflow
+from src.utils import (
+    logger,
+    convert_document_to_markdown,
+    save_to_markdown,
+    determine_top_k,
+    determine_reranking_top_n
+)
 
-
+ 
 UPLOAD_FOLDER = "uploads/"
 PERSIST_DIRECTORY = "./chroma_db"
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
@@ -29,28 +36,34 @@ st.set_page_config(
 st.title("Agentic RAG Chatbot")
 
 with st.sidebar:
-    st.header("PDF Upload")
-    uploaded_file = st.file_uploader("Upload your PDF", type=["pdf"])
-    st.info("Supported file type: PDF")
-    process_button = st.button("Process PDF")
+    st.header("Upload")
+    uploaded_file = st.file_uploader("Upload Document", type=["pdf", "xlsx", "docx", "txt"])
+    process_button = st.button("Process Document")
 
     if uploaded_file and process_button:    
-        with st.spinner("Processing PDF..."):  
+        with st.spinner("Processing Document..."):  
             file_path = os.path.join(UPLOAD_FOLDER, uploaded_file.name)
             with open(file_path, "wb") as f:
                 f.write(uploaded_file.getbuffer())
-
+            convert_to_md = convert_document_to_markdown(file_path)
+            file_path_md = save_to_markdown(convert_to_md, file_path)
             doc_processor = DocumentProcessor()
-            chunks = doc_processor.load_and_split_pdf(file_path)
+            chunks = doc_processor.load_and_split_pdf(file_path_md)
 
             vector_store_manager = VectorStoreManager()
-            vector_store = vector_store_manager.index_documents(chunks, uploaded_file.name, PERSIST_DIRECTORY)
+            vector_store = vector_store_manager.index_documents(chunks)
                     
             st.session_state.vector_store = vector_store
-            st.success("PDF processed and indexed successfully!")
+            st.success("Document processed and indexed successfully!")
+            top_k = determine_top_k(len(chunks))
+            top_n = determine_reranking_top_n(top_k)
                     
             retriever_manager = RetrieverManager(vector_store)
-            retriever_tool = retriever_manager.create_retriever(chunks)
+            retriever_tool = retriever_manager.create_retriever(
+                documents=chunks,
+                top_n=top_n,
+                k=top_k
+            )
             st.session_state.retriever = retriever_tool
             st.success("Retriever tool created successfully!")
             rag_workflow = RAGWorkflow(retriever_tool)

requirements.txt

@@ -1,12 +1,12 @@
-langchain
-langgraph
-langchain-huggingface
-langchain-groq
-langchain-community
-scikit-learn
-langchain-chroma
-pypdf==5.1.0
-tiktoken  
-rank_bm25
-fastembed
-flashrank
+markitdown[all]==0.1.1
+langchain==0.3.24 
+langchain-groq==0.3.2 
+langchain-community==0.3.23 
+langgraph==0.4.3 
+scikit-learn==1.6.1 
+tiktoken==0.9.0 
+rank_bm25==0.2.2 
+fastembed==0.6.1 
+ flashrank==0.2.10 
+langchain-unstructured==0.1.6 
+unstructured==0.17.2 

src/indexing/document_processing.py

@@ -1,5 +1,6 @@
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.document_loaders import PyPDFLoader
+from langchain_community.document_loaders import UnstructuredMarkdownLoader
+from src.utils import logger
 
 class DocumentProcessor:
     def __init__(self, chunk_size=500, chunk_overlap=100):
@@ -10,7 +11,9 @@ class DocumentProcessor:
 
     def load_and_split_pdf(self, file_path: str):
         """Load PDF and split into chunks"""
-        loader = PyPDFLoader(file_path)
+        logger.info(f"Loading and splitting PDF: {file_path}")
+        loader = UnstructuredMarkdownLoader(file_path)
         docs = loader.load()
         chunks = self.text_splitter.split_documents(docs)
+        logger.info(f"Loaded and split PDF into {len(chunks)} chunks")
         return chunks

src/indexing/vectore_store.py

@@ -1,27 +1,21 @@
-from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.vectorstores import SKLearnVectorStore
-from langchain_chroma import Chroma
+from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
 
 class VectorStoreManager:
-    def __init__(self, embedding_model="intfloat/multilingual-e5-small"):
-        self.embeddings = HuggingFaceEmbeddings(model_name=embedding_model)
+    def __init__(self, embedding_model="BAAI/bge-base-en-v1.5"):
+        self.embeddings = FastEmbedEmbeddings(model_name=embedding_model)
         
-    def create_vector_store(self, collection_name, presist_directory):
+    def create_vector_store(self):
         """Create a new vector store"""
-#        vector_store = SKLearnVectorStore.from_documents(
-#            documents=documents,
-#            embedding=self.embeddings,
-#        )
-        vector_store = Chroma(
-            collection_name=collection_name,
-            embedding_function=self.embeddings,
-            persist_directory=presist_directory,  # Where to save data locally, remove if not necessary
+        vector_store = SKLearnVectorStore.from_documents(
+            metric="cosine",
+            embedding=self.embeddings,
         )
-        
+ 
         return vector_store
     
-    def index_documents(self, documents, collection_name, presist_directory):
+    def index_documents(self, documents):
         """Index documents into vector store"""
-        vector_store = self.create_vector_store(collection_name, presist_directory)
+        vector_store = self.create_vector_store()
         vector_store.add_documents(documents=documents)
         return vector_store 

src/tools_retrieval/retriever.py

@@ -10,25 +10,31 @@ class RetrieverManager:
     def __init__(self, vector_store):
         self.vector_store = vector_store
         
-    def create_base_retriever(self, search_type="similarity", k=3):
+    def create_base_retriever(self, search_type: str ="similarity", k: int = 3):
         """Create basic vector store retriever"""
         return self.vector_store.as_retriever(
             search_type=search_type, 
             search_kwargs={"k": k}
         )
     
-    def create_ensemble_retriever(self, texts, vector_weight=0.5, keyword_weight=0.5):
+    def create_ensemble_retriever(
+            self, 
+            texts, 
+            k: int = 3, 
+            vector_weight: float = 0.5, 
+            keyword_weight: float =0.5
+        ):
         """Create ensemble retriever combining vector and keyword search"""
-        vector_retriever = self.create_base_retriever()
+        vector_retriever = self.create_base_retriever(k=k)
         keyword_retriever = BM25Retriever.from_documents(texts)
-        keyword_retriever.k = 3
+        keyword_retriever.k = k
         
         return EnsembleRetriever(
             retrievers=[vector_retriever, keyword_retriever],
             weights=[vector_weight, keyword_weight]
         )
     
-    def create_compression_retriever(self, base_retriever, top_n=5):
+    def create_compression_retriever(self, base_retriever, top_n: int):
         """Create compression retriever with reranking"""
         compressor = FlashrankRerank(top_n=top_n)
         return ContextualCompressionRetriever(
@@ -36,9 +42,9 @@ class RetrieverManager:
             base_retriever=base_retriever
         )
 
-    def create_retriever(self, documents):
-        base_retriever = self.create_ensemble_retriever(documents)
-        compression_retriever = self.create_compression_retriever(base_retriever=base_retriever)
+    def create_retriever(self, documents, top_n: int, k: int = 3, ):
+        base_retriever = self.create_ensemble_retriever(texts=documents, k=k)
+        compression_retriever = self.create_compression_retriever(base_retriever=base_retriever, top_n=top_n)
         return create_retriever_tool(
             compression_retriever,
             "retrieve_docs",

src/utils.py

@@ -0,0 +1,95 @@
+import os
+import logging
+from typing import Any
+from pathlib import Path
+from markitdown import MarkItDown
+
+def setup_logging():
+    """Sets up the logging configuration."""
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(levelname)s - %(message)s',
+        handlers=[
+            logging.FileHandler("app.log", encoding="utf-8"),
+            logging.StreamHandler()
+        ]
+    )
+    return logging.getLogger(__name__)
+
+logger = setup_logging()
+
+def extract_filename(filepath: Path) -> str:
+  """Extracts the filename without extension.
+
+  Args:
+    filepath: The complete path to the file.
+
+  Returns:
+    The filename without extension.
+  """
+  logger.info(f"Extracting filename from {filepath}")
+  return os.path.splitext(os.path.basename(filepath))[0]  # More concise way to get filename
+
+def convert_document_to_markdown(filepath: Path) -> str:
+    """Converts a document to markdown.
+
+    Args:
+        filepath: The path to the document file.
+
+    Returns:
+        The raw markdown content.
+    """
+    logger.info(f"Converting document to markdown: {filepath}")
+    md = MarkItDown(enable_plugins=False)  # Set to True to enable plugins if needed
+    result = md.convert(filepath)
+    return result.markdown
+
+def save_to_markdown(text: Any, path: Path) -> str:
+    """Saves text content to a markdown file.
+
+    Args:
+        text: The text or markdown content to save.
+        path: The complete path to the markdown file.
+
+    Returns:
+        The path to the saved markdown file as a string.
+    """
+
+    filename = extract_filename(path)
+    filepath = f'{filename}.md'  # Create the full filepath
+    with open(filepath, 'w', encoding='utf-8') as f:
+        f.write(text)
+    logger.info(f"Markdown file saved successfully at {filepath}")
+    return filepath  # Return the filepath
+
+def determine_top_k(num_chunks: int) -> int:
+    """Determines the top_k value based on the number of chunks.
+
+    Args:
+        num_chunks: The total number of chunks.
+
+    Returns:
+        The appropriate top_k value.
+    """
+    if num_chunks <= 5:
+        top_k = num_chunks
+    else:
+        top_k = 5
+    logger.info(f"Determined top_k: {top_k} based on num_chunks: {num_chunks}")
+    return top_k
+
+def determine_reranking_top_n(top_k: int) -> int:
+    """Determines the top_n value for reranking based on top_k.
+    Args:
+        top_k: The number of top results to consider.
+    Returns:
+        The appropriate top_n value for reranking.
+    """
+    total_top_k = top_k * 2
+
+    if total_top_k <= 5:
+        top_n = round(total_top_k / 2) + 1
+    else:
+        top_n = 6
+    logger.info(f"Determined top_n: {top_n} based on top_k: {top_k}")
+    return top_n