update the handle of the complaxe query

prepare_env.py

@@ -1,16 +1,22 @@
 # -*- coding: utf-8 -*-
 """
-Enhanced RAG Pipeline for vaccine assistant - Fixed version with max iterations control
-Handles agent creation and question answering with sequential citation numbering
+Environment preparation script for vaccine assistant - Improved version
+Creates vector stores and retrieval tools with better descriptions for efficient agent routing
 """
 
+import os
 import json
 import re
-from llama_index.core import PromptTemplate
-from llama_index.core.agent import ReActAgent
-from llama_index.llms.google_genai import GoogleGenAI
-from langdetect import detect
-import os
+import nest_asyncio
+from typing import List
+from langchain_community.vectorstores import Chroma
+from langchain_core.documents import Document
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.retrievers import BM25Retriever, EnsembleRetriever
+from langchain.retrievers.multi_query import MultiQueryRetriever
+from langchain_google_genai import ChatGoogleGenerativeAI
+from llama_index.core.tools import FunctionTool
+from llama_index.core.schema import TextNode
 
 
 def extract_source_ids(response_text):
@@ -47,13 +53,8 @@ def extract_source_ids(response_text):
         ids = [id_str.strip() for id_str in citation.split(',')]
         all_ids.extend(ids)
     
-    # Get unique source IDs while preserving order
-    seen = set()
-    source_ids = []
-    for id_str in all_ids:
-        if id_str not in seen:
-            seen.add(id_str)
-            source_ids.append(id_str)
+    # Get unique source IDs
+    source_ids = list(set(all_ids))
     
     if not source_ids:
         print("Warning: No valid source IDs found after filtering.")
@@ -62,301 +63,321 @@ def extract_source_ids(response_text):
     return source_ids
 
 
-def convert_citations_to_sequential(response_text, source_id_to_number_map):
-    """
-    Convert source IDs in response text to sequential numbers.
+def setup_models():
+    """Initialize embedding model and LLM"""
+    # Initialize embedding model
+    embedding_function = HuggingFaceEmbeddings(
+        model_name="intfloat/multilingual-e5-base"
+    )
     
-    Args:
-        response_text (str): The response text with source ID citations
-        source_id_to_number_map (dict): Mapping from source IDs to sequential numbers
-        
-    Returns:
-        str: Response text with sequential number citations
-    """
-    def replace_citation(match):
-        citation_content = match.group(1)
-        # Handle multiple IDs in one citation (comma-separated)
-        ids = [id_str.strip() for id_str in citation_content.split(',')]
-        
-        # Convert each ID to its sequential number
-        numbers = []
-        for id_str in ids:
-            if id_str in source_id_to_number_map:
-                numbers.append(str(source_id_to_number_map[id_str]))
-        
-        # Return the formatted citation with sequential numbers
-        if len(numbers) == 1:
-            return f"[{numbers[0]}]"
-        elif len(numbers) > 1:
-            return f"[{','.join(numbers)}]"
-        else:
-            return match.group(0)  # Return original if no mapping found
+    # Initialize LLM with better parameters for focused responses
+    genai_api_key = os.getenv('GOOGLE_API_KEY')
+    llm = ChatGoogleGenerativeAI(
+        model="gemini-2.0-flash",
+        google_api_key=genai_api_key,
+        temperature=0.1  # Lower temperature for more focused responses
+    )
     
-    # Replace all citations in the text
-    sequential_response = re.sub(r'\[([^\[\]]+)\]', replace_citation, response_text)
-    return sequential_response
+    return embedding_function, llm
 
 
-def create_safe_custom_prompt(tools, llm):
-    """Create a safe version that won't have formatting conflicts"""
+def create_vectorstore_from_json(json_path: str, collection_name: str, embedding_function):
+    """Create vector store from JSON chunks"""
+    # Load the chunks.json
+    with open(json_path, "r", encoding="utf-8") as f:
+        chunks_data = json.load(f)
     
-    custom_instructions = """
-## MEDICAL ASSISTANT ROLE
-You are a helpful and knowledgeable AI-powered vaccine assistant designed to support doctors in clinical decision-making.
-You provide evidence-based guidance using only information from official vaccine medical documents.
-Answer the doctor's question accurately and concisely using only the provided information.
-
-## CRITICAL RULES FOR EFFICIENCY
-
-### Tool Usage Strategy
-1. **MAXIMUM 3 TOOL CALLS**: You must provide a complete answer within 3 tool calls maximum.
-2. **Smart Tool Selection**: Choose the most relevant tool first based on the question topic.
-3. **Comparative Questions**: For questions comparing documents/protocols:
-   - First tool call: Get information from primary source (e.g., Algerian guide)
-   - Second tool call: Get information from secondary source (e.g., WHO document)
-   - Third tool call: Only if absolutely necessary for missing details
-4. **Stop Early**: If you have sufficient information after 1-2 tool calls, provide your answer immediately.
-
-### Citation and Sourcing
-1. For each fact in your response, include an inline citation in the format [Source] immediately following the information, e.g., [e795ebd28318886c0b1a5395ac30ad90].
-2. Do NOT use 'Source:' in the citation format; use only the Source in square brackets.
-3. If a fact is supported by multiple sources, use adjacent citations: [source1][source2]
-4. Use ONLY the provided information and never include facts from your general knowledge.
-
-### Content Formatting
-1. When rendering tables:
-   - Convert HTML tables into clean Markdown format
-   - Preserve all original headers and data rows exactly
-   - Include the citation in the table caption, e.g., 'Table: Vaccination Schedule [Source]'
-2. For lists, maintain the original bullet points/numbering and include citations.
-3. Present information concisely but ensure clinical accuracy is never compromised.
-
-### Answer Completeness Guidelines
-- If you find relevant information from 1-2 sources, synthesize and provide a complete answer
-- Don't keep searching for more sources unless critical information is missing
-- For comparative questions, clearly structure your answer with sections for each source
-- If information is not available in the documents, clearly state this limitation
-
----
-
-"""
-    
-    # Get the exact original template first
-    temp_agent = ReActAgent.from_tools(tools, llm=llm, verbose=False)
-    original_prompts = temp_agent.get_prompts()
-    original_template = original_prompts["agent_worker:system_prompt"].template
-
-    # Add instructions at the very beginning
-    safe_template = f"{custom_instructions}{original_template}"
-
-    # Create new prompt with same metadata as original
-    original_prompt = original_prompts["agent_worker:system_prompt"]
+    documents = []
+    for element in chunks_data:
+        text = element["text"]
+        metadata = {
+            "language": "fra",
+            "source": element["filename"],
+            "filetype": element["filetype"],
+            "element_id": element["element_id"]
+        }
 
-    try:
-        new_prompt = PromptTemplate(
-            template=safe_template,
-            template_vars=original_prompt.template_vars,
-            metadata=original_prompt.metadata if hasattr(original_prompt, 'metadata') else None
-        )
-        return new_prompt
-    except:
-        # Even safer fallback
-        return PromptTemplate(template=safe_template)
+        if "TableElement" == element["type"]:
+            metadata["table_text_as_html"] = element["table_text_as_html"]
 
+        doc = Document(page_content=text, metadata=metadata)
+        documents.append(doc)
 
-def create_agent(tools, llm):
-    """Create the ReAct agent with custom prompt and controlled max iterations"""
-    
-    # Create agent with controlled max iterations (reduced from default 10 to 5)
-    agent = ReActAgent.from_tools(
-        tools,
-        llm=llm,
-        verbose=True,
-        max_iterations=5,  # Reduced max iterations
+    # Create vector store
+    vectorstore = Chroma.from_documents(
+        documents=documents,
+        embedding=embedding_function,
+        collection_name=collection_name,
+        persist_directory="chroma_db_multilingual"
     )
+    return vectorstore, documents
 
-    # Create and apply safe custom prompt
-    try:
-        safe_custom_prompt = create_safe_custom_prompt(tools, llm)
-        agent.update_prompts({"agent_worker:system_prompt": safe_custom_prompt})
-        print("✅ Successfully updated with safe custom prompt and max_iterations=5")
-    except Exception as e:
-        print(f"❌ Safe prompt update failed: {e}")
-        print("⚠️  Using original agent without modifications")
-
-    return agent
 
-
-def initialize_rag_pipeline(tools):
-    """Initialize the RAG pipeline with tools"""
+def create_retriever(vectorstore, docs, llm):
+    """Create ensemble retriever with vector and BM25 search"""
+    # Vector retriever
+    vector_retriever = vectorstore.as_retriever(
+        search_type="similarity",
+        search_kwargs={"k": 4}  # Reduced from 6 to 4 for efficiency
+    )
     
-    # Initialize LlamaIndex LLM with specific parameters to improve efficiency
-    llama_index_llm = GoogleGenAI(
-        model="models/gemini-2.0-flash",
-        api_key=os.getenv('GOOGLE_API_KEY'),
-        temperature=0.1,  # Lower temperature for more focused responses
+    # BM25 retriever
+    bm25_retriever = BM25Retriever.from_documents(docs)
+    bm25_retriever.k = 2
+    
+    # Ensemble retriever
+    ensemble_retriever = EnsembleRetriever(
+        retrievers=[vector_retriever, bm25_retriever],
+        weights=[0.5, 0.5]
     )
     
-    # Create agent
-    agent = create_agent(tools, llama_index_llm)
+    # Multi-query expanding retriever (with reduced complexity for efficiency)
+    expanding_retriever = MultiQueryRetriever.from_llm(
+        retriever=ensemble_retriever,
+        llm=llm
+    )
     
-    return agent
+    return expanding_retriever
+
+
+def convert_chromadb_to_llamaindex_nodes(chromadb_documents: List) -> List[TextNode]:
+    """Convert ChromaDB Document objects to LlamaIndex TextNode objects"""
+    nodes = []
+    for i, doc in enumerate(chromadb_documents):
+        try:
+            text = doc.page_content
+            metadata = doc.metadata.copy()
+            element_id = metadata.get("element_id", f"doc_{i}")
+            source = metadata.get("source", "unknown")
+            node_id = f"{source}_{element_id}"
+
+            node = TextNode(
+                text=text,
+                metadata=metadata,
+                id_=node_id
+            )
+            nodes.append(node)
+        except Exception as e:
+            continue
+    return nodes
+
+
+def section_tool_wrapper(retriever, section_path_chunks, query):
+    """Generic section tool wrapper with improved efficiency"""
+    try:
+        retrieved_docs = retriever.get_relevant_documents(query)
+        nodes_from_retrieved_docs = convert_chromadb_to_llamaindex_nodes(retrieved_docs)
 
+        if not nodes_from_retrieved_docs:
+            return "No relevant documents found for the query."
 
-def process_question(agent, question: str) -> str:
-    """Process a question through the RAG pipeline with timeout handling"""
-    try:
-        # Add timeout/retry logic
-        response = agent.chat(question)
-        return response.response
-    except Exception as e:
-        error_msg = str(e)
-        print(f"Error processing question: {error_msg}")
-        
-        # Handle specific "max iterations" error
-        if "max iterations" in error_msg.lower() or "reached max" in error_msg.lower():
-            return ("I apologize, but I was unable to find a complete answer within the allowed search attempts. "
-                   "This might be because the specific comparison you're asking about requires information "
-                   "that spans multiple sections of the documents. Could you please rephrase your question "
-                   "to be more specific about which aspect of the difference you're most interested in?")
+        chunk_ids = [node.metadata['element_id'] for node in retrieved_docs]
+        with open(section_path_chunks, "r", encoding="utf-8") as f:
+            chunks_data = json.load(f)
+
+        chunks_unique = [node for node in chunks_data if node.get('element_id', 'Unknown') in chunk_ids]
+        combined_text = []
         
-        return f"Error processing your question: {error_msg}"
+        # Limit the number of chunks to avoid overwhelming the context
+        max_chunks = 8  # Reasonable limit
+        for chu in chunks_unique[:max_chunks]:
+            if "TableElement" == chu["type"]:
+                text = f"[{chu['element_id']}]\n CONTENT: \n{chu['text']}\n HTML: \n {chu['table_text_as_html']}  \n\n"
+                combined_text.append(text)
+            else:
+                for element in chu["elements"]:
+                    text = f"[{element['element_id']}]\n CONTENT: \n{element['text']} \n\n"
+                    combined_text.append(text)
+
+        result = "\n---\n".join(combined_text)
+        print(f"Retrieved {len(nodes_from_retrieved_docs)} documents for query: {query[:50]}...")
+        return result
+    except Exception as e:
+        print(f"Error in section tool: {e}")
+        return f"Error retrieving documents: {str(e)}"
 
 
-def aswer_language_detection(response_text: str) -> str:
-    """
-    Detect the language of the response text.
+def create_section_tools(embedding_function, llm):
+    """Create all section-specific retrieval tools with improved descriptions"""
     
-    Args:
-        response_text (str): The response text to analyze.
-        
-    Returns:
-        str: Detected language code (e.g., 'en', 'fr', etc.)
-    """
-    try:
-        # Detect the language of the first 5 words of the response
-        first_line = " ".join(response_text.split()[:5])
-        first_line = re.sub(r'\[.*?\]', '', first_line)  # Remove citations
-        answer_language = detect(first_line)
-        if answer_language not in ['en', 'ar', 'fr']:
-            answer_language = 'en'
-    except:
-        answer_language = 'en'
+    # Define section paths
+    section_paths = {
+        'one': 'section_one_chunks.json',
+        'two': 'section_two_chunks.json', 
+        'three': 'section_three_chunks.json',
+        'four': 'section_four_chunks.json',
+        'five': 'section_five_chunks.json',
+        'six': 'section_six_chunks.json',
+        'seven': 'section_seven_chunks.json',
+        'eight': 'section_eight_chunks.json',
+        'nine': 'section_nine_chunks.json',
+        'ten': 'section_ten_chunks.json'
+    }
+    
+    # Create retrievers for each section
+    section_retrievers = {}
+    for section, path in section_paths.items():
+        if os.path.exists(f'./data/{path}'):
+            vstore, docs = create_vectorstore_from_json(f'./data/{path}', f"Guide_2023_{section}", embedding_function)
+            section_retrievers[section] = create_retriever(vstore, docs, llm)
     
-    return answer_language
+    # Create main guide retriever
+    guide_path = './data/Guide-pratique-de-mise-en-oeuvre-du-calendrier-national-de-vaccination-2023.json'
+    if os.path.exists(guide_path):
+        guide_vstore, guide_docs = create_vectorstore_from_json(guide_path, "Guide_2023_multilingual", embedding_function)
+        guide_retriever = create_retriever(guide_vstore, guide_docs, llm)
+    else:
+        guide_retriever = None
+
+    # Primary + Secondary Document Paths
+    immunization_path = './data/Immunization_in_Practice_WHO_eng_2015.json'
+
+    # WHO Immunization in Practice Tool
+    if os.path.exists(immunization_path):
+        immunization_vstore, immunization_docs = create_vectorstore_from_json(
+            immunization_path,
+            "Immunization_in_Practice_WHO_eng_2015",
+            embedding_function
+        )
+        immunization_retriever = create_retriever(immunization_vstore, immunization_docs, llm)
+    else:
+        immunization_retriever = None
 
+    # Tool Functions with Improved Efficiency Focus
 
-def process_question_with_sequential_citations(agent, question: str, chunks_directory="./data/") -> dict:
-    """
-    Process a question through the RAG pipeline and return response with sequential citation numbers.
-    Enhanced with better error handling for max iterations.
-    
-    Args:
-        agent: The initialized RAG agent
-        question (str): The user's question
-        chunks_directory (str): Path to the directory containing JSON files
+    def guide_retrieval_tool(query: str) -> str:
+        """
+        **PRIMARY TOOL - USE FIRST FOR MOST QUESTIONS**
         
-    Returns:
-        dict: {
-            "response": str,  # Response with sequential citation numbers [1], [2], etc.
-            "cited_elements_json": str,  # JSON array of cited elements in order
-            "unique_ids": list,  # Original source IDs in order
-            "citation_mapping": dict  # Mapping from source ID to citation number
-        }
-    """
-    try:
-        # Get the response from the agent with improved error handling
-        response = agent.chat(question)
-        response_text = response.response
+        Comprehensive search across the entire Algerian National Vaccination Guide (2023).
         
-        # Check if the response indicates max iterations was reached
-        if "max iterations" in response_text.lower() or len(response_text.strip()) == 0:
-            # Provide a more helpful fallback response
-            response_text = ("I apologize, but I encountered difficulties processing your comparative question "
-                           "within the allowed search attempts. For questions comparing different protocols "
-                           "or documents, please try asking about each aspect separately. For example, "
-                           "first ask about the Algerian definition of Diphtheria, then ask about the WHO definition.")
+        **When to use this tool:**
+        - General vaccination questions
+        - Disease definitions and descriptions  
+        - Vaccine schedules and protocols
+        - Comparative questions needing Algerian perspective
+        - Any question about Algeria's vaccination program
         
-        # Extract source IDs from the response (preserving order)
-        unique_ids = extract_source_ids(response_text)
+        **Keywords that indicate this tool:** Algeria, Algerian, national, calendrier, vaccination, PEV, diseases (diphteria, polio, measles, etc.)
         
-        # Create mapping from source ID to sequential number
-        source_id_to_number = {source_id: i + 1 for i, source_id in enumerate(unique_ids)}
+        Args:
+            query (str): Any vaccination-related question about Algeria's national program
+            
+        Returns:
+            str: Comprehensive information from the Algerian guide with citations
+        """
+        if not guide_retriever:
+            return "Guide retriever not available"
+        return section_tool_wrapper(guide_retriever, guide_path, query)
+
+    def immunization_tool(query: str) -> str:
+        """
+        **SECONDARY TOOL - USE FOR WHO/INTERNATIONAL PERSPECTIVE**
         
-        # Convert citations to sequential numbers
-        sequential_response = convert_citations_to_sequential(response_text, source_id_to_number)
+        WHO Immunization in Practice 2015 - Global best practices and international standards.
         
-        # Load all chunks data to find cited elements
-        all_chunks_data = []
-        min_chunks_files = ["Guide-pratique-de-mise-en-oeuvre-du-calendrier-national-de-vaccination-2023.json", 
-                           "Immunization_in_Practice_WHO_eng_2015.json"]
+        **When to use this tool:**
+        - Questions specifically asking about WHO recommendations
+        - International/global immunization practices
+        - Comparative questions needing WHO perspective
+        - Technical immunization procedures and best practices
         
-        for json_file in min_chunks_files:
-            json_path = os.path.join(chunks_directory, json_file)
-            try:
-                with open(json_path, "r", encoding="utf-8") as f:
-                    chunks_data = json.load(f)
-                    all_chunks_data.extend(chunks_data)
-            except Exception as e:
-                print(f"Warning: Could not load {json_file}: {e}")
+        **Keywords that indicate this tool:** WHO, international, global, best practices, standards
         
-        # Get cited elements in the same order as the sequential citations
-        cited_elements_ordered = []
-        for source_id in unique_ids:  # This preserves the order
-            for element in all_chunks_data:
-                if element.get("type") == 'TableElement':
-                    if element.get("element_id") == source_id:
-                        cited_elements_ordered.append(element)
-                        break
-                else:
-                    if "elements" in element:
-                        for nested_element in element["elements"]:
-                            if nested_element.get("element_id") == source_id:
-                                cited_elements_ordered.append(nested_element)
-                                break
-                        else:
-                            continue
-                        break
+        Args:
+            query (str): Question about international immunization practices or WHO recommendations
+            
+        Returns:
+            str: WHO guidance and international best practices with citations
+        """
+        if not immunization_retriever:
+            return "Immunization in Practice retriever not available"
+        return section_tool_wrapper(immunization_retriever, immunization_path, query)
+
+    # Section-Specific Tools (USE ONLY IF QUESTION IS VERY SPECIFIC TO THE SECTION)
+
+    def section_two_tool(query: str) -> str:
+        """
+        **DISEASE-SPECIFIC TOOL**
         
-        # Convert to JSON
-        cited_elements_json = json.dumps(cited_elements_ordered, ensure_ascii=False, indent=2)
-        answer_language = aswer_language_detection(response_text)
+        Section 2: Vaccine-preventable diseases - definitions, symptoms, transmission, complications.
         
-        return {
-            "response": sequential_response,
-            "cited_elements_json": cited_elements_json,
-            "unique_ids": unique_ids,
-            "citation_mapping": source_id_to_number,
-            "answer_language": answer_language
-        }
+        **Use ONLY for specific disease definition questions like:**
+        - "What is diphtheria?" 
+        - "Define measles according to Algerian protocol"
+        - "Symptoms of polio"
         
-    except Exception as e:
-        error_msg = str(e)
-        print(f"Error processing question: {error_msg}")
+        **Keywords:** definition, symptoms, transmission, complications, disease characteristics
         
-        # Create appropriate fallback response based on error type
-        if "max iterations" in error_msg.lower() or "reached max" in error_msg.lower():
-            fallback_response = ("I apologize, but I was unable to complete the comparison within the allowed search attempts. "
-                               "For complex comparative questions like yours about the differences between Algerian and WHO "
-                               "definitions of Diphtheria, please try asking about each source separately: \n\n"
-                               "1. First ask: 'What is the definition of Diphtheria in the Algerian vaccination guide?'\n"
-                               "2. Then ask: 'What is the definition of Diphtheria in the WHO document?'\n\n"
-                               "This will help me provide you with more focused and complete information.")
-        else:
-            fallback_response = f"I encountered an error while processing your question: {error_msg}"
+        Args:
+            query (str): Specific question about disease definitions or characteristics
+            
+        Returns:
+            str: Disease-specific medical information with citations
+        """
+        if 'two' not in section_retrievers:
+            return "Section 2 retriever not available"
+        return section_tool_wrapper(section_retrievers['two'], f'./data/{section_paths["two"]}', query)
+
+    def section_three_tool(query: str) -> str:
+        """
+        **VACCINE-SPECIFIC TOOL**
         
-        return {
-            "response": fallback_response,
-            "cited_elements_json": "[]",
-            "unique_ids": [],
-            "citation_mapping": {},
-            "answer_language": "en"
-        }
+        Section 3: Vaccine details - types, composition, administration methods.
+        
+        **Use ONLY for specific vaccine technical questions like:**
+        - "What type of vaccine is used for diphtheria?"
+        - "How is the MMR vaccine administered?"
+        - "Vaccine composition and dosage"
+        
+        **Keywords:** vaccine type, composition, administration, dosage, technical details
+        
+        Args:
+            query (str): Technical question about specific vaccines
+            
+        Returns:
+            str: Technical vaccine information with citations
+        """
+        if 'three' not in section_retrievers:
+            return "Section 3 retriever not available"
+        return section_tool_wrapper(section_retrievers['three'], f'./data/{section_paths["three"]}', query)
+
+    # Create FunctionTool objects with focused selection
+    tools = [
+        # Primary tools - most commonly used
+        FunctionTool.from_defaults(
+            name="algerian_guide_search",
+            fn=guide_retrieval_tool,
+            description="PRIMARY TOOL: Search the complete Algerian National Vaccination Guide for any vaccination-related question"
+        ),
+        FunctionTool.from_defaults(
+            name="who_immunization_search", 
+            fn=immunization_tool,
+            description="SECONDARY TOOL: Search WHO Immunization in Practice for international standards and WHO recommendations"
+        ),
+        # Specialized tools - use only when very specific
+        FunctionTool.from_defaults(
+            name="disease_definitions_search",
+            fn=section_two_tool,
+            description="SPECIALIZED: Search for specific disease definitions, symptoms, and characteristics"
+        ),
+        FunctionTool.from_defaults(
+            name="vaccine_technical_search",
+            fn=section_three_tool, 
+            description="SPECIALIZED: Search for technical vaccine details, composition, and administration methods"
+        ),
+    ]
+    
+    return tools
 
 
-def process_question_with_citations(agent, question: str, chunks_directory="./data/") -> dict:
-    """
-    Legacy function - maintained for backward compatibility.
-    Now calls the new sequential citation function.
-    """
-    return process_question_with_sequential_citations(agent, question, chunks_directory)
+def prepare_environment():
+    """Main function to prepare the environment and return tools"""
+    print("Setting up models...")
+    embedding_function, llm = setup_models()
+    
+    print("Creating section tools...")
+    tools = create_section_tools(embedding_function, llm)
+    
+    print("Environment prepared successfully!")
+    return tools, llm

rag_pipeline.py

@@ -1,22 +1,16 @@
 # -*- coding: utf-8 -*-
 """
-Environment preparation script for vaccine assistant - Improved version
-Creates vector stores and retrieval tools with better descriptions for efficient agent routing
+Enhanced RAG Pipeline for vaccine assistant - Fixed version with max iterations control
+Handles agent creation and question answering with sequential citation numbering
 """
 
-import os
 import json
 import re
-import nest_asyncio
-from typing import List
-from langchain_community.vectorstores import Chroma
-from langchain_core.documents import Document
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.retrievers import BM25Retriever, EnsembleRetriever
-from langchain.retrievers.multi_query import MultiQueryRetriever
-from langchain_google_genai import ChatGoogleGenerativeAI
-from llama_index.core.tools import FunctionTool
-from llama_index.core.schema import TextNode
+from llama_index.core import PromptTemplate
+from llama_index.core.agent import ReActAgent
+from llama_index.llms.google_genai import GoogleGenAI
+from langdetect import detect
+import os
 
 
 def extract_source_ids(response_text):
@@ -53,8 +47,13 @@ def extract_source_ids(response_text):
         ids = [id_str.strip() for id_str in citation.split(',')]
         all_ids.extend(ids)
     
-    # Get unique source IDs
-    source_ids = list(set(all_ids))
+    # Get unique source IDs while preserving order
+    seen = set()
+    source_ids = []
+    for id_str in all_ids:
+        if id_str not in seen:
+            seen.add(id_str)
+            source_ids.append(id_str)
     
     if not source_ids:
         print("Warning: No valid source IDs found after filtering.")
@@ -63,321 +62,301 @@ def extract_source_ids(response_text):
     return source_ids
 
 
-def setup_models():
-    """Initialize embedding model and LLM"""
-    # Initialize embedding model
-    embedding_function = HuggingFaceEmbeddings(
-        model_name="intfloat/multilingual-e5-base"
-    )
+def convert_citations_to_sequential(response_text, source_id_to_number_map):
+    """
+    Convert source IDs in response text to sequential numbers.
     
-    # Initialize LLM with better parameters for focused responses
-    genai_api_key = os.getenv('GOOGLE_API_KEY')
-    llm = ChatGoogleGenerativeAI(
-        model="gemini-2.0-flash",
-        google_api_key=genai_api_key,
-        temperature=0.1  # Lower temperature for more focused responses
-    )
+    Args:
+        response_text (str): The response text with source ID citations
+        source_id_to_number_map (dict): Mapping from source IDs to sequential numbers
+        
+    Returns:
+        str: Response text with sequential number citations
+    """
+    def replace_citation(match):
+        citation_content = match.group(1)
+        # Handle multiple IDs in one citation (comma-separated)
+        ids = [id_str.strip() for id_str in citation_content.split(',')]
+        
+        # Convert each ID to its sequential number
+        numbers = []
+        for id_str in ids:
+            if id_str in source_id_to_number_map:
+                numbers.append(str(source_id_to_number_map[id_str]))
+        
+        # Return the formatted citation with sequential numbers
+        if len(numbers) == 1:
+            return f"[{numbers[0]}]"
+        elif len(numbers) > 1:
+            return f"[{','.join(numbers)}]"
+        else:
+            return match.group(0)  # Return original if no mapping found
     
-    return embedding_function, llm
+    # Replace all citations in the text
+    sequential_response = re.sub(r'\[([^\[\]]+)\]', replace_citation, response_text)
+    return sequential_response
 
 
-def create_vectorstore_from_json(json_path: str, collection_name: str, embedding_function):
-    """Create vector store from JSON chunks"""
-    # Load the chunks.json
-    with open(json_path, "r", encoding="utf-8") as f:
-        chunks_data = json.load(f)
+def create_safe_custom_prompt(tools, llm):
+    """Create a safe version that won't have formatting conflicts"""
     
-    documents = []
-    for element in chunks_data:
-        text = element["text"]
-        metadata = {
-            "language": "fra",
-            "source": element["filename"],
-            "filetype": element["filetype"],
-            "element_id": element["element_id"]
-        }
+    custom_instructions = """
+## MEDICAL ASSISTANT ROLE
+You are a helpful and knowledgeable AI-powered vaccine assistant designed to support doctors in clinical decision-making.
+You provide evidence-based guidance using only information from official vaccine medical documents.
+Answer the doctor's question accurately and concisely using only the provided information.
+
+## CRITICAL RULES FOR EFFICIENCY
+
+### Tool Usage Strategy
+1. **MAXIMUM 3 TOOL CALLS**: You must provide a complete answer within 3 tool calls maximum.
+2. **Smart Tool Selection**: Choose the most relevant tool first based on the question topic.
+3. **Comparative Questions**: For questions comparing documents/protocols:
+   - First tool call: Get information from primary source (e.g., Algerian guide)
+   - Second tool call: Get information from secondary source (e.g., WHO document)
+   - Third tool call: Only if absolutely necessary for missing details
+4. **Stop Early**: If you have sufficient information after 1-2 tool calls, provide your answer immediately.
+
+### Citation and Sourcing
+1. For each fact in your response, include an inline citation in the format [Source] immediately following the information, e.g., [e795ebd28318886c0b1a5395ac30ad90].
+2. Do NOT use 'Source:' in the citation format; use only the Source in square brackets.
+3. If a fact is supported by multiple sources, use adjacent citations: [source1][source2]
+4. Use ONLY the provided information and never include facts from your general knowledge.
+
+### Content Formatting
+1. When rendering tables:
+   - Convert HTML tables into clean Markdown format
+   - Preserve all original headers and data rows exactly
+   - Include the citation in the table caption, e.g., 'Table: Vaccination Schedule [Source]'
+2. For lists, maintain the original bullet points/numbering and include citations.
+3. Present information concisely but ensure clinical accuracy is never compromised.
+
+### Answer Completeness Guidelines
+- If you find relevant information from 1-2 sources, synthesize and provide a complete answer
+- Don't keep searching for more sources unless critical information is missing
+- For comparative questions, clearly structure your answer with sections for each source
+- If information is not available in the documents, clearly state this limitation
+
+---
 
-        if "TableElement" == element["type"]:
-            metadata["table_text_as_html"] = element["table_text_as_html"]
-
-        doc = Document(page_content=text, metadata=metadata)
-        documents.append(doc)
+"""
+    
+    # Get the exact original template first
+    temp_agent = ReActAgent.from_tools(tools, llm=llm, verbose=False)
+    original_prompts = temp_agent.get_prompts()
+    original_template = original_prompts["agent_worker:system_prompt"].template
 
-    # Create vector store
-    vectorstore = Chroma.from_documents(
-        documents=documents,
-        embedding=embedding_function,
-        collection_name=collection_name,
-        persist_directory="chroma_db_multilingual"
-    )
-    return vectorstore, documents
+    # Add instructions at the very beginning
+    safe_template = f"{custom_instructions}{original_template}"
 
+    # Create new prompt with same metadata as original
+    original_prompt = original_prompts["agent_worker:system_prompt"]
 
-def create_retriever(vectorstore, docs, llm):
-    """Create ensemble retriever with vector and BM25 search"""
-    # Vector retriever
-    vector_retriever = vectorstore.as_retriever(
-        search_type="similarity",
-        search_kwargs={"k": 4}  # Reduced from 6 to 4 for efficiency
-    )
-    
-    # BM25 retriever
-    bm25_retriever = BM25Retriever.from_documents(docs)
-    bm25_retriever.k = 2
-    
-    # Ensemble retriever
-    ensemble_retriever = EnsembleRetriever(
-        retrievers=[vector_retriever, bm25_retriever],
-        weights=[0.5, 0.5]
-    )
-    
-    # Multi-query expanding retriever (with reduced complexity for efficiency)
-    expanding_retriever = MultiQueryRetriever.from_llm(
-        retriever=ensemble_retriever,
-        llm=llm
-    )
-    
-    return expanding_retriever
-
-
-def convert_chromadb_to_llamaindex_nodes(chromadb_documents: List) -> List[TextNode]:
-    """Convert ChromaDB Document objects to LlamaIndex TextNode objects"""
-    nodes = []
-    for i, doc in enumerate(chromadb_documents):
-        try:
-            text = doc.page_content
-            metadata = doc.metadata.copy()
-            element_id = metadata.get("element_id", f"doc_{i}")
-            source = metadata.get("source", "unknown")
-            node_id = f"{source}_{element_id}"
-
-            node = TextNode(
-                text=text,
-                metadata=metadata,
-                id_=node_id
-            )
-            nodes.append(node)
-        except Exception as e:
-            continue
-    return nodes
-
-
-def section_tool_wrapper(retriever, section_path_chunks, query):
-    """Generic section tool wrapper with improved efficiency"""
     try:
-        retrieved_docs = retriever.get_relevant_documents(query)
-        nodes_from_retrieved_docs = convert_chromadb_to_llamaindex_nodes(retrieved_docs)
+        new_prompt = PromptTemplate(
+            template=safe_template,
+            template_vars=original_prompt.template_vars,
+            metadata=original_prompt.metadata if hasattr(original_prompt, 'metadata') else None
+        )
+        return new_prompt
+    except:
+        # Even safer fallback
+        return PromptTemplate(template=safe_template)
 
-        if not nodes_from_retrieved_docs:
-            return "No relevant documents found for the query."
 
-        chunk_ids = [node.metadata['element_id'] for node in retrieved_docs]
-        with open(section_path_chunks, "r", encoding="utf-8") as f:
-            chunks_data = json.load(f)
+def create_agent(tools, llm):
+    """Create the ReAct agent with custom prompt and controlled max iterations"""
+    
+    # Create agent with controlled max iterations (reduced from default 10 to 5)
+    agent = ReActAgent.from_tools(
+        tools,
+        llm=llm,
+        verbose=True,
+        max_iterations=5,  # Reduced max iterations
+    )
 
-        chunks_unique = [node for node in chunks_data if node.get('element_id', 'Unknown') in chunk_ids]
-        combined_text = []
-        
-        # Limit the number of chunks to avoid overwhelming the context
-        max_chunks = 8  # Reasonable limit
-        for chu in chunks_unique[:max_chunks]:
-            if "TableElement" == chu["type"]:
-                text = f"[{chu['element_id']}]\n CONTENT: \n{chu['text']}\n HTML: \n {chu['table_text_as_html']}  \n\n"
-                combined_text.append(text)
-            else:
-                for element in chu["elements"]:
-                    text = f"[{element['element_id']}]\n CONTENT: \n{element['text']} \n\n"
-                    combined_text.append(text)
-
-        result = "\n---\n".join(combined_text)
-        print(f"Retrieved {len(nodes_from_retrieved_docs)} documents for query: {query[:50]}...")
-        return result
+    # Create and apply safe custom prompt
+    try:
+        safe_custom_prompt = create_safe_custom_prompt(tools, llm)
+        agent.update_prompts({"agent_worker:system_prompt": safe_custom_prompt})
+        print("✅ Successfully updated with safe custom prompt and max_iterations=5")
     except Exception as e:
-        print(f"Error in section tool: {e}")
-        return f"Error retrieving documents: {str(e)}"
+        print(f"❌ Safe prompt update failed: {e}")
+        print("⚠️  Using original agent without modifications")
+
+    return agent
 
 
-def create_section_tools(embedding_function, llm):
-    """Create all section-specific retrieval tools with improved descriptions"""
+def initialize_rag_pipeline(tools):
+    """Initialize the RAG pipeline with tools"""
     
-    # Define section paths
-    section_paths = {
-        'one': 'section_one_chunks.json',
-        'two': 'section_two_chunks.json', 
-        'three': 'section_three_chunks.json',
-        'four': 'section_four_chunks.json',
-        'five': 'section_five_chunks.json',
-        'six': 'section_six_chunks.json',
-        'seven': 'section_seven_chunks.json',
-        'eight': 'section_eight_chunks.json',
-        'nine': 'section_nine_chunks.json',
-        'ten': 'section_ten_chunks.json'
-    }
+    # Initialize LlamaIndex LLM with specific parameters to improve efficiency
+    llama_index_llm = GoogleGenAI(
+        model="models/gemini-2.0-flash",
+        api_key=os.getenv('GOOGLE_API_KEY'),
+        temperature=0.1,  # Lower temperature for more focused responses
+    )
     
-    # Create retrievers for each section
-    section_retrievers = {}
-    for section, path in section_paths.items():
-        if os.path.exists(f'./data/{path}'):
-            vstore, docs = create_vectorstore_from_json(f'./data/{path}', f"Guide_2023_{section}", embedding_function)
-            section_retrievers[section] = create_retriever(vstore, docs, llm)
+    # Create agent
+    agent = create_agent(tools, llama_index_llm)
     
-    # Create main guide retriever
-    guide_path = './data/Guide-pratique-de-mise-en-oeuvre-du-calendrier-national-de-vaccination-2023.json'
-    if os.path.exists(guide_path):
-        guide_vstore, guide_docs = create_vectorstore_from_json(guide_path, "Guide_2023_multilingual", embedding_function)
-        guide_retriever = create_retriever(guide_vstore, guide_docs, llm)
-    else:
-        guide_retriever = None
-
-    # Primary + Secondary Document Paths
-    immunization_path = './data/Immunization_in_Practice_WHO_eng_2015.json'
-
-    # WHO Immunization in Practice Tool
-    if os.path.exists(immunization_path):
-        immunization_vstore, immunization_docs = create_vectorstore_from_json(
-            immunization_path,
-            "Immunization_in_Practice_WHO_eng_2015",
-            embedding_function
-        )
-        immunization_retriever = create_retriever(immunization_vstore, immunization_docs, llm)
-    else:
-        immunization_retriever = None
+    return agent
 
-    # Tool Functions with Improved Efficiency Focus
 
-    def guide_retrieval_tool(query: str) -> str:
-        """
-        **PRIMARY TOOL - USE FIRST FOR MOST QUESTIONS**
+def process_question(agent, question: str) -> str:
+    """Process a question through the RAG pipeline with timeout handling"""
+    try:
+        # Add timeout/retry logic
+        response = agent.chat(question)
+        return response.response
+    except Exception as e:
+        error_msg = str(e)
+        print(f"Error processing question: {error_msg}")
         
-        Comprehensive search across the entire Algerian National Vaccination Guide (2023).
+        # Handle specific "max iterations" error
+        if "max iterations" in error_msg.lower() or "reached max" in error_msg.lower():
+            return ("I apologize, but I was unable to find a complete answer within the allowed search attempts. "
+                   "This might be because the specific comparison you're asking about requires information "
+                   "that spans multiple sections of the documents. Could you please rephrase your question "
+                   "to be more specific about which aspect of the difference you're most interested in?")
         
-        **When to use this tool:**
-        - General vaccination questions
-        - Disease definitions and descriptions  
-        - Vaccine schedules and protocols
-        - Comparative questions needing Algerian perspective
-        - Any question about Algeria's vaccination program
+        return f"Error processing your question: {error_msg}"
+
+
+def aswer_language_detection(response_text: str) -> str:
+    """
+    Detect the language of the response text.
+    
+    Args:
+        response_text (str): The response text to analyze.
         
-        **Keywords that indicate this tool:** Algeria, Algerian, national, calendrier, vaccination, PEV, diseases (diphteria, polio, measles, etc.)
+    Returns:
+        str: Detected language code (e.g., 'en', 'fr', etc.)
+    """
+    try:
+        # Detect the language of the first 5 words of the response
+        first_line = " ".join(response_text.split()[:5])
+        first_line = re.sub(r'\[.*?\]', '', first_line)  # Remove citations
+        answer_language = detect(first_line)
+        if answer_language not in ['en', 'ar', 'fr']:
+            answer_language = 'en'
+    except:
+        answer_language = 'en'
+    
+    return answer_language
+
+
+def process_question_with_sequential_citations(agent, question: str, chunks_directory="./data/") -> dict:
+    """
+    Process a question through the RAG pipeline and return response with sequential citation numbers.
+    Enhanced with better error handling for max iterations.
+    
+    Args:
+        agent: The initialized RAG agent
+        question (str): The user's question
+        chunks_directory (str): Path to the directory containing JSON files
         
-        Args:
-            query (str): Any vaccination-related question about Algeria's national program
-            
-        Returns:
-            str: Comprehensive information from the Algerian guide with citations
-        """
-        if not guide_retriever:
-            return "Guide retriever not available"
-        return section_tool_wrapper(guide_retriever, guide_path, query)
-
-    def immunization_tool(query: str) -> str:
-        """
-        **SECONDARY TOOL - USE FOR WHO/INTERNATIONAL PERSPECTIVE**
+    Returns:
+        dict: {
+            "response": str,  # Response with sequential citation numbers [1], [2], etc.
+            "cited_elements_json": str,  # JSON array of cited elements in order
+            "unique_ids": list,  # Original source IDs in order
+            "citation_mapping": dict  # Mapping from source ID to citation number
+        }
+    """
+    try:
+        # Get the response from the agent with improved error handling
+        response = agent.chat(question)
+        response_text = response.response
         
-        WHO Immunization in Practice 2015 - Global best practices and international standards.
+        # Check if the response indicates max iterations was reached
+        if "max iterations" in response_text.lower() or len(response_text.strip()) == 0:
+            # Provide a more helpful fallback response
+            response_text = ("I apologize, but I encountered difficulties processing your comparative question "
+                           "within the allowed search attempts. For questions comparing different protocols "
+                           "or documents, please try asking about each aspect separately. For example, "
+                           "first ask about the Algerian definition of Diphtheria, then ask about the WHO definition.")
         
-        **When to use this tool:**
-        - Questions specifically asking about WHO recommendations
-        - International/global immunization practices
-        - Comparative questions needing WHO perspective
-        - Technical immunization procedures and best practices
+        # Extract source IDs from the response (preserving order)
+        unique_ids = extract_source_ids(response_text)
         
-        **Keywords that indicate this tool:** WHO, international, global, best practices, standards
+        # Create mapping from source ID to sequential number
+        source_id_to_number = {source_id: i + 1 for i, source_id in enumerate(unique_ids)}
         
-        Args:
-            query (str): Question about international immunization practices or WHO recommendations
-            
-        Returns:
-            str: WHO guidance and international best practices with citations
-        """
-        if not immunization_retriever:
-            return "Immunization in Practice retriever not available"
-        return section_tool_wrapper(immunization_retriever, immunization_path, query)
-
-    # Section-Specific Tools (USE ONLY IF QUESTION IS VERY SPECIFIC TO THE SECTION)
-
-    def section_two_tool(query: str) -> str:
-        """
-        **DISEASE-SPECIFIC TOOL**
+        # Convert citations to sequential numbers
+        sequential_response = convert_citations_to_sequential(response_text, source_id_to_number)
         
-        Section 2: Vaccine-preventable diseases - definitions, symptoms, transmission, complications.
+        # Load all chunks data to find cited elements
+        all_chunks_data = []
+        min_chunks_files = ["Guide-pratique-de-mise-en-oeuvre-du-calendrier-national-de-vaccination-2023.json", 
+                           "Immunization_in_Practice_WHO_eng_2015.json"]
         
-        **Use ONLY for specific disease definition questions like:**
-        - "What is diphtheria?" 
-        - "Define measles according to Algerian protocol"
-        - "Symptoms of polio"
+        for json_file in min_chunks_files:
+            json_path = os.path.join(chunks_directory, json_file)
+            try:
+                with open(json_path, "r", encoding="utf-8") as f:
+                    chunks_data = json.load(f)
+                    all_chunks_data.extend(chunks_data)
+            except Exception as e:
+                print(f"Warning: Could not load {json_file}: {e}")
         
-        **Keywords:** definition, symptoms, transmission, complications, disease characteristics
+        # Get cited elements in the same order as the sequential citations
+        cited_elements_ordered = []
+        for source_id in unique_ids:  # This preserves the order
+            for element in all_chunks_data:
+                if element.get("type") == 'TableElement':
+                    if element.get("element_id") == source_id:
+                        cited_elements_ordered.append(element)
+                        break
+                else:
+                    if "elements" in element:
+                        for nested_element in element["elements"]:
+                            if nested_element.get("element_id") == source_id:
+                                cited_elements_ordered.append(nested_element)
+                                break
+                        else:
+                            continue
+                        break
         
-        Args:
-            query (str): Specific question about disease definitions or characteristics
-            
-        Returns:
-            str: Disease-specific medical information with citations
-        """
-        if 'two' not in section_retrievers:
-            return "Section 2 retriever not available"
-        return section_tool_wrapper(section_retrievers['two'], f'./data/{section_paths["two"]}', query)
-
-    def section_three_tool(query: str) -> str:
-        """
-        **VACCINE-SPECIFIC TOOL**
+        # Convert to JSON
+        cited_elements_json = json.dumps(cited_elements_ordered, ensure_ascii=False, indent=2)
+        answer_language = aswer_language_detection(response_text)
         
-        Section 3: Vaccine details - types, composition, administration methods.
+        return {
+            "response": sequential_response,
+            "cited_elements_json": cited_elements_json,
+            "unique_ids": unique_ids,
+            "citation_mapping": source_id_to_number,
+            "answer_language": answer_language
+        }
         
-        **Use ONLY for specific vaccine technical questions like:**
-        - "What type of vaccine is used for diphtheria?"
-        - "How is the MMR vaccine administered?"
-        - "Vaccine composition and dosage"
+    except Exception as e:
+        error_msg = str(e)
+        print(f"Error processing question: {error_msg}")
         
-        **Keywords:** vaccine type, composition, administration, dosage, technical details
+        # Create appropriate fallback response based on error type
+        if "max iterations" in error_msg.lower() or "reached max" in error_msg.lower():
+            fallback_response = ("I apologize, but I was unable to complete the comparison within the allowed search attempts. "
+                               "For complex comparative questions like yours about the differences between Algerian and WHO "
+                               "definitions of Diphtheria, please try asking about each source separately: \n\n"
+                               "1. First ask: 'What is the definition of Diphtheria in the Algerian vaccination guide?'\n"
+                               "2. Then ask: 'What is the definition of Diphtheria in the WHO document?'\n\n"
+                               "This will help me provide you with more focused and complete information.")
+        else:
+            fallback_response = f"I encountered an error while processing your question: {error_msg}"
         
-        Args:
-            query (str): Technical question about specific vaccines
-            
-        Returns:
-            str: Technical vaccine information with citations
-        """
-        if 'three' not in section_retrievers:
-            return "Section 3 retriever not available"
-        return section_tool_wrapper(section_retrievers['three'], f'./data/{section_paths["three"]}', query)
-
-    # Create FunctionTool objects with focused selection
-    tools = [
-        # Primary tools - most commonly used
-        FunctionTool.from_defaults(
-            name="algerian_guide_search",
-            fn=guide_retrieval_tool,
-            description="PRIMARY TOOL: Search the complete Algerian National Vaccination Guide for any vaccination-related question"
-        ),
-        FunctionTool.from_defaults(
-            name="who_immunization_search", 
-            fn=immunization_tool,
-            description="SECONDARY TOOL: Search WHO Immunization in Practice for international standards and WHO recommendations"
-        ),
-        # Specialized tools - use only when very specific
-        FunctionTool.from_defaults(
-            name="disease_definitions_search",
-            fn=section_two_tool,
-            description="SPECIALIZED: Search for specific disease definitions, symptoms, and characteristics"
-        ),
-        FunctionTool.from_defaults(
-            name="vaccine_technical_search",
-            fn=section_three_tool, 
-            description="SPECIALIZED: Search for technical vaccine details, composition, and administration methods"
-        ),
-    ]
-    
-    return tools
+        return {
+            "response": fallback_response,
+            "cited_elements_json": "[]",
+            "unique_ids": [],
+            "citation_mapping": {},
+            "answer_language": "en"
+        }
 
 
-def prepare_environment():
-    """Main function to prepare the environment and return tools"""
-    print("Setting up models...")
-    embedding_function, llm = setup_models()
-    
-    print("Creating section tools...")
-    tools = create_section_tools(embedding_function, llm)
-    
-    print("Environment prepared successfully!")
-    return tools, llm
+def process_question_with_citations(agent, question: str, chunks_directory="./data/") -> dict:
+    """
+    Legacy function - maintained for backward compatibility.
+    Now calls the new sequential citation function.
+    """
+    return process_question_with_sequential_citations(agent, question, chunks_directory)