Upload 4 files

retriever/chat_manager.py

@@ -1,57 +1,114 @@
 from datetime import datetime
 import logging
 from typing import List
-from globals import app_config
 
-def chat_response(query: str, selected_docs: List[str], history: List[dict]) -> List[dict]:
-    """
-    Generate a chat response based on the user's query and selected documents.
 
-    Args:
-        query (str): The user's query.
-        selected_docs (List[str]): List of selected document filenames from the dropdown.
-        history (List[dict]): The chat history as a list of {'role': str, 'content': str} dictionaries.
+class ChatManager:
+    def __init__(self, documentManager, llmManager):
+        """
+        Initialize the ChatManager.
+        """
+        self.doc_manager = documentManager
+        self.llm_manager = llmManager
 
-    Returns:
-        List[dict]: Updated chat history with the new response in 'messages' format.
-    """
-    timestamp = datetime.now().strftime("%H:%M:%S")
+        logging.info("ChatManager initialized")
 
-    # Handle empty query
-    if not query:
-        return history + [{"role": "assistant", "content": "Please enter a query."}]
+    def generate_chat_response(self, query: str, selected_docs: List[str], history: List[dict]) -> List[dict]:
+        """
+        Generate a chat response based on the user's query and selected documents.
 
-    # Handle no selected documents
-    if not selected_docs:
-        return history + [{"role": "assistant", "content": "Please select at least one document."}]
+        Args:
+            query (str): The user's query.
+            selected_docs (List[str]): List of selected document filenames from the dropdown.
+            history (List[dict]): The chat history as a list of {'role': str, 'content': str} dictionaries.
 
-    # Retrieve the top 5 chunks based on the query and selected documents
-    top_k_results = app_config.doc_manager.retrieve_top_k(query, selected_docs, k=5)
-    
-    if not top_k_results:
-        return history + [
-            {"role": "user", "content": f"{query}"},
-            {"role": "assistant", "content": "No relevant information found in the selected documents."}
-        ]
+        Returns:
+            List[dict]: Updated chat history with the new response in 'messages' format.
+        """
+        timestamp = datetime.now().strftime("%H:%M:%S")
+        logging.info(f"Generating chat response for query: {query} at {timestamp}")
+
+        # Handle empty query
+        if not query:
+            logging.warning("Empty query received")
+            return history + [{"role": "assistant", "content": "Please enter a query."}]
+
+        # Handle no selected documents
+        if not selected_docs:
+            logging.warning("No documents selected")
+            return history + [{"role": "assistant", "content": "Please select at least one document."}]
+
+        # Retrieve the top 5 chunks based on the query and selected documents
+        try:
+            top_k_results = self.doc_manager.retrieve_top_k(query, selected_docs, k=5)
+        except Exception as e:
+            logging.error(f"Error retrieving chunks: {str(e)}")
+            return history + [
+                {"role": "user", "content": f"{query}"},
+                {"role": "assistant", "content": f"Error retrieving chunks: {str(e)}"}
+            ]
+
+        if not top_k_results:
+            logging.info("No relevant chunks found")
+            return history + [
+                {"role": "user", "content": f"{query}"},
+                {"role": "assistant", "content": "No relevant information found in the selected documents."}
+            ]
+
+        # Send the top K results to the LLM to generate a response
+        try:
+            llm_response, source_docs = self.llm_manager.generate_response(query, top_k_results)
+        except Exception as e:
+            logging.error(f"Error generating LLM response: {str(e)}")
+            return history + [
+                {"role": "user", "content": f"{query}"},
+                {"role": "assistant", "content": f"Error generating response: {str(e)}"}
+            ]
 
-    # Send the top K results to the LLM to generate a response
-    try:
-        llm_response, source_docs = app_config.gen_llm.generate_response(query, top_k_results)
-    except Exception as e:
+        # Format the response
+        response = llm_response
+        # Uncomment to include source docs in response (optional)
+        # for i, doc in enumerate(source_docs, 1):
+        #     doc_id = doc.metadata.get('doc_id', 'Unknown')
+        #     filename = next((name for name, d_id in self.doc_manager.document_ids.items() if d_id == doc_id), 'Unknown')
+        #     response += f"\n{i}. {filename}: {doc.page_content[:100]}..."
+
+        logging.info("Chat response generated successfully")
+        # Return updated history with new user query and LLM response
         return history + [
             {"role": "user", "content": f"{query}"},
-            {"role": "assistant", "content": f"Error generating response: {str(e)}"}
+            {"role": "assistant", "content": response}
         ]
 
-    # Format the response (uncomment and adapt if you want to include source docs)
-    response = llm_response
-    # for i, doc in enumerate(source_docs, 1):
-    #     doc_id = doc.metadata.get('doc_id', 'Unknown')
-    #     filename = next((name for name, d_id in app_config.doc_manager.document_ids.items() if d_id == doc_id), 'Unknown')
-    #     response += f"\n{i}. {filename}: {doc.page_content[:100]}..."
-
-    # Return updated history with new user query and LLM response
-    return history + [
-        {"role": "user", "content": f"{query}"},
-        {"role": "assistant", "content": response}
-    ]
+    def generate_summary(self, chunks: any, summary_type: str = "medium") -> str:
+        """
+        Generate a summary of the selected documents.
+
+        Args:
+            selected_docs (List[str]): List of selected document filenames.
+            summary_type (str): Type of summary ("small", "medium", "detailed").
+            k (int): Number of chunks to retrieve from DocumentManager.
+            include_toc (bool): Whether to include the table of contents (if available).
+
+        Returns:
+            str: Generated summary.
+
+        Raises:
+            ValueError: If summary_type is invalid or DocumentManager/LLM is not available.
+        """
+        if summary_type not in ["small", "medium", "detailed"]:
+            raise ValueError("summary_type must be 'small', 'medium', or 'detailed'")
+
+        if not chunks:
+            logging.warning("No documents selected for summarization")
+            return "Please select at least one document."
+
+        
+        llm_summary_response = self.llm_manager.generate_summary_v0(chunks = chunks)
+        #logging.info(f" Summary response {llm_summary_response}")
+
+        return llm_summary_response
+    
+    def generate_sample_questions(self, chunks: any):
+        questions = self.llm_manager.generate_questions(chunks = chunks)
+        return questions

retriever/document_manager.py

@@ -24,7 +24,7 @@ class DocumentManager:
         """
         try:
             if file is None:
-                return "No file uploaded", [], None, None
+                return "No file uploaded", None, None
 
             logging.info(f"Processing file: {file}")
 
@@ -51,7 +51,6 @@ class DocumentManager:
 
             return (
                 f"Successfully loaded {filename} with {len(page_list)} pages",
-                page_list,
                 filename,
                 doc_id
             )
@@ -104,11 +103,19 @@ class DocumentManager:
         top_k_results = all_results[:k]
 
         # Log the list of retrieved documents
-        logging.info(f"Result from search :{all_results} ")
+        #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

retriever/llm_manager.py

@@ -5,6 +5,8 @@ from langchain_groq import ChatGroq
 from langchain.chains import RetrievalQA
 from langchain_core.documents import Document
 from langchain_core.retrievers import BaseRetriever
+from langchain.chains.summarize import load_summarize_chain
+from langchain.prompts import PromptTemplate
 
 class LLMManager:
     DEFAULT_MODEL = "gemma2-9b-it"  # Set the default model name
@@ -30,7 +32,7 @@ class LLMManager:
         Raises:
             ValueError: If GROQ_API_KEY is not set.
         """
-        api_key = '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' #os.getenv("GROQ_API_KEY")
+        api_key = os.getenv("GROQ_API_KEY")
         if not api_key:
             raise ValueError("GROQ_API_KEY is not set. Please add it in your environment variables.")
         
@@ -113,4 +115,195 @@ class LLMManager:
             return response, source_docs
         except Exception as e:
             logging.error(f"Error during QA chain invocation: {str(e)}")
-            raise e
+            raise e
+
+    def generate_summary_v0(self, chunks: any):
+        logging.info("Generating summary ...")
+        
+        # Limit the number of chunks (for example, top 30 chunks)
+        limited_chunks = chunks[:30]
+        
+        # Combine text from the selected chunks
+        full_text = "\n".join(chunk['text'] for chunk in limited_chunks)
+        text_length = len(full_text)
+        logging.info(f"Total text length (characters): {text_length}")
+        
+        # Define a maximum character limit to fit in a 1024-token context.
+        # For many models, roughly 3200 characters is a safe limit.
+        MAX_CHAR_LIMIT = 3200
+        if text_length > MAX_CHAR_LIMIT:
+            logging.warning(f"Input text too long ({text_length} chars), truncating to {MAX_CHAR_LIMIT} chars.")
+            full_text = full_text[:MAX_CHAR_LIMIT]
+        
+        # Define a custom prompt to instruct concise summarization in bullet points.
+        custom_prompt_template = """
+            You are an expert summarizer. Summarize the following text into a concise summary using bullet points.
+            Ensure that the final summary is no longer than 20-30 bullet points and fits within 15-20 lines.
+            Focus only on the most critical points.
+
+            Text to summarize:
+            {text}
+
+            Summary:
+            """
+        prompt = PromptTemplate(input_variables=["text"], template=custom_prompt_template)
+        
+        # Use the 'stuff' chain type to send a single LLM request with our custom prompt.
+        chain = load_summarize_chain(self.generation_llm, chain_type="stuff", prompt=prompt)
+        
+        # Wrap the full text in a single Document object (chain expects a list of Documents)
+        docs = [Document(page_content=full_text)]
+        
+        # Generate the summary
+        summary = chain.invoke(docs)
+        return summary['output_text']
+    
+    def generate_questions(self, chunks: any):
+        logging.info("Generating sample questions ...")
+        
+        # Use the top 30 chunks or fewer
+        limited_chunks = chunks[:30]
+        
+        # Combine text from chunks
+        full_text = "\n".join(chunk['text'] for chunk in limited_chunks)
+        text_length = len(full_text)
+        logging.info(f"Total text length for questions: {text_length}")
+        
+        MAX_CHAR_LIMIT = 3200
+        if text_length > MAX_CHAR_LIMIT:
+            logging.warning(f"Input text too long ({text_length} chars), truncating to {MAX_CHAR_LIMIT} chars.")
+            full_text = full_text[:MAX_CHAR_LIMIT]
+        
+        # Prompt template for generating questions
+        question_prompt_template = """
+        You are an AI expert at creating questions from documents.
+
+        Based on the text below, generate not less than 20 insightful and highly relevant sample questions that a user might ask to better understand the content.
+
+        **Instructions:**
+        - Questions must be specific to the document's content and context.
+        - Avoid generic questions like 'What is this document about?'
+        - Do not include numbers, prefixes (e.g., '1.', '2.'), or explanations (e.g., '(Clarifies...)').
+        - Each question should be a single, clear sentence ending with a question mark.
+        - Focus on key concepts, processes, components, or use cases mentioned in the text.
+
+        Text:
+        {text}
+
+        Output format:
+        What is the purpose of the Communication Server in Collateral Management?
+        How does the system handle data encryption for secure communication?
+        ...
+        """
+        prompt = PromptTemplate(input_variables=["text"], template=question_prompt_template)
+        
+        chain = load_summarize_chain(self.generation_llm, chain_type="stuff", prompt=prompt)
+        docs = [Document(page_content=full_text)]
+
+        try:
+            result = chain.invoke(docs)
+            question_output = result.get("output_text", "").strip()
+            
+            # Clean and parse the output into a list of questions
+            questions = []
+            for line in question_output.split("\n"):
+                # Remove any leading/trailing whitespace, numbers, or bullet points
+                cleaned_line = line.strip().strip("-*1234567890. ").rstrip(".")
+                # Remove any explanation in parentheses
+                cleaned_line = cleaned_line.split("(")[0].strip()
+                # Ensure the line is a valid question (ends with '?' and is not empty)
+                if cleaned_line and cleaned_line.endswith("?"):
+                    questions.append(cleaned_line)
+            
+            # Limit to 10 questions
+            questions = questions[:10]
+            logging.info(f"Generated questions: {questions}")
+            return questions
+        except Exception as e:
+            logging.error(f"Error generating questions: {e}")
+            return []
+    
+    def generate_summary(self, chunks: Any, toc_text: Any, summary_type: str = "medium") -> str:
+        """
+        Generate a summary of the document using LangChain's summarization chains.
+
+        Args:
+            vector_store_manager: Instance of VectorStoreManager with a FAISS vector store.
+            summary_type (str): Type of summary ("small", "medium", "detailed").
+            k (int): Number of chunks to retrieve from the vector store.
+            include_toc (bool): Whether to include the table of contents (if available).
+
+        Returns:
+            str: Generated summary.
+
+        Raises:
+            ValueError: If summary_type is invalid or vector store is not initialized.
+        """
+
+        # Define chunk retrieval parameters based on summary type
+        if summary_type == "small":
+            k = min(k, 3)  # Fewer chunks for small summary
+            chain_type = "stuff"  # Use stuff for small summaries
+            word_count = "50-100"
+        elif summary_type == "medium":
+            k = min(k, 10)
+            chain_type = "map_reduce"  # Use map-reduce for medium summaries
+            word_count = "200-400"
+        else:  # detailed
+            k = min(k, 20)
+            chain_type = "map_reduce"  # Use map-reduce for detailed summaries
+            word_count = "500-1000"
+
+        # Define prompts
+        if chain_type == "stuff":
+            prompt = PromptTemplate(
+                input_variables=["text"],
+                template=(
+                    "Generate a {summary_type} summary ({word_count} words) of the following document excerpts. "
+                    "Focus on key points and ensure clarity. Stick strictly to the provided text:\n\n"
+                    "{toc_prompt}{text}"
+                ).format(
+                    summary_type=summary_type,
+                    word_count=word_count,
+                    toc_prompt="Table of Contents:\n{toc_text}\n\n" if toc_text else ""
+                )
+            )
+            chain = load_summarize_chain(
+                llm=self.generation_llm,
+                chain_type="stuff",
+                prompt=prompt
+            )
+        else:  # map_reduce
+            map_prompt = PromptTemplate(
+                input_variables=["text"],
+                template=(
+                    "Summarize the following document excerpt in 1-2 sentences, focusing on key points. "
+                    "Consider the document's structure from this table of contents:\n\n"
+                    "Table of Contents:\n{toc_text}\n\nExcerpt:\n{text}"
+                ).format(toc_text=toc_text if toc_text else "Not provided")
+            )
+            combine_prompt = PromptTemplate(
+                input_variables=["text"],
+                template=(
+                    "Combine the following summaries into a cohesive {summary_type} summary "
+                    "({word_count} words) of the document. Ensure clarity, avoid redundancy, and "
+                    "organize by key themes or sections if applicable:\n\n{text}"
+                ).format(summary_type=summary_type, word_count=word_count)
+            )
+            chain = load_summarize_chain(
+                llm=self.generation_llm,
+                chain_type="map_reduce",
+                map_prompt=map_prompt,
+                combine_prompt=combine_prompt,
+                return_intermediate_steps=False
+            )
+
+        # Run the chain
+        try:
+            logging.info(f"Generating {summary_type} summary with {len(chunks)} chunks")
+            summary = chain.run(chunks)
+            logging.info(f"{summary_type.capitalize()} summary generated successfully")
+            return summary
+        except Exception as e:
+            logging.error(f"Error generating summary: {str(e)}")
+            return f"Error generating summary: {str(e)}"

retriever/vector_store_manager.py

@@ -63,7 +63,7 @@ class VectorStoreManager:
         self.vector_store.save_local(self.embedding_path)
         logging.info(f"Vector store updated and saved to {self.embedding_path}")
 
-    def search(self, query, doc_id, k=4):
+    def search(self, query, doc_id, k=10):
         """
         Search the vector store for relevant chunks, filtered by doc_id.