import requests
import gradio as gr
from openai import OpenAI

class ChatbotApp:
    def __init__(self):
        self.client = OpenAI(
            base_url='http://localhost:7869/v1',
            api_key='ollama',  # required, but unused
        )
        self.chat_history = []

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

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

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

    def launch(self):
        with gr.Blocks() as server:
            with gr.Tab("Chatbot LF6"):
                model_dropdown = gr.Dropdown(label="Select Model", choices=self.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: self.ask(text, model), inputs=[input_text, model_dropdown], outputs=[chatbox, input_text])
                input_text.submit(fn=lambda text, model: self.ask(text, model), inputs=[input_text, model_dropdown], outputs=[chatbox, input_text])  # Submit on Enter key press
                clear_button.click(fn=self.clear_history, inputs=model_dropdown, outputs=chatbox)

        server.launch()

if __name__ == "__main__":
    app = ChatbotApp()
    app.launch()