fix handleing a complicated queations

prepare_env.py

@@ -1,22 +1,16 @@
 # -*- coding: utf-8 -*-
 """
-Environment preparation script for vaccine assistant
-Creates vector stores and retrieval tools
+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,13 +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))
-    
-    # Filter out any non-UUID-like IDs (if needed)
-    # This is now optional as we're handling various source ID formats
-    # uuid_pattern = r'^[0-9a-f]{8}-?[0-9a-f]{4}-?[0-9a-f]{4}-?[0-9a-f]{4}-?[0-9a-f]{12}$'
-    # source_ids = [source_id for source_id in source_ids if re.match(uuid_pattern, source_id, re.IGNORECASE)]
+    # 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.")
@@ -68,412 +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"
-    )
-    
-    # Initialize LLM
-    genai_api_key = os.getenv('GOOGLE_API_KEY')
-    llm = ChatGoogleGenerativeAI(
-        model="gemini-2.0-flash",
-        google_api_key=genai_api_key
-    )
+def convert_citations_to_sequential(response_text, source_id_to_number_map):
+    """
+    Convert source IDs in response text to sequential numbers.
     
-    return embedding_function, llm
-
-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)
+    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
     
-    documents = []
-    for element in chunks_data:
-        text = element["text"]
-        metadata = {
-            "language": "fra",
-            "source": element["filename"],
-            "filetype": element["filetype"],
-            "element_id": element["element_id"]
-        }
-
-        if "TableElement" == element["type"]:
-            metadata["table_text_as_html"] = element["table_text_as_html"]
+    # Replace all citations in the text
+    sequential_response = re.sub(r'\[([^\[\]]+)\]', replace_citation, response_text)
+    return sequential_response
 
-        doc = Document(page_content=text, metadata=metadata)
-        documents.append(doc)
 
-    # Create vector store
-    vectorstore = Chroma.from_documents(
-        documents=documents,
-        embedding=embedding_function,
-        collection_name=collection_name,
-        persist_directory="chroma_db_multilingual"
-    )
-    return vectorstore, documents
-
-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": 6}
-    )
-    
-    # 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
-    expanding_retriever = MultiQueryRetriever.from_llm(
-        retriever=ensemble_retriever,
-        llm=llm
-    )
+def create_safe_custom_prompt(tools, llm):
+    """Create a safe version that won't have formatting conflicts"""
     
-    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"""
-    try:
-        retrieved_docs = retriever.get_relevant_documents(query)
-        nodes_from_retrieved_docs = convert_chromadb_to_llamaindex_nodes(retrieved_docs)
+    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 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)
-
-        chunks_unique = [node for node in chunks_data if node.get('element_id', 'Unknown') in chunk_ids]
-        combined_text = []
-        
-        for chu in chunks_unique:
-            if "TableElement" == chu["type"]:
-                text = f"[Source: {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"[Source: {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 create_section_tools(embedding_function, llm):
-    """Create all section-specific retrieval 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'
-    }
-    
-    # Create retrievers for each section
-    section_retrievers = {}
-    for section, path in section_paths.items():
-        if os.path.exists(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 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
-    # General-purpose tool (entire Algerian guide)
-    def guide_retrieval_tool(query: str) -> str:
-        """
-        General-purpose retrieval tool for the entire Algerian National Vaccination Guide (2023).
-
-        Use this tool when a query spans multiple sections or cannot be routed confidently to a specific tool.
-        This is the fallback and all-encompassing tool to retrieve any vaccination-related information 
-        from the national guide.
-
-        Secondary source: The WHO Immunization Guide can be queried separately via `immunization_tool`.
-
-        Args:
-            query (str): A general or complex question related to vaccination policy, schedules, or practice.
-
-        Returns:
-            str: Synthesized response based on the full Algerian guide.
-        """
-        if not guide_retriever:
-            return "Guide retriever not available"
-        return section_tool_wrapper(guide_retriever, guide_path, query)
-
-    # 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
-
-    def immunization_tool(query: str) -> str:
-        """
-        WHO Immunization in Practice 2015 retrieval tool.
-
-        Use this tool to provide global best practices and operational guidance on immunization,
-        especially when context or clarification is needed beyond the Algerian national guide.
-        This can serve as a secondary source for training, logistics, and procedural reference.
-
-        Args:
-            query (str): A question related to immunization practice in general.
-
-        Returns:
-            str: Retrieved guidance from the WHO Immunization in Practice manual (2015).
-        """
-        if not immunization_retriever:
-            return "Immunization in Practice retriever not available"
-        return section_tool_wrapper(immunization_retriever, immunization_path, query)
-
-
-
-    # Section-Specific Tools (Primary: Algerian National Vaccination Guide)
-
-    def section_one_tool(query: str) -> str:
-        """
-        Section 1: Programme Élargi de Vaccination (PEV)
-
-        Use this tool to retrieve information about the Algerian immunization program:
-        its objectives, historical background, strengths and weaknesses, and justification
-        for calendar updates.
-
-        Primary source: Algerian National Vaccination Guide, Section 1.
-        Secondary source for operational benchmarks: WHO Immunization in Practice (optional).
-
-        Args:
-            query (str): A question about Algeria’s national immunization strategy.
-
-        Returns:
-            str: Relevant content from Section 1 of the guide.
-        """
-        return section_tool_wrapper(section_retrievers['one'], section_paths['one'], query)
-
-
-    def section_two_tool(query: str) -> str:
-        """
-        Section 2: Maladies Ciblées par la Vaccination
-
-        Use this tool for questions about the diseases targeted by the national vaccination calendar:
-        symptoms, transmission, complications, and prevention strategies.
-
-        Primary source: Algerian National Guide, Section 2.
-        Secondary source: WHO guide may support contextual insights.
-
-        Args:
-            query (str): A question about a vaccine-preventable disease (e.g. polio, rougeole).
-
-        Returns:
-            str: Disease-specific guidance from Section 2.
-        """
-        return section_tool_wrapper(section_retrievers['two'], section_paths['two'], query)
-
-
-    def section_three_tool(query: str) -> str:
-        """
-        Section 3: Vaccins du Calendrier
-
-        Use this tool to retrieve technical and procedural information about the vaccines used in the calendar:
-        names, contents, administration method, and dosing details.
-
-        Args:
-            query (str): A question about a specific vaccine's type or method of use.
-
-        Returns:
-            str: Vaccine information from Section 3.
-        """
-        return section_tool_wrapper(section_retrievers['three'], section_paths['three'], query)
-
-
-    def section_four_tool(query: str) -> str:
-        """
-        Section 4: Rattrapage Vaccinal
-
-        Use this tool to determine catch-up strategies for children who missed or delayed one or more doses.
-        It provides age-adjusted rescheduling rules and justifications.
-
-        Args:
-            query (str): A question about how to manage missed vaccinations.
-
-        Returns:
-            str: Catch-up guidelines from Section 4.
-        """
-        return section_tool_wrapper(section_retrievers['four'], section_paths['four'], query)
-
-
-    def section_five_tool(query: str) -> str:
-        """
-        Section 5: Vaccination des Populations Particulières
+    # 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
 
-        Use this tool to retrieve recommendations for specific medical contexts:
-        preterm infants, immunocompromised children, allergies (e.g. eggs), and chronic diseases.
+    # Add instructions at the very beginning
+    safe_template = f"{custom_instructions}{original_template}"
 
-        Args:
-            query (str): A question about vaccination adaptations for vulnerable groups.
+    # Create new prompt with same metadata as original
+    original_prompt = original_prompts["agent_worker:system_prompt"]
 
-        Returns:
-            str: Guidelines from Section 5.
-        """
-        return section_tool_wrapper(section_retrievers['five'], section_paths['five'], query)
-
-
-    def section_six_tool(query: str) -> str:
-        """
-        Section 6: Chaîne du Froid
-
-        Use this tool for logistics, storage conditions, temperature monitoring,
-        and emergency procedures in case of cold chain failure.
-
-        Args:
-            query (str): A question about how vaccines should be stored and transported.
-
-        Returns:
-            str: Operational cold chain standards from Section 6.
-        """
-        return section_tool_wrapper(section_retrievers['six'], section_paths['six'], query)
-
-
-    def section_seven_tool(query: str) -> str:
-        """
-        Section 7: Sécurité des Injections
-
-        Use this tool to ensure injection safety: handling equipment, preventing needle-stick injuries,
-        and disposing of biomedical waste.
-
-        Args:
-            query (str): A question about safe injection practices.
-
-        Returns:
-            str: Procedures and guidelines from Section 7.
-        """
-        return section_tool_wrapper(section_retrievers['seven'], section_paths['seven'], query)
-
-
-    def section_eight_tool(query: str) -> str:
-        """
-        Section 8: Tenue d'une Séance de Vaccination & Vaccinovigilance
-
-        Use this tool to plan and monitor vaccination sessions, including material preparation,
-        injection recording, and handling of adverse events post-immunization (AEFI).
-
-        Args:
-            query (str): A question about session operations or vaccine side effect monitoring.
-
-        Returns:
-            str: Guidelines from Section 8.
-        """
-        return section_tool_wrapper(section_retrievers['eight'], section_paths['eight'], query)
+    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)
 
 
-    def section_nine_tool(query: str) -> str:
-        """
-        Section 9: Planification des Séances de Vaccination
+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
+    )
 
-        Use this tool to support logistical planning: mapping, resource estimation,
-        scheduling, and stock management.
+    # 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")
 
-        Args:
-            query (str): A question about planning and organizing vaccination sessions.
+    return agent
 
-        Returns:
-            str: Recommendations from Section 9.
-        """
-        return section_tool_wrapper(section_retrievers['nine'], section_paths['nine'], query)
 
+def initialize_rag_pipeline(tools):
+    """Initialize the RAG pipeline with tools"""
+    
+    # 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 agent
+    agent = create_agent(tools, llama_index_llm)
+    
+    return agent
 
-    def section_ten_tool(query: str) -> str:
-        """
-        Section 10: Mobilisation Sociale
 
-        Use this tool for strategies to increase public engagement, combat vaccine hesitancy,
-        and manage misinformation.
+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?")
+        
+        return f"Error processing your question: {error_msg}"
 
-        Args:
-            query (str): A question about public communication and trust-building around vaccines.
 
-        Returns:
-            str: Social mobilization approaches from Section 10.
-        """
-        return section_tool_wrapper(section_retrievers['ten'], section_paths['ten'], query)
+def aswer_language_detection(response_text: str) -> str:
+    """
+    Detect the language of the response text.
     
-    # Create FunctionTool objects
-    tools = [
-        FunctionTool.from_defaults(name="Guide_vector_tool", fn=guide_retrieval_tool),
-        FunctionTool.from_defaults(name="Immunization_in_Practice_tool", fn=immunization_tool),
-        # Section-specific tools
-        FunctionTool.from_defaults(name="section_one_vector_query_tool", fn=section_one_tool),
-        FunctionTool.from_defaults(name="section_two_vector_query_tool", fn=section_two_tool),
-        FunctionTool.from_defaults(name="section_three_vector_query_tool", fn=section_three_tool),
-        FunctionTool.from_defaults(name="section_four_vector_query_tool", fn=section_four_tool),
-        FunctionTool.from_defaults(name="section_five_vector_query_tool", fn=section_five_tool),
-        FunctionTool.from_defaults(name="section_six_vector_query_tool", fn=section_six_tool),
-        FunctionTool.from_defaults(name="section_seven_vector_query_tool", fn=section_seven_tool),
-        FunctionTool.from_defaults(name="section_eight_vector_query_tool", fn=section_eight_tool),
-        FunctionTool.from_defaults(name="section_nine_vector_query_tool", fn=section_nine_tool),
-        FunctionTool.from_defaults(name="section_ten_vector_query_tool", fn=section_ten_tool),
-    ]
+    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'
     
-    return tools
+    return answer_language
 
-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)
+
+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.
     
-    print("Environment prepared successfully!")
-    return tools, llm
+    Args:
+        agent: The initialized RAG agent
+        question (str): The user's question
+        chunks_directory (str): Path to the directory containing JSON files
+        
+    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
+        
+        # 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.")
+        
+        # Extract source IDs from the response (preserving order)
+        unique_ids = extract_source_ids(response_text)
+        
+        # Create mapping from source ID to sequential number
+        source_id_to_number = {source_id: i + 1 for i, source_id in enumerate(unique_ids)}
+        
+        # Convert citations to sequential numbers
+        sequential_response = convert_citations_to_sequential(response_text, source_id_to_number)
+        
+        # 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"]
+        
+        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}")
+        
+        # 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
+        
+        # Convert to JSON
+        cited_elements_json = json.dumps(cited_elements_ordered, ensure_ascii=False, indent=2)
+        answer_language = aswer_language_detection(response_text)
+        
+        return {
+            "response": sequential_response,
+            "cited_elements_json": cited_elements_json,
+            "unique_ids": unique_ids,
+            "citation_mapping": source_id_to_number,
+            "answer_language": answer_language
+        }
+        
+    except Exception as e:
+        error_msg = str(e)
+        print(f"Error processing question: {error_msg}")
+        
+        # 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}"
+        
+        return {
+            "response": fallback_response,
+            "cited_elements_json": "[]",
+            "unique_ids": [],
+            "citation_mapping": {},
+            "answer_language": "en"
+        }
+
+
+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)

rag_pipeline.py

@@ -1,16 +1,22 @@
 # -*- coding: utf-8 -*-
 """
-Enhanced RAG Pipeline for vaccine assistant
-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,332 +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_custom_prompt():
-    """Create custom prompt with medical assistant instructions"""
+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.
-
-## IMPORTANT REQUIREMENTS
-
-### 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 the following format:
-   - Use adjacent citations: [e795ebd28318886c0b1a5395ac30ad90][21a932b2340bb16707763f57f0ad2]
-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.
-
-## Tools
-
-You have access to a wide variety of tools. You are responsible for using the tools in any sequence you deem appropriate to complete the task at hand.
-This may require breaking the task into subtasks and using different tools to complete each subtask.
-
-You have access to the following tools:
-{tool_desc}
-
-## Output Format
-
-Please answer in the same language as the question and use the following format:
-
-```
-Thought: The current language of the user is: (user's language). I need to use a tool to help me answer the question.
-Action: tool name (one of {tool_names}) if using a tool.
-Action Input: the input to the tool, in a JSON format representing the kwargs (e.g. {{"input": "hello world", "num_beams": 5}})
-```
-
-Please ALWAYS start with a Thought.
-
-NEVER surround your response with markdown code markers. You may use code markers within your response if you need to.
-
-Please use a valid JSON format for the Action Input. Do NOT do this {{"input": "hello world", "num_beams": 5}}.
-
-If this format is used, the tool will respond in the following format:
-
-```
-Observation: tool response
-```
+    documents = []
+    for element in chunks_data:
+        text = element["text"]
+        metadata = {
+            "language": "fra",
+            "source": element["filename"],
+            "filetype": element["filetype"],
+            "element_id": element["element_id"]
+        }
 
-You should keep repeating the above format till you have enough information to answer the question without using any more tools. At that point, you MUST respond in one of the following two formats:
+        if "TableElement" == element["type"]:
+            metadata["table_text_as_html"] = element["table_text_as_html"]
 
-```
-Thought: I can answer without using any more tools. I'll use the user's language to answer. Remember to include proper citations
-Answer: [your answer here with proper citations (In the same language as the user's question)]
-```
+        doc = Document(page_content=text, metadata=metadata)
+        documents.append(doc)
 
-```
-Thought: I cannot answer the question with the provided tools.
-Answer: [your answer here (In the same language as the user's question)]
-```
+    # Create vector store
+    vectorstore = Chroma.from_documents(
+        documents=documents,
+        embedding=embedding_function,
+        collection_name=collection_name,
+        persist_directory="chroma_db_multilingual"
+    )
+    return vectorstore, documents
 
-## Current Conversation
 
-Below is the current conversation consisting of interleaving human and assistant messages.
-"""
+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
+    )
     
-    try:
-        custom_prompt = PromptTemplate(
-            template=custom_instructions,
-            template_vars=["tool_desc", "tool_names"]
-        )
-        return custom_prompt
-    except:
-        # Fallback to simple template
-        return PromptTemplate(template=custom_instructions)
-
-def create_safe_custom_prompt(tools, llm):
-    """Create a safe version that won't have formatting conflicts"""
+    # BM25 retriever
+    bm25_retriever = BM25Retriever.from_documents(docs)
+    bm25_retriever.k = 2
     
-    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.
-
-## IMPORTANT REQUIREMENTS
-
-### 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 the following format:
-   - Use adjacent citations: [e795ebd28318886c0b1a5395ac30ad90][21a932b2340bb16707763f57f0ad2]
-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.
-
----
-
-"""
+    # Ensemble retriever
+    ensemble_retriever = EnsembleRetriever(
+        retrievers=[vector_retriever, bm25_retriever],
+        weights=[0.5, 0.5]
+    )
     
-    # 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"]
-
+    # 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:
-        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)
+        retrieved_docs = retriever.get_relevant_documents(query)
+        nodes_from_retrieved_docs = convert_chromadb_to_llamaindex_nodes(retrieved_docs)
 
-def create_agent(tools, llm):
-    """Create the ReAct agent with custom prompt"""
-    
-    # Create agent
-    agent = ReActAgent.from_tools(
-        tools,
-        llm=llm,
-        verbose=True,
-    )
+        if not nodes_from_retrieved_docs:
+            return "No relevant documents found for the query."
 
-    # 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")
+        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 = []
+        
+        # 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"❌ Safe prompt update failed: {e}")
-        print("⚠️  Using original agent without modifications")
+        print(f"Error in section tool: {e}")
+        return f"Error retrieving documents: {str(e)}"
 
-    return agent
 
-def initialize_rag_pipeline(tools):
-    """Initialize the RAG pipeline with tools"""
+def create_section_tools(embedding_function, llm):
+    """Create all section-specific retrieval tools with improved descriptions"""
     
-    # Initialize LlamaIndex LLM
-    llama_index_llm = GoogleGenAI(
-        model="models/gemini-2.0-flash",
-        api_key=os.getenv('GOOGLE_API_KEY'),
-    )
+    # 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 agent
-    agent = create_agent(tools, llama_index_llm)
+    # 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 agent
+    # 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
 
-def process_question(agent, question: str) -> str:
-    """Process a question through the RAG pipeline"""
-    try:
-        response = agent.chat(question)
-        return response.response
-    except Exception as e:
-        print(f"Error processing question: {e}")
-        return f"Error processing your question: {str(e)}"
+    # Tool Functions with Improved Efficiency Focus
 
-def aswer_language_detection(response_text: str) -> str:
-    """
-    Detect the language of the response text.
-    
-    Args:
-        response_text (str): The response text to analyze.
+    def guide_retrieval_tool(query: str) -> str:
+        """
+        **PRIMARY TOOL - USE FIRST FOR MOST QUESTIONS**
         
-    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'
-    
-    finally:
-        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.
-    
-    Args:
-        agent: The initialized RAG agent
-        question (str): The user's question
-        chunks_directory (str): Path to the directory containing JSON files
+        Comprehensive search across the entire Algerian National Vaccination Guide (2023).
         
-    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
-        response = agent.chat(question)
-        response_text = response.response
+        **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)
-        aswer_language= aswer_language_detection(response_text)
-        return {
-            "response": sequential_response,
-            "cited_elements_json": cited_elements_json,
-            "unique_ids": unique_ids,
-            "citation_mapping": source_id_to_number,
-            "answer_language":aswer_language
-        }
+        Section 2: Vaccine-preventable diseases - definitions, symptoms, transmission, complications.
         
-    except Exception as e:
-        print(f"Error processing question: {e}")
-        return {
-            "response": response_text if 'response_text' in locals() else "Error occurred",
-            "cited_elements_json": "[]",
-            "unique_ids": [],
-            "citation_mapping": {},
-            "answer_language": "en"  # Default to English if not specified
-        }
+        **Use ONLY for specific disease definition questions like:**
+        - "What is diphtheria?" 
+        - "Define measles according to Algerian protocol"
+        - "Symptoms of polio"
+        
+        **Keywords:** definition, symptoms, transmission, complications, disease characteristics
+        
+        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**
+        
+        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