import gradio as gr
from langchain.agents.openai_assistant import OpenAIAssistantRunnable
from openai import OpenAI
import os
#sk
assit_public =  os.getenv('assit_public')
Secret =  os.getenv('Secret')
assit_Public_to_Compliance = os.getenv('assit_Public_to_Compliance')
Confidential_HR_Director = os.getenv('Confidential_HR_Director') 
apikey =os.getenv('openAI') 
#apikey = os.getenv('openAI')  #key from client
#agentkey = os.getenv('asstID')
os.environ["OPENAI_API_KEY"] = apikey
# assit_public 
# Secret 
# assit_Public_to_Compliance 
# Confidential_HR_Director
def get_answer(message, history, additional_input):
    if additional_input == "Public":
        assist = assit_public   
    elif additional_input == "Secret":
        assist = Secret   
    elif additional_input == "Public_to_Compliance":
        assist = assit_Public_to_Compliance   
    elif additional_input == "HR_Director":
        assist = Confidential_HR_Director  
    else:
        return "Invalid knowledge base selected."

    # # Use the selected tool to get the answer
    # response = agent.chat(question) #, tools=[selected_tool]
    interpreter_assistant = OpenAIAssistantRunnable(assistant_id= assist)
    output = interpreter_assistant.invoke({"content": message})
    response = output[0].content[0].text.value
    return response
css = """
label[data-testid="block-label"] {
    display: none !important;
}
footer {
    display: none !important;
}
"""

js_func = """
function refresh() {
    const url = new URL(window.location);
    if (url.searchParams.get('__theme') !== 'dark') {
        url.searchParams.set('__theme', 'dark');
        window.location.href = url.href;
    }
}
"""
with gr.Blocks(css=css, js = js_func, theme="soft") as demo:
    chatbot = gr.Chatbot(placeholder="<strong>Your Personal AI</strong><br>Ask Me Anything")
    additional_input = gr.Dropdown(choices=["Public", "Secret", "Public_to_Compliance","HR_Director"], label="Select Knowledge Base")

    gr.ChatInterface(
        fn=get_answer,
        type="messages",
        chatbot=chatbot,
        additional_inputs=[additional_input]  # Add the additional input to the ChatInterface
    )

demo.launch(show_error=True,debug=True)