import gradio as gr
import asyncio
import ollama
import json

def chat(message):
    messages = [{"role": "user", "content": message}]
    response = ollama.chat(model='llama3', messages=messages, stream=True)
    for chunk in response:
        yield chunk['message']['content']

def gradio_chat(message, history):
    response = ""
    for chunk in chat(message):
        response += chunk
        yield response

iface = gr.ChatInterface(
    fn=gradio_chat,
    title="Diet Plan Generator",
    description="Enter your preferences to generate a personalized diet plan.",
    examples=[
        "Generate a meal plan for weight loss with a daily activity level being no activity or very less exercise. Target calories: 2000 kcal. Macro Distribution: 80 g protein, 220 g of carbs and 32 g of fat. The foods should mainly belong to Indian cuisine and should be strictly vegetarian.",
        "Create a high-protein diet plan for muscle gain. Target calories: 2500 kcal. Macro Distribution: 150 g protein, 250 g of carbs and 70 g of fat. Include Mediterranean cuisine options.",
    ],
    retry_btn=None,
    undo_btn="Delete Previous",
    clear_btn="Clear",
)

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