main.py

# -*- coding: utf-8 -*-
"""
FastAPI server for vaccine assistant
Main entry point for the application
"""

from fastapi import FastAPI, Query, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from langchain_google_genai import GoogleGenerativeAI
import os
from dotenv import load_dotenv
import logging

# Import our modules
from prepare_env import prepare_environment
from rag_pipeline import initialize_rag_pipeline, process_question,process_question_with_citations

# Load environment variables
load_dotenv()

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(
    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 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 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()
        
        # 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("/")
async def root():
    """Health check endpoint"""
    return {
        "message": "Vaccine Assistant API is running",
        "status": "healthy",
        "version": "1.0.0"
    }

@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")
    
    try:
        logger.info(f"Processing question: {question[:100]}...")
        
        # Process the question through RAG pipeline
        answer = process_question_with_citations(agent, question)
        
        logger.info("Question processed successfully")
        
        return {
            "question": question,
            "answer": answer,
            "status": "success"
        }
        
    except Exception as 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.get("/generate_title")
def generate_title(first_question: str = Query(..., description="The first question to generate a title from")):
    # Initialize the LLM - using the same model as in prepare_env.py
    llm = GoogleGenerativeAI(
        model="gemini-2.0-flash",
        google_api_key=os.getenv("GOOGLE_API_KEY")
    )
    
    prompt = f"""Analyze this question and generate a very short title (3-5 words max):
    1. If it's medical/vaccine-related: Create a professional clinical title
    2. If non-medical: Create a general topic title
    3. If unclear or greeting: Use "General Inquiry"
    
    Always return just the title text, nothing else.
    
    Question: {first_question}
    
    Title:"""
    
    title = llm.invoke(prompt)
    return {"title": title.strip()}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)