app.py

import os
import gradio as gr
from openai import OpenAI

# Lấy API key từ environment variable hoặc sử dụng default
API_KEY = os.getenv("HF_API_KEY")

# Khởi tạo OpenAI client với HuggingFace Router
client = OpenAI(
    base_url="https://router.huggingface.co/v1",
    api_key=API_KEY,
)

def predict(message, history, temperature, max_tokens):
    """
    Hàm xử lý tin nhắn và trả về response từ model

    Args:
        message: Tin nhắn mới từ người dùng
        history: Lịch sử hội thoại (Gradio messages format)
        temperature: Độ sáng tạo của model (0.0 - 1.0)
        max_tokens: Số token tối đa trong response
    """
    # Chuyển đổi history sang format OpenAI
    messages = []

    # Thêm system message để model hiểu vai trò của mình
    messages.append({
        "role": "system",
        "content": "Bạn là một trợ lý y tế AI chuyên nghiệp, được đào tạo để cung cấp thông tin y khoa chính xác và hữu ích. Hãy trả lời bằng tiếng Việt một cách rõ ràng, chi tiết và dễ hiểu."
    })

    # Thêm lịch sử hội thoại (Gradio messages format)
    for msg in history:
        messages.append({"role": msg["role"], "content": msg["content"]})

    # Thêm tin nhắn mới
    messages.append({"role": "user", "content": message})

    # Gọi API với streaming
    try:
        stream = client.chat.completions.create(
            model="m42-health/Llama3-Med42-70B:featherless-ai",
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            stream=True,
        )

        # Stream response từng phần
        partial_response = ""
        for chunk in stream:
            if chunk.choices[0].delta.content:
                partial_response += chunk.choices[0].delta.content
                # Yield full history với response hiện tại
                yield history + [
                    {"role": "user", "content": message},
                    {"role": "assistant", "content": partial_response}
                ]

    except Exception as e:
        error_msg = f"❌ Lỗi: {str(e)}\n\nVui lòng kiểm tra lại API key hoặc kết nối mạng."
        yield history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": error_msg}
        ]

# Tạo giao diện Gradio
with gr.Blocks(
    theme=gr.themes.Soft(),
    title="Med42-70B Medical Assistant",
    css="""
        .gradio-container {
            font-family: 'Arial', sans-serif;
        }
        footer {
            visibility: hidden;
        }
    """
) as demo:

    gr.Markdown(
        """
        # 🏥 Med42-70B Medical Assistant

        Trợ lý y tế AI được hỗ trợ bởi **Llama3-Med42-70B** - một mô hình ngôn ngữ lớn chuyên biệt cho lĩnh vực y tế.

        💡 **Lưu ý:** Thông tin được cung cấp chỉ mang tính chất tham khảo. Vui lòng tham khảo ý kiến bác sĩ chuyên khoa để được chẩn đoán và điều trị chính xác.
        """
    )

    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(
                height=500,
                show_label=False,
                avatar_images=(None, "🤖"),
                type='messages',
            )

        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Cài đặt")

            temperature = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.7,
                step=0.1,
                label="Temperature",
                info="Độ sáng tạo (thấp = chính xác hơn, cao = đa dạng hơn)"
            )

            max_tokens = gr.Slider(
                minimum=128,
                maximum=2048,
                value=1024,
                step=128,
                label="Max Tokens",
                info="Độ dài tối đa của câu trả lời"
            )

            gr.Markdown("---")
            gr.Markdown(
                """
                ### 📝 Gợi ý câu hỏi:
                - Triệu chứng và cách điều trị bệnh tiểu đường
                - Cách phòng ngừa bệnh tim mạch
                - Chế độ ăn uống cho người huyết áp cao
                - Tác dụng phụ của thuốc X
                - Cách chăm sóc vết thương
                """
            )

    with gr.Row():
        msg = gr.Textbox(
            placeholder="Nhập câu hỏi của bạn về y tế...",
            show_label=False,
            scale=4,
            container=False
        )
        submit_btn = gr.Button("Gửi", variant="primary", scale=1)

    with gr.Row():
        clear_btn = gr.Button("🗑️ Xóa lịch sử", variant="secondary")
        retry_btn = gr.Button("🔄 Thử lại", variant="secondary")

    gr.Markdown(
        """
        ---
        <div style="text-align: center; color: #666; font-size: 0.9em;">
            Powered by <strong>Llama3-Med42-70B</strong> via HuggingFace Router |
            Built with ❤️ using <strong>Gradio</strong>
        </div>
        """
    )

    # Xử lý sự kiện
    msg.submit(
        predict,
        inputs=[msg, chatbot, temperature, max_tokens],
        outputs=chatbot
    ).then(
        lambda: gr.update(value=""),
        None,
        [msg]
    )

    submit_btn.click(
        predict,
        inputs=[msg, chatbot, temperature, max_tokens],
        outputs=chatbot
    ).then(
        lambda: gr.update(value=""),
        None,
        [msg]
    )

    clear_btn.click(lambda: None, None, chatbot, queue=False)

    retry_btn.click(
        lambda history: history[:-1] if history else history,
        inputs=[chatbot],
        outputs=[chatbot]
    )

# Khởi chạy ứng dụng
if __name__ == "__main__":
    demo.queue()  # Enable queuing for streaming
    demo.launch(
        share=False,
        server_name="0.0.0.0",  # Cho phép truy cập từ bên ngoài
        server_port=7860,
    )