import streamlit as st
from app import client, default_embedding_function
import pandas as pd
from generate_kb import generate_knowledge_box_from_url
from utils import get_chroma_client

# Title of the app
st.title("Create a knowledge box from CSV file")

# File uploader widget
uploaded_file = st.file_uploader("Choose a CSV file", type=["csv"])
df = None

collection_name = st.text_input(label="empy collection name")
if st.button("create empty knowledge box"):
    client = get_chroma_client()
    collection = client.create_collection(name=collection_name)
    st.success("collection created")
    st.write(collection)


if uploaded_file is not None:
    try:
        df = pd.read_csv(uploaded_file)
        st.write("DataFrame:")
        st.write(df)
    except Exception as e:
        st.error(str(e))


if uploaded_file is not None:
    st.text("dont use spaces but underscores _ in your new name")
    kb_name = st.text_input(label="new knowledge base name")
    if st.button("Generate new knowledge box"):
        urls = df.values.tolist()
        res = generate_knowledge_box_from_url(
            client=client,
            urls=urls,
            kb_name=kb_name,
            embedding_fct=default_embedding_function,
            chunk_size=2_000,
        )
        st.json(res)