Upload 2 files

chunker.py

@@ -6,6 +6,8 @@ from typing import List, Dict, Optional
 
 from pypdf import PdfReader
 import docx as python_docx
+# ADDED: Import pandas to handle CSV/XLSX files
+import pandas as pd
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 # --- Logging Setup ---
@@ -18,8 +20,7 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# --- Text Extraction Helper Functions ---
-# Note: These are duplicated from llm_handling.py to make this a standalone script.
+# --- Text Extraction Helper Functions (MODIFIED) ---
 def extract_text_from_file(file_path: str, file_type: str) -> Optional[str]:
     logger.info(f"Extracting text from {file_type.upper()} file: {os.path.basename(file_path)}")
     text_content = None
@@ -33,6 +34,18 @@ def extract_text_from_file(file_path: str, file_type: str) -> Optional[str]:
         elif file_type == 'txt':
             with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                 text_content = f.read()
+        # ADDED: Logic for CSV and XLSX files
+        elif file_type in ['csv', 'xlsx']:
+            df = pd.read_excel(file_path) if file_type == 'xlsx' else pd.read_csv(file_path)
+            if df.empty:
+                return ""
+            # Convert each row into a descriptive string format
+            text_chunks = []
+            for index, row in df.iterrows():
+                row_text = f"Row {index + 1}: "
+                row_text += ", ".join([f"{col}: {val}" for col, val in row.items() if pd.notna(val)])
+                text_chunks.append(row_text)
+            text_content = "\n".join(text_chunks)
         else:
             logger.warning(f"Unsupported file type for text extraction: {file_type} for file {os.path.basename(file_path)}")
             return None
@@ -45,18 +58,16 @@ def extract_text_from_file(file_path: str, file_type: str) -> Optional[str]:
         logger.error(f"Error extracting text from {os.path.basename(file_path)} ({file_type.upper()}): {e}", exc_info=True)
         return None
 
-SUPPORTED_EXTENSIONS = {
-    'pdf': lambda path: extract_text_from_file(path, 'pdf'),
-    'docx': lambda path: extract_text_from_file(path, 'docx'),
-    'txt': lambda path: extract_text_from_file(path, 'txt'),
-}
+# MODIFIED: Added 'csv' and 'xlsx' to the list of supported extensions
+SUPPORTED_EXTENSIONS = ['pdf', 'docx', 'txt', 'csv', 'xlsx']
+
 
 def process_sources_and_create_chunks(
     sources_dir: str,
     output_file: str,
     chunk_size: int = 1000,
     chunk_overlap: int = 150,
-    text_output_dir: Optional[str] = None  # MODIFIED: Added optional parameter
+    text_output_dir: Optional[str] = None
 ) -> None:
     """
     Scans a directory for source files, extracts text, splits it into chunks,
@@ -69,7 +80,6 @@ def process_sources_and_create_chunks(
 
     logger.info(f"Starting chunking process. Sources: '{sources_dir}', Output: '{output_file}'")
     
-    # MODIFIED: Create text output directory if provided
     if text_output_dir:
         os.makedirs(text_output_dir, exist_ok=True)
         logger.info(f"Will save raw extracted text to: '{text_output_dir}'")
@@ -90,10 +100,10 @@ def process_sources_and_create_chunks(
             continue
 
         logger.info(f"Processing source file: {filename}")
-        text_content = SUPPORTED_EXTENSIONS[file_ext](file_path)
+        # MODIFIED: Simplified the call to the unified extraction function
+        text_content = extract_text_from_file(file_path, file_ext)
 
         if text_content:
-            # MODIFIED: Save the raw text to a file if directory is specified
             if text_output_dir:
                 try:
                     text_output_path = os.path.join(text_output_dir, f"{filename}.txt")
@@ -143,7 +153,7 @@ def main():
         '--sources-dir',
         type=str,
         required=True,
-        help="The directory containing source files (PDFs, DOCX, TXT)."
+        help="The directory containing source files (PDFs, DOCX, TXT, CSV, XLSX)."
     )
     parser.add_argument(
         '--output-file',
@@ -151,7 +161,6 @@ def main():
         required=True,
         help="The full path for the output JSON file containing the chunks."
     )
-    # MODIFIED: Added new optional argument
     parser.add_argument(
         '--text-output-dir',
         type=str,
@@ -179,7 +188,7 @@ def main():
             output_file=args.output_file,
             chunk_size=args.chunk_size,
             chunk_overlap=args.chunk_overlap,
-            text_output_dir=args.text_output_dir  # MODIFIED: Pass argument
+            text_output_dir=args.text_output_dir
         )
     except Exception as e:
         logger.critical(f"A critical error occurred during the chunking process: {e}", exc_info=True)

llm_handling.py

@@ -14,6 +14,8 @@ import torch
 from sentence_transformers import SentenceTransformer
 from pypdf import PdfReader
 import docx as python_docx
+# ADDED: Import pandas to handle CSV/XLSX files
+import pandas as pd
 
 from llama_index.core.llms import ChatMessage
 from llama_index.llms.groq import Groq as LlamaIndexGroqClient
@@ -27,7 +29,6 @@ from langchain.callbacks.manager import CallbackManagerForRetrieverRun
 from langchain.schema.runnable import RunnablePassthrough, RunnableParallel
 from langchain.schema.output_parser import StrOutputParser
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-# MODIFIED: Import the new prompt
 from system_prompts import RAG_SYSTEM_PROMPT, FALLBACK_SYSTEM_PROMPT, QA_FORMATTER_PROMPT
 
 logger = logging.getLogger(__name__)
@@ -43,7 +44,6 @@ if not GROQ_API_KEY:
     logger.critical("CRITICAL: BOT_API_KEY environment variable not found. Services will fail.")
 
 FALLBACK_LLM_MODEL_NAME = os.getenv("GROQ_FALLBACK_MODEL", "llama-3.1-70b-versatile")
-# ADDED: New constant for the auxiliary model
 AUXILIARY_LLM_MODEL_NAME = os.getenv("GROQ_AUXILIARY_MODEL", "llama3-8b-8192")
 _MODULE_BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 RAG_FAISS_INDEX_SUBDIR_NAME = "faiss_index"
@@ -61,7 +61,8 @@ RAG_DEFAULT_RETRIEVER_K = int(os.getenv("RAG_RETRIEVER_K", 3))
 GDRIVE_SOURCES_ENABLED = os.getenv("GDRIVE_SOURCES_ENABLED", "False").lower() == "true"
 GDRIVE_FOLDER_ID_OR_URL = os.getenv("GDRIVE_FOLDER_URL")
 
-# --- Text Extraction Helper Function ---
+
+# --- Text Extraction Helper Function (MODIFIED) ---
 def extract_text_from_file(file_path: str, file_type: str) -> Optional[str]:
     logger.info(f"Extracting text from {file_type.upper()} file: {os.path.basename(file_path)}")
     try:
@@ -74,13 +75,28 @@ def extract_text_from_file(file_path: str, file_type: str) -> Optional[str]:
         elif file_type == 'txt':
             with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                 return f.read()
+        # ADDED: Logic for CSV and XLSX files
+        elif file_type in ['csv', 'xlsx']:
+            df = pd.read_excel(file_path) if file_type == 'xlsx' else pd.read_csv(file_path)
+            if df.empty:
+                return ""
+            # Convert each row into a descriptive string format
+            text_chunks = []
+            for index, row in df.iterrows():
+                row_text = f"Row {index + 1}: "
+                row_text += ", ".join([f"{col}: {val}" for col, val in row.items() if pd.notna(val)])
+                text_chunks.append(row_text)
+            return "\n".join(text_chunks)
+            
         logger.warning(f"Unsupported file type for text extraction: {file_type}")
         return None
     except Exception as e:
         logger.error(f"Error extracting text from {os.path.basename(file_path)}: {e}", exc_info=True)
         return None
 
-FAISS_RAG_SUPPORTED_EXTENSIONS = {'pdf': 'pdf', 'docx': 'docx', 'txt': 'txt'}
+# MODIFIED: Added 'csv' and 'xlsx' to the list of supported extensions for RAG
+FAISS_RAG_SUPPORTED_EXTENSIONS = {'pdf': 'pdf', 'docx': 'docx', 'txt': 'txt', 'csv': 'csv', 'xlsx': 'xlsx'}
+
 
 # --- FAISS RAG System ---
 class FAISSRetrieverWithScore(BaseRetriever):
@@ -180,18 +196,17 @@ class KnowledgeRAG:
                         processed_files_this_build.append(filename)
 
         if not all_docs_for_vectorstore:
-            self.logger.warning(f"No processable PDF/DOCX/TXT documents found in '{source_folder_path}'. RAG index will only contain other sources if available.")
+            self.logger.warning(f"No processable PDF/DOCX/TXT/CSV/XLSX documents found in '{source_folder_path}'. RAG index will only contain other sources if available.")
 
 
         self.processed_source_files = processed_files_this_build
         
-        # This print statement is kept for console visibility on startup/rebuild
         print("\n--- Document Files Used for RAG Index ---")
         if self.processed_source_files:
             for filename in self.processed_source_files:
                 print(f"- {filename}")
         else:
-            print("No PDF/DOCX/TXT source files were processed for the RAG index.")
+            print("No PDF/DOCX/TXT/CSV/XLSX source files were processed for the RAG index.")
         print("---------------------------------------\n")
 
         if not all_docs_for_vectorstore:
@@ -242,7 +257,6 @@ class KnowledgeRAG:
 
     def invoke(self, query: str, top_k: Optional[int] = None) -> Dict[str, Any]:
         if not self.rag_chain:
-            # MODIFIED: Changed severity
             self.logger.warning("RAG system not fully initialized. Cannot invoke.")
             return {"answer": "The provided bibliography does not contain specific information on this topic.", "source": "system_error", "cited_source_details": []}
 
@@ -263,7 +277,6 @@ class KnowledgeRAG:
 
             context_str = self.format_docs(retrieved_docs)
             
-            # MODIFIED: Added full logging as per user request
             print(f"\n--- RAG INVOKE ---")
             print(f"QUESTION: {query}")
             print(f"CONTEXT:\n{context_str}")
@@ -314,13 +327,11 @@ class KnowledgeRAG:
             self.retriever.k = k_to_use
         
         try:
-            # Check for docs first to avoid streaming "no info" message
             retrieved_docs = self.retriever.get_relevant_documents(query)
             if not retrieved_docs:
                 yield "The provided bibliography does not contain specific information on this topic."
                 return
             
-            # MODIFIED: Added full logging for streaming as per user request
             context_str = self.format_docs(retrieved_docs)
             print(f"\n--- RAG STREAM ---")
             print(f"QUESTION: {query}")
@@ -368,8 +379,6 @@ class GroqBot:
             messages.append(ChatMessage(role="system", content=f"**Potentially Relevant Q&A Information from other sources:**\n{qa_info}"))
         messages.append(ChatMessage(role="user", content=f"**Current User Query:**\n{current_query}"))
         
-        # MODIFIED: Added full logging as per user request
-        # The conversion to dict is necessary because ChatMessage is not directly JSON serializable
         messages_for_print = [msg.dict() for msg in messages]
         print(f"\n--- FALLBACK STREAM ---")
         print(f"MESSAGES SENT TO LLM:\n{json.dumps(messages_for_print, indent=2)}")
@@ -384,14 +393,9 @@ class GroqBot:
             self.logger.error(f"Groq API error in get_response (Fallback): {e}", exc_info=True)
             yield "I am currently unable to process this request due to a technical issue."
 
-# ADDED: New function for formatting QA answers
 def get_answer_from_context(question: str, context: str, system_prompt: str) -> str:
-    """
-    Calls the LLM with a specific question and context from a QA source (CSV/XLSX).
-    """
     logger.info(f"Formatting answer for question '{question[:50]}...' using QA context.")
     try:
-        # Use the auxiliary model for this task for speed and cost-efficiency
         formatter_llm = ChatGroq(
             temperature=0.1,
             groq_api_key=GROQ_API_KEY,
@@ -402,7 +406,6 @@ def get_answer_from_context(question: str, context: str, system_prompt: str) ->
         
         chain = prompt_template | formatter_llm | StrOutputParser()
         
-        # MODIFIED: Added full logging as per user request
         print(f"\n--- QA FORMATTER ---")
         print(f"QUESTION: {question}")
         print(f"CONTEXT:\n{context}")
@@ -421,14 +424,9 @@ def get_answer_from_context(question: str, context: str, system_prompt: str) ->
         logger.error(f"Error in get_answer_from_context: {e}", exc_info=True)
         return "Sorry, I was unable to formulate an answer based on the available information."
 
-# ADDED: New function for streaming QA answers
 def stream_answer_from_context(question: str, context: str, system_prompt: str) -> Iterator[str]:
-    """
-    Calls the LLM with a specific question and context from a QA source and streams the response.
-    """
     logger.info(f"Streaming formatted answer for question '{question[:50]}...' using QA context.")
     try:
-        # Use the auxiliary model for this task for speed and cost-efficiency
         formatter_llm = ChatGroq(
             temperature=0.1,
             groq_api_key=GROQ_API_KEY,
@@ -439,7 +437,6 @@ def stream_answer_from_context(question: str, context: str, system_prompt: str)
         
         chain = prompt_template | formatter_llm | StrOutputParser()
         
-        # MODIFIED: Added full logging as per user request
         print(f"\n--- QA FORMATTER (STREAM) ---")
         print(f"QUESTION: {question}")
         print(f"CONTEXT:\n{context}")
@@ -493,7 +490,6 @@ def download_and_unzip_gdrive_folder(folder_id_or_url: str, target_dir: str) ->
             logger.info(f"Successfully moved GDrive contents to {target_dir}")
             return True
         except Exception as e:
-            # MODIFIED: Corrected self.logger to logger
             logger.error(f"Error during GDrive download/processing: {e}", exc_info=True)
             return False
 
@@ -556,15 +552,10 @@ def initialize_and_get_rag_system(force_rebuild: bool = False) -> Optional[Knowl
 groq_bot_instance = GroqBot()
 
 def get_auxiliary_chat_response(messages: List[Dict]) -> str:
-    """
-    Handles requests for auxiliary tasks like generating titles or follow-up questions.
-    Uses a separate, smaller model for efficiency.
-    """
     logger.info(f"Routing auxiliary request to model: {AUXILIARY_LLM_MODEL_NAME}")
     try:
-        # Initialize a dedicated client for this call to use the specific auxiliary model
         aux_client = ChatGroq(
-            temperature=0.2, # A bit more creative than RAG, but still grounded
+            temperature=0.2,
             groq_api_key=GROQ_API_KEY,
             model_name=AUXILIARY_LLM_MODEL_NAME
         )