import gradio as gr
import os
import openai
import chromadb
from chromadb.config import Settings
from langchain.embeddings import OpenAIEmbeddings


css="""
#col-container {max-width: 700px; margin-left: auto; margin-right: auto;}
"""

title = """
<div style="text-align: center;max-width: 700px;">
    <h1>Ask P360 Help • OpenAI</h1>
    <p style="text-align: center;">Ask questions about the P360 documentation.
</div>
"""

chroma_client = chromadb.Client(
    Settings(
        chroma_db_impl='duckdb+parquet',
        persist_directory='./vectorstore/',
    )
)

db = chroma_client.get_collection(
    name='help-360',
)

embeddings_generator = OpenAIEmbeddings(
    openai_api_key=os.getenv('AZURE_OPENAI_KEY'),
    openai_api_base=os.getenv('AZURE_OPENAI_ENDPOINT'),
    openai_api_type='azure',
    openai_api_version='2023-05-15',
    deployment='embedding-test'
)


def respond(query):
    '''
    OpenAI GPT response.
    '''
    query_embedding = embeddings_generator.embed_query(query)
    results = db.query(
        query_embeddings=query_embedding,
        n_results=1,
    )
    relevant_help = results['documents'][0][0]

    openai.api_type = 'azure'
    openai.api_base = os.getenv('AZURE_OPENAI_ENDPOINT')
    openai.api_key = os.getenv('AZURE_OPENAI_KEY')
    openai.api_version = '2023-05-15'

    response = openai.ChatCompletion.create(
        engine='entest-gpt-35-turbo',
        messages=[
            {"role": "system", "content": "You are a helpful assistant in the P360 document archiving system. Only provide helpful answers if you have the necessary information, otherwise answer with I have no information about that."},
            {"role": "user", "content": "What do you know about P360?"},
            {"role": "assistant", "content": relevant_help},
            {"role": "user", "content": query},
        ]
    )
    return response['choices'][0]['message']['content']


with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML(title)        
        question = gr.Textbox(label='Question', placeholder='Type your question and hit Enter ')
        answer = gr.Textbox(label='Answer').style(height=350)
        question.submit(respond, [question], [answer])


demo.launch()