import pickle
import gradio as gr
from config.rag_config import RAGConfig
from src.rag_pipeline import RAGPipeline

config = RAGConfig()

# 加载向量库
with open(config.vector_db_path, "rb") as f:
    data = pickle.load(f)

docs, doc_embeddings = data["texts"], data["embeddings"]
pipeline = RAGPipeline(config, docs, doc_embeddings)

def answer_question(query, threshold):
    pipeline.config.similarity_threshold = threshold
    answer, retrieved = pipeline.ask(query)
    context = "\n\n".join([f"Score: {s:.4f}\n{t}" for t, s in retrieved])
    return answer, context

demo = gr.Interface(
    fn=answer_question,
    inputs=[
        gr.Textbox(label="Enter your question"),
        gr.Slider(0.0, 1.0, value=0.4, step=0.05, label="Similarity Threshold")
    ],
    outputs=[
        gr.Textbox(label="Answer"),
        gr.Textbox(label="Retrieved Contexts")
    ],
    title="📘 Multi-PDF RAG System"
)

if __name__ == "__main__":
    demo.launch()