from fastapi import FastAPI, Query
from prepare_env import prepare_environment_and_retriever
from rag_pipeline import full_rag_pipeline
from langchain_google_genai import GoogleGenerativeAI
import os
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()


app = FastAPI()

# Prepare the environment and load the vector store
expanding_retriever = prepare_environment_and_retriever()


@app.get("/ask")
def ask_question(question: str, with_citations: bool = Query(False, description="Include citations in the response")):
    response = full_rag_pipeline(question, expanding_retriever,clean_all_citations=with_citations)
    return {"question": question, "answer": response}
@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)