import gradio as gr
from RagWithConfidenceScore import RagWithScore  # 

# Initialize the RAG system
rag_system = RagWithScore()

# Load or create the vector store
rag_system.load_and_process_documents()

# Define the function to handle user queries
def answer_financial_query(query):
    # Use the RAG system to answer the question
    result = rag_system.answer_question(query)
    
    # Format the output
    answer = result["answer"]
    confidence_score = result["confidence_score"]
    confidence_level = result["confidence_level"]
    sources = "\n\n".join([doc.page_content for doc in result["source_documents"]])

    return answer, f"{confidence_score:.2f}", confidence_level, sources
    
    # Return the results
    # return {
    #     "Answer": answer,
    #     "Confidence Score": f"{confidence_score:.2f}",
    #     "Confidence Level": confidence_level,
    #     "Source Documents": sources
    # }

# Create a Gradio interface
interface = gr.Interface(
    fn=answer_financial_query,  # Function to call
    inputs=gr.Textbox(lines=2, placeholder="Enter your financial query here..."),  # Input component
    outputs=[  # Output components
        gr.Textbox(label="Answer", lines=8),
        gr.Textbox(label="Confidence Score"),
        gr.Textbox(label="Confidence Level")
        # gr.Textbox(label="Source Documents", lines=10)
    ],
    title="Financial RAG System",
    description="Ask questions about financial data and get answers powered by Retrieval-Augmented Generation (RAG).",
    examples=[
        ["What is the current revenue growth rate?"],
        ["Explain the concept of EBITDA."],
        ["What are the key financial risks mentioned in the report?"],
        ["How has the debt-to-equity ratio changed over the last two years?"]
    ],
    cache_examples=False
)

# Launch the interface
interface.launch(share=True)