import os, uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from langchain_google_genai import GoogleGenerativeAIEmbeddings
import google.generativeai as genai
from langchain_community.vectorstores import FAISS
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate
from dotenv import load_dotenv

app = FastAPI()
class Question(BaseModel):
    query: str

load_dotenv()
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
conversational_model = ChatGoogleGenerativeAI(model="gemini-2.5-pro", temperature=0.7)

def get_conversational_chain():    
    prompt_template = """
    You are a helpful assistant tasked with extracting accurate answers **only from the given context**.

    If the question is about matching (e.g., "Who is referred to as X?"), select the correct match from the context.

    If the correct answer is **not present** in the context, respond exactly with:
    "উত্তর প্রসঙ্গে নেই" (The answer is not in the context.)

    ---

    প্রসঙ্গ (Context):
    {context}

    প্রশ্ন (Question):
    {question}

    উত্তর (Answer):
    """
    
    prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
    return load_qa_chain(conversational_model, chain_type="stuff", prompt=prompt)


@app.post("/ask")
async def create_items(que: Question):
    try:
        db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
        docs_and_scores = db.similarity_search_with_score(que.query, k=100)
        docs = [doc for doc, score in docs_and_scores]
        chain = get_conversational_chain()
        response = chain({"input_documents": docs, "question": que.query})
        return {"answer": response["output_text"]}
    except Exception as e:
        return {"error": str(e)}

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=5656)