Update app.py

app.py

@@ -1,62 +1,48 @@
 # app.py
 import gradio as gr
 import spaces
-from transformers import pipeline
-import os
-#from PyPDF2 import PdfReader
-from dotenv import load_dotenv
-
-load_dotenv()
+from agent import generate_response
 
 @spaces.GPU
-class ChatBot:
-    def __init__(self):
-        self.llm = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.1", token=os.getenv("HF_TOKEN"))
-        self.context = ""
-
-    def read_meal_plans(self, folder="meal_plans"):
-        text = ""
-        for file in os.listdir(folder):
-            if file.endswith(".pdf"):
-                reader = PdfReader(os.path.join(folder, file))
-                for page in reader.pages:
-                    text += page.extract_text() + "\n"
-        return text
-
-    def reply(self, message, history, preferences):
-        diet, goal, allergens = preferences
-        if not self.context:
-            mealplan_text = self.read_meal_plans()
-            self.context = f"Meal Plans: {mealplan_text}\nUser Preferences: Diet={diet}, Goal={goal}, Allergens={allergens}"
-
-        prompt = f"{self.context}\nUser: {message}\nAI:"
-        response = self.llm(prompt, max_new_tokens=100, do_sample=True, temperature=0.7)[0]['generated_text'].split("AI:")[-1].strip()
-        return response
-
-bot = ChatBot()
-
-def chat(message, history, diet, goal, allergens):
-    return bot.reply(message, history, (diet, goal, allergens))
-
-diet_choices = ["Vegetarian", "Vegan", "Keto", "Paleo", "No Preference"]
-goal_choices = ["Weight Loss", "Muscle Gain", "Maintenance"]
-allergen_choices = ["Nuts", "Dairy", "Gluten", "Soy", "Eggs"]
-
-with gr.Blocks() as demo:
-    gr.Markdown("# 🥗 AI Meal Plan Assistant")
-    with gr.Row():
-        diet = gr.Dropdown(diet_choices, label="Diet Type")
-        goal = gr.Dropdown(goal_choices, label="Goal")
-        allergens = gr.CheckboxGroup(allergen_choices, label="Allergies")
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(placeholder="Ask me for meal ideas...", label="Message")
-    send = gr.Button("Send")
-
-    def user_input(message, chat_history):
-        response = chat(message, chat_history, diet.value, goal.value, allergens.value)
-        chat_history.append((message, response))
-        return chat_history, ""
-
-    send.click(user_input, [msg, chatbot], [chatbot, msg])
-
-demo.launch()
+def main():
+    with gr.Blocks(title="Meal Plan Chatbot") as demo:
+        gr.Markdown("# 🍎 MealBot: Your AI Nutritionist")
+
+        with gr.Row():
+            diet_type = gr.CheckboxGroup(
+                label="Diet Type",
+                choices=["Vegan", "Keto", "Vegetarian", "Paleo", "High-Protein"],
+                value=["Vegetarian"]
+            )
+            goal = gr.CheckboxGroup(
+                label="Your Goal",
+                choices=["Weight Loss", "Muscle Gain", "Maintenance"],
+                value=["Weight Loss"]
+            )
+            duration = gr.CheckboxGroup(
+                label="Duration",
+                choices=["1 week", "4 weeks", "8 weeks"],
+                value=["4 weeks"]
+            )
+
+        chatbot = gr.Chatbot()
+        msg = gr.Textbox(placeholder="Ask me anything about meal plans...")
+        clear = gr.Button("Clear")
+
+        def user_respond(message, chat_history):
+            prefs = {
+                "diet": diet_type.value,
+                "goal": goal.value,
+                "duration": duration.value
+            }
+            response = generate_response(message, prefs)
+            chat_history.append((message, response))
+            return "", chat_history
+
+        msg.submit(user_respond, [msg, chatbot], [msg, chatbot])
+        clear.click(lambda: None, None, chatbot, queue=False)
+
+    demo.launch()
+
+if __name__ == "__main__":
+    main()