import streamlit as st
import chromadb
from chromadb.utils import embedding_functions
from sentence_transformers import SentenceTransformer

client = chromadb.PersistentClient(path="./chromadb/")
MODEL_NAME: str = "mixedbread-ai/mxbai-embed-large-v1" # ~ 0.5 gb
COLLECTION_NAME: str = "scheme"
EMBEDDING_FUNC = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=MODEL_NAME)
schemer = client.get_collection(
    name=COLLECTION_NAME,
    embedding_function=EMBEDDING_FUNC,
)
DATA_AVAL: bool = schemer.count() > 0
APP_NAME: str = "Groove-GPT"

st.title(APP_NAME)
st.write("Data Avaliable: ", DATA_AVAL)
user_question: str = st.text_area("Enter your groovy questions here")
access_key: str = st.text_input("Enter your gpt key here", type="password")

if st.button('Query Database'):
    if (user_question == ""):
        st.stop()
    else:
        st.header("Results")
        # Perform the Chromadb query.
        results = schemer.query(
            query_texts=[user_question],
            n_results=10,
            include = ['documents']
        )
        st.write(results["documents"])