from fastapi import FastAPI
from rag_config import *

app = FastAPI()

# Charger le RAG une seule fois au démarrage
rag = RAG(CONFIG=CONFIG)

rag.reset_index()

@app.post("/rag")
def rag_query(data: dict):

    query = data["query"]

    context = rag.get_retrieval(query=query, number_of_hits=3)

    prompt = f"""
You are a medical assistant.

Use the following context to answer the question.

IMPORTANT RULES:
- Do not mention the context.
- Do not mention figures or sections.
- Do not say "according to the context".
- Give a clear explanation as if you are speaking to a patient.

Context:
{context}

Question:
{query}

Answer:
"""

    response = rag.foundation_model.generate_response_with_context(
        prompt=prompt,
        context=context
    )

    return {"response": response[0]["generated_text"]}