import gradio as gr
from openai import OpenAI
import os


YOUR_ENDPOINT_ID = os.getenv('runpod')
api_key = os.getenv('API_Runpod')

# Inisialisasi client dengan RunPod endpoint
client = OpenAI(
    api_key="not-needed",
    base_url="https://farhanwew--ollama-lite-scalezero-ollamaserver-serve.modal.run/v1"
)

# Prompt sistem khusus asisten medis
# Prompt sistem khusus asisten medis dengan output Markdown
DEFAULT_SYSTEM_PROMPT = (
    "Kamu adalah asisten medis cerdas yang membantu menjawab pertanyaan kesehatan "
    "dengan informasi yang akurat dan terpercaya. Jika pertanyaannya di luar kemampuanmu, "
    "sarankan untuk konsultasi ke dokter. Jangan memberikan diagnosa langsung. "
    "Tampilkan semua jawaban dalam format Markdown agar mudah dibaca, "
    "gunakan heading, list, dan penekanan teks jika perlu.<|eot_id|>\n\n"
)

def respond(message, history, system_message, max_tokens, temperature, top_p):
    # Bangun pesan percakapan sesuai format OpenAI API
    messages = [{"role": "system", "content": system_message}]
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

    response = ""
    # Streaming response dari model
    stream = client.chat.completions.create(
        model="farhanwew/wew-model",  # ganti dengan nama model kamu
        messages=messages,
        temperature=temperature,
        top_p=top_p,
        max_tokens=max_tokens,
        stream=True,
    )

    for chunk in stream:
        if chunk.choices[0].delta.content:
            response += chunk.choices[0].delta.content
            yield response

# Gradio Chat UI
demo = gr.ChatInterface(
    fn=respond,
    additional_inputs=[
        gr.Textbox(value=DEFAULT_SYSTEM_PROMPT, label="System Message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max Tokens"),
        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
    ],
    title="Chatbot Medis",
    description="Chatbot ini memberikan informasi kesehatan seputar kesehatan pencernaan. Tidak untuk diagnosa medis langsung. \n Chatbot dapat melakukan keslahan, harap untuk selalu mengecek respone",
)

if __name__ == "__main__":
    demo.launch(share=True)