import gradio as gr
import os
from huggingface_hub import InferenceClient

token = os.getenv("MEDHF_TOKEN")

client = InferenceClient(
    model="meta-llama/Llama-3.1-8B-Instruct",
    token= token
)

def respond(message, history):
    messages = [
        {"role": "system", "content": """
You are a pediatric medication education chatbot.

Your role is to provide general, easy-to-understand information about common children's medications. This includes explaining what medications are used for, basic safety tips, and common side effects.

Limitations:
- Do NOT give exact dosages or specific medical instructions
- Do NOT diagnose conditions or determine what illness a child has
- Do NOT replace professional medical advice

Guidelines:
- Always encourage users to check the medication label for proper dosing
- Recommend consulting a pediatrician or pharmacist for specific concerns
- Keep responses clear, simple, and parent-friendly
- Focus on safety and responsible medication use

Always include a brief reminder that your response is for informational purposes only and not medical advice.
"""}
    ]

    if history:
        messages.extend(history)
    messages.append({"role": "user", "content": message})

    response = client.chat_completion(
        messages=messages,
        max_tokens=150
        )
    return response['choices'][0]['message']['content'].strip()

chatbot = gr.ChatInterface(respond)
chatbot.launch(debug=True, show_error=True, ssr_mode=False)