src/retrievers/create_retreivers.py

import os
import logging
from typing import List
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from src.utils.asyncHandler import asyncHandler
from src.constants import EMBEDDING_MODEL
from langchain_classic.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever
from langchain_classic.retrievers.contextual_compression import ContextualCompressionRetriever
from langchain_community.document_compressors import FlashrankRerank
from unstructured.partition.pdf import partition_pdf
from unstructured.chunking.title import chunk_by_title
from langchain_core.documents import Document
from src.entity.config_entity import RetreiverConfig

class CompatibleEmbeddings(HuggingFaceEmbeddings):
    def __call__(self, text: str):
        return self.embed_query(text)

embedding_model = CompatibleEmbeddings(model=EMBEDDING_MODEL)

class Retreiver:
    def __init__(self, retreiver_config: RetreiverConfig):
        self.retreiver_config = retreiver_config

    @asyncHandler
    async def partition_document(self, file_path: str):
        logging.info(f"Partitioning document: {file_path}")
        
        elements = partition_pdf(
            filename=file_path,
            strategy=self.retreiver_config.partition_strategy,
            infer_table_structure=True,
            extract_image_block_types=["Image"],
            extract_image_block_to_payload=True
        )
        
        logging.info(f"Extracted {len(elements)} elements")
        return elements
    
    @asyncHandler
    async def create_chunks_by_title(self, elements):
        logging.info("Creating smart chunks...")
        
        chunks = chunk_by_title(
            elements,
            max_characters=self.retreiver_config.max_characters,
            new_after_n_chars=self.retreiver_config.new_after_n_chars,
            combine_text_under_n_chars=self.retreiver_config.combine_text_under_n_chars
        )
        
        if not chunks and elements:
            logging.warning("chunk_by_title returned 0 chunks, falling back to raw elements.")
            chunks = elements
            
        logging.info(f"Created {len(chunks)} chunks")
        return chunks

    @asyncHandler
    async def separate_content_types(self, chunk):
        extracted_text = chunk.text if hasattr(chunk, 'text') and chunk.text is not None else ""
        if not extracted_text and chunk is not None:
            try:
                temp_text = str(chunk)
                if temp_text is not None:
                    extracted_text = temp_text
            except TypeError:
                pass

        content_data = {
            'text': extracted_text,
            'tables': [],
            'images': [],
            'types': ['text']
        }
        
        elements_to_process = []
        if hasattr(chunk, 'metadata') and hasattr(chunk.metadata, 'orig_elements') and chunk.metadata.orig_elements is not None:
            elements_to_process = chunk.metadata.orig_elements
        else:
            elements_to_process = [chunk]
            
        for element in elements_to_process:
            element_type = type(element).__name__
            
            if element_type == 'Table':
                content_data['types'].append('table')
                table_html = getattr(element.metadata, 'text_as_html', element.text) if hasattr(element, 'metadata') else element.text
                content_data['tables'].append(table_html)
            
            elif element_type == 'Image':
                if hasattr(element, 'metadata') and hasattr(element.metadata, 'image_base64'):
                    content_data['types'].append('image')
                    content_data['images'].append(element.metadata.image_base64)
        
        content_data['types'] = list(set(content_data['types']))
        return content_data

    @asyncHandler
    async def get_documents(self, chunks, ingested_file_path: str):
        documents = []
        for chunk in chunks:
            content_data = await self.separate_content_types(chunk)
            doc = Document(
                page_content=content_data['text'],
                metadata={
                    'types': content_data['types'],
                    'tables': content_data['tables'],
                    'images': content_data['images'],
                    'has_images': len(content_data['images']) > 0,
                    'source': ingested_file_path
                }
            )
            documents.append(doc)
        return documents
    
    @asyncHandler
    async def save_to_vector_store(self, documents):
        if not documents:
            logging.warning("No documents provided to save to vector store. Skipping FAISS creation.")
            return None
            
        logging.info(f"Saving {len(documents)} documents to FAISS at {self.retreiver_config.vector_store_path}")
        os.makedirs(os.path.dirname(self.retreiver_config.vector_store_path), exist_ok=True)
        
        vector_store = FAISS.from_documents(documents, embedding_model)
        vector_store.save_local(self.retreiver_config.vector_store_path)
        return self.retreiver_config.vector_store_path
    
  

    @asyncHandler
    async def create_retreiver(self, vectorstore):
        logging.info("Extracting documents from vectorstore for BM25...")
        documents = list(vectorstore.docstore._dict.values())
        base_k = max(self.retreiver_config.k * 2, 20)
        vector_retriever = vectorstore.as_retriever(search_kwargs={"k": base_k})
        valid_documents = [doc for doc in documents if doc.page_content and doc.page_content.strip()]
        if not valid_documents:
            logging.info("No documents with text content found in vectorstore docstore. Returning vector retriever.")
            return vector_retriever
        bm25_retriever = BM25Retriever.from_documents(valid_documents)
        bm25_retriever.k = base_k
        hybrid_retriever = EnsembleRetriever(
            retrievers=[vector_retriever, bm25_retriever],
            weights=self.retreiver_config.ensemble_weights
        )
        return hybrid_retriever

    
    @asyncHandler
    async def get_all_documents(self, vector_store_paths: List[str]):
        documents = []
        for path in vector_store_paths:
            if os.path.exists(path):
                vectorstore = FAISS.load_local(
                    path, 
                    embedding_model, 
                    allow_dangerous_deserialization=True
                )
                for doc in vectorstore.docstore._dict.values():
                    documents.append({
                        "page_content": doc.page_content,
                        "metadata": doc.metadata
                    })
        return documents

    @asyncHandler
    async def merge_vector_stores(self, vector_store_paths: List[str]):
        logging.info(f"Merging {len(vector_store_paths)} vector stores")
        
        individual_retrievers = []
        for path in vector_store_paths:
            if os.path.exists(path):
                vectorstore = FAISS.load_local(
                    path, 
                    embedding_model, 
                    allow_dangerous_deserialization=True
                )
                retriever = await self.create_retreiver(vectorstore)
                individual_retrievers.append(retriever)
        
        if not individual_retrievers:
            logging.warning("No valid vector stores found to merge")
            return None
            
        weights = [1.0 / len(individual_retrievers)] * len(individual_retrievers)
        
        hybrid_retriever = EnsembleRetriever(
            retrievers=individual_retrievers,
            weights=weights
        )


        compressor = FlashrankRerank(top_n=self.retreiver_config.k)

        compression_retriever = ContextualCompressionRetriever(
            base_compressor=compressor,
            base_retriever=hybrid_retriever
        )

        return compression_retriever