import gradio as gr
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from groq import Groq

# Load FAISS index
vector_store = FAISS.load_local("faiss_index", HuggingFaceEmbeddings())

# Inisialisasi API Groq
client = Groq(api_key="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")

def retrieve_and_generate(query):
    # Retrieve top 3 documents
    docs = vector_store.similarity_search(query, k=3)
    context = "\n\n".join([doc.page_content for doc in docs])

    # Generate response with LLM
    response = client.chat.completions.create(
        model="mixtral-8x7b-32768",
        messages=[
            {"role": "system", "content": "Anda adalah asisten AI yang menjawab pertanyaan tentang RoboHome berdasarkan dokumen ini."},
            {"role": "user", "content": f"{context}\n\nPertanyaan: {query}"}
        ],
        temperature=0.7,
        max_tokens=200
    )

    return response.choices[0].message.content

# UI dengan Gradio
iface = gr.Interface(
    fn=retrieve_and_generate,
    inputs=gr.Textbox(label="Ajukan pertanyaan tentang RoboHome"),
    outputs=gr.Textbox(label="Jawaban"),
    title="RoboHome RAG Chatbot",
    description="Chatbot ini menjawab pertanyaan berdasarkan dokumentasi RoboHome.",
)

iface.launch()