retriever/document_manager.py

import logging
import os
from typing import Any, Dict, List
import uuid
from data.document_loader import DocumentLoader
from data.pdf_reader import PDFReader
from retriever.chunk_documents import chunk_documents
from retriever.vector_store_manager import VectorStoreManager

class DocumentManager:
    def __init__(self):
        self.doc_loader = DocumentLoader()
        self.pdf_reader = PDFReader()
        self.vector_manager = VectorStoreManager()
        self.uploaded_documents = {}
        self.chunked_documents = {}
        self.document_ids = {}
        logging.info("DocumentManager initialized")

    def process_document(self, file):
        """
        Process an uploaded file: load, read PDF, chunk, and store in vector store.
        Returns: (status_message, page_list, filename, doc_id)
        """
        try:
            if file is None:
                return "No file uploaded", None, None

            logging.info(f"Processing file: {file}")

            # Load and validate file
            file_path = self.doc_loader.load_file(file)
            filename = os.path.basename(file_path)

            # Read PDF content
            page_list = self.pdf_reader.read_pdf(file_path)

            # Store the uploaded document
            self.uploaded_documents[filename] = file_path

            # Generate a unique document ID
            doc_id = str(uuid.uuid4())
            self.document_ids[filename] = doc_id

            # Chunk the pages
            chunks = chunk_documents(page_list, doc_id, chunk_size=2000, chunk_overlap=300)
            self.chunked_documents[filename] = chunks

            # Add chunks to vector store
            self.vector_manager.add_documents(chunks)

            return (
                f"Successfully loaded {filename} with {len(page_list)} pages",
                filename,
                doc_id
            )

        except Exception as e:
            logging.error(f"Error processing document: {str(e)}")
            return f"Error: {str(e)}", [], None, None

    def get_uploaded_documents(self):
        """Return the list of uploaded document filenames."""
        return list(self.uploaded_documents.keys())

    def get_chunks(self, filename):
        """Return chunks for a given filename."""
        return self.chunked_documents.get(filename, [])

    def get_document_id(self, filename):
        """Return the document ID for a given filename."""
        return self.document_ids.get(filename, None)
    
    def retrieve_top_k(self, query: str, selected_docs: List[str], k: int = 5) -> List[Dict[str, Any]]:
        """
        Retrieve the top K chunks across the selected documents based on the user's query.

        Args:
            query (str): The user's query.
            selected_docs (List[str]): List of selected document filenames from the dropdown.
            k (int): Number of top results to return (default is 5).

        Returns:
            List[Dict[str, Any]]: List of top K chunks with their text, metadata, and scores.
        """
        if not selected_docs:
            logging.warning("No documents selected for retrieval")
            return []

        all_results = []
        for filename in selected_docs:
            doc_id = self.get_document_id(filename)
            if not doc_id:
                logging.warning(f"No document ID found for filename: {filename}")
                continue

            # Search for relevant chunks within this document
            results = self.vector_manager.search(query, doc_id, k=k)
            all_results.extend(results)

        # Sort all results by score in descending order and take the top K
        all_results.sort(key=lambda x: x['score'], reverse=True)
        top_k_results = all_results[:k]

        # Log the list of retrieved documents
        #logging.info(f"Result from search :{all_results} ")
        logging.info(f"Retrieved top {k} documents:")
        for i, result in enumerate(top_k_results, 1):
            doc_id = result['metadata'].get('doc_id', 'Unknown')
            filename = next((name for name, d_id in self.document_ids.items() if d_id == doc_id), 'Unknown')
            logging.info(f"{i}. Filename: {filename}, Doc ID: {doc_id}, Score: {result['score']:.4f}, Text: {result['text'][:200]}...")

        return top_k_results
    
    def retrieve_summary_chunks(self, query: str, doc_id : str, k: int = 10):
        logging.info(f"Retrieving {k} chunks for summary: {query}, Document Id: {doc_id}")
        results = self.vector_manager.search(query, doc_id, k=k)
        top_k_results = results[:k]
        logging.info(f"Retrieved {len(top_k_results)} chunks for summary")
        
        return top_k_results