update the system

prepare_env.py

@@ -1,10 +1,16 @@
 import json
 import os
+import glob
+from typing import List, Optional
 from dotenv import load_dotenv
+import logging
 
 # Load environment variables from .env file
 load_dotenv()
+
 from langchain_core.documents import Document
+from langchain_core.output_parsers import BaseOutputParser
+from langchain_core.prompts import PromptTemplate
 from langchain_community.vectorstores import Chroma
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.retrievers import BM25Retriever
@@ -12,78 +18,278 @@ from langchain.retrievers import EnsembleRetriever
 from langchain.retrievers.multi_query import MultiQueryRetriever
 from langchain_google_genai import GoogleGenerativeAI
 
-def prepare_environment_and_retriever(
-    chunks_path="./chunks.json",
-    model_name="intfloat/multilingual-e5-base",
-    collection_name="Guide_2023_e5_multilingual",
-    persist_directory="chroma_db_multilingual",
-    k_vector=6,
-    k_sparse=2,
-    weights=[0.5, 0.5],
-    llm_model_name="gemini-2.0-flash"
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class LineListOutputParser(BaseOutputParser[List[str]]):
+    """Custom output parser for a list of lines with better error handling."""
+    
+    def parse(self, text: str) -> List[str]:
+        """Parse the LLM output into a list of queries."""
+        try:
+            lines = text.strip().split("\n")
+            # Remove empty lines and clean up
+            cleaned_lines = []
+            for line in lines:
+                cleaned = line.strip()
+                if cleaned and not cleaned.startswith("#") and len(cleaned) > 5:
+                    # Remove numbering if present (e.g., "1. ", "- ", etc.)
+                    if cleaned[0].isdigit() and ". " in cleaned:
+                        cleaned = cleaned.split(". ", 1)[1]
+                    elif cleaned.startswith("- "):
+                        cleaned = cleaned[2:]
+                    cleaned_lines.append(cleaned)
+            
+            # Ensure we have at least one query
+            if not cleaned_lines:
+                cleaned_lines = [text.strip()]
+                
+            return cleaned_lines
+        except Exception as e:
+            logger.warning(f"Error parsing output: {e}. Returning original text.")
+            return [text.strip()] if text.strip() else [""]
+
+def create_custom_multi_query_retriever(
+    base_retriever,
+    llm,
+    num_queries: int = 5,
+    include_original: bool = True
 ):
-    # Load the chunks.json
-    with open(chunks_path, "r", encoding="utf-8") as f:
-        chunks_data = json.load(f)
+    """Create a custom MultiQueryRetriever with improved prompt."""
+    
+    # Custom prompt template for better query generation
+#     query_prompt = PromptTemplate(
+#         input_variables=["question"],
+#         template="""You are an AI assistant specialized in generating diverse search queries. 
+# Your task is to generate {num_queries} different versions of the given user question to retrieve relevant documents from a knowledge base.
+
+# Guidelines:
+# - Create variations that capture different aspects and perspectives of the question
+# - Use synonyms and alternative phrasings
+# - Consider different levels of specificity (broader and narrower)
+# - Focus on the core intent while varying the expression
+# - Each query should be a complete, well-formed question or statement
+
+# Original question: {question}
+
+# Generate {num_queries} alternative queries (one per line):""".replace("{num_queries}", str(num_queries))
+#     )
+    
+    # Create the MultiQueryRetriever with custom components
+    multi_query_retriever = MultiQueryRetriever.from_llm(
+        retriever=base_retriever,
+        llm=llm,
+        include_original=include_original
+    )
+    
+    # # Override the output parser
+    # multi_query_retriever.output_parser = LineListOutputParser()
+    
+    return multi_query_retriever
+
+def validate_environment():
+    """Validate that required environment variables are set."""
+    required_vars = ["GOOGLE_API_KEY"]
+    missing_vars = [var for var in required_vars if not os.getenv(var)]
+    
+    if missing_vars:
+        raise ValueError(f"Missing required environment variables: {missing_vars}")
+    
+    logger.info("✅ Environment variables validated.")
 
+def load_documents_from_json(chunks_directory: str) -> List[Document]:
+    """Load documents from JSON files with better error handling."""
+    json_files = glob.glob(os.path.join(chunks_directory, "*.json"))
+    
+    if not json_files:
+        raise ValueError(f"No JSON files found in directory: {chunks_directory}")
+    
+    logger.info(f"Found {len(json_files)} JSON files: {[os.path.basename(f) for f in json_files]}")
+    
     documents = []
+    total_processed = 0
 
-    for element in chunks_data:
-        text = element["text"]
-        metadata = {
-            "source": element["filename"],
-            "filetype": element["filetype"],
-            "element_id": element["element_id"]
-        }
+    for json_file in json_files:
+        try:
+            logger.info(f"Processing: {os.path.basename(json_file)}")
+            
+            with open(json_file, "r", encoding="utf-8") as f:
+                chunks_data = json.load(f)
 
-        if element.get("type") == "TableElement":
-            metadata["table_text_as_html"] = element["table_text_as_html"]
+            file_doc_count = 0
+            for element in chunks_data:
+                try:
+                    text = element.get("text", "").strip()
+                    if not text:  # Skip empty text
+                        continue
+                        
+                    metadata = {
+                        "source": element.get("filename", "unknown"),
+                        "filetype": element.get("filetype", "unknown"),
+                        "element_id": element.get("element_id", "unknown"),
+                        "json_source": os.path.basename(json_file)
+                    }
 
-        doc = Document(page_content=text, metadata=metadata)
-        documents.append(doc)
+                    # Add table-specific metadata if present
+                    if element.get("type") == "TableElement" and element.get("table_text_as_html"):
+                        metadata["table_text_as_html"] = element["table_text_as_html"]
+                        # metadata["element_type"] = "table"
+                    else:
+                        metadata["element_type"] = element.get("type", "text")
 
-    # Create the embedding function
+                    doc = Document(page_content=text, metadata=metadata)
+                    documents.append(doc)
+                    file_doc_count += 1
+                    
+                except Exception as e:
+                    logger.warning(f"Error processing element in {json_file}: {e}")
+                    continue
+            
+            logger.info(f"  → Loaded {file_doc_count} documents from {os.path.basename(json_file)}")
+            total_processed += file_doc_count
+            
+        except Exception as e:
+            logger.error(f"Error processing file {json_file}: {e}")
+            continue
+
+    if not documents:
+        raise ValueError("No valid documents were loaded from any JSON files.")
+    
+    logger.info(f"✅ Total loaded: {len(documents)} documents from {len(json_files)} JSON files.")
+    return documents
+
+def prepare_environment_and_retriever(
+    chunks_directory: str = "./data/",
+    model_name: str = "intfloat/multilingual-e5-base",
+    collection_name: str = "Guide_2023_e5_multilingual",
+    persist_directory: str = "chroma_db_multilingual",
+    k_vector: int = 6,
+    k_sparse: int = 2,
+    ensemble_weights: List[float] = [0.5, 0.5],
+    llm_model_name: str = "gemini-2.0-flash-exp",
+    num_query_variations: int = 5,
+    include_original_query: bool = True,
+    temperature: float = 0.1
+):
+    """
+    Prepare the complete retrieval environment with MultiQueryRetriever.
+    
+    Args:
+        chunks_directory: Directory containing JSON files with document chunks
+        model_name: HuggingFace embedding model name
+        collection_name: Chroma collection name
+        persist_directory: Directory to persist Chroma database
+        k_vector: Number of documents to retrieve from vector search
+        k_sparse: Number of documents to retrieve from BM25 search
+        ensemble_weights: Weights for ensemble retriever [vector, sparse]
+        llm_model_name: Google Gemini model name for query expansion
+        num_query_variations: Number of query variations to generate
+        include_original_query: Whether to include original query in search
+        temperature: LLM temperature for query generation
+    
+    Returns:
+        MultiQueryRetriever: Configured retriever ready for use
+    """
+    
+    # Validate environment
+    validate_environment()
+    
+    # Load documents
+    documents = load_documents_from_json(chunks_directory)
+    
+    # Create embedding function
+    logger.info(f"Creating embeddings with model: {model_name}")
     embedding_function = HuggingFaceEmbeddings(
-        model_name=model_name
+        model_name=model_name,
     )
 
-    # Create and persist the vector store
-    vectorstore = Chroma.from_documents(
-        documents=documents,
-        embedding=embedding_function,
-        collection_name=collection_name,
-        persist_directory=persist_directory
-    )
-    # vectorstore.persist()
-    print("✅ Stored with multilingual embeddings.")
+    # Create or load vector store
+    logger.info("Creating/loading vector store...")
+    try:
+        # Try to load existing vectorstore first
+        if os.path.exists(persist_directory):
+            vectorstore = Chroma(
+                collection_name=collection_name,
+                embedding_function=embedding_function,
+                persist_directory=persist_directory
+            )
+            logger.info("✅ Loaded existing vector store.")
+        else:
+            # Create new vectorstore
+            vectorstore = Chroma.from_documents(
+                documents=documents,
+                embedding=embedding_function,
+                collection_name=collection_name,
+                persist_directory=persist_directory
+            )
+            logger.info("✅ Created new vector store with multilingual embeddings.")
+    except Exception as e:
+        logger.warning(f"Error with persistent storage: {e}. Creating in-memory store.")
+        vectorstore = Chroma.from_documents(
+            documents=documents,
+            embedding=embedding_function,
+            collection_name=collection_name
+        )
 
-    # Build retrievers
-    retriever_multilingual = vectorstore.as_retriever(
+    # Create base retrievers
+    logger.info("Setting up retrievers...")
+    
+    # Vector retriever
+    vector_retriever = vectorstore.as_retriever(
         search_type="similarity",
         search_kwargs={"k": k_vector}
     )
 
+    # BM25 (sparse) retriever
     bm25_retriever = BM25Retriever.from_documents(documents)
     bm25_retriever.k = k_sparse
 
     # Ensemble retriever (combining vector + sparse search)
     ensemble_retriever = EnsembleRetriever(
-        retrievers=[retriever_multilingual, bm25_retriever],
-        weights=weights
+        retrievers=[vector_retriever, bm25_retriever],
+        weights=ensemble_weights
     )
+    logger.info(f"✅ Ensemble retriever created with weights: {ensemble_weights}")
 
     # Language model for multi-query expansion
-    # Using GoogleGenerativeAI instead of ChatGoogleGenerativeAI
-    llm = GoogleGenerativeAI(
-        model=llm_model_name,
-        google_api_key=os.getenv("GOOGLE_API_KEY")
-    )
-
-    expanding_retriever = MultiQueryRetriever.from_llm(
-        retriever=ensemble_retriever,
-        llm=llm
-    )
-
-    print("✅ Retrieval system ready (vector + sparse + ensemble + multi-query).")
+    logger.info(f"Initializing LLM: {llm_model_name}")
+    try:
+        llm = GoogleGenerativeAI(
+            model=llm_model_name,
+            google_api_key=os.getenv("GOOGLE_API_KEY"),
+            temperature=temperature,
+            max_output_tokens=1000  # Reasonable limit for query generation
+        )
+        
+        # Test the LLM with a simple call
+        test_response = llm.invoke("Generate a simple test query about artificial intelligence.")
+        logger.info("✅ LLM connection verified.")
+        
+    except Exception as e:
+        logger.error(f"Error initializing LLM: {e}")
+        raise
 
-    return expanding_retriever  # Return the final retriever
+    # Create MultiQueryRetriever with custom configuration
+    logger.info("Creating MultiQueryRetriever...")
+    try:
+        multi_query_retriever = create_custom_multi_query_retriever(
+            base_retriever=ensemble_retriever,
+            llm=llm,
+            num_queries=num_query_variations,
+            include_original=include_original_query
+        )
+        
+        logger.info(f"✅ MultiQueryRetriever ready:")
+        logger.info(f"  - Vector search: top-{k_vector}")
+        logger.info(f"  - Sparse search: top-{k_sparse}")
+        logger.info(f"  - Ensemble weights: {ensemble_weights}")
+        logger.info(f"  - Query variations: {num_query_variations}")
+        logger.info(f"  - Include original: {include_original_query}")
+        
+        return multi_query_retriever
+        
+    except Exception as e:
+        logger.error(f"Error creating MultiQueryRetriever: {e}")
+        logger.info("Falling back to ensemble retriever without query expansion.")
+        return ensemble_retriever

rag_pipeline.py

@@ -1,15 +1,16 @@
 import json
 import re
+import glob
+import os
 from langchain_google_genai import GoogleGenerativeAI
 from langchain_core.documents import Document
 from langdetect import detect
-import os
 from dotenv import load_dotenv
 
 # Load environment variables from .env file
 load_dotenv()
 
-def generate_rag_response(query, retrieved_documents, model="gemini-2.0-flash"):
+def generate_rag_response(query, retrieved_documents, model="gemini-2.0-flash-exp"):
     """
     Perform Retrieval-Augmented Generation (RAG) using Google's Gemini.
     Args:
@@ -192,14 +193,14 @@ def format_response_with_sequential_citations(response_text, unique_ids, clean_a
     
     return formatted_response.strip()
 
-def retrieve_documents_and_prepare_inputs(query, expanding_retriever, chunks_path="./chunks.json"):
+def retrieve_documents_and_prepare_inputs(query, expanding_retriever, chunks_directory="./data/"):
     """
     Retrieve relevant documents and prepare them for the RAG generation.
 
     Args:
         query (str): The user's query.
         expanding_retriever: The retriever object (e.g., returned by prepare_environment_and_retriever).
-        chunks_path (str): Path to the chunks.json file.
+        chunks_directory (str): Path to the directory containing JSON files.
 
     Returns:
         tuple: (source_texts_for_rag, retrieved_elements_full)
@@ -208,23 +209,43 @@ def retrieve_documents_and_prepare_inputs(query, expanding_retriever, chunks_pat
     retrieved_docs = expanding_retriever.get_relevant_documents(query)
 
     retrieved_chunk_ids = [doc.metadata["element_id"] for doc in retrieved_docs]
-
-    # Load all chunks
-    with open(chunks_path, "r", encoding="utf-8") as f:
-        chunks_data = json.load(f)
+    
+    # Get unique filenames from retrieved documents
+    needed_filenames = set(doc.metadata["source"] for doc in retrieved_docs)
+    
+    # Convert PDF filenames to JSON filenames (e.g., "file.pdf" -> "file.json")
+    needed_json_files = []
+    for filename in needed_filenames:
+        # Remove extension and add .json
+        base_name = os.path.splitext(filename)[0]
+        json_filename = f"{base_name}.json"
+        json_path = os.path.join(chunks_directory, json_filename)
+        if os.path.exists(json_path):
+            needed_json_files.append(json_path)
+        else:
+            print(f"Warning: JSON file not found: {json_path}")
+
+    # Load only the needed JSON files
+    all_chunks_data = []
+    for json_file in needed_json_files:
+        print(f"Loading: {os.path.basename(json_file)}")
+        with open(json_file, "r", encoding="utf-8") as f:
+            chunks_data = json.load(f)
+            all_chunks_data.extend(chunks_data)
 
     source_retrieved_texts = []
     retrieved_elements_full = []
 
-    for chu in chunks_data:
+    for chu in all_chunks_data:
         if chu["element_id"] in retrieved_chunk_ids:
             if chu.get("type") == "TableElement":
                 text = (
-                    f"[Source ID: {chu['elements']['element_id']}]\n"
+                    f"[Source ID: {chu['element_id']}]\n"
                     f"CONTENT:\n{chu['text']}\n"
                     f"HTML:\n{chu['table_text_as_html']}\n\n"
                 )
                 source_retrieved_texts.append(text)
+                retrieved_elements_full.append(chu)
             else:
                 for element in chu.get("elements", []):
                     text = (
@@ -236,15 +257,16 @@ def retrieve_documents_and_prepare_inputs(query, expanding_retriever, chunks_pat
 
     return source_retrieved_texts, retrieved_elements_full
 
-def full_rag_pipeline(query, expanding_retriever, chunks_path="./chunks.json", model="gemini-2.0-flash", clean_all_citations=False):
+def full_rag_pipeline(query, expanding_retriever, chunks_directory="./data/", model="gemini-2.0-flash-exp", clean_all_citations=False):
     """
     Full RAG pipeline from query to RAG response + extracted sources.
 
     Args:
         query (str): The user's query.
         expanding_retriever: The retriever object.
-        chunks_path (str): Path to the chunks.json.
+        chunks_directory (str): Path to the directory containing JSON files.
         model (str): Gemini model.
+        clean_all_citations (bool): Whether to remove all citations from response.
 
     Returns:
         dict: {
@@ -253,12 +275,11 @@ def full_rag_pipeline(query, expanding_retriever, chunks_path="./chunks.json", m
             "answer_language": str
         }
     """
-    source_texts, retrieved_elements = retrieve_documents_and_prepare_inputs(query, expanding_retriever, chunks_path)
+    source_texts, retrieved_elements = retrieve_documents_and_prepare_inputs(query, expanding_retriever, chunks_directory)
 
     # Step 1: RAG
     response_text = generate_rag_response(query, source_texts, model=model)
 
-
     # Step 2: Extract cited sources
     unique_ids = extract_source_ids(response_text)