update the code with new simpler version

main.py

@@ -1,138 +1,135 @@
 # -*- coding: utf-8 -*-
 """
-Agentic RAG API Server
-A FastAPI server that provides an intelligent vaccination assistant using agentic RAG.
+FastAPI server for vaccine assistant
+Main entry point for the application
 """
 
-import asyncio
-from fastapi import FastAPI, HTTPException, Query
+from fastapi import FastAPI, Query, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
+import os
+from dotenv import load_dotenv
+import logging
 
-from config import Config
-from models import QuestionRequest, QuestionResponse, HealthResponse
-from rag_system import AgenticRAGSystem
+# Import our modules
+from prepare_env import prepare_environment
+from rag_pipeline import initialize_rag_pipeline, process_question
 
-# Initialize the system
-config = Config()
-rag_system = AgenticRAGSystem(config)
+# Load environment variables
+load_dotenv()
 
-# Global variable for initialization status
-is_initialized = False
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# FastAPI app
+# Initialize FastAPI app
 app = FastAPI(
-    title="Agentic RAG Vaccination Assistant API",
-    description="An intelligent vaccination assistant using agentic RAG for medical professionals",
+    title="Vaccine Assistant API",
+    description="AI-powered vaccine assistant for medical professionals",
     version="1.0.0"
 )
 
 # Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Configure this properly for production
+    allow_origins=["*"],  # Configure appropriately for production
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
+# Global variables for the agent (initialized on startup)
+agent = None
+
 @app.on_event("startup")
 async def startup_event():
-    """Initialize the RAG system on startup"""
-    global is_initialized
-    print("🚀 Starting Agentic RAG API server...")
-    
-    # Run initialization in a thread pool to avoid blocking
-    import concurrent.futures
-    
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        loop = asyncio.get_event_loop()
-        success = await loop.run_in_executor(executor, rag_system.initialize_system)
+    """Initialize the RAG pipeline on startup"""
+    global agent
+    try:
+        logger.info("Starting up vaccine assistant...")
+        
+        # Check for required environment variables
+        if not os.getenv('GOOGLE_API_KEY'):
+            logger.warning("GOOGLE_API_KEY not found in environment variables")
+        
+        # Prepare environment and tools
+        logger.info("Preparing environment and tools...")
+        tools, llm = prepare_environment()
         
-    if success:
-        is_initialized = True
-        print("✅ System initialized successfully")
-    else:
-        print("❌ System initialization failed")
+        # Initialize RAG pipeline
+        logger.info("Initializing RAG pipeline...")
+        agent = initialize_rag_pipeline(tools)
+        
+        logger.info("✅ Vaccine assistant startup completed successfully!")
+        
+    except Exception as e:
+        logger.error(f"❌ Error during startup: {e}")
+        raise e
 
-@app.get("/", response_model=HealthResponse)
+@app.get("/")
 async def root():
-    """Root endpoint"""
-    return HealthResponse(
-        status="online",
-        message="Agentic RAG Vaccination Assistant API is running"
-    )
-
-@app.get("/health", response_model=HealthResponse)
-async def health_check():
     """Health check endpoint"""
-    if is_initialized:
-        return HealthResponse(
-            status="healthy",
-            message="System is initialized and ready to serve requests"
-        )
-    else:
-        return HealthResponse(
-            status="initializing",
-            message="System is still initializing. Please wait."
-        )
+    return {
+        "message": "Vaccine Assistant API is running",
+        "status": "healthy",
+        "version": "1.0.0"
+    }
 
-@app.get("/ask", response_model=QuestionResponse)
-async def ask_question(
-    question: str = Query(..., description="The question to ask the vaccination assistant"),
-    with_citations: bool = Query(False, description="Include citations in the response")
-):
-    """Ask a question to the vaccination assistant"""
-    if not is_initialized:
-        raise HTTPException(
-            status_code=503,
-            detail="System is still initializing. Please wait and try again."
-        )
+@app.get("/ask")
+async def ask_question(question: str = Query(..., description="The medical question to ask")):
+    """
+    Main endpoint for asking questions to the vaccine assistant
+    
+    Args:
+        question: The medical question related to vaccines
+        
+    Returns:
+        JSON response with question and answer
+    """
+    global agent
+    
+    if agent is None:
+        raise HTTPException(status_code=503, detail="Agent not initialized. Please try again later.")
     
     if not question.strip():
-        raise HTTPException(
-            status_code=400,
-            detail="Question cannot be empty"
-        )
+        raise HTTPException(status_code=400, detail="Question cannot be empty")
     
     try:
-        # Run the question processing in a thread pool to avoid blocking
-        import concurrent.futures
+        logger.info(f"Processing question: {question[:100]}...")
         
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            loop = asyncio.get_event_loop()
-            answer = await loop.run_in_executor(
-                executor, 
-                rag_system.ask_question, 
-                question, 
-                with_citations
-            )
+        # Process the question through RAG pipeline
+        answer = process_question(agent, question)
         
-        return QuestionResponse(
-            question=question,
-            answer=answer,
-            status="success"
-        )
+        logger.info("Question processed successfully")
+        
+        return {
+            "question": question,
+            "answer": answer,
+            "status": "success"
+        }
         
     except Exception as e:
-        print(f"Error processing question: {e}")
-        raise HTTPException(
-            status_code=500,
-            detail=f"An error occurred while processing your question: {str(e)}"
-        )
+        logger.error(f"Error processing question: {e}")
+        raise HTTPException(status_code=500, detail=f"Error processing question: {str(e)}")
+
+@app.get("/health")
+async def health_check():
+    """Detailed health check endpoint"""
+    global agent
+    
+    health_status = {
+        "status": "healthy" if agent is not None else "unhealthy",
+        "agent_initialized": agent is not None,
+        "google_api_key_configured": bool(os.getenv('GOOGLE_API_KEY')),
+        "version": "1.0.0"
+    }
+    
+    if agent is None:
+        health_status["status"] = "unhealthy"
+        health_status["message"] = "Agent not initialized"
+    
+    return health_status
 
-@app.post("/ask", response_model=QuestionResponse)
-async def ask_question_post(request: QuestionRequest):
-    """Ask a question to the vaccination assistant (POST version)"""
-    return await ask_question(request.question, request.with_citations)
 
 if __name__ == "__main__":
     import uvicorn
-    
-    print("Starting Agentic RAG API server...")
-    uvicorn.run(
-        "main:app",
-        host="0.0.0.0",
-        port=8000,
-        reload=True,
-        log_level="info"
-    )
+    uvicorn.run(app, host="0.0.0.0", port=7860)

models.py

@@ -1,14 +0,0 @@
-from pydantic import BaseModel
-
-class QuestionRequest(BaseModel):
-    question: str
-    with_citations: bool = False
-
-class QuestionResponse(BaseModel):
-    question: str
-    answer: str
-    status: str = "success"
-
-class HealthResponse(BaseModel):
-    status: str
-    message: str

prepare_env.py

@@ -0,0 +1,254 @@
+# -*- coding: utf-8 -*-
+"""
+Environment preparation script for vaccine assistant
+Creates vector stores and retrieval tools
+"""
+
+import os
+import json
+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
+
+# Apply nest_asyncio for compatibility
+nest_asyncio.apply()
+
+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', 'AIzaSyBho3W4W9fR7wHUJbX18JKH-12wDSD7pWg')
+    llm = ChatGoogleGenerativeAI(
+        model="gemini-2.0-flash",
+        google_api_key=genai_api_key
+    )
+    
+    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)
+    
+    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"]
+
+        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
+    )
+    
+    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)
+
+        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['elements']['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': './data/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(path, f"Guide_2023_{section}", embedding_function)
+            section_retrievers[section] = create_retriever(vstore, docs, llm)
+    
+    # # Create main guide retriever
+    # guide_path = '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
+    
+    # # Define tool functions
+    # def guide_retrieval_tool(query: str) -> str:
+    #     """General-purpose retrieval tool for the entire Algerian National Vaccination Guide"""
+    #     if not guide_retriever:
+    #         return "Guide retriever not available"
+    #     return section_tool_wrapper(guide_retriever, guide_path, query)
+    
+    # def section_one_tool(query: str) -> str:
+    #     """Section 1: Programme Élargi de Vaccination"""
+    #     return section_tool_wrapper(section_retrievers['one'], section_paths['one'], query)
+    
+    # def section_two_tool(query: str) -> str:
+    #     """Section 2: Maladies Ciblées"""
+    #     return section_tool_wrapper(section_retrievers['two'], section_paths['two'], query)
+    
+    # def section_three_tool(query: str) -> str:
+    #     """Section 3: Vaccins du Calendrier"""
+    #     return section_tool_wrapper(section_retrievers['three'], section_paths['three'], query)
+    
+    # def section_four_tool(query: str) -> str:
+    #     """Section 4: Rattrapage Vaccinal"""
+    #     return section_tool_wrapper(section_retrievers['four'], section_paths['four'], query)
+    
+    # def section_five_tool(query: str) -> str:
+    #     """Section 5: Populations Particulières"""
+    #     return section_tool_wrapper(section_retrievers['five'], section_paths['five'], query)
+    
+    # def section_six_tool(query: str) -> str:
+    #     """Section 6: Chaîne du Froid"""
+    #     return section_tool_wrapper(section_retrievers['six'], section_paths['six'], query)
+    
+    # def section_seven_tool(query: str) -> str:
+    #     """Section 7: Sécurité des Injections"""
+    #     return section_tool_wrapper(section_retrievers['seven'], section_paths['seven'], query)
+    
+    # def section_eight_tool(query: str) -> str:
+    #     """Section 8: Séance de Vaccination & Vaccinovigilance"""
+    #     return section_tool_wrapper(section_retrievers['eight'], section_paths['eight'], query)
+    
+    # def section_nine_tool(query: str) -> str:
+    #     """Section 9: Planification des Séances de Vaccination"""
+    #     return section_tool_wrapper(section_retrievers['nine'], section_paths['nine'], query)
+    
+    def section_ten_tool(query: str) -> str:
+        """Section 10: Mobilisation Sociale"""
+        return section_tool_wrapper(section_retrievers['ten'], section_paths['ten'], query)
+    
+    # Create FunctionTool objects
+    tools = [
+        # FunctionTool.from_defaults(name="Guide_vector_tool", fn=guide_retrieval_tool),
+        # 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),
+    ]
+    
+    return tools
+
+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

@@ -0,0 +1,189 @@
+# -*- coding: utf-8 -*-
+"""
+RAG Pipeline for vaccine assistant
+Handles agent creation and question answering
+"""
+
+from llama_index.core import PromptTemplate
+from llama_index.core.agent import ReActAgent
+from llama_index.llms.google_genai import GoogleGenAI
+import os
+
+def create_custom_prompt():
+    """Create custom prompt with medical assistant instructions"""
+    
+    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
+```
+
+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:
+
+```
+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)]
+```
+
+```
+Thought: I cannot answer the question with the provided tools.
+Answer: [your answer here (In the same language as the user's question)]
+```
+
+## Current Conversation
+
+Below is the current conversation consisting of interleaving human and assistant messages.
+"""
+    
+    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"""
+    
+    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.
+
+---
+
+"""
+    
+    # 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"]
+
+    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 create_agent(tools, llm):
+    """Create the ReAct agent with custom prompt"""
+    
+    # Create agent
+    agent = ReActAgent.from_tools(
+        tools,
+        llm=llm,
+        verbose=True,
+    )
+
+    # 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")
+    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"""
+    
+    # Initialize LlamaIndex LLM
+    llama_index_llm = GoogleGenAI(
+        model="models/gemini-2.0-flash",
+        api_key=os.getenv('GOOGLE_API_KEY', 'AIzaSyDsbC8H6e08TKDwa5WPE3SiBA39e20K4co'),
+    )
+    
+    # Create agent
+    agent = create_agent(tools, llama_index_llm)
+    
+    return agent
+
+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)}"

rag_system.py

@@ -1,345 +0,0 @@
-import os
-import json
-from typing import List, Dict, Any, Optional
-
-# LlamaIndex imports
-from llama_index.core import Settings, PromptTemplate
-from llama_index.llms.google_genai import GoogleGenAI
-from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-from llama_index.core.agent import ReActAgent
-from llama_index.core.tools import FunctionTool
-from llama_index.core.schema import TextNode
-
-# LangChain imports
-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 config import Config
-
-class AgenticRAGSystem:
-    def __init__(self, config: Config):
-        self.config = config
-        self.setup_llm_and_embeddings()
-        self.guide_retriever = None
-        self.section_retrievers = {}
-        self.agent = None
-        
-    def setup_llm_and_embeddings(self):
-        """Initialize LLM and embedding models"""
-        # LlamaIndex settings
-        Settings.llm = GoogleGenAI(
-            model=self.config.LLM_MODEL,
-            api_key=self.config.GOOGLE_API_KEY_1,
-        )
-        Settings.embed_model = HuggingFaceEmbedding(
-            model_name=self.config.EMBEDDING_MODEL
-        )
-        
-        # LangChain components
-        self.embedding_function = HuggingFaceEmbeddings(
-            model_name=self.config.EMBEDDING_MODEL
-        )
-        self.llm = ChatGoogleGenerativeAI(
-            model="gemini-2.0-flash",
-            google_api_key=self.config.GOOGLE_API_KEY_2
-        )
-        
-    def create_vectorstore_from_json(self, json_path: str, collection_name: str):
-        """Create vector store from JSON chunks"""
-        if not os.path.exists(json_path):
-            raise FileNotFoundError(f"JSON file not found: {json_path}")
-            
-        with open(json_path, "r", encoding="utf-8") as f:
-            chunks_data = json.load(f)
-        
-        documents = []
-        for element in chunks_data:
-            text = element.get("text", "").strip()
-            if not text:
-                continue
-                
-            metadata = {
-                "language": "fra",
-                "source": element.get("filename", "unknown"),
-                "filetype": element.get("filetype", "unknown"),
-                "element_id": element.get("element_id", "unknown")
-            }
-            
-            if element.get("type") == "TableElement":
-                metadata["table_text_as_html"] = element.get("table_text_as_html", "")
-            
-            doc = Document(page_content=text, metadata=metadata)
-            documents.append(doc)
-        
-        vectorstore = Chroma.from_documents(
-            documents=documents,
-            embedding=self.embedding_function,
-            collection_name=collection_name,
-            persist_directory=self.config.CHROMA_DB_PATH
-        )
-        
-        return vectorstore, documents
-    
-    def create_retriever(self, vectorstore, docs):
-        """Create ensemble retriever with semantic and BM25 search"""
-        retriever_multilingual = vectorstore.as_retriever(
-            search_type="similarity",
-            search_kwargs={"k": 6}
-        )
-        
-        bm25_retriever = BM25Retriever.from_documents(docs)
-        bm25_retriever.k = 2
-        
-        ensemble_retriever = EnsembleRetriever(
-            retrievers=[retriever_multilingual, bm25_retriever],
-            weights=[0.5, 0.5]
-        )
-        
-        expanding_retriever = MultiQueryRetriever.from_llm(
-            retriever=ensemble_retriever,
-            llm=self.llm
-        )
-        
-        return expanding_retriever
-    
-    def convert_chromadb_to_llamaindex_nodes(self, chromadb_documents: List) -> List[TextNode]:
-        """Convert ChromaDB documents 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:
-                print(f"Error converting document {i}: {e}")
-                continue
-        
-        return nodes
-    
-    def section_tool_wrapper(self, retriever, section_path_chunks, query):
-        """Generic wrapper for section-specific tools"""
-        try:
-            retrieved_docs = retriever.get_relevant_documents(query)
-            
-            if not retrieved_docs:
-                return "No relevant documents found for the query."
-            
-            chunk_ids = [doc.metadata.get('element_id') for doc in retrieved_docs]
-            
-            if not os.path.exists(section_path_chunks):
-                return f"Section data file not found: {section_path_chunks}"
-            
-            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 chu.get("type") == "TableElement":
-                    text = f"[Source: {chu.get('element_id', 'Unknown')}]\nCONTENT:\n{chu.get('text', '')}\nHTML:\n{chu.get('table_text_as_html', '')}\n\n"
-                    combined_text.append(text)
-                else:
-                    elements = chu.get("elements", [chu])  # Handle both formats
-                    for element in elements:
-                        text = f"[Source: {element.get('element_id', 'Unknown')}]\nCONTENT:\n{element.get('text', '')}\n\n"
-                        combined_text.append(text)
-            
-            result = "\n---\n".join(combined_text)
-            print(f"Retrieved {len(retrieved_docs)} documents for query: {query[:50]}...")
-            return result
-            
-        except Exception as e:
-            print(f"Error in section tool wrapper: {e}")
-            return f"Error retrieving documents: {str(e)}"
-    
-    def initialize_system(self):
-        """Initialize all retrievers and create the agent"""
-        try:
-            # File paths - make these configurable
-            json_files = {
-                # "guide": "Guide-pratique-de-mise-en-oeuvre-du-calendrier-national-de-vaccination-2023.json",
-                # "section_1": "section_one_chunks.json",
-                # "section_2": "section_two_chunks.json",
-                # "section_3": "section_three_chunks.json",
-                # "section_4": "section_four_chunks.json",
-                # "section_5": "section_five_chunks.json",
-                # "section_6": "section_six_chunks.json",
-                # "section_7": "section_seven_chunks.json",
-                # "section_8": "section_eight_chunks.json",
-                # "section_9": "section_nine_chunks.json",
-                "section_10": "section_ten_chunks.json",
-            }
-            
-            # Check if files exist
-            for name, filepath in json_files.items():
-                full_path = os.path.join(self.config.BASE_PATH, filepath)
-                if not os.path.exists(full_path):
-                    print(f"Warning: {name} file not found at {full_path}")
-            
-            # Initialize main guide retriever
-            guide_path = os.path.join(self.config.BASE_PATH, json_files["guide"])
-            if os.path.exists(guide_path):
-                guide_vstore, guide_doc = self.create_vectorstore_from_json(guide_path, "Guide_2023_multilingual")
-                self.guide_retriever = self.create_retriever(guide_vstore, guide_doc)
-            
-            # Initialize section retrievers
-            for i in range(1, 11):
-                section_key = f"section_{i}"
-                section_path = os.path.join(self.config.BASE_PATH, json_files[section_key])
-                if os.path.exists(section_path):
-                    vstore, doc = self.create_vectorstore_from_json(section_path, f"Section_{i}_multilingual")
-                    self.section_retrievers[section_key] = {
-                        'retriever': self.create_retriever(vstore, doc),
-                        'path': section_path
-                    }
-            
-            # Create tools
-            tools = self.create_tools()
-            
-            # Create agent
-            self.agent = ReActAgent.from_tools(
-                tools,
-                llm=Settings.llm,
-                verbose=True
-            )
-            
-            # Apply custom prompt
-            self.apply_custom_prompt()
-            
-            print("✅ Agentic RAG system initialized successfully")
-            return True
-            
-        except Exception as e:
-            print(f"❌ Failed to initialize system: {e}")
-            return False
-    
-    def create_tools(self):
-        """Create all the function tools"""
-        tools = []
-        
-        # # Main guide tool
-        # if self.guide_retriever:
-        #     def guide_tool(query: str) -> str:
-        #         """General-purpose retrieval tool for the Algerian National Vaccination Guide (2023)"""
-        #         return self.section_tool_wrapper(
-        #             self.guide_retriever, 
-        #             os.path.join(self.config.BASE_PATH, "Guide-pratique-de-mise-en-oeuvre-du-calendrier-national-de-vaccination-2023.json"),
-        #             query
-        #         )
-            
-        #     tools.append(FunctionTool.from_defaults(name="Guide_vector_tool", fn=guide_tool))
-        
-        # Section tools
-        section_descriptions = {
-            # "section_1": "Programme Élargi de Vaccination - General national immunization program in Algeria",
-            # "section_2": "Maladies Ciblées - Diseases targeted by the national vaccination calendar",
-            # "section_3": "Vaccins du Calendrier - Vaccines themselves: types, administration methods, compositions",
-            # "section_4": "Rattrapage Vaccinal - Catch-up vaccination procedures and schedules",
-            # "section_5": "Populations Particulières - Vaccination of special populations (premature, immunosuppressed, etc.)",
-            # "section_6": "Chaîne du Froid - Vaccine cold chain logistics and storage",
-            # "section_7": "Sécurité des Injections - Safe injection practices",
-            # "section_8": "Séance de Vaccination & Vaccinovigilance - Vaccination sessions and adverse event monitoring",
-            # "section_9": "Planification des Séances - Planning of vaccination sessions",
-            "section_10": "Mobilisation Sociale - Community mobilization and vaccine hesitancy"
-        }
-        
-        for section_key, description in section_descriptions.items():
-            if section_key in self.section_retrievers:
-                def create_section_tool(section_data, desc):
-                    def section_tool(query: str) -> str:
-                        return self.section_tool_wrapper(
-                            section_data['retriever'],
-                            section_data['path'],
-                            query
-                        )
-                    section_tool.__doc__ = f"Handles queries about {desc}"
-                    return section_tool
-                
-                section_tool_func = create_section_tool(self.section_retrievers[section_key], description)
-                tools.append(FunctionTool.from_defaults(
-                    name=f"{section_key}_vector_query_tool",
-                    fn=section_tool_func
-                ))
-        
-        return tools
-    
-    def apply_custom_prompt(self):
-        """Apply custom instructions to the agent"""
-        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.
-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.
-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
-2. For lists, maintain the original bullet points/numbering and include citations.
-3. Present information concisely but ensure clinical accuracy is never compromised.
-"""
-        
-        try:
-            # Create safe custom prompt
-            temp_agent = ReActAgent.from_tools([], llm=Settings.llm, verbose=False)
-            original_prompts = temp_agent.get_prompts()
-            original_template = original_prompts["agent_worker:system_prompt"].template
-            
-            safe_template = f"""{custom_instructions}
-
----
-
-{original_template}"""
-            
-            original_prompt = original_prompts["agent_worker:system_prompt"]
-            new_prompt = PromptTemplate(
-                template=safe_template,
-                template_vars=original_prompt.template_vars,
-                metadata=getattr(original_prompt, 'metadata', None)
-            )
-            
-            self.agent.update_prompts({"agent_worker:system_prompt": new_prompt})
-            print("✅ Successfully updated with custom prompt")
-            
-        except Exception as e:
-            print(f"❌ Custom prompt update failed: {e}")
-    
-    def ask_question(self, question: str, with_citations: bool = False) -> str:
-        """Process a question using the agentic RAG system"""
-        if not self.agent:
-            raise ValueError("Agent not initialized. Call initialize_system() first.")
-        
-        try:
-            response = self.agent.chat(question)
-            answer = response.response
-            
-            if not with_citations:
-                # Simple regex to remove citations if not wanted
-                import re
-                answer = re.sub(r'\[[\w\d-]+\]', '', answer)
-            
-            return answer
-            
-        except Exception as e:
-            print(f"Error processing question: {e}")
-            return f"Sorry, I encountered an error while processing your question: {str(e)}"