from indexer import (
    create_vector_database,
    get_llm,
    get_prompt_template,
)
import gradio as gr


def format_contexts(contexts):
    return "\n".join(
        [
            f"Reference {i+1}:\n{doc.metadata['question']}\n{doc.metadata['answer']}"
            for i, doc in enumerate(contexts)
        ]
    )


class CustomRAG:
    def __init__(self, vector_db, llm, prompt_template):
        self.vector_db = vector_db
        self.llm = llm
        self.prompt_template = prompt_template

    def run(self, query):
        retriever = self.vector_db.as_retriever(search_kwargs={"k": 3})
        contexts = retriever.invoke(query)
        formatted_context = format_contexts(contexts)
        prompt = self.prompt_template.format(context=formatted_context, question=query)
        return self.llm.invoke(prompt), contexts


def answer_question(query):
    llm = get_llm("google/flan-t5-base")
    # llm = get_llm("FreedomIntelligence/HuatuoGPT-o1-7B")
    vector_database = create_vector_database("sentence-transformers/all-MiniLM-L6-v2")
    prompt_template = get_prompt_template()
    rag = CustomRAG(
        vector_database,
        llm,
        prompt_template,
    )
    response, _ = rag.run(query)

    return response


demo = gr.Interface(
    fn=answer_question,
    inputs=[
        gr.Textbox(
            label="Describe your medical concern",
            placeholder="e.g. I've been feeling tired and dizzy lately.",
            lines=3,
        ),
    ],
    outputs="text",
    title="Medical Assistant – RAG",
    description=(
        "Get helpful insights based on your described symptoms. "
        "This assistant uses medical reference data to provide informative responses. "
        "Note: This is not a substitute for professional medical advice."
    ),
)

demo.launch()