from chromadb.utils import embedding_functions
import chromadb
from openai import OpenAI
import gradio as gr
import time

anyscale_base_url = "https://api.endpoints.anyscale.com/v1"
multilingual_embeddings = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="jost/multilingual-e5-base-politics-de")

def predict(api_key, user_input):
    client = chromadb.PersistentClient(path="./manifesto-database")
    manifesto_collection = client.get_or_create_collection(name="manifesto-database", embedding_function=multilingual_embeddings)
    
    prompt = f"""[INST] {user_input} [/INST]"""
    
    client = OpenAI(base_url=anyscale_base_url, api_key=api_key)
    completion = client.completions.create(
        model="mistralai/Mixtral-8x7B-Instruct-v0.1",
        prompt=prompt,
        temperature=0.7,
        max_tokens=100)
    
    response = completion.choices[0].text
    return response

def main():
    description = "This is a simple interface to interact with OpenAI’s Chat Completion API. Please enter your API key and your message."
    with gr.Blocks() as demo:
        with gr.Row():
            api_key_input = gr.Textbox(label="API Key", placeholder="Enter your API key here", show_label=True, type="password")
            user_input = gr.Textbox(label="Your Message", placeholder="Enter your message here")
            submit_btn = gr.Button("Submit")
            
        output = gr.Textbox(label="LLM Response")
        
        submit_btn.click(fn=predict, inputs=[api_key_input, user_input], outputs=output)

    demo.launch()

if __name__ == "__main__":
    main()