import gradio as gr
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
import openai
import pickle
from utils import rag_based_generation, search_with_metadata_and_reranking


embedder = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")

index = faiss.read_index('faiss_index.index')


chunked_documents = np.load('chunked_documents.npy', allow_pickle=True)

def rag_interface(query, author=None, category=None):
    try:
        filters = {}
        if author:
            filters['author'] = author
        if category:
            filters['categories'] = category
        generated_answer = rag_based_generation(query, index, chunked_documents, embedder, filters=filters, top_k=5)
        return generated_answer
    except Exception as e:
        import traceback
        return f"An error occurred: {str(e)}\n{traceback.format_exc()}"


rag_demo = gr.Interface(
    fn=rag_interface,
    inputs=[gr.Textbox(label="Query"), gr.Textbox(label="Author (optional)"), gr.Textbox(label="Category (optional)")],
    outputs="text",
    title="Search for Articles!"
)


rag_demo.launch()