import gradio as gr
import chromadb
import pandas as pd

client = chromadb.Client()
collection = client.create_collection("bolivian-recipes")
df = pd.read_parquet("hf://datasets/asoria/bolivian-recipes@~parquet/default/other/0000.parquet")
text_column = "preparation"
ids = [str(i) for i in range(df.shape[0])]
documents = df[text_column].to_list()
metadatas = df.drop(text_column, axis=1).to_dict("records")
collection.add(ids=ids, documents=documents, metadatas=metadatas)


with gr.Blocks() as demo:
    gr.Markdown(" ## PandasAI demo using datasets library")
    gr.Markdown(" pandasai library https://github.com/gventuri/pandas-ai")
    gr.Markdown(" datasets library https://huggingface.co/docs/datasets")

    dataset = gr.Textbox(label="dataset", placeholder="mstz/iris", value="mstz/iris")
    config = gr.Textbox(label="config", placeholder="iris", value="iris")
    split = gr.Textbox(label="split", placeholder="train", value="train")
    prompt = gr.Textbox(label="prompt (str)", placeholder="How many records do I have?. Give me the median of sepal_width. Show me the first 5 rows.")
    cached_responses_table = gr.DataFrame()
    get_result_button = gr.Button("Submit")

    def get_result(dataset, config, split, prompt) -> str:
        result = collection.query(query_texts=[prompt], n_results=4)
        return {
                    cached_responses_table: gr.update(value=result)
        }

    get_result_button.click(get_result, inputs=[dataset, config, split, prompt], outputs=cached_responses_table)

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