Update chatbot.py

chatbot.py

@@ -1,52 +1,49 @@
 import os
 import requests
 from dotenv import load_dotenv
+from messages import krishna_blessings
+from ayush_messages import ayush_surprises
 
 # Load environment variables (Hugging Face Space secrets)
 load_dotenv()
 HUGGINGFACE_API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")
 
-def generate_krishna_image(prompt, size="512x512"):
-    try:
-        headers = {
-            "Authorization": f"Bearer {HUGGINGFACE_API_TOKEN}",
-            "Content-Type": "application/json"
-        }
-        payload = {
-            "inputs": f"{prompt}, vibrant colors, cartoon style, Little Krishna theme with peacock feathers and Vrindavan background",
-            "parameters": {
-                "num_inference_steps": 50,
-                "guidance_scale": 7.5
-            }
+def get_krishna_response(user_input):
+    """
+    Generate a response from Little Krishna based on user input.
+    - If the input contains specific keywords, return predefined messages.
+    - Otherwise, use Hugging Face's Inference API with the Gemma model.
+    """
+    if "birthday" in user_input.lower():
+        return ayush_surprises["birthday"]
+    if "shy" in user_input.lower():
+        return krishna_blessings["shy"]
+    if "hello" in user_input.lower() or "hi" in user_input.lower():
+        return krishna_blessings["greeting"]
+
+    headers = {
+        "Authorization": f"Bearer {HUGGINGFACE_API_TOKEN}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "inputs": f"You are Little Krishna, playful and wise. Respond to this in a fun, Krishna-like way: '{user_input}'",
+        "parameters": {
+            "max_length": 60,
+            "temperature": 0.9
         }
+    }
+    try:
         response = requests.post(
-            "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5",
+            "https://api-inference.huggingface.co/models/google/gemma-2b",
             headers=headers,
             json=payload
         )
         if response.status_code == 200:
-            with open("temp_image.png", "wb") as f:
-                f.write(response.content)
-            return "https://via.placeholder.com/512x512.png?text=Stable+Diffusion+Image"
+            result = response.json()
+            return result[0]["generated_text"].strip()
         else:
-            print(f"Error generating image: {response.text}")
-            return None
+            print(f"Error with Gemma: {response.text}")
+            return "Sorry, I couldn’t generate a response. Let’s try something else!"
     except Exception as e:
-        print(f"Error generating image: {str(e)}")
-        return None
-
-def generate_comic_strip():
-    panels = [
-        "Little Krishna seeing a shy girl reading alone in Vrindavan",
-        "Little Krishna appearing with a playful smile, holding a flute",
-        "Little Krishna handing the girl a book, surrounded by cows",
-        "Little Krishna and the girl smiling together, with peacocks dancing"
-    ]
-    comic_images = []
-    for panel_prompt in panels:
-        image_url = generate_krishna_image(panel_prompt)
-        if image_url:
-            comic_images.append(image_url)
-        else:
-            comic_images.append("https://via.placeholder.com/512x512.png?text=Error+Generating+Image")
-    return comic_images
+        print(f"Error connecting to Hugging Face Inference API: {str(e)}")
+        return "Hare Krishna! I seem to be lost in Vrindavan’s magic—let’s try again!"