import gradio as gr
from langchain.chains import RetrievalQA
from langchain.embeddings import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.vectorstores import Qdrant
from openai.error import InvalidRequestError
from qdrant_client import QdrantClient
from config import DB_CONFIG


PERSIST_DIR_NAME = "nvdajp-book"


def get_retrieval_qa() -> RetrievalQA:
    embeddings = OpenAIEmbeddings()
    db_url, db_api_key, db_collection_name = DB_CONFIG
    client = QdrantClient(url=db_url, api_key=db_api_key)
    db = Qdrant(client=client, collection_name=db_collection_name, embeddings=embeddings)
    retriever = db.as_retriever()
    return RetrievalQA.from_chain_type(
        llm=OpenAI(temperature=0), chain_type="stuff", retriever=retriever, return_source_documents=True,
    )


def get_related_url(metadata):
    urls = set()
    for m in metadata:
        # p = m['source']
        url = m["url"]
        if url in urls:
            continue
        urls.add(url)
        category = m["category"]
        # print(m)
        yield f'<p>URL: <a href="{url}">{url}</a> (category: {category})</p>'


def main(query: str):
    qa = get_retrieval_qa()
    try:
        result = qa(query)
    except InvalidRequestError as e:
        return "回答が見つかりませんでした。別な質問をしてみてください", str(e)
    else:
        metadata = [s.metadata for s in result["source_documents"]]
        html = "<div>" + "\n".join(get_related_url(metadata)) + "</div>"

    return result["result"], html


nvdajp_book_qa = gr.Interface(
    fn=main,
    inputs=[gr.Textbox(label="query")],
    outputs=[gr.Textbox(label="answer"), gr.outputs.HTML()],
)


nvdajp_book_qa.launch()