import requests
import gradio as gr
from openai import OpenAI

client = OpenAI(
    base_url = 'http://localhost:7869/v1',
    api_key='ollama', # required, but unused
)

chat_history = []

def get_models():
    response = requests.get("http://localhost:7869/api/tags")
    models = response.json()["models"]
    return [model["name"] for model in models]

def ask(text, model):
    global chat_history
    chat_history.append({"role": "user", "content": text})
    response = client.chat.completions.create(
        model=model,
        messages=chat_history
    )
    response_text = response.choices[0].message.content
    chat_history.append({"role": "assistant", "content": response_text})
    formatted_history = "\n".join([f"**{msg['role'].capitalize()}**: {msg['content']}" for msg in chat_history])
    print("Formatted History: ", formatted_history)
    return formatted_history, ""

def clear_history(model):
    global chat_history
    chat_history = []
    # Reset the conversation on the server side
    client.chat.completions.create(
        model=model,
        messages=[{"role": "system", "content": "Reset conversation"}]
    )
    return ""
    

with gr.Blocks() as server:
    with gr.Tab("Chatbot LF6"):
        model_dropdown = gr.Dropdown(label="Select Model", choices=get_models())
        chatbox = gr.Textbox(label="Chat History", lines=10, interactive=False)
        input_text = gr.Textbox(label="Input Text", placeholder="Enter your text here...")
        submit_button = gr.Button("Submit")
        clear_button = gr.Button("Clear")
        
        submit_button.click(fn=lambda text, model: ask(text, model), inputs=[input_text, model_dropdown], outputs=[chatbox, input_text])
        input_text.submit(fn=lambda text, model: ask(text, model), inputs=[input_text, model_dropdown], outputs=[chatbox, input_text])  # Submit on Enter key press
        clear_button.click(fn=clear_history, inputs=model_dropdown, outputs=chatbox)

server.launch()