import gradio as gr
import requests
import os

# Load API settings from environment variables
OPENWEBUI_URL = os.getenv("OPENWEBUI_URL")
OPENWEBUI_API_KEY = os.getenv("OPENWEBUI_API_KEY")
MODEL_NAME = os.getenv("MODEL_NAME", "phi3:latest")  # Default to "phi3:latest" if not set


# Function to send messages to OpenWebUI
def chat_with_model(message, history):
    if not OPENWEBUI_URL or not OPENWEBUI_API_KEY:
        return history + [("System", "Error: Missing API credentials. Please set environment variables.")]

    headers = {
        "Authorization": f"Bearer {OPENWEBUI_API_KEY}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "model": MODEL_NAME,
        "messages": [{"role": "user", "content": message}]
    }
    
    response = requests.post(OPENWEBUI_URL, json=payload, headers=headers)
    
    if response.status_code == 200:
        bot_reply = response.json().get("choices", [{}])[0].get("message", {}).get("content", "No response.")
    else:
        bot_reply = f"Error: {response.text}"
    
    history.append((message, bot_reply))
    return history

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## Chat with Ollama via OpenWebUI")
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Your Message")
    submit = gr.Button("Send")

    submit.click(chat_with_model, inputs=[msg, chatbot], outputs=chatbot)

demo.launch()