# Import necessary libraries
import gradio as gr
from haystack.document_stores import InMemoryDocumentStore
from haystack.nodes import DensePassageRetriever, FARMReader
from haystack.pipelines import ExtractiveQAPipeline

# 1. Initialize Document Store
document_store = InMemoryDocumentStore(embedding_dim=768)

# 2. Add Documents
documents = [
    {"content": "Haystack is an open-source NLP framework for search.", "meta": {"source": "Introduction"}},
    {"content": "You can use Hugging Face models in Haystack pipelines.", "meta": {"source": "Hugging Face"}},
    {"content": "The DensePassageRetriever is a key component of Haystack.", "meta": {"source": "Retrievers"}}
]
document_store.write_documents(documents)

# 3. Set up Retriever
retriever = DensePassageRetriever(
    document_store=document_store,
    query_embedding_model="facebook/dpr-question_encoder-single-nq-base",
    passage_embedding_model="facebook/dpr-ctx_encoder-single-nq-base"
)
document_store.update_embeddings(retriever)

# 4. Set up Reader
reader = FARMReader(model_name_or_path="deepset/roberta-base-squad2", use_gpu=False)

# 5. Create QA Pipeline
qa_pipeline = ExtractiveQAPipeline(reader=reader, retriever=retriever)

# 6. Define Prediction Function
def ask_question(query):
    results = qa_pipeline.run(query=query, params={"Retriever": {"top_k": 3}, "Reader": {"top_k": 1}})
    if results["answers"]:
        answer = results["answers"][0].answer
        context = results["answers"][0].context
        source = results["answers"][0].meta.get("source", "Unknown Source")
        return f"**Answer:** {answer}\n\n**Context:** {context}\n\n**Source:** {source}"
    else:
        return "No relevant answer found. Please refine your question."

# 7. Set up Gradio Interface
interface = gr.Interface(
    fn=ask_question,
    inputs=gr.Textbox(lines=2, label="Ask a Question"),
    outputs="text",
    title="AI Search with Haystack",
    description="Ask any question about the content in the document set."
)

# 8. Launch Application
if __name__ == "__main__":
    interface.launch()